用於高效測試時擴展的前綴滑動

摘要

測試時擴展利用額外的測試時計算來提升效能，例如讓語言模型在解決問題時推理更久。由於模型透過完整注意力將整個推理軌跡保留在記憶體中，需要長時間思考的困難任務可能變得過於昂貴。然而，我們發現當模型持續推理時，大多數中間推理 token 的重要性會逐漸降低。這使我們質疑保留它們是否值得付出成本。基於此洞察，我們提出 Prefix Sliding（前綴滑動），在推理過程中丟棄那些既不屬於前綴、也不屬於最後數千個 token 視窗的 token。前綴包含提供給模型的關鍵指令與工具，而最近的 token 則是模型當前正在進行的推理。這使得無論模型推理多久，總記憶體需求都受到上限限制，從而實現高效的長時程測試時擴展。在無需訓練的情況下，Prefix Sliding 能讓現有模型速度提升 3 倍，同時維持效能。使用強化學習搭配 Prefix Sliding 進行訓練，則能透過將擴展能力延伸到超過十萬個 token 的推理軌跡，來達成更佳的效能。消融實驗顯示，Prefix Sliding 的表現優於對中間 token 進行摘要或使用傳統滑動視窗的方法。我們的程式碼位於 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