研究者らは、投機的なデコードを行わずに、LLM 重みに直接 3 倍の推論高速化を組み込みました。

エージェント AI ワークフローでは長い推論チェーンのコストとレイテンシが増大するため、メリーランド大学、ローレンス リバモア国立研究所、コロンビア大学、TogetterAI のチームは、3 倍のスループット向上をモデルの重みに直接組み込む方法を発見しました。

別の製図モデルを必要とする投機的デコードとは異なり、このアプローチでは追加のインフラストラクチャは必要ありません。モデルの既存のアーキテクチャに 1 つの特別なトークンを追加するだけです。

次のトークンの予測の限界

次のトークン予測 (前方パスごとに 1 トークンのテキストを生成する) では、モデルが数千のトークンを生成する必要がある場合、スループットの上限が非常に高価になります。 このボトルネックは、頻繁に数千もの「」を生成する推論モデルで特に問題になります。思考の連鎖」トークンが最終応答を生成する前に生成されるため、ユーザー エクスペリエンスが遅くなり、コストが高くなります。

マルチトークン予測 (MTP) は、言語モデルが単一の前方パスで複数のトークンを同時に生成できるようにする代替トレーニング パラダイムを提供します。たとえば、すぐ次のトークンだけでなく、トークンのブロックを一度に予測するようにモデルをトレーニングできます。

メリーランド大学のコンピュータ サイエンスの博士号取得候補者であり、この論文の共著者である John Kirchenbauer 氏は VentureBeat に対し、エージェント ワークフローへの移行に伴い、焦点は全体のスループットからシングル ユーザーの速度に移っていると語った。 「今日では、超長時間の思考トレースが標準であり、エージェントの外部ループによってそのコストがさらに増大しているため、レイテンシは、ハードウェア ユニットあたり 1 秒あたりの総トークン数 (tps/GPU) と同じくらい、全体的なサービス効率の要素として重要になりつつあります」と Kirchenbauer 氏は述べています。同氏は、標準的なバッチ処理による次のトークン予測はすでに全体的なスループットにとって最適であるが、新しいアプローチは「[s] たった 1 人のユーザーのクエリで GPU を飽和させ、その 1 人のユーザーのレイテンシを短縮します。」

他の方法も存在しますが、欠点があります。 「投機的デコードと、次トークン予測 (NTP) に代わる効率重視の代替手段としての拡散 LLM は、どちらも遅延重視の高速化技術であることは注目に値します」と Kirchenbauer 氏は述べています。ただし、投機的なデコードには、補助的な「製図」モデルの展開と管理が必要であり、製図と検証により多くの絶対的なコンピューティングが費やされます。一方、MTP は「同様のトレードオフを活用しており、提供がより簡単で、それ自体が科学的に興味深いものです。」

ただし、現在の MTP パラダイムには制限があります。 MTP の言語モデルをトレーニングする標準的な目的には、その予測をデータセットからのグラウンド トゥルース テキストと比較することが含まれます。落とし穴は、この標準トレーニングでは、一連のトークン間の結合関係を考慮するのではなく、特定の位置にあるトークンの確率を個別に予測するようにモデルに学習させることです。

モデルがこの標準的な方法を使用して一度に複数のトークンを予測しようとすると、2 つの大きな問題が発生します。 1つ目は文法的な不一致です。たとえば、モデルが「動物園の飼育員が餌を与えた」という接頭辞に続く 2 つの単語を予測する場合、モデルは個別にサンプリングして、「パンダの竹」や「ライオンの肉」ではなく、「パンダの肉」や「ライオンの竹」のような不一致のフレーズを生成する可能性があります。

2番目の問題は、退化的反復です。一般的なテキストは予測できないため、標準データセットに対して 100 位置先のトークンを予測しようとするモデルは、英語で最も一般的な単語である「the」を予測するだけになります。その結果、モデルは遠い将来のポジションに対して「…the the the…」のような無意味な出力を行うことになります。

自己蒸留によるマルチトークン予測

複数のトークンの生成の問題を解決するために、研究者らは学生と教師のスキームを使用する新しいトレーニング手法を提案しています。スチューデント モデルは、複数のトークンを予測することを学習するモデルであり、決定論的なマルチトークン ブロックを生成します。教師モデルは、強力な標準のネクストトークン予測言語モデルとして機能し、そのブロックを評価します。教師は批評家の役割を果たし、生徒が提案したシーケンスがどの程度可能性があり、一貫性があるかを計算します。生徒が「ライオン竹」のような不一致のフレーズを提案すると、教師はそれに高い損失を割り当て、その構文を避けるように生徒に教えます。

学生モデルは静的なテキストを単に暗記するわけではないため、このパラダイムはポリシーに基づく強化学習からインスピレーションを得ています。単一のフォワード パスで完全なロールアウト (RL の用語で一連のアクション) を瞬時に並行して生成し、教師がそれがどれだけ優れているかに基づいて報酬を受け取ります。 トレーニング ペアが事前に固定されている静的な教師ありメソッドとは異なり、ここでのフィードバックは動的であり、生徒自身の出力からリアルタイムで生成されます。 強力な教師はトークンの一貫性も検証します。これにより、生徒モデルが単語の繰り返しなどの退化した出力を学習するのを防ぎます。

開発者にとって、このアプローチの利点はそのシンプルさにあります。 「特別なトークンの追加を除いて、アーキテクチャにまったく変更はありません」とキルヒェンバウアー氏は語った。モデルの既存の埋め込み行列内の未使用のスロットを採用して、 マスク トークンとは、逐次操作を並列操作に変換する技術です。 「あらゆる標準的な次のトークン予測言語モデルは、この方法で適応させることができます。MoE、ウィンドウ化されたアテンション、SSM レイヤーなどの内部実装はそのまま残されており、適応を妨げるものはありません。」

これは、エンジニアリング チームにとって、パイプラインを再構築することなく、すでに運用されているモデルに適応を適用できることを意味します。

複数のトークンを同時に生成すると、推論時の応答の精度が損なわれる可能性があります。出力の品質を犠牲にすることなく生成速度を最大化するために、著者らは ConfAdapt と呼ばれる適応デコード戦略を導入しました。

ConfAdapt は、各ステップで信頼度のしきい値 (90% など) を評価します。モデルはトークンのブロックを生成しますが、この信頼性の高いしきい値以上のトークンのみを保持します。今後のテキストが非常に予測可能または構造的である場合、モデルの信頼度は非常に高くなります。大量のトークンを一度に受け入れて出力するため、簡単なトークンの計算時間を大幅に節約できます。次に、コストのかかる単一トークンのパスを、より多くの計算量を必要とするより困難なトークンに集中させます。

マルチトークン予測をテストする

トレーニング パラダイムが実際にどのように機能するかを確認するために、研究者らはその方法を一般的なオープンウェイトの指導調整モデルに適用しました。彼らは、強力な汎用モデル Llama-3.1-8B-Magpie と、コスト重視の企業展開によく選ばれる小型で効率的な Qwen3-4B-Instruct-2507 をテストしました。どちらのモデルも、推論トレースに大きく依存する小学校の算数の合成問題のデータセットである MetaMathQA に基づいて調整されました。

実験により、速度と精度の間に明確なスイートスポットがあることが明らかになりました。 ConfAdapt 戦略を使用すると、Llama-3.1-8B モデルは、数学ベンチマークでの精度の低下が 3% 未満で、3 倍の高速化を達成しました。 Qwen3-4B モデルは、同じ 3 倍の高速化を達成しましたが、精度はわずかに 7% 低下しました。より積極的な設定では 5 倍の速度向上を達成できますが、精度はさらに低下します。

これが現実世界のタスクにどのように変換されるかは、予測可能性に依存します。同氏は、「ConfAdaptアプローチはドメイン内に固有のエントロピーに合わせて加速を自然に調整するため、モデルが次に何が起こるかを正確に『知っている』場合、それを1回のパスで出力できる」と述べ、不確実な出力に対してはより多くのステップを使用しながら、予測可能なタスクでは大幅な高速化につながると述べた。

高速化は、マルチトークン予測トレーニング フェーズに含まれていないドメイン間でも転送されました。これには、数学や推論などのトレーニング データと同じドメイン内のタスクと、創造的な執筆や要約などの自由形式のタスクが含まれます。

この転移学習にもかかわらず、特殊なタスクのためにこれらのモデルを展開する企業は、それに完全に依存すべきではありません。 「私たちの推奨は、特別な産業分野からのサンプルを使用して MTP 用にモデルを調整/適応させることです」と Kirchenbauer 氏は述べています。 「展開ドメインからのプロンプトを使用して MTP 適応を実行すると、おそらく最高のパフォーマンスが達成されます。」

互換性と今後の展開に向けたサービスの提供

研究チームが発表したのは、 ハグフェイスで訓練されたモデル そしてすぐにリリースされます MTP フレームワークのコード。 これらのモデルを vLLM または SGLang に統合するインフラストラクチャ チームは、バッチ処理と KV キャッシュの処理方法の変更を考慮する必要がありますが、それは 1 回限りのエンジニアリング投資であり、継続的な負担ではありません。 しかし、キルヒェンバウアー氏は「統合に対する明確な障壁は存在しない」とみており、チームが「統合への最短経路を特定するために一部のシステム専門家と協力している」ことを認めた。

リリースされたモデルをテストしたいチームに対するキルヒェンバウアー氏のアドバイスは、まず数を数えたりフレーズを繰り返すなどのおもちゃのプロンプトから始めて、実際の ConfAdapt の効果を確認し、次に最良の結果を得るために特定のデプロイメント ドメインのサンプルを使用してモデルを適応させます。 「全体として、私たちのアプローチを本番環境にすぐに実装できれば、低遅延のエージェント モデルの構築と展開のライフサイクルが簡素化される可能性があると期待しています」と Kirchenbauer 氏は結論付けました。 「NTP モデルの既存の高速化技術は、推論ハーネスとロジックのみにほぼ重点を置いていますが、私たちのアプローチは、複雑さの一部をモデル自体に組み込むだけであり、既存の作業を大幅に補完するものになります。」