Helping Diffusion Language Models Explore More at Test Time

Diffusion language models do not generate text in a single left-to-right pass. Instead, they repeatedly denoise a sequence, deciding which tokens to keep, which positions to revisit, and when the current prediction is reliable enough to commit. These remasking decisions are central to output quality: an overly aggressive commitment can lock the model into a weak reasoning path, while excessive remasking can waste inference compute.

This paper studies how to scale diffusion language models at test time. Existing certainty-aware methods commonly use instantaneous confidence or entropy to decide which tokens should be committed. Such local signals can be misleading on difficult reasoning tasks. A token may look highly confident at one denoising step even though the broader chain of reasoning is heading in the wrong direction. Once that path is established, repeated sampling may simply reproduce the model’s preferred trajectory instead of exploring alternatives.

Core approach

• A two-stage compute policy: The system first runs a fast base sampler on all questions. Extra computation is then directed toward the cases that the initial strategy does not solve.
• Tracking prediction changes: Rather than examining only the current confidence, the method counts how often each token’s predicted value changes across denoising steps. This provides a trajectory-level signal about where the model is still exploring.
• Frequency-aware remasking: Tokens with more frequent prediction changes receive a different remasking treatment. The goal is to preserve opportunities for alternative reasoning paths instead of allowing the sample to collapse too early onto one favored solution.
• Adaptive use of expensive decoding: The more involved strategy is reserved for hard examples, rather than imposing the same additional cost on every input.

The experiments use LLaDA-8B-Instruct and Dream-7B-Instruct. On MATH-500, the method improves absolute accuracy over a strong adaptive allocation baseline by up to 3.73 percentage points for LLaDA and 7.33 points for Dream. On HumanEval, the reported absolute coverage gains reach 11.39 and 10.98 percentage points, respectively. The authors also report that, among examples still unsolved by the base certainty-aware strategy after 100 samples, frequency-aware remasking recovers solutions for 3.75% to 18.75% of cases.

Why it matters

The contribution is not simply to sample more often. It changes how additional test-time compute is spent. The history of prediction changes becomes a complementary signal to confidence, helping the decoder identify positions where the model may need more exploration. This is particularly relevant to diffusion language models, whose iterative generation process exposes a natural opportunity to revise intermediate decisions without retraining the model.

The reported evidence is limited to the two evaluated models and the listed mathematics and coding benchmarks. The stability of the frequency signal across other tasks, as well as the practical latency and compute cost of the extra remasking, will require further study. Still, the work suggests a useful division of labor for DLM inference: commit quickly when the problem is easy, but keep the reasoning trajectory flexible when the first solution fails.

Source: Hugging Face Daily Papers