LLMの自己回帰生成を並列化する拡散モデル基盤 Uno──ベースモデル比最大 3 倍の高速化と精度維持を両立

論文概要

Large Language Models (LLMs) における従来の次トークン予測 (Next-Token Prediction: NTP) による自己回帰 (Autoregressive: AR) 生成は、逐次的な処理がボトルネックとなり推論速度が制限される。本論文では、ARモデルの分布を維持しつつ、拡散モデルを用いて複数トークンを並列に生成する手法「diffusion-augmented LLMs」を提案している。開発されたモデル「Uno」は、ベースとなるARモデルと比較して最大 3 倍のスループット向上を達成しつつ、精度を完全に維持することに成功した。

関連研究

従来の推論高速化手法としては、別個のドラフトモデルを用いる Speculative Decoding や、生成品質のトレードオフを伴う Diffusion LLMs (d-LLMs) が存在した。しかし、Speculative Decoding は別モデルの管理コストがかかり、従来の拡散型LLMは基盤となるARモデルの性能を劣化させる課題があった。本研究は、パラメータの分離と Diffusion Distillation により、これらの限界を克服する。

新規性と貢献

最大の貢献は、ARウェイトと軽量な拡散ウェイトを明確に分離した点にある。標準的なNTP目的関数でAR重みを学習し、追加の Diffusion Distillation 段階によってわずかなオーバーヘッドで拡散重みを最適化する。また、独立したドラフトモデルを不要とするサンプラファミリ $\Psi$-Spec を導入し、あらゆるバッチサイズで既存のSpeculative Decoding手法を凌駕するスループットを実現した点も学術的・実用的なブレイクスルーである。

提案手法の詳細

本手法では、モデルのパラメータを2つのセットに分割する。1つ目は標準的なARタスク用の重み、2つ目は複数トークンを同時に生成するための軽量な拡散重みである。なぜこの設計にしたかといえば、既存の強力なオープンウェイトARモデルの資産をそのまま活用しつつ、並列生成の能力を後付けで効率よく統合するためである。さらに、固定されたコンテキスト長でロスレスな加速と推論時スケーリングを可能にする $\Psi$-Spec サンプラが組み込まれている。

評価・考察

評価実験において、8B規模の Uno モデルは、26B規模のオープンな d-LLM である DiffusionGemma や、プロプライエタリな Mercury 2 を、エージェントによるツール利用、コーディング、長文脈推論の全評価ベンチマークで上回った。また、評価されたすべてのバッチサイズにおいて、主要な Speculative Decoding 手法よりも高いスループットを記録し、ベースARモデルに対して最大 $3\times$ の高速化を実証している。

応用例と今後の展望

実務インパクトとして、APIプロバイダや大規模なコンタクトセンターシステムにおいて、LLMの推論コスト削減とレイテンシ改善に直接寄与する。特に、数百リクエストを同時に処理するクラウドインフラストラクチャ運用企業や、複雑なマルチステップ推論を行うAIエージェント開発において、ハードウェア増強なしのスループット向上手段として有効である。今後は、より多様なアーキテクチャや大規模モデルへの適用拡大が期待される。

結論

Uno は、拡散モデルを補完的に組み込むことで、ARモデルの優れたテキスト生成品質を損なうことなく並列生成と大幅な高速化を両立できることを示した。コードとチェックポイントが一般公開されており、今後のオープンソースLLMのエコシステムに大きな影響を与える。

注釈

• 次トークン予測 (NTP): 文脈から直後の1単語（トークン）を予測して順次出力していく、現在の主流の言語モデル学習・生成方式。
• Diffusion Distillation: 拡散モデルの生成プロセスを効率化・蒸留するための手法。