Uncovering Latent Reasoning Strategies in Language Models
2026-07-20 • Machine Learning
Machine LearningArtificial Intelligence
AI summaryⓘ
The authors study how a trained language model that uses multiple hidden ways to solve problems can be broken down into clear, distinct strategies. They create a system with two parts: one that picks a strategy based on the input, and another that generates the answer using that strategy. A common issue is that the generator can ignore the strategy part, so the authors design a training method that makes the model focus on differences where the original model was unsure. This helps reveal meaningful strategy codes without changing the model's overall answers. They also introduce tasks to test their method and show it finds strategy codes that match known solving methods.
language modellatent-variable factorizationvariational inferenceposterior collapsestrategy decompositionalgorithmic tasksroutinginformation gaintoken surprisal
Authors
Awni Altabaa, John Lafferty
Abstract
A language model $p_θ(y \mid x)$ trained on reasoning tasks learns to solve problems via multiple distinct strategies, yet these strategies are implicit and entangled within the model's response distribution. We study the problem of decomposing the response distribution of a given pretrained language model into a structured, strategy-conditioned representation. Specifically, we learn a latent-variable factorization $p_θ(y \mid x) \leadsto (r_φ(z \mid x), g_φ(y \mid x,z))$, where a router $r$ maps each input to a distribution over latent strategies $z$ and a generator $g$ produces the response conditioned on that strategy. A key challenge is that the generator, initialized from the base model, already represents $p_θ(y \mid x)$ without using $z$. Standard variational inference therefore gives the model no incentive to route information through $z$ and can yield a severe form of posterior collapse. To address this, we propose a variational objective that measures fractional information gain relative to the base model's response loss and concentrates reconstruction pressure on tokens with high base model surprisal, encouraging $z$ to encode strategy-relevant response variation. We introduce a benchmark of multi-strategy algorithmic tasks and show that this objective recovers latent codes aligned with distinct reference strategies while preserving the base model's response distribution.