# Unlocking Lossless Speedups in LLMs via Discrete Diffusion

> Diffusion-augmented LLMs decouple generation quality from speed, enabling lossless parallel token generation with up to 3x faster inference via rejection sampling.

- **Source:** [arXiv](https://arxiv.org/abs/2609.04010)
- **Published:** 2026-09-09
- **Permalink:** https://picx.dev/p/E38hkk
- **Whiteboard:** https://picx.dev/p/E38hkk/image

## Summary

## 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.

---

_Markdown view of https://picx.dev/p/E38hkk, served by PicX — AI-generated visual whiteboard summaries of research papers._