A Tiny 856K-Parameter Policy Speeds Up the Diffusion LLM LLaDA-8B

Diffusion language models promise something classic autoregressive models cannot do: generate many tokens in a single pass. In practice, everything hinges on one question: which of the predicted tokens can be safely committed right now, and which should stay masked until the next step. A new controlled study broke this question into its components and ended up training a policy of just 856 thousand parameters. On LLaDA-8B-Instruct, it produced the best accuracy–NFE (number of model function evaluations) frontier and the highest throughput on GSM8K.

The policy as decoding distillation

Learned unmasking policies have appeared before, and they do improve the trade-off between speed and quality. But earlier work differed along several axes at once: how the training data is labeled, what inputs the policy sees, and how large it is. As a result, comparing the approaches was nearly impossible, and it was unclear which design choice actually matters.

The authors proposed a unified framework. Here, a supervised learned policy distills a decoding procedure through a "teacher", that is, the labeling. They then studied four things separately:

• what the teacher counts as a token that can be committed;
• on which generation states the policy imitates the teacher;
• what the policy observes as input;
• how much NFE must drop to pay back the policy's own overhead.

Where naive labeling breaks down

The most obvious labeling method works like this: a token is considered ready to commit if it matches what ends up in the final answer. The study showed that at low NFE, this criterion is too lenient. The policy learns to accept tokens that happen to coincide with the final result but actually pull generation away from the reference path. When compute is scarce, the cost of such mistakes grows.

The picture on inputs is also mixed. The model's hidden states alone give a weak signal for the commit decision. As a complement to the prediction's confidence, however, they work: together, the two sources of information give more than either does alone.

The authors also separately計算 the cost of the policy itself. Any add-on to a model adds compute, and the gain from fewer steps can be eaten up by maintaining it. By their estimate, an NFE reduction of roughly 10% is enough to offset the overhead of policies up to 7.5 million parameters across all configurations tested.

Trajectory instead of matching the final answer

The final recipe grew out of these observations. The policy is trained on trajectory-contraction objectives: a commit is accepted only if it stays on the reference decoding trajectory. The policy's inputs are the predictions' confidence and the hidden states.

Key results:

• policy size: 856K parameters, less than 0.02% of the 8 billion parameters of LLaDA-8B-Instruct;
• GSM8K: the best accuracy–NFE frontier and the highest throughput among the compared approaches;
• MATH-500 and HumanEval: competitive results at moderate target speedup levels;
• Dream-7B-Instruct: the policy mostly transfers to another model without a full redesign.

A caveat from the results themselves: on MATH-500 and HumanEval, the advantage is claimed only at moderate targets, which means that with aggressive step reduction, the lead is not guaranteed.

What to take away in practice

For teams running diffusion LLMs in production, the main takeaway is an engineering one. Faster decoding does not necessarily require a heavy add-on: a policy under a million parameters pays for itself if it cuts NFE by at least about ten percent. That gives a handy budgeting guideline: if the gain in steps is smaller, the policy's overhead will consume it entirely.

The second recommendation concerns labeling. If you are distilling a decoding strategy, do not treat "matches the final answer" as a sufficient criterion, especially in the low-compute regime. Anchoring to a reference trajectory provides stricter training signals. And hidden states are worth adding to confidence rather than using as the sole input.

Finally, the transfer to Dream-7B-Instruct suggests that such policies can be made relatively universal within the family of masked diffusion models. For an industry where diffusion language models are competing with autoregressive ones on tasks with strict latency requirements, this is an argument for cheap, portable add-ons instead of retraining the models themselves.