Learning When to Stop: Adaptive Latent Reasoning via Reinforcement Learning
This work addresses the challenge of computational efficiency in AI reasoning for applications requiring faster or cheaper inference, though it is incremental as it builds on existing latent reasoning methods.
The paper tackled the problem of reducing reasoning length in Transformer language models by developing adaptive-length latent reasoning models with a reinforcement-learning method, achieving a 52% drop in total reasoning length without accuracy loss on the GSM8K-Aug dataset.
Latent reasoning represents a new development in Transformer language models 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.