Unlocking Lossless Speedups in LLMs via Discrete Diffusion

Unlocking Lossless Speedups in LLMs via Discrete Diffusion. By Subham Sekhar Sahoo and sixteen others, from the Institute of Foundation Models, the University of Illinois Urbana-Champaign, and three other institutions. Published on the third of September. Its code repository is worked out: about ninety source files, with a readme and a licence.

Large language models owe much of their success to next-token prediction. The same autoregressive structure that makes that objective work also makes generation slow, because tokens have to be produced one after another. That sequential dependency is the bottleneck: each token waits on the one before it. The two familiar escape routes each carry a cost. Speculative decoding needs a separate draft model. Diffusion language models draw tokens in parallel, but at the price of the quality of the autoregressive model. One option adds machinery, and the other gives up the thing that made the model good in the first place.

The proposal is a new class of model, the diffusion-augmented LLM. It defines an ordinary autoregressive model distribution, and uses diffusion only as a way to draw several tokens from that distribution at once. The distribution stays autoregressive. Diffusion is the sampling engine, not the model of language.

The parameters split into two sets. The autoregressive weights are trained with the standard next-token objective. A lightweight set of diffusion weights is trained to generate multiple tokens simultaneously, and it is learned in a simple phase called Diffusion Distillation, which adds negligible overhead to existing LLM training pipelines.

The second contribution is the sampler. Psi-Spec is a family of samplers, and it provides lossless acceleration, meaning the output still comes from the underlying autoregressive distribution rather than an approximation of it. It also allows inference-time scaling at a fixed context length. The resulting models are called Uno. They can be trained from scratch, or built by augmenting existing open-weight autoregressive LLMs.

Uno achieves higher throughput than leading speculative-decoding methods at every batch size evaluated, and delivers speedups of up to three times over its base autoregressive model. That holds including at the largest batch size the device supports.

On quality, the eight billion parameter Uno model outperforms the leading open diffusion LLM, the twenty-six billion parameter DiffusionGemma, and also the proprietary Mercury 2, on every benchmark evaluated in agentic tool use, coding and long-context reasoning. Code and checkpoints are released.

The trick is where the line is drawn. Previous routes to parallel generation either added a second model or changed what the model was. Uno keeps the autoregressive model as it is and trains a small companion that knows how to sample from it several tokens at a time. The speed comes from the sampler, and the quality stays with the original. Because the companion weights can be added to an existing open-weight model, the approach does not demand a fresh start. And because the phase that trains them adds negligible overhead, it does not demand a large new budget either.