Accelerating Diffusion LLMs
A lightweight, learned policy speeds up diffusion-based LLM decoding by deciding which tokens are already “final” and when to stop generation. The authors train a tiny MLP filter on token confidence signals and add an End-of-Text Prediction that halts decoding as soon as [EoT] is reliably produced. On LLaDA-8B-Instruct, this reaches large throughput gains with minimal or no accuracy loss.
Ask this paper
Problem and insight: Semi-autoregressive diffusion LLMs parallelize token updates, but static heuristics keep remasking already-correct tokens. The paper defines an oracle strategy, Extremely Greedy Parallel, that unmasks tokens immediately upon correct prediction and shows big headroom for speedup.
Method: Learn2PD filter: Train a 2-layer MLP filter fθ on token confidence patterns to predict “finalize or remask” per position. Only the filter is trained with BCE loss; the dLLM stays frozen. Inference applies a threshold τ to the filter’s logits to commit tokens.
Stop early with EoTP: End-of-Text Prediction halts once [EoT] is decoded, avoiding long tails filled with [EoT]. Appendix B notes about 89.59% of extra compute at length 1024 comes from post-EoT padding.
Results: On GSM8K, MATH, HumanEval, and MBPP, Learn2PD alone yields 3–12× speedup depending on length; Learn2PD+EoTP reaches 22.58× at length 1024 on GSM8K with accuracy preserved or slightly improved. Combining with KV cache further boosts throughput to 57.51× with small accuracy tradeoffs. Longer sequences benefit more; Table 4 shows acceleration grows from 3.36× at length 128 to 22.58× at 1024.
Engineering notes: The filter is tiny and quick to train: for block size 32 it has ~2k parameters, trained in minutes on a single T4 after a short data collection pass. Overhead at inference is negligible relative to gains. Method is orthogonal to KV caching and slotting into existing dLLM decoders is straightforward.