Unlocking Lossless Speedups in LLMs via Discrete Diffusion
Subham Sekhar Sahoo and colleagues introduce Uno, a class of diffusion-augmented LLMs that keeps an autoregressive model's exact distribution while drawing several tokens per step from it, giving lossless speedups without a draft model.
Ask this paper
Parameter split, not architecture swap: AR weights train under standard next-token prediction; a lightweight set of diffusion weights is added by a cheap Diffusion Distillation phase, so an existing open-weight AR LLM can be upgraded rather than retrained.
Lossless by construction: the Psi-Spec sampler family draws multiple tokens in parallel from the AR distribution itself, so unlike diffusion LLMs there is no quality tax, and unlike speculative decoding there is no separate draft model to serve.
Throughput at every batch size: Uno beats leading speculative-decoding methods at every evaluated batch size and reaches up to 3x over the base AR model, including at the largest batch the device supports, which is where speculative methods usually lose their edge.
Small model beats big diffusion baselines: the 8B Uno outperforms the 26B DiffusionGemma and proprietary Mercury 2 across agentic tool use, coding and long-context reasoning benchmarks.
Why it matters: agent loops are latency-bound on serial decode. A drop-in that preserves the base model's distribution while tripling throughput is the rare efficiency win that requires no eval re-litigation.
Abstract
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 $3\times$ 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/