Bringing Stability to Diffusion: Decomposing and Reducing Variance of Training Masked Diffusion Models
Authors: Mengni Jia
Mengyu Zhou
Yihao Liu
xiaoxi jiang
guanjunjiang
Venue: ICLR 2026 Poster (2026)
Keywords: Diffusion LLMs
Masked Diffusion Models
Training Variance
Training Stability
Mask Schedule
Mask Sampling
Primary area: foundation or frontier models, including LLMs
TL;DR: We stabilize MDM training by deriving a variance decomposition and introducing two core methods: P-POTS, which is Pareto-optimal among all unbiased t-samplers, and MIRROR, which complements it. Experiments yield clear gains on final performance.

Abstract: Masked diffusion models (MDMs) are a promising alternative to autoregressive models (ARMs), but they suffer from **inherently** much higher training variance. High variance leads to noisier gradient estimates and unstable optimization, so even equally strong pretrained MDMs and ARMs that are competitive at initialization often diverge after task-specific training, with MDMs falling far behind. Currently, there has been no theoretical explanation or systematic solution. In this paper, we derive **the first decomposition** of MDM training variance into three sources: {A} masking pattern noise, {B} masking rate noise, and {C} data noise -- while ARMs are only affected by {C}. This cleanly explains the fundamental training gap. Building on this foundation, we design six variance-reduction methods, including two core methods: (1) P-POTS, a **Pareto-optimal** $t$-sampler that minimizes training variance by sampling harder $t$ values more often with appropriately smaller update steps, and (2) MIRROR, which uses negatively correlated samples to reduce {A}. Experiments show that, compared to standard MDM training, our methods improve accuracy by **7–8\%** on complex reasoning tasks, while simultaneously reducing run-to-run variability to **near ARM levels**, substantially narrowing the gap with strong ARM baselines; in most settings, even the best baseline method runs remain below the worst run of our method.

Published as a conference paper at ICLR 2026
BRINGING STABILITY TO DIFFUSION: DECOMPOSING AND REDUCING VARIANCE OF TRAIN- ING MASKED DIFFUSION MODELS
Mengni Jia', Mengyu Zhou?, Yihao Liu", Xiaoxi Jiang'. Guanjun Jiang?
University of Cambridge, ?Qwn Large Model Application Team, Alibaba, 'Peking University
ABSTRACT
Masked diffusion models (MDMs) promising alternative to autoregressive models (ARMs), but they suffer from inherently much higher training variance. High variance leads noisier gradient estimates and unstable optimization, even pretrained MDMs and ARMs are competitive at initialization, they often di- verge after task-specific training, with MDMs falling far behind. Currently, there has been no theoretical explanation or systematic solution. In this paper, we derive the first decomposition of MDM training variance into three sources: AA mask- ing pattern noise, (B) masking rate noise, and (C) data noise while ARMs are only affected by C. This cleanly explains the fundamental training gap. Build- ing on this foundation, we design six variance-reduccion methods, including two core methods: (1) P-POTS, Pareto-optimal f-sampler that minimizes training variance by sampling harder t values more often with appropriately smaller up- date steps, and (2) MIRROR, which uses negatively correlated samples to reduce A. Experiments show that, compared to standard MDM training, methods improve accuracy by 7-8% on complex reasoning tasks, while simultaneously reducing run-to-run variability to near ARM levels, substantially narrowing the gap with strong ARM baselines; in most settings, even the best baseline method runs remain below the worst of our method. The code will be released
INTRODUCTION
Stundard
P-POTS MIRROR
PUF
PDF
P-POTS
Training
Tmining Dainee
Sumple Xg from Ptaes (+)
Sample from U10,1]
Sample Xy from Pdana
Sample from
MIRROR
0.3
0.3
Mask with probability
Mask with probability complementaniy
Prompt
Response
[MASK]
Variance Source
Graphic illustration three sources of training variance in MDMs. The left panel illustrates standard MDM training, compared to our core methods (P-POTS and MIRROR) on the right.
Figure I:
Recently, masked diffusion models (MDMs) Austin al. 2021 Lou et al. 2023 Sahoo et al.
2024) have been viewed as strong alternative architecture for both large language models (Dream- 7B Ye al. 2025), LLaDA-8B (Nie et al. 2025), DiffuCoder (Gong et al. 2025) and multi-modal models (Dimple-7B Yuu et al. 2025). MMaDA-8B (Yang et al. (025))) They operate by masking input tokens at randomly sampled masking rate / and learning to reconstruct them. By moving

---

Published as a conference paper at ICLR 2026
beyond sequential token modeling, they potentially address limitations of auto-regressive models (ARMs), such as lack of parallelism, exposure bias, and reversal curse (Berglund et al. 2023
However, the noise-based training paradigm that powers MDMs also makes their training unstable. Numerous studies have reported that diffusion training suffers from high variance (Meng et al. 202T
Xu et al. 2023 Jeha et al. 2024 Kim et al.) 2024 Zhu et al. 2025 Arriola et al.) (022).. We refer to this phenomenon as the training divergence from ARMs, which manifests in two aspects:
Across-run variability. Even under identical training setups, MDMs can converge to very dif- ferent solutions:
since parameter updates scale with loss gradients, fluctuations in losses translate directly into erratic updates. By contrast, ARMs are more predictable: with the same data order, they converge reliably to the same solution across runs.
Within-run sub-optimality. Although pretrained MDMs and ARMs show comparable abilities at initialization, we consistently observe a sharp divergence after post-training: once fine-tuned on the same dataset, ARMs always outperform MDMs by wide margins, even in cases where MDMs begin from a stronger point. This highlights that, even within single run, the variance gap makes MDMs use supervision resources less efficiently, as noisier updates slow convergence.
To address this problem, several prior works have proposed mitigation strategies. However, these efforts remain isolated and ad-hoc, lacking a unified theoretical framework for explaining the high training variance of MDMs. For example, Zhu et al.] (2025) evaluates both the policy and reference ELBOs on the same time-steps and masking patterns so that their Monte Carlo errors cancel out at comparison, but their method is limited to preference optimization. Arriola et al. (025)) proposes adaptively re- re-fitting masking-rates samplers at each evaluation, selecting from fixed number of candidate clipped intervals next evaluation. While this strategy can adjust the learning process accordingly. its reliance on such coarse candidate intervals remains heuristic. These limitations call for more principled solutions.
C [0, 1] to determine the best interval for training until the
Motivated by this, we revisit the basic definitions to develop systematic explanation and theoreti- cally grounded solutions. Specifically, Sahoo et al. (2024) derived a principled training objective:
CMDM t- 100 (o t, t,), to ot
log Σ E=1
(ab 1:b
.t,, x))
a
(1)
for interpolating MDMs, where the loss lo is a function of the clean data x. the masking rate f, and the noised (masked) input 2, obtained from the forward process q( x, t). To understand the variance origins in this training objective, we derive the first systematic variance decomposition (proof in A. Mathematically, this yields:
Var (lo) Exzo [Vrrz, (o zo, ()) + Ezo (ga (20,, 1) mo)] + Var.o (Belg0(20n:) 1) (2)
Masking Pattern Noise A
Masking Rate Noise ①
Data Noise C
with (o( f) =E,, [4e Ho, d). This decomposition relies on no hard assumptions and emerges naturally, so it both explains existing methods and lays a foundation for understanding how they may complement one another. As shown in Figure the MDM training variance is decomposed into three factors: AA Masking Pattern Noise (randomness from 2; even if xo and f are fixed), B Masking Rate Noise (variability the expected loss g(zo, t) across different / for the same rg), and C Data Noise (variability across To. as some are inherently easier or harder to predict). In contrast, ARMs are subject to only data noise ( Therefore, the high training variance of MDMs is explained by the extra sources (A) and B,, highlighting the need to reduce them to stabilize MDM training.
To address this challenge, we introduce six techniques, including two core methods (P-POTS and MIRROR in 3 and four additional ones. Specifically, P-POTS is the Pareto optimal (-sampler among all unbiased ones in minimizing A, B and C: it focuses training regions while doin-weighting their updates to prevent destabilizing optimization. MIRROR, on the other hand, constructs two complementary noised inputs and employs their negative correlation to explicitly reduce A Crucially, these two methods exhibit strong synergy, leading us to recommend P-POTS alone for cost-efficiency and P.POTS+MIRROK for maximal performance.
high-variance t-
Experiments validate our claims. On GSMSK, the accuracy of standard MDM training ranges from 50.6% to 53.7%, while MIRRTS+MIRROR improves to 38.6%-62.00% On HiTab, the accuracy
2

---

Published as a conference paper at ICLR 2026
On these textual datasets, method allows MDMs to match or even surpass ARMs, with run-to-run variability reduced to near-ARM levels. On text-to-image-MM, the CLIP score range narrows from 28.61-34.28 to 34.10-35.27.
range rises from 52.9% 62.6% to
Our major contributions can be summarized as follows:
We derive the first systematic variance decomposition for MDM training as a foundational frame- work (Eq. and link prior w orks relevant variance factors they actually address (
We introduce six variance-reduuctin methods, with P-POTS Pareto optimal and P-POTS+MIR- ROR giving the best results. Each method is theoretically grounded with proofs or analysis, simple implement (requiring little or no hyperparameter tuning) and robust across datasets.
We validate our methods both textual and multi-modal data across various masked diffusion models, demonstrating their effectiveness and generality in stabilizing MDM training.
2 RELATED WORK
The inherently high training variance in diffusion models has been empirically observed. Several remedies have been proposed, yet most of them target continuous diffusion models that are widely used in image generation, and relatively few of them directly apply masked diffusion models, the model architecture behind dLLMSs (diffusion LLMs). Moreover, due to the lack a theoretical de- composition of the variance sources, the methods remain largely isolated and heuristic. Concretely:
Meng al. 2022) proposes using denoising models trained at antithetic levels of Gaussian noise, which reduces mask pattern noise A but limited to continuous diffusion models.
Xu et al. (2023) computes the denoising score as a solf-normalized average over reference batch. It reduces AA) and (B. but introduces systematic bias with finite reference batches and suffer weight degeneracy in high-dimensional or poorly covered posteriors.
apportacc-ssmplee weighted
Jeha et al. (2024) uses a kth order Taylor expansion of the denoising score as a control variate and subtracting it (with an optimally fitted coefficient) from the Monte Carlo estimate. However, it de- pends on closed-form Gaussian moment calculations that don't generalize to non-Gaussian noise (e.g., masked diffusion models). Moreover, the method's reliance low-order Taylor approxi- mations leads to ineffective variance reduction under high noise or rapidly varying derivatives.
Kim et al. (024) dynamically adjusts the probability of sampling each diffusion timestep based online estimates of gradient variance. This reduces mask ratio noise BB) but does not apply importance reweighting, which biases their loss estimates; moreover, the sampler is restricted to a simple Beta(a, b), which may fail to capture the true optimal sampling distribution.
Zhu al. 2025) formalizes the variance of ELBO-based gradient estimates in masked diffusion language models and employs optimal Monte Carlo budget allocation and antithetic sampling, reducing the mask pattern noise (A) but is limited in preference optimization.
Arriola al. (2025) select t-sampling interval by evaluating candidate sub-intervals and choos- ing the one that minimizes evaluation variance, which reduces BB) While they justify clipping away extreme is (e.g., [0.45, 0.95)), our analysis suggests extremes can be worth em- phasizing, so their gains likely come from clipping that sometimes approximates the optimal p(t).
3
METHOD
Let's recap the standard training algorithm for masked diffusion models (MDMs). Let zo denote the concatenation prompt and response tokens and P denote the number of response tokens, then:
1. Sample clean data za from the data distribution.
2. Sample masking rate t [4(0,1] 1].
3. For each eligible token i, sample U, ~ 4((0, 1]. and construct 2, by replacing (i(() with [MASK] U, < t, so each eligible token is independently masked with probability t.
4. Compute the sample loss 1o(20, 1, xx)
Den ((xa(1) [MASK]] log po (x0 (a))
3

---

Published as a conference paper at ICLR 2026
5. Average lo (2o. (,o.t,20 21) across the batch to obtain lo and update 0 via backpropagation.
Which tokens are eligible for masking depends on the training stage: in pretraining both prompts and responses eligible, while in supervised fine-tuning only responses are. This algorithm is special case of discrete diffusion (Austin et al. (022)) - MDMs with a = t (see Eq. D. Nearly all pretrained MDMs adopt this setting. Beyond this, MDMs can model both text and image tokens Yang et al. 2025), making it a unified framework for both language and multi-modal modeling.
To minimize the training objective in Eq. [ gradient-based methods such stochastic gradient descent (SGD) are typically used (Robbins & Monro 1951). However, these methods require an empirical estimate of Eq. | which is usually the batch mean of losses 1 -i=1 ((=0)) t,, (8)
8). For ARMs, this estimator is generally reliable, since the variance (Eq. 2 arises only from O However, MDM training, we question its reliability, since extra sources (A) and (B come into play.
While the standard MDM training procedure gives an unbiased estimator of Eq. its variance can be significantly high. This raises central question can we construct an alternative estimator that
remains unbiased but achieves lower variance, providing a more stable basis for optimization?
In the rest of this section, we introduce six methods two core methods and then four others to address the problem. A qualitative analysis of their combinations is provided in A.6
3.1 CORE METHODS
3.1.1 P-POTS: PARAMETRIC- PARETO OPTIMAL A-SAMGLING
P-POTS (Parametric-Parte Pareto Optimal Timestep Sampling) replaces the masking rate sampler t ~ U10, 1] with data-fitted non-uniform one to reduce A B O.
To estimate Eq. standard training uses t ~ [0, 1] directly: unbiased but high variance ($3) Instead, one could use another t-sampler that remains unbiased but reduces variance via importance weights (Kahn & Marshall 1953).
We now derive the optimal unbiased (-sampler in the view of minimizing ABC+ If we instead sample ptt), we can re-weight to form p(t) le- This estimator has the same expectation as t~ U(0,, 1] and is thus still unbiased in estimating the training objective as shown below:
Etcp, de p(1) (z t, m2) = L.' p(t) p(t) a())tt ('' g(t) dt E,NU, (l0),
(3)
With Eq. 3 establishing the unbiasedness, we next turn to minimizing its variance from xo, t and It, which is exactly B in Eq. 2 Define v(t) Varza,,, (lg 6) and g Ezg. (le t). Under regularity conditions, Fubini's theorem (Rudin) 1987) allows us to change the order of expectation between & and To, Te. The variance can then be written as (detailed proof in A.552.))
p(t) 1 lo
/''
g(t)² + uu() dt -t-6 (f/ g(t) dt
A +
B + Va
(4)
p(t)
To minimize Eq. over all t-samplers, one Lagrange multipliers (agrange 1811 yielding
the optimal choice p" pt(t)= = fo v +(0)+ de gg(()²² + v(t). In minimizing A, B and C. this
+v(0)
is Careto-optimal among all unbiased t-samplers of U|0, 1]: no other choice improves them jointl..
Algorithm of P-POTS. Since g(t) and v(t) are unknown, P-POTS estimates them empirically:
(4) )._,, draw b values [t,}g-, and, for each (z) ,t,). draw e masked before training, for each of {ap
samples (1 (b.j.h) J5 with losses ligj == -6((0)) , 11 (.3.k))) We then compute
f==1
0j ac
ligky 0j
Σ
(Cij. Σ (eg D. -)
+ 0
for j L,...,b. b. The weights { represent an empirical scatter of the underlying optimal p" (t). To capture its structure, we design model called EPR
Root) as:
p EPR (1) V at b(1 1)9 A exp(2kl"),
a, A, 0,r,q m.
(5)
4

---

Published as a conference paper at ICLR 2026
gtt] 0x13expl2.77 (
6t0 0011 0.11 exp(4.61(120) vCt 1.79 + 0.72 1//01
g00 1..13 expr2.220 (0
vit) 3.38 phitt 0.00 -j--i
vtti 5.47 0.00 53--4
(a) OpenSeience
(b) GSM8K
(c) HiTab
Figure 2: Empirical (p, 1 1 (scatter) and the fitted curve p (t) vg(()² v(t) (line) three datasets: (a)
OpenScience. (b) GSM8K, and (c) HiTab. The equations each subplot show the fitted forms g(t) and ut).
which together characterize the sampling distrihution across masking rates
with FER ((t) A cxp(xt")) and e EE = b(1 -t--tt t)" + +at", and fit it via KL divergence. By the EPV rule that suggests 10 data points per fitted parameter (Peduzzi et al. 1996). we set 70 and slightly raise a and e 15, which proved reliable. As shown in Figure 2 EPR fits 170 almost perfectly with only seven parameters, suggesting captures the characteristic form of optimal p" (t)-
We call a hard to train if / gg((t) (t) is large, typically at higher (Figure - Our P-POTS allocates more samples to these regions, providing extra training on them, while the reweighting factor 1/p"(t) keeps their contribution small that their noisy signals do not dominate and hurt global optimization. In this way, p" (t) acts an adaptive sampler, directing training effort where most needed and using resources more efficiently.
Intuition behind P-POTS.
Intuition behind EPR. To motivate our forms of EPR (t) and g EPR (t) in Eq. 5 we consider loss behavior at different The mean loss g(t) reflects average task difficulty at t. Its growth is not just because more tokens are masked, but because entire reasoning chains are broken. Once chain is cut, the model must guess at multiple points, and errors compound multiplicatively (or equivalently, add in log space), motivating exponential form with delayed ignition. The design of m > ensures BPR (0) = 0, since the average loss gt)) t is 0 and must be local minimum.
Meanwhile, empirically, we find that vtt) falls when is small, rises after a local minimum, and falls again after a local maximum thus exhibiting up two turning points. We explain this behavior as follows. At small I, a loss spike occurs if a critical set of r key tokens is masked simultaneously (~~ !"). At large loss remains high unless variance of a rare event is approximately its probability, and multiple nearly independent rare events contribute additively, these mechanisms together generate the observed curvature Based on this explanation, we propose the model e(t) = at" + b(1 t)9.
of q key tokens survives (~ ()9). Since the
While one may wonder it suffices to naively {p, 1 with polynomials, we argue this prob- lematic for three reasons. First, as discussed in the next paragraph, p(t) faces a drift problem: its shape may evolve during training. Although this drift is typically slow, it can become severe if the training loss remains high (e.g., MMaDA on tot-ttoo-mageMM11 Consequently, polynomials quickly lose effectiveness if the true functional form of p" (t) is not well understood. Experiments in A.8 show that EPR performs better this case. Second, polynomials require manual tuning of degree, while P- P-POTS requires no tuning at all. Third, they lack interpretability and ignore known structure.
Drift of pl).. While p(l) reduces variance, it may become outdated as model evolves during training- phenomenon refer to the drift of p(t). In practice, P-POTS only fits p² once before training and then replaces the uniform sampler. This incurs negligible overhead, but leaves the sampler static and therefore potentially exposed drift. Empirically, however, the drift tends to slow, since our experiments show strong gains even without correction; see $4 An adaptive strategy, such periodically re-estimating p(t) to re-estimate p(t), could address this. Given the strong empirical performance of P-POTS, we leave such validation to future work.
3.1.2 MIRROR: VARIANCE REDUCTION WITH MIRRORED MASKS
MIRROR (Variance reduction with MIRRORed masks) targets A by creating two masked samples complementarily from the same 2n and f. With the notation introduced in $ it works follows:
5

---

Published as a conference paper at ICLR 2026
1. Sample clean data To from the data distribution.
2. Sample masking rate t ~ 21 1].
3. After sampling zo and f, for each eligible token i, sample, and generate two noisy samples:
zy by masking zo (i) if U; <t
.े by masking 2/ (1) U, > -
4. Compute I - E (1) [MASK]] log (zo xx]) for E {1,2)
5. Average == 2 (41 12) across batch for backpropagation.
Compared to standard training that averages 4 from a single x/ across the batch, MIRROR uses I, which remains unbiased in estimating the the training objective (Eq. [ It uses the same number of d0 samples, thereby avoiding overfitting. While it doubles training cost due to an extra forward pass, it is particularly effective long-response datasets and shows strong synergy with P-POTS.
Intuition behind MIRROR. MIRROR can be understood as a hedging strategy between two assets. To reduce (A) one needs a loss estimate from z with lower variance that arises solely from masking patterns. Suppose the token sequence is masked at n positions, if the easier tokens are masked (e.g., uninformative prefix tokens) versus the harder ones (e.g., answer tokens), the resulting losses can differ significantly. To reduce this variance, MIRROR samples two noised inputs as described above. Due to the "complementary" masking pattern design, the two noised samples tend to yield negatively correlated losses. Consequently, regardless of whether x correspond to an easier or harder case, =? provides the complementary one, and the averaged loss from them still remains a reliable estimate.
Reduced by at Least Half. Given my and t. the only randomness comes from Zt. Since and 22 are identically distributed, define o² = Var(I) Var(l2) and p = Corr(1,, 12) S Then
Var(I)- = Var 5(4 + +12)) = (Var((1) + Var((2) - Cov((1,12)) 4 20 (1 (1+
o² 2 (1+
The effectiveness of variance reduction depends on p = Corr(4,, (2). Since E 1-,, p Var(!) < Var(4,,, so MIRROR is never worse than standard training. Because 2, and x are constructed complementarily, P tends be negative, which already reduces AA by at least half. For t < 0.5, the two masks have no overlap, producing even stronger negative correlation between 1 and l and thus greater variance reduction. This explains its synergy with P-POTS (discussed later in $3.1.3..
Comparison to MultiSample-k. We refer to as MultiSample-k the method that estimates lo æo, 8) by drawing independent masked samples x! (1)
,..., 27 from the same masking strategy U, < and averaging = E/ Io. t). Our comparison focuses on:
Negative Correlation: MultiSample-2 uses independent masks, so 2ov(4,, 0. MIRROR introduces natural complementarity. leading to Cov(l1, (2) < 0 and better variance reduction.
Token Coverage: Only masked tokens contribute to gradient updates. Intuitively, masking more tokens reduces pattern noise, as each token carries less individual weight when more contribute. Standard training masks each token with probability I, giving expected coverage tP. MIRROR uses complementary masks, which together maximize two-view coverage as Pr(i E Mi U Pr(U, < t) + Pr(U, > 1--t) t) = min (1, 2t). By contrast, MultiSample2 does not have double coverage: Pr(i Mi U M2) = 1 Pr(i o M1 O M2) = 1 (1 t)² 2t t < 2t.
3.1.3 SYNERGY OF P-POTS AND MIRROR
P-POTS and MIRROR are compatible because they rely on non-interacting assumptions - accurate modeling of ptt) and negative correlation between (ar), =²?). However, why their combination yields synergies beyond simple additivity (1 > 2) worth further investigation.
To answer this question, we return to Eq. the variance in estimating the training objective (gp
in Eq. 1.. With p(t) & (t)
it becomes (Jo' V + ut) -(o g*()) dt
Jo
Since this value no longer depends on p(t), for fixed (t) and u(t), it cannot be further reduced by altering p(t). However, it can still be lowered by changing (t) and vol) during training.

---

Published as a conference paper at ICLR 2026
Suppose u(t) 0 and Jg v(t) dt = V is fixed. The variance depends how close So Vgg(()++ v(t dt and Sn g(t) dt are. To bring them closer, v(t) should be concentrated where g(t) large, as g²)(t) + v(t) is less sensitive v(t) in this regime. MIRROR helps here by sup- pressing e(t) most strongly at moderate t, leaving v(t) relatively larger where g(t) is high. Thus, MIRROR naturally shifts u(t) towards the direction favored by P-POTS. Together they reinforce each other: MIRROR clears moderate-f regions, while P-POTS emphasizes low- and high-t. This synergy remarkable, as naive method combinations rarely deliver such gains (A.6).
3.2
OTHER TECHNIQUES
We introduce several other variance-reduction techniques baseline methods in our experiments.
3.2.1
ISAD: IMPORTANCE SAMPLING ON ANSWER DELIMITERS
ISAD shift masking toward answer delimiters, targeting A, by the sampling distribution:
if token $ rare ids, (t) min (1, t A), if token E rare ids,
with A (0, 1). To keep the estimator unbiased, token-level losses are reweighted by 1/gg(t), The intuition is to guide the model on when to output answer tokens for better downstream performance.
3.2.2 SYRM: SYNTAX-AND -EDAN-ONND-- MASK
SyRM designed for structured data such as HTML tables, source code, or graphs. We partition tokens into two groups: (i) response tokens R, and (ii) syntax tokens C' in the prompt. For example, an HTML table, C includes tags such as <table>, <tr>, <t>>>,th>, and <td>. SyRM modifies the eligibility set for masking RUC, so that both response tokens and syntax tokens may be masked.
Rationale, Assuming (i) tokens within each group are homogeneous in first and second moments of token-level losses, (ïi syntax tokens predicted more stably than response tokens, and (iii) reasoning errors couple response tokens more strongly, we show that SyRM reduces moasking-pattenn noise A at the cost of small, bounded bias (optimum shift). Overall, the mean-squared error of estimating training loss can still decrease. Intuitively, by treating eesy-to-predict syntax tokens as auxiliary targets, SyRM stabilizes training and mitigates mask-pattern noise, while still preserving the primary learning signal on response tokens (see A.5.4 for the rigorous proof).
3.2.3 STRATS: STRATIFIED T-SAMPLING
Instead of drawing i.i.d. t, ~ O( 1] for batch size n, partition [0, 1] into l equal strata [j/k, (j + 1)/k] j O,..., 1-1 and draw [+] values of 4 per stratum, with leftovers assigned randomly.
This approach was first introduced in [ahoo et al. (2024) with k = n. In this work, we prove that stratified sampling yields unbiased estimator and reduces B via inter-stratum variance. Further- more, we recommend & = [n] instead of k = n as general strategy (see Appendix A.5.5). Although StraTS uses stratified sampling for the t-sampler instead of p" pt() ($3.1.1) and is not op- timal for minimizing Azo + Bro. it still outperforms baselines and is easy to implement.
3.2.4 EMA: BIN- WISE EMA CONTROL VARIATE
To reduce B we reduce its variance arising from heteroscedastic losses across masking rates by maintaining exponential moving average (EMA) of losses within bins of t [0. 1]. Each bin tracks baseline that is updated online, yielding a piccewise-constant approximation of the expected loss against t. By using these bin-wise estimates control variates, cancel seseline--rlatte noise more effectively than with fixed scaling, while preserving unbiasedness gradient estimates.
This approach requires choosing the number of bins and EMA rate 77, which balances stability and adaptivity. From our ASS-minniizatton analysis in/6.5.6, we derive m batch size per GPU s 0.1 (train data size), which serves practical rule for selecting and n).

---

Published as a conference paper at ICLR 2026
4
EXPERIMENTS
(a)
(d)
(c)
(I)
(g)
(h)
Figure 3: Images generated by RMMDA--B-Mixch trained with P-POTS+MIRROR (top row) and the stan- dard method (bottom row) under the same seed. The columns correspond to the prompts: 1) A woman with blonde hair is wearing colorful crochet headband and a black top with a floral patterned shawl. She is holding a toy sword with a green handle and yellow ball on top; and small porch stands
rustic wooden farmhouse with weathered roof field dry grass. The sky filled with fluffy white clouds, flying in the distance. The landscape is devoid of any other structures or signs of life, giving the scene a serene and isolated feel; 3) A car is on fire under bridge with smoke billowing out. The fire is intense, with tlames visible the car's underside. The bridge has metal railing and is located near forested area. There a street lamp above bridge, and scene appears at night or during dusk; 4) close-up up grey fringed purse with a tassel detail on the front. The purse is placed on a surface with a soft, blurred background; more case studies can be found in A.11
single bird
Limited by resources, and given that current MDMs perform comparably strong ARM baselines Nie et al. 2025 Yang et al.) 2025), we focus our validation on supervised fine-tuning, though our methods naturally extend to pretraining. We evaluate methods using the setup below:
Setting. We train each model three times independently using seeds 42, 731 and 20231, holding model, data and all hyper-parameters fixed, and evaluate the trained models by greedy decoding. This setting ensures difference across runs arise solely from training, not inference.
We evaluate GSM8K (Cobbe et al. (222I) and HiTab (Cheng et al. 2022) as well as one image genera- tion dataset, (ext-to-image-M1 -2 (Hate 2024). OpenScience a multi-domain, knowledge-intensiv dataset spanning humanities, law, and STEM; GSM8K targets mathematical reasoning; and HiTab focuses on tabular question answering. Together with the multimodal sert-to-image-222 2M benchmark, diverse data modalities and domains are covered. Details of data processing and training/inference hyper-parameters are provided in Appendix A.7 We report pass@1 accuracy OpenScience, GSM8K and HiTab, and CLIP score for text-to-image-2M.
Tasks.
three textual datasets OpenScience (NVIDIA Corporation 2025
Models. We evaluate two MDMs total. On textual datasets, we evaluate one MDM, LLaDA-88 Instruct (Nie al. (025), with three ARMs, Diven..--BBIInsteen (Team 2024), Qwen3-8B (Tem)) 2025 and Llama3.1 1...3.11888 SB-Instruct (Meta Al 2024b). For (ext-to-image-2MM we evaluate one multi- modal MDM, MMaDA-800 BAEaDAA-B-MAA -8--M Yang et al. 2025).
We compare our methods to two baselines: (i) standard training ($3) and (ii) the clipped noise schedule of Arriola et al. (2025). which has been reported to achieve the best performance perplexity across most standard schedules (Chang et al., 2022), including linear, logarithmic, square root, square, and cosine. We therefore adopt it as representative baseline.
Baselines.
8

---

Published as a conference paper at ICLR 2026
Table 1: Per-seed performance comparison of MDMs and ARMs. The best performance within MDMs training methods is bolded. Training repeated under three seeds, model, data and all hyper- perr--aameeerrs fixed, and we evaluate trained models by greedy decoding, so difference across runs are solely from training rather than inference. Avg is the mean over available runs.
OpenScience
GSM8K
HiTab
Method (Seed)
42
731
20231
Avg
42
731
20231
Avg
42
731
20231
Avg
LLaDA- 8B IB-trutrttt
P-POTS+Mirror
50.80
53.80
53.00
52.53
62.02
58.61
60.96
60.53
66.02
66.67
68.62
67.10
P-POTS
45.34
47.60
47.47
46.80
59.59
59.36
56.79
58.58
60.56
64.7
58.80
61.37
MIRROR
45.63
47.25
46.25
46.38
56.56
51.40
53.15
53.70
64.72
65.04
63.68
64.48
ISAD
45.40
45.40
55.27
55.27
64.39
64.39
SyRM
59.79
59.79
StraTS
44.69
44.69
55.57
55.57
61.86
61.86
EMA
47.08
47.08
52.39
52.39
61,66
61.66
Clipped
47.90
49.20
50.90
49.33
55.57
54.60
58.38
56.18
58.90
60.80
62.20
60.63
Standard
43.57
46.50
46.50
45.52
50.64
53.22
53.68
52.51
52.96
60.43
62.64
58.68
Auto-Regres Regressive Large Language Models
Dwwe2.55-B-Innttrun
47.85
46.30
45.10
46.42
56.20
54.66
53.60
54.82
46.55
45.74
45.20
45.83
Qwen3-8B-Base
53.05
51.70
50.10
51.62
72.40
70.58
69.80
70.93
49.70
48.84
48.20)
48.91
Llama3.1-8 ----BIsssuutt
44.95
47.07
47.76
46.59
47.60
46.40
49.10
47.70
65.02
66.28
64.88
65.39
Table 2: Per-seed performance of MMaDA-8B-
MB- on text-to-image-2N We report CLIP scores under three training seeds with identical model, data, and papameparameters. Avg is the mean CLIP over runs.
text-to-image-2M
Method (Seed)
Avg
42
731
20231
P-POTS+Mirror
35.27
34.10
34.39
34.59
Standard
28.61
34.28
31.19
31.36
P-POTS
31.51
30.02
33.41
31.65
Mirror
34.88
33.77
33.57
34.07
4.1 KEY INSIGHTS
MDMs exhibit uniquely high training variance, which makes gradient updates noisy and optimiza- tion less stable. As result, even pretrained MDMs are shown achieve comparable performance with strong ARMs, after task-specific training, they often under-perform by a wide margin. This in- stability is evident in Table under standard MDM training HiTab, accuracy varies substantially across random seeds (about 52% ~ 63%), and the mean performance across datasets frequently re- mains below that of ARMs. Our method targets this root cause by reducing MDM training variance. Figure4ssoww that P-POTS+MIRROR (green) converges to lower final loss with smoother dynam- ies than the Standard method (red), indicating more stable training. This improved stability shows consistent improvement and closes the training gap effectively: on all three textual benchmarks, even the best Standard/Clipped run remains below the worst run our method (OpenScience: best Standard 46.50 vs. worst P-POTS+MIRRON 66.02). On text-to-image-2M4 the two methods are essentially comparable (best Standard 34.28 vs. worst ours 10). This highlights the importance of variance reduction in delivering more consis- tent and reliable performance across runs. Overall, MDMs trained with our method can match, and in some cases surpass, strong ARMs, while run-to-run variability drops to near ARM levels.
50.80; GSM8K: 53.68 vs. 58.61; HiTab: 62.64
Moreover, ablation studies in A.8 confirm synergy between P-POTS and MIRROR. On GSM8K and OpenScience, compared to standard training, the accuracy gains from P.POTS+MIRROR (8.02% and 7.01 respectively) surpass the sum of the individual gains from P-POTS and MIRROR ((6.07+ 1.19)% 7.26% and (1.28 0.86)% 2.14% ), while HiTab they are nearly additive (8.42%

---

Published as a conference paper at ICLR 2026
OpenScience
text-to-image-2M
6
start: 8.474
train/loss
Start: 1.453
Standard
train/loss 8
Itart: 8.313
783
end: 1. 1.352
MURT end:
2
91
0
start:
1.428
lend: 1:134
100
300
400
100 200
300
400
500
600
700
Step
Step
Figure 4: Training loss comparison between P-POTS+MIRROR (green) and Standard (red) on OpenScience (left) and coxt-t--image-22 2M (right). The annotations indicate the average loss at the beginning and end of training over the first and last steps. Overall, P-POTS+MIRROR achieves more stable loss trajectories with consistently lower end losses compared to the Standard baseline.
and (2.69 + 5.80)% % = 8.49% ). We emphasize that this is nontrivial, because as discussed in A.6 naive combinations of methods do not typically yield such improvements.
MIRROR is especially effective on long-response datasets due to its design. As shown in Tables and 2 its gains on OpenScience, HiTab and text-to-image-2MM are substantially higher than on GSM8K. This aligns with its design: MIRROR targets masking pattern noise A. which becomes more severe with longer responses due to the greater number of possible masking patterns.
MDMs achieve comparable performance ARMs with much shorter outputs. In our experiments, MDMs often generate responses as short as 5% of ARM outputs and consistently shorter than train- ing labels (sec exact values in A.7.I. Truncating ARMs to the same length severcly hurts their per- formance because outputs are often cut off mid-reasoning. One possible explanation that MDMs decode by repeatedly unmasking block of [MASK] tokens whose length is predetermined by the user-passed
giving the model an implicit sense of how many tokens it is al- lowed to generate. In contrast, ARMs generate token by token without an explicit stopping criterion until hitting EOS, which can lead to verbosity or redundancy.
On text-to-image-MM, the drift in p is more severe; see A.8.2 We attribute this larger loss scales (and thus gradients) that move f faster. Although EPR targets a U-shape, we use it as a stable compromise because P-POTS fits p(t) only once pre-training: empirically, EPR is more robust than fitting the initial S-shape, so we recommend it as a default. In principle, periodically refitting p(t) (e.g., at each eval) could further improve P-POTS and P-POTS+MIRROR.
4.2 CASE STUDY
In Figure 3
standard training often produces blurry images or misses fine details. We suspect two causes. First, rare detail-critical tokens are masked too infrequently under uniform p(f). so the model under-trains on details. Consequently, features such as the "crochet headband" and "toy sword" in case (b) are missing or distorted. Second, cross-entropy on masked image tokens can show averaging effect. When multiple outputs are possible, the model may learn to predict their mean; at inference with low temperature, argmax decoding turns this into unstable choices, leading blur. This may explain case (d), where fire and smoke looks unrealistic. By contrast, P-POTS+MIRROR may address these problems. P-POTS emphasizes harder regions, while MIRROR averages two negatively correlated samples and reduces the tendency to hedge. Together, they could improve the gradient signal-to-noke toise ratio, giving sharper textures and more faithful objects, e.g. correct headband/sword in (a) and more realistic flames in (c).
CONCLUSION
This work derives systematic decomposition of MDM training variance as a fundamental frame- work for variance reduction, and proposes six corresponding techniques. Among them, P-POTS provably Pareto-optimal for minimizing loss variance, while MIRROR is particularly beneficial on long-response datasets; together, they yield strong synergy. Our methods reduce both run-to-run variability and within-run suboptimality, making one training run sufficient for reliable performance. Looking forward, expect our decomposition and methods to serve as building blocks for scal- able MDM training, and to inspire future work toward a unified framework for variance control in masked diffusion models.
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Marianne Arriola, Aaron Gokaslan, Justin T Chiu, Zhihan Yang, Zhixuan Qi, Jiaqi Han, Sub- ham Sekhar Sahoo, and Volodymyr Kuleshov. Block diffusion: Interpolating between autore- gressive and diffusion language models. arXiv preprint urXi:2253300957 2025.
Jacob Austin, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne Van Den Berg. Structured denoising diffusion models in discrete state-spaces. Advances in neural information processing systems, 34:1 191117988199999 1-19993, 2021.
Lukas Berglund, Meg Tong, Max Kaufmann, Mikita Balesni, Asa Cooper Stickland, Tomasz Kor- bak, and Owain Evans. The reversal curse: Lims trained on" is b" fail learn" is a". arXiv preprint arXiv:2309. 1228888,09.1228 2023.
Huiwen Chang, Han Zhang, Lu Jiang, Ce Liu, and William I. Freeman. Maskgit: Masked generative image transformer. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 11315-11325, 2022.
Zhoujun Cheng. Haoyu Dong, Zhiruo Wang. Ran Jia, Jiaqi Guo, Yan Gao, Shi Han, Jian-Guang Lou, and Dongmei Zhang. Hitab: A hierarchical table dataset for question answering and natural language generation. Proceedings of the 60th Annual Meeting of the Association for Compu- tational Linguistics (Volume 1: Long Papers), pp. 1094-1110, 2022.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichir Nakano, et al. Training verifiers to solve math word problems. arXiv preprint arXiv:21 14168, 2021.
Shansan Gong, Ruixiang Zhang, Huangjie Zheng, Jiatao Gu, Navdeep Jaitly, Lingpeng Kong, and Yizhe Zhang. Diffucoder: Understanding and improving masked diffusion models for code gen- eration. arXiv preprint arXiv:2506 20639, 2025.
Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, et al. The llama 3 herd of models. arXiv preprint arXiv:2407.21788 2024.
Daya Guo, Dejian Yang. Haowei Zhang, Junxiao Song, Peiyi Wang, Qihao Zhu, Runxin Xu, Ruoyu Zhang, Shirong Ma, Xiao Bi, et al. Deepseek-rl: Incentivizing reasoning capability in llms via reinforcement learning. arXiv preprint MwXiv:2501.12944 2025.
Jacky Hate. text-to-image-2m (revision e64fca4), 2024. URL https://huggingface.eu
Paul Jeha, Will Grathwohl, Michael Riis Andersen, Carl Henrik Ek, and Jes Frellsen. Variance reduction of diffusion model's gradients with taylor approximation- control variate. arXiv preprint arXiv:2408. 12270, 2024.
Herman Kahn and Theodore E. Marshall. Methods of reducing sample variance in monte carlo computations. Journal of the Operations Research Society of America, 1(5):263-276 1953.
Myunsoo Kim, Donghyeon Ki, Seong-Woone Shim, and Byung-Jun Lee. Adaptive non-uniform timestep sampling for diffusion model training, arXiv preprint arXiv:2411.09998, 2024.
Joseph-Louis Lagrange. Mécanique Analytique. Courcier, 1811.
Aaron Lou, Chenlin Meng, and Stefano Ermon. Discrete diffusion modeling by estimating the ratios of the data distribution. arXiv preprint XrXiv:23 10. 16834, 2023.
Chenlin Meng, Yang Song, Wenzhe Li, and Stefano Ermon. Estimating high order gradients of the data distribution by denoising. Advances in Neural Information Processing Systems, 34:25359- 25369, 2021.
Meta AI. Introducing Ilama 3.1: Our most capable models to date. Meta AI Blog, 2024a. URL https: //ai .mtta. com/ haam/blogr wl meta 11ama-3-1/

---

Published as a conference paper at ICLR 2026
Meta AL 1at-aa..---bbbnittree HuggingFace model card. 2024b.
URL https://
huggingface http://aac..comtttaaaaaaae- wwt.c////mtttai Llama- 1aammaa3.1-- 3.1-8B- Instruct
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. arXiv preprint arXiv:2502.09992, 2025.
NVIDIA Corporation. Openscience: A multi-domain synthetic dataset for seneral-purpost reason- ing. httpps://huggingface co dstetttssss//////////aaa.aaassst /nvii a OpenScience 2025, Version 1; licensed under CC BY 4.0; data creation date May 15 2025.
Peter Peduzzi, John Concato, Elizabeth Kemper, Theodore R. Holford, and Alvan R. Feinstein. A simulation study of the number of events per variable in logistic regression analysis. Journal of Clinical Epidemiology, 49(12):1373-1333 1379, 1996. ISSN 0895-4356. doi: https://oi.org// 1066800895-436666600006333 URL attps://viit wttp://wwww..ccccciiere.t... wwtp://ttt.ococcecnneeeeeeeeeee.
Herbert Robbins and Sutton Monro. A stochastic approximation method. The Annals of Mathemat-
ical Statistics, 22(3):4004477, 1951.
Walter Rudin. Real and Complex Analysis. McGraw-Hill 3 edition, 1987.
Subham Sahoo, Marianne Arriola, Yair Schiff, Aaron Gokaslan, Edgar Marroquin, Justin Chiu, Alexander Rush, and Volodymyr Kuleshov. Simple and effective masked diffusion language models. Advances in Neural Information Processing Systems, 37: 37:130133-13018,, 2024.
John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, and Picter Abbeel. High- dimensional continuous control using generalized advantage estimation. arXiv preprint arXiv:1506.024433 024388 2015.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv: (oXiv:1707.064777 2017.
Zhihong Shao, Peiyi Wang. Qihao Zhu, Runxin Xu, Junxiao Song, Xiao Bi. Haowei Zhang, Mingchuan Zhang, YK Li, Yang Wu, et al. Deepseekmath: Pushing the limits of mathemati- reasoning in open language models. arXiv preprint arXiv:2402.03300, 2024.
Qwen Team. Qwen2.5: A party of foundation models, September 2024. URL URL httphhttpss///eettt github. wttp:////lllg/goo qwen2.5/
Qwen Team. Qwen3 technical report, 2025. URL
Yilun Xu, Shangyuan Tong, and Tommi Jaakkola. Stable target field for reduced variance score estimation in diffusion models. arXiv preprint rxXiv:2302.00670. 2023.
An Yang, Baosong Yang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Zhou, Chengpeng Li, Chengyuan Li, Dayiheng Liu, Fei Huang, Guanting Dong. Haoran Wei, Huan Lin, Jialong Tang, Jialin Wang. Jian Yang, Jianhong Tu, Jianwei Zhang. Jianxin Ma, Jin Xu, Jingren Zhou, Jinze Bai, Jinzheng He. Junyang Lin, Kai Dang, Keming Lu, Keqin Chen, Kexin Yang, Mei Li, Mingfeng Xue, Na Ni, Pei Zhang, Peng Wang, Ru Peng, Rui Men, Ruize Gao, Runji Lin, Shijie Wang. Shuai Bai, Sinan Tan, Tianhang Zhu, Tianhao Li, Tianyu Liu, Wenbin Ge, Xiaodong Deng, Xiaohuan Zhou, Xingzhang Ren, Xinyu Zhang, Xipin Wei, Xuancheng Ren, Yang Fan, Yang Yao, Yichang Zhang, Yu Wan, Yunfei Chu, Yuqiong Liu, Zeyu Cui. Zhenru Zhang, and Zhihao Fan. Qwen2 2024.
technical report. arXiv preprint arXiv:2407.106:
Ling Yang, Ye Tian, Bowen Li, Xinchen Zhang, Ke Shen, Yunhai Tong, and Mengdi Wang. Mmada: Multimodal large diffusion language models. arXiv preprint arXiv:2505. 15809, 2025.
Jiacheng Ye, Zhihui Xie, Lin Zheng, Jiahui Gao, Zirui Wu, Xin Jiang. Zhenguo Li, and Lingpeng Kong. Dream 7b: Diffusion large language models. arXiv preprint arXiv:2508. /5487, 2025.
Runpeng Yu, Xinyin Ma, and Xinchao Wang. Dimple: Discrete diffusion multimodal large language model with parallel decoding. arXiv preprint arXiv: 16990, 2025.
12

---

Published as a conference paper at ICLR 2026
Fengqi Zhu, Rongzhen Wang, Shen Nie, Xiaolu Zhang, Chunwei Wu, Jun Hu, Jun Zhou, Jianfei Chen, Yankai Lin, Ji-Rong Wen, et al. Llada 1.5: Variance-reduced preference optimization for large language diffusion models, 2025.
13

---

Published as a conference paper at ICLR 2026
A
APPENDIX
CONTENTS
Statement
14
A.2
Preliminaries
15
A.2.1
Stochastie Gradient Descent on Convex Objectives
15
2.2.
Discrete Diffusion Models
15
A.2.3
Proximal Policy Optimization
17
A.2.4
Group Relative Policy Optimization
17
A.3
Training Variance Decomposition
18
4
Algorithms
18
A.4.1
Bin- Wise EMA
18
A.5
Proofs
19
A.5.1
Training Variance Decomposition
19
A.5.2
P-POTS: Pareto-Optimal Property
20
A.5.3
Importance Sampling
21
A.5,4
SyRM
22
A.5.5
Stratified Sampler
26
A.5.6
Bin-Wise EMA: Variance Reduction & Hyperparameter Selection
27
A.6
Method Combination
28
A.6.1
Stratified Sampling VS EMA
29
A.6.2
ISAD VS SyRM
30
A.7
Experimental Details
30
A.7.1
Hyperparameter Settings
30
A.7.2
How Training Variance is Estimated
31
A.8
Complete Experimental Results
32
A.8.1
Per-Run Experimental Results LLaDA
32
A.8.2
Per-Run Experimental Results on MMaDA
33
A.9
Extension Reinforcement Learning
33
A. 10 Future Work
34
A.Il Case Study
34
A.11.1 Case Study: Image Generation
34
A.1 STATEMENT
Reproducibility. All models and datasets used in this work are publicly available, as listed in Details of dataset splits and hyperparameter configurations are provided in Appendix [A.7.] while the evaluation protocols for inference results are described in Appendix 4 All theoretical proofs and assumptions are presented in Appendix A.5
LLM Usgae. We used large language models (LLMs) as a general-purpost assist tool three ways: (i) polishing the writing for grammar and clarity, (ii) helping identify potentially relevant
14

---

Published as a conference paper at ICLR 2026
related work (with all citations manually verified by the authors), and (iii) providing code assistance for implementing ideas developed by the authors. The research ideas, methodology, experiments, and final paper content remain the responsibility of the authors.
A.2 PRELIMINARIES
A.2.1 STOCHASTIC GRADIENT DESCENT ON CONVEX OBJECTIVES
We consider the optimization problem
min f(3), BEC
(6)
where C' C IR is closed convex set and C R is convex. In machine learning, typically represents the population risk, i.e., the expected loss with respect to the true but unknown data distribution. A common example is the expected cross-entropy loss of the model under parameter relative to the true data distribution.
Gradient-based optimization methods generally rely on estimates of the gradient of the expected loss. Using larger batch size reduces the variance of this estimate but increases the computational cost per iteration. Stochastic Gradient Descent (SGD) strikes balance by sampling batch of fixed size at random and using the average loss gradient over this batch an unbiased estimator of the population gradient. Convergence of SGD can be guaranteed in principle, provided that certain standard assumptions are satisfied:
Theorem 1 Suppose 3 is minimiser of f over closed convex set C C RP. Suppose
supsec || < R < do and
sup E BEC
sup (Eaf((;
1
11. <0.
Then if 77, = ())[(L((L))) the output / of the stochastic gradient descent algorithm satisfies
2LR
- I
VE
From a statistic point of view, the batch mean is merely an estimator of the population gradient, and thus inevitably subject estimation variance. In MDM training, where the expected loss is subject to three sources of variance, this issue becomes particularly severe. is therefore important to develop estimators that are more variance-efficient.
Note.
A.2.2 DISCRETE DIFFUSION MODELS
Let Xo E [1....,K)) denote initial discrete data vector, such length token sequence, where each element zá takes values from finite vocabulary of size K. A discrete diffusion model defines forward corruption process that transforms Xo into sequence of progressively noisier variables X1, Xy over 7'' timesteps.
At each timestep I.... ; T. let (k) represents the kth token of x; for {1, n) then each token x (k) is corrupted into xzkk) using predefined transition matrix Q₂ E k Formally, the transition probabilities are defined as:
(<x((k) - Ik- = = (Qrk
where Qt is a stochastic matrix with each column sum 1. Let v(x) E RK be one-hot vector with 1 at position For simplicity, abbreviate token position information and use x, for x, (k)
for any E (I,..., KJ. Then we can write the transition equation as
q( xد-1) = v (x.)).vvaa)))))
(8)
We then marginalize intermediate steps and compute the transition probability from Zo 24:
15

---

Published as a conference paper at ICLR 2026
g( mo) = (=)))))))))== 1....o))
(9)
The posterior distribution conditioned on is then:
g(<t-1 (14, 1,0 ( (zo)
(xv (1-
v vT(
(10)
For the reverse (denoising) process, a neural network (0 vx) = Categorical(#1-1 Po (1()))
trained to undo this corruption, where Pa(xt) I denotes the logits produced by the model.
(o xx) trained to estimate this posterior. Since this posterior defined on predefined forward processes only (the transition matrix Qr that controls the noisy level each time step t), the design of transition matrix crucial in the success of discrete diffusion models. A well-designed noise schedule ensures the model encounters smooth curriculum of denoising tasks, leading better final performance.
Early works propose uniform noise design:
1 ax
Qc a + -1-
This design means at each time step, a token has probability a + B, to remain and 3 probability to transfer other other K - 1 eigenvalues equal to Thus, the signal of zn decays like ag...Otg If this goes to 0 too quickly, you lose all information.
tokens. This matrix has one eigenvalue
1(with eigenvector 1) and all
Among the general discrete diffusion model framework, masked diffusion models consistently achieves the best performance (Austin et al.] (O0T) et 2003). masked diffusion, a specific design is employed in which the K category serves an absorbing state, commonly correspond- ing to special token such as [MASK]. ]. Under this configuration, the transition matrix Q: is defined as:
Qq
0
0
1
Q2
1 at
Q
at
0t
where az E [0, denotes the masking ratio timestep This design ensures that each token is
independently replaced with the [MASK] token with probability 1 1-a,, and retained with probability dee Note the [MASK] token is an absorbing state. By isolating noise into dedicated [MASK] token, the model preserves clean, low-dimensional signal subspace, which it can focus on when reconstructing X,-1-
The training objective maximizes variational lower bound (ELBO) logpe(zo), introduced in Eq. T Importantly, Sahoo et al. (2024) has shown that Eq T
1. is invariant the specific choice of the noise schedule ag and
2. serves as an upper bound of
-E -xpan()) [x)),
thus enabling principled and simplified training framework for masked diffusion models.
Restricting attention to interpolating masked diffusion models-i.ed those whose forward process q interpolates between clean data E V and target distribution Cat(- :)) Sahoo al. (2024) derived principled training objective Eq. [ that is invariant to noisy schedules. Recent models such L.LaDA and MMaDA base their training this training objective and adopt log-linear noise schedule defined by:
0₁ exp(log(1 - t)) 1,tt,
which means that at each timestep f, every token independently replaced with a [MASK] token with probability t.
16

---

Published as a conference paper at ICLR 2026
A.2.3 PROXIMAL POLICY OPTIMIZATION
Proximal Policy Optimization (PPO) can be derived as a first-order approximation the Trust Re- gion Policy Optimization (TRPO) problem by replacing the hard KL-constraint with KL-penalty. We begin with the constrained optimization
max
Es (e) Ad]
s.t. Ey DKL Tem ( st) mo(. | S 6,
(11)
where
(0) To(a: (t)) Teas (at st)
and Åg is an estimator of the advantage at time
Introducing a Lagrange multiplier 3 for the KL-constraint yields the penalized surrogate objective
KLPEN (0) Es (r(()) A. DRL (Teee (+1 se) To [ss]]]]]
(12)
In practice, PPO proceeds by:
1. Collect trajectories {(5t, oe, under the old policy Team
Compute advantage estimates A, (e.g., via GAE).
3. For K epochs, perform minibatch gradient ascent on I. KLPEN (0):
0 + n V KLPEN (0).
Optionally adjust keep the KL close to target
a4 a+BB B,
D 1.50.
B +
a B,
DDKL < 8/1.5,
B,
otherwise,
where Dra. E, [RL (T [[]]]]
5. Set Oata +- and repeat.
A.2.4 GROUP RELATIVE POLICY OPTIMIZATION
Group Relative Policy Optimization (GRPO) extends PPO to settings where we have a collection of G groups (or tasks), each potentially requiring its own policy behavior while sharing parameters. Let G = (1...., G}} denote the set of groups, with weightings (wg bggg such that Σg Wg = 1. For each group define
(a St. :) Te (ae st,g)' and the group-specific KL divergence
(0)
Ag = advantage estimate for group g at time t,
DKa (0) Erng DKL Tome (- sts g) TT(- St =)]g
Introducing a common penalty coefficient 3. the GRPO surrogate objective is L CRPO
Σ
Wi Beng (0) Aq DKL (Tam (: se.g) ma(- st. 9.]
(13)
gEg
In some cases, one may also include an inter-group regularizer to encourage policies for different groups to remain similar:
Σ With EOND DKL (Tp(- s, || ne(- ss,
g.heg a
where 7 2 0 and { are pairwise weights.
The GRPO algorithm proceeds analogously to PPO:
17

---

Published as a conference paper at ICLR 2026
1. For each g E 9, collect trajectories under Tewa ( (1)).
Compute group-specific advantages {Ág].
3. For K epochs, perform minibatch gradient ascent L (G)
0 - 0 - (RPO
4. Optionally adjust 3 (and 7, if used) to keep each D (and inter-group KLs) near desired targets.
5. Update Oold 0 and repeat.
A.3 TRAINING VARIANCE DECOMPOSITION
The proof of Eq. 2 is deferred to A.5.I Here we clarify why this decomposition is well-motivated and how it supports the P-POTS derivation 3.1.1.
Effect of Alternative decompositions. Our decomposition is obtained by iteratively conditioning on the random variables (1,,,,, different conditioning orders give different decompositions. The order adopted in this paper + t - Ty is the only interpretable one, because it aligns with the MDM training procedure and matches ARMs' variance decomposition. Importantly, the resulting form of P-POTS in $ $3.111 invariant to the conditioning order. This is because the optimal sampler p" (t) depends only on E f). not on how the total variance is partitioned.
Loss Variance vs. Gradient Variance. While gradient variance is most directly tied to optimiza- tion stability, we focus on loss variance for three reasons. (1) it is is optimizer-agnostic: adaptive optimizers rescale gradients, making "raw" gradient variance dependent
the optimizer, whereas Var(I) directly improves the SNR of the Monte Carlo estimate of the training objective. (2) Gra- dients are nigh-dimensional (e.g., 8B-model has gradient of dimension 8B), so any notion of "gradient variance" requires scalarization (e.g.. g) Yet this not theoretically neutral: dif- ferent parameters may prefer different sampling distributions. Any scalarization therefore introduce extra design choices and breaks the clean optimality claim. In contrast, the loss is already a unified scalar with clear interpretation unambiguous. (3) A t-sampler can reliably control only the conditional variance of the loss: by first-order linearization, Var(g t) s Jr Var(1 J," + (terms from model sensitivity), where Jr is the Jacobian w.e.t parameters. This highlights that part of gradient variance comes from model sensitivity (the Jacobian), which cannot be affected by t-sampler. The only reliably controllable component is Var(I t) (holding fixed over short horizons).
the objective whose expectation
estimate, so its variance is
A.4 ALGORITTMS
A.4.1 BIN-WISE EMA
We partition the interval E [0, 1] into m equal bins and fix an EMA learning rate For each bin j. we maintain four EMA-tracked quantities:
(r) R [&A t bin j], [n(() v B(b,), MEH()) v B(Cb,), MиH(j) so E(G)
and bin baseline estimate by. initialized zero. On each training step:
Draw a batch { sample mask ratios 4; ~ U[00, 1]. run the forward diffusion, and compute per-sample losses &,.
Assign each 4 to bin via
jj-1 j m m
j f(1(1) if E
For each sample, form an online estimate the optimal control-variate coefficient
MIH(j) pz(i)uH(i Cj Mn (j) pn |n (j)² -
where is small damping constant.
18

---

Published as a conference paper at ICLR 2026
Compute the adjusted loss
(:th) = -gg -bg.
detaching bj so that gradients flow only through Ls and backpropagate on e
Update all four EMA statisties for bin j:
mL(s) - (1 - (1-n)p((j)>>) + nli
[2 < (1 11-17 -n))u(()) nbj,
MT,(() < (1 11-11 --1) + ,((bb),,
Мин (j) < (1 MH (5) + m(b))
Finally, update the baseline itself via
bj (1-1) npli.
This procedure produces a phecewise-constane approximation b of the true expected loss b(t) while dynamically choosing, per bin, the coefficient y that maximally cancels out baseline-related noise- yielding tighter variance reduction than a fixed scaling egg. e = Additionally, this "per-bin" EMA is updated per GPU batch and is analogous to standard moving-ave -veerage updates used elsewhere in deep learning.
A.5 PROOFS
A.5.1 TRAINING VARIANCE DECOMPOSITION
The law of total variance states that
Var(Y) B[Var(Y X)] Var(B[y X]).
(14)
Let Y 1(rrg, 21) where Ty ~ Pauca (., U10 1], ~ q(1r ro. t) are three random variables. We apply the law of total variance twice:
First Decomposition over 20
View To as the outer random variable. By the law of total variance,
Varzo, (Y) Exp [Vare,,, (Y [o)] + Varz, [,, (Y xx)]].
Define
go(Zo, t) = Ez, (0o xx) wo, t), ha(zo) = Et,,, (x Øol = E, [go t) xo].
Then the second term becomes
Var [ho [ho Varao (Brlgo(zo: t)])
C data variance.
Second decomposition inside Er,, [aar,,,, (Y zo)]
We apply the law of total variance again, this time conditioning on t: Varg,, w, (Y 20) E, [Var., (Y) æo, t)] + Var, [Ez, (Y Io. t)]
Note that Ez, (Y Fo, t) = go (ro, t), so:
Eup [Var, (Y xo)] = Ezz [TE, [Varz, (fo xo, ()] + E.D
1) xo))
Ez.. [Var,, (Ro Fo, t)] + Er, [Var (xo, t) ao)]
A mask pattern noise
mask rate noise
Combine Everything
Putting steps and together, we have:
Var (Co) E Exg. [Var (Lo 20. t)] Bzg [Var:(ge(an) (go [00]] Vara, (IE, (2 t)])
which is exactly the claimed decomposition Eq.
19

---

Published as a conference paper at ICLR 2026
A.5.2 P-POTS: PARETO-OPTIMAL PROPERTY
Consider the reweighted estimator
p(t) Lo lo t, xz), t p(t)..
As shown in Eq. 3 this estimator is unbiased for the training objective in Eq. detailed proof that
We now give
p(t) (t) u(t So (g) (s)) ds
the theoretically optimal distribution that minimizes A B B+C
Recall Var[X] = |=| - E[ Then
Varrzo
1 40( t, (, p(ej
Expanding yields
E 2, pP(e) (p()) (z t. B (p()) ta(=o,t,,r)
Assuming absolute convergence and finite expectations, Fubini's theorem allows us to exchange the order of integration, giving
1 So''e p(t) Sam-/be 1z(so, 24) dr,dro dt ( lo t, xz) du dr dt
2
Define
g(t) E Exg,er [o t),
u(t) Varzo. No t).
Substituting yields
L. (2(1) v0(1) dt (f(
p(t)
The second term is independent of p(t). Thus the optimization problem reduces to
min Jlp] (-' (²(() + v(()
dt, p(t) 0, S p(()
p(t)
Lagrangian method. Form the Lagrangian:
Clp, A1 -- 0² (
( v(t) p(1)
+ sp(() dl X,
where À enforces normalization. Taking the variation w.r.t. p(t) gives
g2(1) +(() p(t)2
= = p1(() V o (()+)((())
Normalization requires
vー= L' Vg o(s) ds.
Final result. Therefore, the unique minimizer is
v
('(t)= SO (
which achieves the Pareto-optimal balance by minimizing A
20

---

Published as a conference paper at ICLR 2026
A.5.3 IMPORTANCE SAMPLING
Part 1: Unbiasedness of the Importance-Sampling Estimator
Suppose we wish to estimate
I 1= [() da,
but drawing samples from p(x) directly is difficult. Instead, we draw 21...., ZN i.id. from proposal distribution g(z), which satisfies ((x) 0 whenever f(x)p((x) * 0, and construct the estimator
1 1 Σ w(x)) f(x;), iN
p(x) where w(x) g(x)
We check its expectation under the sampling law q:
[N[[N] Bng (xx))) (by linearity and identical distribution)
(X, (drop index)
/ p(a) q(x) f(z)g(x) dr / f(z)p(z)dz= I.
Hence, IN is an unbiased estimator of
Part Choice of q Minimize Variance
The variance of IN under q
Var, (frz(IN) = N Var, Varg(w(X)f(([)) = N (E)(((()))(())))-- 12).
Since is fixed, minimizing Var, (Ix) amounts to minimizing
Eg [)) = p(z)²
g(x) f(x)2 dx,
subject the constraint f (((x)) dr = Introduce Lagrange multiplier X, and consider the functional
C(g]= - / f()) dr + A ( ((00) dr (((o)doo))
q(x)
Taking the functional derivative with respect to ((x) and setting to zero:
f(x)
+ 0
q(x) x p(x) f(x)).
g(x)2
Thus, the optimal proposal distribution is
[f(x))((z)
q (x) [T((w) plu)
which, for nonnegative f. minimizes the variance to its theoretical minimum meiiuum-zero continuous-sampling limit.
zero in the ideal
21

---

Published as a conference paper at ICLR 2026
A.5.4 SYRM
SyRM is designed for structured data such HTML tables, souree code, or graphs. We categorize tokens into two groups: 1) response tokens R, and 2) syntax tokens C in the prompt. For example, in html table, the syntax tokens are <table>, <tr>, <th>, <td>, and similar tags. SyRM modifies the eligibility set for masking as follows:
cligible == response tokens U syntax tokens in prompt
Definition of Eligible Tokens
We define a token position as eligible if can be masked during the forward process. According to LLaDA((Nee et al. 2025), during pretraining, all tokens in Xg eligible. During supervised fine- tuning, where Xo is composed of concatenated prompt and response tokens, only response tokens eligible. Thus, in this stage, prompt tokens remain unmasked while response tokens are masked independently.
Under the following assumptions,
Tokens within each group are homogeneous in first and second moments of token-level losses.
Syntax tokens are stabler predicted than response tokens.
Reasoning errors couple response tokens more strongly than syntax tokens.
we show below that SyRM rigorously reduces the masking pattern noise A Although this intro- duces some bias (optimum shift), we show this bias is bounded, and the total MSE of estimating training loss is reduced due to significant variance reduction. Intuitively, by including syntax tokens in prompts that are both easy to predict and helpful for guiding the model's 's reasoning as an auxiliary loss, SyRM reduces AA) mask pattern noise.
Setup
Number of eligible positions: P PED N.
Mask variables: M; ~ Bernoulli(t) with 0 < f 1.
Per-token loss: random variables Y, satisfying
E = pr:
Var(V,) =
V)) = Pi (i#5).
Batch loss:
L
M.Y
PL 2m1
The M,'s are independent of Y,, and the Mj's are mutually independent.
Theorem
1
2
Var [1t]
4
Σ
Pij.
p21
12t
p2
ISKCjEP
Proof Let Din M,Y,, so L = L
Then
Var(S) Σ
Var(MYY) +2
Σ
Cov(MY,,,Y))
(15)
1=1
ISiSSSP
where
Var(M,Y,) - 2 = [R( 12 107 + 1)p2.
(16)
and
22

---

Published as a conference paper at ICLR 2026
M,Y) M -[MMY] ( +111)) Pppps P (17)
Combine Eqs. 15
16 and
17
P
Var(S)= = Σ
of
2 + 111 1) Σ
11 212
Σ
Pij-
(18)
t=1
KSKjSS
Convert Eq. 18 back to L = S.
Var[L 14
Var(S)
(P1)2
1
2
Σ
o
Pij.
p2t
p2t
P2
ISiCjSP
Theorem 2
Let Sc be the SyRM strategy and SR the baseline strategy that masks only response tokens. Define the maakk-ptteerd noise of a strategy as A(S) := Var(L S). For m E {C,R) with corresponding cligible-set size Pm and token statistics pi,m; {, Pij.m
Pm B[L/t]
[[(1- t)/t]
Pm Σ
A(S,.) =
P2
Σ
o
m.
P2
P mm
+
P2
Σ ISi<iSPm
Pij.m
Proof Take the conditional variance from Theorem and integrate over Note that token-wiss wise parameters naturally depend the chosen strategy, let P, os Ms, Psy depend on {C, R).
Theorem Partition the eligible tokens into
1. response tokens R (count Pr).
2. syntax prompt tokens C' (count Pc),
so P = PR + Pc. Under the assumptions
(A1) All tokens in R share un, oi tokens in C share pc, o
(A2) ot < oi
(A3) PCC & PRR.
(A4) |prcl & VPRRPCC.
(1 -1-60 (PR 1)PRR (1 -), where 00/02 PCc/PRR.
(A5) 2wR <
o7,B
WR PR/P P. and = [[[[[/
we have
A(Sc) A(Sa).
Proof
Denote ::=E [
0. By Theorem 2,
B 2
Pn -
A(SR)
OR PR
(B
1)p2
PaR
(19)
PR
PR
23

---

Published as a conference paper at ICLR 2026
Ps. In the SyRM strategy, both R and C are eligible. Denote wn p WC Pc > max PR. Po). We group terms accordingly, yielding
Pc. where P PA +
RR PRB
oc PR -1 PoB P PRRt P
Pc pcct p PRC www.tuc.cttuuuuueqqe.
2PRPc
A(Sc)
(20)
where
0R
PR 1
dc PaB
Pc
AR
PRR:
Ac =
+
PCC.
PRB
PR
Pc
Thus,
A(Sc) - A(SR) = wCCAcc - Ar) + 2URWCPRO-
Claim 3.1
An > Ac
Proof
By the Cauchy-Schwarz, inequality,
|pRCl VPRRPCo.
From Eqs. 19 and
20
we further obtain:
AR Ac
PR
R
2
Pc
2
B
PR -1
PR
PRR
Pc -
Pc
PCc
(21)
Define:
2
0C
PCC
(as in Assumption A, assume a, < 1),
PRR
then:
AR AR-Ac - BB
Pn
1
Pe
a
+PRR
PR-1 -1-1
PR
B
Pc -1
Pe
2 RR
1-a
Pn.
-DRR"
1-B PR - 1
Pn.
(22)
Therefore,
An > Ac
Using the inequality
Ipncl S VPRRPOC S 3pnn.
obtain:
o2
PR
A(Sc) A(Sn) & wc
- 0:)
(1 (8)
PRn
22wnwe3pnn.
Pn
PR
=:D
Since D > 0 and 3 < 1, setting
D
2w 3 <
PRR
ensures that
A(Sc) A(Sn) < 0.
24

---

Published as a conference paper at ICLR 2026
Discussion of Assumptions
A1-A2. Syntax tokens C' are homogeneous and more stably predicted, hence lower mean and variance.
A3. Reasoning errors couple response tokens more strongly than syntax tokens, implying pan > PCC-
A4. Cauchy-Schwand Schwwrz bounds the cross-covariance.
A5. Given A2-A3, this technical condition typically satisfied in practice and ensures the cross- term cannot offset the variance reduction provided by including C.
Theorem 4
LsyRM
Lresp + (1 a) Leoord
Proof 4 Consider single forward pass with Pn response tokens and Po syntax tokens (disjoint sets, so P P PR + Pc). All eligible positions are masked with the same rate ~ Uniform(0, 1). Define the per-step loss
= PI
M,Y,
M₁ ~ Bernoullit)),
where P is the number of eligible tokens under the chosen strategy. For the SyRM strategy, PSyRM
Pn + Pc, and
LSyRM
1
Σ M,Y,
PSYRM
iERUC
PR Σ M.Y. +
Pc
Pot Σ M,Y
PSYRM PR iER
PSyRM
jEC
= Trrep + (1 a)) Tecondss
where
PR a PR + Pc
Pc PR + Po
1-a
Hence,
LSyRM Lreu (1 a) Lewnds
which is a purely algebraic identity requiring no additional assumptions.
Theorem Optimum-shift bias inevitable.
Proof Define the objectives Jresp (0) = TE[ [reep)) Jooord (0) = EE Lccord), JSyRM (0) (O) + (1 a) Jeoord(O)) (0). Let the response-only optimum Oresp satisfy V VeJreap (0īosp) 0. If the same point VoJcomne (0ieap) + 0, then
V (0 = (1 - a) V VaJcoord (0) (0tesp) + O,
6t resp is not a stationary point of JSyRM- Under gradient- gradient-based optimization, the parameters will so continue move and eventually converge different optimum S # t Hence, whenever VaJcoard (0resp) + 0, the two optima necessarily diverge; coincidence occurs iff VaJeoord (0; (0jesp) =
0.
Theorem
If the Hessian of JsyRM is positive definite,
1 CV
||0š,RM
0ispll 5 Amni (SSRR))
V Jecord ( (9ies))
25

---

Published as a conference paper at ICLR 2026
Proof Perform first-order Taylor expansion of syRM around fesp:
0 V J.syRM = V (JyyR(((esp)) (0 0 @resp):
where (S) is the Hessian evaluated at some point g between 03yRM and Fesp. Rearranging gives:
fGyRM Presp (1 ax) HyyRM ( Jamond
(E) is positive definite, we can bound the Euclidean norm using its smallest eigenvalue Amin( (HyyRM):
1
[05MMM
0ieupl S
V Veoord (f (eesp))
Amin (HSyRM)
Explanation:
Weight factor - The smaller the proportion of syntax tokens, the tighter the bound.
Curvature Amin bmin(HsyRMM)) A steeper loss (larger curvature) suppresses the shift.
Task conflict || Va (9 (/sp If the syntax gradient is already small at the resp-optimal point, the shift will naturally be small.
A.5.5 STRATIFIED SAMPLER
Specifically, when processing batch of data points, instead of independently drawing mask ratios t ~ 24(0, 1], one partitions the interval [0, 1] into & uniform strata: [0, 1/k],, [1/k, 2/k]..... [[k 1)/k, 1), and samples one masking ratio from each stratum. This approach, known as the stratified sampler, is known to produce low discrepancy. Here, & is the number of strata, which must be chosen before training.
This technique is discussed in Sahoo et al. (2024), where they use k: (i.e., one sample per stratum). Here, we provide a rigorous analysis and offer a practical guideline for selecting
A principled Method: Unbiased and Reduce Variance
Let
(o (, 2x) wo. t]. G(() go (wo, () = Ee (2
For fixed Zo, we are interested in the expected value of G(t) under t ~ Unif(0, 1], i.e.
# = EG)((t))
During training, this estimated by the empirical mean over a batch:
G(L))
If t/ ~ i.i.d. Unif(0, 1] for i 1,...... then by standard results,
Varla [o] = o. where of = Var,[G(th
This corresponds the B term in the variance decomposition:
Exg (Var
For stratified sampling, partition [0, 1] into & equal-length intervals:
j k. for j = I,
and sample 1g Unif(T;) independently within each stratum. Then the estimator becomes
c(ty).
26

---

Published as a conference paper at ICLR 2026
This remains unbiased estimator of μ4, and its conditional variance is:
1
Varláa mo)
V o wT
12
where we define the within-stratum variance:
Vanee,
Decomposing the total variance of into within- and between-stratun variance gives:
71 of v + of), where of = Vars[ux), l = Eee [G(()-
Therefore, the variance the stratified estimator satisfies:
Var [ zul
0
k bb
which is smaller than the SRS variance:
1 Varses - 0,
n2
The reduction is exactly
1 k
%
Practical Suggestion for Choosing
To choose an appropriate number of strata, we aim balance two objectives:
1. Reduce the overall estimation variance (which favors using larger number of strata rt);
2. Ensure the estimation of within-stratum variance oe-denoted by oğ-is sufficiently stable (which favors having more samples per stratum, i.e., m int(B/n)) B/n)).
While setting B results in highly fine-grained strata, may lead unstable estimates within each stratum due to too few samples. To balance this trade-off, we propose the following optimal number of strata:
Nogt O B,
where B is the total sampling budget, and m int B B/)) denotes the number of samples per stratum.
A.5.6 BIN- WISE EMA: VARIANCE REDUCTION HYPERPARAMETER SELECTION
Bine EMA Reduces B as a Control-Variate Strategy To estimate a random variable X with mean yx [[ samples ((x,, 1 the control variate estimator
control variate Y with known mean byy E(y]. Given i.i.d.
pCV = X Cov(X,Y) of
reduces the variance from estimator gives unbiased estimate of loss gradients.
3 to 2 (1 0²), where
Cov(X,Y)
Importantly, this form of axax
In our context, X represents the loss gradient for the current batch, while is the exponential moving average by of past gradients within the same t-bin. Assuming strong heteroscedasticity over t E [o, 1], partitioning t into bins and applying EMA within each bin maximizes |pl, thus achieving significant variance reduction. Several alternative control strategies are listed in Table 3 Among these, we adopt bin-wise EMA for its strong variance reduction under heteroscedastic conditions and minimal implementation burden, which also helps prevent overfitting.
27

---

Published as a conference paper at ICLR 2026
Variant
Assumptions Smooth {(t), no sharp changes
Tuning and Resources
Characteristics
Spline
Requires degree/node selec-
Captures trends well but sen- sitive to under-/over-
tion; low memory usage
Kernel
Flexible; handles arbitrary her-
Only bandwidth /; stores O(1hh) samples; higher mem-
Most robust; no shape sumptions; tuning and imple- mentation
eroscedasticity
ory cost than bin-wise EMA
complex
Table 3: Comparison of spline and kernel variants
Choose the Number of Bins in EMA: A Mathematical Perspective Suppose we have m data points and choose b bins; then each bin will contain approximately m/b samples estimate the piecewise-constan value of b(t) over an interval of width 1/b.
We use MSE optimization analysis to determine principled choice of given m. Specifically, we aim to estimate the function
(t) = Fz,, (Co (No: f, mr) ma, 4]
on [0, 1] by piecewise-constant approximation Do (6) with b equal-width bins of size A 1/b, using the sample average within each bin.
The total mean squared error is decomposed as
MSE B
E[((os
((t)) + (ar(
Blas
Variance
Bias Term
Assuming g(t) is twice differentiable within each bin, we can apply Taylor approxi-
mation:
Bias² = CG4 C1
b
where Ci x max (I)]-
Variance Term Each bin contains m/b samples, hence:
C2 C2b
Var
n m
where C2 x Varmask [e].
Total MSE
Combining both terms:
C1
C2b
MSE
b
m
Optimal
Taking the derivative w.r.f b and setting it to zero:
[40mm²
1/5
MSE B bC, db
0 b
xR 1/5
C2
Selection for
For 1). it's useful to think in terms of the e-folding time in(--pp L. that
7.
each EMA bin effectively averages over the past k x (batch size per GPU) losses. Larger adapts faster but fewer past samples to estimate b(t); smaller v is stabler but cause the EMA to lag (and suffer from domain-shift issues). In our experiments, we set:
x batch size per GPU) RJ 0.1 (train data size).
X
A.6 METHOО COMBINATION
One might wonder whether all methods proposed in this work could be combined into a single ultra-strong approach. While such combination may sound attractive, it is not straightforward task: different methods rely on distinct assumptions, and naive aggregation can easily introduce redundancy bias:
28

---

Published as a conference paper at ICLR 2026
1. Different Assumptions. Each method relies specific assumptions. For example, the effectiveness of P-POTS depends on accurate modeling of g(t) and v(t), while Bin- Bin-Wise EMA requires heteroskedasticity of the loss I with respect to t. Blindly combining them risks violating these assumptions.
2. Redundancy. Each method in this work already effective on its own. If two methods target the same source, the first may remove most of the variance (e.g., 80%), while the second may introduce estimation bias that offset any residual variance reduction, resulting in worse performance than using a single method.
Instead, we emphasize the synergy between P-POTS and MIRROR, which our experiments confirm be stable. In practice, we recommend P-POTS+MIRRON for maximal performance gains, and P-POTS alone for cost efficiency. Other techniques may still be preferable in specific contexts due to their simplicity or ease of implementation. Rather than combining everything, we provide guidance below on how to select among them.
A.6.1 STRATIFIED SAMPLING VS EMA
Both methods rely on the assumption that the loss is heteroskedastic in t: after averaging over æo and It, different values of yield losses of different magnitudes. This generally holds, since larger t masks more tokens, making reconstruction harder and the loss higher, as shown in Figure 5a
(heteraccccasslckky)
Shi Shankage discreaaccy gao with
pimplo 101
022
0.6
2.8
1.0
10'
35'
E leve!
Butchaaee soutd
(a)
(b)
Figure 5: The left panel shows the heteroscodasticity of the t-loss, and the right panel illustrates the shrinking difference between IID and stratified sampling as batch size B increases, where we use the KS statistic to measure their maximum deviation.
However, the two methods reduce B fundamentally different ways:
Batch Size: Stratified sampling enforces fixed proportions of t-values in each batch, which is particularly beneficial when the batch size small. As the batch size grows, its advantage over simple independent sampling diminishes, shown in Figure 5b) By contrast, EMA accumulates information continuously across training; its effectiveness depends on the choice of m and 7), but batch size.
Stratified sampling cannot fully eliminate variance, within-stratum variance always remains. In contrast, EMA can, in principle, reduce B to zero if the control variate is sufficiently correlated.
Theoretical Limit:
Note. We do not recommend combining EMA with stratified sampling. Each method effective on its own, and their joint use may introduce estimation bias that offset any residual variance reduction.
In practice, we suggest using stratified sampling as the default strategy for reducing B and switch- ing EMA only when the batch size is sufficiently large.
29

---

Published as a conference paper at ICLR 2026
A.6.2 ISAD VS SYRM
These methods are -ataset-specific. We recommend choosing based on the dataset type: ISAD designed for QA datasets with explicit answer delimiters, while SyRM is designed for datasets containing syntax tokens in prompts.
A.7 EXPERIMENTAL DETAILS
A.7.1 HYPERPARAMETER SETTINGS
We detail here the hyperparameters used for training and inference of masked diffusion models (MDMs) and auto-regressive models (ARMs).
Chain-of- Thhought (CoT) Reasoning.
Dataset
CoT Handling
All except text-to-image-22 2M & Hilab
Responses include chain-of- tf-tthuhsss
HiTab
No CoT provided: generated via rejection sampling
Train/ Vaann/ValiaationTesst Splits.
Dataset
Train/Validation Split
Test Split
OpenScience
5,000 samples, 9:1 (4,500/500) Official train split
5001st sttss00tthhh examples Official test split Official test split st01st-5100m examples
GSM8K
HiTab
Official train split
text-to-image-2M
5,000 samples, 9:1 (4,500/500)
pext-to-image-i consists of two subsets, 1024_10k and 512.2M. We sample 2,500 examples from each subset to construct the training and validation splits, and 500 examples from each subset to form the test split.
General Hyperparameter.
Parameter
Value
Learning rate
5 x 10-5
Scheduler
Linear
Global batch size
32
Epochs
5
Max sequence length
4096
Warmup steps
0
Note that after constructing the training data as described in Train/Validation/Test Splits., we fur- ther filter out samples with a total sequence length (prompt + response) exceeding 4096, matching the fixed context length used in the pretraining stage LLaDA. As a result, the final dataset may be smaller than the original.
Method-Specific Hyperparameters (MDMs Only).
Technique
Setting
Importance sampling
A
0.2
Stratified sampler
#strata bbtch size 6
EMA bins
10
EMA learning rate
0.01
Checkpoint Selection.
30

---

Published as a conference paper at ICLR 2026
Criterion Evaluation strategy
Setting Early stopping (patience 3)
Random Seeds (MDMs).
Seeds used
42, 731, 20231
How Data Cleaned in HiTab
Raw answers are extracted from [boxed(} delimiters, then cleaned by removing surrounding whitespace, unwrapping ftext blocks, converting escaped percent signs { to %, deleting backslashes and trimming trailing period
Inference Settings for LLaDA.
We the decoding lemperature to 0 for all ARMs and MDMs, ensuring that variance in down- stream task performance originates solely from training.
Dataset
Train Data Response Max Length
ARMs
MDMs
GSM8K
444
512
128-128-32
HiTab
3571
3968
512-256-16
Open
4605
5120
256-256-256
The ARM column indicates the maximum number of new tokens, while the MDM column specifies generation length-steps-bloct length, as we adopt block diffusion for inference (Arriola al. 2025).
For ARMs, we determine the maximum new tokens as:
max new tokens = Tresponse max length (1 + s) 164,,
where s a safety factor (set to 0.1), and the result is upscaled to the nearest multiple of 64.
For MDMs, hyperparameters are selected according to dataset type.
Note. shown in the table, ARMs require significantly longer generation lengths at inference. This arises from their auto-regressive nature: given a budget of max new tokens, they auto-regressively generate token by token, without explicit awareness of how many tokens are ultimately needed. In contrast, MDMs substantially fewer tokens-even shorter than the training responses-yet still produce accurate reasoning without truncation. This efficiency arises because MDMs predict all masked tokens in parallel and are constrained by the generation length. For example, with gen length = 512, the model predicts all 512 tokens for the masked input and iteratively applies the remasking-net maaiinn scheduler until completion. By design. MDMs avoid unnecessarily verbose outputs, which we argue constitutes another advantage of MDMs over ARMs.
Inference Settings for MMADA.
Dataset
generation length
steps
block 1024
temperature (
mask schedule cosine
text-to-image-MM
1024
15
A.7.2 HOw TRAINING VARIANCE IS ESTIMATED
During training, maintain online estimator of the per-sample loss variance using importance- weighted statistics. Since losses decrease during training, this metric captures both fluctuations and the overall downward trend, it is approximate empirical proxy for Eq. 2
For each batch, the por-token (or per-sample) losses are detached, flattened, and reweighted by the corresponding importance weights w. We then update three running sums: the weighted mean of losses (S1), ) the weighted mean of squared losses (S2), and the weighted mean of squared weighted losses (S12). These sufficient statistics allow to compute unbiased estimates of the second mo-
31

---

Published as a conference paper at ICLR 2026
ment m2, the squared mean 2², and finally the unbiased variance estimator
Varl = m2 -
(23)
where m2 S2 and 1² - n(n-1) Here, n denotes the total number of accumulated samples across updates. If imporaaeee sampling on t is enabled, the weights w are drawn from the fitted distribution p(t) ensure unbiased estimation; otherwise, unit weights are used.
At the end of training, the accumulated statistics are queried once to produce the final unbiased variance estimate. These values are then substituted into Eq. 23 obtain the final estimate of training variance.
Intuitively, the use of an unbiased estimator is necessary because the naive sample variance can be downward biased when computed incrementally with importance weights. By correcting for this bias, our estimate better reflects the true variance of the training process, independent of batch size or logging frequency.
A.8 COMPLETE EXPERIMENTAL RESULTS
A.8.1
PER-RUN EXPERIMENTAL RESULTS ON LLADA
Seed 42
Seed 731
Seed 20231
Dataset Method
Perf
Time (h)
Var
Perf
Time (h)
Var
Perf
Time (h)
Var
OpenScience
Baseline
43.57
11.29
2.1404
46.50
11.32
2.0614
46.50
11.40
2.0743
P-POTS+Mirror
50.80
14.70
1.5550
53.80
14.07
1.5672
53.00
14.00
1.5776
P-POTS
45.34
13.63
1.6319
47.60
12.98
1.6417
47.47
12.88
1.6622
Mírror
45.63
12.46
2.0800
47.25
12.47
2.0392
46.25
12.47
2.0205
GSM8K
Baseline
50.64
1.55
0.2526
53.22
1.62
0.2524
53.68
1.61
0.2673
P-POTS+Mirror
62.02
3.64
0.1906
58.61
3.73
0.1949
60.96
3.88
0.1922
P-POTS
59.59
1.67
0.1726
59.36
1.77
0.1715
56.79
1.67
0.1722
Mirror
56.56
3.69
0.3167
51.40
3.37
0.3042
53.15
3.38
0.3162
HiTab
Baseline
52.96
8.22
1.0063
60.43
8.22
1.0116
62.64
8.26
1.0173
P-POTS+Minor
66.02
17.80
0.9903
66.67
16.78
0.9892
68.62
16.83
0.9847
P-POTS
60.56
10.19
0.7498
64.74
10.66
0.7536
58.80
10.99
0.7503
Mirror
64.72
16.08
1.3052
65.04
16.27
1.3282
63.68
16.29
1.3480
Table 4; Complete results -pefoormance, training time (in hours), and variance across datasets (OpenScience, GSM8K, HiTab).
Compared to only reporting mean and standard deviations across runs, this complete results table provides a more direct view training instability: under standard training, the final performance vary substantially, e.g., from 52.96% to 62.64% on HiTab. In contrast, method achieves more consistent outcomes across runs, demonstrating greater reliability.
Tad show that methods with lower reported training variance generally associated with better downstream performance. Interestingly, MIRROR sometimes reports higher variance than stable results in evaluation. Since variance measure serves only
the baseline, yet achieves as proxy that captures both the downward trend and short-term fluctuations of training losses, we conjecture that the stronger overall loss reduction under MIRROR may interact with this proxy in way that inflates the reported variance, even while true training stability improves.
32

---

Published as a conference paper at ICLR 2026
Table 5: Average training variance with training hours across datasets (All experiments were trained on 2 H100 GPUs; details training variance approximation provided in A.7.2)
OpenScience
GSM8K
HiTab
Baseline
2.0920 (11.33)
0.7723 (1.59)
3.0352 (8.23)
P-POTS+MIRROR
1,5666 (14,26)
0.5777 (3.75)
2.9642 (17. (17.13)
P-POTS
1.6453 (13.17)
0.5163 (1.71)
2.2537 (10.61)
MIRROR
2.0466 (12.47)
0.9371 (3.48)
3,9814 (16.21)
ISAD
2.1480 (11.48)
0.2525 (1.61)
1,0016 (10.51)
SyRM
0.9233 (8.23)
StraTS
2.0851 (11.46)
0.2568 (1.59)
1,0083 (10.19)
EMA
2. 2,1400 (11.46)
0.2535 (1.62)
1.0018 (8.40)
A.8.2 PER-RUN EXPERIMENTAL RESULTS ON MMADA
Seed 42
Seed 731
Seed 20231
Dataset / Method
Perf
Time (h)
Var
Perf
Time (h)
Var
Perf
Time (h)
Var
Baseline MIRROR
28.61
3.21
1.4149
34.28
3.09
1,4245
31.19
3.31
1.3848
34.88
6.88
1.0308
33.77
6.13
1.0808
33.57
6.74
1.0458
text-to-image-2M
EPR
P-POTS
31.51
3.28
2.6203
30.02
3.33
2.6150
33.41
3.37
2.6236
P-POTS+MIRROR
35.27
6.97
2.4044
34.10
6.79
2.3504
34.39
6.87
2.3370
text-to-image-2M
polynomial (degree 7)
P-POTS
30.00
3.28
2.8395
33.35
3.27
2.8398
28.66
3.56
2.6187
P-POTS+MIRROR
34.51
5.96
2.3538
33.78
6.43
2.7319
34.78
6.71
2.3199
Table 6: Performance (Perf), training time (in hours), and variance (Var) on Lext-to-image-2MMe 2M-EPP and text- to-image-2M-pol across methods and random seeds (42, 731, 20231).
1:24
1.3
174
120
65
18
(a)
(b)
(c)
Figure 6: Shape of p(t) under different situations: (a) shape of empirical and fitted p(t) by EPR before training: (b) by polynomial of degree 7 before training: (c) fitted by EPR after training for one epoch using polynomial (degree 7). These functions fitted by minimizing the KI. divergence, only their shapes matched; they later renormalized in training yield valid PDFs.
As shown in Figure even when the model is trained with the S-shaped distribution, after one epoch it quickly reverts back to the U-shape. This suggests that the EPR form is stable and persistent structure across datasets.
A.9 EXTENSION TO REINFORCEMENT LEARNING
Applying standard RL algorithms such as PPO and GRPO to MDMs is challenging, because it requires computing logelikelssoods over transitions between consecutive states. In ARMs, states are token sequences. Since tokens are generated one by one, by the chain rule P(I 11; P(z; Ies), the transition probability for each token 2; is naturally po(x; xzz). In contrast, MDMs generate tokens in parallel, which induces complex marginalization between intermediate states It-1 and 2,, Training stability is further challenged by token-level rewards derived from high- variance sequence returns.
33

---

Published as a conference paper at ICLR 2026
Algorithm UniGRPO Policy Gradient Optimization
Require: Reference model Tief, prompt distribution D, number of completions per prompt G, num- ber of inner updates diffusion steps '''
1: Initialize policy <- 7ref
2: while not converged do
3: Told 7T
4: Sample a prompt ~ D
5: Sample G completions o; To ()) for E [G]
6:
For each o;, compute reward r; and advantage A! (ToMd Sample starting timestep to M00 1)
7:
8:
Generate μ uniformly spaced timesteps & fu-1 from [o, 27
9:
for gradient update iterations n to μ do
10:
if n = then
Sample starting mask ratio r1 2(0, 1) and compute initial timestep t [r]]] else
12:
13:
Compute I,, = (T (1) +4
14:
end
15:
Construct input (g. masked 0:) using timestep (with @ always unmasked) For no Told, Trei estimate log-probabilities of masked tokens in o, at in
16:
17:
Compute UniGRPO objective 24 and update Ty via gradient descent
18:
end for
end while
20:
return T
Our methods can be naturally extended to these problems. example, consider UniGRPO (Yngg et al. 2025 (Algorithm T.. whose training objective defined as
UnitReO (0)
(.g.)) f (11)), (prEjua))
led
(iin(F(,(()) A.,, clip(r)..()), 1 1 1 +e) 3DKu(Tb ll Hide))
(24)
G
We can incorporate our core methods through the following modifications:
P-POTS:
Replace lines 7-8 with sampling μ timesteps from p" (t), sorted tor..... .... fu-1-
(2) for each ln and average their MIRROR: At line 15, construct complementary masks of" 0; log-probabilities before updating TT.
At present, this integration is conjecture: confirming its effectiveness requires empirical validation, which we leave as future work.
A.10 FUTURE WORK
While our current results highlight both the importance variance reduction and the effectiveness of our proposed methods, several open directions remain. First, we aim deepen understanding of the p-drift problem and to validate the effectiveness periodically re-estimating or gradually down-weighting p(t), suggested in 13.1.1 which far remains conjectural. Second, since rein- forcement learning on MDMs face the difficulty of likelihood evaluation, is worthwhile to explore whether our methods can also reduce estimation variance in this setting (see A.9).
A.11
CASE STUDY
A.11.1
CASE STUDY: IMAGE GENERATION
Below are images generated by MMaDA-8B -B-MiXCoT trained with P-POTS+MIRROR (left) and standard method (right).
34

---

Published as a conference paper at ICLR 2026
h
35

---

Published as a conference paper at ICLR 2026
Dioinas DOI ANDTTIGGG s.oottttia 6 L.
OIGIOU D10000
36


Reviews:
[{"review_id": "pyjNQm1kcH", "reviewer": "Reviewer_oCXk", "summary": "The paper introduces a theoretical and practical framework to understand and address instability in masked diffusion model (MDM) training. The authors first derive a variance decomposition framework that attributes training variance to three distinct sources: (A) masking pattern noise, (B) masking ratio noise, and (C) data noise. Based on this decomposition, they propose six variance-reduction methods. Experiments on multiple text reasoning datasets (GSM8K, HiTab, OpenScience) and a text-to-image benchmark show consistent improvements in both performance and stability. Overall, the paper is technically strong, though somewhat dense and in need of a tiny bit of clearer writing.", "questions": "See Weakness.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Strong theoretical foundation. The paper provides a clear and principled variance decomposition for masked diffusion model (MDM) training, unifying prior ad-hoc stabilization methods under a single theoretical framework. It then builds directly on this foundation by proposing six targeted variance-reduction techniques to mitigate the identified sources of instability.\n2. Comprehensive empirical validation. The experiments cover both language and multimodal domains, demonstrating that the proposed methods are broadly effective and improve training stability across diverse settings.", "weaknesses": "1. Narrow comparison to ARMs. The study includes only two autoregressive baselines from the same family. Incorporating additional ARM baselines, especially models with different architectures or training paradigms, would help clarify whether the observed variance gap is a general phenomenon or specific to the chosen comparison set.\n2. Limited model diversity and scaling analysis. While the empirical results are solid, they are restricted to a single MDM backbone (LLaDA-8B-Instruct). Evaluating the proposed methods across different model sizes and architectures would strengthen the claims of generality and potentially reveal whether the improvements follow any scaling trends within or across model families.\n3. Writing and presentation. The exposition could be tightened to improve readability; the current density of equations and notation can make the paper feel more complex than necessary. Additionally, some figures could benefit from more informative captions, for instance, Figure 3 presents image generation results but omits the corresponding prompts.", "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": "51X5CVdK08", "reviewer": "Reviewer_ag9a", "summary": "The paper analyzes training instability in masked diffusion models by decomposing the loss variance into three parts: (A) randomness from the mask pattern, (B) randomness from the mask ratio/timestep, and (C) data variance. To reduce these sources, it introduces P-POTS (a timestep sampler), Mirror (a complementary mask), and various other techniques. Across multiple benchmarks, the proposed approaches stabilize training and improve final results.", "questions": "The variance decomposition is not unique. It depends on the conditioning order when you iteratively apply the law of total variance. Could the authors discuss how your conclusions change under alternative decompositions, and whether those alternative decompositions could lead to other interesting approaches?", "limitations": "", "rating": 4, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "Overall, I think the paper studies an important problem. The proposed fixes P-Pot and Mirror are clearly argued and have demonstrated practical usefulness in terms of accuracy and training stability.", "weaknesses": "1. The loss-variance decomposition is insightful, but I believe for training stability it would be more insightful to analyze gradient variance. It would be nice to see how reductions in the proposed loss variances translate to reduced gradient variances and more stable optimization.\n2. MIRROR roughly doubles the cost on some benchmarks compared to the baselines, which is quite expensive. Would MIRROR still be the best choice under a fixed time budget, which is a more practical scenario?\n3. I am not familiar with MDMs and thus cannot comment on the commonly reported numbers. But for ARMs, the Qwen-2.5-7B-Instruct and Qwen-3-8B numbers appear below the commonly reported results. Could the authors clarify their evaluation approachs?", "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": "MDBz9VNjWQ", "reviewer": "Reviewer_WDaB", "summary": "This paper try to decompose the MDM's training noise. The noise is divided into three parts: data noise and two about masking.\nThe core analyses are about inter-group and inner-group decoupling by the law of the total variance. Based on the aforementioned analyses, two main approaches are proposed, P-POTS and MIRROR.", "questions": "- Selected benchmarks are all about QA reasoning tasks. How about general QA tasks and knowledge-intensive benchmarks, e.g., Graph QA, OCR detection?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- Intuitive and sharp derivation of the theorem provides a mathemetically elegant  and practical explanation.\n- Numerical pre-experiments seems robust and adheren to the expected Pareto frontier.", "weaknesses": "- Limited experiments about generalization and comparison. The ablaition experiments are mixed in the table of comparison. The included MDM baselines are too limited.\n- Parts of the error bars are missing, and meanwhile, the error bars reported are too large to convince audience that the methods are consistently performing well as it's tested in the main table.", "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": "tXs7tOmcn1", "reviewer": "Reviewer_J4Ut", "summary": "This paper tackles the critical problem of high training variance in Masked Diffusion Models (MDMs), which causes optimization instability and performance degradation compared to Auto-Regressive Models (ARMs). The paper derives the first systematic decomposition of MDM training variance, identifying three distinct sources: A) Masking Pattern Noise, B) Masking Ratio Noise, and C) Data Noise. Based on this decomposition, the authors propose six methods, including two core techniques: P-POTS and MIRROR.  Experiments on complex reasoning and multimodal tasks show that the proposed methods, particularly the combination \"P-POTS+MIRROR,\" dramatically improve performance (e.g., +7-8% accuracy on GSM8K) and, critically, reduce run-to-run performance variability.", "questions": "1. Based on my understanding, LLaDA and Dream seem to use slightly different formulations to train MDM. However, this paper only conducts experiments on LLaDA. Thus, are the proposed methods applicable to other baselines, such as Dream?\n2. This paper discusses two types of LM, MDM and ARM, where MDM is fully bidirectional. How about the combined version, such as BlockDiffusion[1]? Does the author have any insights on this case, such as the training variance or whether proposed methods are still applicable?\n3. Though experimental results are quite impressive, I still recommend that the author provide train-from-scratch results in the future since the training variance in pretraining might be more salient.\n\n\n[1] Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models, in ICLR 2025", "limitations": "", "rating": 8, "confidence": 2, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "1. This paper provides a clean, intuitive, and theoretically sound framework for understanding why MDMs are unstable.\n2. The proposed core methods (P-POTS and MIRROR) are well-motivated and directly derived from the theoretical analysis. \n3. The significant improvement in the MDM's performance makes the proposed method not only theoretically sound but also practically effective.", "weaknesses": "This is a strong paper without major weaknesses. One suggestion on the presentation is that although this paper proposes two core methods to reduce the instability, it introduces six methods in total. I think this may harm the readability of this paper and distract the reader from understanding the core contribution. Therefore, the author may reconsider the paper structure.", "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)