Unlocking Lossless Speedups in LLMs via Discrete Diffusion
Abstract
Diffusion-augmented autoregressive language models use parallel token sampling via distilled diffusion weights and a specialized sampler to accelerate inference without quality loss or draft models.
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/
Community
Isn't the core framework here exactly the same as Orthrus, released 4 months ago? https://arxiv.org/abs/2605.12825
Thanks for brining this to our attention.
The core idea of our method Uno is to achieve AR and diffusion generation keeping the architecture (with its attention) unchanged, but the paper that you shared changes the architecture by adding diffusion attention heads.
Not only that, it uses bidirectional attention for the diffusion blocks.
There are a lot of similarities, yes, but i respectfully disagree with this characterization that the core ideas are the same.
We’ll add a discussion in the next revision
Thanks. I am referring to the core framework, not the exact implementation. I don't think "unchanged" vs. "changed" architecture is an accurate distinction: both preserve a frozen Transformer backbone and augment it with trainable diffusion parameters—LoRA in Uno and separate attention in Orthrus. Also, causal vs bidirectional attention is just simply a design choice within the diffusion path.
At the framework level: keep the AR weights frozen, add and train a diffusion pathway, reuse the AR KV cache, draft multiple tokens in parallel, and verify them with the AR weights for lossless generation. That is the same core framework.
Your claim that Uno is the first diffusion-based model to achieve lossless generation without a separate drafter is not correct.
Diffusion LLMs have two major issues wrt AR LLMs:
- Poor quality,
- Slow inference at high batch sizes.
To overcome these limitations, 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 next-token prediction loss to define an AR distribution.
- Diffusion weights: Trained to generate multiple tokens in parallel from the AR distribution.
under a single unified LLM architecture.
We name our model Uno. To sample from it, we propose $\Psi$-Spec sampler, which enables provably lossless multi-token prediction from the AR distribution.
- Unlike speculative decoding methods,
Unodoes not require a separately trained draft model and achieves higher throughput across all batch sizes. - Unlike self-speculative decoding approaches,
Unois lossless.
Results
Our model Uno
- Achieves higher throughput than speculative-decoding methods at every batch size, and
- Delivers up to 3× 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.
Lossless is doing a lot of work in that title. Tokens per second on a clean benchmark doesn't tell me what happens at batch 32 on an H100 with KV cache pressure — that's where AR models are actually fast. The real question is denoising steps: how many passes does discrete diffusion need to match AR quality? If it's more than two or three, the parallel gain gets eaten by sampling overhead and you're back to trading latency for throughput. I'd want the wall-clock number at a realistic batch size before I believe "lossless speedup" survives contact with prod.
Models citing this paper 5
Browse 5 models citing this paperDatasets citing this paper 0
No dataset linking this paper
Spaces citing this paper 0
No Space linking this paper