Source-linked AI summary
Random Remasking Scales Masked Diffusion Inference
Zhaoyi Zhou, Frank Wu, Andrea Zanette
TL;DR
Masked diffusion decoding can freeze errors, and existing theory does not readily explain iterative correction with an imperfect denoiser. The paper models inference as a remask–denoise Markov chain on clean sequences and proposes R2D inference with TCT training. R2D scales on logical tasks and works with pretrained models on HumanEval+, while TCT further improves logical-task accuracy.
Problem
Method
The paper models inference as a clean-sequence Markov chain using a remask–denoise kernel, motivating test-time R2D inference and trajectory-aware TCT training.
Source: §1 · p. 2Results
Takeaways & Limitations
R2D’s early-exit mechanism is more compute-efficient, and the approach also applies to pretrained diffusion models on practical language tasks.
Source: §6 · p. 9
Abstract
from arXiv · showhide
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+.
1 INTRODUCTION
Iterative correction can address irreversible decoding errors, but existing theory does not explain its benefits with an imperfect denoiser. The paper analyzes MDM inference as a clean-sequence remask–denoise chain and introduces R2D and TCT, reporting scaling on logical tasks and gains on language tasks.
1 INTRODUCTION
- Existing theory analyzes intermediate marginals under a perfect denoiser, not the actual generation trajectory produced by an imperfect denoiser.The paper frames that trajectory as the theoretically relevant object.Source: §1 · p. 1
- R2D with standard training raises Sudoku-Hard accuracy from 19.7% to 93.5% and 3-SAT accuracy from 86% to 98.5%; TCT raises them further to 97.5% and 99.3%, respectively.On HumanEval+, R2D with off-the-shelf LLaDA2.1-mini surpasses iterative correction baselines and official Q-mode inference.Source: §1 · p. 2
- The paper models MDM inference as a Markov chain on clean sequences, with standard decoding and correction represented by one remask–denoise kernel.Its single-step KL-divergence update bound identifies conditions under which the chain is more likely to move toward the data distribution.Source: §1 · pp. 1–2
- R2D is a reward-free, finetuning-free inference procedure motivated by the update bound, while TCT trains the denoiser with trajectory-aware supervision at every position.The analysis suggests a non-increasing timestep schedule and iterating at each step until the clean-sequence prediction stabilizes.Source: §1 · p. 2
- The paper reports that test-time scaling depends on token coupling and that TCT further boosts R2D’s self-correction ability.Random remasking is described as more favorable for tasks with local constraints, where revealed information can help reconstruct masked tokens.Source: §1 · p. 2
2 RELATED WORK
Prior iterative-correction methods either alter pretrained-model inference or train models to revise their own outputs. The paper distinguishes its framing by treating iterative remasking as a reasoning mechanism and a test-time scaling axis, rather than only as error repair.
2 RELATED WORK
- Training-free approaches modify inference, while trained approaches fine-tune or augment denoisers to recognize and revise their own mistakes.ReMDM remasks decoded tokens before re-prediction; ProSeCo trains a corrector that alternates unmasking and correction.Source: §2 · p. 3
- Reward-guided MDM scaling uses external quality signals, whereas the inference algorithms studied here are reward-free.The closest cited complete-sequence Markov-chain framing is Phunyaphibarn & Sung (2026), whose sampler has limited convergence analysis.Source: §2 · p. 3
3 PRELIMINARIES
The preliminaries define masked sequences and the token-prediction behavior of an absorbing MDM. Standard inference progressively unmasks a fully masked sequence and leaves revealed tokens fixed.
3 PRELIMINARIES
- A clean sequence is a length-L discrete sequence over vocabulary V; corruption adds a mask token and uses a monotonically decreasing noise schedule.The schedule has α0 = 1 and α1 = 0, and clean sequences are drawn from pdata.Source: §3 · p. 3
- An MDM predicts clean-token distributions at masked positions, with standard training using a weighted cross-entropy loss over those positions.The model’s conditionally independent per-position outputs induce a joint distribution used as its posterior over clean sequences.Source: §3 · p. 3
- Standard inference starts fully masked, follows a decreasing time grid, and reveals selected positions using sampled or arg-max token predictions.Once revealed, a token remains fixed; this paper departs from that convention by allowing correction of previously revealed tokens.Source: §3 · pp. 3–4
4 METHOD
The section reframes MDM inference as a Markov chain over clean sequences, enabling analysis of trajectories with an imperfect denoiser. It develops adaptive R2D inference and TCT to address scheduling and model error.
4 METHOD
- MDM inference can be represented as repeated denoise–remask transitions over clean sequences, with pdata stationary under the true kernel.This places standard decoding and iterative correction within one clean-sequence-space framework.Source: §4 · p. 4
- For convergence near pdata, remasking must be strictly Markovian, so R2D uniformly remasks tokens in clean x0 rather than only previous masks.The latter depends on mask history inaccessible from x0.Source: §4 · p. 4
4.1 SCHEDULE: RANDOM REMASK-DENOISE (R2D) INFERENCE
- R2D discretizes noise levels and holds each fixed, iterating the model kernel until predictions stabilize, yielding an adaptive schedule.Stability is heuristically detected when argmax clean predictions remain unchanged for τ consecutive denoise steps.Source: §4.1 · pp. 5–6
- R2D's standard variant preserves revealed tokens, while its correction variant overwrites them; both are test-time procedures for pretrained MDM denoisers.The standard variant also has error-correction ability through random remasking.Source: §4.1 · p. 6
4.2 TASK STRUCTURE: WHEN IS THE DENOISING ERROR SMALL?
- Sparse local constraints, as in Sudoku and 3-SAT, can make denoising error small after only some tokens are revealed, supporting operation at larger noise levels.Each Sudoku constraint spans nine cells, while each 3-SAT clause has three literals.Source: §4.2 · p. 6
4.3 MODEL QUALITY: TRAJECTORY CORRECTION TRAINING (TCT)
- TCT trains on remasked states from the chain's own marginal and applies clean-target token cross-entropy to masked and revealed positions, penalizing prior mistakes.Its training trajectory omits the correction inner loop, uses a fixed stage budget, and applies a low-confidence reveal policy.Source: §4.3 · pp. 6–7
5 EXPERIMENTS
The experiments compare R2D and TCT with iterative-correction baselines on logical and coding tasks, measure test-time scaling and compute, and ablate key design choices. R2D scales on logical tasks, performs strongly on HumanEval+, and exits early; the ablations examine remasking, early exit, decoding, and TCT training choices.
5 EXPERIMENTS
- The experiments compare R2D standard and correction variants with ReMDM and ProSeCo, alongside other correction-training and non-correction-training variants.Logical tasks use a 6.8M-parameter model trained from scratch; Sudoku-Hard contains 2.3M puzzles, with the last 1,000 used for testing.Source: §5 · p. 7, §5.1 · p. 7
5.1 LOGICAL TASKS
- TCT-trained R2D reaches 97.5% on Sudoku-Hard and 99.3% on 3-SAT, outperforming ReMDM (correction) by 5.2 and 2.8 points at the same budget.Standard-trained R2D already surpasses other baselines on both tasks; 3-SAT has higher peak performance than Sudoku-Hard.Source: §5.1 · p. 7, Table 1 · p. 8
5.2 CODING BENCHMARK
- R2D reaches 85.4% on HumanEval+, surpassing the published Q-mode accuracy of 82.9% and iterative-correction baselines.The evaluation uses off-the-shelf LLaDA2.1-mini, a 16B model already trained with self-correction ability.Source: §5.2 · p. 8, Figure 2 · p. 8
5.3 TEST-TIME SCALING
- R2D scales near-log-linearly, rising from 19.7% to 93.5% on Sudoku-Hard at M = 1 to 256 and from 86.0% to 98.5% on 3-SAT at M = 32.ReMDM improves with budget but stays below R2D at every budget; ProSeCo plateaus early at low performance.Source: §5.3 · p. 8
5.4 COMPUTE COMPARISON
- R2D uses about 1/4 and 1/2 of ReMDM’s compute while reaching higher accuracy, because it exits early when predictions stabilize.Figure 3 reports averaged NFE per sample at the Table 1 budget; ReMDM and ProSeCo spend their full budgets.Source: §5.4 · pp. 8–9, Figure 3 · p. 9
5.5 ABLATION STUDY
- Uniform remasking scales from 65.8% to 97.5% across a 16× budget increase, while low-confidence remasking moves only from 68.7% to 69.6%.Removing early exit preserves scaling but nearly quadruples NFE at M = 256 and lowers accuracy; stochastic decoding loses 10.4 accuracy points.Source: §5.5 · p. 9
- Top-k position selection is TCT’s load-bearing choice: random positions lower Sudoku-Hard accuracy from 97.5% to 85.5%, while stochastic reveal costs 0.7 points.The authors attribute the position-selection result to top-k keeping the training trajectory close to a valid partial solution.Source: §5.5 · p. 9
6 CONCLUSION
The paper presents a unified Markov-chain view of masked diffusion inference and uses its contraction analysis to motivate R2D and TCT. Together, they produce near-log-linear scaling on logical tasks, while R2D also performs strongly on HumanEval+ and uses early exit for compute efficiency.
6 CONCLUSION
- The paper frames standard decoding and self-correction as one remask–denoise kernel acting on clean sequence space.The authors identify applying R2D and TCT to domains such as biological sequence design as a promising future direction.Source: §6 · p. 9
A PROOFS
The appendix proofs establish the propositions underlying the paper’s analysis. They use Bayes’ rule, KL-divergence factorizations, and Jensen’s inequality.
A.1 PROOF FOR PROPOSITION 1
- A.1 expands the noisy marginal, applies Bayes’ rule, and sums over noisy states to establish the proposition.Source: §A.1 · p. 12
A.2 PROOF FOR PROPOSITION 2
- A.2 equates two factorizations of a joint KL divergence, then drops the nonnegative remainder to obtain the bound.The remainder is itself a KL divergence.Source: §A.2 · p. 12
A.3 PROOF FOR PROPOSITION 3
- A.3 applies Jensen’s inequality, averages over clean data, and subtracts the data entropy to obtain the claim.Source: §A.3 · p. 12
B TRAJECTORY CORRECTION TRAINING ALGORITHM
TCT trains on model-generated trajectories using a streaming batch of chains. The appendix also reports how inference error bars are estimated.
B TRAJECTORY CORRECTION TRAINING ALGORITHM
- TCT updates the correction loss on partially masked sequences, then refills chains at stage budget K or reveals top-k masked positions using argmax predictions.Chains advance in a streaming fashion across gradient steps.Source: B TRAJECTORY CORRECTION TRAINING… · pp. 12–13
C.1 ERROR BARS OF R2D AND REMDM INFERENCE.
- Inference error bars use three training seeds and three inference seeds for each training checkpoint.Source: §C.1 · p. 13
C.2 COMPUTE AND LATENCY
The section compares correction methods on Sudoku-Hard, 3-SAT, and HumanEval+ in terms of accuracy and compute. It also describes the latency trade-off as the correction budget changes.
C.2 COMPUTE AND LATENCY
- At the largest budget, R2D latency is about 77× standard decoding on Sudoku-Hard and 26× on 3-SAT, while M = 1 costs 1.31 s and 0.36 s, respectively.The correction budget M can be selected along the scaling curves.Source: §C.2 · p. 13
C.3 COMPUTE AND LATENCY ON HUMANEVAL+
- Table 3 pairs Figure 2’s accuracy with the cost of three correction methods; R2D reaches top accuracy using roughly half of ReMDM’s compute, due to early exit.Source: §C.3 · p. 13
- On HumanEval+, R2D reaches the top accuracy among three correction methods while spending roughly half of ReMDM’s compute.Source: §C.3 · p. 13
C.4 REMDM IMPLEMENTATION CHOICE UNDER TCT TRAINING
The section compares R2D with ReMDM under different inference and training choices, then ablates R2D’s design. Its comparisons report accuracy alongside scaling or compute costs.
C.4 REMDM IMPLEMENTATION CHOICE UNDER TCT TRAINING
- R2D scales near-log-linearly and stays above ReMDM at every correction budget across Sudoku-Hard and 3-SAT, under both training settings.Source: Figure 4 · p. 14
- For TCT-trained ReMDM, argmax sampling is used instead of the official stochastic sampling; the reported correction configuration uses carry-over unmasking with argmax.Stochastic sampling is weaker on both tasks, while 3-SAT favors carry-over under argmax sampling.Source: §C.4 · pp. 13–14
- 97.5% accuracy in 49.4 s/sample for R2D versus 92.3% in 187.9 s/sample for ReMDM on Sudoku-Hard.R2D is the only compared method that exits early.Source: Table 2 · p. 15
- On HumanEval+, R2D reaches the top accuracy among three correction methods while spending roughly half of ReMDM’s compute.Source: Table 3 · p. 15
D.1 R2D INFERENCE DESIGN
- Uniform remasking rises from 65.8% to 97.5% accuracy across a 16× budget increase, while low-confidence remasking moves only from 68.7% to 69.6%.Disabling early exit nearly quadruples NFE at M = 256 and still lowers accuracy; stochastic sampling loses 10.4 accuracy points.Source: §D.1 · p. 15
D.2 TCT TRAINING DESIGN
TCT aligns its reveal choices with R2D inference and keeps training trajectories near valid partial solutions. Ablations identify confidence-based position selection as the load-bearing choice, while argmax reveal has a smaller effect.
D.2 TCT TRAINING DESIGN
- Argmax reveal during TCT matches R2D's argmax inference operator, while top-k confidence reveal keeps training trajectories close to correct rather than letting them degrade into noise.The reveal choices are deliberate design choices, not arbitrary departures from theory.Source: §D.2 · pp. 15–16
- The correct Sudoku grid is an exact fixed point across masking ratios, but on 3-SAT token preservation is weakest near completion, making carry-over essential there but not on Sudoku-Hard.Table 4 compares carry-over and sampling variants at M = 256 for Sudoku-Hard and M = 32 for 3-SAT; argmax dominates, and carry-over is close to free on Sudoku-Hard but decisive on 3-SAT.Source: Table 5 · p. 16, Table 4 · p. 16
- Random position selection costs 12.0 points on Sudoku-Hard, whereas stochastic reveal costs 0.7 points, making position selection the load-bearing TCT choice.Top-k confidence selection keeps revealed prefixes mostly correct, so training predicts remaining masks from valid partial solutions rather than corrupted grids.Source: §D.2 · p. 16
E ADDITIONAL EXPERIMENT DETAILS
The additional experiments specify task-specific model and inference settings for Sudoku-Hard, 3-SAT, and HumanEval+. The logical-task ablations examine inference design and TCT reveal choices, while HumanEval+ evaluates inference algorithms on an off-the-shelf model.
E.1 SUDOKU-HARD
- On Sudoku-Hard, low-confidence remasking almost removes scaling, removing early exit preserves scaling at markedly higher NFE, and stochastic sampling costs accuracy at every budget while using more NFE.The denoiser is a 6.8M-parameter Qwen2-style model trained with TCT; inference uses correction frequency 1, patience 5, every grid point, and no early commitment.Source: §E.1 · pp. 16–17, Table 6 · p. 17
E.2 3-SAT
- The 3-SAT denoiser is a 6.8M-parameter Qwen2-style model trained from scratch with TCT on 1M synthetic instances using K = 5 correction phases.Its context length and vocabulary follow the SATLIB encoding of 20 variables and 91 clauses; inference uses patience 5, correction frequency 1, every grid point, and no early commitment.Source: §E.2 · p. 17
E.3 HUMANEVAL+
- HumanEval+ is inference-only: off-the-shelf LLaDA2.1-mini is evaluated without TCT or finetuning, so methods vary only in inference algorithm.All 164 problems use base HumanEval prompts and are scored against the plus test suite; correction is block-wise on the current 32-token block, with prior blocks and future masks untouched.Source: §E.3 · pp. 17–18