Improving Discrete Diffusion Unmasking Policies Beyond Explicit Reference Policies
Authors: Chunsan Hong
Seonho An
Min-Soo Kim
Jong Chul Ye
Venue: ICLR 2026 Poster (2026)
Keywords: discrete diffusion models
masked diffusion models
reinforcement learning
Primary area: generative models
TL;DR: We learn a unmasking policy model for masked diffusion models via a KL-regularized MDP (GRPO) that comes with convergence and KL-tightening guarantees

Abstract: Masked diffusion models (MDMs) have recently emerged as a novel framework for language modeling. MDMs generate sentences by iteratively denoising masked sequences, filling in [MASK] tokens step by step. Although MDMs support any-order sampling, performance is highly sensitive to the choice of which position to unmask next. Prior work typically relies on rule-based schedules (e.g., max-confidence, max-margin), which provide ad hoc improvements. In contrast, we replace these heuristics with a learned scheduler. Specifically, we cast denoising as a KL–regularized Markov decision process (MDP) with an explicit reference policy and optimize a regularized objective that admits policy-improvement and convergence guarantees under standard assumptions. We prove that the optimized policy under this framework generates samples that more closely match the data distribution than heuristic schedules. Empirically, across four benchmarks, our learned policy consistently outperforms max-confidence: for example, on SUDOKU, where unmasking order is critical, it yields a 22% gain over random and a 12% gain over max-confidence.

Published as a conference paper at ICLR 2026
IMPROVING DISCRETE DIFFUSION UNMASKING POLICIES BEYOND EXPLICIT REFERENCE POLICIES
Chunsan Hong', Sconho An², Min-Soo Kim², Chul Ye! 'Graduate School of AI, KAIST School of Computing, KAIST
ABSTRACT
Masked diffusion models (MDMs) have recently emerged as a novel framework for language modeling. MDMs generate sentences by iteratively denoising masked sequences, filling in [MASK] tokens step by step. Although MDMs support any- order sampling, performance is highly sensitive to the choice of which position unmask next. Prior work typically relies rule-based schedules (e.g., max- confidence, max-margin), which provide ad hoc improvements. In contrast, we replace these heuristics with learned scheduler. Specifically, we cast denoising as KL-regularized Markov decision process (MDP) with an explicit reference policy and optimize regularized objective that admits policy-improvement and convergence guarantees under standard assumptions. We prove that the optimized policy under this framework generates samples that more closely match the data distribution than heuristic schedules. Empirically, across four benchmarks, our learned policy consistently outperforms max- max-confidence: for example, on SUDOKU, where unmasking order is critical, it yields 20.1% gain over random and a 11.2% gain over max-confidence. Code available at hittt::::: //tpt/uggttttt.t// com/ chunsanHong/UPPS
INTRODUCTION
Diffusion models have achieved remarkable success in domains such as image generation (Ho et al. 2020; Song et al.) (202T). by defining noise-adding forward process and denoising reverse process in continuous space, and learning the reverse dynamics within this framework. Recent studies (Nie) al., 2025 Ye et al. 2025 Lou et al. 2024) have advanced this line of work into the discrete domain, applying it to language modeling through masked diffusion models (MDMs). In the discrete setting, MDMs define a forward process that replaces tokens in text with a [MASK] token and corresponding reverse process; training typically maximizes evidence lower bound (ELBO) (Lou et al., 2024) to fit the reverse dynamics. Small-scale MDMs (Lou et al., 2024; Sahoo et al. 2024) have matched surpassed autoregressive models (ARMs), and recent large-scale studies (Ye et al.
2025
Nie et al. 2025) report results competitive with strong ARMs such as LLaMA-3 (Grattafiori
al. 2024).
Despite a shared template, diffusion discrete versus continuous spaces differs substan- filly-especiallly at inference. Continuous-space diffusion models (Song al. [021] Karras et al..
2022)
perform denoising by integrating a stochastic differential equation (SDE) or an ordinary differ-
ential equation (ODE). In contrast, discrete-space diffusion models (Sahoo et al. 2024 Nie al.
denoise by iteratively sampling concrete replacements for [MASK] tokens. These mechanisms lead to different update geometries: continuous models accumulate numerous small changes across all dimensions, whereas discrete models enact localized, high-magnitude jumps at few coordinates. Consequently, effective denoising in discrete space requíres its own methodology. Just as advances in SDE solvers unlocked large gains for continuous diffusion, improved denoising strategies likely to yield substantial performance gains for MDMs.
2025)
This raises a natural question: what matters most in MDM denoising? Among several factors, focus on the choice of which position to unmask next. Kim et al. (2025) provide a theoretical rationale solve any- generation; i.e.,
for this choice, proving that no polynomial-tim cannot train MDMs that exactly recover the real data distribution for all masked sentences. However, they empirically show that selecting unmasking tokens via max-confidence (Chang al. 2022)
algorithm

---

Published as a conference paper at ICLR 2026
max-margin (Kim al. 2025) inference time can bypass sub-hard instances and even outperform
ARMs. Reflecting this, contemporary large-scale MDMs such as LLaDA or Dream-7B (Nic et al.
2025
Ye al. 2025) commonly adopt max-confidence.
In this paper, we aim to move beyond heuristic references (c.g., max-confidence) by learning improved unmasking policies. To this end, we cast MDM denoising Markov decision process (MDP) and train an unmasking position-selection policy using the group relative policy optimization (GRPO) (Saoo et al.. 2024). Theoretically, we prove that under this framework, the optimized policy produces samples closer to the true data distribution than the strong heuristic references. Moreover, we introduce three tractable surrogate objectives and provide propositions that justify optimizing them in place of the ideal, yet intractable. objective. Practically, our learned policy substitutes for max-confidence and related heuristics, achieving improved accuracy across four benchmarks. For example, on SUDOKU, where unmasking order is crucial, it delivers a 20.15 1 % improvement over random and 11.2% improvement over max- confidence.
Related works. Reinforcement learning (RL) has emerged as effective tool to fine-tune large lan- guage models toward complex objectives or preferences (Ouyang et al., 2022 Guo et al. 2025] Shao et al. 2024 Rafailov et al., 2023). In MDMs, Zekri & Boullé (2025) apply policy gradient methods to a pretrained discrete diffusion model so as to maximize a task-specific reward at the last denoising step. Zhao et al. (025) adopts the GRPO framework to improve the reasoning ability by rewarding only the final answer accuracy. Zhu et al. 2025) propose a variance reduced preference optimization (VRPO) for MDMs, which shares the concept of direct preference optimization (DPO) (Rafailov)
cet al. 2023) in ARMs. While the aforementioned works utilize RL to post-train the MDM itself, they differ from ours, where we train the unmasking policy model. Recently, Huang al. (2025) introduce DCOLT, which also trains an unmasking policy via reinforcement learning. Our approach differs in that we formulate denoising as a re-regular MDP with an explicit reference policy and optimize a regularized objective that admits policy- imnrovemmmo and convergence guarantees under standard assumptions. Furthermore, while Huang et al. (2025) jointly applies RI. train both the unmasking policy and the underlying MDMs for maximum accuracy, our work focuses solely optimizing the unmasking policy for frozen MDMs, thereby reducing computational cost and mitigating the risk of reward over-optimization.
2 PRELIMINARY
Notations. Let the distribution Pdata be the data distribution over sequences of length L. with and let the set of all possible sequences be X' VL. We define pL
vocabulary V {1.....m) m). as the Kronecker delta distribution that assigns all probability mass to the sentence in which every token masked. We denote the kernel forms of Paata and pr. by P data and P,. respectively, with for the 1-th step sequence with tokens. We denote the mask by MI and the set of sequences with n masks by Xr = {x ( e
entries [P [Pdata] Ø (a)) We write an {a xx xy ........ xh) ..,.,,,))
(VU (M)) E/= 1(x' = M) = n), so that x,, E Xn has n masks and we refer a [x] for the i-th mask index of æ. For example, if the number of masks in æ is n, then a² [x]] E {I,..., L}. x=(([) M and E (1..... .n.. For brevity, we will omit in ot [x]] for the rest of the paper. We denote MDMs as f and their model distribution as wy. We denote Markov transition kernel T x X - [0. where T(æ, x²) is the probability of æ transform into x' such that > Zx''x T(x, x') = for all æ.
Masked diffusion models. Masked Diffusion Models (Shi et al. 2024 Sahoo et al., 2024) define a mask-adding forward process and corresponding reverse process in discrete space, and learn the model Ty to mimic the true reverse process. The forward process gradually masks tokens independently one-by-onc in Do until the sequence is fully masked at æn with n L (or =1 where := n/ L). For fort= = n/LL E (0, 1), the sequence æ,, is partially masked, with each being masked with probability t or remaining unmasked with probability 1- t. The reverse process recovers the data distribution by iteratively predicting masked tokens as t decreases from 1 to O. While previous works (Campbell al. 2022) Sahoo et al. 2024) propose the time-dependent MDM training framework, the recent works (Zheng et al. 2025 Ou et al. 2025) propose the aine-agnostic model where time is implicitly included the number masks. Then, the time-agnostic loss for MDM is defined as follows:
(0) 2 -EE, -En,zo, st H
f 1
Ml log To [=f]
(I)
2

---

Published as a conference paper at ICLR 2026
Zheng et al. 2025); ) Ou et al. (2025) have shown that LMDM is the upper bounds of the negative log- legelikelihood. Zheng et al. (2025) propose the simple generation algorithm, named the First-Hitting sampler, beyond a continuous time Markov chain (CTMC) that follows the conventional reverse process. First-Hitting sampler randomly selects the index i for unmasking and samples the next token sequentially 1 21-1
Z%o- This is proven converge Pdata- Recent large-scale MDMs (Nic et al., 22025) Ye et al., (025) currently adopt time-agnostic modeling and further employ sampling strategies such as the max-confidence beyond the First-Hitting sampler.
Markov transition kernel of random sampler and training hard problem of MDM. Large-scale MDMs first select the masked token to unmask and sample from the categorical distribution. See Algorithm in Appendix B.2 for the detailed sampling algorithm. We here formulate general Markov transition kernels defined by various samplers of MDM.
T 9)
In
(9) (2 x)
(2)
i=1
where
( x) { g(aa(zn) =0(x0(((n)) if xep" = x-aa 0 otherwise,
[ n
(3)
where Σx= g(a ((aa|arn))) (.n) = x x--a = {x+,, xx?
}. 2 M by definition,
and no To (2= ( is the categorical distribution of model prediction of the i-th mask of a.
The non-homogeneous Markov kernel for First-Hitting sampler (Zheng et al. 2025) for time-agnostic MDMs can be defined as choosing by
1
g(a (x.) == Irand (a" (x
(5)
In algorithmic manner, this is equal to first sample ~ Uf(1..... n). then sample za ~ Te (x Ixn). and obtain 2n-1 by replacing i-th mask of æn with zad In this case, (II Thrand) PL ~ P Pdata if Ty is perfectly estimated (Zheng et al.) 2025). However, Kim et al. (2025) have proven that learning my for every combination of masked sequence cannot be solved by a polynomial algorithm, and there exist hard subproblems. This can interpreted as there exists some points that true reverse process
and model's s reverse process are different, such that resulting in (II L provide simplified example to help understand:
Trandd Pt + Pdata. We
Example 1. Consider the following simplified zebra (Einstein) puzzle. There are two houses, and the task is determine the person living in each house and their favorite food. The possible persons are Robert and Tom, and the possible foods are pizza and hamburger. The clues are given as follows: 1) The person living in the first house is Robert, and 2) Tom likes pizza.
Now consider the MDM directly answer the questions, filling the mask of sentence: "House I: name: [MASK], food: [MASK], House 2: name: [MASK], food: [MASK]". Obviously, if MDM tries to answer the name of the person who lives in the first house, the problem becomes trivial. However, it tries to fill attributes for the second house first before filling attributes of the first house, the problem becomes much harder. As the number of questions and clues grows, this kind of discrepancy would be much larger. Indeed, the zebra puzzle is also known as NP-hard problem (Shah al., 2024 Yato & Seta 2003). Note that this is a highly simplified example to help understand the problem of MDM training, and refer to Appendix A for more details and Kim al. 2025) cxplanation.
for complete
Practical approach to avoid such hard problems. Meanwhile, Kim al. 2025) empirically shows that unmasking the token with max-confidence or max-margin helps to avoid such hard problems, leading to better generations. Therefore, in practice, rather than randomly sampling the token for unmasking, large-scale discrete diffusion models (Nie et al., 2025; Ye et al., 2025) utilize such rule-based samplings. The weighted kernel related to the max-confidence sampling can defined with g( (aa1 xx,)) := feant defined as follows:
(a' (m)) 1(a' arg max max To(= c(a,))), a²
(6)
cEV
3

---

Published as a conference paper at ICLR 2026
or in other stochastic forms,
Icont Ocont~(a'(2xn) Eeen exp(Te(af c(xn)/T)
s.t. lim Icontr 7-10+
Icont,
Ew Eeev exp( = [T)
(Top-K(a'lan,) 1 1(a' argtopk(max Te (x clan))) s.t. 9Top-1
Jeont
K
cEV
where argtopk denotes the indices from the top-k sampler. Similarly, the weight kernel for max- margin sampler can be defined with Imargin whose value is 1 where is the mask index with the token.
maximum margin between the first-confident and
3 IMPROVING UNMASKING POLICY BEYOND A REFERENCE POLICY
While twisted Markov transition kernels driven by heuristic unmasking policies-such as Joont or Imargin deliver substantial gains, we posit that stronger policy g can guide MDM denoising more effectively. Motivated by this, we learn instead of defining in a heuristic manner to find an optimal path. We further cast the learning problem of 96 as KL-regularized MDP, where strong heuristic sampler serves as a reference policy, enabling stability and policy improvement guarantees.
deeiiver
3.1 EXISTENCE OF OPTIMAL PATH BEYOND 9conf
We test various in GSM8K (Cobbe et al. 2021), the mathematical questions benchmark, probe whether max- confidence is merely strong or whether substantially better unmasking paths exist. We compare three schedulers: Iconf, Iraud, and 9Tp--K- Here, Jcont is deterministic, while [lup-K smoothly interpolates between random and max-confidence. For controlled comparison to analyze the effect of generation ordering and following the default setting used in LLaDA (Nie it al. 2025), we conducted token sampling with argmax op- erator, i.e., Sn arg max, Ty (z I x,).
0.9
Pasc@/N
Max-Confidence
Top-5
Random
Number of paths (N)
Figure 1: @ass@N on GSM8K.
For each test question, we draw N independent trajectories using either Irand or @Top-5 and report @ass@N, the fraction of questions for which at least one trajectory is correct. Fig- ure shows three consistent trends:
The dashed red line is the single-trajectory accuracy of max-confidence Jconf.
Jeonf significantly
outperforms Grand when N = 1; (ii) gTop-5 dominates grand across N, highlighting the strength of confidence-based ordering; (iii) as N grows, both curves surpass the max-confidence line by a clear margin, and 9Top-5 exceeds 0.92 when N 10.
These results carry two implications. First, max-confidence indeed a strong single-path baseline. Second, the rise of @ass@N above the max-confidence line indicates the presence of higher-yield unmasking trajectories than the deterministic max-confidence order. Brute-force search for such paths is computationally prohibitive and requires unobserved oracle feedback, motivating our learned policy Io that aims to discover high-reward paths without enumerating many trajectories.
3.2 REFORMULATING THE PROBLEM AS REINFORCEMENT LEARNING
Motivated by the strength of goonf and the existence of higher-yield unmasking paths, we set two goals: (i) increase the probability that an MDM produces the correct answer, and (ii) move the induced sampling distribution closer to Pasta than strong heuristic reference Jref- We therefore learn stochastic position-selection policy Jo rather than fixing g heuristically. The policy is parameterized as softmax over masked positions using learnable scorer (x:) E R24:
exp (ap(ihs (Zn))
go(a" [..)
Σ exp (I
We cast unmasking as a regularized Markov decision process (MDP) with verifiable terminal rewards. Each episode begins from prompt-answer pair where the answer is fully masked, {æy, ql with
4

---

Published as a conference paper at ICLR 2026
Pe referring to the answer distribution and æy M' and terminates when all masks removed at n = 0. yielding Ty E X = V At state zn the action space the set of masked indices Ar, {a' [z........ a" ["[x,]): ;; the policy (- æx) selects index an E Ar.. to unmask. Transitions factor through the fixed MDM denoiser Tg:
Pas (In-1 xx)) = The (En, Ln-1) go(an (n) Te( Dn, an);
where (T nn an) is the token distribution for position On- The agent thus controls where to unmask, while the environment dynamics induced by <p. Rewards verifiable: external checker evaluates the fully unmasked output Zo and returns r(q. æa) E 1).
To stabilize and accelerate training, we introduce strong reference policy gret (e.g., Top-K) and optimize a theoretical KL-regularized GRPO-style objective (Mroueh 2025). We optimize directly over the terminal-oulpul distribution induced by the policy: all expectations taken with respect to Po(m (mo q). and the regularizer compares terminal-output marginals. Accordingly, our theoretical loss is an output-level GRPO objective without clipping:
max Ea-pe Exo (119) Preau (o
Pao ( q)
A(q, 20) DKI. (Pas (20lg) PPror (2o (q))
q)
(10)
where A(q, 20) is standardized advantage. Let (g(g) = Bno (:ap), (r/.)//q.. æoo)
[0, 1] and stdg(q)
V [Fam)ppp(1 [(((2)) - rg(()))), then A(q.zo)
(r(q. xo)
T Tgeai (q))/(stdgee.) doold (q) + e). This objective follows the output-level GRPO loss used in DeepSeek-
R1 (uuo al. (2225) rather than the original token-wise GRPO used in DeepSeekMath (Shao et al..
2024). The regularizer keeps go close to fref. acting as a trust-region that mitigates instability while
enabling improvement over strong baseline; selecting an appropriate goer provides high-accuracy starting point and accelerates convergence.
3.3 THEORETICAL ANALYSIS
In this section, we provide how go optimizing Eq. 10) can guarantee the convergence of the policy model and superior performance over the reference policy 9ref- Mrouch (2025) have proven that policy model learning output-level GRPO loss (Eq. (I0) can converge to its optimal performance
Tg', which is higher than N "rer"
Theorem (Restatement of GRPO convergence theorem (Mroneh, 2025 define for 0:
Assume 1 > Tgget 0,
hepee (Ta)
(11)
I-T exp( 1 B/ (1
1 1
Foret
Then, for the local optimizer &n max Fq. 10 where dold = Φn-1, the probability of success satisfies the following fixed point iteration ie we have almost surely for all q 1,
Toon (q) he, (Tae.. (q)
(12)
and T "oo (q)= = (gr. Also, let rg be a fixed point of he and assume that have We (Tg )1 < Given that hep and Wi heepper are continuous in [0, 1), following local fixed point convergence is established:
lim Tasn Tg"
(13)
and rg- Taed'
Therefore, setting gvot as a strong proxy model is important, as we can always get a better policy Iø.. Moreover, in the following we prove that policy model & maximizing Eq. (10) can restore Pasta better than gref:
Theorem 2 (Reference-Fe KL Tightening for MDM Policy Improvement). Assume there exists an ideal MDM 0 satisfying MDM LMDM((") (0*) for Pdata where data is composed of prompt and oracle

---

Published as a conference paper at ICLR 2026
answer {q. a). Assume select well- reward r(q. a) E {0. 1) such that Pdata (q. a) 0
r(q.a) = and 0 < gag < 1. Now consider incomplete MDM such that LMDM (0) > 0, and -ø-
that satisfies Eq.
(13) for 0. Let gret be a policy model such that, for every q Paier (mo q) >
whenever pe- (xo
0, i.e., supp(po-(( q)) S supp(Pana (- q)) Then, the following
inequality holds:
DeI (200) PP (xo (1) <DI. (20 Pais (2))
(14)
Or equivalently,
DKL (P (a/g)) (wo ( DKkl)))>KKK(Paaaa (Pdataaala)))lppppg (alq)
(2r q))
(15)
Proof sketch. Mrouch 2025 provides policy dynamic under training Eq. (10). Using the result given
in Lemma ii Appendix, we can express Iø- with 9ref and thereby derive KL divergence. Refer to Appendix for detailed proof and reason why the support assumption holds.
The interpretation of Theorem 2 as follows. If we choose binary reward r(q, a) E {0, } satisfying Pdaa (q. a) > 0 = aq.aa) 1 and 0 < ge < 1, then the theorem guarantees that the KL divergence is strictly tightened compared to the reference policy. Consequently, one may view any reward that satisfies these conditions learning dynamics that bring the sampling distribution closer to Pdatu Under the assumption that the dataset is well constructed, such binary rewards naturally correspond to properties such correctness, the absence of reasoning errors, grammatical validity, or even combinations thereof, any of which encourage the model move its output distribution closer the true data distribution.
"good" reward in the sense that it provably induces
Our provided Theorem 2 guarantees the closer sampling to Pdata or ideal MDM 0" than gref- Consequently, under the assumptions of two theorems, initializing with geo Iref and repeatedly optimizing Is using Eq. (IO) yields convergent policy 9ø- that (i) attains a higher expected reward than gref and (ii) generates samples that are closer to the real data distribution than those from gref-
3.4 PRACTICAL REALIZATION OF UNMASKING POLICY OPTIMIZATION
Here, we provide practical and tractable training objectives with various Jret- In real-world setting where L and [V] are high, Eq. I0 is not directly learnable since Pa. (æo (q) is intractable in MDMs. It requires marginalizing over every trajectory to Zu. i.e., Ps. (2o q) = E La.....r ær.
Ps Po(zzoz/ (q), making
both the main objective (reward maximization) and the KL term in Eq. (10) intractable. To address this, provide surrogate loss that alternate the main objective.
Proposition 1 (Output-Token) Level Gradient Alignment (informal)). For MDM and unmasking policy model Io where Poo (In=1 xx) = (a xn) T x | æn an) and Pao (Do (4) Eai....e, P Po (ZooL q), consider the output level loss Loutput and token-wise policy loss Lroken'
Pgo (Eo
9) q)
Coutpue (Ø) Earup Ex Poo (119)
A(q, 27))
(16)
Posaad (#o
ga(an xn)
Lroken (Ø) Eqrip 1E Man Popsa (19)
Alq mo)].
(17)
Doas (an xnn)
n=
Assume that optimize o carefulll such that & RS dotd the two gradients approximately equal:
V Erotenn Re V VoLoutput
(18)
and V Ltoken = V &ouupput at first optimization step where o = Pold-
Proof sketch. Although the result can be readily inferred from standard RL, we provide a rigorous proof for our setting: sparse-reward, finite-horizon, episodic, and non-stationary MDPs where we maximize GRPO-style objective. The full proof is in Appendix C.2
While Eq. 16) optimizes Pas (xolg) Exg..... Po( which is intractable, our surrogate
loss in Eq. 17 optimizes go (an x/n) Doms (an xx,)) which is tractable once trajectory is sampled.

---

Published as a conference paper at ICLR 2026
Trainable 1: Layer Transfomeer th
Feat
[Mask]
Frozes
Feal
Frozen
Fest
Trinable
0.8
[Makk]
Tiaesf
MDM
Feat
Layer
MLP
Pribs
Feat
s-aayer
MLP MLP
0.2
The structure of our unmasking policy model. The model is composed of single Transformer layer and 3-layer MLP. In inference time, the model proceeds as follows: 1) Given an input sentence, we run the base MDMs transformer and extract feature. 2) This feature branches into two paths: first, the base MDMs continue to produce a prediction of token distribution, and second, fed to the policy model's transformer layer. 3) We concatenate the extracted feature with the base MDM's Top-K probabilities. 4) The concatenated feature then processed by 3-layer MLP to yield policy over unmasking positions. The shared structure with frozen MDM supports memory--fficient training (e.g., LLaDA: 8B, policy model: 134M), where we only update the unmasking policy model. The details and memory-efficient training algorithm are provided in Appendix B.2
Figure 2:
Table 1: Various realization of Luro Iref D) varying Jref-
Reference policy
Mlx--Confieence
Softmax Realization
Top-K Realization @Tbp-K
9nd
Icoat
Divergence D(p (Pg..))
Cross-Fatropy Loss
Grop Peulaalee
E Segs
Parametrization
exp (b. (m)))
excal (x))).)))
9o (a' z))
Mal baiae) exp (x))
Σ Jee Mask nnckow) (xp((Ag(m))ar
Model initialization do
Model Pretrained with Cross-Entropy Loss w.E.t. Goonf
Random Initialization
Now, we detail how we train the policy model efficiently. Our unmasking policy model structure is shown in Figure g and the full training algorithm table can be found in Appendix B.2
Final Training Objective. By using Proposition T. clipping, and Monte Carlo approximation, we transform the theoretical loss (Eq. (TU)) into tractable unmasking policy optimization (UPO) loss as follows:
1
G
Jo (af9) (9) (x (g)
Luro(gs Fret, D) = pappp(110 Pot ( ( i PPowold
Σ m=1
min
Ag:
q-pa
G
Donr (al
(g) (g)
clip(
I (an
e, + +))a))
B
DPpo IP
Poara (an (o!!))
(g)
(19)
where D(pso I (Pac)) should be empirical and tractable estimator of DKL. (pg, (wo q)) || P (wro k4) and clipped loss has been proven to attain global optimality (Huang et al. 2024). Note that (g) in superscript of Eq. (19p is the group index.
Realizations of LUPO via the choice of fref- We consider three reference policies and 9Top-K-Ta and summarize their associated divergence terms in Table [ In Table
Icont, Feont we provide
the tractable surrogate LKI. that is justified by Propositions and and in Appendix C.3 Since we are sampling the trajectory from old policy model while training, we design the estimator LKT. that can be measured by the probability of the trajectories, Pa. (E):1/9), Pge.n (æo.zlq), and Patas (Zo.L q). For softmax realization, since fcontr assigns nonzero probability to all trajectories, we can replace DKL with the tractable surrogate LKL. We parameterize Jo as the softmax distribution over all masked positions by default (see Eq. (9)) for softmax realization. For pretraining, although one could initialize by matching 90u to Iconf via KL minimization, we observe stronger performance with cross- eross-entropy. For 9Top-K. there exists some sampled trajectory where Pgrop- (#u.L lq) 0 since 9Top assigns zero probability to mask index out of Top-K. In this regard, we reparameterize P as the softmax distribution of his over the Top-K K set only. This reparametrization gives two
7

---

Published as a conference paper at ICLR 2026
Sudoku
GSMSK
0.80
0.35
Accuracy
Accuracy
8.70
0.5
Max-Cont Realizarion
Softmax Realization
0.00
Random Initiaizzation
Top-K Realization
400
1901
2406
1300
600
100
1200
169t
Step
Step
Figure 3: Training accuracy on SUDOKU and GSM8K under LUPO (go. 9ref. D) for different choices of gref. Reference realizations are summarized in Table "Random initialization" denotes training randomly initialized model with Curo(9a: 20).
benefits: 1) we can utilize LKI. since Pyrep is non-zero for all trajectory sampled from Pgemop-x' and 2) we do not need to pretrain policy model since it is same as 91op-K when randomly initialized. In contrast, the determinism of Jconf renders LKL ill-posed whenever 9o assigns zero mass to the chosen index, derivations and detailed explanations are provided in Appendix B.1
train with CE(gcout l| 9e) while retaining the regularized-MDP perspective. Full
4 EXPERIMENTAL RESULTS
4.1 EXPERIMENTAL. SETTING
Models and benchmarks. We evaluate on the widely used large-scale MDM LLADA-8B INSTRUCT (Nie et al., (2255) and four datasets: (i) logic-puzzle benchmarks Secchmarks-SUDoKK (Zhao
ct al., 2025) and ZEBRA Shah et al., 2024) and (ii) mathermatical-rrasoiing bench-
marks- GSM8K (Cobbe et al. 2021) and MATH500 (Lightman et al. 2023).
(Nie et al. 2025) for all experiments. For logic puzzles, the model receives prompt and masked solution; then denoises by unmasking one position at a time. The reward is the fraction slots predicted correctly (e.g., in SUDOKU, the proportion of blank cells filled with the correct digits). For mathematical reasoning, the model is given a problem and denoises a length- 128 masked sequence. Following the LLADA generation protocol, by each baselines sequentially over the bins. The reward is if the final answer matches the ground truth and otherwise. In addition, we incorporate the of rog-prbbabilities produced by (o()) as an auxiliary reward signal. Detailed explanation and settings such as training aper-parameters are provided in Appendix We employ the dense reward rather than the binary reward because it leads to faster convergence and higher accuracy in practice. Nonetheless, we also verify that our proposed method outperforms the max-confidence baseline under the binary reward and report these results in Appendix D.I
Generation protocols and rewards. We utilize
partition the 128 positions into four bins of size 32 and perform sampling
4.2 TRAINING DYNAMICS ACROSS REFERENCE POLICIES
We train our policy model on GSM8K and SUDOKU by optimizing Luro (arrolge-rre 9ref D) with three reference realizations from Table Jeouf; Dooutt and 9Top-K. K- We also report ablation that drops the divergence term, i.e. Luro (go, 2, Ø), with randomly initialized policy, which is conceptually equivalent to DCOLT. The result is shown Figure The strongest configuration depends on the task: on SUDOKU, max-confidence realization achieves the best accuracy, while on GSM8K, Top-K realization performs best. We conjecture that GSM8K, restricting exploration to the Top-K indices stabilizes learning under LUPO due the much larger number of masked tokens. Across the two benchmarks, our optimal configuration outperforms the randomly initialized model that is, DCOLT (green bine)-by margins of at least 2% and 3%, respectively.
8

---

Published as a conference paper at ICLR 2026
Table Accuracy of each sampling strategy. Best per row in bold.
Sampler
Benchmark
Random Zheng al. 2025
Margin Kim et al. 2025
Entropy
Confidence Chang al. [0222]
Yc al. 2025
Ours
SUDOKU
0.616
0.713
0.671
0.705
0.817
Zebra
0.339
0.346
0.351
0.337
0.362
GSM8K
0.612
0.671
0.667
0.684
0.703
Math500
0.196
0.284
0.266
0.272
0.284
Table Accuracy on GSM8K of post-trained LLaDA with diffu-GRPO by sampling strategy.
Random
Margin
Entropy
Confidence
Zheng et al. 2025
Kim al. 2025
Ye et al. 2025
CChng et al. 2022
Ours
0.638
0.750
0.740
0.751
0.764
4.3
MAIN RESULT
Baselines and setup. We compare our learned unmasking policy with four widely used sched- ulers: random (Zheng et al. 2025), margin (Kim et al. 2025), entropy (Ye et al. (2025), and confidence (Chang et al. 2022), For SUDOKU and ZEBRA, we utilize max-confidence realization, and for GSM8K and MATH500, we use Top-K realization where == 5.
Results. Table 2 shows that our policy matches or surpasses all heuristics on every benehmark. For SUDOKU and ZEBRA, where use Jconf the reference, our model exceeds this reference: SUDOKU improves from 70.5% to 81.7%, and ZEBRA from 33.7% to 36.2%. On GSM8K, trained with Top- K reference, our policy reaches 70.3% and exceeds the strong mox-ccofieecce baseline at 68.4%. On MATH500, also trained with Top-K K reference, our policy attains 28, 28.4%, tying the best baseline while improving over max-confidence 27.2%. Overall, max-confidence remains strong single-path heuristic, yet a learned policy trained in our regularized framework consistently surpasses it.
4.4 ABLATION STUDY
Compatibility with diffu-GRPO. Zhao et al. (2025) introduces diffu-GRPO, an RL method that post-trains the MDM itself via GRPO. Our approach also uses RL but targets different component: we optimize the unmasking policy while keeping the base MDM fixed. The two methods are therefore complementary and can be combined. To demonstrate this, we first post-train LLADA- SB-INSTRUCT on GSM8K with diffu-GRPO, then train our policy model with Top-K realization, EUpo(Geno 9Top-3-p-- 3, L. As reported Table 3 our unmasking policy model yields additional gains on top of the diffu-GRPO model: + +12.6% over random and +1.3% over max-com max-confidence.
Zhao
Power regularization. We compare the mean and standard deviation of the group reward while optimizing Curo Gref D) in two setups: for SUDOKU, we train pretrained policy with/without CE to Iconf, and for GSM8K, we train a randomly initialized Je with/without LKL to 9Top-K- As shown in Figure 4 adding the divergence term yields higher final accuracy and, at similar mean rewards-even near --maintains larger group-reward standard deviation. This helps the policy model avoid premature convergence under the GRPO framework and reach higher convergence point. We conjecture that, for GSM8K, LKL to 9Top-K force the probability of the trajectory sampled from IoTop-K to be all equal and prevent early path collapse. For SUDOKU, suppose that CE toward Jconf counteracts collapse onto suboptimal order by pulling the policy back to high-confidence alternatives, keeping multiple competitive paths active and maintaining variance.

---

Published as a conference paper at ICLR 2026
SUDOKU
GSMSK
Groep Roward Mean
0.52
Reward Std.
018-
Promaned CIt
Groep Rowand Mean
Std.
0.54
dreeddedd weoring
0.50
R.12
0.42
0.78
005
0.38
630
1200 1958
3400
3000
1240
1840
2000
3000
800
2006
3200
4058
800
1600
2000
4006
Step
Step
Step
Skep
Figure 4: Mean and standard deviation of group reward during training under LUPO- For SUDOKU, we compare pretrained policy trained with/without CE to Seonf, and for GSM8K, we compare randomly initialized I6Top- trained with/without KL to 9Top-K.
5 CONCLUSION
We introduced KL-regularized MDP formulation for learning the unmasking policy in masked diffusion models. Our analysis provides two guarantees: convergence to fixed point that improves expected reward over the reference policy, and a KL tightening showing that the femmnna--oupple distribution is closer to the real data than the reference. To make the framework practical, we derived tractable surrogate loss and proposed several reference realizations. Across SUDOKU, ZEBRA, A, GSMSK, and MATH500, the learned unmasking policy model consistently matches or surpasses strong heuristics.
Discussion. Our learned unmasking policy improves performance across benchmarks, though the gains on GSM8K and MATH500 are relatively smaller than on SUDOKU. We conjecture that this stems from two differences between the domains. First, SUDOKU exhibits relatively clear and consis- tent ordering structure- example, certain positions are statistically more constraining-wheeee such regularity is much weaker in mathematical reasoning tasks. Second, while SUDOKU requires learning ordering patterns despite having far fewer training examples. These factors likely make more challenging to extract reliable ordering cues in GSMSK, resulting in modest gains.
small digit space, GSM8K requires learning over the full vocabulary,
Future directions. Our work demonstrates that GRPO-based training can yield task-specific
unmasking policies that outperform max-coonidence and other rule-based heuristics. As discussed above, an important next step is to develop more generalizable policies that remain effective on larger and more diverse language datasets; explore this direction further in Appendix D.2 with additional experiments. Another promising direction is to design more stabilized policy optimization methods. On SUDOKU, the full reparametrization Is outperforms Jonex suggesting that near-optimal training benefits from occasionally selecting low-confidence positions. In contrast, the much larger search space in GSM8K favors the stability of the Top-K Top-K-restricted policy. Developing an optimization scheme that enables effective exploration of all masked positions, even in large search spaces, constitutes an important avenue for future work.
ACKNOWLEDGMENTS
This work was supported by the National Research Foundation of Korea under Grant RS-2024- 00336454, by the Institute of Information & Communications Technology Planning & Evalua- tion(IITP) grant funded by the Korea government(MSTT) (RS-2025-023044) ship(KAIST) & RS-2024-00457882. AI Research Hub Project), and by AI Computing Infrastructure Enhancement (GPU Rental Support) User Support Program funded by the Ministry of Science and ICT (MSIT), Republic of Korea 12Q2-25-200217).
AT Star Fellow-
ETHICS STATEMENT
We follow the ICLR Code of Ethics. Our work trains a lightweight unmasking policy while keeping the base MDM frozen, and evaluates only on public benchmarks (SUDOKU, ZEBRA, GSM8K, MATH500) that contain no personally identifiable information; human-subject data were collected and no IRB approval was required. As our method optimizes only a lightweight scheduling policy to maximize verifiable terminal reward and does not modify the base MDM's parameters training
10

---

Published as a conference paper at ICLR 2026
data, the safety and faimesss properties are largely inherited from the underlying model. Absent defects in the base model, the incremental risk is limited to potential amplification of existing behaviors; in practice, this can be mitigated by pairing scheduler with standard safety filters and auditing the combined system on downstream tasks.
REPRODUCIBILITY STATEMENT
The training algorithm is detailed in Algorithm 2 the overall experimental setup is described in Section 4 and the hyperparameters and prompt templates used for evaluation are provided in
Appendix E. Model architecture and the policy inference procedure are summarized in Appendix B Complete proofs of all theorems and propositions stated the main paper appear in Appendix
REFERENCES
Jacob Austin, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne Van Den Berg. Structured denoising diffusion models in discrete state-spaces. In Advances in Neural Information Processing Systems (NeurIPS), volume 34, pp. 17981-17999 2021.
Andrew Campbell, Joe Benton, Valentin De Bortoli, Thomas Rainforth, George Deligiannidis, and Arnaud Doucet. continuous time framework for discrete denoising models. In Advances in Neural Information Processing Systems (NeurIPS), volume 35, pp.
2022.
Huiwen Chang.
Han Zhang, Lu Jiang, Ce Liu, and William T. Freeman. Maskgit: Masked generative image transformer. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2022.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems, 2021. CRLhttps::1 /tt.iiv.//aaxxv.ooxrr abs/2210.14163
Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, Amy Yang. Angela Fan, Anirudh Goyal, Anthony Hartshorn, Aobo Yang, Archi Mitra, Archie Sravankumar, Artem Korenev, Arthur Hinsvark, Arun Rao, Aston Zhang, Aurelien Rodriguez, Austen Gregerson, Ava Spataru, Baptiste Roziere, Bethany Biron, Binh Tang, Bobbie Chern, Charlotte Caucheteux, Chaya Nayak, Chloe Bi, Chris Marra, Chris McConnell, Christian Keller, Christophe Touret, Chunyang Wu, Corinne Wong, Cristian Canton Ferrer, Cyrus Nikolaidis, Damien Allonsius, Daniel Song, Danielle Pintz, Danny Livshits, Danny Wyatt, David Esiobu, Dhruv Choudhary, Dhruv Mahajan, Diego Garcia-Olano, Diego Perino, Dieuwke Hupkes, Egor Lakomkin, Ehab AlBadawy, Elina Lobanova, Emily Dinan, Eric Michael Smith, Filip Radenovic, Francisco Guzmán, Frank Zhang, Gabriel Synnaeve, Gabrielle Lee, Georgia Lewis Anderson, Govind Thattai, Graeme Nail, Gregoire Mialon, Guan Pang, Guillem Cucurell, Hailey Nguyen, Hannah Korevaar, Hu Xu, Hugo Touvron, Iliyan Zarov, Imanol Arrieta Ibarra, Isabel Kloumann, Ishan Misra, Ivan Evtimov, Jack Zhang, Jade Copet, Jaewon Jan Geffert, Jana Vranes, Jason Park, Jay Mahadeokar, Jeet Shah, Jelmer van der Linde, Jennifer Billock, Jenny Hong, Jenya Lee, Jeremy Fu, Jianfeng Chi, Jianyu Huang, Jiawen Liu, Jie Wang, Jiecao Yu, Joanna Bitton, Joe Spisak, Jongsoo Park, Joseph Rocca, Joshua Johnstun, Joshua Saxe, Junteng Jia, Kalyan Vasuden Alwala, Karthik Prasad, Kartikeya Upasani, Kate Ke Li, Kenneth Heafield, Kevin Stone, Khalid El-Arini, Krithika Iyer, Kshitiz Malik, Kuenley Chiu, Kunal Bhalla, Kushal Lakhotia. Lauren Rantala- Yeary, Laurens van der Maaten, Lawrence Chen, Liang Tan, Liz Jenkins, Louis Martin, Lovish Madaan, Lubo Malo, Lukas Blecher, Lukas Landzaat, Luke de Oliveira, Madeline Muzzi, Mahesh Pasupuleti, Mannat Singh, Manohar Paluri, Marcin Kardas, Maria Tsimpoukelli, Mathew Oldham, Mathieu Rita, Maya Pavlova, Melanie Kambadur, Mike Lewis, Min Si, Mitesh Kumar Singh, Mona Hassan, Naman Goyal, Narjes Torabi, Nikolay Bashlykov, Nikolay Bogoychev, Niladri Chatterji. Ning Zhang, Olivier Duchenne, Onur Çelebi, Patrick Alrassy, Pengchuan Zhang, Pengwei Li, Petar Vasic, Peter Weng. Prajjwal Bhargava, Pratik Dubal, Praveen Krishnan, Punit Singh Koura, Puxin Xu, Qing He, Qingxiao Dong, Ragavan Srinivasan, Raj Ganapathy, Ramon Calderer, Ricardo Silveira Cabral, Robert Stojnic, Roberta Raileanu, Rohan Maheswari, Rohit Girdhar, Rohit Patel, Romain Sauvestre, Ronnie
Plawiak,
11

---

Published as a conference paper at ICLR 2026
Polidoro, Roshan Sumbaly, Ross Taylor, Ruan Silva, Rui Hou, Rui Wang, Saghar Hosseini, Sahana Chennabasappa, Sanjay Singh, Sean Bell, Seohyun Sonia Kim, Sergey Edunov, Shaoliang Nie, Sharan Narang, Sharath Raparthy, Sheng Shen. Shengye Wan, Shruti Bhosale, Shun Zhang, Simon Vandenhende, Soumya Batra, Spencer Whitman, Sten Sootla, Stephane Collot, Suchin Gururangan, Sydney Borodinsky, 'aamar Herman, Tara Fowler, Tarek Sheasha, Thomas Georgiou, Thomas Scialom, Tobias Speckbacher, Todor Mihaylov, Tong Xiao, Ujiwal Karn, Vedanuj Goswami, Vibhor Gupta, Vignesh Ramanathan, Viktor Kerkez, Vincent Gonguet, Virginie Do Vish Vogeti, Vítor Albiero, Vladan Petrovic, Weiwei Chu, Wenhan Xiong, Wenyin Fu, Whitney Meers, Xavier Martinet, Xiaodong Wang, Xiaofang Wang, Xiaoqing Ellen Tan, Xide Xia, Xinfeng Xie, Xuchao Jia, Xuewei Wang, Yaelle Goldschlag. Yashesh Gaur, Yasmine Babaei, Yi Wen, Yiwen Song, Yuchen Zhang. Yue Li, Yuning Mao, Zacharie Delpierre Coudert, Zheng Yan, Zhengxing Chen, Zoe Papakipos, Aaditya Singh, Aayushi Srivastava, Abha Jain, Adam Kelsey, Adam Shajnfeld, Adithya Gangidi, Adolfo Victoria, Ahuva Goldstand, Ajay Menon, Ajay Sharma, Alex Boesenberg, Alexci Baevski, Allie Feinstein, Amanda Kallet, Amit Sangani, Amos Teo, Anam Yunus, Andrei Lupu, Andres Alvarado, Andrew Caples, Andrew Gu, Andrew Ho, Andrew Poulton, Andrew Ryan, Ankit Ramchandani, Annie Dong, Annie Franco, Anuj Goyal, Aparajita Saraf, Arkabandhu Chowdhury, Ashley Gabriel, Ashwin Bharambe, Assaf Eisenman, Azadeh Yazdan, Beau James, Ben Maurer, Benjamin Leonhardi, Bernie Huang, Beth Loyd, Beto De Paola, Bhargavi Paranjape, Bing Liu, Bo Wu, Boyu Ni. Braden Hancock, Bram Wasti, Brandon Spence, Brani Stojkovic, Brian Gamido, Britt Montalvo, Carl Parker, Carly Burton, Catalina Mejia, Ce Liu, Changhan Wang. Changkyu Kim, Chao Zhou, Chester Hu, Biiig- Chu. Chris Cai, Chris Tindal, Christoph Feichtenhofer, Cynthia Gao, Damon Civin, Dana Beaty, Daniel Kreymer, Daniel Li, David Adkins, David Xu, Davide Testuggine, Delia David, Devi Parikh, Diana Liskovich, Didem Foss, Dingkang Wang, Duc Le. Dustin Holland, Edward Dowling, Eissa Jamil, Elaine Montgomery, Eleonora Presani, Emily Hahn, Emily Wood, Eric-Tuan Le. Erik Brinkman, Esteban Arcaute, Evan Dunbar, Evan Smothers, Fei Sun, Felix Kreuk, Feng Tian, Filippos Kokkinos, Firat Ozgenel, Francesco Caggioni. Frank Kanayet, Frank Seide, Gabriela Medina Florez, Gabriella Schwarz, Gada Badeer, Georgia Swee, Gil Halpem, Grant Herman, Grigory Sizov, Guangyi, Zhang, Guna Lakshminarayanan. Hakan Inan, Hamid Shojanazeri, Han Zou, Hannah Wang, Hanwen Zha, Haroun Habeeb, Harrison Rudolph, Helen Suk, Henry Aspegren, Hunter Goldman, Hongyuan Zhan, Ibrahim Damlaj, Igor Molybog, Igor Tufanov, Ilias Leontiadis, Irina-Elena Veliche, Itai Gat, Jake Weissman, James Geboski, James Kohli, Janice Lam, Japhet Asher, Jean -ean-Baptistte Gaya, Jeff Marcus, Jeff Tang, Jennifer Chan, Jenny Zhen, Jeremy Reizenstein, Jeremy Teboul, Jessica Zhong, Jian Jin, Jingyi Yang, Joe Cummings, Jon Carvill, Jon Shepard, Jonathan McPhie, Jonathan Torres, Josh Ginsburg, Junjie Wang, Kai Wu, Kam Hou U. Karan Saxena, Kartikay Khandelwal, Katayoun Zand, Kathy Matosich, Kaushik Veeraraghavan, Kelly Michelena, Keqian Li, Kiran Jagadeesh, Kun Huang. Kunal Chawla, Kyle Huang, Lailin Chen, Lakshya Garg, Lavender A, Leandro Silva, Lee Bell, Lei Zhang, Liangpeng Guo, Licheng Yu, Liron Moshkovich, Luca Wehrstedt, Madian Khabsa. Manay Avalani, Manish Bhatt, Martynas Mankus, Matan Hasson, Matthew Lennie, Matthias Reso, Maxim Groshev, Maxim Naumov, Maya Lathi, Meghan Keneally, Miao Liu, Michael L. Seltzer, Michal Valko, Michelle Restrepo, Mihir Patel, Mik Vyatskov, Mikayel Samvelyan, Mike Clark, Mike Macey, Mike Wang, Miquel Jubert Hermoso, Mo Metanat, Mohammad Rastegari, Munish Bansal, Nandhini Santhanam, Natascha Parks, Natasha White, Navyata Bawa, Nayan Singhal, Nick Egebo, Nicolas Usunier, Nikhil Mehta, Nikolay Pavlovich Laptev, Ning Dong. Norman Cheng, Oleg Chernoguz, Olivia Hart, Omkar Salpekar, Ozlem Kalinli. Parkin Kent, Parth Parekh, Paul Saab, Pavan Balaji. Pedro Rittner, Philip Bontrager, Pierre Roux, Piotr Dollar, Polina Zvyagina, Prashant Ratanchandani. Pritish Yuvraj, Qian Liang, Rachad Alao, Rachel Rodriguez, Raf Ayub, Raghotham Murthy, Raghu Nayani, Rahul Mitra, Rangaprabhu Parthasarathy, Raymond Li, Rebekkah Hogan, Robin Battey, Rocky Wang, Russ Howes, Ruty Rinott, Sachin Mehta, Sachin Siby, Sai Jayesh Bondu, Samyak Datta, Sara Chugh, Sara Hunt, Sargun Dhillon, Sasha Sidorov. Satadru Pan, Saurabh Mahajan, Saurabh Verma, Seiji Yamamoto, Sharadh Ramaswamy, Shaun Lindsay, Shaun Lindsay, Sheng Feng, Shenghao Lin, Shengxin Cindy Zha, Shishir Patil, Shiva Shankar, Shuqiang Zhang, Shuqiang Zhang, Sinong Wang, Sneha Agarwal, Soji Sajuyigbe, Soumith Chintala, Stephanie Max, Stephen Chen, Steve Kehoe, Steve Satterfield, Sudarshan Govindaprasad, Sumit Gupta, Summer Deng, Sungmin Cho, Sunny Virk, Suraj Subramanian, Sy Choudhury. Sydney Goldman, Tal Remez, Tamar Glaser, Tamara Best, Thilo Koehler, Thomas Robinson, Tianhe Li, Tianjun Zhang, Tim Matthews, Timothy Chou, Tzook Shaked, Varun Vontimitta, Victoria Ajayi, Victoria Montanez, Vijai Mohan, Vinay Satish Kumar, Vishal Mangla, Vlad Ionescu, Vlad Poenaru, Vlad Tiberiu Mihailescu, Vladimir Ivanov, Wei Li, Wenchen Wang, Wenwen Jiang, Wes Bouaziz,
12

---

Published as a conference paper at ICLR 2026
Will Constable, Xiaocheng Tang, Xiaojian Wu, Xiaolan Wang, Xilun Wu, Xinbo Gao, Yaniv Kleinman, Yanjun Chen, Ye Hu, Ye Jia, Ye Qi, Yenda Li, Yilin Zhang, Ying Zhang, Yossi Adi, Youngjin Nam, Yu. Wang, Yu Zhao, Yuchen Hao, Yundi Qian, Yunlu Li, Yuzi He, Zach Rait, Zachary Vito Zef Rosnbrick, Zhaoduo Wen, Zhenyu Yang, Zhiwei Zhao, and Zhiyu Ma. The Ilama 3 herd of models, 2024.
21783
Daya Guo, Dejian Yang. Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang. Xiao Bi, et al. Deepseek-rl: Incentiviziing reasoning capability in llms via reinforcement learning, 2025. URL http:p::/aaari...orggaaas///////
12948
Jonathan Ho, Ajay Jain, and Pieter Abheel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems (NeurIPS), volume 33, pp. 6840-6851, 2020.
Emiel Hoogeboom, Didrik Nielsen, Priyank Jaini, Patrick Forré, and Max Welling. Argmax flows and multinomial diffusion: Learning categorical distributions. In Advances Neural Information Processing Systems (NeurIPS), volume 34, pp. 12454-12465, 2021.
Nai-Chieh Huang, Ping-Chun Hsieh, Kuo-Hao Ho, and I-Chen Wu. Ppo-clip attains global optimality: Towards deeper understandings of clipping. In AAAI Conference Artificial Intelligence (AAAI), volume 38, pp. 12600-12607, 2024.
Zemin Huang, Zhiyang Chen, Zijun Wang, Tiancheng Li, and Guo-Jun Qi. Reinforcing the diffusion chain of lateral thought with diffusion language models. In Advances in Neural Information Processing Systems (NeurIPS), 2025.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmax, 2017. URL httth:t/aaxxiiiii org/ abw. 01144
Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion- based generative models. In Advances in Neural Information Processing Systems (NeurIPS), volume 35, pp. 26565-26577, 2022.
Jaeyeon Kim, Kulin Shah, Vasilis Kontonis, Sham Kakade, and Sitan Chen. Train for the worst, plan for the best: Understanding token ordering in masked diffusions. In International Conference on Machine Learning (ICML), 2025.
Florent Krzakala and Lenka Zdeborová. Hiding quiet solutions in random constraint satisfaction problems. Physical review letters, 102(23))238701, 2009.
Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur Parikh, Chris Alberti, Danielle Epstein, Illia Polosukhin, Jacob Devlin, Kenton Lee, Kristina Toutanova, Llion Jones, Matthew Kelcey, Ming- Kiig-Wei Chang, Andrew M. Dai, Jakob Uszkoreit, Quoc Le, and Slay Petrov. Natural questions: A benchmark for question answering research. Transactions of the Association for Computational Linguistics, 7:452-466, 2019.
Xuanlin Li, Brandon Trabucco, Dong Huk Park, Michael Luo, Sheng Shen, Trevor Darrell, and Yang Gao. Discovering non-monotonic autoregressive orderings with variational inference, 2021. URL
Hunter Lightman, Vineet Kosaraju, Yura Burda, Harri Edwards, Bowen Baker, Teddy Lee, Jan
Leike, John Schulman, Ilya Sutskever, and Karl Cobbe. Let's verify step by step, 2023. URL
Aaron Lou, Chenlin Meng, and Stefano Ermon. Discrete diffusion modeling by estimating the ratios of the data distribution. In International Conference on Machine Learning (ICML), 2024.
Youssef Mroueh. Reinforcement learning with verifiable rewards: Grpo's effective loss, dynamics, and success amplification, 2025. URL / 06639
Shen Nie, Fengqi Zhu, Zebin You, Xiaolu Zhang, Jingyang Ou, Jun Hu, Jun Zhou, Yankai Lin, Ji-Rong Wen, and Chongxuan Li. Large language diffusion models. Advances in Neural Information Processing Systems (NeurlPS), 2025.
13

---

Published as a conference paper at ICLR 2026
Jingyang Ou, Shen Nie, Kaiwen Xue, Fengqi Zhu, Jiacheng Sun, Zhenguo Li, and Chongxuan Li. Your absorbing discrete diffusion secretly models the conditional distributions of clean data. In International Conference Learning Representations (ICLR), 2025.
Long Ouyang, Jeff Wu, Xu Jiang, Diogo Almeida, Carroll L. Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, John Schulman, Jacob Hilton, Fraser Kelton, Luke Miller, Maddie Simens, Amanda Askell, Peter Welinder, Paul Christiano, Jan Leike, and Ryan Lowe. Training language models to follow instructions with human feedback, 2022. URL
htttp:://rxiv..o wwg. 2203. 02155
Rafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D Manning, Stefano Ermon, and Chelsea Finn. Direct preference optimization: Your language model is secretly a reward model. In Advances in Neural Information Processing Systems (NeurlPS), volume 36, pp. 23728-53741 2023.
Subham Sahoo, Marianne Arriola, Yair Schiff, Aaron Gokaslan, Edgar Marroquin, Justin Chiu, Alexander Rush, and Volodymyr Kuleshov. Simple and effective masked diffusion language models. In Advances in Neural Information Processing Systems (NeurIPS), volume 37, pp. 130136-130184 2024.
Kulin Shah, Nishanth Dikkala, Xin Wang. and Rina Panigrahy. Causal language modeling can elicit www.iiiiorrr.....o /bs/
search and reasoning capabilities on logic puzzles, 2024. URL
2409. 10502
Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Xiao Bi, Haowei Zhang. Mingchuan Zhang, Y. K. Li, Y. Wu, and Daya Guo. Deepseekmath: Pushing the limits of mathematical reasoning in open language models, 2024.
22C2.03300
Jiaxin Shi, Kehang Han, Zhe Wang, Arnaud Doucet, and Michalis Titsias. Simplified and generalized masked diffusion for discrete data. In Advances in Neural Information Processing Systems (NeurIPS), volume 37, pp. 103131-103161 2024.
Andy Shih, Dorsa Sadigh, and Stefano Ermon. Training and inference on any-order autoregressive models the right way. In Advances in Neural Information Processing Systems (NeurlPS), volume 35, pp. 2062-22275 2022.
Yang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference Learning Representations (ICLR), 2021.
Yuxuan Song, Zheng Zhang, Cheng Luo, Pengyang Gao, Fan Xia, Hao Luo, Zheng Li, Yuehang Yang, Hongli Yu, Xingwei Qu, Yuwei Fu, Jing Su, Ge Zhang, Wenhao Huang, Mingxuan Wang, Lin Yan, Xiaoying Jia, Jingjing Liu, Wei-Ying Ma, Ya-Qin Zhang, Yonghui Wu, and Hao Zhou. Seed diffusion: A large- scale diffusion language model with high-speed inference, 2025. URL
https::/arxixs org//rgg//aooo..oooo abs 2508. 02193
Richard S Sutton, David McAllester, Satinder Singh, and Yishay Mansour. Policy gradient methods for reinforcement learning with function approximation. In Advances in Neural Information Processing Systems (NeurIPS), volume 12, 1999.
Ronald J Williams. Simple statistical gradient- algorithms for connectionist reinforcement learning. Machine learning, 8(3):229-255) 256, 1992.
An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chengen Huang, Chenxu Lv, Chujie Zheng, Dayiheng Liu, Fan Zhou, Fei Huang, Feng Hu, Hao Ge, Haoran Wei, Huan Lin, Jialong Tang, Jian Yang. Jianhong Tu, Jianwei Zhang, Jianxin Yang. Jiaxi Yang. Jing Zhou, Jingren Zhou, Junyang Lin, Kai Dang, Keqin Bao, Kexin Yang, Le Yu, Lianghao Deng, Mei Li, Mingfeng Xue, Mingze Li, Pei Zhang, Peng Wang, Qin Zhu, Rui Men, Ruize Gao, Shixuan Liu, Shuang Luo, Tianhao Li, Tianyi Tang, Wenbiao Yin, Xingzhang Ren, Xinyu Wang, Xinyu Zhang, Xuancheng Ren, Yang Fan, Yang Su, Yichang Zhang, Yinger Zhang. Yu Wan, Yuqiong Liu, Zekun Wang. Zeyu Cui, Zhenru Zhang, Zhipeng Zhou, and Zihan Qiu. Qwen3 technical report, 2025. URL https: ///axx/avii org abs / 2505 09388
14

---

Published as a conference paper at ICLR 2026
Takayuki Yato and Takahiro Seta. Complexity and completeness of finding another solution and its application to puzzles. IEICE transactions on fundamentals of electronics, communications and computer sciences, 86(5):1052-1060, 2003.
Jiacheng Ye, Zhihui Xie, Lin Zheng, Jiahui Gao, Zirui Wu, Xin Jiang, Zhenguo Li, and Lingpeng Kong. Dream 7b: Diffusion large language models, 2025. URL https::/aaxii..om abs/ 2508. 15487
Oussama Zekri and Nicolas Boullé. Fine-tuning discrete diffusion models with policy gradient methods. In Advances in Neural Information Processing Systems (NeurIPS), 2025.
Siyan Zhao, Devaansh Gupta, Qinqing Zheng, and Aditya Grover. d1: Scaling reasoning in diffusion large language models via reinforcement learning. In Advances in Neural Information Processing Systems (NeurIPS), 2025.
Kaiwen Zheng, Yongxin Chen, Hanzi Mao, Ming- Yu Liu, Jun Zhu, and Qinsheng Zhang. Masked diffusion models are secretly time-agnostic masked models and exploit inaccurate categorical sampling. In International Conference on Learning Representations (ICLR), 2025.
Fengqi Zhu, Rongzhen Wang, Shen Nie, Xiaolu Zhang, Chunwei Wu, Jun Hu, Jun Zhou, Jianfei Chen, Yankai Lin, Ji- Ji-Rong Wen, and Chongxuan Li. Llada 1.5: Variance-reduced preference optimization for large language diffusion models, 2025. URL httt::////ax...oorrbb////...... 19223
15

---

Published as a conference paper at ICLR 2026
CONTENTS
Introduction
2
Preliminary
3
Improving Unmasking Policy Beyond a Reference Policy
3.1
Existence of optimal path Beyond gconf
3.2
Reformulating the problem as reinforcement learning
3.3
Theoretical analysis
3.4
Practical realization of unmasking policy optimization
4
Experimental results
4.1
Experimental Setting
Training dynamics across reference policies
Main result
4.4
Ablation Study
5
Conclusion
10
Related works
17
Training recipe detail
18
B. Full Derivations for LUPO Realizations
18
B.2 Algorithms
20
Detailed Proof
22
proof of KL tightening for MDM policy
22
................................
Surrogate loss of output-level reward maximization
Surrogate loss of output-level KL divergence
28
D Additional experiments
30
D.I unmasking policy optimization binary reward setting
30
D.2 Toward scalable and generalizable unmasking policies
31
D.3 Pass result of unmasking policy model
32
Experimental Details
32
E.
Hyper-parametert and other details
32
reward and prompt template for each benchmark
32
SUDOKU unmasking examples
34
Usage of Large Language Model
34
16

---

Published as a conference paper at ICLR 2026
Truinable 1-Layer Transformer 4
Feat 1
Fear I
Frozen MDM
Haal
[Mask]
Frozen MDM Transf.
Fent 2
Top-K
Feat 2
Trainable
0.8
3-Layer
[Makk]
Feat 3
Layer
MLP
With
Seail
Probs
Fear 3
MLP &MnF
0.2
Dog
Feat 4
Feat 4
Figure 5: The structure of our unmasking policy model. The model is composed of single Transformer layer and a 3-layer MLP. In inference time, the model proceeds as follows: 1) Given an input sentence, we run the base MDMs transformer and extract feature. 2) This feature branches into two paths: first, the base MDMs continue to produce a prediction of token distribution, and second, fed to the policy model's transformer layer. 3) We concatenate the extracted feature with the base MDM's 's Top-K K probabilities. 4) The concatenated feature is then processed by 3-layer MLP to yield a policy over unmasking positions.
A RELATED WORKS
Discrete Diffusion Models. Diffusion probabilistic models have become the state-of-the-arr approach for continuous data domains such as images, audio, and video, by iteratively denoising Gaussian- corrupted inputs (Ho et al., 20200 Song et al. 2021). Extending this success discrete data (e.g. text or sequences) has led to the development of discrete diffusion models built on discrete-state Markov chains rather than continuous noise processes (Austin et al. 202T [Hogebooo al. 202T). Lou itt al. (2024) proposed SEDD, which utilizes theoretical score entropy objective. Mcanwhile, masked diffusion models (MDMs), which utilize an absorbing mask strategy, have emerged as leading class of discrete diffusion models for text generation (Lou et al. [2224] Sahoo et al., 2024 Shi et al. 2024). By using an absorbing masking state, MDMs simplify the corruption process: once token masked, it remains masked until generation, analogous absorbing diffusion in continuous models (Shi et al. 2024). This simplification yields clean and principled training recipe aligned with continuous-time diffusion theory, including evidence lower bound (ELBO) objective (Shi et al. 2024). Beyond continuous- time discrete diffusion, Zheng et al. 20025;; ;
Ou et al. (0025) provides a time-agnostic training framework and simple cross- coos-entropyy loss. These advances have significantly closed the performance gap between MDMs and traditional autoregressive models on language tasks (Ye et al.
2025
Nie et al. 2025) when the scale of the model is similar.
Unmasking Strategy in Masked Diffusion Inference. crucial factor for MDM performance is the unmasking strategy used during inference-i.e., how the model decides which masked tokens to reveal at each generation step. While MDMs have shown superior performance on text generations, earlier works (Ou al. 2025 Shih al. 2022, Li al. 2021) have empirically observed that a random unmasking strategy is suboptimal. Moreover, Kim et al.] (2025) proves that no polynomial-tne -ime algorithm can solve any-order generation; iee., we cannot train MDMs that exactly recovers the real data distribution for all masked sentences. However, they empirically show that selecting unmasking tokens via max-confidence (Chang al.] (202) or max- max-margin (Kim al., 2025) at inference time can bypass sub-hard instances and even outperform ARMs. As result, contemporary large-scale MDMs (Nie et al. 2025 Ye et al, 2025) commonly adopt max-confidence scheduling at inference.
Remark. Detailed Explanation of Hardness of MDM Training. Kim et al. (2025) show that the
conditional prediction tasks induced by Any-order masked diffusion models fall into the planted constrained satisfaction problem (CSP) hard phase (Krzakala & Zdeborová 2009), implying that no polynomial- time algorithm can achieve optimal performance in general. Consequently. MDM training loss, Eq. (I) (the reverse-process KL objective driven by Sahoo et al. (2024 Zheng et al.
(2023)) cannot be driven to zero, which is consistent with their empirical finding that loss remains high text and L&O-NAF-SAT datasets.
17

---

Published as a conference paper at ICLR 2026
B TRAINING RECIPE IN DETAIL
In this section, we explain our memory-efficient training recipe in detail. We provide full derivations of various reference policy realizations (Appendix [B.]]. and algorithms for training and generating (Appendix B.2). For deeper understanding of our memory-efficient training algorithm and model architecture, we first provide notations and model inference in detail.
Policy model inference process. Our rationale behind policy model structure is two: 1) utilizing the MDM's knowledge (extracted feature), which allows for avoiding nom-scratch training and allow memory-efficient training, and 2) utilizing the MDM's token prediction distribution, which can help the model mimic the behavior of feonf. The main paper outlined the policy model's 's inference at high level; here, we provide notation and description of how, given masked sentence Zns the model generates an unmasking trajectory. For reference, we reproduce the model architecture from the main paper (Figure 2 as Figure 5
To recap, the unmasking policy model is parameterized softmax over masked positions using a learnable scorer h (a:)
RR
(a x.)
(20)
exp [1e (2x))a)))
where a' i-th mask of In. As illustrated in Figure the policy model decomposes as (ØMLP, (TT)) where ΦrF is a 1-layer Transformer and @MLP is 3-layer MIP. We similarly write the MDM as (@Tr, MML)) with a feature extractor transformer @Tr and a token prediction head 0MLP. Given the current masked sequence Hn and selected mask a', the MDM computes the token distribution as follows:
H = (TF(2x)
IR lxd
U = [a'] E R"
Z, (u,) R
1n-1
En softmax(za:)e (cEV).
where is the output dimension of the transformer feature. Let Pa softmax(z.:) the token posterior at index a' with dimension [V]. The policy encoder first refines the sequence features,
H ΦTF (H.. E Lxd
U = H,,aa E Rd.
Let (P.)( E R denote the vector of the K largest probabilities of Pat sorted in descending order. We form the concatenated feature
Sa [ūa TopK(pa:)] Ba+K
(21)
and map it to a scalar score via the 3-layer MLP:
[hs(ar))a dDML R.
Finally, the unmasking policy over the masked index set the softmax of these scores (Eq. (20)
B. FULL DERIVATIONS FOR CUPO REALIZATIONS
General formulations. To recap, final objective is given follows:
L
1
Io (a ((())
LUPO(go, Gref, D) = E
min
Ag
qr-po (m) ( Ail Papoola
Joara (a x (0)
Io (2) (0)
clip(
1 - c. 1.1 c)A,)) - D(9 (P
Jtaa (a (xo
(g) (g)
(22)
The divergence term in Eq. 22 functional that, given a prompt q. a full trajectory (Zo.L, (,,)), and models (gs- Gvet: Te), returns scalar estimate. Formally,
D Dos Q Nn.L. ALt -> RR
18

---

Published as a conference paper at ICLR 2026
and its empirical expectation serves tractable estimator of the output-level KL divergence
Dxu(Pae (ao q) Po..t (To q)). Following Proposition 2 and Proposition 3 we can decrease the
output-level KL divergence with LKL. defined follows:
CKI aL.L q) StopGrad(
Prona (Zta.. [q] (1 log
Pus (#F.. (q)
Pera (2e0.
log go (an (.)):
Porer (ata (q)
(23)
where gø: Prot and me is omitted in LKT. for brevity. We convert Pa (30.1.19) into I (an/(a))
TT( 1 an q) in Eq. (23):
( (xn)
g (n (20n) log gref (an (2(n))
(.L() StopGrad
II n=1 Iebna (an m
(1
Σ
Σ log go (an wn).
(24)
We already computed Ioaa (-) while sampling, and ( while computing reward-maximization objective term in CUPO, and Gref does not require network evaluations. Therefore, we can measure
fKL without no additional computation. Furthermore, since Σn= log go(a, 2r,) is summation, we can perform gradient update token-by- by-t-e,-token, which allows tractable and memory training.
Softmax realization.
For softmax realization of gref, we consider Jeanfr given as follows:
Evev exp(eeeep = c(2n//T)
Icont- (a (.n)) (a/(
Eni ExEv exp(mo Exeyeep(o = c/x))/77*
s.t. lim Icont" T-00+
Jconf-
(25)
Here, since Jconf" assigns non-zero probability to every unmasking index, we can estimate the terminal KL divergence with sampled trajectory. We utilize LKL. (Eq. (24) by replacing gvet with Scont'
Top-K realization For Top-K realization of Iref, we consider Scont given as follows:
@Top-K(a|(xa) = K 1(a" e argtopk(max To (o = c(x,))), s.t. 9Top-1 = Jcont. (26)
eEV
where we reparametrize GoTop defined as follows:
exp((h((an))
Poroo x (a (xn))
(27)
Eas Eargtopk max ro(-) exp(hhg(2zn)aa
which is the softmax distribution of h& on the Top-K K set of unmasking positions. This reparametriza- tion offers two advantages: 1) when & is randomly initialized, we can expect Jēt = -Top-K K. and 2) for sampled trajectory, 9Top-K (a, is non-zero for all n E [1, L] since we sample from Gotop- K Therefore, we can train Jorop-K from scratch with LKL. by replacing Jret = 9Top-K.
Max-confidence realization. For max-confidence realization of Jref, we consider Jconf given as follows:
Jeont (a" (x.)) 1(aa arg max ( c(xx))),
(28)
cEV
where Goont is deterministic. The terminal KL divergence DRI(Pos (zo I q) Pacat (xn q)) can be defined where Te stochastic, but we cannot empirically estimate the terminal KL divergence with sampled trajectory since there exists trajectory sampled from Pys is 0 for Pgcons Therefore, we just utilize cross-entropy given as follows:
CF CT.gs (DO.L. 9)
Σ
Σ
Icont (a) log (o(a)o..)
(29)
a=1 aE
Although not theoretical, we can expect to regularize da being far away from Gconts and we have empirically confirmed the performance gain (Section 4.4
19

---

Published as a conference paper at ICLR 2026
Algorithm I Generalized Sampling of MDMs
Require: Sequence length L, vocabulary V {l...., m), pretrained masked diffusion model rT. unmasking policy
1: TL. <-1
2: for 71 + L to 1 do
3: sample On g(a((xn)
unmasking token distribution. e.g., ( {a Hi=))
4: sample Zn To (2 (= (2n)
per-position categorical probabilities
5: In-1 {1 2 22 ... 2n n
n
6: end for
7: Output: æ0
B.2 ALGORIHTM
Sampling algorithm. We provide generalized sampling algorithm of large-scale MDMs in Algo- rithm The sampling trajectory differs by unmasking policy g. and we can perform various sampling by setting g as grands Seonf, Iconf", -Top--K K. or 9d-
Unmasking policy training algorithm. We provide memory-effiiient unmasking policy model training algorithm in Algorithm 2 We have previously mentioned how we utilize the extracted features and token distribution of MDM as imput to the unmasking policy model. We also have provided how the gradient of LUPO can measured by the token-wise gradient of (eo(-))). Summing up all techniques, we can perform training without using GPU memory more than the amount of original MDM occupies. Indeed, we have confirmed that training can be performed with only one A100 40GB GPU.
Specifically, the algorithm is divided into two phases: the sampling phase and the gradient update phase. At the sampling phase, sample various sequences from Pge.n (-1q). In this phase, we gather the extracted feature of the MDM transformer, the token distribution of MDM, the mask token index, and record the probability training where all the gradients of g (-( in LUPO independent of others. Furthermore, since we have gathered all the information mg that is needed to update the gradient in the sampling phase, free the memory of MDM 8 The only parameters requiring forward and backward process in the gradient update phase is the unmasking policy model, which is much smaller than MDM (e.g., LLaDA: 8B, corresponding policy model: 134M). Therefore, training. For more details, refer to the algorithm table.
the trajectory. At the gradient update phase, we can perform mini-batch
perform memory-efficient
20

---

Published as a conference paper at ICLR 2026
Algorithm 2 Memory-efficiene training of the unmasking policy model
Require:
Pretrained masked diffusion model composed of feature extracting transformer @TF and token estimator MIP.. unmasking policy model Is composed of 1-layer transformer ØTF and 3-Layer MI.P @MLP, divergence loss D, reference policy gref, prompt distribution 2, number of inner updates 1
while not converge do Sampling phase
2:
with barch.no.grad(.
3:
Idnia + 96
4:
Sample prompt q Q
5:
Sample G trajectories f 91t g=1 Paoni (+q)), g [G]
(g) G
> while sampling, save H = rIF (x (g)
) A lils TO TT(1|ar) (2)), Iowi (an (2%
6:
> if D LKI. record Precia (x)) a...
(0) |q), Pro
(x (0) a (q) while sampling
7:
8:
For each x compute reward rg and advantage Ag end with
9:
Gradient update phase
10:
for gradient update iterations n = L...., 1 do
with torch. no grad ():
12:
For each 0:1. pre-compute Pga. (2b)
(g)
13:
if D LKL then
Pap (wo.1.
Pag (moLL) Peret (mo.2)
14:
Kig +
Pressa (mo.1) (1 + log
15:
end if end with Make mini-batch (B)_ such that Bi U....U...UU = {1,.... L} and B; n B,
17:
for all i i }
18:
for B E (B.) do
For each n E and g E {1,.., G}. A + ΦIF (
19:
20:
For each B and g {1, G}, construct
(g)
using Eq. (1
21:
x
t ØMLP (s!
Ge (a ((("))
so (af (جاt)
22:
L +- [B|| EnEB.gEG min(
Ioald (a! (11))) Ag. clip(
goota (a!. (n - e, + e)4,)
23:
if D = CE then
L. Ls B EgEG Enen CE(geont (-121 )lga( (()))
end if
26:
if D)
LKI. then
L. L. - 3 E.EE Kig EnEB log gs (g) x()))
27:
28:
end if
29:
Update by performing gradient ascent on CUPO
30:
end for
31:
end for
32:
end while
33:
Output:
21

---

Published as a conference paper at ICLR 2026
C
DETAILED PROOF
In this section, we provide full proofs omitted in the main paper.
C.I PROOF OF KL TIGHTENING FOR MDM POLICY
Theorem 2 (Reference-Kee KI. Tightening for MDM Policy Improvement). Assume there exists an ideal MDM 6* satisfying MDM L for Paaa where data is composed of prompt and oracle answer {q. a}. Assume we select well-defined reward r(q, a) E {0, 1} such that Páaa qq a) > +
r(q.a) - and 0 < gnd < 1. Now consider incomplete MDM 0 such that (MDM(() 0, and 9ø"
that satisfies Eq. (13) for 0. Let fref be policy model such that, for every q Paue (ro ② > 0
whenever po- (mo
q 0, i.e. supp(po-(( q)) E supp(Ps.ac
Q)) Then, the following
inequality holds:
DKL(Po- (zo
Pg.. (z)l)))
(14)
Or equivalently,
(allPaaaaala)llpe.. (oolg)) (Paataaa)))) (x: q))-
(15)
Proof.
We first recall the GRPO policy training dynamics learning Eq. 10 proven by Mroueh
(2025):
Lemma 1 (GRPO Policy Dynamic (Mroueh 2025 Optimal GRPO iterations policies solving Eg. 10 satisfy the following recursion, for n
1
Pgen (Fo
q)
Porat (Zo ())
(30)
Zn-1(q)
(q)
w ( (Tgan_, (q)) 2()) w. (r(rr..., (0on_,(a))) (q)) (1 r(q, rqq (31)
exp
where
Zn-1(q)= = Torer (q) exp ( w wr (a)) + (1 - rg rg)o(())) (q)) exp( 1 w (roe. (a))
(32)
and w.t (r) and w (r) are defined as follows:
1-r
w (r)
r(1 r(1-r)+é w (r)
(33)
r(1 (1- +
Under Theorem the fixed point, substitute the converged policy into the recursion
and compare the two KL terms. The next steps simply expand the KL difference. Then, for the convergence point where limp Tgo. rg
22

---

Published as a conference paper at ICLR 2026
DKL (po- (2o1g) |l Pa (zolg)) DKL (po (201g) || (2o)
(34)
pp- log po (zo (q) log P gron ( q)] - Ep-- [log po- (zulg) log Pa.- (zol))
(35)
Epg-
[log (zu q) log Pernt (2)
(36)
Epp- 1 (T (q)) r(q, x0) we (r (q)) (1 r(q. r(q. x0)) logg ZΦ-
(37)
1 (T (q)) - -log Z
Vap E N, pa- (xo q) > + (q, =1
(38)
3
1
w (r (q))
3
log (g) exp( (r (q)) (1 rge (q)) exp( ". (Tas (q))
(39)
3 w (res. (r (q))
log (oror(q) exp (3 wr g.- (q)) (-rr)(()))) exp(3 ( (q))
(40)
-go. (q) Tg (q) > 0, s.t. ue (q)) 2 0, w. (r (q)
=0,
(41)
where supp(po- (- q)) C supp(Pger q)). Therefore, the inequality in Eq. 14 holds. Accord- ingly, inequality in Eq. (15) holds where pa- = Pdata for ideal MDM 6".
Remark. Why does the support assumption supp(pe- (. q)) S supp(Pa.c (- 4)) hold?
Even Pret deterministic (e.g., max-confidence), the support assumption naturally holds for masked diffusion models. The reason is that the MDM denoiser To xx) assigns strictly positive probability to every token in the vocabulary at each unmasking step. Therefore, for any terminal sequence Xy V", there always exists at least one valid denoising trajectory under Jref that reaches Xo. implying Po.. (xo ②) 0. Formally, the terminal-outper likelihood marginalizes over all
trajectories:
L II
Paie [o ())
Σ
gret(an Xx; g) (x Xn. Ony Q)-
XI: 41.L ==1
Because Ty has full support, the product above is strictly positive for at least one trajectory. To illustrate, take any target sequence xi- Starting from the fully masked state, Jeet selects index at. al: since ( _n +,0a xx)= > 0, the next state x satisfying # In 4,42 is reachable. Repeating this inductively shows that xō reachable regardless of gret. Thus the support assumption holds.
C.2 SURROGATE LOSS OF OUTPUT-LEVEL REWARD MAXIMIZATION
Our goal here to show that the output- level reward maximization term in Eq. (10) can be converted into a tractable token-wie wise reward maximization term. The conclusion of Proposition [ with conventional RI. that have developed with rigorous proofs, yet in our setting, where training the MDM unmasking policy, the assumptions are slightly different. More specifically, finite-horizon MDP with fixed length (text length is determined), sparse reward (only a verifiable reward for the terminal state exists), and non-stationary (for different time steps, states do not overlap). To fill this gap, we rigorously show that Proposition holds.
MDP is
We cannot even directly apply the policy gradient theorem (Sutton et al. 1999). which is essential in the proof. Therefore, we start with justifying how the policy gradient theorem also works in our framework, which can be easily proven:
23

---

Published as a conference paper at ICLR 2026
Lemma 2 (Policy gradient theorem in sparse-reward, finite-horizon, episodic, and non-stationary MDP), Fix a horizon L. For each E {0, (0......,4 L} let the state space be Xn, and assume the spaces pairwise disjoint: X, n 0 for # At step i {1...., L. from state In Xn discrete action on E An is drawn from a stochastic policy go go(0n Tn., and the environment transitions to the next layer's state 1n-1 - according to dynamics æ (z nn On)- Assume the start state is fixed as Lt The trajectory is (Zo:L: 01:L) with joint density
Pso (T xL)
II go (n.) ( n=1
In, an))
Rewards are terminal-only: r(za) at 0 and otherwise. The objective the expected terminal reward
J( Er-pas )r(
Define the value function and state-action value function under g4 by
Vg (x,)) Eir(x)) xx,),
Qgo (2n, a,) ((((())) Dns dn]
Under the above assumptions, the policy gradient depends only the policy gø:
VoJ (ZL)
Σ
Pgo (x)
E Qo (x, a) Vogo(((z)
(42)
Proof. By definition,
(( = Ex ((z)))((2o))
Using the log-derivative trick,
V (o) = Σ ((00) V VoPoo (T x1))
r(00)) Pas (T x1) V log log wt..
Since only the policy Is depends o, have
logPs. (T xL) = Σ
log (an xx) + const,
and therefore
L
V log (x)) =
V log gol In).
Thus,
V Erop,, (1xe) r(z0) V (an xx.)
Taking conditional expectations with respect to (In, an) yields for each E {l...., L} we have
E, ((vv)) V log (an x.))
(43)
= E
Br(x)) æn, an] Ve log (a atm)
(tower property)
Ee.. on
Qgo (I an) V log (an amn)
(44)
Σ Σ
Poo (In, an) Qg (I an) V log (a Φn)
(45)
an E. an anEAn
Σ
Pgo (In)
Σ
go(an xxn) Qa (In, an)
V log Loggga(a x.))
2n E.Y
Oa C.A.
(p(#n, an) (((xn)g(a. (x.))
= Σ
Pgo (Inn)
Σ
Qoo (In, an) Vogo(an xx),
V = Vg)
2nn Exi
Or Edd
24

---

Published as a conference paper at ICLR 2026
where Pys (I.)) Pr, Prompoo (X,, 2n) and Ps. (En, Pgo (2 go 2n,) are the marginals induced by the rollout distribution. Therefore,
[Ø) BT- (1=1) r(zo)
log galgo(an Inn)
Σ
Poo
pas (2r, a) Vogo(ala))
which ends the proof.
This can also be proved another way, just iteratively deriving the V aa Using Lemma we can prove main objective follows:
Proposition (Output-Token) Level Gradient Alignment (informal)). For MDM ag and unmasking policy model 9s where Pg. (In- σn) (a, Fn) <o( Fn. an) and Py. (wo (g)
Exi.... xy Pe L q). consider the output level loss Coutput and token-wise policy loss Ltoken'
Pgo (xn
q))
Loutput Bawin Po (14)
A(q, æo)],
(16)
Posaaa (2o0
q
ga(an En)
Leoken (Ø)
Egrup 16 (anz PPaao (14)
A(q, mo)]
(17)
Ioars
(an xnn)
Assume that optimize o carefullly such that & se Pold. the gradients approximately equal:
&o Re V CoCuutputs
(18)
and VoLtoken = V VoLoutput at first optimization step where & = Pold-
Proof. We derive V V Couputtt and Vo separately. Then, by applying Lemma to V Lontput. obtain form that matches V Vinkon First, expanding foutput shows that it equivalent to maximizing the advantage from the initial state:
Pgo ao
q))
Loutput (Ø) = Equ-p Exor (:19)
A(q, æo)]
(46)
Pooaaa (Do
Poo (xp
q)
-Egppol Σ
Procus (#o 9
A(q,
(47)
Pronaa (xp q)
woEx
= Eguupl
[ Σ
Pyo (wo (q){(q. 200]
(48)
#
Egrpp Σ
Σ
Pao (Xo.. q)A(q, #o)]
(49)
-Equpal
Σ
Pao (xo.1. q)A(q mo)]
(50)
Bp..e...... 201
EEqupa Eoo PPps (11g) [A(q. xo))
(51)
25

---

Published as a conference paper at ICLR 2026
Now derive V Coutput as follows:
V Vouuutu (Ø) V Earip Brsz Pas (1))) [A(q. mo)]
(52)
Barpa Σ
Pao (Fo. q)A((q
(53)
20.00.....0
(r(q. xo) Teaw (9))
& Earpal
Σ
Pao (oo.L q))
(54)
stdgeoua (q) +
r(q. 2b
V Equpa
Σ
Pgo (Tu:L
q))
std std Bolld (q) +
T Todald (q))
R Pawpal
Poo (zopz (
(55)
std (q) +
10..1....22
r(q. ma))
V Equpal
Σ
Pyo (ZoeL
qq)
std std goald (q) +
Φo Fin LL
Doold (q))
Eqrpp
Vol Σ
Peo Poo.L a)]
(56)
stdgead
(q) +
r(q, x0
V Equpa
Σ
Pyo (DuiL q))
stdgeau goold (q) +
(57)
std
20........3
Bmo.1 Puys (19) [r(q. xoo]
V Vo (q)
V Equno
Eqrip
(58)
std goold (q) +
std geas (q) +
where we define state-wise value function as the expectation of terminal reward starting from the state, i.e. Vas (I (q) == E Bxoa- Pap ((1q.=.))))))) [(q, æo)] and we denote value function at the initial state as Vg. (q) E Eegi (+1a) [(q.//qq)))) =. Accordingly, state wise action-valus function can be defined as Qgo (In, an 9)) = Ex... To (= an)) V. (In- Note that Eq. 57 is conceptually
equivalent to reversing the derivation process of REINFORCE with baseline (Willamss 1992). Since we want to show that V Vo Cookenn X VoCoutput. we breakdown V (q) using Lemma 2
V V. 9a (q)
Coooutput (Ø) Eqrpo
(59)
stigeats
(q) +
Boold
Σx Pos (x) 9 Σ Qno (x, a (//)))))))))) Vaga(a x))
Equpp
(60)
std (q) +
En Paoaad (x//g Σ Quean (æ,a Valo(a
~ Eqrop
(61)
std (q) +E
where comes from every state, i.e.. Xo U U... U Ur. Here, Eq. 59
Eq. 61 exactly
holds at the first optimization step where & = Dold. On the other side, we now show that Ltoken is
equivalent to Eq. (61) We start by converting the expectation into probability sum as follows:
ga(an (nn) A(q,zo) Joasa (an xnn)
Leokeu (0) Eqrapa Exo. Panasa (+1a)
Σ
(62)
n=
go(an Fn) Jori (an xnn)
Alq, A(g )
=Equup
Σ Po
II
no (2 Fni an) Ddom (an, Itn))
(63)
26

---

Published as a conference paper at ICLR 2026
For further explanation, want to recap Peeaia (#n-1 1 To In an) Idui (an [nn Note that Preaa > 0 only if a M. and =- # M since Pebiaa (an])r.) for ann # M,, T (n) for all æn-1 satisfying a M by definition. To simplify the derivation process, we consider Lroken, (Ø, q) defined as follows:
Erokom, (
Σ Propia (oo.r. ④))
Jouua (as 4i 20) mo)
(a; xi)
(64)
Prena ( (2)- 9)) Do(a; xia) Doald (a: x;))
Σ #0 a1
II ==1
r(q, xD)
(65)
Σ C0...e 3e G1
II J=i+1
Peena (25-111235, 9))
(Peoosa (xi-1/12; | ② [a xx;))
Paena (0)-1 9) rr(. mn)
(66)
Idhoa (ai x3a)
==1
Σ 20......0 dL
II J=
Peosia (15-11203, 9)
(Tp (0t xss az Dgo(aa ax:))
Pos (mj-11 4) r(q.. xo)
(67)
E d/+1,..,dt
II
Prosia (0.-11123,, ())
j=i+1
Σ
(To (Ta(xi-1/z0, as a)go(a; x.))
Pynau (0)-1//z,, 9)) r(q, xoo
(68)
Σ Mio.. ar.
Paoaaa (25-11223, 9))
j=i+1
ΣΣ
gola; | x.i) x0(xx-11200 4)
Σ
Pasus (z.-1|z33 9) .r(q. 2o)
01
Te 2 do
j=1
G 1....
(69)
Σ Σ
II
Preai (0)-21zz,.
Gil l.....z
j=i+1
Σ
(ol xxi)
Σ
Te Te( 12, a;, 9)
Σ
Panna (I 9)) .r(q. zo)
2i
2.....
==
(70)
Σ
Σ
II f=++1
Poona (
Li ditlnnag
To(2;;11 ai a) Volw (z: (q) ))
Σ gala xz.)
Σ
(71)
1
EmipProend (11) Σ (o(a,120)) Qanan (x), az19)
(72)
27

---

Published as a conference paper at ICLR 2026
where color parentheses by their nesting depth to improve readability of the multi-level factor- ization: outer, middle, inner, and deepest inner. The expansion proceeds by repeatedly pulling out terms that independent of the summation variables, and keeping only dependent factors inside the corresponding sums. Then,
Croken, std g hold (q) +
V Croken (0)
Bquupe
(73)
i=1
E Ban (1)) IC Vogo (aslz))) Qganu (2 agl)
(74)
std stdd (q) te
Ex Paecia (a)/g) Σ Qganes (a,, Vr.all)Vooslc (ogo(a)0,)))
=Equpe
(75)
stcggooi (q) te
Eq. 73 holds where goa (q) in A(q, mro = (r(g. mo) -V Yagola (q)) (c)+e) can be deleted as we have done VoLoutput. Since Eq. (61) and Eq. (75) are identical, we conclude V VoLtoken R V a Coutput and V Vaooke VoCoutput the first optimization step where o Pold-
C.3 SURROGATE LOSS OF OUTPUT-LEVEL KL DIVERGENCE
Our final goal is to provide evidence that by performing gradient ascent LKL, expect to decrease DKL. (Po. (æt (q) Pa.as (wo (q). We provide two propositions here. The first proposition shows that trajectory-wiss KL divergence is the upper bound of output-level KL divergence. The second proposition shows that the gradient of trajectory-wiss KL divergence is equivalent to LKI.
Emym.n m po( consider the terminal KL divergence DEL (Pao (Do q)) ll Panee (zolg)) and trajectory-wise KL divergence DKL. (Pou (0o.119))lpp... (F...l)) defined as follows:
Proposition 2. For MDMs where Pao (2olg)
Pas (2t0
DKL (Pgo (2 Paver (3))lq)) Permpo log
(76)
Pennr (2olg)
log Pso (#0:2
DKL (Pea (Zo.1 PP (Fo:L (q)
Em (m0) 14)
(77)
Pernt (Do.L (q)
Then, DEL (Pos (Lo.L q)) Pouer (20. (q) is upper bound of DKL (P (molgg |l Poucs (2olg)):
DKL. (Pas (wo q) Patar (wo 9)) SDKI. (Pas (wo.r. (//////pllPp (P (2ro.z
(78)
Proof.
Pao (wo
(q)
DKL. (pou (wo ()ll (aro ())= Exoo~ (coll))
log
(79)
Pgrer (20
qq)
log Pau (ato ())
Ev L (00))/a))
(80)
Pernt (Ie
q)
log Pao (wO.L
④))
Pao (#11)
20. ④)
= -Ezo MPoy (mT))la))
log
(81)
Pasut (Vo.L
q)
Pocet (x1:L
2o.
-KL (poo (xO.L I q). Porae (XO:L
q))
(82)
ExorPas (meala) [DLL (Poa (x1.L
Do 9)
lall (x1:L Zo, q))]
(83)
S DRI. (poo (wo.. a)llpgne (2o.r. ())
(84)
Therefore,
we
can
lower DKL (Pas (xo 9)) Parer (xo 9)) by
decreasing
DKL. (20. (#o.t/9))lPpnd (#n.1lq)). Now we show that V. DKL. (por (#o.1.19)llz (4)ll (2o.1.4) = VoLKI.
28

---

Published as a conference paper at ICLR 2026
Proposition 3. Consider the trajectory-wise KL divergence DKL (Pos (DoaL q)llp (ZruL ()) and tractable LKL defined as follows:
(aI, 01.1, ②) =StopGrad
Pap (wo. Paocia (woi (q) (1 + log
Pag (#o.z (q)
Pernt (1tv.2 97 )
log Io (a
(85)
Then, the gradient of DEL. (Pas (200.z Paces (dro.r. q)) and E[CKI. are equal:
V DKL. (Poo (zro.z @llleteer (rouL q)) V - BBBaaL P (KK..)
(86)
Proof.
DEL (pao (zo.z a Paues (2o.1
q))
(87)
Poo (Z0:L
q)
EuairPo (#o lq)
log
(88)
Poier (Fo.L.
q)
Pgo (10:
q)
Pso (1
9))
-Ba L Poosll (20.2.9)
log
(89)
Poesid Po.:L
q))
Poter (0:.1 9)
Pyo (ZO:L
q)
Pgs (x
Øc q)
EE PPadold (#0.2/9)
log
(90)
Proaid (Z0.1.
q
Panee (xi-1
Bi ②
where the change of measure from Pgc to Pread (Eq. continuity Pa. (x) & Preau (xo.1.())
88
to Eq.
89 is justified by the absolute
Vo DEL (Poo (Fo. ()ll (zo.z 1 (q)
(91)
Poo (IFO:L
q)
I.
Poo (x 1
xz. q)
= VolEmonn PPoaws
(.)))
Paosia (wort
Q)
log
Pasd (xi-1
Ci ④
(92)
1=1
Pgu (2ouz ())
b
Pao (xi-1
Izs 9)
E
Ve
log
Paoai (2o.L
②
Piret (x:-1
2in q)
Pas (Zo.L
9)
Pas (x;
Di 9)
log
(93)
Poeai (00.L
q
Poret (
0ei q)
oPo (2 xc 9)
Poo (30:L)
q))
Pao (x
1z q)
log
Poo (I
xz q)
Poeasa (ROLL
9)
Pornt (xi-1
24, q)
Pys (Z0:L
q)
1 q) VePgo (x,-1
(94)
Pyncu (Zu:L
9
Pgo (x,-1 1 12;
1=1
Pgs (ZozL 9)
(1+
Poo (Ii-1
xis q)
V Pop (xt, 23 9)
E
Σ log
Preas (Z0:L q)
Porer (x
1
Di, q)
Pgo (x 1
Ds. q)
(95)
V StopGrad Pgs (oo.r. ())
Pao (vo.r. 91)
Penna (DuiL | qq) -11+ + log
log go (as (x;)
(96)
Pynn (Iv:L q)
E(LLL)
(97)
Remark. Why the absolute continuity Po. (xo.1lq) < Peen (xo.L ()) holds? Recall that we consider two parametrization of our unmasking policy model, Iø and gonex given as follows:
exp
(a xn)) = Σn ンリ exp (hopp
for all mask indices a' e Ax.,
(98)
29

---

Published as a conference paper at ICLR 2026
Table 4: Comparison different unmasking policies across benchmarks. For our method, the dense reward follows the definition in Section 4.1 and Appendix The binary reward assigns only when (i) for SUDOKU, all masked positions are correctly predicted by LLADA, and (ii) for GSM8K, the final answer predicted by LLADA matches the ground truth.
Benchmark
Random
Max-Confidence
Ours (Binary Reward)
Ours (Dense Reward)
SUDOKU
0.616 0.612
0.705
0.801
0.817
GSM8K
0.684
0.701
0.703
Sudoku
GSM8K
Accuracy /Reward
0.3
Accuracy Reward
0.74
0.80
0.72
0:70
Donse Reward
0.68
Binary Reward
800
1600
2400
3200
1000
2000
3000
4000
Step
Step
Figure 6: Average reward per training step when optimizing the unmasking policy using (i) dense reward and (ii) binary reward.
exp((ho(Xxn))
a' E argtopk max mT())
D (a' xxx)=
Ees as Eargtopk aa exp (1)(())))aa))
(99)
otherwise,
where argtopk max TT(-) denotes the set of Top-K mask indices corresponding to the highest- confidence predictions given by the MDM denoiser ne. For the parametrization Jg. absolute continuity holds trivially. This is because I6 assigns non-zero probability to every unmasking index, and ma also assigns non-zero probability to every token. Consequently, every sequence Xg.r receives strictly positive probability under Jø- i.e., Ps. (xa:L) > 0 for all Xp:L E Ao x X1 x ... x Xt. Our remaining question is whether absolute continuity Pae (XorL ()) < Pneas (Xo:L q) also holds for J0nok This condition is indeed satisfied: Paonse (Xo:1. q) and Promoxaa (Xu.L. | ()) are equivalent measures, that is, they are mutually absolutely continuous (PraTw & Pyozaxau and Ps
< Pgenas This is because the set of sequences Xo:L for which Prempx (XO:L) depends entirely on the denoiser TT. and not on the particular choice of Jenexx Formally, recall that Pgomps can be written as:
Pranax (Xo:L (q)
II Jowss (an Nns q) To (Xn-1 Xn, ans q). n=1
(100)
Here, since To is strictly positive at every step, whether a given path Xp:L has zero probability is determined solely by Domps (a, X,, q). Inspecting Eq. (99). see that Dompx (a,, X,. q) becomes zero exactly when an argtopk max Ta X,, q). Thus, whether Pgenex (XO:L q) is zero depends entirely on Ty and independent the specific choice Jono Consequently, Paew. (Xo. q)
and Poong.x.a (Xo:1. q) are equivalent measures, and hence Paomyx (xo:L ( < Pgong Xo:L readily follows.
D ADDITIONAL EXPERIMENTS
D.1 UNMASKING POLICY OPTIMIZATION IN BINARY REWARD SETTING
In the experiments reported in Section we employed the dense reward defined in Section 4.1 and Ap- pendix E.] Specifically, the dense reward corresponds the proportion of correctly predicted masked positions for SUDOKU, and for GSM8K- K, it consists of the binary correctness signal augmented with the sum of log-probabilities from (
30

---

Published as a conference paper at ICLR 2026
Natural Question
GSMSK
QWEN-3 3 LLHHOLD
-2.05
-2.10
0.780
-2.15
0.795
2.20
400
800
1200
1600
1600
2400
Steps
Steps
Figure 7: Training dynamics of UPO when using the average log-likelihood of Qwen-3 32B as the reward. The left plot corresponds to the NATURAL QUESTION benchmark, perproe-purpose natural language generation dataset, and the right plot shows results on GSM8K. In both cases, the log-likes wi-likeeihood increases monotonically throughout training, indicating that the learned unmasking policy improves the generation quality.
In contrast, Theorem and Theorem 2 theoretically show that under the binary reward setting, our Kl.-regularized unmasking policy optimization can outperform the reference policy. Therefore, in this section, we additionally evaluate our method under the theoretically aligned binary reward setting and demonstrate that our KL-regularized unmasking policy continues to surpass the reference policy.
The experimental results are presented in Table 4 and the training dynamics are shown in Figure b On SUDOKU, the dense reward yields slightly higher final accuracy and converges faster than the binary reward; however, in both SUDOKU and GSM8K, our method consistently outperforms the max-confidence baseline. These empirical observations are aligned with the theoretical predictions of Theorem and Theorem 2
D.2 TOWARD SCALABLE AND GENERALIZABLE UNMASKING POLICIES
While the main experiments focus on correctness-based rewards within structured reasoning bench- marks, our next goal to learn an unmasking policy that scales beyond task-specific supervision and remains useful across natural language domains. To this end, we conduct preliminary study replacing the verifiable reward with preference-based signal obtained from a stronger autoregressive model (ARM).
Specifically, we use the average log- ligeliikl of Qwen-3 32B Yang et al. 2025) as the reward signal, keeping all training settings identical to the GSM8K setup except for the reward definition. We evaluate this configuration on two datasets: GSM8K, which enables comparison to correct- metrics, and NATURAL. QUESTIONS (K (Kwiatkowski et al., 2019). large-scale open-domain question-answering benchmark representing unconstrained natural language generation.
As shown in Figure 7 the learned policy consistently improves the reward across both datasets. Since the log-likelihode of a strong ARM is standard proxy for generation quality, this trend indicates that the policy can adapt even without task-specific correctness signals. Moreover, on GSM8K, this preference reward also yields measurable downstream improvement, achieving 70.5% accuracy compared 68.2% from the max-confiee baseline.
These findings provide early but encouraging evidence that unmasking policies may be learned under general preference signals rather than explicit oracle correctness. We view this as step toward scalable, pprpose token ordering strategies aligned with natural language structure rather than tied individual benchmarks or reward designs.
31

---

Published as a conference paper at ICLR 2026
D.3 PASS AT N RESULT OF UNMASKING POLICY MODEL
In Section 3.1 we demonstrated that the @ of 9oop-K can grow significantly as N increases, surpassing the single- trajectory accuracy of Scont- This shows that performance can improve through the diversity induced by exploring multiple denoising paths, Our method can similarly benefit from both single-shot and multi-shop generation, which can be achieved by employing Gumbel-Softmax (Jang et al., 2017]. When using Gumbel-Softmax, sampling is defined as:
0.9
PossBN
0.8
Ours (a=
0.2
Max Confideene
Top-5
Randor
0g = arg max ( (a' a'EAxx
xn) + +
6 ((0,00
Number paihs
10
Figure 8: @ on GSMSK. The dashed red line is the single- trajectory accuracy of max- confidence Gcout, and the dashed green line is the trajectrt accuracy of our unmasking policy model where a = 0 in Gumbel- Softmax.
Therefore. if one wishes to perform single-shot generation, setting a = 0 yields deterministic selection and achieves high accuracy. Conversely, increasing a enables multi-shot generation by injecting diversity into the sampled trajectories. In fact, when is sufficiently large, the resulting behavior becomes theoretically identical to that of Ilop-K. We provide empirical evidence in Figure where the Pass N curves for Derox (Ours) and gTop-K match almost perfectly, confirming this correspondence.
E EXPERIMENTAL DETAILS
E PARAMETRR AND OTHER DETAILS
For all experiments, we set the group size G as 6. We utilize AdamW optimizer with 31 = 0.9, 82 = 0.99 and weight decay is 0.1. We perform gradient clipping where the maximum gradient norm is set as 0.2. For computational efficiency, we utilized Flash Attention 2 and 4-bit quantization for MDM parameters. We set the gradient accumulation steps as 8. For SUDOKU and ZEBRA We utilize 3 40GB A100 GPUs for SUDOKU and ZEBRA and use constant learning rate of 3 x 10-6, We perform a gradient update 16 times for each prompt. Since MATH500 and GSM8K are much harder to learn, 40GB A100 GPUs and use cosine learning rate scheduler starting from 2 x 10 and ends at 10, 000 updates. We perform gradient update times for each prompt. Following Diffu-GRPO (Zhao et al. 2025), we save models for every 100 steps and report the best accuracy. For all experiments, we set 8 = 10-4 For softmax realizations we set 0.05 for Icont" Since ZEBRA benchmark is symbolic dataset that differs from human language, we performed supervised finetuning
more computation resources and set stable hyperparameters. We utilize 8
E.2 REWARD AND PROMPT TEMPLATE FOR EACH BENCHMARK
For SUDOKU and ZEBRA, we set the reward function as the matching rate between the generated sequence and the answer. For example, if there are 8 masks in sudoku puzzle and the model got 5 correct answers, then the reward is 5/8 = 0.625. For MATI550 and GSM8K, the reward is the model got the correct answer and 0 otherwise. We do not use formatting rewards such other GRPO papers since unmasking policy models do not have the ability to change the token probability distribution itself. However, for MATH500 and GSM8K, the reward is too sparse, as it only gives
or 1. Therefore, we utilize additional reward Σ - log no (x [x,, a,). We do not utilize it directly, but we assigned relative reward, e.g., 0.25 for the trajectory with the largest probability, 0 for the trajectory with the smallest probability. We have confirmed that trajectories with higher probability are more likely to get a correct answer, and this reward indeed stabilized the unmasking policy model training process. Similarly, large-scale MDM such as Seed diffusion (Song et al. 2025) post-trains the MDM to fit to the generated sequence with a higher ELBO.
We provide a prompt template for SUDOKU, MATH500, and GSM8K below, and ZEBRA is omitted since it is a symbolic dataset.
We have utilized GitHub code from Zhao al. 2025
32

---

Published as a conference paper at ICLR 2026
Prompt template used for SUDOKU
Please solve the Following 4x4 Sudcku puzzle. The puzzle is provided as 16-character string reading lefh-to-right, top-to-botton where represents empty cells.
Rules:
Fill empty cells with digits 1-4
Rach row must contain digits 1-4 exactly once
Each column must contain digits 1-4 exactly once
Each 2x2 box must contain digits 1-4 exactly once
Important:
Your solution must be COMPLETE 16-character string with only the digits 1-4, representing your final solved grid. Never leave it as
Respond in this exact format:
</reasoning>
<answer>
[First row of 4-character solutien] [Second rew of 4-character solution]
[Third rew of 4-character solution| [Fourth of 4-character soluticn] </answer>
Solve the following Sudoku puzzle: 1
2
3
<answer>
1 3. [M] 2
[M] [M] [M]
2. [M] 3
[M] [M] [M]
</answer>
Prompt template used for GSM8K and MATH500
You are a math expert. You will be given question solve. Solve step by step. Wrap the final answer in (boxed(). Respond in the following format: <reasoning>
Your reascning here </reasoning>
<answer>
\bboxed....
</answer>
<reasoning> ng>
128 Masks
33

---

Published as a conference paper at ICLR 2026
3
2
2
3
2
3
(a) Ours
4
3
2
2
2
2
First Error!
2
2
3
2
2
(b) Max-confidence
Figure 9: (a) Unmasking path of our unmasking policy model and (b) unmasking path of the max- confidence baseline on the same 4x4 Sudoku instance. Blue digits indicate the original clues the puzzle, red digits denote the value unmasked at the current step by the unmasking policy and sampled by the LLADA: ADA, and black digits represent previously filled values. Our model successfully solves the entire puzzle, whereas the max-confidence method begins make mistakes from the point marked "First error!".
F SUDOKU UNMASKING EXAMPLES
To better illustrate the qualitative behavior of our learned unmasking policy, visualize representa- tive unmasking trajectories from the main SUDOKU experiment. In all examples, the underlying MDM samples the token values via argmax over the model's categorical distribution, difference between methods lies in the choice of the unmasking position. As shown in Figures 9 and 10 our learned policy tends prioritize positions whose values already strongly determined by the puzzle constraints, leading to stable infilling and consistent recovery of the correct solution.
the only
In contrast, the mox-idniiee heuristic shows clear failure behaviors. In Figure D chooses reasonable unmasking position, but the predicted value is wrong. This single mistake immediately propagates and breaks the remaining steps. In Figure I0, it selects a position whose correct value is not yet identifiable. The predicted token does not violate any explicit Sudoku rule, but it places the puzzle in a state that cannot lead to a valid solution, highlighted by "Bump!". The next prediction therefore fails.
These visualizations highlight central claim our work: the choice where to unmask as crucial as the values predicted by the MDM, and a learned policy can more reliably identify structurally deterministic positions than hand-crafted schedules.
G USAGE OF LARGE LANGUAGE MODEL
have utilized a large language model assist with grammar and wording for writing the paper.
34

---

Published as a conference paper at ICLR 2026
2
3
2
3
2
4
2
3
(a) Ours
2
3
2
2
4
3
3
4
2
3
2
Bump!
First Erroot
2
3
2
2
3
4
3
3
4
2
(b) Max-confidence
Figure 10: (a) Unmasking path of our unmasking policy model and (b) unmasking path of the max-confidence baseline on the same 4x4 Sudoku instance. Blue digits indicate the original clues in the puzzle, red digits denote the value unmasked the current step by the unmasking policy and sampled by the LLADA, and black digits represent previously filled values. Our model successfully solves the entire puzzle, whereas the maxxcconiieece method reaches point labeled "Bump!" where the unmasking decision does not immediately violate Sudoku constraints but nonetheless leads to an unsatisfiable partial state. Consequently, the very next prediction results in
35


Reviews:
[{"review_id": "r8A3NFAEzk", "reviewer": "Reviewer_DEPM", "summary": "The paper proposes learning an unmasking position-selection policy for masked diffusion language models (MDMs) by casting denoising as a KL-regularized Markov Decision Process and optimizing a GRPO-style objective against a strong reference scheduler (e.g., max-confidence or Top-K). The authors prove (i) fixed-point convergence that improves expected reward over the reference and (ii) a “reference-KL tightening” result showing the terminal output distribution becomes closer to the data than that induced by the reference policy [Thm. 1–2]. To make training tractable, they derive a surrogate token-level objective (LUPO) whose gradient aligns with the intractable output-level objective and instantiate realizations for different references [Eq. (19), Table 1]. Empirically, a lightweight policy model (≈134M) plugged into a frozen 8B MDM improves accuracy across SUDOKU, ZEBRA, GSM8K, and MATH500 over random, margin, entropy, and max-confidence schedules, and adds gains on top of diffu-GRPO post-training.", "questions": "1. In Sec. 2, please (i) formalize the exact training objective under which the “true vs. model reverse-process” divergence occurs, (ii) explain the mechanism by which this leads to inference-time errors in the unmasking policy (with a minimal worked example), and (iii) discuss whether the issue can be mitigated during pretraining (e.g., via alternative noise schedules, calibration terms, or reverse-process consistency regularization) rather than deferred to policy learning at inference.\n\n2. Does the learned policy transfer across MDM backbones (e.g., Dream-7B), or between tasks without re-training? Any negative transfer?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- Clear problem focus: unmasking order materially affects discrete diffusion inference; paper formalizes learning this policy rather than relying on heuristics.\n\n- Theoretical contributions: output-level GRPO formulation with stated convergence and a KL tightening guarantee relative to a reference policy; proofs sketched in main text and detailed in Appendix C.\n\n- Practical surrogate: token-level objective with clipping and an empirical divergence estimator, plus realizations for max-confidence/softmax/Top-K with explicit training recipe.\n\n- Solid empirical signal: consistent wins vs. strong baselines on four benchmarks (e.g., SUDOKU +11.2% over max-confidence; GSM8K +1.9%) and compatibility with diffu-GRPO (GSM8K +1.3% over max-confidence).\n\n- Efficient design: trains a small policy while keeping the 8B MDM frozen; memory-efficient algorithm described (Algorithm 2).", "weaknesses": "- Sec. 2 reverse-process mismatch is under-explained. The section claims that there exist states where the true reverse process and the model’s reverse process diverge, but the accompanying example is hard to follow and it’s unclear what concrete training pathology this implies. Specifically, how this mismatch arises under the stated pretraining objective, how it manifests during inference-time unmasking (e.g., biased token selection vs. compounding calibration error), and whether it is an artifact of modeling choices (masking schedule, likelihood factorization) are not clearly connected.\n\n- Evaluation breadth. Benchmarks are primarily small-format puzzles and math QA; no natural-language generation or long-form tasks to test the generality of unmasking policy beyond structured reasoning.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "oiBRSGg5z1", "reviewer": "Reviewer_cCbK", "summary": "This paper proposes a method to improve the token generation order of pretrained Masked Diffusion Models (MDMs) through RL. Specifically, the authors train a learned unmasking scheduler using GRPO. Compared to heuristic orderings such as confidence-based or margin-based selection, the proposed approach achieves better performance across several benchmarks.", "questions": "* **Pass@K:**\n  Could the authors report *pass@K* results? It would be informative to know whether the proposed method sacrifices diversity (e.g., generating overly similar samples) in pursuit of higher accuracy.\n\n* **NFE (Number of Function Evaluations):**\n  How many diffusion reverse steps are used during inference? Please specify the number of denoising iterations to facilitate fair comparison with other methods.\n\n* **Qualitative Analysis:**\n  It would be helpful to visualize the learned unmasking order on SUDOKU and compare it to existing baselines such as max-margin or max-confidence. A qualitative comparison could clarify what the learned policy actually discovers beyond heuristics.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 1, "strengths": "**Novelty:**\nThe idea of learning the generation order of a non-autoregressive model through reinforcement learning has been previously explored in classical non-autoregressive language models (see [1]). However, applying this idea to **modern diffusion-based generation models** represents an original and meaningful contribution.\n\n**[1]** Insertion-based Decoding with Automatically Inferred Generation Order, [arXiv:1902.01370](https://arxiv.org/abs/1902.01370)", "weaknesses": "**Generalizability of the Method:**\nSince the optimization is based on RLVR (GRPO), the method can only be applied when a **verifiable reward signal** is available. This limits its general applicability to domains where rewards can be explicitly computed.\n\n**Suggestion:**\nCould the authors explore the use of a **preference reward model** to enable training in more general settings? Alternatively, using the **generative perplexity** of a strong language model such as GPT-2 Large as a reward could be an interesting direction for learning generation order without explicit ground-truth supervision.\n\n**Incremental Performance Gain:**\nDespite requiring additional RL-based training, the performance improvements over strong baselines appear modest—except for the SUDOKU task, where the generation order plays a particularly critical role. The paper would benefit from discussing why the proposed method does not yield stronger improvements on other benchmarks.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "h8l4v8HIB1", "reviewer": "Reviewer_BVXC", "summary": "This paper explores learning the unmasking order in MDMs. People have realized that a pretrained MDM has an any-order sampling property, which led them to explore heuristic unmasking policies. These use the information of per-index logits, so they don't fully leverage the information and also rely on heuristics. This work focuses on learning the unmasking policy in fine-tuning the LLaDA using RL methods, achieving better performance than baselines.", "questions": "- In Figure 1, we can see that the Pass@N performance can actually go far beyond the Max-Confidence baseline, indicating that there's indeed a decent sampling path (that has a much higher accuracy even than Max-Confidence). The results in the paper, however, aren't that significant. I'm curious how the authors could discuss the discrepancy, i.e., what was the main factor that reduced this expected huge gap?\n\n- I wondered what the authors meant by `Meanwhile, on GSM8K, the much larger search space makes Top-K preferable for stability. In future work, we plan to develop an unmasking policy optimization algorithm that explores the full mask tokens even when the search\nspace is large.'", "limitations": "", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "First of all, this paper is clearly well-written. The authors explain the background clearly and then the problem that they want to solve in the paper. Also, examples in Sections 2 and 3 assist the reader in grasping the context better. Moreover, the experiment design (loss design, architectural modification) is really clear to follow. I believe this paper takes a step toward finding an effective unmasking strategy beyond prior heuristic strategies.", "weaknesses": "Perhaps the most significant weakness of this paper is the less-than-ideal experimental results: We can see that on several tasks, the performance of the trained model's sampling strategy is comparable to the Margin sampling strategy from (Kim et al, 2025). Given that the method proposed in the paper requires additional training (while the Margin strategy is training-free), it might be questionable to some people that if we really need an approach that requires training.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "O5yUh1N7ux", "reviewer": "Reviewer_9cnV", "summary": "This paper presents an algorithm to train a unmasking-position policy for masked diffusion models. While rule-based unmasking policies based on confidence or probability margin are popular,  the authors propose reinforcement learning (more concretely, GRPO) with KL-divergence regularization starting from such rule-based reference policies. The proposed algorithms have theoretical guarantees under technical assumptions, and show practical improvements on puzzle and math benchmarks.", "questions": "In addition to the points raised in the Weaknesses section, I have the following question:\n- [Q1] As the authors write in +, Huang et al. (2025) applies RL jointly to the masked diffusion model and unmasking policy so the existing method is heavy. Doesn't it mean that, by fixing the diffusion model, the existing method can train only the policy model for lighter computational cost? Can the authors discuss a bit more on the difference and the preposed method's original contribution? (I assume it is basically about initialization and KL regularization though.)", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "Adding KL-divergence regularization in the RL of unmasking policy yields theoretical tractability and improvement in ablations. It also proposes a KL-divergence inequality regarding among the data distribution, reference policy, and optimized policy, which can be more informative than the reward inequality.", "weaknesses": "The weakness of the paper primarily lies in the mismatch of theory and practice.\n- [W1] While the analyzed algorithm in Section 3.2-3.3 is based on 0-1 reward, it is totally different from what is described in Section D.2; the hitting rate in puzzle benchmarks, and more heuristic additional dense reward in math reasoning benchmarks. The readers might think that the practical performance is just based on the heuristic choise of rewards rather than theoretical insights.\n- [W2] Theorem 2 suggests the KL divergence is closer (or equal to) the reference policy's KL divergence, but I feel that it only reflects the reward/rate of correct answer. For example, if there are multiple answers in each query, then there is a case where $L_{\\text{MDM}}>0$ but $r_{\\text{ref}}=1$, with the model only outputting a small subset of possible answers. In such a case, the policy is already converged, and the KL divergence is the same. Though this example does not violate the inequality (as it is not a strict inequality), the message after Theorem 2 (L281) \"Theorem 2 guarantees the closer sampling to $p_{\\text{data}}$ or ideal MDM $θ^*$ than $g_{\\text{ref}}$\" is somewhat misleading to me.\n- [W3] In Theorem 2, the assumption $\\text{supp}(p_{\\theta*}(\\cdot|q))\\subseteq \\text{supp}(p_{g\\text{ref}}(\\cdot|q))$ looks too strong, especially when considering the deterministic nature of the \"confidence\" unmasking policy.\n- [W4] While $g_{\\text{conf}}$ is deterministic, since you have a sampling randomness with $\\pi_\\theta$ in addition to the positional choice $g$, the sampling results should still be random. I am not sure why Max-Confidence is treated as if a constant in Figure 1, which might be misleading.\n- [W5] The proofs are not written in a careful way. For example:\n    - [W5-1] At least in Equations 30, 31, 36, 38, 39, 71, and 73, the parentheses are not appropriately closed. Some are just trivial mistakes but some make the equations ambiguous. At this rate of obvious mistakes left, I'm not sure if the draft is appropriately checked by the authors themselves.\n    - [W5-2] In Eq. 82, why is it a strict inequality rather than \"$\\leq$\"?\n    - [W5-3] The equality from Eq. 92 to Eq. 93 is valid only when $p_{g\\phi}$ is abosolute continuous w.r.t. $p_{g\\phi\\text{old}}$. Is it an additional assumption?\n\nWhile the proposed method consitently improves upon existing methods, the current manuscript has a large room for improvement as discussed above, which makes me lean to rejection.\n\nTypos:\n- L95: $\\sum T_{x'\\in X}$ -> $\\sum_{x'\\in X}T$\n- L888: {1, L} -> {1, ..., L}?\n- L937: Let let -> Let", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)