💻 computer science

JetFlow: Breaking the Scaling Ceiling of Speculative Decoding with Parallel Tree Drafting

JetFlowは、分岐ごとに一貫したトークントリーを生成するための因果的並列ドラフトヘッドを学習させることにより、既存の手法のスケーリングの限界を克服し、受理率を損なうことなく多様なLLMワークロードにおいて大幅な高速化（最大9.64倍）を実現する新しい投機的デコーディングフレームワークである。

原著者： Lanxiang Hu, Zhaoxiang Feng, Yulun Wu, Haoran Yuan, Yujie Zhao, Yu-Yang Qian, Bojun Wang, Daxin Jiang, Yibo Zhu, Tajana Rosing, Hao Zhang

公開日 2026-06-19

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Lanxiang Hu, Zhaoxiang Feng, Yulun Wu, Haoran Yuan, Yujie Zhao, Yu-Yang Qian, Bojun Wang, Daxin Jiang, Yibo Zhu, Tajana Rosing, Hao Zhang

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

あなたは長い物語を書こうとしていると想像してください。しかし、あなたには非常に厳格なエディター（AIモデル）がついており、そのエディターは一度に一単語ずつしか書かせることができません。二つ目の単語を書く前に、エディターが最初の単語を読み、チェックし、ゴーサインを出さなければなりません。三つ目の単語の前に、二つ目の単語をチェックし、という具合に続きます。この「一度に一単語ずつ」というルールは、たとえエディターが非常に賢かったとしても、執筆を非常に遅いものにします。

**投機的デコーディング（Speculative Decoding）**は、これをスピードアップするためのトリックです。エディターが一つ一つの単語をチェックするのを待つ代わりに、あなたは素早く安価な助手（「ドラフター」）を雇い、次の数単語を予測させます。そして、その予測をまとめてエディターに提示します。もしエディターがその予測に同意すれば、あなたは一単語を書く時間で、複数の単語を書くことができます。もしエディターが同意しない場合は、間違った予測を捨てて最初からやり直します。

問題：「スピード vs 正確性」の罠

論文によれば、従来の手法はこの壁に突き当たっていました。彼らはジレンマに直面していました。

「慎重な」助手： ある助手は非常に慎重です。彼らは次の単語を予測し、その予測に基づいた次の単語を予測し……という手順を踏みます。これにより、彼らの予測は非常に正確になります（エディターによく受け入れられます）が、ステップ・バイ・ステップで考えなければならないため、速度は遅くなります。
「速い」助手： 他の助手は超高速です。彼らは単語同士のつながりを考えずに、単語のリストを一気に叫びます。これは非常に早いのですが、リストが意味をなさないことがよくあります（例：「猫が……飛んだ……へ……月……昨日」）。エディターは、これらの単語が物語の流れに合っていないため、ほとんどを拒否することになり、助手のスピードを無駄にしてしまいます。

論文では、これを**「因果関係と効率性のジレンマ（Causality-Efficiency Dilemma）」**と呼んでいます。あなたは通常、正確だが遅いか、あるいは速いが不正確かのどちらかを選ばなければなりませんでした。

解決策：JETFLOW

著者たちは、この罠を打ち破るJETFLOWと呼ばれる新しいシステムを作り出しました。JETFLOWを、**「並列的に考えることができるが、依然として物語の論理を尊重できるスーパー助手」**だと考えてください。

その仕組みは以下の通りです（簡単な比喩を用いて説明します）：

木のメタファー： 物語を「木」だと想像してください。幹はすでに書かれたテキストです。あなたは新しい枝（未来の単語）を伸ばしたいと考えています。
- 従来の「速い」助手は、ランダムに枝を伸ばします。ある枝は「猫が」、別の枝は「犬が」、そして三番目の枝は「月が」と言うかもしれません。彼らはどの経路が正しいのかを知らないため、死んだ枝を育てることに時間を浪費します。
- 従来の「慎重な」助手は、一つの枝を伸ばし、それをチェックしてから、次の枝を伸ばします。これは安全ですが、遅いです。
- JETFLOWは、幹を見て、瞬時に多くの異なる可能性のある枝を描き出します。しかし、ここが魔法のような点です。JETFLOWは、すべての枝が物語のルールに従うように保証します。もしある枝が「猫が」で始まるなら、その特定の枝における次の単語は、猫がやりそうなことでなければなりません。JETFLOWは異なる枝を混同することはありません。

JETFLOWの仕組み

一回のパスで、多くの経路を： JETFLOWは、メインのエディターの隠れた思考（hidden thoughts）を見る特別な「ヘッド（頭部）」を使用します。一瞥するだけで、起こりうる単語のツリー全体を予測します。
流れの尊重： JETFLOWは、特定のパスにおける「前の単語」だけを見るように強制する特別な「マスク（交通ルールのようなもの）」を使用します。これにより、助手が未来を覗き見たり、異なるストーリーラインを混ぜ合わせたりすることを防ぎます。
結果： 助手の予測が物語の流れと論理的に一致しているため、メインのエディター（ターゲットモデル）は、一度に非常に長い文字列を承認することができます。

結果：どれくらい速くなったのか？

論文では、強力なコンピューターチップ（H100 GPU）を使用して、数学の問題、コーディングタスク、およびチャットの会話でJETFLOWをテストしました。

数学の問題： 難解な数学テストにおいて、JETFLOWは標準的な遅い手法よりも9.6倍速くなりました。
チャット： 開放的な会話において、4.5倍速くなりました。
スケーラビリティ（拡張性）： 助手に許容される「予測（予算）」が増えるほど、より速くなります。従来のメソッドは、あまりに多くの単語を予測させようとすると混乱したり速度が低下したりしましたが、JETFLOWはより速く、より効率的に動作し続けました。

まとめ

JETFLOWは、異なる物語の結末を同時に叫ぶことができる助手のチームを雇うようなものですが、彼らはそれぞれの結末が単独で意味を成さなければならないことを理解している、賢いチームです。これにより、品質を損なうことなく、「一度に一単語ずつ」という速度制限を打破し、メインのエディターが膨大なテキストの塊を即座に承認することを可能にしています。

技術サマリー: JETFLOW

1. 問題提起

投機的デコーディング（Speculative Decoding: SD）は、複数のトークンをドラフト（下書き）し、それらを並列に検証することで、自己回帰型の大規模言語モデル（LLM）を加速させます。しかし、SDには根本的なスケーリングの天井が存在します。すなわち、ドラフト予算（ドラフトするトークン数）を増やしても、受理率（acceptance rate）が高く、かつドラftingオーバーヘッドが低い場合にのみ、エンドツーエンドの速度が向上するという問題です。

既存の手法は、以下の**因果関係と効率性のジレンマ（causality-efficiency dilemma）**により、これら2つの要因を同時に最適化することに苦慮しています。

自己回帰型ドラフター（例: EAGLE）: 木構造の投機的デコーディングに適した、パス条件付き（path-conditioned）の高精度な候補を生成します。これにより、高い受理長を実現できます。しかし、そのドラフトコストは木の深さに比例して線形に増加するため、スケーラビリティが制限されます。
双方向ブロック拡散ドラフター（例: DFlash）: すべてのドラフト位置を単一のフォワードパスで生成し、非常に低いトークンあたりのドラフトコストを実現します。しかし、これらは**ブランチ非依存のマージナル（branch-agnostic marginals）**を生成します。これらのマージナルによって形成される木では、個々のトークンは妥当であっても、分岐間で相互に矛盾が生じることがあります。この不整合はドラフト予算を浪費し、実効的な受理率を低下させるため、システムを効果的にスケールさせることを妨げます。

2. 手法: JETFLOW

JETFLOWは、**「1回のフォワードパスによるドラフト効率」と「ブランチごとの因果的条件付け」**を組み合わせることで、因果関係と効率性のジレンマを打破するために設計された、ヘッドベースの投機的デコーディング・フレームワークです。

コアアーキテクチャ

因果的並列ドラフトヘッド（Causal Parallel Draft Head）: JETFLOWは、凍結されたターゲットモデルの上に軽量なドラフトヘッドを学習させます。標準的なブロック拡散ヘッドが各位置を独立して予測するのに対し、JETFLOWのヘッドは、因果的依存関係を保持しながら、単一のフォワードパスで複数のツリーノードを予測します。
ツリー因果アテンションマスク（Tree-Causal Attention Mask）: 因果的一貫性を確保するために、ドラフトヘッドは特殊なアテンションマスクを採用しています。ドラフトツリー内の各ノードは、元の入力プレフィックスおよびツリー内の自身の先祖（ancestors）にアテンションを張ることができますが、子孫（descendants）や無関係な兄弟ブランチにはアテンションを張りません。
- これにより、ブランチごとの因子分解 $q(\pi(v) | x) = \prod_{u \in \pi(v)} q(y_u | x, h^o_x, \pi_{<u})$ が強制されます。
- これにより、ドラフト分布をターゲットモデルの自己回帰的因子分解に適合させ、候補となるブランチがそのパスに沿って選択された特定のトークンに条件付けされることを保証します。
特徴量融合（Feature Fusion）: ドラフトヘッドは、凍結されたターゲットモデルから融合された隠れ状態を再利用し、ターゲットのコンテキスト特徴量をドラフト層のKVキャッシュに注入することで、別個のドラフトモデルを構築するコストをかけずに生成をガイドします。

学習戦略

データ: ドラフトヘッドは、トレーニングコーパスから、あるいはより効果的には、ターゲットモデル自身によって再生成された継続シーケンスを用いて、ターゲットに整合するように学習されます。
損失関数: JETFLOWは蒸留に**前方KLダイバージェンス（Forward KL Divergence）**を利用します。著者らは、前方KLがハードラベルによる教師あり微調整（SFT）や逆KL（Reverse KL）よりも優れていることを見出しました。前方KLは、複数の妥当な継続先におけるターゲットモデルのソフトラベルの好みを保持するため、ツリー展開において極めて重要です。一方、逆KLはモード探索的（mode-seeking）すぎて、確率を過度に集中させてしまいます。
ツリー構築: 推論時、ドラフトヘッドはすべてのツリー深度に対して並列にロジットを生成します。ブランチ・スコアリング関数（デフォルトは累積ドラフト対数確率）がノードをランク付けし、ベストファースト展開アルゴリズムが特定の予算（例：256トークン）まで候補ツリーを構築します。

3. 主な貢献

JETFLOWフレームワーク: 因果的並列ドラフトヘッドを学習させることで、ドラフトコストと受理率を同時に最適化する新しいパラダイムを提示します。これにより、低コストな単一パスでのツリー生成と、高い受理率に必要なパス条件付きの依存関係の維持を両立させます。
スケーラブルなツリー・ドラフティング・アルゴリズム: 投機的デコーディングがより大きなドラフトトークン予算を効果的に活用できるようにするアルゴリズムを開発・評価しました。因果的一貫性を維持することで、JETFLOWは計算量の増加（より大きな予算）を、不整合なブランチへの予算の浪費ではなく、より長い受理プレフィックスへと変換します。
産業グレードの統合: JETFLOWは、サービングエンジンである vLLM に統合されています。この実装には、現実的なサービング負荷の下でツリー検証を効率的に処理するためのカスタムカーネル（例：SM90 paged FlashAttention とツリーマスク・ステージング）が含まれています。

4. 実験結果

著者らは、数学（GSM8K, MATH-500, AIME25）、コーディング（HumanEval, MBPP, LiveCodeBench）、およびチャット（MT-Bench）のベンチマークにおいて、Qwen3-8B（密なモデル）および Qwen3-30B-A3B（MoEモデル）を用いてJETFLOWを評価しました。

低予算領域（16トークン）: JETFLOWは、DFlashと同等の競争力のあるスピードアップ（例：GSM8Kで約4.8倍）を達成しており、短い線形ドラフトが、高確率な継続には十分であることを示しています。
高予算領域（256トークン）: JETFLOWは、予算が増加するにつれてベースラインを大幅に上回ります。
- MATH-500 において、JETFLOWはH100 GPU上で標準的な自己回帰デコーディングに対して 9.64倍 のスピードアップを達成しました（ $\tau \approx 10.7$ 受理トークン）。
- オープンエンド会話（MT-Bench） では、4.58倍 のスピードアップを達成しました。
- 対照的に、双方向ベースライン（DFlash-T）は、ブランチ非依存のツリーによる不整合のため、高い予算において収穫逓減または飽和を示します。
サービング性能: vLLMに統合されたJETFLOWは、小規模から中規模のサービング負荷の下で一貫してベースラインを上回ります。最適な予算は負荷に依存します。低バッチサイズでは、大きな予算（256）が最高のスループット（最大6.75倍のスピードアップ）をもたらしますが、重い負荷の下では、ステップあたりのオーバーヘッドを減らすために小さな予算の方が効率的です。
アブレーションによる知見:
- 因果性は極めて重要: 因果的ヘッドを拡散（diffusion）ヘッドに置き換えると、性能が大幅に低下します（例：MATH-500でのスピードアップが8.36倍から5.46倍へ）。これは、拡散ヘッドが一貫したツリーを構築できないためです。
- 損失関数: 前方KL蒸留は、性能を相対的に約40%低下させる逆KLよりも優れています。
- データ: 再生成されたターゲットモデルのシーケンスが最良の学習シグナルを提供しますが、生のコーパスデータでの学習でも一貫したスピードアップが得られます。

5. 意義と主張

本論文は、JETFLOWがドラフトの効率性と受理の質との間のトレードオフを解決することにより、投機的デコーディングのスケーリングの天井を打破したと主張しています。

天井の打破: 著者らは、従来の手法が、因果関係のために高い逐次コストを支払うか（EAGLE）、あるいは並列効率のために一貫性を犠牲にするか（DFlash）のいずれかを選択せざるを得ず、効果的にスケールできなかったことを証明しました。JETFLOWは、ブロックレベルの因果的条件付けが、安価でありながら因果的に一貫した並列ドラフトを可能にすることを証明しています。
実用的な影響: より大きなドラフト予算（例：256トークン）を使用しても受理率が低下しないことを可能にすることで、JETFLOWはデコーディングのレイテンシを大幅に削減します。この結果は、複雑な推論タスク（数学、コーディング）に特に有効であり、現代のハードウェア（H100, B200）上で大幅なスループット向上を実現するために、現実世界のサービング環境（vLLM）にデプロイ可能であることを示唆しています。
限定的な範囲: 本論文は、静的な予算ポリシーは効果的であるものの、動的なサービング時の予算スケジューリングは今後の課題として残されていることを認めています。主要な貢献は、因果関係と効率性のジレンマによってこれまで制限されていた投機的デコーディングのスケーラビリティを可能にする、因果的並列ドラフトへのアーキテクチャ上の転換にあります。