让大模型学会提前退出，减少无谓计算，提升推理效率。
A transformer architecture alteration to incentivise externalised reasoning
- 在中间层加入提前退出机制，引导模型浅层推理时尽早停止
- 通过强化学习训练，使模型在保持准确率前提下尽可能早退出
- 适合追求推理效率的轻量化部署场景
我们提出一种新的架构改进和后训练流程，使大语言模型成为更冗长的推理者，方法是教会模型在前向传播中尽早截断。我们在现有Transformer架构中引入中间层的提前退出机制，并训练模型在无需深度计算即可预测下一个词时，从较浅层退出。经过校准阶段后，利用强化学习激励模型尽可能早地退出，同时维持任务性能。初步实验结果显示，小型推理模型能够自适应地减少各词的计算量。我们预计，若在适当规模上应用，该方法可最小化推理模型因内部激活而进行非短期规划所造成的多余计算，仅保留复杂难预测词的深层计算。
原文摘要 · Abstract (English)
We propose a new architectural change, and post-training pipeline, for making LLMs more verbose reasoners by teaching a model to truncate forward passes early. We augment an existing transformer architecture with an early-exit mechanism at intermediate layers and train the model to exit at shallower layers when the next token can be predicted without deep computation. After a calibration stage, we incentivise the model to exit as early as possible while maintaining task performance using reinforcement learning. We provide preliminary results to this effect for small reasoning models, showing that they learn to adaptively reduce computations across tokens. We predict that, applied at the right scale, our approach can minimise the amount of excess computation that reasoning models have at their disposal to perform non-myopic planning using their internal activations, reserving this only for difficult-to-predict tokens.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。