効率的なテスト時スケーリングのためのプレフィックススライディング

要旨

テスト時スケーリングは、追加のテスト時計算を用いて性能を向上させる手法であり、例えば問題を解くときに言語モデルにより長く推論させることを指す。モデルはフルアテンションで推論トレース全体をメモリに保持するため、長い思考を必要とする難しいタスクではコストが法外に高くなり得る。しかし、モデルが推論を続けるにつれて、中間の推論トークンのほとんどは重要度を失うことが分かった。これは、それらを保持する価値がコストに見合うのかという疑問を提起する。この知見に基づき、我々は Prefix Sliding を提案する。Prefix Sliding は、推論中に、プレフィックスまたは直近の数千トークンのウィンドウに含まれないトークンを破棄する。プレフィックスにはモデルが利用できる重要な指示とツールが含まれ、直近のトークン群はモデルが現在取り組んでいる推論である。これにより、推論の長さに関係なく総メモリ要件に上限が設けられ、効率的な長期のテスト時スケーリングが可能になる。訓練なしでも、Prefix Sliding は既存モデルの性能を維持したまま3倍高速化できる。強化学習を用いた Prefix Sliding による訓練は、10万トークンを超える推論トレースへのスケーリングを可能にすることで、より良い性能を達成できる。アブレーション実験の結果、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