前缀滑动用于高效的测试时扩展

摘要

测试时扩展利用额外的测试时计算来提升性能，例如让语言模型在解决问题时进行更长时间的推理。由于模型通过全注意力将完整推理轨迹保留在内存中，需要长思考的困难任务可能会产生高昂的成本。然而，我们发现随着模型继续推理，大多数中间推理词元会逐渐失去重要性。这不禁让人质疑保留这些词元是否值得其成本。基于这一见解，我们提出了Prefix Sliding（前缀滑动），即在推理过程中丢弃不属于前缀或最后几千个词元窗口的词元。前缀包含模型可用的关键指令和工具，而最近的词元则是模型当前正在进行的推理内容。这样一来，无论模型推理多久，总内存需求都被限制在一个上限内，从而实现了高效的长时程测试时扩展。无需训练，Prefix Sliding即可使现有模型速度提升3倍，同时保持性能。使用Prefix Sliding并通过强化学习进行训练，能够将推理轨迹扩展到超过十万词元，从而获得更好的性能。消融实验表明，Prefix Sliding优于对中间词元进行总结或使用普通滑动窗口的方法。我们的代码位于 https://github.com/Muennighoff/prefix-sliding

English

Test-time scaling uses extra test-time compute to improve performance, such as letting language models reason longer when solving a problem. As models keep the entire reasoning trace in memory via full attention, hard tasks that need long thinking can be prohibitively expensive. However, we find most intermediate reasoning tokens lose importance as the model continues reasoning. This calls into question whether retaining them is worth the cost. Based on this insight, we propose Prefix Sliding, which discards tokens during reasoning that are not part of the prefix or the window of the last few thousand tokens. The prefix has key instructions and tools available to the model, while the most recent tokens are the current reasoning the model is working on. This caps the total memory requirement regardless of how long the model reasons, allowing for efficient long-horizon test-time scaling. Without training, Prefix Sliding can make existing models 3x faster while maintaining performance. Training with Prefix Sliding using reinforcement learning can achieve better performance by enabling scaling to reasoning traces beyond a hundred thousand tokens. Ablations show Prefix Sliding outperforms summarizing intermediate tokens or vanilla sliding window. Our code is at https://github.com/Muennighoff/prefix-sliding