Unlocking Lossless Speedups in LLMs via Discrete Diffusion

Abstract: 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/

• The paper introduces a new diffusion-augmented language model (Uno), with lightweight adapters that enable up to 405 tokens per second throughput during inference.
• The proposed method uses a dual-pathway architecture with AR weights for verification and diffusion adapters for drafting, retaining the AR language model's distribution by using a rejection-sampling procedure called $\Psi$-Spec.
• When the Uno model was tested on a variety of benchmarks, it achieved substantial speedups at both low- and high-concurrency scenarios, while also preserving the high-quality output of the base autoregressive language model.

Problem formulation and central contribution

“Unlocking Lossless Speedups in LLMs via Discrete Diffusion” (2609.04010) addresses the inference cost imposed by autoregressive (AR) decoding. Standard next-token prediction yields a factorization in which each token depends on the previously generated prefix, forcing one sequential decoding decision per token. This is particularly costly for long reasoning traces, agentic trajectories, and reinforcement-learning rollouts, where generation rather than optimization can dominate total runtime. The problem is also hardware-relevant: decoding is frequently memory-bound, so sequentially loading model weights and updating the KV cache underutilizes accelerator compute.

The paper’s central proposal is a diffusion-augmented LLM: a single model contains an unchanged AR pathway responsible for the target distribution and a lightweight diffusion pathway trained to propose multiple tokens in parallel. The resulting model, called Uno, uses the AR weights for verification and low-rank diffusion adapters for drafting. A rejection-sampling procedure, $\Psi$-Spec, then accepts the longest valid prefix of the diffusion proposal and applies the usual residual correction when a proposal is rejected. Since verification is performed against the original AR distribution, the method is intended to preserve the output distribution exactly rather than trading quality for speed.

The distinction from conventional diffusion LLMs is consequential. Existing d-LLMs modify or replace the AR parameterization and therefore generally introduce a quality–speed tradeoff. In contrast, Uno treats diffusion as an auxiliary proposal mechanism. The paper consequently makes a strong claim: parallel generation can be added to an AR LLM without changing its sampling distribution, without a separate draft model, and without requiring the AR model to be retrained for diffusion generation (2609.04010).

Figure 1: Diffusion distillation augments a standard AR training pipeline with lightweight diffusion weights and combines quality evaluation with throughput measurements across serving regimes.

Diffusion-augmented architecture

Each projection matrix in the Transformer receives a LoRA adapter. The base parameters $\theta_{\mathrm{AR}}$ remain responsible for standard causal next-token prediction, while the adapter parameters $\theta_{\mathrm{diff}}$ are activated only for the diffusion proposal pathway. At inference time, the draft distribution is generated using $\theta_{\mathrm{AR}}+\theta_{\mathrm{diff}}$, whereas the verifier uses $\theta_{\mathrm{AR}}$ alone.

This construction provides three systems-level advantages over a separate speculative drafter. First, the diffusion parameters are substantially smaller than a full auxiliary model. In the Qwen3-8B experiment, the adapters contain approximately $0.35$B trainable parameters, compared with $0.40$B for EAGLE-3 and $1.05$B for DFlash. Second, draft and verification use a shared model architecture and KV cache, reducing peak memory relative to methods that maintain independent drafting and target models. Third, the method can augment an existing open-weight AR checkpoint without modifying its base parameters.

The diffusion pathway retains the NTP output convention rather than adopting the usual same-position denoising parameterization of many diffusion LMs. Given a corrupted block, its output at position $\ell$ predicts the subsequent clean token. This alignment allows the diffusion adapter to imitate the joint sequence of AR conditionals while operating on multiple corrupted positions in parallel.

A technically important implementation detail is gated LoRA. During diffusion distillation, the clean prefix positions use the base AR weights, while the corrupted block positions use both AR and diffusion weights. The same forward pass therefore produces teacher logits for the clean sequence and student logits for the noisy sequence. At inference, the first token of a proposed block is sampled from the base AR distribution and is thus always compatible with the verifier; the remaining tokens are proposed by the diffusion pathway.

Diffusion distillation

The diffusion adapters are trained after the AR model has been trained and frozen. The paper adapts Discrete Consistency Distillation (DCD) to the AR-teacher setting, but avoids explicitly constructing the intermediate probability-flow states ordinarily used by DCD. Instead, a fully corrupted block is mapped directly to a clean block, with the frozen AR model defining the teacher distribution.

Because one-step denoising over a complete long sequence is difficult, training proceeds blockwise. A sequence is partitioned into blocks of size $B$. Each noisy block is conditioned on the preceding clean context, and the student is trained to match the AR teacher’s distribution over the same block. The training context is therefore twice the sequence length rather than $\theta_{\mathrm{AR}}$0 times the sequence length, which the authors identify as a major computational advantage over DFlash-style training.

The training objective combines two terms. The first is a distillation divergence between student and teacher token distributions. The second is a total-variation objective that directly encourages draft distributions to overlap with the verifier distribution. This distinction matters because speculative decoding speed depends not merely on marginal predictive quality but on the probability that a consecutive sequence of proposed tokens will be accepted. The paper reports that the total-variation term is more directly useful for this purpose: in the Qwen ablation, TV-only training achieves an average TPF of $\theta_{\mathrm{AR}}$1, compared with $\theta_{\mathrm{AR}}$2 for either KL-only or unweighted KL plus TV. Reducing the KL coefficient to $\theta_{\mathrm{AR}}$3 improves TPF slightly to $\theta_{\mathrm{AR}}$4.

The authors also find that the block-size curriculum is material. Increasing the block size progressively from 2 through 16 yields an average TPF of $\theta_{\mathrm{AR}}$5, compared with $\theta_{\mathrm{AR}}$6 when the model trains on block size 16 for the corresponding period. This suggests that the adapter must first learn local parallel prediction before being optimized for longer blocks. Applying adapters across all attention and MLP projections is also more effective than concentrating an equal parameter budget in only selected projections.

The distillation cost is small relative to AR training. The from-scratch Uno model is trained on approximately $\theta_{\mathrm{AR}}$7T AR tokens, whereas its diffusion adapters are trained on $\theta_{\mathrm{AR}}$8B tokens. The Qwen augmentation experiment trains $\theta_{\mathrm{AR}}$9B adapter parameters for three epochs, or approximately $\theta_{\mathrm{diff}}$0B tokens, without access to Qwen’s original training corpus. Training takes approximately 32 hours on four eight-GPU H200 nodes in that setting and approximately 60 hours on eight such nodes for the from-scratch model.

$\theta_{\mathrm{diff}}$1-Speculative sampling

$\theta_{\mathrm{diff}}$2-Spec combines a parallel diffusion proposal with exact AR verification. To draft a block of size $\theta_{\mathrm{diff}}$3, the sampler appends $\theta_{\mathrm{diff}}$4 prior samples to the current prefix and performs one-step denoising. The first token is sampled using the base AR pathway; the remaining tokens are sampled independently from the diffusion-conditioned marginals. The resulting candidates are then verified by the AR model using standard speculative-decoding rejection sampling.

For one candidate sequence, the verifier accepts the longest prefix whose token-wise acceptance tests succeed. If a rejection occurs, the rejected position is sampled from the renormalized residual distribution. If all proposed tokens are accepted, the verifier samples one additional continuation token. This is the standard mechanism that preserves the target distribution, and the paper’s losslessness claim follows from applying it to the fixed AR verifier rather than from assuming that the diffusion proposal is exact.

The sampler supports two operating modes. The linear sampler generates one candidate and is optimized for high-batch system throughput, where verification compute is scarce. The tree sampler produces multiple candidate prefixes and verifies them concurrently, exploiting spare compute under batch-size-one or low-concurrency serving. The choice of sampler is therefore not merely an accuracy or acceptance-rate decision; it depends on the hardware regime and the distinction between aggregate throughput and per-request latency.

The authors report a tokens-per-forward-pass range bounded by the two required passes: diffusion drafting and AR verification. In the minimal case, the first token and a verifier correction yield two output tokens over two forward passes; when all $\theta_{\mathrm{diff}}$5 draft tokens are accepted, the verifier can produce $\theta_{\mathrm{diff}}$6 tokens. Additional diffusion denoising steps can also be used for inference-time scaling at fixed context length. However, the paper explicitly leaves open whether sufficiently many denoising steps can exceed the quality of the AR model, in which case AR verification would constrain rather than improve the final output quality.

End-to-end Uno results

The from-scratch Uno model is a dense approximately 8B-parameter decoder-only Transformer with a 262,144-token evaluation context and up to 131,072 generated tokens. Its AR pathway is trained on approximately $\theta_{\mathrm{diff}}$7T tokens through staged context extension to 512K tokens. Rank-128 LoRA adapters with LoRA scaling $\theta_{\mathrm{diff}}$8 are applied to every projection matrix.

Against its own base AR model, Uno improves throughput at every evaluated batch size. At batch size 1, per-request throughput increases from 176 to approximately 383–405 tokens per second, depending on the reported configuration. At the maximum batch size supported by the base AR model, 64, system throughput increases from 3,577 to 5,255 tokens per second, a $\theta_{\mathrm{diff}}$9 improvement. The paper also reports approximately $\theta_{\mathrm{AR}}+\theta_{\mathrm{diff}}$0 higher throughput at batch size 1 and up to $\theta_{\mathrm{AR}}+\theta_{\mathrm{diff}}$1 speedups under selected configurations.

The distinction between batch size 1 and maximum feasible batch size is central to the evaluation. Uno’s advantage is not restricted to an underutilized accelerator. Its linear sampler with $\theta_{\mathrm{AR}}+\theta_{\mathrm{diff}}$2 achieves the best system throughput at batch size 64, whereas the tree sampler with $\theta_{\mathrm{AR}}+\theta_{\mathrm{diff}}$3 is optimized for per-request throughput. The mean TPF is approximately $\theta_{\mathrm{AR}}+\theta_{\mathrm{diff}}$4 for the system-throughput configuration and $\theta_{\mathrm{AR}}+\theta_{\mathrm{diff}}$5 for the per-request configuration.

The accuracy results indicate that the AR verifier remains functionally intact. Uno obtains 90.1 on $\theta_{\mathrm{AR}}+\theta_{\mathrm{diff}}$6-Bench Telecom, 67.1 on Retail, 39.6 on Terminal-Bench v2.1, 68.4 on SWE-bench Verified, and 68.0 on AA-LCR. Its mathematical results include 95.4 on GSM8K, 98.9 on MATH500, 93.0 on AIME 2024, 90.7 on AIME 2025, and 86.3 on AIME 2026. The paper reports a single notable exception in its main comparative claims: AA-Omniscience, where Uno scores 14.3 while Mercury 2 scores 20. The implication is that the lossless mechanism protects the base AR distribution, but it cannot compensate for differences in model scale, training data, or task specialization.

Comparison with diffusion and speculative baselines

Uno is compared with open-weight d-LLMs, proprietary diffusion models, and lossless speculative decoders. Against DiffusionGemma-26B-A4B and Nemotron-Labs-Diffusion-14B, the 8B Uno model achieves higher accuracy on every reported benchmark in the main comparison and also the highest system throughput. Its system throughput is 5,255 tokens per second, compared with 1,136 for DiffusionGemma and 2,794 for Nemotron-Labs-Diffusion. DiffusionGemma retains higher batch-size-one throughput, approximately 836 tokens per second versus approximately 405 for Uno, but does so with substantially lower accuracy on most tasks. This result supports the paper’s more specific claim that lossy d-LLM speedups at low concurrency do not necessarily translate into superior serving throughput or quality under realistic batching.

The comparison with Mercury 2 is similarly asymmetric. Uno exceeds Mercury 2 on the reported agentic tool-use, agentic coding, and long-context reasoning tasks, while trailing on AA-Omniscience. The reported maximum system throughput is approximately 5,255 tokens per second for Uno versus 1,197 for Mercury 2, or roughly a $\theta_{\mathrm{AR}}+\theta_{\mathrm{diff}}$7 difference. This comparison requires caution: Mercury 2 uses undisclosed hardware and quantization, and its public throughput figure was obtained under a different proprietary serving stack. The numerical contrast is therefore informative but not a controlled hardware comparison.

The open-weight Qwen3-8B experiment provides the strongest controlled comparison with lossless speculative methods. With a linear $\theta_{\mathrm{AR}}+\theta_{\mathrm{diff}}$8 configuration, Uno reaches 5,733 tokens per second at the largest supported batch size, compared with 5,351 for DFlash, 4,944 for EAGLE-3, and 3,592 for the base AR model. At batch size 1, the best tree configuration reaches 445 tokens per second for Uno, versus 289 for EAGLE-3 and 370 for DFlash. Thus, the paper reports a maximum $\theta_{\mathrm{AR}}+\theta_{\mathrm{diff}}$9 system-throughput improvement over the base AR model and a $\theta_{\mathrm{AR}}$0 per-request improvement.

The advantage is not attributable solely to acceptance length. Uno’s average accepted tokens per speculative step are higher across the evaluated mathematical, coding, scientific, and instruction-following tasks. Under the system-throughput configuration, Uno averages approximately 3.89 tokens per step, compared with 2.08 for EAGLE-3 and 2.07 for DFlash. Under the per-request configuration, the corresponding values are approximately 5.97, 3.48, and 2.74. Moreover, Uno uses 122.2 GiB peak memory in the system-throughput configuration, compared with 130.0 GiB for EAGLE-3 and 130.1 GiB for DFlash.

The paper’s comparison with lossy Qwen-derived diffusion methods shows a different pattern. Uno generally achieves higher TPF while retaining the Qwen3-8B accuracy, whereas several lossy methods suffer substantial accuracy reductions. For example, on AIME 2024, Uno reports 76.7 accuracy with TPF 4.01, while SDAR reports 10.0 with TPF 2.8 and OPDLM reports 14.7 with TPF 1. On HumanEval, Uno reports 94.8 with TPF 3.67; several lossy methods are both slower and markedly less accurate. These comparisons are heterogeneous because they rely on different checkpoints, training procedures, and in some cases results reported by the respective papers rather than a single reproduced evaluation pipeline.

Robustness to RL post-training and data mismatch

A particularly consequential experiment tests whether diffusion adapters trained before RL remain useful after the AR weights have changed. The authors distill the adapters from an SFT checkpoint, freeze them, and then apply DAPO updates only to the AR weights. After training mathematics, code, tool-use, and web-search experts, the adapters retain their acceleration with only an approximately $\theta_{\mathrm{AR}}$1 average TPF reduction, from 2.25 to 2.10.

The end-to-end RL effect is substantial but uneven. Rollout acceleration produces up to a $\theta_{\mathrm{AR}}$2 reduction in total training time for the mathematics and code experts. Gains are smaller for tool use and search because tool calls and external interactions dominate the runtime. The result supports a specific systems implication: diffusion acceleration can reduce rollout cost without requiring modifications to standard AR RL objectives or policy-optimization algorithms, but the benefit is bounded by the fraction of wall-clock time attributable to model decoding.

The Qwen experiment further shows that adapter training need not use the AR model’s original data distribution. The adapters are trained on OpenThoughts, while the frozen AR weights originate from Qwen3-8B. Fine-tuning Qwen itself on OpenThoughts reduces accuracy by as much as 15 percentage points on some benchmarks, but training only the diffusion adapters on that corpus preserves the original AR model’s performance under lossless verification. This demonstrates that the diffusion dataset primarily needs to support proposal alignment, not reproduce the full pretraining distribution. It does not, however, establish that arbitrary out-of-distribution adapter data will maintain high acceptance rates.

Limitations and open questions

The paper’s losslessness guarantee is algorithmic rather than empirical: it follows from exact rejection sampling against the fixed AR distribution. In practice, numerical nondeterminism, implementation errors, sampling-policy mismatches, and incorrect handling of greedy proposals can invalidate the guarantee. The paper identifies such an issue in the released I-DLM sampler: greedy drafting combined with an unadjusted rejection procedure produces accuracy degradation. This observation reinforces that “lossless” depends on the complete proposal and verification implementation, not merely on the use of LoRA or diffusion.

The diffusion training procedure also has a stated train–test mismatch. DCD is trained on deterministic probability-flow trajectories, whereas $\theta_{\mathrm{AR}}$3-Spec samples stochastic denoising trajectories. The adapter therefore learns a surrogate proposal objective rather than the exact stochastic inference process. The empirical TV objective partly addresses this mismatch, but the paper does not provide a general bound connecting its blockwise TV loss to throughput under arbitrary sampling temperatures, contexts, or model updates.

The current sampler requires two forward passes per decoding iteration. Quadratic samplers could combine drafting and verification, but they introduce quadratic numbers of masked positions and require specialized kernels. The reported throughput therefore includes an avoidable systems overhead, and the relative ranking against baselines may change under optimized kernels.

The results also depend on substantial configuration search. Different sampler configurations are selected for system throughput and per-request throughput, and the optimal block size varies with batch size. The paper evaluates fixed 1K-input/8K-output workloads, which improve comparability but do not fully characterize interactive latency, short responses, heterogeneous request lengths, prompt processing, or continuous batching under production traffic.

Finally, the proposed inference-time scaling mechanism remains unresolved. Additional denoising steps may improve diffusion proposal quality at a fixed context length, but the paper does not establish when this additional computation improves end-task quality, increases acceptance, or ceases to be throughput-positive. It also leaves open whether combining Uno with MTP heads, quadratic sampling, or other architectural speculative mechanisms yields additive gains.

Conclusion

The paper presents a coherent separation between AR modeling and parallel proposal generation. Diffusion adapters are trained to approximate blocks of the frozen AR distribution, while $\theta_{\mathrm{AR}}$4-Spec applies exact AR verification to preserve the target distribution. Across from-scratch and Qwen3-8B settings, Uno reports higher throughput than its AR baseline at batch size 1 and at high concurrency, lower additional parameter and memory costs than separate speculative drafters, and substantially stronger quality than the evaluated lossy d-LLMs.

The strongest empirical result is not the isolated maximum speedup but the persistence of acceleration at realistic batch sizes: Uno reaches 5,255 tokens per second against 3,577 for its base AR model and 5,733 against 3,592 for Qwen3-8B. The method also accelerates RL rollouts without retraining the diffusion adapters after AR post-training. Its remaining technical questions concern stochastic distillation, fused draft–verify kernels, distribution drift under larger policy updates, and the quality–compute behavior of multistep diffusion sampling.