让模型自适应决定推理长度,显著减少计算量而不损失准确率。
Learning When to Stop: Adaptive Latent Reasoning via Reinforcement Learning
- 用强化学习训练模型自主判断何时停止推理,优化推理长度。
- 在GSM8K-Aug数据集上推理长度降低52%,准确率不变。
- 适合追求高效推理的AI研究者与部署工程师。
隐式推理是Transformer语言模型的新进展,相比链式思维推理,能在不损失性能的前提下压缩推理长度。通过直接传递前序最终隐状态至下一序列,隐式推理摆脱了以人类语言标记为推理媒介的限制。本文提出自适应长度的隐式推理模型,并引入后SFT强化学习方法,通过最小化推理长度来优化模型表现,从而进一步降低计算开销,提升隐式推理的压缩能力。在Llama 3.2 1B模型和GSM8K-Aug数据集上的实验表明,总推理长度下降52%且准确率无损。未来工作将扩展至更多模型与数据集,分析训练系数关系,尝试架构变体,并持续推进隐式推理的SFT知识蒸馏。代码与预训练权重已开源:https://github.com/apning/adaptive-latent-reasoning。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。