Random Remasking Scales Masked Diffusion Inference

Zhaoyi Zhou; Frank Wu; Andrea Zanette

cs.LGOct 22, 2026

Random Remasking Scales Masked Diffusion Inference

Authors: Zhaoyi Zhou, Frank Wu, Andrea Zanette

Organizations: Carnegie Mellon University

Abstract

Masked diffusion models (MDMs) generate sequences by progressively unmasking tokens, but errors introduced during decoding are typically irreversible. Recent work has shown that iterative correction can improve sample quality, yet existing theory fails to explain this given a learned, imperfect denoiser model. We re-formulate MDM inference as a Markov chain on the clean sequence space, defined by a remask-denoise kernel that randomly re-introduces masks and then denoises. This helps us analyze the actual denoising trajectory towards the data distribution. The analysis motivates Random Remask-Denoise (R $^2$ D) inference, a purely test-time, finetuning-free, reward-free procedure. To further improve performance, we propose Trajectory Correction Training (TCT), a denoiser-training procedure that exposes the model to its own errors along the generation trajectory. On logical tasks, our approach exhibits a striking near-log-linear test-time scaling, achieving over 97% accuracy on hard Sudoku puzzles and over 99% accuracy on 3-SAT. Our approach also applies to pretrained models on actual language tasks, beating iterative correction baselines and the official Q-mode accuracy of LLaDA2.1-mini on HumanEval+.

Figures & tables

Panel of Figure 1 — Figure 1: Inference scaling on Sudoku-Hard and 3-SAT. $\mathrm{R}^{2}\mathrm{D}$ and ReMDM are run on standard training checkpoints, both using the standard variant; ProSeCo has its own training method. $\mathrm{R}^{2}\mathrm{D}$ scales near-log-linearly on both tasks and stays above ReMDM at every budget, whereas ProSeCo plateaus after a handful of correction steps.

Training method	Inference method	Sudoku-Hard	3-SAT
Standard	Standard	$48.7$	$90.6$
	ReMDM (standard)	$54.7$	$93.5$
	$\mathrm{R}^{2}\mathrm{D}$ (standard, $M=1$ )	$19.7$	$86.0$
	$\mathrm{R}^{2}\mathrm{D}$ (standard)	$93.5$	$98.5$
PUMA	Standard	50.8	79.6
RemeDi	RemeDi	50.2	84.3

Table 1: Accuracy (in

\%

) on Sudoku-Hard and 3-SAT, organized by training method and inference method. ReMDM,

\mathrm{R}^{2}\mathrm{D}

and ProSeCo use the same compute budget (256 steps per token for Sudoku-Hard, 32 steps per token for 3-SAT).

\mathrm{R}^{2}\mathrm{D}

is the primary driver for inference-time scaling, and TCT cooperates well with

\mathrm{R}^{2}\mathrm{D}

to further raise performance.

Figure 2: HumanEval+ accuracy with the off-the-shelf LLaDA2.1-mini (16B). The three colored bars are inference methods we run ourselves, and the two grey bars are the model’s own published Q-mode and S-mode results. $\mathrm{R}^{2}\mathrm{D}$ achieves the top accuracy, surpassing the official Q-mode performance.

Figure 3: Averaged number of function evaluations (NFE) per sample at the budget of Table 1 (lower is better), with accuracy in parentheses (TCT checkpoint for $\mathrm{R}^{2}\mathrm{D}$ and ReMDM). $\mathrm{R}^{2}\mathrm{D}$ exits early, whereas ReMDM and ProSeCo spend their full budget.

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 4 — Figure 4: Inference scaling on Sudoku-Hard (top) and 3-SAT (bottom), under standard training (left) and TCT training (right). The y-axis reports task accuracy and the x-axis the per-token correction budget $M$ on a logarithmic scale; error bars are one standard deviation across three training seeds. The horizontal dashed line marks each panel’s own no-correction baseline. $\mathrm{R}^{2}\mathrm{D}$ scales near-log-linearly in all four panels and stays above ReMDM at every budget. Protocol details are in Appendix E .

Training method	Inference method	Acc.	NFE	s/sample	Acc.	NFE	s/sample
		Sudoku-Hard			3-SAT
Standard	Standard	$48.7{\scriptstyle\pm 1.0}$	81	0.74	$90.6{\scriptstyle\pm 2.3}$	20	0.16
	ReMDM (standard)	$54.7{\scriptstyle\pm 9.7}$	20737	177.2	$93.5{\scriptstyle\pm 0.7}$	641	5.30
	$\mathrm{R}^{2}\mathrm{D}$ (standard, $M=1$ )	$19.7{\scriptstyle\pm 4.7}$	162	1.41	$86.0{\scriptstyle\pm 1.1}$	40	0.42
	$\mathrm{R}^{2}\mathrm{D}$ (standard, $M=M^{\dagger}$ )	$93.5{\scriptstyle\pm 2.8}$	4607	36.9	$98.5{\scriptstyle\pm 0.2}$	336	2.70
ProSeCo	ProSeCo ( $M=M^{\dagger}$ )	55.7	20817	168.5	82.2	660	5.02

Appendix

Table 2: Compute-normalized comparison on Sudoku-Hard and 3-SAT. Each inference method is reported with its accuracy (in

\%

), averaged NFE, and per-sample wall time at batch size

1

on a single NVIDIA GH200 GPU.

\mathrm{R}^{2}\mathrm{D}

is the only method that exits early: on Sudoku-Hard it reaches

\mathbf{97.5}

49.4

s/sample where ReMDM reaches

92.3

187.9

s/sample —

5.2

points better at

3.8\times

less wall time. Measurement protocol in Appendix E .

Method	Accuracy	Avg NFE	Avg wall time
$\mathrm{R}^{2}\mathrm{D}$ (correction)	$\mathbf{85.4}$	757	6.6
ProSeCo	81.7	692	6.0
ReMDM	71.6	1632	14.3

Appendix

Table 3: HumanEval+ accuracy, averaged NFE and averaged wall time (in minute per sample) for the three iterative correction methods of Figure 2 , all evaluated with the off-the-shelf LLaDA2.1-mini.

\mathrm{R}^{2}\mathrm{D}

reaches the top accuracy while spending roughly half of ReMDM’s compute.

Carry-over	Sampling	Sudoku-Hard	3-SAT
Yes	Stochastic	$40.0{\scriptstyle\pm 6.7}$	$9.9{\scriptstyle\pm 3.0}$
No	Stochastic	$46.9{\scriptstyle\pm 12.3}$	$1.4{\scriptstyle\pm 1.8}$
Yes	Argmax	$92.3{\scriptstyle\pm 2.9}$	$\mathbf{96.5}{\scriptstyle\pm 0.8}$
No	Argmax	$\mathbf{92.4}{\scriptstyle\pm 2.6}$	$67.7{\scriptstyle\pm 30.9}$

Appendix

Table 4: ReMDM implementation variants on the TCT checkpoint, comparing carry-over denoising against the sampling operator, at

M=256

on Sudoku-Hard and

M=32

on 3-SAT. Accuracy is in

\%

, mean

\pm

standard deviation over three training seeds

\times

three inference seeds. The sampling operator dominates:

\arg\max

is the operator TCT trains under. Carry-over is close to free on Sudoku-Hard but decisive on 3-SAT. The standard checkpoint is not ablated; there we run the official carry-over

+

stochastic configuration.

Mask ratio	Sudoku-Hard	3-SAT
0.00	1.000	0.951
0.10	1.000	0.988
0.30	1.000	0.997
0.50	1.000	0.998
0.70	1.000	1.000
0.90	1.000	1.000

Appendix

Table 5: Token preservation of the TCT-trained denoiser: The fraction of correct revealed tokens that survive a single forward pass without carry-over, averaged over three training seeds. On Sudoku-Hard the correct grid is an exact fixed point of the denoiser at every masking ratio; On 3-SAT it is not, and preservation is weakest precisely when the sequence is nearly complete, which is why carry-over is essential there and not on Sudoku-Hard.

Run	$M=16$	$M=32$	$M=64$	$M=128$	$M=256$
Default	$65.8{\scriptstyle\pm 3.6}$ (1050)	$\mathbf{75.6}{\scriptstyle\pm 3.1}$ (1477)	$\mathbf{85.4}{\scriptstyle\pm 2.8}$ (2412)	$\mathbf{92.9}{\scriptstyle\pm 1.9}$ (3552)	$\mathbf{97.5}{\scriptstyle\pm 0.8}$ (5742)
Low confidence	$\mathbf{68.7}{\scriptstyle\pm 1.8}$ (751)	$69.0{\scriptstyle\pm 2.4}$ (990)	$69.5{\scriptstyle\pm 2.4}$ (1356)	$69.3{\scriptstyle\pm 2.4}$ (2212)	$69.6{\scriptstyle\pm 2.3}$ (3271)
No early exit	$63.9{\scriptstyle\pm 3.1}$ (1377)	$71.7{\scriptstyle\pm 3.3}$ (2673)	$79.2{\scriptstyle\pm 4.1}$ (5265)	$86.7{\scriptstyle\pm 3.9}$ (10449)	$93.3{\scriptstyle\pm 2.4}$ (20817)
Stochastic	$37.5{\scriptstyle\pm 4.3}$ (1181)	$50.1{\scriptstyle\pm 6.3}$ (2034)	$63.6{\scriptstyle\pm 6.0}$ (3005)	$76.1{\scriptstyle\pm 6.3}$ (4146)	$87.1{\scriptstyle\pm 4.7}$ (6156)

Appendix

Table 6:

\mathrm{R}^{2}\mathrm{D}

(correction) inference design on Sudoku-Hard under TCT training. Each row changes one design element of Algorithm 1 relative to Default . Each cell reports accuracy (in

\%

, mean

\pm

standard deviation over three training seeds) with averaged NFE in parentheses. Low-confidence remasking almost removes the scaling effect; Removing early exit preserves scaling but at markedly higher NFE; Stochastic sampling in place of

\arg\max

costs accuracy at every budget while spending more NFE.

Reveal rule	Sudoku-Hard acc. ( $\%$ )
Top- $k$ confidence $+$ argmax (published)	97.5
Random position $+$ argmax (Line 11 ablation)	85.5
Top- $k$ confidence $+$ stochastic (Line 9 ablation)	96.8

Appendix

Table 7: Reveal rule in TCT training ( Algorithm 2 ) on Sudoku-Hard, replacing the argmax reveal of Line 9 and the top-

k

confidence reveal of Line 11. Position selection is the load-bearing choice, while argmax reveal is a minor lever; Both approximations help rather than hurt.

Hyperparameter	Sudoku-Hard	3-SAT
Vocabulary size	11	42
Hidden size	256	256
FFN intermediate size (SwiGLU)	768	768
Number of layers	8	8
Number of attention heads	8	8
Number of key/value heads	8	8

Appendix

Table 8: Denoiser architecture for Sudoku-Hard and 3-SAT. The two share a Qwen2-style backbone and differ only in vocabulary size and maximum position, which are dictated by the input encoding.

Masked Diffusion Models

Masked diffusion models generate sequences by repeatedly predicting and revising masked tokens, using bidirectional context for parallel sampling.

See all 184 papers

Revise, Don't Freeze: Sampler-Matched Training for Self-Correcting Masked Diffusion Language Models
May 31, 2026Longxuan Yu, Shaorong Zhang, Yu Fu +3Masked Diffusion Models Diffusion Model Fine-Tuning
University of California, Riverside · Microsoft
Learn from Your Mistakes: Self-Correcting Masked Diffusion Models
Feb 12, 2026Yair Schiff, Omer Belhasin, Roy Uziel +6Masked Diffusion Models Diffusion Language Model Decoding
Cornell · NVIDIA
Targeted Remasking: Replacing Token Editing with Token-to-Mask Refinement in Discrete Diffusion Language Models
Apr 6, 2026Lin YaoDiffusion Language Models Diffusion Language Model Decoding
School of Computer Science, Shanghai Jiao Tong University, Shanghai, 200240, China · Zhongguancun Academy, Beijing, 100097, China

May 31, 2026J/K move · Enter open · S save

Revise, Don't Freeze: Sampler-Matched Training for Self-Correcting Masked Diffusion Language Models

Longxuan Yu, Shaorong Zhang, Yu Fu, Hui Liu+2

University of California, Riverside · Microsoft

Masked Diffusion Models Diffusion Model Fine-Tuning Diffusion Language Model Decoding Diffusion Language Models Diffusion Language Model Training

Open paper AskPDF arXiv

Masked diffusion language models (MDLMs) re-predict every position at each denoising step, but standard samplers commit tokens once revealed, leaving this revision capability unused. Existing approaches either add heuristic or learned mechanisms to revise committed tokens, or remask them back to [MASK] before re-predicting; a principled sampler that directly revises visible tokens without auxiliary modules remains underexplored. We introduce D3IM, a parameter-free sampler derived as a corrector-style reverse update that permits direct visible-to-visible revision without additional modules or auxiliary passes. D3IM also reveals a model-side obstacle we term preservation bias: the model tends to reproduce its own wrong committed tokens rather than correct them. We address this with SCOPE (Self-Conditioned On Prediction Errors), a lightweight post-training procedure that simulates D3IM's sampling process. On LLaDA-8B at 64 denoising steps, SCOPE+D3IM improves over the original LLaDA-8B with standard unmasking by +13.0 on GSM8K (68.3%), +4.8 on MATH-500 (23.6%), +15.3 on HumanEval (29.3%), and +10.4 on MBPP (30.8%), with gains that increase as more denoising steps are used on math and HumanEval.

Abstract

Masked Diffusion Models

Explore similar work

Revise, Don't Freeze: Sampler-Matched Training for Self-Correcting Masked Diffusion Language Models

Learn from Your Mistakes: Self-Correcting Masked Diffusion Models

Targeted Remasking: Replacing Token Editing with Token-to-Mask Refinement in Discrete Diffusion Language Models