Table of Contents
Fetching ...

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

Yu-Yang Qian, Junda Su, Lanxiang Hu, Peiyuan Zhang, Zhijie Deng, Peng Zhao, Hao Zhang

TL;DR

This work addresses the accuracy–parallelism trade-off in diffusion LLMs by proposing d3LLM, a framework that combines training-time pseudo-trajectory distillation with inference-time entropy-based multi-block decoding and KV-cache refresh to enable high parallelism without large accuracy loss. It introduces the Accuracy Under Parallelism (AUP) metric to jointly quantify usefulness under increasing parallelism. Empirical results across LLaDA, Dream, and Dream-Coder show substantial throughput gains (up to 5x speedups over autoregressive baselines and 10x over vanilla diffusion) while maintaining competitive accuracy, with d3LLM-LLaDA achieving the strongest AUP across most tasks. The approach demonstrates practical potential for ultra-fast diffusion LLMs and provides a foundation for integrating with stronger foundation models and speculative decoding in future work.

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 performance. 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.

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

TL;DR

This work addresses the accuracy–parallelism trade-off in diffusion LLMs by proposing d3LLM, a framework that combines training-time pseudo-trajectory distillation with inference-time entropy-based multi-block decoding and KV-cache refresh to enable high parallelism without large accuracy loss. It introduces the Accuracy Under Parallelism (AUP) metric to jointly quantify usefulness under increasing parallelism. Empirical results across LLaDA, Dream, and Dream-Coder show substantial throughput gains (up to 5x speedups over autoregressive baselines and 10x over vanilla diffusion) while maintaining competitive accuracy, with d3LLM-LLaDA achieving the strongest AUP across most tasks. The approach demonstrates practical potential for ultra-fast diffusion LLMs and provides a foundation for integrating with stronger foundation models and speculative decoding in future work.

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 performance. 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 speedup over vanilla LLaDA/Dream and 5 speedup over AR models without much accuracy drop. Our code is available at https://github.com/hao-ai-lab/d3LLM.
Paper Structure (20 sections, 2 equations, 10 figures, 9 tables)

This paper contains 20 sections, 2 equations, 10 figures, 9 tables.

Figures (10)

  • Figure 1: Illustration of the AUP metric, where we calculate the weighted area under the accuracy-parallelism curve.
  • Figure 2: Illustration of the trajectory-based distillation recipe in d3LLM, where we combine the pseudo-trajectory of the teacher dLLM model and the ground-truth prompt-response pair to construct a noisy sequence for training the student dLLM.
  • Figure 3: Illustration of the multi-block decoding strategy in d3LLM, where we decode multiple blocks in parallel based on token entropy, and we introduce a KV-cache together with a KV-refresh mechanism to mitigate quality degradation.
  • Figure 4: (a): Accuracy--parallelism curves of LLaDA-based models on MATH dataset. (b) & (c): Radar charts of AUP score.
  • Figure 5: Accuracy--parallelism curves for LLaDA-based models across five benchmark tasks (i.e., GSM8K-CoT, HumanEval, MBPP, MATH, and Long-GSM8K).
  • ...and 5 more figures

Theorems & Definitions (2)

  • Remark 1: Choice of $\alpha$
  • Remark 2: Hardware-Independence