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 183 papers

May 31, 2026

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

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.

Longxuan Yu, Shaorong Zhang, Yu Fu +3

University of California, Riverside · Microsoft

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

Ask PDF

Feb 12, 2026

Learn from Your Mistakes: Self-Correcting Masked Diffusion Models

Masked diffusion models (MDMs) have emerged as a promising alternative to autoregressive models, enabling parallel token generation while achieving competitive performance. Despite these advantages, MDMs face a fundamental limitation: once tokens are unmasked, they remain fixed, leading to error accumulation and ultimately degrading sample quality. We address this by proposing a framework that trains a model to perform both unmasking and correction. By reusing outputs from the MDM denoising network as inputs for corrector training, we train a model to recover from potential mistakes. During generation we apply additional corrective refinement steps between unmasking ones in order to change decoded tokens and improve outputs. We name our training and sampling method Progressive Self-Correction (ProSeCo) for its unique ability to iteratively refine an entire sequence, including already generated tokens. We conduct extensive experimental validation across multiple conditional and unconditional tasks, demonstrating that ProSeCo yields better quality-efficiency trade-offs (up to ~4x faster sampling while still outperforming the best accuracy baseline method) and enables further gains through inference-time compute scaling (up to ~1.2x improvement on benchmarks when allowing for slower, high-NFE decoding).

Yair Schiff, Omer Belhasin, Roy Uziel +6

Cornell · NVIDIA

Masked Diffusion Models Diffusion Language Model Decoding Diffusion Model Inference Acceleration

Ask PDF

Apr 6, 2026

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

Discrete masked diffusion language models such as LLaDA generate text through iterative denoising, where mask tokens are progressively replaced with predicted tokens. LLaDA2.1 introduced a Token-to-Token (T2T) editing mechanism that accelerates generation by directly replacing committed tokens suspected of being incorrect. However, we identify fundamental limitations of T2T editing: it couples error detection with replacement, pollutes the generation context with potentially incorrect tokens, and introduces a train-inference noise mismatch where systematic model-generated errors differ from the random perturbations seen during training. We propose Token-to-Mask (T2M) remasking, a training-free, drop-in replacement for T2T editing that resets suspected erroneous tokens back to the mask state, allowing the diffusion process to re-predict them under cleaner context. We design and empirically validate three complementary error detection strategies -- probability-based, trigger-mirrored, and temporal-difference-based -- and provide a unified theoretical analysis showing that T2M remasking purifies the generation context, converts systematic inference errors back to the model's native mask noise type, and enables delayed commitment for joint multi-position optimization. Comprehensive experiments across 12 benchmarks spanning knowledge, reasoning, mathematics, coding, and instruction following show that T2M generally improves performance on tasks requiring precise token-level output, with the largest gain on mathematics (+5.92% on CMATH). Error analysis on CMATH reveals that the dominant failure mode is last-mile token corruption -- where correct reasoning produces a corrupted final answer -- and that T2M repairs 59.4% of such cases.

Lin Yao

School of Computer Science, Shanghai Jiao Tong University, Shanghai, 200240, China · Zhongguancun Academy, Beijing, 100097, China

Diffusion Language Models Diffusion Language Model Decoding Language Model Error Detection

Ask PDF

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