Test-Time Scaling for Diffusion Language Models with Frequency-Aware Remasking

Test-Time Scaling for Diffusion Language Models with Frequency-Aware Remasking. By Bowen Zuo and three others, from the University of California, Riverside and Indiana University. Published on the seventh of November.

Diffusion language models produce text by iterative denoising. At each step the model decides which tokens to commit and which to remask, and those decisions strongly shape the quality of what comes out. Current certainty-aware strategies make the call from instantaneous confidence or entropy. On difficult reasoning problems, such local signals can mislead.

The authors cast test-time scaling for these models as adaptive allocation of two things at once: how much inference compute a question gets, and which decoding strategy it gets. Easy questions and hard questions need not be treated the same way.

In practice the approach has two stages. A fast base sampler handles the easy questions first. The questions it leaves behind, the hard ones, then receive frequency-aware remasking. This strategy keeps count of how often each token's prediction changes across the denoising steps, and it preferentially commits the tokens that change frequently. Committing the unsettled tokens, rather than waiting for them to settle, pushes samples to explore alternative reasoning paths instead of collapsing onto the trajectory the model already prefers.

Tests use LLaDA-8B-Instruct and Dream-7B-Instruct, and the comparison is a strong adaptive allocation baseline. Improvements are consistent. On MATH-500, absolute accuracy rises by up to three point seven three percentage points for LLaDA and up to seven point three three for Dream.

On HumanEval the measure is coverage. Absolute gains reach up to eleven point three nine percentage points for LLaDA and up to ten point nine eight for Dream. One last test isolates the hardest cases: questions the base certainty-aware strategy still had not solved after one hundred samples. On those, frequency-aware remasking recovers solutions for three point seven five percent to eighteen point seven five percent of examples.

Confidence tells a sampler what the model believes now. Frequency tells it how much that belief has been moving. By deliberately locking in the restless tokens on hard questions, the method trades the model's favourite path for a wider search, and the recovered unsolved questions show that wider search reaching answers the confident route never found.