One-sentence Summary

Researchers from the University of Illinois Urbana-Champaign, Cornell Tech, Harvard University, and Cerebras Systems introduce Uno, a diffusion-augmented LLM class that decouples AR weights from lightweight diffusion weights to draw multiple tokens in parallel, achieving up to $3\times$3× speedups over base AR models without draft models or quality loss, and outperforming larger d-LLMs in agentic tool use, coding, and long-context reasoning.

Key Contributions

• Introduces diffusion-augmented LLMs, a model class that defines an autoregressive distribution while using lightweight diffusion weights to draw multiple tokens in parallel, with these weights trained via a simple Diffusion Distillation phase that adds negligible overhead to existing LLM training pipelines.
• Presents Ψ-Spec, a family of samplers enabling lossless acceleration and inference-time scaling at a fixed context length, requiring no separate draft model and preserving the base AR model’s quality, unlike speculative decoding or lossy diffusion LLMs.
• The resulting Uno models, trainable from scratch or by augmenting open-weight AR LLMs, achieve up to 3× speedups over the base AR model at every evaluated batch size, outperform leading speculative-decoding methods in throughput, and an 8B Uno model surpasses the 26B DiffusionGemma and proprietary Mercury 2 across agentic tool use, coding, and long-context reasoning benchmarks.

Introduction

Large language models (LLMs) achieve strong performance through next-token prediction (NTP), but this objective forces autoregressive (AR) decoding, which generates one token per step. This sequential process becomes a bottleneck for long reasoning traces, increasing latency and slowing reinforcement learning (RL) post-training, where rollout generation dominates runtime. Additionally, decoding is often memory-bound, especially at long context lengths, leaving GPUs underutilized because moving model weights and key-value states limits speed. Existing acceleration methods have notable drawbacks: speculative decoding requires a separate draft model, discrete diffusion models offer lossy speedups that vanish at large batch sizes, and multi-token prediction methods modify the architecture with extra heads.

The authors introduce Uno, a diffusion-augmented LLM that unifies AR and diffusion weights within a single architecture. The core idea is to define a high-quality AR distribution and learn to sample multiple tokens in parallel from that same distribution. They achieve this by augmenting each layer with lightweight diffusion weights alongside standard AR weights, trained via diffusion distillation after freezing the AR weights. This design decouples response quality from generation speed, avoiding the need for a separate draft model or lossy AR-to-diffusion conversion. The proposed Ψ-Spec sampler enables lossless, AR-verified acceleration and inference-time scaling. Uno achieves up to a 2x speedup at the largest batch size supported by the base AR model, accelerating both inference and end-to-end RL training, and outperforms speculative decoding methods like EAGLE-3 and DFlash, as well as open-weight and proprietary d-LLMs, across all evaluated batch sizes.

Method

The authors propose a diffusion-augmented LLM framework that decouples generation quality from generation speed. The architecture augments each layer of a standard autoregressive model with a separate set of diffusion weights dedicated to parallel token generation. This design introduces two distinct pathways: autoregressive weights, which determine response quality and are trained via standard next-token prediction, and diffusion weights, which accelerate inference by generating tokens in parallel. At generation time, both sets of weights draft tokens concurrently, after which the autoregressive weights alone verify the drafts. This separation preserves the established autoregressive training pipeline while enabling lossless acceleration without sacrificing response quality.

The base autoregressive weights, denoted as $\theta_{AR}$θAR​, follow the standard causal masking paradigm where the distribution of the $\ell$ℓ-th token is modeled as:

$$p_{\theta_{AR}}(x^\ell \mid x^{<\ell}) = x_{\theta_{AR}}^{\ell-1}(x)$$pθAR​​(xℓ∣x<ℓ)=xθAR​ℓ−1​(x)

The diffusion weights, denoted as $\theta_{\Delta}$θΔ​, are parameterized as Low-Rank Adaptation adapters attached to each autoregressive weight matrix. This parameterization ensures the draft distribution remains tightly coupled to the verification distribution while adding minimal memory overhead. During drafting, the model utilizes the combined weights $\theta_{AR} + \theta_{\Delta}$θAR​+θΔ​, whereas the verification pathway relies exclusively on $\theta_{AR}$θAR​. The diffusion pathway retains the next-token prediction parameterization, predicting the subsequent clean token given a noisy sequence.

To train the diffusion parameters, the authors introduce a Diffusion Distillation Phase that aligns the single-step diffusion distribution with the autoregressive distribution over a token block. They adapt Discrete Consistency Distillation to approximate the autoregressive distribution through single-step block denoising. The training objective combines a distillation loss and a total variation loss:

$$\mathcal{L}(\theta_{\Delta}; \theta_{AR}, \alpha, \beta) = \mathbb{E}_{x \sim \mathcal{D}, z_1 \sim \pi^L} \left[ \alpha \mathcal{L}_{DCD}(\theta_{\Delta}; \theta_{AR}, x, z_1) + \beta \mathcal{L}_{TV}(\theta_{\Delta}; \theta_{AR}, x, z_1) \right]$$L(θΔ​;θAR​,α,β)=Ex∼D,z1​∼πL​[αLDCD​(θΔ​;θAR​,x,z1​)+βLTV​(θΔ​;θAR​,x,z1​)]

To make single-step denoising tractable, the sequence is partitioned into blocks. The authors employ a gated LoRA mechanism during a single forward pass over the concatenated clean and noisy sequences. This mask disables adapters at clean-sequence positions to compute teacher logits using only $\theta_{AR}$θAR​, and enables them at noisy-sequence positions to compute student logits using both $\theta_{AR}$θAR​ and $\theta_{\Delta}$θΔ​. The blockwise distillation loss minimizes the Kullback-Leibler divergence between the student and teacher distributions. Additionally, the total variation loss minimizes the distance between the diffusion and autoregressive distributions to increase the expected length of the accepted draft prefix during rejection sampling.

The training order of the autoregressive weights and diffusion adapters is structured to support different deployment goals. The authors outline two primary training regimes. If the sole objective is faster inference, autoregressive pre-training and post-training are completed first, after which the autoregressive weights are frozen to train the diffusion adapters. Alternatively, to accelerate both reinforcement learning rollouts and inference, diffusion distillation is applied after supervised fine-tuning but before reinforcement learning post-training.

As shown in the figure below:

This curriculum allows the resulting adapters to accelerate rollout generation during the subsequent reinforcement learning phase. Standard policy-optimization recipes update the base autoregressive weights while keeping the diffusion adapters frozen. Despite the base weights changing during reinforcement learning, the tight coupling maintained by the Low-Rank Adaptation parameterization ensures that the draft distribution does not drift significantly from the verifier distribution, thereby retaining the inference speedups.

During inference, the authors introduce the Ψ-Speculative sampler to draw multiple tokens in parallel while strictly sampling from the autoregressive distribution. The sampler uses the diffusion pathway to propose a block of tokens and performs rejection sampling against the base autoregressive distribution to accept the longest valid prefix. To balance acceptance length and verification cost, the framework supports two candidate sampling strategies. The Linear Sampler generates a single candidate sequence directly from the marginal distribution, optimizing aggregate system throughput at high batch sizes. The Tree Sampler exploits underutilized compute at low batch sizes by sampling multiple candidates and verifying them concurrently as a prefix tree using tree attention. In both cases, the first token is generated using the base autoregressive weights to guarantee exact matching with the verifier distribution, ensuring lossless speculative speedups.

Experiment

The experiments evaluate Uno, a diffusion-augmented LLM that combines autoregressive weights with rank-128 LoRA diffusion adapters, in two settings: training from scratch on proprietary data and augmenting the open-weight Qwen3-8B model with adapters trained on a different data distribution (OpenThoughts). Using a standardized 1K/8K throughput test to control for context length and output size, Uno matches the base AR model's quality while achieving 1.5x to 2.2x higher throughput across batch sizes, outperforms open-weight diffusion baselines (DiffusionGemma, Nemotron-Labs-Diffusion) on all benchmarks, and exceeds the proprietary Mercury 2 in system throughput by roughly 4.6x. When applied to Qwen3-8B, Uno surpasses speculative decoding baselines EAGLE-3 and DFlash in speed across all batch sizes with fewer added parameters and a shared KV cache, while ablations show that using only the total variation loss, a block-size curriculum, and adapters distributed across all layers yield the best tokens-per-forward-pass performance.

Uno consistently outperforms both open-weight diffusion models and the proprietary Mercury 2 across agentic and coding benchmarks, with the largest gains on agentic tasks. Uno also achieves higher system throughput than the open-weight baselines, despite using full attention, while DiffusionGemma is only faster at batch size 1 but with substantially lower accuracy. Uno exceeds Mercury 2 on all agentic and coding benchmarks where both are evaluated. Uno outperforms DiffusionGemma and Nemotron-Labs-Diffusion on every task, with especially large margins on agentic tasks. Uno achieves higher system throughput than both open-weight baselines, despite DiffusionGemma using strided attention. DiffusionGemma is faster than Uno only at batch size 1, but its accuracy is much lower.

Uno_Qwen achieves higher acceptance lengths than EAGLE-3 and DFlash across all tested math benchmarks under both system-optimal and per-request-optimal settings. It also delivers superior system throughput and per-request throughput while using fewer additional parameters and a shared KV cache, reducing peak memory usage compared to baselines. Uno_Qwen consistently reports larger acceptance lengths than EAGLE-3 and DFlash on all math benchmarks for both throughput-optimal configurations. The largest system throughput is achieved at batch size 4, where Uno_Qwen exceeds 5700 tokens per second, outperforming both baselines and achieving a 1.6x speedup over the base autoregressive model. For per-request throughput, Uno_Qwen and EAGLE-3 are the best, but Uno_Qwen achieves a 2.5x speedup over the base model, surpassing the baselines. Uno_Qwen uses a shared draft-verifier KV cache, unlike EAGLE-3 and DFlash which maintain separate caches, leading to lower peak memory usage. Uno_Qwen is strictly faster than DFlash and EAGLE-3 across all batch sizes while requiring fewer additional parameters.

The table compares the lossless Uno method against several lossy diffusion-based acceleration methods across math benchmarks, reporting accuracy and tokens per forward pass. Uno consistently achieves higher accuracy and TPF than most lossy methods, despite prioritizing quality over speed. Lossy methods often show accuracy drops relative to their parent models, while Uno maintains lossless performance. Uno achieves higher TPF than most lossy diffusion methods across the evaluated benchmarks. Lossy methods like TiDAR and Jacobi Forcing show substantial accuracy degradation on AIME-24 and AIME-25 compared to Uno. Uno maintains accuracy above 96% on GSM8K and MATH500, while several lossy methods fall below 90% on these tasks. Fast-dLLM v2 and LLaDA2.1-Flash show competitive accuracy on some benchmarks but are not consistently better than Uno.

Uno outperforms open-weight diffusion models and Mercury 2 on agentic and coding benchmarks, with the largest gains on agentic tasks, while achieving higher system throughput than open-weight baselines despite using full attention. Uno_Qwen also surpasses EAGLE-3 and DFlash in acceptance lengths and throughput across math benchmarks under both system-optimal and per-request-optimal settings, using fewer parameters and a shared KV cache to reduce peak memory. Compared to lossy diffusion-based acceleration methods, Uno maintains lossless accuracy and higher tokens per forward pass on math benchmarks, while several lossy methods show substantial accuracy drops on AIME-24 and AIME-25.