離散拡散によるLLMのロスレス高速化の実現

要旨

大規模言語モデル（LLM）は、その成功の多くを次トークン予測（NTP）に負っているが、自己回帰（AR）構造は、低速で逐次的なトークン生成を必要とする。このボトルネックを克服するため、我々は拡散拡張LLM（diffusion-augmented LLM）を導入する。これは、ARモデルの分布を定義しつつ、拡散を用いてその分布から複数のトークンを並列に抽出する新たなモデルクラスである。我々は、これらのモデルのパラメータを、標準的なNTP目的関数で訓練されるAR重みと、複数のトークンを同時に生成するように訓練される軽量な拡散重みの2つの集合に分離する。拡散重みは、既存のLLM訓練パイプラインに無視できるほどのオーバーヘッドしか加えない単純な拡散蒸留フェーズによって学習される。また、固定コンテキスト長でのロスレスな高速化と推論時スケーリングを可能にするサンプラ群Ψ-Specも導入する。投機的復号（speculative decoding）とは異なり、本手法は別個のドラフトモデルを必要としない。拡散LLM（d-LLM）とは異なり、基盤となるARモデルの品質を犠牲にせずに生成を高速化する。得られるモデルはUnoと呼ばれ、ゼロから訓練するか、既存のオープンウェイトAR LLMを拡張して構築できる。Unoは、評価されたすべてのバッチサイズにおいて主要な投機的復号手法よりも高いスループットを達成し、デバイスがサポートする最大のバッチサイズを含めて、ベースARモデルに対して最大3倍の高速化を実現する。特筆すべきことに、我々の8B Unoモデルは、エージェンティックなツール使用、コーディング、長文脈推論のすべての評価ベンチマークにおいて、主要なオープンd-LLMである26B DiffusionGemmaと、プロプライエタリなMercury 2を上回る。コードとチェックポイントはhttps://s-sahoo.github.io/uno/ で公開する。

English

Large Language Models (LLMs) owe much of their success to next-token prediction (NTP), but their autoregressive (AR) structure requires slow, sequential token generation. To overcome this bottleneck, we introduce diffusion-augmented LLMs, a new class of models that defines an AR model distribution while using diffusion to draw multiple tokens in parallel from that distribution. We decouple the parameters of these models into two sets: AR weights, trained using the standard NTP objective, and lightweight diffusion weights, trained to generate multiple tokens simultaneously. The diffusion weights are learned through a simple Diffusion Distillation phase that adds negligible overhead to existing LLM training pipelines. We also introduce Ψ-Spec, a family of samplers that enables lossless acceleration and inference-time scaling at a fixed context length. Unlike speculative decoding, our method requires no separate draft model. Unlike diffusion LLMs (d-LLMs), it accelerates generation without sacrificing the quality of the underlying AR model. The resulting models, called Uno, can be trained from scratch or built by augmenting existing open-weight AR LLMs. Uno achieves higher throughput than leading speculative-decoding methods at every evaluated batch size and delivers up to 3times speedups over the base AR model, including at the largest batch size supported by the device. Notably, our 8B Uno model outperforms the leading open d-LLM, the 26B DiffusionGemma, and the proprietary Mercury 2 across all evaluated benchmarks in agentic tool use, coding, and long-context reasoning. We release code and checkpoints at: https://s-sahoo.github.io/uno/