PaperlayerView exact revision on arXiv

Source-linked AI summary

Nemotron-Labs-Diffusion: A Tri-Mode Language Model Unifying Autoregressive, Diffusion, and Self-Speculation Decoding

Yonggan Fu, Lexington Whalen, Abhinav Garg, Chengyue Wu, Maksim Khadkevich, Nicolai Oswald, Enze Xie, Daniel Egert, Sharath Turuvekere Sreenivas, Shizhe Diao, Chenhan Yu, Ye Yu, Weijia Chen, Sajad Norouzi, Jingyu Liu, Shiyi Lan, Ligeng Zhu, Jin Wang, Jindong Jiang, Morteza Mardani, Mehran Maghoumi, Song Han, Ante Jukić, Nima Tajbakhsh, Jan Kautz, Pavlo Molchanov

arXiv:2607.05722v1cs.CL

TL;DR

Existing AR and diffusion LMs face a trade-off between sequential inference, parallel decoding, accuracy, and practical efficiency. Nemotron-Labs-Diffusion unifies AR and diffusion objectives in a tri-mode model that switches among AR, diffusion, and self-speculation, achieving strong accuracy and speed while exposing further diffusion headroom. The family also supports flexible deployment across concurrency regimes, with self-speculation favored at low concurrency and AR suited to compute-bound high concurrency.

  • Problem

    AR decoding limits parallelism, while diffusion LMs can lag AR models in accuracy and learning efficiency and lack clear practical advantages over MTP methods.

  • Method

    Nemotron-Labs-Diffusion jointly optimizes AR and diffusion losses in one model and supports AR, parallel diffusion, and self-speculation decoding.

  • Results

    Across benchmarks, the 3B/8B/14B model family outperforms open-source AR and diffusion LMs in accuracy and speed; the 8B model delivers 6× more tokens per forward than Qwen3-8B and 4× higher SPEED-Bench throughput.

  • Takeaways & Limitations

    AR and diffusion can be unified, with self-speculation supporting low-concurrency deployment and AR remaining suitable for compute-bound high-concurrency scenarios.

  • Takeaways & Limitations

    Current confidence-based diffusion samplers remain far below the speed-of-light upper bound, leaving sampler development as an open challenge.

Abstract

from arXiv · show

We introduce Nemotron-Labs-Diffusion, a tri-mode language model (LM) that unifies AR, diffusion, and self-speculation decoding within a single architecture. Trained with a joint AR-diffusion objective, Nemotron-Labs-Diffusion can switch modes to sustain high throughput across deployment settings and concurrency levels. Our study shows that (1) AR and diffusion objectives are complementary: diffusion improves lookahead planning, while AR provides left-to-right linguistic priors. (2) In self-speculation mode, diffusion drafts while AR verifies, outperforming multi-token prediction (MTP) methods in both acceptance rate and real-device efficiency. (3) A speed-of-light analysis further demonstrates diffusion's long-term potential, with up to 76.5% more tokens per forward pass than self-speculation under an optimal sampler. Scaling to 3B, 8B, and 14B parameters, our Nemotron-Labs-Diffusion family, including base, instruct, and vision-language models, consistently outperforms state-of-the-art open-source AR and diffusion LMs in both accuracy and speed. For example, Nemotron-Labs-Diffusion-8B decodes 6x more tokens per forward than Qwen3-8B with comparable accuracy, translating to 4x higher throughput on SPEED-Bench with SGLang on a GB200 GPU.

1. Introduction

Nemotron-Labs-Diffusion addresses the throughput limits and accuracy concerns of existing AR and diffusion LMs by unifying their objectives and decoding modes. The resulting tri-mode models combine AR priors with diffusion parallelism and show strong accuracy, speed, deployment flexibility, and future potential.

  • AR decoding limits inference parallelism through strictly sequential token-by-token generation, especially in low-batch-size deployment.
  • Diffusion LMs enable parallel generation but can lag AR models in accuracy and learning efficiency, while often lacking practical advantages over MTP methods.
  • Nemotron-Labs-Diffusion jointly optimizes AR and diffusion losses and supports AR, parallel diffusion, and self-speculation decoding within one model.
  • The model family spans 3B, 8B, and 14B base, instruct, and vision-language variants, outperforming open-source AR and diffusion LMs in accuracy and inference speed across benchmarks.
  • 6× more tokens per forward than Qwen3-8B translates to 4× higher SPEED-Bench throughput on an NVIDIA GB200 GPU with comparable or better general-benchmark accuracy.
  • AR and diffusion objectives are complementary: diffusion supports parallel generation and lookahead planning, while AR supplies left-to-right linguistic priors across deployment regimes.
  • Under an optimal sampler, diffusion correctly predicts over 76.5% more tokens per forward pass than self-speculation, revealing unused parallelism in current samplers.
  • Self-speculation uses diffusion to draft multi-token candidates and AR to verify them, yielding higher acceptance rates and favorable throughput trade-offs than prior MTP approaches.

2. Tri-Mode LM Training

Nemotron-Labs-Diffusion trains AR and diffusion objectives jointly, using AR priors to support linguistically coherent parallel diffusion decoding. Its training pipeline combines block-wise diffusion, two-stage initialization, global loss averaging, and carefully weighted objectives, producing substantial accuracy gains while preserving AR performance.

  • Joint objectives: AR and diffusion objectives are complementary: AR supplies left-to-right linguistic priors, while diffusion supports lookahead planning and parallel prediction.The model uses a weighted combination of AR next-token loss and block-wise diffusion denoising loss.
  • Diffusion formulation: Block-wise diffusion denoises contiguous token blocks bidirectionally within each block while remaining causal across blocks.Previously generated blocks can reuse their KV cache during inference.
  • Optimization: The joint loss uses AR coefficient 1 and diffusion weight α=0.3, selected to balance loss magnitudes and maintain AR accuracy.Two-stage training first strengthens AR capabilities, then enables joint AR-diffusion optimization.
  • Optimization: Global loss averaging improves average accuracy by 2.12% by weighting token contributions across the batch and reducing variance from variable masking ratios.The ablation pipeline also adds DP-rank varying masking ratios and two-stage training.
  • Ablations: Two-stage training yields a 5.74% gain, while adding AR loss yields the largest single improvement of 7.48% in diffusion accuracy.The full pipeline improves baseline average accuracy by 16.05%.
  • Mutual effects: Both modes peak at α=0.3, and diffusion training preserves or slightly improves AR accuracy, supporting a complementary rather than zero-sum relationship.Reported AR improvements are 0.14% for the base model and 0.43% for the instruct model.

3. Tri-Mode LM Inference

Nemotron-Labs-Diffusion unifies AR, diffusion, and self-speculation decoding in one model, enabling distinct inference strategies for different deployment conditions. Diffusion drafts tokens in parallel, while AR verification supports speculative decoding with cached prefixes.

  • The joint AR–diffusion objective enables three inference modes: left-to-right AR, parallel diffusion, and self-speculation decoding.The model can switch among these modes within a single architecture.
  • AR decoding: AR decoding generates tokens sequentially with causal attention and is preferred under high concurrency.
  • Diffusion decoding: Diffusion decoding initializes masked blocks and iteratively denoises multiple positions in parallel using confidence-based decisions.Completed blocks refresh their KV cache before decoding proceeds.
  • Self-speculation: Self-speculation uses diffusion to draft k tokens, then AR to verify the draft in a second forward pass.The verifier accepts the longest matching prefix and can add one token at the first rejected position, producing 1 to k+1 tokens per step.
  • Self-speculation: LoRA enhances linear self-speculation by aligning the diffusion drafter with the unchanged AR verifier.The training loss combines distribution matching and token-level cross-entropy on accepted positions plus the first rejection.
  • Self-speculation: Quadratic self-speculation performs drafting and verification simultaneously in one forward pass with a structured attention mask, but incurs quadratic decoding cost.

4. Speed-of-Light Analysis

The speed-of-light analysis estimates diffusion’s maximum parallel-decoding potential by comparing parallel acceptance with a serial-denoising target. Diffusion approaches this ceiling in acceptance rate, while real tokens per forward remain lower for linear self-speculation because it verifies separately and accepts only prefixes.

  • The speed-of-light ceiling measures diffusion’s intrinsic parallel-decoding capability without an AR verifier.It provides a reference for parallelism left unused by current samplers.
  • SOL construction: Serial denoising fills each block one position per forward pass to define the oracle target sequence used by the SOL analysis.
  • SOL construction: Greedy parallel acceptance commits masked positions matching the serial target, with a highest-confidence fallback when no position matches.It is fast but can be inexact because simultaneous commits alter later conditional distributions.
  • SOL construction: Recursive dynamic compaction searches for the largest safe subset of matching positions to reproduce the serial target exactly.
  • SOL results: 7.60× average SOL acceptance grows from 2.89× at block length 4 to 7.60× at block length 32, exceeding 10× for multilingual and coding content.
  • SOL results: 6.82× versus 7.60× overall acceptance compares linear self-speculation with SOL, but real TPF is 3.41× versus 6.02×, a 76.5% gap.Two forward passes and prefix-only acceptance explain the lower real TPF of linear self-speculation.

5. Nemotron-Labs-Diffusion Family

The Nemotron-Labs-Diffusion family spans 3B, 8B, and 14B models, including base, instruct, and vision-language variants. Its vision-language extension adds visual components while retaining the joint training objective and tri-mode inference capability.

  • The model family includes 3B, 8B, and 14B base, instruct, and vision-language models.
  • Text models: Base models use two-stage training: pure AR continuous pretraining followed by joint AR–diffusion training.The supplied training setup uses 1T tokens in Stage 1 and 300B tokens in Stage 2.
  • Text models: Instruct models are produced by supervised fine-tuning base models with the joint AR–diffusion objective.The described SFT uses 45B tokens.
  • Vision-language models: The vision-language model adds a vision encoder and multimodal projector to the diffusion LM backbone while inheriting joint training and tri-mode inference.
  • Evaluation: The family is benchmarked against state-of-the-art AR and diffusion instruct models across scientific QA, instruction following, coding, and mathematics reasoning.
  • Vision-language models: The asymmetric dual-stream layout removes never-masked vision tokens from the noisy stream to reduce unnecessary attention computation.The clean stream retains vision tokens for visual context and AR conditioning.

6. Evaluation and Analysis

Across benchmarks and model scales, Nemotron-Labs-Diffusion improves accuracy–efficiency trade-offs in AR, diffusion, and self-speculation modes. Linear self-speculation is the preferred practical mode, while diffusion and quadratic self-speculation expose additional parallel-decoding potential.

  • Benchmark coverage: Nemotron-Labs-Diffusion-8B achieves higher accuracy and efficiency than state-of-the-art AR and diffusion instruct models across all evaluated modes.The evaluation covers scientific QA, instruction following, coding, and math reasoning.
  • Instruct-model benchmarks: +0.86% average accuracy over Qwen3-8B is achieved in AR mode, while diffusion mode reaches 2.57× TPF and +0.43% average accuracy.Compared with SDAR-8B Chat, diffusion mode also improves average accuracy by +9.09%.
  • Self-speculation: 5.99× TPF is achieved by LoRA-tuned linear self-speculation with accuracy comparable to diffusion mode, whereas quadratic self-speculation reaches 6.38× TPF but lower real-device efficiency.Linear self-speculation is therefore used by default.
  • Sampler analysis: 1.3× TPF at the same accuracy or +10.6% accuracy at the same TPF is obtained with the trained sampler, shifting the accuracy–TPF Pareto frontier upward.This result is reported for the instruct model with block size 32 across denoising thresholds.
  • Scaling: 4.36× and 5.96× TPF are achieved by linear self-speculation at 3B and 14B, respectively, while accuracy exceeds Qwen3-4B/14B by +1.77%/+1.19%.TPF broadly increases with model scale, consistent with stronger future-prediction abilities.
  • Real-device efficiency: 3.3× speedup over AR is achieved on GB200, while linear self-speculation also delivers 2.4×/2.3×/1.8× speedup over Eagle3 on GB200/RTX Pro 6000/DGX Spark.The reported mechanism is diffusion drafting paired with AR verification, which produces higher acceptance length than auxiliary-head MTP.

7. Related Work

Related work develops diffusion language models for parallel generation and explores their scaling, modalities, and acceleration. However, practical speedups remain constrained by attention, caching, and denoising-process limitations.

  • Diffusion language models: Diffusion language models perform non-AR decoding to enable parallel token generation and include continuous, discrete, and masked variants.Masked diffusion models such as LLaDA and Dream have been scaled substantially.
  • Capabilities and extensions: Prior work reports diffusion advantages in data-constrained learning, planning, and text embedding, alongside larger-scale, domain-specific, reinforcement-learning, and multimodal extensions.These directions broaden diffusion LMs beyond general text generation.
  • Acceleration: Caching strategies and confidence-based sampling address diffusion acceleration, but bidirectional attention and one-token-per-step denoising still limit achievable speedups.These constraints motivate further work on practical parallel decoding.

8. Insights and Future Directions

Nemotron-Labs-Diffusion unifies AR, diffusion, and self-speculation through joint training rather than treating AR and diffusion as competing paradigms. The resulting design links AR priors, parallel prediction, and verification within one model.

  • Tri-mode generation: Joint AR–diffusion optimization produces AR, diffusion, and self-speculation modes without mode-specific architectural modifications.The model supports both causal and non-AR parallel token prediction within one architecture.
  • Complementary objectives: AR and diffusion losses are complementary: AR supplies left-to-right linguistic priors, while diffusion preserves or slightly improves AR accuracy through better future planning.The two objectives mutually benefit each other and peak at α=0.3.
  • Self-speculation: Self-speculation uses diffusion to generate multi-token drafts and AR verification to ensure correctness, outperforming MTP methods in acceptance rates and efficiency.The approach avoids relying solely on auxiliary prediction heads.

4. Variance reduction is critical for diffusion

The analysis identifies both practical constraints and future opportunities for diffusion decoding. Current methods leave substantial headroom, while verification and generation remain limited by prefix-wise and token-level processing.

  • Current efficiency: Linear self-speculation is currently the most efficient mode, while quadratic self-speculation offers higher TPF per step but depends on improved infrastructure support.The comparison reflects current implementation and kernel efficiency.
  • Speed-of-light potential: 76.5% more tokens per forward pass than linear self-speculation is the speed-of-light potential under an optimal diffusion sampler.This indicates a higher theoretical upper bound for diffusion decoding than prefix-based speculative decoding.
  • Sampler limitations: Current confidence-based samplers remain far from the speed-of-light upper bound, motivating optimized samplers and training schemes for more aggressive parallel sampling.The stated direction is to improve reliable token identification and acceptance policy learning.
  • Verification boundary: Prefix-wise AR verification does not fully exploit non-AR diffusion drafts, motivating verification of multiple non-contiguous drafted tokens.Diffusion-mode verification is proposed as one possible direction.
  • Parallelism boundary: Diffusion generation remains mainly token-level despite parallel prediction, leaving segment-level and paragraph-level parallelism as an open direction.Future training algorithms could target broader global planning parallelism.

A. Diffusion Sampler Details

The sampler is a lightweight module built over a frozen diffusion backbone to improve diffusion decoding decisions. It is trained from denoising trajectories using confidence-based and hybrid policies with per-position supervision.

  • Architecture and features: The sampler operates on a frozen backbone with negligible parameter overhead of approximately 0.06%, or 4.8M parameters relative to the 8B backbone.It is illustrated in Figure 10.
  • Architecture and features: The sampler is a four-layer bidirectional Transformer with hidden dimension d=384 and a per-position sigmoid prediction head.Each position uses a 144-dimensional feature representation.
  • Architecture and features: Each position combines PCA-compressed embeddings of the top-3 predictions with distribution statistics including top-1 probability, margin, top-3 mass, and entropy.The model’s own top-1 prediction embedding is the dominant semantic feature.
  • Trajectory collection: Sampler training uses approximately 20M denoising trajectories from math, code, STEM, and chat data at block lengths B∈{8,32}.The trajectories are collected from Nemotron-Labs-Diffusion-8B.
  • Trajectory collection: The trajectories use standard confidence decoding or a hybrid policy that commits already-correct top-1 predictions before applying confidence decoding.Per-position labels indicate whether the current top-1 prediction matches the token ultimately committed under the same policy.
  • Evaluation: The sampler uses early stopping, and its accuracy–TPF gains over confidence thresholding are reported in Section 6.1 and Figure 8.The supplied passages identify the comparison but do not provide its numerical gains.

B. LoRA-Enhanced Linear SS

The supplied passages identify a visualization of enhanced linear self-speculation with LoRA. They do not describe the training procedure or its measured effects.

  • The paper provides a visualization of enhanced linear self-speculation with LoRA in Figure 11.
  • The supplied passages do not report quantitative results for the LoRA-enhanced linear self-speculation method.
  • Figure 11 concerns the linear self-speculation mode rather than the diffusion sampler itself.

C. Quadratic SS Details

Quadratic self-speculation interleaves speculative tokens with fresh masks so one forward pass can verify prior drafts while generating new drafts. The design preserves fresh speculative outputs even when verification stops early.

  • Input layout: Quadratic self-speculation reuses k speculative tokens and inserts k fresh mask tokens after each one, producing k^2 inserted masks.The next autoregressive token is immediately verified.
  • Joint drafting and verification: A structured attention mask produces causal predictions for verifying the previous draft and parallel predictions for newly inserted masks in the same forward pass.The newly masked predictions become the next iteration’s draft tokens.
  • Verification: Verification accepts the longest prefix satisfying the criterion, commits accepted tokens, and stops at the first mismatch.
  • Joint drafting and verification: The interleaved layout keeps newly inserted mask positions productive even after an early verification failure, yielding k tokens to verify per iteration.
  • Verification signals: Verification can use causal AR predictions, diffusion denoising predictions, or an interpolated AR–diffusion ensemble controlled by λ∈[0,1].The diffusion verifier uses the denoising logits at the first newly inserted mask position following each speculative token.
  • Visualization: Figure 12 illustrates quadratic self-speculation with simultaneous drafting and verification.
Loading 2607.05722v1…