Accelerating Diffusion LLMs
Free while signed in. Answers cite the passages they came from.
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.
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.
Get next week’s papers.
The same picks and the same summaries, in your inbox. Free, and 176 issues deep.
Subscribe on Substack