Test-Time Scaling for Diffusion Language Models with Frequency-Aware Remasking
Abstract
Diffusion language models (DLMs) generate text through iterative denoising, where decisions about which tokens to commit or remask strongly affect output quality. Existing certainty-aware strategies rely on instantaneous confidence or entropy, but these local signals can be misleading on difficult reasoning problems. We formulate DLM test-time scaling as adaptive allocation of both inference compute and decoding strategies. Our two-stage approach uses a fast base sampler to solve easy questions, then applies frequency-aware remasking to the remaining hard questions. This strategy tracks how often each token's prediction changes across denoising steps and preferentially commits frequently changing tokens, encouraging samples to explore alternative reasoning paths rather than collapse onto the model's preferred trajectory. Experiments with LLaDA-8B-Instruct and Dream-7B-Instruct show consistent improvements over a strong adaptive allocation baseline, with absolute accuracy gains of up to 3.73 percentage points (LLaDA) and 7.33 percentage points (Dream) on MATH-500, and absolute coverage gains of up to 11.39 percentage points (LLaDA) and 10.98 percentage points (Dream) on HumanEval. On questions unsolved by the base certainty-aware strategy after 100 samples, frequency-aware remasking recovers solutions for 3.75% to 18.75% of examples.
Community
Hi everyone. This is a paper regarding DLM test-time scaling with a remasking strategy. The remasking strategy applies a different signal other than model confidence and is used aiming to boost generation diversity.
This is an automated message from the Librarian Bot. I found the following papers similar to this paper.
The following papers were recommended by the Semantic Scholar API
- Fork-dLLM: Avoiding the Flexibility Trap in Diffusion Language Models (2026)
- Pivot-SD: Efficient Self-Distillation for Masked Diffusion Language Models (2026)
- CoDR: Training-Free Confidence-Drift Remasking for Diffusion Language Models (2026)
- Batch-Paced: Fixed-Budget Batch Decoding in Discrete Diffusion Language Models (2026)
- Learning When to Commit: Schedule Distillation for Fast Diffusion Language Models (2026)
- TeCoP: Accelerating Diffusion Large Language Model Decoding via Tentative Context Pipelining (2026)
- TraceQuery: learning what to recompute from full-denoising trajectory in diffusion language models (2026)
Please give a thumbs up to this comment if you found it helpful!
If you want recommendations for any Paper on Hugging Face checkout this Space
You can directly ask Librarian Bot for paper recommendations by tagging it in a comment: @librarian-bot recommend
Committing the tokens whose predictions keep changing is a counterintuitive choice, and it makes sense as a way to stop samples from collapsing onto one path. Saving the extra compute for the hard questions is also practical. How does the gain change as the number of denoising steps grows?
Thank you for raising this question. I actually did not delve deeply into it. So I ran some experiments on the GSM8k hard set (the last table set in the paper). I used a no-block decoding setting with steps = {16, 32, 64, 128} and a generation length of 128. The results are
| Denoising steps | Low-confidence coverage (%) | SAFE coverage (%) | Distinct answers per question (Low-confidence / SAFE) |
|---|---|---|---|
| 16 | 19.4 | 22.6 | 3.3 / 4.4 |
| 32 | 14.5 | 25.8 | 2.4 / 5.4 |
| 64 | 14.5 | 24.2 | 2.0 / 6.9 |
| 128 | 16.1 | 29.0 | 1.7 / 7.1 |
In this setting, the gap widens as the number of steps grows. Low-confidence remasking does not benefit from more steps, and its samples collapse onto fewer distinct answers, while SAFE keeps producing more distinct answers and its coverage trends upward. These are single-seed results on a small hard subset, so I would read this as a trend and not a precise law.
Get this paper in your agent:
hf papers read 2611.10486 Don't have the latest CLI?
curl -LsSf https://hf.co/cli/install.sh | bash Models citing this paper 0
No model linking this paper
Datasets citing this paper 0
No dataset linking this paper
Spaces citing this paper 0
No Space linking this paper
Collections including this paper 0
No Collection including this paper