Technical Summary: Nemotron-Labs-Diffusion

Problem Statement

Autoregressive (AR) language models (LMs) are fundamentally limited by their strictly sequential, token-by-token decoding process, which restricts inference parallelism and leads to resource under-utilization, particularly in low-batch-size scenarios. While diffusion LMs have emerged as a promising alternative by enabling parallel generation through multi-token decoding per forward pass, they currently lag behind AR models in accuracy and learning efficiency. Existing diffusion LMs often require substantially more data to reach comparable performance and struggle to outperform Multi-Token Prediction (MTP) methods in practical efficiency-accuracy trade-offs. Furthermore, the relationship between AR and diffusion paradigms remains unclear: it is uncertain whether they are competing approaches or if they can be harmonized to leverage the strengths of both.

Methodology

The authors introduce Nemotron-Labs-Diffusion, a tri-mode language model family that unifies AR, diffusion, and self-speculation decoding within a single architecture. The core methodology involves a joint training framework and a flexible inference strategy.

1. Joint AR-Diffusion Training

The model is trained using a weighted combination of an AR next-token loss and a block-wise diffusion denoising loss:
$$\mathcal{L}(\theta) = \mathcal{L}_{AR}(\theta) + \alpha \mathcal{L}_{diff}(\theta)$$

• Two-Stage Training: To establish strong left-to-right linguistic priors, the authors employ a two-stage strategy. Stage 1 trains exclusively on the AR objective ($\alpha=0$). Stage 2 introduces diffusion supervision ($\alpha=0.3$) to integrate parallel decoding capabilities without overwriting AR priors.
• Global Loss Averaging: To stabilize optimization given the variable number of masked tokens in diffusion training, the authors utilize a global loss averaging strategy across the batch, rather than sequence-wise averaging, to prevent gradient variance from samples with few noisy tokens.
• Attention Pattern: The model utilizes a dual-stream input (noisy and clean views) with a structured attention mask. Crucially, the clean stream enforces a strictly causal mask, allowing the AR objective to be computed on clean tokens without label leakage, while the noisy stream uses bidirectional attention within blocks for parallel denoising.

2. Tri-Mode Inference

The trained model supports three distinct decoding modes, selectable at inference time:

• AR Mode: Standard left-to-right generation using causal attention. This mode is optimized for high-concurrency cloud serving.
• Diffusion Mode: Parallel block-wise denoising. The model iteratively denoises multiple tokens per step based on confidence thresholds. A lightweight, trained sampler is used to predict whether top-1 predictions are correct, optimizing the trade-off between tokens per forward (TPF) and error rate.
• Self-Speculation Mode: A hybrid approach where the diffusion pathway drafts multiple candidate tokens in parallel, and the AR pathway verifies them.
  • Linear Self-Speculation: Diffusion drafts $k$ tokens; AR verifies them in a second forward pass. The authors enhance this by tuning a LoRA adapter on the diffusion drafter to better align with the AR verifier.
  • Quadratic Self-Speculation: A variant using quadratic decoding to prepare for all possible acceptance positions within a single forward pass, though it currently faces efficiency challenges due to kernel optimization.

3. Speed-of-Light (SOL) Analysis

The authors propose a "Speed-of-Light" analysis to quantify the theoretical upper bound of diffusion decoding. By using a recursive dynamic compaction method to identify the largest safe subset of matching positions against a serial denoising target, they estimate the maximum achievable TPF without AR verification. This analysis reveals that current confidence-based sampling leaves significant parallelism on the table.

Key Contributions

1. Unified Tri-Mode Architecture: The first model to seamlessly integrate AR, diffusion, and self-speculation decoding, allowing dynamic switching based on deployment needs (e.g., high concurrency vs. low concurrency).
2. Complementarity of Objectives: Demonstrates that AR and diffusion objectives are complementary rather than competing. AR provides strong linguistic priors and future planning, while diffusion unlocks parallel generation. The joint objective preserves or slightly improves AR accuracy while enabling high-speed diffusion modes.
3. Superiority over MTP: Shows that self-speculation (diffusion drafting + AR verification) outperforms existing MTP methods (like Eagle3) in both acceptance rates and real-device efficiency, particularly at low concurrency.
4. Scalability and Versatility: The model family scales from 3B to 14B parameters and extends to Vision-Language Models (VLMs), consistently outperforming state-of-the-art open-source AR and diffusion LMs.

Results

• Accuracy and Efficiency: The Nemotron-Labs-Diffusion-8B model achieves comparable or better accuracy than Qwen3-8B across general, coding, and math benchmarks while decoding 6× more tokens per forward pass in diffusion mode.
• Throughput Gains: On the SPEED-Bench, the model achieves 4× higher throughput than Qwen3-8B on NVIDIA GB200 GPUs using SGLang.
• Self-Speculation Performance: Linear self-speculation with LoRA tuning achieves an average TPF of 5.99× (8B model) and 6.38× (14B model), significantly outperforming baselines.
• SOL Potential: The SOL analysis indicates that diffusion decoding has the potential to predict 76.5% more tokens per forward pass than current self-speculation methods under an optimal sampler, highlighting substantial future headroom.
• VLM Performance: The VLM variant maintains competitive accuracy with AR baselines while offering 2.46×–3.15× TPF in diffusion mode and up to 7.45× TPF in linear self-speculation mode for long responses.

Significance and Claims

The paper claims that Nemotron-Labs-Diffusion represents a paradigm shift where AR and diffusion are not mutually exclusive but can be unified to create a more flexible and efficient inference engine. The authors assert that:

• Harmonization is Possible: Unifying AR and diffusion objectives is a promising pathway that avoids the accuracy penalties typically associated with pure diffusion models.
• Deployment Flexibility: The tri-mode design allows a single model to serve as a drop-in replacement for AR models in high-concurrency settings while offering significant speedups in low-concurrency personal AI scenarios via self-speculation.
• Future Potential: The "Speed-of-Light" analysis suggests that the full potential of diffusion decoding is not yet realized by current samplers. The gap between current performance and the SOL ceiling indicates that future work in sampler optimization could unlock even greater parallelism, potentially surpassing the limitations of prefix-based speculative decoding.

The work concludes that while current infrastructure favors linear self-speculation, the theoretical upper bounds of diffusion decoding suggest a long-term trajectory toward highly parallel, non-prefix generation that could fundamentally alter LLM inference efficiency.