Summary

• Core Contribution: Introduces diffusion-augmented LLMs, a new class of models that decouple generation quality (autoregressive weights) from generation speed (lightweight diffusion weights), enabling lossless parallel token generation.
• Key Innovation: The Ψ-Spec sampler enables lossless acceleration via AR-verified rejection sampling, requiring no separate draft model and preserving the base AR model's output distribution exactly.
• Main Results: Achieves up to 3× speedup over base AR models, with speedups persisting at the largest feasible batch sizes (up to 2× at batch 64), outperforming speculative decoding (EAGLE-3, DFlash) and diffusion LLMs (DiffusionGemma, Nemotron-Labs-Diffusion, Mercury 2).
• Versatility: Can be trained from scratch or created by augmenting existing open-weight AR LLMs (e.g., Qwen3-8B), with diffusion weights trainable on different data distributions than AR weights.
• Practical Impact: Accelerates both inference and RL post-training (up to 40% end-to-end training speedup), with a single shared KV cache and minimal additional parameters (0.35B for 8B models).

Introduction and Theoretical Foundation

Background

Large Language Models (LLMs) rely on next-token prediction (NTP), which requires sequential, one-token-at-a-time generation. This autoregressive (AR) structure is increasingly costly as reasoning traces grow longer, creating bottlenecks in serving latency and RL post-training (where rollout generation dominates runtime). Language contains predictable collocations that could be generated in blocks, but standard LLMs cannot exploit this redundancy. Decoding is often memory-bound, especially at long context lengths, leaving GPUs underutilized.

Limitations of Existing Approaches

• Speculative decoding (Leviathan et al., 2023; Chen et al., 2023) requires a separate, well-aligned draft model and maintains multiple KV caches.
• Discrete diffusion LLMs (d-LLMs) support parallel generation but face a quality-speed tradeoff relative to AR models, and their speedups vanish at large batch sizes.
• Multi-token prediction (MTP) methods modify the LLM architecture with additional prediction heads.

Theoretical Foundation

The paper builds on:

• Autoregressive factorization: $\log p_{\theta}(\mathbf{x}) = \sum_{\ell=1}^{L} \log p_{\theta}(\mathbf{x}^{\ell} | \mathbf{x}^{<\ell})$
• Discrete diffusion with interpolating forward process (Sahoo et al., 2024a): $\mathbf{z}_t^{\ell} \sim \operatorname{Cat}(\cdot; \alpha_t \mathbf{x}^{\ell} + (1-\alpha_t)\boldsymbol{\pi})$
• Ψ-samplers (Deschenaux et al., 2026) for discrete diffusion sampling
• Discrete Consistency Distillation (DCD) (Sahoo et al., 2025a) for compressing multistep diffusion into few-step generators

Methodology

Architecture: Diffusion-Augmented LLMs

Each layer contains two decoupled weight sets:

• AR weights ($\theta_{\mathrm{AR}}$): Trained with standard NTP objective, determine response quality
• Diffusion weights ($\theta_{\Delta}$): Lightweight LoRA adapters (rank-128, LoRA-α=256), trained for parallel generation

The diffusion pathway uses $\theta_{\mathrm{AR}} + \theta_{\Delta}$ for drafting, while verification uses only $\theta_{\mathrm{AR}}$.

Training: Diffusion Distillation Phase

The training objective combines two losses:

$$\mathcal{L}(\theta_{\Delta}; \theta_{\mathrm{AR}}, \alpha, \beta) = \mathbb{E}_{\mathbf{x} \sim \mathcal{D}, \mathbf{z}_1 \sim \boldsymbol{\pi}^L} \left[\alpha \mathcal{L}_{\mathrm{DCD}}(\theta_{\Delta}; \theta_{\mathrm{AR}}, \mathbf{x}, \mathbf{z}_1) + \beta \mathcal{L}_{\mathrm{TV}}(\theta_{\Delta}; \theta_{\mathrm{AR}}, \mathbf{x}, \mathbf{z}_1)\right]$$

• Blockwise DCD loss: Distills the AR distribution into a single-step diffusion generator, using block-causal attention masks and gated LoRA to compute teacher/student logits in one forward pass
• Total Variation loss: Minimizes blockwise TV distance between diffusion and AR distributions to increase acceptance rates

Ψ-Speculative Sampler

• Linear Sampler: Samples single candidate sequence, optimized for system throughput at high batch sizes
• Tree Sampler: Samples multiple candidates via tree-based selection (top-K at each position), optimized for per-request throughput at low batch sizes

The sampler uses a single diffusion forward pass to draft B tokens, then verifies via standard speculative decoding rejection sampling, preserving the AR distribution exactly. TPF bounded by $1 \leq \mathrm{TPF} \leq \frac{B+1}{2}$.

Empirical Validation / Results

Uno (8B, trained from scratch)

Key Results (Table 1):

• Outperforms DiffusionGemma (26B-A4B) and Nemotron-Labs-Diffusion (14B) on ALL benchmarks
• Outperforms proprietary Mercury 2 on all agentic, coding, and long-context benchmarks
• System throughput: 5255 tokens/s (vs. 1197 for Mercury 2, 1136 for DiffusionGemma, 2794 for Nemotron)
• Up to 2.2× speedup at batch size 1, 1.5× at batch size 64 (largest supported by base AR)

Uno_Qwen (Qwen3-8B augmented)

Comparison with Speculative Decoding (Table 2):

• System throughput: 5733 tokens/s (vs. 4944 EAGLE-3, 5351 DFlash)
• Per-request throughput: 445 tokens/s (vs. 284 EAGLE-3, 370 DFlash)
• Fewer additional parameters (0.35B vs. 0.40B EAGLE-3, 1.05B DFlash)
• Lower peak memory (122.2 GiB vs. 130.0 GiB)

Comparison with Lossy Diffusion Methods (Table 3):

• Achieves higher TPF than most methods (e.g., 3.56-4.21 TPF vs. 1.5-7.3 for lossy methods)
• Maintains high accuracy without the degradation seen in lossy methods

Ablations

• TV loss alone (TPF 2.39) outperforms DCD + TV (2.23) and DCD alone (2.23)
• Increasing block size curriculum improves TPF (2.65 → 2.71)
• Distributing adapters across all layers is more effective than concentrating them
• LoRA rank 256 improves TPF over rank 128

Theoretical and Practical Implications

Theoretical Significance

• Lossless acceleration: Provably preserves the AR model's output distribution via rejection sampling, unlike lossy d-LLMs
• Unified architecture: Shows AR and diffusion pathways can coexist in a single model with decoupled parameters
• Training efficiency: Diffusion weights require orders of magnitude fewer tokens than AR weights (7B vs. 23T tokens)

Practical Implications

• Drop-in replacement: Can replace speculative decoding (EAGLE-3, DFlash) and self-speculative decoding (TiDAR) without separate draft models
• RL acceleration: Up to 40% end-to-end RL training speedup, with speedups retained even as AR weights are updated during RL
• Memory efficiency: Single shared KV cache reduces peak memory vs. methods with separate draft/verifier caches
• Serving efficiency: Speedups persist at all batch sizes, including the largest feasible, making it practical for agentic workloads with high concurrency

Conclusion

The paper introduces diffusion-augmented LLMs as a new design paradigm that unifies AR quality with diffusion parallelism. Key takeaways:

• Uno achieves lossless speedups up to 3× over base AR models, outperforming both speculative decoding and d-LLM baselines
• The framework supports both from-scratch training and augmentation of existing open-weight models
• Diffusion weights remain effective even as AR weights are updated during RL post-training
• Future directions include combining with MTP methods, quadratic samplers for single-pass draft+verify, and systematic exploration of inference-time scaling via additional denoising steps

The approach expands the design space of AR language models by showing that quality and speed can be decoupled within a single architecture, with the diffusion pathway serving as a lightweight, trainable accelerator that preserves the base model's distribution exactly.