Training large language models to think more clearly and accurately
Principled Thoughts for Latent Recursive LLM Systems
Artificial IntelligenceComputation and LanguageMachine Learning
Summary
Large language models usually learn only to produce the right final answer, but this often causes confusion in their internal thought process. The authors show that training models by focusing just on the final answer can cause problems like mixing up different questions or keeping irrelevant details. They introduce a new training method called REST that guides the model’s internal thoughts to be clear, simple, separate, and stable. This method improves accuracy on tests in math, science, medicine, and coding, while making it easier to understand what the model is thinking.
What this means in practice
- •For software developers: Improve code generation tools by training models to keep more accurate and distinct internal reasoning states for better final outputs.
- •For medical ai teams: Build medical question-answering systems whose intermediate reasoning is less prone to confusion and easier to interpret while increasing accuracy.
Authors
Fahd Seddik, Fatemeh Fard
Abstract
Large language models can reason in continuous space instead of decoded text, by recurring on their own hidden states or by passing those states between agents, while training supervises only the Cross-Entropy (CE) of the final decoded answer and does not constrain the thought. Theoretical and empirical analyses establish and confirm four failures of CE-only training that lead to a lower probability of the correct answer such as collapsing thoughts across distinct questions and retaining irrelevant information. We introduce REST (REpresentation-Supervised Thoughts), a training objective that turns four properties of a valid thought representation (causality, minimality, separability, and stability) into differentiable losses added to CE. We instantiate it in latent single-agent and multi-agent systems, without architectural changes or added parameters at inference. Across 7 benchmarks spanning mathematics, science, medicine, and code generation, with the same training data, compute, and latent budget, REST increases accuracy over CE-only training across agent settings and model sizes by up to 7.5 percentage points and convergence on a final answer by 30\%. Furthermore, REST thoughts encode more of what is required to achieve the correct answer, and decoding them better recovers the intended output of the agent, which makes latent communication easier to interpret. Project Website: https://fard-lab.github.io/REST