New Breakthrough in Diffusion Language Models: Frequency-Aware Remasking Boosts Reasoning

Diffusion Language Models (DLMs) generate text through iterative denoising, where decisions about which tokens to commit or remask directly affect output quality. Existing deterministic perception strategies rely on instantaneous confidence or entropy, but these local signals can be misleading on difficult reasoning problems. A recent paper proposes formalizing test-time scaling of DLMs as adaptive allocation of inference compute and decoding strategy, and introduces a frequency-aware remasking strategy that achieves significant improvements across multiple benchmarks.

Two-Stage Adaptive Allocation Framework

The study designs a two-stage approach: first, a fast base sampler solves easy problems, then the remaining hard problems are handed to the frequency-aware remasking strategy. The core insight behind this allocation is that different problems require vastly different amounts of compute — easy problems need not consume excessive sampling budget, while hard problems require more refined decoding control.

Compared to existing methods, this framework does not treat all problems uniformly but dynamically adjusts strategy based on problem difficulty. The base sampler quickly filters out easy cases, allowing frequency-aware remasking to concentrate resources on truly difficult reasoning tasks.

Test-time scaling: Refers to improving model performance at inference time by increasing compute (e.g., more sampling iterations, longer denoising steps, or more complex decoding strategies) without changing model parameters. In autoregressive language models, common forms include chain-of-thought prompting, majority voting, and tree search. For diffusion language models, test-time scaling also involves dynamic adjustment of decoding strategies, such as adaptively allocating denoising steps or remasking frequency. The two-stage framework in this paper is a typical practice of routing test-time compute based on problem difficulty, belonging to adaptive allocation of inference compute.

Principles of Frequency-Aware Remasking

Traditional deterministic perception strategies decide whether to commit at each denoising step based on the current token's confidence or entropy. However, such local signals are unreliable in complex reasoning — the model may exhibit high confidence at a certain step, yet the final answer is still wrong.

Frequency-aware remasking instead tracks the frequency at which each token's prediction changes across denoising steps. Specifically, it prioritizes committing tokens that change frequently, thereby encouraging samples to explore alternative reasoning paths rather than collapsing onto the model's preferred trajectory. The intuition behind this mechanism is: if a token's prediction repeatedly wavers, it indicates model uncertainty at that position, and forcing a commit may lock in an error; conversely, allowing these tokens to continue changing gives the model more opportunities to find the correct path.

Remasking and deterministic perception: In the iterative denoising process of diffusion language models, each step generates prediction distributions for all uncommitted tokens. Remasking refers to re-marking certain committed tokens as undetermined (i.e., re-masking) before the next denoising step, giving them a chance to be re-predicted. Deterministic perception strategies typically decide which tokens should be committed (fixed) or remasked based on the current step's confidence or entropy. Although local signals are computationally efficient, they can easily lock in errors due to model overconfidence in multi-step reasoning.

Frequency-aware signals and the exploration-exploitation trade-off: Frequency-aware remasking does not rely on single-step confidence but instead counts the number of times (or frequency) each token's predicted class changes across multiple denoising steps. High-frequency changes mean the model remains uncertain at that position, possibly corresponding to multiple candidate reasoning paths. Prioritizing high-frequency tokens is equivalent to forcing the model to "choose" a path early in decoding, thereby avoiding late-stage collapse onto a single trajectory; while retaining low-frequency tokens allows stable information to gradually solidify. This mechanism essentially dynamically balances exploration (retaining uncertain tokens to try alternative paths) and exploitation (committing stable tokens to advance generation), sharing similar intuition with exploration strategies in reinforcement learning.

Experimental Results: Dual Improvement in Accuracy and Coverage

The researchers conducted experiments on two models: LLaDA-8B-Instruct and Dream-7B-Instruct. Compared to strong adaptive allocation baselines, this method achieved absolute accuracy improvements of up to 3.73 percentage points (LLaDA) and 7.33 percentage points (Dream) on MATH-500, and absolute coverage improvements of up to 11.39 percentage points (LLaDA) and 10.98 percentage points (Dream) on HumanEval.

More interestingly, for problems that the base deterministic perception strategy still failed to solve after 100 samples, frequency-aware remasking was able to recover 3.75% to 18.75% of these cases. This indicates that the method not only improves average performance but can also tackle hard problems that baseline methods completely fail on.

Implications for DLM Reasoning Optimization

The significance of this work lies in revealing the importance of test-time compute allocation. As diffusion language models scale up, how to efficiently utilize inference compute becomes a key issue. Frequency-aware remasking provides a lightweight decoding strategy adjustment that can improve reasoning capabilities without retraining the model.

Furthermore, this method is complementary to existing deterministic perception strategies and can be used in combination. The design of the two-stage framework also reminds us that in practical deployment, dynamically routing different decoding strategies based on problem difficulty may be more effective than a single strategy. In the future, this frequency-aware idea may be extended to other generative models or more complex reasoning tasks.

One-line Summary

A two-stage framework combined with frequency-aware remasking dynamically allocates inference compute to improve diffusion language model performance.

Paragraph Summary

This paper proposes formalizing test-time scaling of diffusion language models as adaptive allocation of inference compute and decoding strategy, and designs a two-stage framework: first use a fast base sampler to solve easy problems, then hand hard problems to a frequency-aware remasking strategy. Frequency-aware remasking tracks the frequency of token prediction changes across denoising steps, prioritizing committing frequently changing tokens to encourage exploration of alternative reasoning paths and avoid misleading local confidence signals. On LLaDA-8B-Instruct and Dream-7B-Instruct, the method improves accuracy on MATH-500 by up to 3.73 and 7.33 percentage points, and coverage on HumanEval by up to 11.39 and 10.98 percentage points, while recovering 3.75% to 18.75% of cases that the baseline still failed after 100 samples.