Orthrus-Qwen3: Diffusionベースの投機的デコードによるLLM推論の高速化

ほとんどの最新 Large Language Model (LLM) のアーキテクチャである自己回帰 (AR) Transformer の課題は、その逐次的な性質にあります。トークンを一つずつ生成することは計算コストが高く、実際のアプリケーションではボトルネックになりがちです。この問題を解決するために、研究者は投機的デコードを検討しています。投機的デコードでは、より小さく高速な「ドラフタ」モデルが複数トークン先を予測し、より大きな「ターゲット」モデルがそれらを並列で検証します。

Orthrus-Qwen3 は、この問題に対して新しいアプローチを提示します。凍結された Qwen3 バックボーンの各層に、学習可能な拡散注意モジュールを直接組み込むのです。従来の投機的デコードとは異なり、Orthrus は外部のドラフタモデルや別個の KV キャッシュを必要とせず、精度を犠牲にすることなく大幅な速度向上を実現します。

Orthrus の仕組み: 拡散注意メカニズム

Orthrus は、標準的な AR Transformer に対して各層に学習可能な拡散注意モジュールを注入することで変更を加えます。重要なイノベーションは、ベースモデル（凍結された AR Transformer）と拡散ヘッドが 単一の KV キャッシュ を共有するように設計されている点です。

このプロセスは二段階で行われます:

  1. 拡散パス: 拡散ヘッドが $K=32$ トークンを並列に投影します。これは単一ステップのデノイジングプロセスで、複数の候補トークンを予測します。
  2. 検証パス: AR ヘッドが第二のパスでこれらのトークンを検証します。ベースモデルの元の出力分布と一致する最長のプレフィックスを受け入れます。

ベースモデルの重みが凍結されているため、出力分布は元の Qwen3 モデルと 同一 であることが保証されます。これにより、ユーザーは同等の品質の応答を得られますが、はるかに高速に処理できます。

パフォーマンスベンチマークと利点

著者らによると、Orthrus-Qwen3 は従来の拡散 LLM および既存の投機的デコード手法に比べ、スループットと効率の両面で大幅な改善を示しています。

スループットと速度

  • Tokens Per Forward (TPF): Orthrus は最大 7.8 倍 の TPF を達成し、MATH-500 ベンチマークでは約 6 倍 の実時間速度向上を示しています。

拡散 LLM との比較

従来の拡散 LLM（例: Dream、Fast-dLLM-v2、Mercury）は、並列生成を可能にするためにモデルのベース重みを変更することが多く、精度が低下しがちです。たとえば Fast-dLLM-v2 は MATH-500 で 11 ポイント の精度低下を示しました。対照的に Orthrus はバックボーンを凍結するため、精度は Qwen3-8B と完全に一致 します。

投機的デコードとの比較

EAGLE-3 や DFlash といった手法と比較した場合、Orthrus は以下のようなアーキテクチャ上の利点を提供します:

  • 外部ドラフタ不要: 別モデルの初期化や同期が不要なため、Time to First Token (TTFT) のペナルティが排除されます。
  • メモリ効率: KV のオーバーヘッドは $O(1)$（約 4.5 MiB のフラットサイズ）に抑えられます。
  • 高い受容率: MATH-500 ベンチマークで Orthrus は受容長 11.7 を記録し、DFlash の 7.9、EAGLE-3 の 3.5 を大きく上回ります。

訓練および実装の詳細

Orthrus の実装は、訓練要件の面で非常に効率的です。パラメータの 16% のみが学習対象となり、8 台の H200 GPU を用いて 24 時間 未満で 1B 未満のトークンで訓練が完了しました。

研究者は、受容率向上のために KL ディスティレーション が Cross-Entropy (CE) よりも効果的であること、そして 単一ステップデノイジング（6.35 TPF）が多ステップデノイジング（3.53 TPF）を上回ることを確認しています。

制限事項と考慮点

Orthrus-Qwen3 の結果は印象的ですが、モデルは凍結されたベースモデルに依存しています。したがって、元の Qwen3 が持つバイアス、幻覚、知識ギャップをすべて引き継ぎます。さらに、現在の評価は Qwen3 と greedy / rejection サンプリング手法に限定されています。

コミュニティでこの手法の可能性が議論される中、DeepSeek-V3 やローカル LLM 実行向けの量子化 GGUF バージョンへの適用への関心が高まっています。もし成功すれば、高スケール AI プロバイダーやローカルユーザーにとって、レイテンシと混雑の大幅な削減につながるでしょう。

Sources