arXiv (NLP)AI
JetFlow:並列ツリードラフティングで推測デコーディングのスケーリング限界を突破
JetFlow: Breaking the Scaling Ceiling of Speculative Decoding with Parallel Tree Drafting
この記事についてAIに質問する →
日本語要約青い用語にマウスを合わせると解説が表示されます
大規模言語モデル(LLM)の推論速度向上は、生成型AIの実用化において重要な課題です。推測デコーディング(SD)は複数のトークンを同時にドラフトして検証することで自己回帰型LLMの処理を加速させる技術として注目されていますが、従来の手法にはスケーリングにおける根本的な制限がありました。ドラフト予算を増やしても、受理率が低下したりドラフティングのオーバーヘッドが増加したりすると、期待される速度向上が得られなくなるという課題に直面していたのです。
この問題の根底には「因果性と効率性のジレンマ」がありました。自己回帰的なドラフターは高い受理長を実現できますが、ツリー構造が深くなるにつれて計算コストが増加します。一方、双方向ブロック拡散ドラフターは単一パスですべての位置を生成できるものの、枝ごとに独立した確率分布を用いるため、個々には妥当でも相互に矛盾した候補ツリーが生成され、予算を浪費してしまいます。
新たに提案されたJetFlowは、この両者の長所を組み合わせたヘッドベースのSD框架です。凍結された対象モデルから抽出した融合隠れ状態を使用してカジュアルな並列ドラフトヘッドを訓練することで、候補ツリーのスコアが目的モデルの自己回帰分解と一致する設計を実現しました。これにより、より大きなドラフト予算をより長い受理プレフィックスに変換し、エンドツーエンドの速度向上を実現できます。
H100 GPUにおいて、JetFlowはMATH-500ベンチマークで最大9.64倍、会話ワークロードで4.58倍の速度向上を達成し、Qwen3モデルの密集型およびMoE型両方で従来のSD手法を一貫して上回ります。vLLM統合による実運用環境下での潜在性も実証されており、コードとモデルはGitHubで公開されています。