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

Uno uses diffusion-augmented LoRA adapters to enable lossless, parallel token generation in LLMs.

How can we accelerate LLM generation by using discrete diffusion to propose multiple tokens in parallel without sacrificing the accuracy of the original autoregressive model?

Large Language Models (LLMs) are bottlenecked by autoregressive decoding, which forces the model to generate tokens one by one, leaving modern parallel hardware underutilized. Uno addresses this by augmenting each layer of a standard LLM with lightweight diffusion weights, allowing the model to draft blocks of tokens in parallel while keeping the original autoregressive weights frozen for verification. This approach achieves lossless acceleration, matching the base model's output quality while delivering up to 3× higher throughput across all batch sizes.

Paper Primer

The core mechanism is a decoupled architecture: the frozen autoregressive (AR) weights handle response quality and verification, while rank-128 LoRA adapters (the diffusion weights) are trained via Discrete Consistency Distillation to predict token blocks in parallel. During inference, the model drafts a block of tokens using the diffusion pathway and then performs rejection sampling against the AR distribution to ensure the output remains identical to standard sequential decoding.

Uno provides lossless acceleration that persists at high batch sizes, unlike existing diffusion-based LLMs.

In the 1K/8K throughput test, Uno achieves 1.5× speedup at the largest batch size supported by the base AR model and 2.2× at batch size 1.

Uno outperforms leading open-weight and proprietary diffusion-augmented models on agentic and reasoning benchmarks.

The 8B Uno model surpasses the 26B DiffusionGemma and proprietary Mercury 2 on agentic tool use, coding, and long-context reasoning tasks. 4.6× higher maximum system throughput than Mercury 2.

Why use diffusion adapters instead of a separate draft model like standard speculative decoding?

Standard speculative decoding requires maintaining and running a separate, smaller draft model. Uno integrates the drafting capability directly into the base model's architecture via LoRA adapters, eliminating the need for a separate model while maintaining the same lossless verification guarantees.

Does the speedup hold up if the base AR model is further fine-tuned or updated via RL?

Yes. Because the diffusion adapters are trained to match the AR distribution, they retain their effectiveness even when the AR weights are updated during RL post-training, with only a nominal 6% decrease in tokens-per-forward-pass.

Uno demonstrates that parallel generation can be unified into existing LLM architectures without sacrificing quality or requiring separate draft models, effectively turning inference-time acceleration into a standard, drop-in LoRA adaptation.

Introduction and Motivation

We expose why sequential token generation limits LLM speed and motivate parallel diffusion as a solution.

Large language models achieve impressive capabilities by training on next‑token prediction, but the resulting autoregressive (AR) architecture forces generation one token at a time, creating a latency bottleneck.

The AR design requires each decoding step to wait for the previous token, so inference cannot fully exploit parallel hardware.

Language contains predictable collocations and repetitive structures that could be emitted together, yet current LLMs cannot exploit this because they are forced to emit one token per step.

The sequential nature of AR generation is the primary constraint on inference latency.

Foundations of Discrete Diffusion

We introduce discrete diffusion, a parallel token‑block generation mechanism to speed up LLM inference.

We work with one‑hot column vectors $v\in\{0,1\}^K$ whose entries sum to 1, denoted $V$, and with the categorical distribution $\text{Cat}(\cdot;\pi)$ over $K$ classes. The $K$‑th class is reserved for a special [MASK] token, represented by $m\in V$.

Autoregressive (AR) language models factorize $p_\theta(x)$ as $\sum_{\ell=1}^{L}\log p_\theta(x_\ell\mid x_{<\ell})$, which forces generation to proceed token by token.

Speculative decoding speeds up inference by letting a smaller draft model emit a block of candidate tokens, then verifying the block in a single forward pass of the larger base model.

Instead of corrupting a sequence with continuous Gaussian noise, we replace tokens with categorical noise that can be undone in a few discrete steps, enabling parallel block generation.

Compute the mixed probability vector: $\alpha_t x_\ell + (1-\alpha_t)\pi = 0.5\cdot(0,1,0,0) + 0.5\cdot(0.25,0.25,0.25,0.25) = (0.125,0.625,0.125,0.125)$.

Sample $z_t$ from $\text{Cat}$ with these probabilities; suppose we obtain $z_t = (0,1,0,0)$ (the same token) with probability $0.625$.

If we repeat the step with $\alpha_t = 0.2$, the mixed vector becomes $(0.05,0.25,0.05,0.05)$, making the uniform prior dominate.

Discrete diffusion lets the model trade off fidelity (high $\alpha_t$) against randomness (low $\alpha_t$) in a single categorical draw, which can be parallelized across all positions.

How does discrete diffusion differ from the continuous Gaussian diffusion used in image models?

Continuous diffusion adds real‑valued Gaussian noise to each pixel, requiring a denoising network that predicts means and variances. Discrete diffusion replaces each token with a categorical distribution, staying entirely within the vocabulary space and allowing the reverse step to be a simple softmax prediction.

The paper later introduces $\Psi$‑samplers, a family of discrete‑diffusion samplers that augment the basic reverse posterior with predictor‑corrector steps, improving sample quality.

Discrete consistency distillation (DCD) compresses a multi‑step diffusion trajectory into a few deterministic steps by minimizing the KL divergence between teacher and student denoisers, formalized as the LDCD loss.

The Uno Framework

Diffusion-augmented LLMs add parallel diffusion weights to each layer for lossless speedup.

The autoregressive bottleneck makes token‑by‑token generation a latency limiter. By introducing a parallel diffusion pathway we can draft many tokens at once, then let the AR model verify them, preserving quality while cutting wall‑clock time.

Each layer keeps the original AR weights for quality and adds a tiny diffusion adapter that drafts tokens in parallel; the AR weights then accept or reject the drafts, guaranteeing the same output distribution.

Block 1: feed $[a,b,c,d,\langle\text{MASK}\rangle,\langle\text{MASK}\rangle,\langle\text{MASK}\rangle,\langle\text{MASK}\rangle]$ to the model; the diffusion adapters generate draft tokens $[a',b',c',d']$.

Block 2: the same process yields drafts $[e',f',g',h']$ for the second half.

The AR verifier checks each draft token against the AR distribution; suppose the first three drafts in block 1 are accepted, the fourth is rejected.

Accepted prefix after block 1 is $[a',b',c']$; the verifier then resumes normal AR generation for the remaining positions.

Blockwise diffusion lets the model propose a whole chunk in one forward pass, but only the longest prefix that matches the AR distribution is kept, turning a parallel draft into a safe sequential output.

Parallel draft followed by AR verification.

How does this differ from standard speculative decoding that uses a separate draft model?

Standard speculative decoding trains an entirely independent draft model, so its distribution can drift far from the verifier’s, requiring extra acceptance checks. Here the draft pathway is a LoRA‑adapted copy of the same base model, sharing $\theta_{AR}$; the only difference is the low‑rank adapters, which keep the two distributions tightly coupled and enable lossless verification.

**Figure 1.** *(Top)* Training overview for diffusion-augmented LLMs. Gray cells indicate AR-weight training, while the blue cell indicates diffusion-weight training. *(Bottom Left)* System throughput of Uno, the base AR model, and the baselines; see Sec. 5.1 for details. *(Bottom Right)* Performance across agentic and long-context reasoning benchmarks.

The Ψ-Speculative Sampler

The $\Psi$‑Spec sampler drafts token blocks via diffusion and verifies them with AR to speed generation.

Autoregressive (AR) generation proceeds token by token, which limits throughput. The $\Psi$‑Spec sampler tackles this bottleneck by proposing whole token blocks in parallel and then pruning them with the AR verifier.

The sampler first sketches a rough draft of $B$ tokens with diffusion (like a quick sketch), then an AR verifier polishes the draft by keeping the longest prefix that matches the AR distribution (like an editor discarding mismatched lines).

How does $\Psi$‑Spec differ from standard speculative decoding?

Standard speculative decoding drafts tokens one‑by‑one using the same AR model, while $\Psi$‑Spec drafts an entire block with a diffusion model, uses AR only for the first token’s logits, and then verifies the whole block with the frozen AR verifier. This decouples drafting from verification and enables parallel candidate generation.

Start from the current sequence “\<s\>”. Append $B\!-\!1=2$ random tokens sampled from the prior, yielding a noisy block $z_{t}$ of length 3.

Apply the $\Psi$‑Spec transition to obtain a less‑noisy block $z_{s}$ (e.g., “the cat sits”).

Compute AR logits for the first token (“the”) using frozen AR weights; compute diffusion‑augmented logits for the remaining two positions.

Generate $K^{B-1}=2^{2}=4$ candidate sequences, rank them by joint log‑probability, and retain the top $V=2$ prefixes.

Run the AR verifier on each prefix; accept the longest prefix that the verifier deems valid (e.g., “the cat”).

The diffusion draft provides a diverse set of plausible continuations, while the AR verifier guarantees that any accepted prefix respects the original autoregressive distribution.

Append $B\!-\!1$ random tokens to the current context and form the noisy block $z_{t}$.

Denoise $z_{t}$ with the diffusion model to obtain $z_{s}$ via the $\Psi$‑Spec transition.

Compute logits: use AR weights for token 1, and gated LoRA‑augmented diffusion weights for tokens 2…$B$.

Sample a candidate set $C$ from the joint proposal distribution (linear or tree‑based).

Verify each candidate prefix with the frozen AR model, keeping the longest prefix accepted.

Output the accepted prefix; if all $B$ tokens are accepted, generate one additional token from the AR verifier.

Linear vs. Tree sampler implementations (pseudo‑code).

Tokens‑Per‑Forward‑pass (TPF) measures how many output tokens are produced per forward pass. Because the first draft token is always accepted, the bound is $1 \le \text{TPF} \le \frac{B+1}{2}$.

Empirical Results

LLMs are limited by sequential generation; we add a parallel diffusion path to speed generation without losing accuracy.

The central premise is that sequential token generation limits LLM speed, and a parallel diffusion pathway can accelerate generation without sacrificing quality.

We evaluate two experimental settings: (1) training autoregressive ($AR$) weights from scratch on proprietary data while training diffusion weights on the same distribution, and (2) augmenting the open‑weight Qwen3‑8B model with diffusion weights trained on the OpenThoughts dataset.

Small trainable modules are injected into every weight matrix so the diffusion pathway can be learned while keeping the original $AR$ parameters fixed.

Two samplers are used: a Linear sampler that drafts a fixed block of $B$ tokens per forward pass, and a Tree sampler that balances drafting and verification with parameters $(B,K,V)$.

Throughput is measured with the fixed “1K/8K throughput test” ($m\!=\!1024$ input tokens, $n\!=\!8192$ output tokens). For each method we compute the average $TPF$ (tokens per forward pass) across benchmarks and run $\lceil n\rceil$ decoding steps, constraining each step to accept the measured $TPF$ tokens.

Uno achieves up to $2.2\times$ higher per‑request throughput than the base $AR$ model.

Table 7 shows Uno is $2.2\times$ faster at batch size 1 and $1.5\times$ faster at batch size 64.

RL training of four expert models (math, code, tool use, web search) benefits from the diffusion adapters: updating only the $AR$ weights yields up to $40\%$ end‑to‑end training speedup, with a modest $6\%$ drop in $TPF$.

When initializing $AR$ weights from the open‑source Qwen3‑8B checkpoint (UnoQwen), diffusion adapters trained on OpenThoughts still provide lossless speedups, even though the base model’s accuracy degrades (Table 9).

Compared to lossless baselines (EAGLE‑3, DFlash), UnoQwen attains $>5700$ tokens/s system throughput—a $1.6\times$ speedup over the base $AR$ model—and $2.5\times$ higher per‑request throughput, while using a shared KV cache that reduces peak memory.

Against lossy diffusion methods, UnoQwen consistently records higher $TPF$ values, confirming that our approach preserves accuracy while accelerating generation.

Ablation studies (Table 12) reveal that using only the total‑variation loss yields $TPF=2.39$, that increasing block size from $4$ to $16$ improves $TPF$ from $2.65$ to $2.71$, and that a LoRA rank‑to‑$\alpha$ ratio of $64$ offers the best trade‑off between $TPF$ and inference cost.

Related Work

We compare Uno to prior speculative, self‑speculative, and diffusion‑based generation approaches.

Speculative decoding, self‑speculative decoding, and diffusion‑based generation have each been proposed to accelerate large language models, but they differ in how they integrate drafting and verification.

Speculative decoding method that trains a smaller draft model to propose token blocks, then verifies them with the target model.

Speculative decoding approach that also uses a draft model but is engineered for lower memory and latency, yielding faster inference than other draft‑based methods.

Lossy speedup method that fine‑tunes a base autoregressive model with gated LoRA adapters; claims lossless sampling but reproductions show otherwise.

Lossy speedup technique that modifies the base AR model’s weights to enable parallel token proposals.

Lossy speedup method that applies AR‑Trust techniques to accelerate generation.

Lossy diffusion method that trades model fidelity for faster token generation.

Lossy diffusion approach evaluated in Table 3; details are not expanded in the paper.

Method that trains separate math and coding models from Qwen2.5‑Math‑7B‑Instruct and Qwen2.5‑Coder‑7B‑Instruct.

Diffusion‑augmented model built on Qwen2.5‑7B‑Instruct, offering lossy speedups.

Lossy speedup method that operates in “S Mode” to accelerate generation.

Technique that merges drafting and verification into a single forward pass by inserting masked placeholders for future continuations.

Lossless speedup approaches that modify the transformer architecture to predict multiple future tokens.

**Table 1.** Comparison of UnoQwen against various lossy speedup methods across different benchmarks.

Overall, Uno’s single‑architecture design avoids the memory and latency penalties of separate draft models, delivering lossless speedups that scale with batch size.

Comparison with I-DLM

We compare I‑DLM variants with Uno, highlighting speed and accuracy trade‑offs.

This subsection expands the experimental suite introduced earlier, reporting additional benchmarks, sampler settings, and RL‑post‑training results. We first present the headline throughput advantage of Uno before diving into detailed comparisons with I‑DLM variants.

Uno attains up to 4.6× higher system throughput than the AR baseline.

Table 7 shows Uno reaching $5255$ tokens / sec versus $1136$ tokens / sec for AR.

I‑DLM is a lossless diffusion variant that preserves the causal attention pattern while enabling parallel block proposals.

**Table 8.** TPFs for the SFT checkpoint and its RL-post-trained counterpart, evaluated using the diffusion weights trained for the SFT checkpoint. **Even after extensive RL post-training, the SFT diffusion adapters retained their speedup**, with only a 6% reduction in TPFs.

**Table 10.** I-DLM and `Uno_Qwen`^1ep evaluation with the linear sampler at B = 16 and temp = 1. We compare the I-DLM checkpoints (Ours and Official) under $\Psi$-spec samplers with $\pi$ = m; accuracy is reported with 95% confidence intervals across seeds, using two seeds for MMLU-Pro and ten seeds for all other benchmarks.

**Table 11.** TPFs for the official I-DLM (R-ISD) model with our $\Psi$-spec samplers across block sizes $B$.

Overall, Uno outperforms both I‑DLM variants in speed while maintaining comparable accuracy across all benchmarks, confirming the advantage of the proposed parallel diffusion pathway.

Ablation Studies

We quantify how LoRA rank, loss weighting, curriculum, projection choice, and scaling affect generation speed.

We conduct four ablations to isolate the contribution of each design choice in Uno.

Increasing LoRA rank from 128 to 256 raises average TPF from 2.39 to 2.47.

Table 12 shows the TV‑only objective with rank 128 yields 2.39 average TPF, while rank 256 yields 2.47.

A loss weighting of 0.01 × KL + 1 × TV slightly outperforms pure TV loss.

Table 12 reports average TPF 2.51 for KL + TV (0.01) versus 2.39 for TV‑only.

Using a progressive block‑size curriculum (Curriculum A) improves average TPF to 2.71 versus 2.65 for a fixed‑size curriculum (Curriculum B).

Table 13 reports overall average TPF 2.71 for Curriculum A and 2.65 for Curriculum B.

**Table 13.** TPF for two three-epoch TV-only block curricula. The notation $e@B$ denotes $e$ epochs of training at block size $B$. Both models first train on $0.5@2$, $0.5@4$. The standard curriculum then progressively increases the block size, whereas the fixed-$B = 16$ curriculum trains for the remaining two epochs entirely at $B = 16$. Models are evaluated with the linear sampler at $B = 16$ and temp = 1.

Applying LoRA only to attention projections (Q,K,V,O) attains higher average TPF (2.73) than applying LoRA to all projections (2.47).

Table 14 lists 2.73 for Q,K,V,O versus 2.47 for all projections.

**Table 14.** TPF for `Uno_Qwen`^lep LoRA target-projection ablations with `alpha_LoRA`/`r_LoRA` = 2. Ranks are chosen to approximately maintain LoRA parameter parity. All checkpoints are evaluated with the linear sampler at B=16 and temp = 1.

One‑epoch training peaks at $\alpha_{\text{LoRA}}/r = 64$ (average TPF 2.76), whereas three‑epoch training peaks at $\alpha_{\text{LoRA}}/r = 16$ (average TPF 2.97).

Table 15 shows the highest one‑epoch TPF of 2.76 at ratio 64 and the highest three‑epoch TPF of 2.97 at ratio 16.

**Table 15.** TPF for one-epoch and three-epoch, TV-only `Uno_Qwen`^1ep models with rank-128 LoRA adapters applied to all projections, across LoRA scaling values. All checkpoints are evaluated with the linear sampler at B = 16 and temp = 1; averages are unweighted across the 12 benchmarks.

This table compares performance metrics across various benchmarks (Math, Coding, Science and Knowledge, Instruction Following) for two different configurations: "0.5@2, 0.5@4, 0.5@6, 0.5@8, 0.5@12, 0.5@16" and "0.5@2, 0.5@4, 2@16".

Additional Performance Metrics

We analyze DFlash’s thinking mode impact and compare UnoQwen, EAGLE‑3, and DFlash across token‑per‑step metrics.

We first examine DFlash with its “thinking” mode toggled on and off at temperature = 1. Disabling thinking roughly doubles the tokens‑per‑step metric while causing a pronounced drop in accuracy.

**Table.** Performance metrics across different $\alpha_{LoRA}$ and $\alpha_{LoRA}/r$ configurations for various benchmarks categorized by Math, Coding, Science and Knowledge, and Instruction Following.

Because the quality gap does not violate losslessness—thinking merely alters the chat template and thus the target‑model distribution—we adopt the thinking mode for all subsequent experiments.

C.8 details the baseline settings for the extended comparison: DFlash is evaluated with block sizes B ∈ {8, 16}, while EAGLE‑3 uses a linear block size B = 8 and a tree configuration (depth D = 7, B = 8) with top‑k = 10 and verification budget V = 60 tokens.

**Table 16.** DFlash thinking ablation at temp = 1, top-p = 0.95, and top-k = 50. Accuracy in percent. AL denotes average acceptance length, i.e., tokens per verification step.

Across the three samplers, TPS values increase from the smallest block (B = 4) to the largest (B = 16), confirming that larger blocks better exploit parallelism.

Table 18 presents median per‑request throughput (tokens / second / stream) for the 1K‑input/8K‑output Qwen3‑8B workload at temperature = 1, alongside per‑user throughput figures.

Supplementary Details

Additional diffusion background and related samplers for the appendix.

This appendix supplies extra background on diffusion‑based language models and the samplers that build on them. The material is not required for the main narrative but clarifies the mechanisms referenced later.

Masked Diffusion (MDM) uses a one‑hot mask token $m$ as a fixed prior $\pi = m$. During the forward process each token either stays unchanged or becomes $m$, after which it remains masked forever. The learned reverse posterior $p_{\theta}(s|t)$ replaces the unknown clean sequence $x$ with the denoiser $x_{\theta}(z_t,t)$, but once a token is unmasked it cannot be remasked, which can cause compounding errors at inference.

Uniform‑State Diffusion Models (USDMs) adopt a uniform prior $\pi = \tfrac{1}{K}$, allowing every token position to be updated continuously. Their reverse posterior $q_{\text{USDM}}(s|t)$ is a categorical distribution that mixes the current noisy state $z_t$, the denoiser $x_{\theta}(z_t,t)$, and the uniform prior, enabling self‑correction during inference. The training objective is the Negative Evidence Lower Bound (NELBO), which reduces to a weighted denoising loss.

Self‑Speculative Decoding removes the need for a separate draft model by using a single network for both proposal and verification. The model predicts a block of future tokens in parallel, then verifies the longest prefix that matches the autoregressive (AR) model, guaranteeing exact‑match equivalence to greedy decoding. This scheme requires modifying the base AR model’s weights, so it is not lossless.

$\Psi$‑Samplers are predictor‑corrector samplers that interpolate between the reverse posterior $q_{s|t}$ and the forward noising process $q$, controlled by a correction strength $\kappa_t\in[0,1]$. When $\kappa_t=1$ the sampler reduces to ancestral sampling; lower $\kappa_t$ values allow the corrector to revise earlier decisions, e.g., remasking tokens for MDMs or re‑assigning probabilities for USDMs. Consequently, sample quality improves with additional steps because errors can be explicitly corrected.

Discrete Consistency Distillation (DCD) leverages the fact that uniform‑state diffusion can be viewed as the arg‑max projection of an underlying Gaussian diffusion. By constructing deterministic probability‑flow ODE trajectories in continuous space and mapping them back to discrete tokens, a student model $x_{\theta}$ is trained to match a teacher $x_{\theta^0}$ via a KL divergence loss. The process compresses many‑step USDM samplers into few‑step ones, but the deterministic training trajectories differ from the stochastic sampling trajectories, limiting effectiveness.

Sampler Extensions

Details of the $\Psi$‑Spec sampling process and the Uno algorithm.

Section B.1 describes how the $\Psi$‑Spec transition turns a noisy block $z_t$ into a less‑noisy block $z_s$ by sampling each token in parallel from its marginal distribution.

Because the denoiser is parameterized as a next‑token‑prediction (NTP) model, the clean‑token distribution for the first position uses only the base autoregressive parameters $\theta_{AR}$, while the remaining positions also incorporate diffusion adapters $\theta_{\Delta}$.

Section B.2 shows that after removing time conditioning and shifting logits left, the marginal distribution at step 0 separates the first token (pure AR) from the rest (AR + $\Delta$), yielding a closed‑form expression for $\Psi_{0,\ell}$.

Lossless $\Psi$‑Spec decoding for USDMs (Linear Sampler)

Compared with speculative decoding, Uno’s sampler never discards a draft block; it either accepts the whole block or replaces the tail with a single token drawn from the verifier, guaranteeing lossless generation.

Read the original paper

Open the simplified reader on Paperglide

Browse all simplified papers