吴恩达も参加したスタンフォードの新研究:中間の推論トークンを捨て、長い思考が最大3倍高速化

推論チェーンが長くなるにつれて、モデルはその全履歴を覚えておく必要があるのだろうか?Prefix Slidingはタスクのプレフィックスと最新のウィンドウだけを保持し、長い思考が推論チェーン全体を背負わずに済むようにする。

推論チェーンが長くなるにつれて、モデルが覚えておくべき情報はどれほどだろうか?

標準的なフルアテンションのやり方では、その答えはほぼすべてである。

これまでに計算したこと、中間ステップ、蓄積された推論履歴は、すべてその後のアテンション範囲に残り続ける。推論が長くなるほど、新しいトークンを生成するたびにコストも増加する。

問題は、モデルが本当にこれらの過ぎ去った中間ステップを振り返る必要があるのかどうかだ。

最近、スタンフォード大学、ワシントン大学などの機関による共同研究がここに着目した。吴恩达、Yejin Choi、Percy Liang らが参加している。

チームは Prefix Sliding を提案し、タスクのプレフィックスと最新の推論ウィンドウだけを長期的に保持し、より古い中間トークンを後続のアテンション計算から徐々に除外する。

最終的に、既存のモデルを再トレーニングすることなく、長い思考で最大約3倍の高速化を実現強化学習に使用すると、単一の推論軌跡を10万トークン以上に拡張できる。

图片

論文タイトル:

Prefix Sliding for efficient test-time scaling

論文アドレス:

https://arxiv.org/abs/2608.26070

コードアドレス:

https://github.com/Muennighoff/prefix-sliding

图片

アテンションが推論の両端に集中する理由

フルアテンションのコストは推論チェーンの長さに比例して増加し続けるが、モデルによるこれらの履歴トークンの使用は均一ではない。

チームがQwen3-1.7BのAIME25における推論軌跡を分析したところ、アテンションは明らかに「両端が高く、中央が低い」というパターンを示した。

最初のPrefixと最近生成されたトークン群により多くの注意が向けられ、中央の多数の推論トークンへのアテンションは全体的に低い水準に留まる。

圖片

〓 長い推論ではアテンションは主にプレフィックスと最近生成されたトークンに集中する

Prefixにはシステム指示、タスク記述、ツール情報が保存されており、特に最初の数トークン(特に最初の4トークン)はアテンションシンクの役割も果たしている。最近生成されたトークンは、モデルが現在処理している問題に対応する。

論文では簡単な式 ((42 + 84) × 4) - 5 を例に挙げている:42 + 84を計算した後、後続の計算ではその結果だけが必要になる可能性があり、それまでの完全な推論プロセスを保持する必要はないかもしれない。

これがPrefix Slidingの直接的な動機となる:後続の計算が常に完全な推論履歴を携帯する必要はない。

图片

Prefix Slidingはどうやって中間トークンを捨てるのか?

Prefix Slidingはタスクのプレフィックスを固定で保持し、最近のW個の推論トークンでスライディングウィンドウを構成する。生成が進むにつれて、より古い中間トークンは後続のアテンション計算に使用されるKVキャッシュから徐々に削除される。

Prefixが100トークン、ウィンドウが4096だとすると、完全な推論が10万トークンに達しても、後続のアテンション計算は最大約4196トークンだけ保持すればよい。

ウィンドウが満杯になると、単一ステップのコストは完全な推論チェーンの長さに依存しなくなる。

图片

〓 Prefixは固定で保持され、最近の推論トークンはウィンドウとともに前方へ移動する

通常のスライディングウィンドウでは最初のタスク情報がウィンドウから徐々に外れてしまうが、Prefix SlidingはPrefixを常に保持する。

チームは最終的にContinue PEを採用した。付録のAIME25での実験では、Reset PEと同等の性能を示し、新しい位置エンコーディングを再適用する必要がなく、既存のキャッシュ表現を直接再利用できるため、計算がより効率的である。

チームはNVIDIA Hopper向けに専用のFlashAttentionカーネルも実装した。

Prefixとスライディングウィンドウの外側に完全にあるタイルは直接スキップし、有効領域と部分的に重なるタイルには要素単位のマスクを適用することで、無駄なロードと計算を減らし、実際の速度も通常のスライディングウィンドウ実装に近い。

图片

〓 ウィンドウが満杯になると、Prefix Slidingの生成スループットは安定する

图片

トレーニングなしで最大3倍高速化、RLで10万トークンに拡張

追加トレーニングなしの設定では、著者らはQwen3-1.7Bをメインモデルとして、AIME25、GPQA、MATH500で評価した。

Prefix Slidingの利点は、同じ思考時間内により多くのトークンを生成できることにあり、個々のトークンの品質が高いからではない。

4096トークンのウィンドウを例にとると、Prefix SlidingはAIME25、GPQA、MATH500でフルアテンションと同等の性能を示す一方、長いシーケンスではスループットがはるかに高く、128Kでは5224 tok/sに達したのに対し、フルアテンションは448 tok/sしかなかった。

图片

〓 同じ思考時間で、Prefix Slidingはより多くの推論トークンを生成できる

ここでいう「3倍」は、同等のタスク性能における思考時間の比較であり、上記のスループット数値を直接比較したものではない。

图片

〓 同じ思考時間で、Prefix Slidingはより多くの推論トークンを生成できる

強化学習フェーズでは、Prefix Slidingはロールアウトを10万トークン以上に拡張できる。超長軌跡全体に対する完全な逆伝播を避けるため、著者らは切り詰められた逆伝播を採用した。

スライディングウィンドウが複数層にまたがる理論上の受容野はW×Lに達するが、著者らは既存の分析を引用して、情報ボトルネックの影響により実際の有効受容野は約1.5×Wに近いと指摘し、したがって最後のウィンドウに対する逆伝播では、その前の限られたコンテキストを提供するだけでよい。

图片

〓 超長推論軌跡では、チャンク化または切り詰められた逆伝播を使用できる

例えば10万トークンの軌跡でウィンドウが2048の場合、トレーニング側は最後の8192トークンだけを受け取り、最初の6144トークンはコンテキストとして、最後の2048トークンでRL損失を計算する。

実験では、ウィンドウの1倍だけを保持すると大きなKL乖離が生じるが、2倍に拡大すると大幅に低下し、4倍は8倍とほぼ同等になるため、メイン実験では4倍ウィンドウを採用した。

圖片

〓 より多くの履歴トークンを渡すと、切り詰められた逆伝播のKL乖離は急速に低下する

付録では、DeepSeek-R1-Distill-Qwen-7Bでも小規模な非同期RL実験を実施した。

結果は、トレーニング側に32768トークンを入力し、Prefix Slidingが最後の8192トークンに対してのみ逆伝播する設定で、トレーニング報酬とAIME24の性能はフルアテンションと同等であることを示した。

图片

〓 7Bモデルでは、切り詰められた逆伝播はフルアテンションと同等の性能を示す

著者らは、より大規模で長い推論チェーンでの性能はさらなる検証が必要であると明言している。

同程度のメモリ予算では、フルアテンションの最大長は8192トークンであるのに対し、Prefix Slidingは最大生成長を104Kまで段階的に拡大し、より高いトレーニング報酬を得た。

图片

〓 同程度のメモリ予算で、Prefix Slidingはより長い強化学習軌跡を支える

图片

中間トークンをすべて捨てられるわけではない

Last-k、Summary、通常のスライディングウィンドウと比較して、Prefix Slidingは性能と効率の総合バランスが最も良い。

Last-kは保持トークンを繰り返し処理する必要があり、通常のスライディングウィンドウはタスクのプレフィックスを徐々に失う。Summaryは追加の要約生成と多数のハイパーパラメータを必要とし、実際のサンプルでは、モデルが要約内の既存の推論を無視して再び解き始めることもある。

图片

〓 Prefix SlidingとLast-k、Summary、通常のスライディングウィンドウの性能-効率比較

トレーニングなしのLiveCodeBench実験では、モデルが先にコードを書き、数千トークンのコメントで思考を続け、コードに戻るときには以前の内容がウィンドウから滑り落ちている可能性がある。

少なくとも16384トークンのウィンドウが必要で、フルアテンションに匹敵する。

著者らは、Prefix Slidingを強化学習に使用すると、モデルがこのようなコメント推論行動を適応させ、より短いウィンドウを許容するようになる可能性があると指摘している。

图片

〓 LiveCodeBenchでは、フルアテンションに匹敵するには少なくとも16384トークンのウィンドウが必要

短いタスクでも恩恵が得られるとは限らない。HealthBenchは平均約2086トークンしか生成せず、ウィンドウが2048であるため、多くのサンプルが実際にはスライディング段階に入らず、加速の余地は限られている。

エージェントのシナリオでは、一度の長すぎるウェブページやファイル出力が限られたウィンドウを埋め尽くし、モデルが完全なコンテンツを同時に読み取れなくなる可能性もある。

多ターン対話も新たなコンテキスト管理の問題を生む:後続のユーザー指示をPrefixに統合して長期的に保持すべきか、それともウィンドウから滑り落ちることを許すべきかについて、論文は現時点で結論を出していない。

Prefix Slidingは、超長いPrefixがPrefill段階で生じるKVキャッシュのオーバーヘッドを削減することもない。

論文の直接的な実証比較は、既存の事前学習済みTransformerに直接適用でき、単一ステップの生成コストが有界である手法に主に限定されており、他のモデルアーキテクチャ、準二次複雑度の手法、ハイブリッドスライディングウィンドウモデルをカバーしていない。

結語

Prefix Slidingは、長い推論の効率問題をコンテキスト管理の問題として具体化する:生成済みの履歴情報をいつまで保持する必要があるのか。

どの情報をPrefixに長期的に保持すべきか、どの情報をウィンドウとともに外に出してもよいか、より細かいコンテキスト管理戦略が必要である。コードタスク、エージェントシナリオにおける大規模なツール出力、多ターン対話では、追加のメモリメカニズムが必要かどうかも未定である。

現在の結果は、固定プレフィックス+スライディングウィンドウが超長推論コストを大幅に削減できることを示すが、より大きなモデルやより複雑な長期タスクへの適用範囲は今後の検証が必要である。

関連記事

分享網址
AINews·AI 新聞聚合平台
© 2026 AINews. All rights reserved.