d3LLM: Ultra-Fast Diffusion LLM using Pseudo-Trajectory Distillation

AI summary

Overview

Research area: Machine learning — efficient inference for diffusion large language models (dLLMs), covering distillation-based training, parallel decoding strategies, and evaluation metrics for the accuracy–efficiency trade-off.

Technical level: Advanced. The paper combines trajectory-level knowledge distillation, entropy-thresholded block decoding, KV-cache management, and a custom weighted-area evaluation metric; it assumes familiarity with autoregressive versus diffusion decoding, masking objectives, and throughput metrics such as tokens per forward (TPF) and tokens per second (TPS).

Scope: The paper proposes d3LLM, a training-plus-inference recipe that pushes the accuracy–parallelism frontier of open-source diffusion LLMs (LLaDA, Dream, Dream-Coder) and introduces AUP, a metric that jointly scores accuracy and parallelism, validated across 10 benchmark task settings on H100 and A100 GPUs.

What This Paper Is About

Diffusion LLMs can decode tokens in parallel and in random order, but open-source versions such as LLaDA and Dream achieve only around 20 tokens per second and often trail similarly sized autoregressive (AR) models in accuracy. Existing acceleration methods sit on a fixed accuracy–parallelism curve: some push parallelism and lose accuracy (D2F), others preserve accuracy but limit parallelism (Fast-dLLM-v2). The paper's goal is to move that frontier outward — teaching a dLLM which tokens can be safely decoded early during training, and decoding across multiple blocks in parallel at inference — while introducing a metric (AUP) that captures both sides of the trade-off instead of reporting only one.

Key Contributions

1. Pseudo-trajectory distillation. A training recipe that uses the teacher dLLM's own decoding order (its "pseudo-trajectory") as intermediate supervision, combined with a noisy sequence constructed from the ground-truth prompt–response pair, so the student learns which tokens can be unmasked confidently at early steps. It is paired with a curriculum schedule over mask ratio (0.0 → 0.8) and decoding window size (16 → 32).

2. Entropy-based multi-block decoding with KV-cache refresh. At inference, the method decodes the current and future blocks simultaneously, prioritizing low-entropy (high-confidence) tokens, with blocks transitioning through five states (Inactive, Activated, Fully-Activated, Completed but Stabilizing, Completed) and periodic KV-cache refresh plus early stopping on the EOS token.

3. The AUP metric (Accuracy Under Parallelism). A weighted area under the accuracy–parallelism curve, where the weight penalizes accuracy drops relative to the best accuracy achieved on that task, with a minimum accuracy threshold of y_min = y_1 − 5 and default penalty factor α = 3.

4. Empirical validation across three base models. d3LLM-LLaDA, d3LLM-Dream, and d3LLM-Coder are derived from LLaDA, Dream, and Dream-Coder using the same recipe; the paper reports the highest AUP on 9 out of 10 benchmark task settings.

Main Findings

• Highest AUP in nearly all settings: d3LLM achieves the highest AUP score on 9 out of 10 tasks. d3LLM-LLaDA wins all five of its benchmarks; d3LLM-Dream wins 4 of 5, with Fast-dLLM-v2 taking the MATH task for Dream-based models (AUP 126.7, accuracy 48.7%) — which the authors attribute to Fast-dLLM-v2 being fine-tuned from Qwen-2.5-7B with an additional 1B tokens from the LLaMA–Nemotron post-training dataset, versus d3LLM-Dream's 65M additional tokens distilled from vanilla Dream.

• Large AUP margins over prior dLLM acceleration methods: On GSM8K-CoT (0-shot), d3LLM-LLaDA scores 637.7 AUP versus 358.1 for dParallel-LLaDA; on MATH (4-shot), 107.6 versus 64.5 for dParallel-LLaDA; on HumanEval (0-shot), 96.6 versus 83.7; on Long-GSM8K (5-shot), 441.1 versus 309.1.

• Throughput gains on real hardware: d3LLM-LLaDA reaches 288.9 TPS on H100 (5.0× over Qwen-2.5-7B-it) and 183.3 TPS on A100 (3.6×), a 10.3× speedup over vanilla LLaDA (27.9 TPS on H100). d3LLM-Dream reaches 235.3 TPS on H100 (4.1×) and 128.2 TPS on A100 (2.5×), an 8.5× improvement over vanilla Dream (27.6 TPS on H100).

• Accuracy is largely preserved but not always: d3LLM-LLaDA on GSM8K-CoT scores 73.1% accuracy at 9.11 TPF versus LLaDA's 72.6% at 1.00 TPF. On Long-GSM8K, however, d3LLM-LLaDA reaches 74.2% accuracy versus LLaDA's 78.6%, i.e., a measurable accuracy drop even as the AUP score rises to 441.1.

• Component-level gains from the training recipe: Pseudo-trajectory distillation improves TPF by 18% compared to random masking; the curriculum noise schedule adds approximately 12% TPF over a fixed mask ratio; the curriculum window size adds roughly 8% TPF over a constant window. In the ablation table, adding pseudo-trajectory distillation moves TPF from 6.41 to 7.55 (a 17.8% improvement) at similar accuracy (72.2% to 72.1%).

• Component-level gains from the inference strategy: Multi-block decoding increases TPF by 30%; the KV-cache refresh mechanism improves TPS by approximately 35% in long-context scenarios; early stopping on EOS yields a 5% TPF improvement.

• Curriculum hyperparameters matter: A fixed noise level (t = 0.5) gives 7.49 TPF and 72.8% accuracy, while curriculum 0.0 → 0.8 gives 9.11 TPF and 73.1% accuracy — a 21.6% TPF improvement and a 22.2% AUP improvement. A fixed window size (k = 32) gives 8.22 TPF at 69.8% accuracy, while curriculum 16 → 32 gives 9.11 TPF at 73.1% accuracy, a 19.0% AUP improvement; starting from 0 → 32 yields lower accuracy and AUP.

• AUP rankings are robust to the penalty function: Across nine configurations spanning exponential, power, and linear weighting with varying α (α = 1, 3, 5 for exponential and power; α = 1, 5, 10 for linear), the ranking of LLaDA-based methods is unchanged, with d3LLM-LLaDA ranked #1 in every case.

• Coding scenario: The authors state they are the first to develop an efficient dLLM-coder comparable to AR coders, with 8× speedup over vanilla Dream-Coder; the supporting results are deferred to Appendix A and are not included in the provided content.

• Closed- versus open-source context: The introduction notes that closed-source diffusion models such as Mercury, Gemini Diffusion, and Seed Diffusion sometimes exceed 1000 tokens per second, while open-source LLaDA and Dream achieve only around 20 tokens per second.

Methodology in Plain English

The authors start from the observation that diffusion LLMs are not a single point but a curve: you can always decode more tokens per forward pass if you accept lower accuracy. Their fix has two halves.

On the training side, standard dLLM training masks tokens at random, which gives the model no signal about which tokens are "easy" to commit to early. The authors instead run the teacher dLLM and record the order in which it unmasks tokens, discarding the actual content of its answer. That ordering is combined with the correct ground-truth answer to build a noisy training sequence, so the student is trained to predict the masked ground-truth tokens in an order resembling the teacher's. To keep training stable, they ramp up difficulty: the mask ratio climbs from 0.0 to 0.8 and the decoding window grows from 16 to 32 tokens.

On the inference side, rather than finishing one block before starting the next, the model works on several blocks at once and only commits the tokens it is most confident about (lowest entropy). A block becomes "Activated" when its predecessor is 10% done (conservative, entropy-thresholded decoding), becomes "Fully-Activated" when the predecessor hits 95% (at least one token decoded per forward pass regardless of threshold), and passes through a brief "Completed but Stabilizing" phase where the KV cache is refreshed with full forward passes before the block is finalized. Decoding also stops as soon as an EOS token appears.

To measure all this, they define AUP: plot accuracy against parallelism (TPF), then take a weighted area under that curve, where low-accuracy regions are exponentially down-weighted. A model that gains speed without losing accuracy is rewarded heavily; one that buys speed with accuracy collapse barely moves the score. The paper notes AUP cannot be a training objective because TPF is a runtime statistic and the curve is built from discrete threshold choices, so it is used purely for evaluation.

Why This Matters

Impact on research: The paper reframes dLLM evaluation. Single-number reporting (TPS alone or accuracy alone) hides the trade-off that defines diffusion LLMs, and AUP gives the community a way to compare methods that sits at different operating points. The pseudo-trajectory distillation idea — using a teacher's decoding order rather than its output content as supervision — is a reusable technique for any masked or diffusion-style generation model.

Real-world applications:

• Latency-sensitive chat and assistants: up to 5.0× speedup over Qwen-2.5-7B-it on H100 for GSM8K-CoT makes interactive drafting of reasoning steps more viable on fixed hardware.
• Code generation and completion: the d3LLM-Coder variant targets coding, which the paper reports at 8× speedup over vanilla Dream-Coder, relevant for IDE autocomplete and code assistants.
• Long-context reasoning: Long-GSM8K (prompt length ≈ 1000) is where the KV-cache refresh mechanism is specifically credited with roughly 35% TPS gain, relevant for document-level or retrieval-heavy workloads.
• Throughput-constrained serving: at batch size 1 on a single GPU, the reported 288.9 TPS (H100) and 183.3 TPS (A100) figures speak to per-request serving rather than datacenter-scale batching.

Industry relevance: The gap the paper highlights — closed-source diffusion models sometimes exceeding 1000 tokens per second while open-source LLaDA and Dream sit near 20 tokens per second — is exactly the gap that matters for deploying open-weight diffusion LLMs in products. A recipe that works on top of three existing open base models (LLaDA, Dream, Dream-Coder) without changing the architecture lowers the barrier to adoption.

Future Directions

• Closing the remaining accuracy gap. d3LLM-LLaDA drops to 74.2% on Long-GSM8K versus 78.6% for vanilla LLaDA, and Fast-dLLM-v2 still beats d3LLM-Dream on MATH. Reducing these residual regressions — especially in long-context and mathematical settings — is the natural next target.

• Data-efficiency or data-scale trade-offs. The authors attribute Fast-dLLM-v2's MATH win to 1B additional fine-tuning tokens versus d3LLM-Dream's 65M. Whether d3LLM's recipe scales with more distillation data, or whether its data advantage is structural, is left open.

• Making parallelism-aware objectives trainable. The paper explicitly states AUP cannot serve as a training objective because TPF is a runtime statistic and the metric is non-differentiable over discrete confidence thresholds. Finding a differentiable surrogate would let training optimize the quantity the paper argues is the right thing to measure.

• Generalization beyond the three base models and the evaluated tasks. The content covers LLaDA, Dream, and Dream-Coder on reasoning, math, and code benchmarks; the paper's own concluding sentence mentions extending beyond what is truncated, and the d3LLM-Coder details are deferred to Appendix A, so the breadth of the coding claims is not fully verifiable from the provided content. Extending to multimodal and further long-context settings is not evaluated here.

Target Audience

Researchers and engineers working on efficient LLM inference, particularly those interested in diffusion and masked-generation language models, parallel decoding, and knowledge distillation. It is also relevant to practitioners choosing between AR and diffusion backbones for latency-constrained serving, and to evaluation researchers interested in metrics that couple quality and efficiency rather than reporting them separately. Readers need a working understanding of transformer decoding, masking objectives, and throughput metrics; the paper is not an introductory treatment.

Authors’ abstract

Diffusion large language models (dLLMs) offer capabilities beyond those of autoregressive (AR) LLMs, such as parallel decoding and random-order generation. However, realizing these benefits in practice is non-trivial, as dLLMs inherently face an accuracy-parallelism trade-off. Despite increasing interest, existing methods typically focus on only one-side of the coin, targeting either efficiency or accuracy. To address this limitation, we propose d3LLM (Pseudo-Distilled Diffusion Large Language Model), striking a balance between accuracy and parallelism: (i) during training, we introduce pseudo-trajectory distillation to teach the model which tokens can be decoded confidently at early steps, thereby improving parallelism; (ii) during inference, we employ entropy-based multi-block decoding with a KV-cache refresh mechanism to achieve high parallelism while maintaining accuracy. To better evaluate dLLMs, we also introduce AUP (Accuracy Under Parallelism), a new metric that jointly measures accuracy and parallelism. Experiments demonstrate that our d3LLM achieves up to 10$\times$ speedup over vanilla LLaDA/Dream, and 5$\times$ speedup over AR models without much accuracy drop. Our code is available at https://github.com/hao-ai-lab/d3LLM.