PaperlayerView exact revision on arXiv

Source-linked AI summary

Unlocking Lossless Speedups in LLMs via Discrete Diffusion

Subham Sekhar Sahoo, Lingjie Chen, Khiem Pham, Jonathan Geuter, Chaitanya Dwivedi, Varad Pimpalkhute, Yash Akhauri, Alexander Moreno, Mikhail Yurochkin, Zhenting Wang, Mostafa Elhoushi, Nolan Dey, Shane Bergsma, Joel Hestness, John Thickstun, Eric Xing, Zhengzhong Liu

arXiv:2609.04010v1cs.LG

TL;DR

LLMs rely on next-token prediction, but its autoregressive structure makes generation sequential and costly. The paper adds lightweight diffusion weights and Ψ-Spec samplers to draw multiple tokens in parallel from the AR distribution. Uno reaches up to a 3× speedup over the base AR model while retaining up to a 2× speedup at the largest supported batch size.

  • Problem

    Next-token prediction gives LLMs strong capabilities but forces one-token-at-a-time decoding, creating costly sequential inference and slowing rollout-heavy RL post-training.

  • Method

    Diffusion-augmented LLMs decouple NTP-trained AR weights from lightweight diffusion weights, then use Diffusion Distillation and Ψ-Spec to generate token blocks in parallel from the AR distribution.

  • Results

    Uno achieves up to a 3× speedup over the base AR model across evaluated batch sizes, retaining up to a 2× speedup at the largest supported batch size.

  • Takeaways & Limitations

    Uno combines AR generation quality with diffusion parallelism in one architecture, offering a drop-in alternative to speculative methods without a separate draft model.

Abstract

from arXiv · show

Large Language Models (LLMs) owe much of their success to next-token prediction (NTP), but their autoregressive (AR) structure requires slow, sequential token generation. To overcome this bottleneck, we introduce diffusion-augmented LLMs, a new class of models that defines an AR model distribution while using diffusion to draw multiple tokens in parallel from that distribution. We decouple the parameters of these models into two sets: AR weights, trained using the standard NTP objective, and lightweight diffusion weights, trained to generate multiple tokens simultaneously. The diffusion weights are learned through a simple Diffusion Distillation phase that adds negligible overhead to existing LLM training pipelines. We also introduce $Ψ$-Spec, a family of samplers that enables lossless acceleration and inference-time scaling at a fixed context length. Unlike speculative decoding, our method requires no separate draft model. Unlike diffusion LLMs (d-LLMs), it accelerates generation without sacrificing the quality of the underlying AR model. The resulting models, called Uno, can be trained from scratch or built by augmenting existing open-weight AR LLMs. Uno achieves higher throughput than leading speculative-decoding methods at every evaluated batch size and delivers up to $3\times$ speedups over the base AR model, including at the largest batch size supported by the device. Notably, our 8B Uno model outperforms the leading open d-LLM, the 26B DiffusionGemma, and the proprietary Mercury 2 across all evaluated benchmarks in agentic tool use, coding, and long-context reasoning. We release code and checkpoints at: https://s-sahoo.github.io/uno/

1 Introduction

Standard next-token prediction gives LLMs strong capabilities but forces sequential decoding, while diffusion-augmented LLMs learn to sample multiple tokens in parallel from the same high-quality AR distribution. Uno separates quality and speed parameters and uses Ψ-Spec samplers for lossless acceleration without a separate draft model.

  • Motivation: NTP-based LLMs generate one token per decoding step, making inference increasingly costly as reasoning traces grow longer.The sequential constraint also slows reinforcement-learning post-training because rollout generation dominates runtime.
  • Motivation: Predictable token sequences could be generated in blocks, amortizing memory transfers and improving accelerator utilization.Sequential decoding is especially inefficient when inference is memory bound at long context lengths.
  • Approach: Diffusion-augmented LLMs decouple AR weights that determine response quality from lightweight diffusion weights optimized for parallel generation.The AR weights are trained with NTP, then frozen while diffusion weights are learned through Diffusion Distillation with negligible training overhead.
  • Approach: Ψ-Spec samples multiple tokens in parallel from the AR distribution, providing lossless acceleration without a separate draft model or lossy AR-to-diffusion conversion.The resulting model is positioned as a drop-in alternative to speculative and self-speculative decoding.
  • Results: Uno improves the speed-quality frontier over speculative-decoding systems and leading open-weight and proprietary d-LLMs across evaluated batch sizes.Its speedup persists at the largest batch size supported by the device, supporting both low-latency generation and high-throughput serving.

2 Background

Autoregressive models factorize sequence probabilities but require sequential token decisions, whereas discrete diffusion corrupts sequences toward a prior and learns to reverse that corruption. Uniform-state diffusion and Ψ-samplers provide the discrete diffusion foundation used for parallel, few-step generation.

  • Autoregressive modeling: An AR language model factorizes the joint sequence distribution with a causal Transformer, but generating length L requires L sequential decoding decisions.Each token can only be generated after its preceding prefix has been produced.
  • Speculative decoding: Speculative decoding uses a smaller draft model to generate candidate blocks that a larger base model verifies in parallel.Its central efficiency opportunity is that sequence likelihood evaluation can be parallel even when base-model sampling is sequential.
  • Discrete diffusion: Discrete diffusion corrupts clean sequences into a simple prior and trains a denoising model to predict clean tokens from noisy states.The forward process mixes each token with the prior using a monotonically decreasing noise schedule α_t.
  • Discrete diffusion: Uniform-state diffusion uses a uniform prior rather than a mask prior and is selected for self-correction, few-step generation, and inference-time scaling.Masked diffusion instead uses the special [MASK] token as its prior.
  • Sampling and distillation: Ψ-samplers generalize posterior sampling with predictor-corrector capabilities, using κ_t to control correction strength.The resulting transition kernel maps a noisier state z_t to a less noisy state z_s.
  • Sampling and distillation: Discrete Consistency Distillation compresses multistep uniform-state diffusion into a few-step generator through deterministic discrete trajectories.Increasing the gap between adjacent distillation states teaches the student to take larger denoising steps.

3 Diffusion-augmented LLMs

Diffusion-augmented LLMs separate response-quality parameters from lightweight parallel-generation parameters, then use Ψ-Spec to draft token blocks and AR verification to preserve the base distribution. The framework trains diffusion weights by distillation and supports lossless acceleration plus inference-time scaling.

  • 3.1 Autoregressive and Diffusion Pathways: Diffusion-augmented LLMs augment each AR layer with separate diffusion weights dedicated to parallel token generation.AR weights determine response quality, while diffusion weights accelerate generation.
  • 3.1 Autoregressive and Diffusion Pathways: The AR weights retain the standard LLM training pipeline, while diffusion weights are trained separately to generate tokens in parallel.The base AR parameters remain frozen during diffusion fine-tuning, and only LoRA diffusion parameters are updated.
  • 3.2 Diffusion Distillation Phase: The training objective combines Discrete Consistency Distillation with a total-variation loss that encourages longer prefixes to pass AR rejection sampling.Minimizing total variation increases the probability of consecutive token acceptance and the expected accepted draft length.
  • 3.2 Diffusion Distillation Phase: Diffusion Distillation trains a one-step, blockwise generator to approximate the AR distribution over multiple tokens.Fully corrupted sequences are denoised directly to clean sequences, with long sequences partitioned into blocks.
  • 4 Ψ-Spec Sampling: Ψ-Spec proposes token blocks in parallel through diffusion and uses AR rejection sampling to accept the longest valid prefix.The sampler is designed to draw multiple tokens while preserving the AR model’s output distribution.
  • 4 Ψ-Spec Sampling: Increasing denoising steps provides an inference-time scaling axis at fixed context length, potentially motivating removal of AR verification when diffusion quality surpasses AR quality.The passage frames this as a central question rather than an established outcome.

5 Experiments

Experiments evaluate Uno across controlled throughput tests, serving regimes, benchmark suites, baselines, and ablations. Uno preserves qualitative AR performance while improving throughput across batch sizes, competing favorably with speculative and diffusion-based methods and accelerating RL training.

  • Throughput Analysis: The 1K/8K throughput test fixes inputs at 1024 random tokens and outputs at 8192 tokens, equalizing effective input and output lengths across methods.The evaluation averages tokens per forward pass before setting decoding steps for diffusion and speculative methods.
  • Uno vs. Base AR: Uno matches the base AR model qualitatively while achieving 1.5× higher throughput at the largest supported batch size and approximately 2.2× at batch size 1.The largest base-AR-supported batch size is 64 in the 1K/8K throughput test.
  • Comparison with Open-Weight Diffusion Language Models: Uno outperforms open-weight DiffusionGemma and Nemotron-Labs-Diffusion on all evaluated tasks while achieving the highest system throughput.DiffusionGemma is faster at batch size 1 but has substantially lower accuracy, whereas Nemotron-Labs-Diffusion trails Uno in both throughput and output quality.
  • Comparison with Proprietary Diffusion Language Models: The 8B Uno outperforms Mercury 2 across agentic tool use, coding, and long-context reasoning, while achieving approximately 4.6× higher maximum system throughput.Uno trails Mercury 2 on one Science and Knowledge benchmark; the throughput comparison uses different hardware and undisclosed Mercury 2 quantization.
  • Faster RL Training: Diffusion weights enable up to a 40% end-to-end RL training speedup, with smaller gains when tool calls dominate runtime.During RL, only base AR weights are updated while diffusion weights remain frozen and accelerate rollouts.
  • Open-Weight Initialization: Uno can augment Qwen3-8B using diffusion weights trained on OpenThoughts while retaining lossless speedups on the evaluated benchmarks.The passage notes that training Qwen3-8B on this dataset degrades quality, while the diffusion weights still enable lossless acceleration.
  • Comparison with Lossless Speculative Decoding Methods: UnoQwen exceeds 5700 tokens per second at the largest supported batch size, outperforming DFlash and EAGLE-3 and reaching 1.6× the base AR throughput.At batch size 1, UnoQwen reaches a 2.5× speedup over the base AR model and is Pareto-dominant across batch sizes.
  • Ablations: Ablations improve TPF through loss weighting, curriculum scheduling, and diffusion-adapter placement, with curriculum increasing TPF from 2.65 to 2.71.Training blocks increasing from 4 to 16 every half epoch outperform a fixed block size of 16; adapter placement across all layers is also more effective.

6 Related Work

Uno is positioned against speculative, diffusion, self-speculative, and multi-token prediction methods as a lossless approach that avoids a separate draft model while addressing speed-quality tradeoffs.

  • Speculative Decoding: Speculative decoding methods provide lossless speedups but require separately trained draft models and extensive architectural design choices.EAGLE-3 and DFlash use smaller AR or diffusion drafters, respectively.
  • Diffusion Models: Large-scale d-LLMs can be faster at small batch sizes but generally trail AR models in quality, with lossy gains diminishing at larger batches.These limitations also complicate post-training because rollout generation slows and RL recipes require modification.
  • Comparative Evaluation: Uno’s evaluation compares its lossless method with lossy diffusion methods using benchmark accuracy and tokens per forward pass.The comparison includes accuracy degradation relative to parent AR models and reports UnoQwen with a specified tree-sampler configuration.
  • Multi-Token Prediction: Uno outperforms speculative decoding methods and consequently multi-token prediction methods while preserving lossless acceleration.MTP methods modify the transformer architecture to predict future tokens.
  • Quadratic Samplers: Quadratic samplers combine drafting and verification in one pass by inserting k^2 masked positions for possible future continuations.This contrasts with Uno’s sampler, which uses separate drafting and verification forward passes.
  • Concurrent Work: Concurrent methods such as OPDLM, FLARE, and I-DLM provide lossy speedups by fine-tuning base AR weights, while an I-DLM variant is evaluated with Uno’s Ψ-Spec sampler for lossless speedups.The authors report that released I-DLM components did not reproduce lossless speedups independently.

7 Conclusion

The paper concludes that diffusion-augmented LLMs unify autoregressive quality and diffusion-based parallel generation in one architecture. Uno supports lossless acceleration and can be trained from scratch or added to existing AR models, while diffusion choices and distillation retain important scope limitations.

  • 7 Conclusion: Diffusion-augmented LLMs unify an autoregressive pathway parameterized by θ_AR and a diffusion pathway parameterized by θ_AR+θ_Δ.The diffusion weights require orders of magnitude fewer training tokens than the AR weights.
  • 7 Conclusion: Uno can be trained from scratch or created by augmenting an existing AR model with lightweight diffusion weights for faster generation.This supports both standalone training and retrofitting of open-weight AR models.
  • 7 Conclusion: 3× speedup over the base AR model is achieved across evaluated batch sizes, with up to 2× retained at the largest supported batch size.The gains accelerate inference and RL post-training, where rollout generation is a major computational component.
  • A.1 Masked Diffusion: Masked diffusion cannot remask unmasked tokens, allowing denoising errors to compound during inference.The limitation arises because the denoising model imperfectly models the clean data.
  • A.2 Uniform-State Diffusion: Uniform-state diffusion permits continuous updates of all token positions, supporting self-correction, few-step generation, and inference-time scaling.Its uniform prior is contrasted with masked diffusion’s fixed tokens after unmasking.
  • A.3 Ψ-Samplers: Ψ-samplers interpolate between reverse denoising and forward noising, allowing earlier decisions to be revised; κ_t=1 recovers ancestral sampling.The practical sampler substitutes the denoiser prediction for the unavailable clean sequence.
  • A.5 Discrete Consistency Distillation: Discrete Consistency Distillation trains few-step generators by matching a student’s prediction on noisier states to a fixed teacher’s clean-token distribution.The teacher is periodically updated from the student while denoising steps progressively increase.
  • A.5 Discrete Consistency Distillation: DCD’s effectiveness is limited because it trains on deterministic probability-flow trajectories that differ from stochastic sampling trajectories.This train-test mismatch is explicitly identified as a limitation.

C.1.1 Benchmarks

The evaluation uses benchmark suites, multiple sampled generations, specialized RL-trained experts, and fixed decoding configurations to assess Uno across its targeted capabilities.

  • C.1.1 Benchmarks: Table 4 lists the benchmarks used in Uno’s evaluation and defines generations as responses sampled per example for average pass@1.Average pass@1 is computed over multiple generations for each benchmark.
  • C.1.1 Benchmarks: UnoQwen uses block sizes B∈{4,8,16}, while tree verification uses B=16, tree size V=60, and candidate top-K=32.The reported configurations also specify temperature and top-p/top-k decoding settings.
  • C.1.1 Benchmarks: The RL setup trains specialized mathematics, code-generation, tool-use, and browse-search experts from a supervised fine-tuning checkpoint.These experts use the DAPO reinforcement learning algorithm.
  • C.1.1 Benchmarks: Table 5 catalogs the open-source model checkpoints used in the evaluations.The checkpoint comparison covers lossless and lossy baseline families described in the evaluation setup.

C.3 Uno Evaluations

Uno’s sampler configurations trade off system and per-user throughput, while its diffusion weights retain speedups after RL post-training and can accelerate generation despite distribution mismatch.

  • The Linear sampler with B = 4 achieves the highest system throughput at batch size 64, while Tree (B,K,V ) = (16,32,32) achieves the highest per-user throughput at batch size 1.
  • SFT diffusion adapters retain their speedup after extensive RL post-training, with only a 6% reduction in TPF.
  • OpenThoughts fine-tuning reduces accuracy by up to 15% points, yet its diffusion weights still accelerate generation despite distribution mismatch.
  • Table 9 compares Qwen3-8B with a variant fine-tuned on OpenThoughts across a subset of benchmarks.

C.5 UnoQwen vs I-DLM (R-ISD) Comparison

The comparison finds that I-DLM’s sampler can degrade the causal model’s accuracy, whereas Ψ-Spec preserves losslessness for I-DLM adapters but delivers lower speedups than Uno.

  • I-DLM’s sampler degrades accuracy on numerous benchmarks relative to the Qwen AR model, while Uno1ep does not.
  • Greedy drafting without an adjusted rejection-sampling procedure causes I-DLM’s sampler to lose its claimed losslessness property.
  • Ψ-Spec remains lossless with both evaluated I-DLM adapter variants, but both provide considerably lower speedups than Uno.
  • I-DLM’s TPF increases significantly from B = 4 to B = 8 and appears to saturate beyond B = 8.

C.6.1 LoRA Rank and Loss Ablations

The ablations show that increasing LoRA rank improves TPF at higher inference cost, while a lightly weighted KL term combined with TV performs best among tested losses.

  • Average TPF increases from 2.39 to 2.47 when LoRA rank rises from 128 to 256, while trainable parameters increase from 349M to 698M.
  • A 0.01×KL plus 1×TV objective slightly outperforms a pure TV loss objective.

C.6.2 Training Block Size Curriculum

Incrementally increasing training block size produces higher TPF than switching early to consistently large blocks.

  • Average TPF drops from 2.71 to 2.65 when training at B = 16 for the remaining two epochs instead of progressively increasing block size.
  • The standard curriculum progressively increases block size after 0.5@2 and 0.5@4, whereas the fixed-B = 16 curriculum remains at B = 16 for two epochs.

C.6.3 LoRA Adapter Ablations

The ablations examine LoRA placement and scaling for Uno, finding that the preferred αLoRA/rLoRA ratio depends on training horizon.

  • LoRA target projections: LoRA is applied to transformer projection matrices, with ablations covering all projections, attention-only projections, Q/K projections, Q-only, and O-only settings.The default applies LoRA to attention Q, K, V, and O matrices plus the MLP gate, up, and down projections.
  • LoRA target projections: The ablation evaluates Qwen LoRA target projections at B=16 and temp = 1 using the linear sampler.Ranks are chosen to approximately maintain LoRA parameter parity.
  • LoRA scaling: The LoRA update adds a low-rank adapter to the base projection as W = Wbase + αLoRA BA.The adapter uses rank-rLoRA factors B and A, with αLoRA controlling their overall contribution.
  • LoRA scaling: αLoRA/rLoRA = 16 is optimal at three epochs, whereas αLoRA/rLoRA = 64 is preferable at one epoch.The selected ratios are 16 for UnoQwen and 64 for Uno.

C.7 DFlash Thinking Mode Ablation

The DFlash ablation compares thinking-mode settings and sampler configurations, showing that disabling thinking increases tokens per step but substantially reduces accuracy.

  • Thinking mode: Disabling DFlash thinking mode increases tokens per step but reduces accuracy from around 76% to roughly 55%.The evaluation uses temp = 1; the authors subsequently enable thinking mode.
  • Thinking mode: The reported quality difference reflects a changed chat template and therefore a different target-model distribution, rather than contradicting losslessness.The lossless speculative decoder preserves the target distribution used in each setting.
  • Sampler settings: DFlash is evaluated with block sizes B ∈{8,16}, while EAGLE-3 uses B = 8 or its standard depth-7 tree with top-k = 10 and V = 60.The EAGLE-3 tree configuration corresponds to block size B = 8 and follows the official settings.
  • Sampler settings: Extended results compare UnoQwen, EAGLE-3, and thinking-enabled DFlash across sampler settings and temperatures 0 and 1.Table 17 reports tokens per step across these configurations without filtering the values.
Loading 2609.04010v1…