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

TL;DR

d3LLM introduces a novel training and inference approach for diffusion-based large language models, balancing accuracy and parallelism to enable faster decoding without significant performance loss.

Contribution

The paper proposes pseudo-trajectory distillation and entropy-based multi-block decoding for diffusion LLMs, achieving high parallelism and accuracy balance.

Findings

Up to 10x speedup over vanilla LLaDA/Dream.

5x faster than autoregressive models with minimal accuracy loss.

Introduces AUP metric for joint accuracy and parallelism evaluation.

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.