Learning When to Stop: Adaptive Latent Reasoning via Reinforcement Learning (2511.21581v1)

Abstract: Latent reasoning represents a new development in Transformer LLMs that has shown potential in compressing reasoning lengths compared to chain-of-thought reasoning. By directly passing the information-rich previous final latent state into the next sequence, latent reasoning removes the restriction to human language tokens as the medium for reasoning. We develop adaptive-length latent reasoning models and introduce a post-SFT reinforcement-learning methodology to optimize latent reasoning length by minimizing reasoning length while maintaining accuracy. This, in turn, further reduces compute usage and raises the bar on the compressive capabilities of latent reasoning models. Experiments on the Llama 3.2 1B model and the GSM8K-Aug dataset show a $52\%$ drop in total reasoning length with no penalty to accuracy. In future work, we plan to extend to additional models and datasets, analyze relationships between training coefficients, experiment with architecture variations, and continue our knowledge distillation for latent reasoning SFT efforts. We make our code and pretrained weights available at https://github.com/apning/adaptive-latent-reasoning.

• The paper introduces a novel RL-based framework that adaptively determines latent reasoning steps according to question complexity.
• It leverages a two-stage training process—supervised fine-tuning followed by Group Relative Policy Optimization—to reduce reasoning tokens by up to 61.99%.
• The approach dynamically prevents overthinking on simple problems while allocating robust compute for complex queries, enhancing efficiency and reliability.

Adaptive Latent Reasoning in LLMs via Reinforcement Learning

Introduction and Motivation

Latent reasoning has recently emerged as a notable alternative to chain-of-thought (CoT) reasoning in LLMs, addressing the inefficiencies inherent in token-level, language-constrained intermediate reasoning. The paper "Learning When to Stop: Adaptive Latent Reasoning via Reinforcement Learning" (2511.21581) investigates whether models can be trained to self-terminate latent reasoning steps adaptively per-query, leveraging reinforcement learning (RL) to optimize the number of reasoning steps needed without sacrificing accuracy.

Standard CoT approaches—despite widespread empirical success—overcommit compute by mandating uniform, language-token-based step sequences that are often unnecessarily long and ill-aligned with question difficulty. Latent reasoning sidesteps this limitation by transforming the internal hidden state between steps, bypassing tokenization. Prior work such as COCONUT (Hao et al., 9 Dec 2024) and CODI (Shen et al., 28 Feb 2025) established fixed-length latent reasoning, yielding significant compression over CoT but with a static structure and only partial retention of accuracy. The present work innovates by enabling variable, input-adaptive latent step counts, thus allowing for problem-dependent allocation of compute and advancing toward truly efficient neuro-symbolic reasoning in LLMs.

Figure 1: Comparison between chain-of-thought and latent reasoning architectures, with adaptive-endpoint logic highlighted for latent reasoning.

Methodology

Model Architecture

The adaptive latent reasoning framework builds on a Transformer backbone, with a 2-layer MLP recurrent filter (matching CODI's architecture), special <START> and <END> tokens, and introduces a linear binary classification head to signal stopping criteria at each latent step. This architecture is trained in two stages:

• Supervised Fine-Tuning (SFT): Initial latent reasoning capabilities are distilled via the CODI approach, using a large corpus (GSM8K-Aug [whynlp2025gsm8kaug]) with CoT-labeled samples modified to support latent state trajectory supervision.
• Reinforcement Learning: Post-SFT, RL (specifically, Group Relative Policy Optimization [shao2024deepseekmathpushinglimitsmathematical]) fine-tunes the model’s adaptive step controller. The RL objective is non-trivial: it is constructed to minimize steps (for easy samples) and optionally increase steps (for hard samples), but only when accuracy is unaffected or potentially improved.

Reward shaping includes:

• Correctness reward: Binary indication of answer correctness.
• Format penalty: Guards against malformed output.
• Relative length penalty/reward: Penalizes excess steps on easy inputs, rewards longer reasoning when harder inputs demand it, dynamically using group-correctness frequencies for context-sensitive balancing.

The result is a model where, per input, the number of latent thinking steps is decided endogenously, optimizing a compromise between compute and reliability.

Quantitative Results

Comprehensive experiments were conducted on Llama 3.2 1B [meta2024llama321b] with the GSM8K-Aug arithmetic word problem dataset. Key model variants include pure CoT SFT, CoT-free SFT, static-length latent SFT (Latent-6), and their RL-refined counterparts. The most significant findings are:

• Latent reasoning models with RL achieved up to 52.94% to 61.99% reduction in reasoning token count (vs. CoT) without material loss of accuracy.
• For the Latent-6 + RL model, the average number of reasoning tokens plummeted from 8 to 3.76 with a minor accuracy boost (from 49.73% to 50.11%).
• The strictest compression (Latent-6-by-1 + RL) cut 61.99% of steps at a cost of 3.26% accuracy.
• The adaptive latent models, post-RL, consistently avoided overthinking on easy instances but preserved the computational bandwidth for more challenging samples as needed.

  Figure 2: Histogram of reasoning lengths for CoT SFT vs. Latent-6 + RL, illustrating substantial compression in the RL variant.

  

  Figure 3: Histogram of reasoning lengths for Latent-6-by-1 before and after RL, highlighting adaptive shrinkage after RL.

These distributions confirm that adaptive latent reasoning allows the model to "think less" on simple problems and "think more" on complex ones, with reduced standard deviation and overall shift toward efficiency. Crucially, the RL reward was carefully engineered to prevent degeneracies where the model would sacrifice correctness for brevity.

Implications and Theoretical Perspectives

The introduction of adaptive latent reasoning marks a meaningful step in the evolution of computational efficiency for autoregressive models. Unlike traditional static-latent systems, the model dynamically budgets compute, resembling amortized inference. Practically, this means substantial reductions in inference-time resource needs, particularly as dataset and task complexity scale, without a penalty to reliability.

Theoretically, this work reinforces the notion that explicit internal control signals (i.e., a latent "stop" action) can be learned and aligned with high-level task goals via RL. It extends the view of LLMs as modular, controlled dynamical systems, not mere sequential prediction engines. The employment of group-wise RL objectives and careful penalty/reward partitioning underscores the importance of reward design in training these emergent behaviors.

Future directions identified include further scaling to diverse tasks and model sizes, hyperparameter ablation (e.g., impact of $\lambda_\text{penalty}$, $\lambda_\text{reward}$, and $p_\text{cutoff}$), and investigation of alternative latent/architecture designs. The (so far unsuccessful) experiments on Meaned Reasoning Loss and Intermediate Block Loss for knowledge distillation (see Appendix and Figure 4, Figure 5, Figure 6) also suggest ripe ground for continued exploration.

Figure 4: Schematic contrasting CODI knowledge distillation with meaned reasoning loss, outlining different supervision points for latent state alignment.

Figure 5: Illustration of Intermediate Block Loss, applying distillation at transformer middle layers rather than all or final layers.

Figure 6: Test accuracy outcomes for knowledge distillation variants; meaned and intermediate block styles lag conventional CODI.

Conclusion

"Learning When to Stop: Adaptive Latent Reasoning via Reinforcement Learning" (2511.21581) demonstrates that variable-length latent reasoning, when paired with a well-designed RL objective, can both maintain accuracy and drastically cut inference cost relative to CoT and fixed-latent models. The findings highlight a tangible path toward scalable, efficient LLMs capable of input-sensitive compute allocation.

The approach also deepens understanding of controllable neural reasoning trajectories and provides a blueprint for future research in efficient, high-fidelity inference across linguistic and non-linguistic modalities. Further work is necessary to generalize beyond arithmetic QA, optimize RL reward couplings, and refine knowledge distillation protocols for broader adaptability.