AI-Papers
- モデルの重みを通常の次トークン予測で学ぶ自己回帰重みと、並列生成用の軽量な拡散重みに分離する「diffusion-augmented LLM」を提案
- サンプラー族Ψ-Specにより、別のドラフトモデルを用意せずに元の自己回帰分布と同じ出力を保ったまま最大3倍の高速化を達成
- 8BのUnoが26BのDiffusionGemmaや商用のMercury 2をエージェント的なツール利用やコード生成で上回ったと報告
研究の背景
現在のLarge Language Model（LLM、大規模言語モデル）は、1トークンずつ順番に予測する自己回帰生成を基本にしています。品質の面では非常に強力な一方、生成するトークンの数だけモデルを呼び出す必要があるため、長い応答ほど待ち時間が伸びるという構造的な弱点を抱えています。
この逐次のボトルネックを崩す方向として、これまで大きく2つの流れがありました。1つは投機的デコード（speculative decoding）で、小さなドラフトモデルに先の複数トークンを推測させ、本体モデルがまとめて検証する方式です。出力は元のモデルと一致するロスレスな高速化になりますが、ドラフトモデルを別途用意して運用する手間とメモリコストがかかります。
もう1つは離散拡散に基づく拡散LLMで、複数トークンを同時に埋めていくため並列性は高いものの、自己回帰モデルとは異なる分布を学ぶことになり、元のモデルの品質を犠牲にしがちでした。Cornell大学やCarnegie Mellon大学などの研究者らによる本論文「Unlocking Lossless Speedups in LLMs via Discrete Diffusion」は、この2つの弱点を同時に避けることを狙っています。
提案手法の考え方
提案されたモデルUnoの中心にあるのは、パラメータの役割分担です。モデルの重みを、通常どおり次トークン予測で学習する自己回帰重みと、並列トークン生成を担う軽量な拡散重みに分けます。品質を決める部分は自己回帰学習のまま据え置き、速度を稼ぐ部分だけを後付けする、という設計です。
拡散重みは「Diffusion Distillation」と呼ばれる蒸留フェーズで学習されます。既存のLLM学習に対する追加コストは小さく、ゼロから拡散LLMを訓練し直す必要がありません。追加されるパラメータ数もわずかで、推論時のGPUメモリ使用量は比較対象の中で最も小さいと報告されています。
図1に示すように、1つのモデルの内部で自己回帰重みと拡散重みが役割を分け合い、生成時には両者を組み合わせて複数トークンを同時に出します。ドラフト用の別モデルを持たずに済むのは、この「1モデル2種類の重み」という構成によるものです。
Ψ-Specによるロスレス生成
並列に生成したトークンをそのまま採用すると、出力の分布は元の自己回帰モデルからずれてしまいます。そこで論文は、Ψ-Specと名付けたサンプラー族を導入しました。これは自己回帰分布から複数トークンを並列にサンプリングするための一群の手続きで、生成される系列の分布が元のモデルと一致するように設計されています。
この性質により、精度を落とさずに速度だけを取り出せます。KVキャッシュの一部を捨てて速度を得るRandom Attentionのような近似ベースの高速化とは異なり、Unoは出力の同一性を保ったまま計算を削る点に特徴があります。さらにΨ-Specは、文脈長を固定したまま推論時により多くの計算を投じて品質や速度を調整する、推論時スケーリングの余地も持たせています。
実験結果
評価では、ベースとなる自己回帰モデルに対して全バッチサイズで最大3倍のスループット向上が確認されました。ドラフトモデルを使うEagle3やDFlashといった投機的デコード手法との比較でも、どのバッチサイズでもUnoのスループットが上回り、速度と品質のトレードオフでPareto優位に立ったと報告されています。投機的デコードに対する高速化は約2.5倍とされています。
品質面では、8BパラメータのUnoが、26BのDiffusionGemmaや商用モデルのMercury 2を、エージェント的なツール利用、コード生成、長文脈推論のタスクで上回りました。主要な比較の整理は次のとおりです。
手法 | ドラフトモデル | 出力の同一性 | 報告された速度 |
|---|---|---|---|
ベースの自己回帰モデル | 不要 | 基準 | 1倍（基準） |
投機的デコード（Eagle3、DFlash） | 必要 | 保たれる | Unoに劣る |
拡散LLM（DiffusionGemmaなど） | 不要 | 保たれない | 並列だが品質低下 |
Uno | 不要 | 保たれる | 最大3倍 |
大きいモデルを小さいモデルが上回ったという結果は、拡散側の学習が単なる速度装置にとどまらず、生成の質にも寄与している可能性を示しています。ただし比較対象の学習データや訓練規模は揃っていないため、この一点だけでアーキテクチャの優劣を結論づけるのは早いでしょう。
まとめと今後の展望
Unoの実用的な意味は、既存の自己回帰LLMを作り直さずに、軽い蒸留フェーズを足すだけで推論コストを下げられる点にあります。ドラフトモデルの管理が不要になることは、サービング環境の単純化という面でも効いてきます。
一方で、Diffusion Distillationに必要な計算量や、モデル規模・タスク種別による効果の差については、公開された情報だけでは判断が難しい部分が残ります。並列に出せるトークン数は文脈やタスクの予測しやすさに左右されるため、実運用での平均的な加速率は条件次第で変わると考えるのが妥当です。
コードとチェックポイントはGitHubとHugging Faceで公開されており、手元の環境で検証しやすい形になっています。推論コストの削減は多くの応用に直結するため、他のモデルファミリーへの適用がどこまで進むかが次の焦点になりそうです。
論文情報: "Unlocking Lossless Speedups in LLMs via Discrete Diffusion"（Subham Sekhar Sahoo et al., 2026） arXiv:2609.04010, CC BY 4.0
