Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffusions
Authors: Jaeyeon Kim
Kulin Shah
Vasilis Kontonis
Sham M. Kakade
Sitan Chen
Venue: ICML 2025 oral (2025)
Keywords: Discrete Diffusion models
Masked Diffusion Models
Diffusion Models
Learning Theory
Inference-Time Strategy
Primary area: deep_learning->generative_models_and_autoencoders

Abstract: In recent years, masked diffusion models (MDMs) have emerged as a promising alternative approach for generative modeling over discrete domains. Compared to autoregressive models (ARMs), MDMs trade off complexity at training time with flexibility at inference time. At training time, they must learn to solve an exponentially large number of infilling problems, but at inference time, they can decode tokens in essentially arbitrary order. In this work we closely examine these two competing effects. On the training front, we theoretically and empirically demonstrate that MDMs indeed train on computationally intractable subproblems compared to their autoregressive counterparts. On the inference front, we show that a suitable strategy for adaptively choosing the token decoding order significantly enhances the capabilities of MDMs, allowing them to sidestep hard subproblems. On logic puzzles like Sudoku, we show that adaptive inference can boost solving accuracy in pretrained MDMs from $<7$\% to $\approx 90$\%, even outperforming ARMs that were explicitly trained via teacher forcing to learn the right order of decoding.

Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffusions
Jaeyeon Kim
Kulin Shah 22
Vasilis Kontonis 2
Sham Kakade Sitan Chen'
Abstract
variant, distinguished by simple and principled objective: to generate samples, learn to reverse a noise process which independently and randomly masks tokens.
In recent years, masked diffusion models (MDMs) have emerged as a promising alternative approach for generative modeling over discrete domains. Compared autoregressive models (ARMs). MDMs trade off complexity at training time with flexibility at inference time. At training time, they must learn to solve an exponentially large number of infilling problems, but at inference time, they can decode tokens in essentially arbitrary order. In this work, we closely examine these two compet- ing effects. On the training front, we theoretically and empirically demonstrate that MDMs indeed train on computationally intractable subproblems compared to their autoregressive counterparts. On the inference front, we show that a suitable strat- egy for adaptively choosing the token decoding order significantly enhances the capabilities of MDMs, allowing them to sidestep hard subprob- lems. On logic puzzles like Sudoku, we show that adaptive inference can boost solving accuracy in pretrained MDMs from 7% to zs 90%, even outperforming ARMs with 7x as many parame- ters and that were explicitly trained via teacher forcing to learn the right order of decoding.
In many applications, such as language modeling, masked diffusion models (MDMs) still underperform compared to autoregressive models (ARMs) (Nie et al., 2024; Zheng et al., 2024), which instead learn to reverse noise process that unmasks tokens sequentially from left to right. However, recent studies suggest that MDMs may offer advantages in areas where ARMs fall short, including reasoning (Nie et al., 2024; Kitouni et al., 2025), planning (Ye et al., 2024), and infilling (Gong et al., 2024). This raises a key question: what are the strengths and limitations of MDMs compared to ARMs, and under what conditions MDMs be scaled to challenge the dominance of ARMs in discrete generative modeling?
To understand these questions, we turn a microscope to two key competing factors when weighing the merits of MDMs over ARMs:
Complexity training time: By design, the predic- tion task that MDMs are trained on is more challenging. Whereas ARMs seek to predict the next token given un- masked prefix, MDMs seek to predict token conditioned on a set of unmasked tokens in arbitrary positions.
Flexibility at inference time: On the other hand, the sampling paths taken by an MDM are less rigid. The order in which tokens are decoded at inference time is random instead of fixed left-to-right. In fact, even more is possible: MDMs can actually be used to decode in any order.
1. Introduction
While diffusion models (Ho et al., 2020; Song et al., 2021) are now the dominant approach for generative modeling in continuous domains like image, video, and audio, efforts to extend this methodology to discrete domains like text and proteins (Austin et al., 2021: Lou et al., 2024: Hoogeboom al., 2021b) remain nascent. Among numerous proposals, masked diffusion models (MDMs) (Lou et al., 2024; Sahoo et al., 2025; Shi et al., 2024) have emerged as leading
Therefore, we ask:
Are the benefits of flexibility for MDMs enough to outweigh the drawhacks of training complexity?
In this work, we provide dual perspectives on this question.
Eqqua contribution 'Harvard University University of Texas Austin. Correspondence to: Kulin Shah <kulin- @ @utuuuuutxssd........
(1) Training for the worst. First, we provide theoretical and empirical evidence that the overhead imposed by train- ing complexity quantifiably impacts MDMs' performance.
Proceedings of the 42nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
We prove that even for simple, benign models of data, there are noise levels at which a large fraction, but not all, of

---

Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffuusioss
the corresponding subproblems solved by MDMs are com- putationally intractable. We then show this imbalance in computational complexity across s subproblems persists even in real-world text data (Fig. 2, left).
MDM training
(2) Planning for the best. While the above might appear to be bad news for MDMs, in the second part of this paper we answer our guiding question in the affirmative by build- ing upon the observation (Zheng et al., 2023) that MDMs which can perfectly solve all masking subproblems can be used to decode in any order.
M
M
M
M
M
In place of vanilla MDM inference whereby tokens are unmasked in random order, we consider adaptive strategies that carefully select which token to unmask next. Our key insight is that this adaptivity makes it possible to sidestep the hard subproblems from training (Fig. 1). In fact, we find that even without modifying how MDMs are trained, the resulting models' logits contain enough information to determine the right order in which to unmask.
MDM inferences (Vanilla vs. Adaptive)
t 0
t
M
M
Vanilla
Our main empirical result is to show that the performance of MDMs pretrained on logic puzzle data dramatically im- proves when one goes from vanilla to adaptive inference. For example, Sudoku puzzles, a simple adaptive strategy (Section 4.1) improves the accuracy of MDMs from < 7% to almost 90%. Remarkably, this not only outperforms vanilla ARMs, but even bespoke ARMs trained to learn the right decoding order via supervised teacher forcing (Shah al., 2024; Lehnert al., 2024) (Table 2).
M
Adaptive
M
b
Figure I. (Top) MDM training can seen as learning multi- ple masked prediction problems, where some are harder to learn, leading to performance imbalance (Section 3). (Bottom) During inference, adaptive MDM can avoid difficult problem instances, improving performance (Section 4).
Organization. In Section 2, we provide preliminaries on MDMs and set notation. In Section 3, we examine MDM training and demonstrate the imbalance in computational intractability across subproblems. In Section 4, we consider adaptive inference in MDMs and investigate its impact on likelihood modeling across various tasks.
oi (e))eo) and az is the predefined noise schedule satisfying an z 1, R O and exg E Rm++1 denotes a one-hot vector corresponding to the value of token w Cat(TT) denotes the categorical distribution given by π E A". In other words, for each i-th coordinate, ar is masked to the mask token O with probability 1 - a and unchanged otherwise.
2. Masked Diffusion Models (MDM)
In this section, we explain the framework of Masked Dif- fusion Models (Shi et al., 2024; Sahoo et al., 2025) and its interpretation as an order-agnostic learner. MDMs gradu- ally add noise to the true discrete data and learn the marginal distribution of the induced reverse process. Below, we for- mulate the forward and reverse processes for MDMs.
Reverse process. The reverse process of the above for- ward process is denoted using Zsj (x, (re x0) and is given by (2,,(5 where
I aeje (ar |2,, 20) for any s t,
(Ca((e.)) 1-ax
o #
Let the distribution Pdata on [1....,m] mm3 be the data dis- tribution over sequences of length L and with vocabulary {1, m... We use 0 to denote the "mask" token.
Cat
On + T-os
a.-
T-ai
ezo
x =
The reverse transition probability (. 20) is approx- imated using go (0(2) (ox)) - dafe(af 2 vo + po( !)) where po(xe, t) is a denoising network trained to predict the marginal on xo via an ELBO- loss. To be pre- cise, gaje (xf Fe, 2o < po (x,, t)) indicates the conditional probability where po 1) is placed in the position of xp
Forward process. For given To Pdata and noise level E (0, 1}. the forward process (- 2o) is coordinate-indeeendee masking process via gejo (2)(2))))) ) =
IE (1) where (a) xn)) Catea +(1-
2

---

Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffussons
within Uale(at) To).
possible masks. Hence, we will refer to the MDM training as order-agnostic training.
Lu ! dj
E Paata
Σ
- log pa t)dt.
On the other hand, Autoregressive Models (ARMs) learn to solve smaller set of infilling problems (L infilling prob- lems in ARMs as opposed to exp(L) infilling problems in MDM) by predicting jtlt token x' given all previous tokens x0, ..... x'=1. This prediction problem equivalent to pre- dicting x' by masking at positions (i...., L-13. Therefore, we write it as
1 at
21 gejo (mr) =0
Here, dj des and òzg,0 is the indicator function; the summation is computed over coordinates i s.t. = O. In practice, a time-embedding-ff free architecture for the denois- ing network, i.e., Po t) = (xa),, is usually employed as z4 implicitly contains information about & via the number of masked tokens.
dt
L-1
log (oro) = Σ
log po (061oo L 1HD- (2)
The reverse sampling process starts from the fully masked sentence x1 = (0, (0,..... 0). At given noise level E (0, I]. suppose have partially masked sequence sy. For prede- termined noise level < t, we sample I, go())aa). This process is repeated recursively from t = 1 to = 0.
ARMs are trained to predict tokens sequentially from left to right in all sequences. We refer to this as left-to-right training. In general, one can also consider predicting tokens sequentially under some fixed, known permutation of the sequence; we refer to this as order-aware training.
2.1. Reformulating the training and inference of MDMs
In this section, we first discuss vanilla order-agnostic train- ing of MDMs and compare it with "left feftttt--iihh" order train- ing of autoregressive models in Section 1. Then, we reformulate vanilla MDM inference in Section 2.1.2 to set the stage for the upcoming discussion.
2.1.2.
INFERENCE OF MDMS
The MDM inference can be decomposed into two steps: (a) randomly selecting set of positions to unmask and (b) assigning token values to each position via the denoising network pe. More precisely, we can reformulate the reverse process F3 go follows.
2.
TRAINING OF MDMs
Recent works (Zheng al., 2024; Ou et al., 2024) have observed that the learning problem of MDM is equivalent to masked language model. Building upon their analysis, we reformulate the loss Lg to show that La is a linear com- bination of the loss for all possible infilling masks. We first define xnp [M] as masked sequence, obtained from original sequence Zo where indices in the mask set M (regarded as subset of [Z] A (1, L)) are replaced with mask token O.
Vanilla MDM inference
(a) Sample a set of masked tokens S S {i xr} - 0}, P(i S) -.-0.
(b) For each i E S, sample xं ~ po(x=
Therefore, the inference in MDM is implemented by ran- domly selecting S and then filling each token value accord- ing to the posterior probability po (o (x)).
Proposition 2.1. Assume ag 1, a1 = and de- noising network pa is time-embedding free. Then Lo s
-E Spann [ogg po (xo) and
3. MDMs train on hard problems
Lo
Σ
1
1
E
[lg po (z)
[o]l
In this section, we theoretically and empirically demonstrate that a large portion of masking subproblems (= x [M]) can be difficult to learn. For intuition, consider solving a masked prediction problem (= [M[) on text data like masking an arbitrary sentence the middle of a document and predicting the correct word for specific position in that sentence. It is reasonable that this task should be more complex, even for humans, than left-to-right prediction, and in this section, we place this intuition on a rigorous footing.
|M
20 zarepaaaa
MEIL)JEM
(1)
where |M] is the size of the set M and (a M)) indicates the conditional probability of the i-th coordinate from (xt)..
The proof of the above proposition is given in Appendix E. As the MDM loss is linear combination of the loss for all possible infilling mask M, the minimizer of the loss Lo learns to solve every masking problem. In other words, the optimal predictor Pa is the posterior marginal of the i-th token, conditioned on no [M] for all masks M. The training objective of MDM aims to predict zo from xn [M] across
In Section 3.1, we show several examples of simple, non-pan con-pathological distributions for which: (1) the mask- ing problems encountered during order-aware training are computationally tractable, yet (2) many of the ones en-
3

---

L&O distributions contain two types of tokens: (1) latent to- kens and (2) bservation tokens. Intuitively, latent tokens are Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffuusioss
AR HDW Iilcmmer coso
oicemerccteer nieamerunt
Task
(Sa)
(wrr(
Tog(Ulkelhhod
Val. Loss
3.171
3.212
3.245
5.-
2.-
Prediction Error
14-
22-
lgg(FLOPs)
Position
Figure Left: MDMs train on hard problems (Section 3.2). x-axis and y-axis correspond to log(PL)Ps) and - log (ag() MDM (Blue) is worse than ARM (Orange) in likelihood modeling, Most masking problems (Other lines) that MDM is trained on are harder than those encountered by ARM, as indicated by small log-likelhhods. Right: Task error imbalance (Section 3.3). MDM's s performance varies across different tasks. For text data (top right), this is indicated by validation loss. For T&&O-NA--SAT (hottom right), MDM performs well on the masking problems for observation positions (light region) but struggles with latent positions (dark region).
resppcivve respectively.
countered during order-agnostic training are computation- ally intractable. In Section 3.2, we empirically show that text data also exhibits this gap between the com- putational complexity of order-aware and order-agnostic training. In Section 3.3, we reveal that this discrep- ancy in computational complexity manifests empirically in performance imbalance across tasks: as predicted by the theory. MDMs trained on data from such distributions exhibits small errors on easy subproblems but suffers from large errors on harder ones.
tokens in the sequence, indexed by T(1, ( (NV) that serve as "seeds" that provide randomness in the se- quence; the remaining tokens, called observation tokens (indexed by x xN + 1), (N + .., aN + P)), are de- termined as (possibly randomized) functions of the latent tokens via O1...., Op.
Note that by design, order-aware training, e.g. by permuting the sequence so that becomes the identity permutation and then performing autoregressive training, is computationally tractable: predicting (1) given ==1 ,.... ". (i- is trivial when S N as the tokens are independent, and computa- tionally tractable when i >> N because x*(( only depends (NV
3.1. Benign distributions with hard masking problems
x: (1) ,...,, (NN) and is efficiently learnable by assump- tion. In contrast, below we will show examples where if one performs order-agnostic training à la MDMs, one will run into hard masking problems with high probability.
We now describe a simple model of data under which we explore the computational complexity of masking problems. (L&O) distribu- tion is a data distribution Pdata over sequence of length L with alphabet size m (precisely, Pdata is over {O..... m)4 is specified by a permutation æ over indices {1, 2, ........... L). number of latent tokens N, number of observation tokens P such that N + P L. prior distribution Ppeior of latent vari- ables over (1...., m) and efficiently learnable observation functions O,...., Op: {l, m + A(!O...., m,)),
Definition 3.1. A ceteett-aandobbeerattti
First note that if the observations (O,,...,Op are given by a cryptographic hash function, then the masking problem of predicting (2=*(1, (L)) given (x+(++1)) ....., I. is computationally intractable by design because it requires inverting the hash function. While this is a well-known folklore observation regarding the role of token ordering in language modeling, it is not entirely satisfying because this construction worst-case in nature in real-world data, one rarely trains on sequences given by cryptographic hash functions. Furthermore, only estab- lishes hardness for a specific masking pattern which need not be encountered in the course of running the reverse process.
(Latent tokens) For i = L....,,,,..,...... N, sample (1) indepen- dently from the prior distribution Pprior of the latents.
(Observation tokens) For j j- = 1,... P. sample ((N+)) independently from O, (x ...., x. (M)).
We provide several simple instances of L&O distributions that address these issues: instead of leveraging delicate cryp- tographic constructions, they are average-case in nature and furthermore we can establish hardness for typical masking
Here efficiently learnable is in the standard PAC sense: given polynomially many examples of the form (2, y) where x² and y eu (z), there is efficient algorithm that can w.h.p. learn to approximate O2 in expectation over x²"
4

---

Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffussios
problems encountered along the reverse process.
likelihood model log
given as follows:
In all these examples, the hardness results we establish hold even if the algorithm knows all of the parameters of Pdatt as well as the observation functions OI...., Op. Due to space constraints, here we focus on the following example, deferring two others to Apps. B.1 and B.2.
4-1 log po = Σ log pe (z x1) |000 pof L 10--1) 13]) (3)
In other words, the т T-learner predicts the token at position ((i) given the clean tokens z (0(0) ..... 20 (i-1)) and masked To (L-1)) If i is the identity permutation, this reduces to the standard (left-to-right) autoregressive model. Note that the MDM loss encodes a æ--earner for every permutation because the MDM loss (1) is equivalent to the average loss of those r-learners over IT sampled from ULif(S
Example 3.2 (Sparse predicate observations). Consider the following class of L&O distributions. Given arity k 2 2, fix predicate function 9 .......m}+ - {0. 1}. Con- sider the set of all ordered subsets of 21,2....,A N of size k and set the total number of observation latents P equal to the size of this set (hence P = N!/(N - k)! = (N 11)......0 (N - k + 1)). To sample a new sequence, we first sample latent tokens x=(11) ,...,," from the prior distribution Pprior and an observation latent corresponding to a k-sized subset S is given by g g((x=() YeEs). In other words, each observation latent corresponds to a k-sized sub- set S of {1, (1,2...., N} and the corresponding observation function Os(x=(1),,... ....),.,,,,,), is given by gi (z)(() Jes)).
tokens 20 (6)
Co -.g, Pag--Paaaa
Σ log pe (o0 (1000
zo1TTi. .i, L 131]
where Sr denotes the set of all permutations over {0, (0......,Z 1}. The proof of the above equivalence is given in Appendix F. Therefore, by measuring the 'hardness' of each r-learner, we can probe differences in hardness be- tween arbitrary masking problems and left-to-right masking problems.
Proposition 3.3. Let x be a sample from an L&O dis- tribution Pdata with sparse predicate observations as de- fined in Example 3.2, with arity k and predicate g satis- fying Assumption B. and let 7 be the probability that g is satisfied by a random assignment from (l....,m}. m]* Let DKS and Deand be some constants associated with the predicate function 9 (see Definition 12). Suppose each token in x is independently masked with probability o, and M is the set of indices for the masked tokens. If 1->- DDks/ OKkskkkk> /Nrk-1 aa sa-1-- Deona/ /aak-/N then under the IRSB cavity prediction (see Conjecture B. B.13), with probability (1) over the randomness of the mask- ing, no polynomial- time algorithm can solve the resulting subproblem of predicting any of the masked tokens among (=(1) x (N) given x [MM
Experimental setup. We use the Slimpajama dataset (Soboleva et al., 2023) to evaluate the performance of train- ing in different orders. To train a r-learner, we employ a transformer with causal attention and use permuted data т (zo) as input. By varying π while maintaining all other training configurations (e.g.. model, optimization), we can use the resulting likelihood (computed using Equation (3)) as a metric to capture the hardness of subproblems solved by the m-learner.
In our experiments, the sequence length L is approximately 103, so repeating the above for each π is infeasible. Instead, we sample π ~ Unif(S;) and examine the scaling law of the r-learner's likelihood. We leverage the codebase from (Nie al., 2024), where the baseline scaling laws of MDM and ARM were introduced. Moreover, given that RoPE has an inductive bias towards left-to-right ordering, we employ leamaale positional embedding layer for all experiments to correct this. Consequently, we also re-run the baseline results, where RoPE was employed. To investigate how the distance between и and the identity permutation affects the scaling law, we sample πT from other distributions interpo- lating between (S.) and the point mass at the identical permutation. Further experimental details are provided in Appendix C.I.
The complete proof of the proposition is given in Ap- pendix B.4. We also provide proof outline in Appendix B.3 for a comprehensive understanding.
3.2. Empirical evidence of hardness via likelihoods
Recent studies (Nie et al., 2024; Zheng et al., 2024) have shown that masked diffusion models (MDMs) underperform compared to autoregressive models (ARMs) on natural text data. In this section, we provide evidence that this perfor- mance gap is primarily due to the order-agnostic training of MDMs. Since natural text follows a left-to-right token order, we demonstrate that as training deviates from this order, model performance gradually deteriorates.
As shown in Fig. 2, the scaling law for rllemnner with uniformly random ι is worse than that of an ARM. This elucidates the inherent hardness of masking problems po zo [M] beyond left-to-right prediction and also explains why MDM, which is trained simultaneously on
Results.
To understand the importance of the order during the train- ing, we use the following setting: Given a permutation π of indices (0, 1,.... I. - 1), define -laarnr to be a
5

---

Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffusions
all π Sz, is worse than ARM in likelihood modeling- Additionally, as gets closer the identity permutation, the scaling laws also get closer to ARM (7-leameer-oloser and π r-learmer-mucct the common belief that ARM is good fit for text data as it inherently follows a left-to-right ordering.
next section, we show that while vanilla MDM inference can run into such subproblems, simple modifications at the infer- ence stage can effectively circumvent these issues, resulting in dramatic, training-free performance improvements.
in Fig. 2). This also supports
4. MDMs can plan around hard problems
That said, it should also be noted that even though MDMs are trained on exponentially more masking problems than ARM (O(22) versus L), its performance is not signifi- cantly worse than x-learners. We attribute this to the blexs- ing of task diversity, multi-task training can benefit both the optimization dynamics (Kim et al., 2024) and validation performance (Tripuraneni et al., 2021; Maurer et al., 2016; Ruder, 2017) due to positive transfers across tasks.
We previously argued that due to the complex nature of masking subproblems, MDM must perform poorly cer- tain ones (= Therefore, during vanilla MDM infer- ence, MDM inevitably encounters such difficult subprob- lems at Step (b). While this might suggest that we need to fundamentally revisit how MDMs are trained, in this section we show that, surprisingly, simple modifications at the inference stage-without any further training- can sidestep these issues and lead to significant performance improvements.
3.3. Error is imbalanced across masking problems
In previous sections, we have demonstrated that the hard- ness of different masking problems po(2 wo [M]) can vary significantly, potentially hindering the MDM's learn-
MDM offers multiple sampling paths. The vanilla MDM inference (Algorithm 1) aim to align the intermediate distributions with the forward process, as used in continu- ous diffusion. However, unlike continuous diffusion, the reverse process of MDM allows multiple valid sampling paths (different orders of unmasking the tokens) that match the starting distribution of the forward process of MDM.
ing. In this section, we provide empirical evidence that the MDM's final performance exhibits similar imbalance across subproblems. Details are provided in App. C.2.
L&O-NAE-SAT. Consider an L&O distribution with π given by the identity permutation and where each observa- tion O2 is deterministically given by NAE(Z,, Big Tig) A 1 - 1 [x; = Ii, = zigl for some randomly chosen (pre- fixed) triples (f, 12, da) E [N]. For an MDM trained on this distribution, we measure the error achieves on each task log pa (2o [M]] via Ezo log pe (2o0|/o0
We first show that when we have an ideal MDM that per- fectly solves all masking problems, i.e., pa (z) [ [o] Pasta (zr) [2 [M]), then using any sampling path (unmask- ing the tokens in any order) results in the same distribu- tion. Consider the following sampler: For every step, S is a set with one index selected agnostically (with- out following any distribution). For any clean sam- plc zo generated by this sampler, note that po(zo) I po (=o =({)
(M))
log Pdata (xatl=oo/z [=0)== p| where Pdata ( (oo[oo(( denotes the Bayes-optimal predictor. Technically, we do not have access to this, so instead we train another MDM for much larger number of iterations and use this as proxy. Fig. 2 reveals that prediction tasks for latent positions (light re- gion) exhibit larger errors compared to those for observation positions (dark region).
1J) I-))) by chain rule, and this is equal to I1 Pdata 7(1))
................................
1nl)
(zo) Therefore, other choices of S, not necessarily following Algorithm 1, still capture the true likelihood.
In practice, unlike this ideal case, MDM does not perform equally well on all subproblems, as shown in Section 3.3. Consequently, different sampling paths result in varying likelihood modeling abilities. Motivated by this observation, we consider adaptive inference for MDMs:
Text. Here we revisit the text experiment from Section 3.2. Since we do not have access to the Bayes-optimal predictor, we use the metric 1
Exze<Paota D log (I xo -.)).
This captures the accumulation of error across since ()) Pdata Paata (#o| [M] minimizes this metric. Fig. 2 shows clear gap between different subproblems.
Adaptive MDM inference
subproblems Pe (=)
(a) Sample a set of masked tokens S = F (0, C
{ x = 0}.
(b) For each i E S, sample oj ~ (1(x=111)))
The theoretical and empirical evidence demonstrates that MDMs perform better in estimating (0 [M] for some subproblems M than for others. We therefore want to avoid encountering hard subproblems M at inference time. In the
Instead of selecting S randomly, adaptive MDM inference leverages an oracle F(O, vx) to select S strategically to
6

---

Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffussins
Gen PPL (arrllaa inference! PPL (Adaptive inforeco) 5-4
Table 2. Comparison of accuracy for solving the Sudoku puzzle.
Entrapy (Adeptive
Method
Param
Accuracy
22
Generative Perplexity
ARM (w/o ordering)
9.73%
42M
20-
Entropy
ARM (with ordering)
87.18%
-6-
MDM (vanilla)
6.88%
MDM (Top- K probability)
6M
18.51 %
16-
MDM (Top- K prob. margin)
89.49%
14
Top-K probability (Zheng et al., 2023). Suppose we want to unmask K positions at time step t, i.e., select [S] = K. In the Top- strategy, the uncertainty of a po- sition is estimated by the maximum probability assigned to any value in the vocabulary. More precisely. the cer- tainty at position is max Po(x* = j]æ,) and F(0, xx) = Top K (max po
352 500
1000
1250
1500
1752
2000
Sampilng Steps
Figure 3. Generative Perplexity. We compare the resulting gener- ative perplexity (GenPPL) of adaptive vanilla MDM inference. We employ a pretrained 170M MDM and LLaMA-7B (Touvron al., 2023) inference and evaluation, respectively. Adaptive MDM inference (Blue) leads to substantial reduction in genera- tive perplexity, while maintaining the entropy.
Top-K K strategy is good proxy for many tasks and works well in practice (Zheng et al., 2023; Ye et al., 2024; Wang et al., 2024). However, this approach can often provide mis- leading estimates of uncertainty. Consider when an MDM is confused between two token values, thus assigning them almost equal but high probabilities. In this case, Top-K strategy may still choose to unmask this position, despite its uncertainty. To mitigate this issue, we propose the following alternative strategy.
avoid hard masking problems. This naturally raises the question of how to design an effective oracle F.
In the following sections, demonstrate that adaptive MDM inference with careful choices of F enhance MDM's likelihood matching ability. In other words, a pretrained MDM, even if it performs poorly on certain hard subprob- lems, still contains sufficient information to avoid them when paired with an effective oracle F.
Top-K probability margin. In this strategy, the uncer- tainty of a position is instead estimated using the ab- solute difference between the two most probable values at position i. More precisely, if ji and j/ are the two most probable values in vocabulary according to pa (x*)2))) in position i, the certainty in the position is given by [pe(z = j1(=,) - pe = f2|=x1) and F(0, xx) Top K (IPo = =1(=))) po = jo xz)). When multi- ple values have similar probabilities at position, Top-K probability margin will provide a better estimate of the un- certainty of a position, and when there is a single best choice of value then Top-K K and Top-K K probability margin work similarly.
4.1. Effective design of ordering oracle
We introduce two different oracles, Top-K and Top-K K prob- ability margin. Intuitively, both strategies are based on the idea that 5 should be selected based on how "certain" the model is about each position. We caution that these strategies should not be confused with notions like nucleus sampling in ARMs (Holtzman et al., 2019); the oracles we describe are for selecting the position of the next token to decode, rather than the value, and thus are only meaningful in the context of MDMs.
Table 1. L&O-NAE-SAT. Adaptive MDM inference achieves bet- ter likelihood matching than vanilla MDM inference. Note that naive guessing leads to 75% accuracy, indicating that vanilla infer- ence performs similarly or worse than naive guessing.
4.2. Adaptive MDM inference
In this section, we experimentally validate that adaptive MDM inference helps MDMs avoid hard subproblems, lead- ing to better likelihood matching. We first show our results L&O-NAE-SAT and text data, before turning primary application to logic puzzles.
(N, P)
Vanilla inference
Adaptive inference
(25, 275)
78.06%
93.76%
(30, 270)
75.70%
93.54%
(40, 260)
74.60%
92.21%
L&O-NAE-SAT and text data. For the L&O-NAE-SAT distribution defined in Section 3.3, we evaluate the effec- tiveness of adaptive inference by measuring the accuracy in predicting the observation tokens. Table 1 in the appendix
(50, 250)
67.94% 94%
90.01%
(100, 200)
62.84%
88.91%
7

---

Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffussins
reveals clear improvement over vanilla inference. For the text dataset, we evaluate using the standard metric of generative perplexity, by which likelihood is measured by a large language model. We also compute the entropy of the generated samples to ensure both inference strategies exhibit similar levels of diversity. As shown in Fig. 3, we observe a substantial decrease in generative perplexity using adaptive inference. We defer further experimental details to Appendix D.1.
To do so, we will compare the performance of adaptive MDM inference to that of ARM on Sudoku and Zebra puz- zles. For these puzzles, the natural order of generation is not only different from left-to-right, but it is also sequence- dependent. For such tasks, prior works have shown that ARMs struggle if the information about the order is not pro- vided during the training (Shah et al., 2024; Lehnert et al., 2024). Therefore, to obtain a strong baseline, we not only consider an ARM trained without the order information but also consider an ARM trained with the order information for each sequence in the training data. Note that the latter is much stronger baseline than the former as one can hope to teach the model to figure out the correct order by some form of supervised teacher forcing (as performed in Shah et al. (2024); Lehnert et al. (2024)), eliminating the issue of finding the right order in an unsupervised manner.
Logic puzzles. We consider two different types of logic puzzles: Sudoku and Zebra (Einstein) puzzles. Intuitively, for Sudoku, some empty (masked) cells are significantly easier to predict than others and we want to choose the cells that are easier to predict during the inference. We evaluate the effectiveness of adaptive MDM inference over vanilla MDM inference in selecting such cells.?
We compare ARMs and MDMs for Sudoku in Table 2 and Zebra puzzles in Table 3. We observe that for both, Top-K K probability margin- based adaptive MDM inference not only outperforms the ARM trained without ordering information, but even oupprforrs the ARM trained with ordering infor- mation! This shows that the unsupervised way of finding the correct order and solving such logic puzzles using adaptive MDM inference outperforms the supervised way of finding the correct order and solving such puzzles using an ARM, and is significantly less computationally intensive.
To measure the performance of an inference method, we use the percentage of correctly solved puzzles. For both puzzles, we use train and test datasets from (Shah et al., 2024). For the Sudoku puzzle (Table 2) we observe that adaptive MDM inference, particular Top-K probability margin, obtains substantially higher accuracy (89.49%) compared vanilla MDM inference (6.88%). Additionally, Top-K probabil- ity margin obtains higher accuracy (89.49%) than Top-K (18.51%). As mentioned in Section 1. this is because Top- K probability margin more reliably estimates uncertainty when multiple competing values are close in probability at a given position, as often the case in Sudoku. For the Zebra puzzle, as shown in Table 3, we observe consis- tent result: Top-K (98.5%) and Top-K probability margin (98.3%) outperform vanilla MDM inference (76.9%).
4.4. Adaptive MDM inference on text benchmarks
To examine the effect of different inference strategies on text benchmarks, we adapted LLaDA, the 8B MDM model from (Nie et al., 2025). We compare three inference strategies: Vanilla, Top-K K probability, and Top- K probability margin. The results are presented in Table 4.
Table Comparison accuracy for solving the Zebra puzzle.
We see that both adaptive MDM inference strategies, Top- K probability and Top- probability margin, , consistently outperform vanilla MDM inference. Notably, Top-K K proba- bility margin demonstrates a clear advantage over Top- in challenging tasks like HumanEval- Mulliiliee, HumanEval- Split Line, and Math. This is because Top-k Prob. Margin provides a more reliable estimate of uncertainty when multi- ple tokens have similar probabilities, frequent occurrence in these difficult tasks. These results further underscore the potential for developing new, sophisticated adaptive infer- ence strategies for various tasks.
Method
# Param
Accuracy
ARM (w/o ordering)
80.31 %
42M
ARM (with ordering)
91.17 %
MDM (vanilla)
76.9 %
MDM (Top-K probability)
19M
98.5 %
MDM (Top-K prob. margin)
98.3 %
4.3. Eliciting sequence-dependent reasoning paths using adaptive MDM inference in logic puzzles
In this section, we study the effectiveness of adaptive MDM inference in finding the right reasoning/generation order for tasks where every sequence has a different "natural" order.
4.5. Easy to hard generalization
In the previous section we showed that when the training and inference sequences come from the same distribution, order-agnostic training of MDMs combined with adaptive inference can perform very well on logic puzzles. To evalu- ate if the model has learned the correct way of solving the puzzles and test the robustness of adaptive inference, we
A prior work (Ye et al., 2024) reported that a 6M MDM with Top- K inference achieves 100% accuracy on Sudoku. Given that a 6M MDM with Top-K only achieves 18.51% on our dataset (Table 2). this suggests that the Sudoku dataset in (Ye al., 2024) is significantly easier than ours.
8

---

Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffussios
Table 4. Performance of different inference strategies for LLaDDa 8B model on coding and math tasks.
Method
HumanEval-MMult Hinnnevvl-Singge
HumanEval-Split
Math
MMLU
ROCStories
Vanilla
31.8%
16.5%
14.2%
28.5%
33.2%
21.23%
Top-k
32.9%
20.8%
18.4%
31.3%
36.5%
21.10%
Top-k Margin
33.5%
25.4%
22.3% %
34.3%
35.4%
21.41 %
An important direction for future work is to explore settings beyond logic puzzles where adaptive inference can help MDMs match surpass ARMs. For these, it may be crucial to go beyond the relatively simple adaptive strategies like Top- and Top-K K probability margin considered here.
Table 5. Comparison of accuracy for solving the hard Sudokus,
Method
Param
Accuracy
ARM (with ordering)
42M
32.57 %
MDM (random)
3.62 %
MDM (Top-K probability)
6M
9.44 %
Acknowledgements. JK thanks Kiwhan Song for discus- sions about MDM training. KS and VK are supported by the NSF AI Institute for Foundations of Machine Learning (IFML). KS thanks Nishanth Dikkala for the initial discus- sions about the project. SC is supported by the Harvard Dean's Competitive Fund for Promising Scholarship and thanks Brice Huang and Sidhanth Mohanty for enlighten- ing discussions about computational tradeoffs for planted CSPs.
MDM (Top-K prob. margin)
49.88 %
also test the MDMs on harder puzzles than the ones from training, for Sudoku.
We keep the training dataset the same as proposed in Shah et al. (2024). Shah et al. (2024) created this dataset from Radcliffe (2020) by selecting the puzzles that can solved using fixed strategies and do not require backtracking- based search. We use the remaining puzzles in Radcliffe (2020) our hard dataset. Hence, these puzzles all use strategy not seen during training and/or backtracking to obtain the correct solution.
Impact statement
This paper advances the understanding of discrete diffu- sion models, contributing to the broader field of Machine Learning. There are many potential societal consequences of our work, none of which we feel must be specifically highlighted here.
We measure the accuracy of MDMs and ARMs on the hard test set and present the results in Table We see that the Top-K probability margin-based adaptive MDM inference strategy (49.88 (8%%) again significantly outperforms ARMs trained with order information (32.57%). In particular, al- though the accuracy drops for both methods due to the more challenging test set, MDMs with adaptive inference appear to be more robust to this distribution shift than ARMs. We believe this is due to the fact that MDMs try to solve a sig- nificantly higher number of infilling problems than ARMs (exp(T.) compared to T.) and therefore are able to extract knowledge about the problem more efficiently than ARMs.
References
Alaoui, A. E. and Gamarnik, D. Hardness of sampling solutions from the symmetric binary perceptron. arXiv preprint arXiv:2407.16627 2024.
Alekhnovich, M. More average case vs approximation complexity. In 44th Annual IEEE Symposium Foun- dations of Computer Science, 2003. Proceedings., pp. 298-307. IEEE, 2003.
5. Conclusion
Aubin, B., Perkins, W., and Zdeborová, L. Storage capacity in symmetric binary perceptrons. Journal of Physics A: Mathematical and Theoretical, 52929):294003, 2019.
In this work, we examined the impact of token ordering on training and inference in MDMs. We provided theoretical and experimental evidence that MDMs train on hard mask- ing problems. We also demonstrated that adaptive inference strategies can be used to sidestep these hard problems. For logic puzzles, we find that this leads to dramatic improve- ments in performance not just over vanilla MDMs, but even over ARMs trained with teacher forcing to learn the right order of decoding.
Austin, J., Johnson, D. D., Ho, J., Tarlow, D., and van den Berg, R. Structured denoising diffusion models in discrete state-spaces. NeruPPS, 2021.
Bormashenko, O. A coupling argument for the random transposition walk. arXiv preprint arXiv: 1109.3915, 2011.
9

---

Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Difffusioss
Chang, H., Zhang, H., Jiang, L., Liu, C., and Freeman, W. T. Maskgit: Masked generative image transformer. CVPR, 2022.
Hoogeboom, E., Nielsen, D., Jaini, P., Forré, P., and Welling, M. Argmax flows and multinomial diffusion: Learning categorical distributions. NeurIPS, 2021b.
Chen, H. and Ying, L. Convergence analysis of discrete dif- fusion model: Exact implementation through uniformiza- tion, arXiv preprint arXiv: 2402.08095, 2024.
Kaplan, J., McCandlish, S., Henighan, T., Brown, T. B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., and Amodei, D. Scaling laws for neural language models. arXiv preprint arXiv:2001.088: 08361, 2020.
Chen, X., Chi, R. A., Wang, X., and Zhou, D. Premise order matters in reasoning with large language models. arXiv preprint SiLV:22002008939, 2024.
Kim, J., Kwon, S., Choi, J. Y., Park, J., Cho, J., Lee, J. D., and Ryu, E. K. Task diversity shortens the icl plateau. arXiv preprint arXiv:2410.05448 2024.
Decelle, A., Krzakala, F., Moore, C., and Zdeborová, L. Asymptotic analysis of the stochastic block model for modular networks and its algorithmic applications. Phys. Rev. E, 84:066106, Dec 2011.
Kitouni, O., Nolte, N. S., Williams, A., Rabbat, M., Boucha- court, D., and Ibrahim, M. The factorization curse: Which tokens you predict underlie the reversal curse and more. Advances in Neural Information Processing Systems, 37: 1 112329-112353 12355, 2025.
Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K. BERT: Pre-training of deep bidirectional transformers for lan- guage understanding. In Proceedings of the 2019 Confer- ence of the North American Chapter of the Association for Computational Linguistics: Human Language Technolo- gies, Volume I (Long and Short Papers), pp. 1171-4186, 2019.
Krzakala, F. and Zdeborová, L.. Hiding quiet solutions in random constraint satisfaction problems. Physical review letters, 102((3):238701, 2009.
Lehnert, L., Sukhbaatar, S., Su, D., Zheng, Q., Mc P., Rabbat, M., and Tian, Y. Beyond a Better planning with transformers via search dynamics bootstrapping. 2024.
Gamarnik, D. The overlap gap property: A topological bar- rier to optimizing over random structures. Proceedings of the National Academy of Sciences, 1 118441))2108849218,, 2021.
Liao, Y., Jiang, X., and Liu, Q. Probabiliistially masked language model capable of autoregressive generation in arbitrary word order. Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pp. 263-27222222 Association for Computational Linguistics, 2020.
Golovneva, O., Allen-Zhu, Z., Weston, J., and Sukhbaatar, S. Reverse training to nurse the reversal curse. arXiv preprint arXiv:2403. 13799, 2024.
Liu, A., Broadrick, O., Niepert, M., and Broeck, G. V. d. Dis- crete copula diffusion. arXiv preprint arXiv:2410.01949, 2024a.
Gong, S., Agarwal, S., Zhang, Y., Ye, J., Zheng, L., Li, M., An, C., Zhao, P., Bi, W., Han, J., et Scaling diffu- sion language models via adaptation from autoregressive models. arXiv preprint 1rXiv:2410.1 17891, 2024.
Liu, S., Mohanty, S., and Raghavendra, P. On statistical inference when fixed points of belief propagation are unstable In 2021 IEEE 62nd Annual Symposium on Foundations of Computer Science (FOCS), pp. 395-405, IEEE Computer Society, 2022.
Ho. J., Jain, A., and Abbeel, P. Denoising diffuuson proba- bilistic models. Advances in neural information process- ing systems, 33:6840-683 66000,6811110 2020.
Liu, S., Nam, J., Campbell, A., Stärk, H., Xu, Y., Jaakkola, T., and Göomez-Bombarelee R. Think while you generate: Discrete diffusion with planned denoising. arXiv preprint arXh::2410.0626. 2024b.
Hoffmann, J., Borgeaud, S., Mensch, A., Buchatskaya, E., Cai, T., Rutherford, E., Casas, D. d. L., Hendricks, L. A., Welbl, J., Clark, A., et al. Training compute-optimal large language models. arXiv preprint arXiv:2203115556, 2022.
Loshchilov, and Hutter, Decoupled weight decay regu- larization. arXiv preprint 1711::111.0050000 2017.
Holtzman, A., Buys, J., Du, L., Forbes, M., and Choi, Y. The curious case of neural text degeneration. arXiv preprint arXiv:1904.09751, 2019.
Lou, A., Meng, C., and Ermon, S. Discrete diffusion model- ing by estimating the ratios of the data distribution. ICML, 2024.
Hoogeboom, E., Gritsenko, A. A., Bastings, J., Poole, B., Berg, R. v. d., and Salimans, T. Autoregressive diffusion models. arXiv preprint arXiv:2 1 10.02037, 2021a.
Maurer, A., Pontil, M., and Romera-Paredes, B. The benefit multitask representation learning. JMLR, 17(81):1 1-11)) 2016.
10

---

Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Difffusions
Montanari, A. Estimating random variables from random sparse observations. European Transactions Telecom- munications, 19(4)38554403 2008.
Shih, A., Sadigh, D., and Ermon, S. Training and infer- ence any-order autoregressive models the right way. NeurIPS, 2022.
Nie, S., Zhu, F., Du, C., Pang. T., Liu, Q. Zeng, G., Lin, M., and Li, C. Scaling up masked diffusion models text. arXiv preprint ar arXiv:2410. 18500.88544,4 2024.
Soboleva, D., Al-Khateeb, F., Myers, R., Steeves, J. R., Hes- tness, J., and Dey, N. Slimpajama: A 627b token cleaned and deduplicated version of redpajama, June 2023.
Nie, S., Zhu, F., You, Z., Zhang, X., Ou, J., Hu, J., Zhou, J., Lin, Y., Jen, and Li, C. Large language diffusion models. arXiv preprint UrXiv:2502.09992, 2025.
Sohl-Dickstein, J., Weiss, F. A., Maheswaranathan, N., and Ganguli, S. Deep unsupervised learning using nonequi- librium thermodynamics. ICML, 2015.
Ou, J., Nie, S., Xue, K., Zhu, F., Sun, J., Li, Z.,, and Li, C. Your absorbing discrete diffusion secretly models the conditional distributions of clean data. arXiv preprint arXiv:2406.03736, 2024.
Song, Y., Sohl-Dickstein, J., Kingma. D. P., Kumar, A., Er- mon, S., and Poole, B. Score-based generative modeling through stochastic differential equations. ICLR, 2021.
Touvron, H., Martin, L., Stone, K., Albert, P., Almahairi, A., Babaei, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., Bikel, D., Blecher, I.., Ferrer, C. C.,, Chen, M., Cucurull, G., Esiobu, D., Fernandes, J., Fu, J., Fu, W., Fuller, B., Gao, C., Goswami, V., Goyal, N., Hartshorn, A., Hosseimi, S., Hou, R., Inan, H., Kardas, M., Kerkez, V., Khabsa, M., Kloumann, I., Korenev, A., Koura, P. S., Lachaux, M.-A., Lavril, T., Lee, J., Liskovich, D., Lu, Y., Mao, Y., Martinet, X., Mihaylov, T., Mishra, P., Molybog, I. Nie. Y., Poulton, A., Reizenstein, J., Rungta, R., Saladi, K., Schelten, A., Silva, R., Smith, E. M., Subramanian, R., Tan, E.,, Tang. B., Taylor, R., Williams, A., Kuan, J. X., Xu, P., Yan, Z., Zarov. I., Zhang, Y.. Fan, A., Kambadur, M., Narang, S., Rodriguez, A., Stojnic, R., Edunov, S., and Scialom, T. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv: 2307.09288 2023.
Papadopoulos, V., Wenger, J., and Hongler, C. Arrows of time for large language models. arXiv preprint arXiv:2401. 17505, 2024.
Peng. F. Z., Bezemek, Z., Patel, S., Yao, S., Rector- Brooks, J., Tong, A., and Chatterjee, P. Path planning for masked diffusion model sampling. arXiv preprint arXiv:2502.035.2 2025.
Radcliffe, D. G. 3 million sudoku puzzles with rat- ings, 2020. URL
kww.la. com/
dsv/1495975.
Rector-Brooks, J., Hasan, M., Peng, Z., Quinn, Z., Liu, C., Mittal, S., Dziri, N., Bronstein, M., Bengio, Y., Chatterjce, P., et al. Steering masked discrete diffusion models via discrete denoising posterior prediction. arXiv preprint arXiv:2410.08134 2024.
Tripuraneni, N., Jin, C., and Jordan, M. I. Provable meta- learning of linear representations. ICML, 2021.
Ruder, S. An overview of multi-task task learning in deep neural networks. arXiv 1706.05096 2017.
Varma, H., Nagaraj. D., and Shanmugam, K. Glauber gen- erative model: Discrete diffusion models via binary clas- sification. arXiv preprint arXiv: 2405. 70355, 2024.
Sahoo, S., Arriola, M., Schiff, Y., Gokaslan, A., Marroquin, E., Chiu, J., Rush, A., and Kuleshov, V. Simple and effective masked diffusion language models. Advances in Neural Information Processing Systems, 37:130136- 130184, 2025.
Wang, X., Zheng, Z., Ye, F,, Xue, D., Huang, S., and Gu, Q. Diffusion language models are versatile protein learners. ICML, 2024.
Schiff, Y., Sahoo, S. S., Phung, H., Wang. G., Boshar, S., Dalla-torre, H., de Almeida, B. P., Rush, A., Pierrot, T., and Kuleshov, V. Simple guidance mechanisms for dis- crete diffusion models. arXiv preprint arXiv:2412.10199 10193, 2024.
Xu, M., Geffner, T., Kreis, K., Nie, W., Xu, Y., Leskovec, J., Ermon, S., and Vahdat, A. Energy-based diffusion language models for text generation. arxiv preprint arXiv: 2410.21357, 2024.
Shah, K., Dikkala, N., Wang, X., and Panigrahy, R. Causal language modeling can elicit search and rea- soning capabilities on logic puzzles. arXiv preprint arXiv:2409. 10500, 2024.
Ye, J., Gao, J., Gong. S., Zheng, L., Jiang, X., Li, Z., and Kong, L. Beyond autoregression: Discrete diffusion for complex reasoning and planning. arXiv preprint arXiv: 2410..41577 2024.
Shi, J.. Han, K., Wang, Z., Doucet, A., and Titsias, M. K. Simplified and generalized masked diffusion for discrete data. NeurlPS, 2024.
Zhang, P., Zeng, G., Wang, T., and Lu, W. Tinyllama: An open-source small language model. arXiv preprint arXiv: 2401.02385, 2024.

---

Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffussins
Zheng, K., Chen, Y., Mao, H., Liu, M.-Y., Zhu, J., and Zhang, Q. Masked diffusion models are secretly time- agnostic masked models and exploit inaccurate categori- 2024.
cal sampling. arXiv preprint arXiv:2409,00
Zheng, L.,, Yuan, J., Yu, L., and Kong, L. A reparameter- ized discrete diffusion model for text generation. arXiv preprint arXiv:2302.057777 2023.
12

---

Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffusions
A. Related works
Discrete diffusion models. (Continuous) diffusion models were originally built on continuous-space Markov chains with Gaussian transition kernels (Sohl-Dickstein et al., 2015; Ho al., 2020). This was later extended to continuous time through the theory of stochastic differential equations (Song et al., 2021). In similar vein, discrete diffusion models have emerged from discrete- seretee--pceeaee Markov chains (Hoogeboom et al., 2021b). Specifically, (Austin et al., 2021) introduced D3PM with various types of transition matrices. Later, Lou et al. (2024) proposed SEDD, incorporating a theoretically and practically robust score-entropy objective. Additionally, Varma ct al. (2024); Liu et al. (2024b) introduced novel modeling strategies that classify tokens in a noisy sequence as either signal (coming from clean data) or noise (arising from the forward process). In particular, Liu et al. (2024b) uses this to give a planner that adaptively determines which tokens to denoise. While this is similar spirit to our general discussion about devising adaptive inference strategies, we emphasize that their approach specific discrete diffusions for which the forward process scrambles the token values, rather than masking them.
common choice due to its better performance than other kernels. Building on this, Sahoo et al. (2025); Shi et al. (2024) aligned its framework with continuous diffusion, resulting in a simple and principled training recipe, referring to it as Masked Diffission Model. Subsequent studies have explored various aspects of MDM. Gong et al. (2024) efficiently trained MDM via adaptation from autoregressive models, scaling MDM up to 7B parameters. Zheng et al. (2024) interpreted MDMs as order-agnostic learners and proposed a first-hitting sampler based on this insight. Ye et al. (2024); Gong et al. (2024) demonstrated that MDM outperforms autoregressive models in reasoning and planning tasks, emphasizing its impact on downstream applications. Nie et al. (2024) examined the scaling laws of MDM, while Xu et al. (2024); Liu et al. (2024a) identified limitations in capturing coordinate dependencies when the number of sampling steps is small and proposed additional modeling strategies to address this issue. Schiff et al. (2024) studied conditional generation using MDM and Rector-Browd et al. (2024) tackled the challenge of controlling generated data distributions through steering methodologies. Chen & Ying (2024) provided a theoretical analysis showing that sampling error is small given accurate score function estimation.
Masked diffusion models. Meanwhile, the absorbing transition kernel has gained popularity
Any-order reasoning. Even though language tasks generally have natural order of "left-to-right" token generation, in many tasks like planning, reasoning, and combinatorial optimization, the natural order of token generation can be quite different from "left-to-right.. Even though prominent performance on various tasks, many works (Golovneva et al., 2024; Chen et al., 2024; Kitouni et al., 2025) have shown that this performance is tied to the training order of the tasks and therefore can cause brittleness from it. For example, Chen et al. (2024) showed that simply permuting the premise order on math tasks causes a performance drop of 30%. The reason behind such brittleness regarding the ordering is the inherent "left-to-right" nature of the autoregressive models. Several works (Liao et al., 2020) have tried to address this issue in the autoregressive framework. In partieular, (Papadopoulos et al., 2024) highlighted the significance of left-to-right ordering in natural language by comparing its likelihood to that of the reverse (right-to-left) ordering.
seteeegressivv-bseee language models achieve impressive
Recently, discrete diffusion models have emerged as a promising approach for discrete data apart from autoregressive models. Additionally, the prder-agnostic training of discrete diffusion models opens up the multiple sampling paths during the inference but it also faces some challenges during the training therefore, they seem a promising approach to elicit any order reasoning. Zheng et al. (2023) proposed different ways of implementing an adaptive inference strategy for MDM but concrete understanding of why such an adaptive inference strategy is needed is still lacking. In this work, we explore various aspects of vanilla MDM training and how adaptive MDM inference can mitigate the issues raised by vanilla MDM training and elicit any order reasoning.
We also want to mention the concurrent work by Peng et al. (2025) that proposes an alternative adaptive inference strategy by seleeting F(0, 22)) based on the BERT model or the denoiser itself. In particular, Peng et al. (2025) uses the BERT model or the denoiser to obtain the uncertainty of token and then uses Top-K K to decide the positions to unmask it. In contrast to their work, we disentangle the impact of token ordering on MDM training vs. MDM inference and provide a more complete understanding of the motivations for and benefits of adaptive inference. Additionally, our results indicate drawbacks to using Top-K strategy as opposed to Top-K margin in deciding which tokens to unmask when there are multiple values with high probabilities.
Beyond autoregressive models. Efforts to learn the natural language using non non-autoregressive modeling began with BERT (Devlin et al., 2019). Non-causal approaches can take advantage of the understanding the text data representation.
13

---

Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffuusioss
(Chang et al., 2022) adopted a similar approach for learning image representations. Building on these intuitions, (Shih al., 2022; Hoogehoom et al., 2021a) proposed any-order modeling, which allows model to generate in any desired order. Shih et al. (2022) made the same observation that any-order models by default have to solve exponentially more masking problems than autoregressive models. However, whereas our work shows that learning in the face of this challenging task diversity can benefit the model at inference time, their work sought to alleviate complexity at training time by reducing the number of masking problems that need to be solved.
B. Technical details from Section 3
Notations. Throughout this section, we use x' to denote the i-th coordinate of the vector and =(f) to denote the j-th example. The i-th coordinate of the vector z(j) is denoted by z(j)..
B.1. Additional example: sparse parity observations
Example B.1 (Noisy sparse parity observations). Let m = E N, and N log N & P No.49k, Fix noise rate 1j >> 0 as well as strings 2(1)..... ((,)...,,,(P))) (PP) sampled independently and uniformly at random from the set of k-sparse strings in {0, 1).1.... For each j E [P], define (z(x)) to be the distribution which places mass 1 17 on 1 (resp. 2) and mass n on 2 (resp. 1) if Σ,
2. E, z(i) is odd (resp. even). Note that for l-- = O(1), each of these observations is efficiently learnable by brute-force.
Below we show that for a certain range of masking fractions, a constant fraction of the masking problems for the correspond- ing L&O distributions are computationally hard under the Sparse Learning Parity with Noise assumption (Alekhnovich, 2003). Formally we have:
be arbitrary absolute constant, and let l = 1/poly(N) (N) be sufficiently large. Let be sample from a L&O distribution Pdata with noisy parity observations as defined in Example B. Suppose each token is independently masked with probability a, and M is the set of indices for the masked tokens. If1 - 1/N S a S 1 then under the Sparse Learning Parity with Noise (SLPN) assumption (see Definition B.3), with constant probability over M, no polynomial-time algorithm can solve the resulting masking problem of predicting any of the masked tokens among (1) xT.) given x x.M).
Proposition B.2. Let O < a < a<<ll
We note that it is important for to take the observations to be sparse parities and to leverage the Sparse Learning Parity with Noise assumption. If instead we used dense parities and invoked the standard Learning Parity with Noise (LPN) assumption, we would still get the hardness of masking problems, but the observations themselves would be hard to learn, assuming LPN. This result is based on the following standard hardness assumption:
Definition B.3 (Sparse Learning Parity with Noise). Given input dimension N, noise parameter 0 < < 1/2, and sample size P. an instance of the Sparse Learning Parity with Noise (SLPN) problem is generated as follows:
Nature samples random bitstring x from {0, 1
We observe P examples of the form (a(1), v(i) where z(1) is sampled independently and uniformly at random from k-sparse bitstrings in (0, 1) and y is given by 6 + (z((), a) (mod 2), where e; is with probability n and 0 otherwise.
Given the examples ((x(1), (()) the goal is to recover
The SLPN assumption is that for any P = (1-p)) for constant 0 < p < 1, and any sufficiently large inverse polynomial noise rate 7, no poly(N) (N)-time algorithm can recover æ with high probability.
Proof of Proposition B.2. With probability at least 1 - (1- I/N) 2 9(1), all of the variable tokens =() for i < N are masked. Independently, the number of unmasked tokens among the observation tokens O, is distributed as Bin (P, 1 at), so by a Chernoff bound, with probability at least 1 RRP/N33 = 1 1/poly(N) we have that at least P/4N Q(N log N) observation tokens are unmasked. The masking problem in this case amounts to an instance of SLPN with input dimension N and sample size in [9t[N log N), 0(v9.49k) Because of the lower bound on the sample size, prediction of xM is cermmatio-thhooretialll possible. Because of the upper bound on the sample size, the SLPN assumption makes it computationally hard. As a result, estimating the posterior mean on any entry of M given the unmasked tokens is computationally hard as claimed.
14

---

Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffusions
B.2. Additional example: random slab observations
Example B.4 (Random slab observations). Let m and = N22 for constant 7 > Fix slab width and vectors e(.)....,,(( x(P) sampled independently from (O, I). For each j E (P], define the corresponding observation (z(z) be deterministically if (=()), 2r 1) S ByN, and deterministically 0 otherwise.
In (Alaoui & Gamarnik, 2024), it was shown that stable algorithms (Definition B.7), which encompass many powerful methods for statistical inference like low-degree polynomial estimators, MCMC, and algorithmic stochastic localiza- tion (Gamarnik, 2021), are unable to sample from the posterior distribution over a random bitstring conditioned on it satisfying (z(),, x) & 3/N for any ((V) number of constraints , x(P'), provided p' is not too large that the support of the posterior empty. This ensemble is the seul-studied symmetric perceptron (Aubin et al., 2019). The following
a direct reinterpretation of the result of (Alaoui & Gamarnik, 2024):
Proposition B.5. Let Pdata be L&O distribution with random slab observations defined in Example B.4, with parameter > O and slab width 3 >> O. There exists a constant > O such that for any absolute constant O < c < cg. if 1 - 1-c3Nj21 N2N I a 1 - cN/P and y > Cg. the following holds. Let Pasta denote the distribution given by independently masking every coordinate in Pataa with probubility a. Then any (1 on masked diffusion, which takes as input a sample x' from Paataa and, with probability 1 1 o(1) outputs a Wasserstein- approximate" sample from Pdata conditioned on the unmasked tokens in x', must run in super-polynomial time.
algorithm, even one not based
The upshot of this is that any stable, polynomial-time masked diffusion sampler will, with non-negligible probability, encounter computationally hard masking problem at some point during the reverse process.
For the proof, we first formally define the (planted) symmetric Ising perceptron model:
Definition B.6. Let 3 > O. The planted symmetric Ising perceptron model is defined as follows:
Nature samples uniformly at random from {=1}
For each j = 1, 1,... P L..N we sample =(j) independently from (0 Ix) conditioned on satisfying (=(()) s BVN.
The goal is to sample from the posterior on conditioned on these observations {=(i)){E
Next, we formalize the notion of stable algorithmss
Definition B.7. Given a matrix Z ~ (O, 1)@PXN, define Z. = t2 + -1 for independent Z' ~ (0, 1)8PXN randomized algorithm A which takes as input Z E PxN and outputs an element of [=1] N is said to be ---bbal if law law((())))
lim
As discussed at depth in (Gamarnik, 2021), many algorithms like low-degree polynomial estimators and Langevin dynamics are stable.
Theorem B.8 (Theorem 2.1 in (Alaoui Gamarnik, 2024)*). For any constant 3 >> 0, there exists 03 > O such that the following holds for all constants < a < cg. For iN S 52log (o(loog(n)/ne² (n)/n²), any t--stable randomized algorithm A which takes as input Z (=(1),...,,(P))))) z(P)) and outputs an element of {+1] will fail to sample from the posterior on conditioned on Z. in the symmetric Ising perceptron model to Wasserstein error of (NN).
Proof of Proposition B.5. By a union bound, with probability at least 1 - (1- > 1- N2 cp// over a draw x' ~ P'data- all of the 7() tokens are masked. The number of unmasked tokens in x' among the observations O, is distributed as Bin(P, - a.. By a Chernoff bound, this is in /3, N/4] with at least constant probability. The claim then follows immediately from Theorem B.8 above.
B.3. Proof outline of Proposition 3.3
To understand the proof idea, we consider the case where all the latent tokens are masked and some of the observation tokens are unmasked. In this case, the prediction task reduces to learning to recover the latent tokens that are consistent
'Here the notion of approximation is o(1 )-closeness in Wasserstein-2 distance.
"Note that while the theorem statement in (Alaoui & Gamarnik, 2024) refers to the non-planted version of the symmetric binary perceptron, the first step in their proof is to argue that these two models are mutually contiguous the regime of interest.
15

---

Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Difffusions
1.0
random init
0.9
0.5
overlap
0.7
0.5
0.5
0.4
55
D //k k
Figure 4. Overlap achieved by belief propagation initialized at ground truth versus random for planted CSP with & = 3, m = 3, and 9= NAE, for N 10000 and varying choices of average degree D. Dks/ K can shown analytically be 64, consistent with the phase transition depicted. Plot suggests Deand/ K x 50. By Prop. 3.3 this implics a range of masking fractions at which 21(1) fraction of masking problems are computationally hard.
with the observations. Intuitively, each observation provides some constraints and the task is to recover an assignment that satisfies the constraints. This is reminiscent of Constraint Satisfaction Problems (CSPs). Indeed, to show the hardness result, we use the rich theory developed for planted CSPs at the intersection of statistical physics and average-case complexity.
In a planted CSP, there is an unknown randomly sampled vector y of length N and, one is given randomly chosen Boolean constraints which y promised to satisfy, and the goal is to recover y as best as possible (see Definition B.9). Prior works have shown the hardness of efficiently learning to solve the planted CSP problem (Krzakala & Zdeborová, 2009; Alaoui & Gamarnik, 2024). We show the hardness of masking problems in L&O distributions based on these results. Consider the ground truth latent tokens the random vector y and each observation as constraint. In this case, the problem of learning to recover the latent tokens from the observation tokens reduces to recovery for the planted CSP.
There are precise predictions for the values of vocabulary size m and the number of observations for which the information- theoretically best possible overlap and the best overlap achievable by any computationally efficient algorithm are different. We show that these predictions directly translate to predictions about when masking problems become computationally intractable:
As simple example, let consider sparse predicate observations with k 2 and g(x',=") = 1[xx + x"). These can be formally related to the well-studied problem of planted m-coloring. In the planted m-coloring, random graph of average degree D is sampled consistent with an unknown vertex coloring and the goal is to estimate the coloring as well as possible (Krzakala & Zdeborová, 2009), as measured by the overlap of the output of the algorithm to the ground-truth coloring (see Definition B.9). As a corollary of our main result, we show that when all the latent tokens (1) x (N)
are masked and a few unmasked observation tokens provide the information of the form 9(x x(1) = 1 m + x mt] for j S N, then solving the masking problem can be reduced to solving planted coloring.
For planted m-coloring, when = 5 the thresholds in Proposition 3.3 are given by Dks/2 = 16 and Deona/ 2 x 13.23 (Krzakala & Zdeborová, 2009) (the factor of 2 here is simply because the observations correspond to ordered subsets of size 2). For general predicates and arities, there is an established recipe for numerically computing Dks and Deond based on the behavior of the belief propagation algorithm (see the discussion in Appendix B.4). As an example, in Fig. 4, we execute this recipe for m = 3, k = 3, and given by the Not-All-Equal predicate NAE(a,, ,," ,,'""" = 1 1/x x" x") to obtain thresholds that can be plugged into Proposition 3.3.
Additional examples of the hardness. The above setup can also be generalized to capture Bayesian constraint satisfaction problems (Montanari, 2008; Liu et al., 2022), one notable example of which is the stochastic block model (Decelle et al., 201 There are analogous predictions for the onset of hardness of inference, which can likewise be translated to hardness of masking problems for seemingly benign L&O distributions. In Appendix B.I and B.2, we give two more examples of L&O distributions for which awder- training is tractable yet order-agnostic training of the MDM computationally
16

---

Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffussins
hard.
First, we consider L&O distributions whose observations are sparse, noisy parities in the latents and deduce hardness for order-agnostic training from the Sparse Learning Parity with Noise assumption (Alekhnovich, 2003). We then consider L&O distributions whose observations are generalized linear models in the latents, and deduce hardness for large class of efficient algorithms from existing results on Lipschitz hardness (Alaoui & Gamarnik, 2024) for the symmetric binary perceptron (Aubin et al., 2019).
B.4. Proof of Proposition 3.3: sparse predicate observations
Here we formally define the relevant notions needed to formalize our claim about hardness in Proposition 3.3.
Definition B.9 (Planted CSPs), Given arity k E N, vocabulary/alphahet size m E N, predicate {1,..., .m)* - {0. 1). latent dimension N, and clause density P/N, ,, the corresponding planted constraint satisfaction problem is defined as follows: Nature samples an unknown assignment a uniformly random from (1....,m, ,m,". N and then for each ordered k-tuple S of distinct elements from [M], we observe the clause S independently with probability 6/Nk-1 if g(als) 1.
To measure the quality of an algorithm for recovering a given the observations, define the overlap between an estimate ē and the ground truth a by d(a,ō) - min. E, 1 j;; x(ē;) where Sy denotes the set of all permutations of {0, ,N
.N- - 1}. Define the average degree to be kP/N. N. i... the expected number of variables that share at least one
clause with given variable.
We begin by defining the central algorithm driving statistical physics predictions about hardness for random constraint satisfaction problems: belief propagation (BP).
Definition B.10 (BP update rules). Belief propagation an algorithm that iteratively updates set of messages
{MS M MMs)" where i, range over all pairs of variable indices [N]] and observations S 9 At time
++ 1, the messages computed via
MS (1 + 1] x
II
MST-+[1
(4)
THETYS
MS. S +1 + 1] x
Σ
g c) II
M [t]
(5)
EE
oE(I....maJ
jiijeS
where o U, c E (1....,m) mms assigns e to entry i and a to the remaining entries.
set of messages can be used to estimate the marginals of the posterior on a conditioned on the observations as follows. The marginal on the i-th variable has probability mass function over {1.... m) proportional to [TIræer MS. Given
of marginals, natural way to extract an estimate for is round to the color in (1, (1....,m) m] at which the probability mass function is largest.
Throughout we will make the following assumption that ensures that the trivial messages MS 1/m and MS 1/m are fixed point, sometimes called the paramagnetic fixed point, for the iteration above:
Assumption B.11. The quantity E=E(1,...m)\), g( U; e) is constant across all c E f1....,m) and i E [k]].
Definition B.12. Given k m. g, the -Siign threshold Dks is defined to be the largest average degree for which BP is locally stable around the paramagnetic fixed point, that is, starting from a small perturbation of the paramagnetic fixed point, it converges to the paramagnetic fixed point. More formally, Dks is the largest average degree at which the Jacobian of the BP operator (MS ( [1]} + (MSi (t + 1}} has spectral radius less than
The condensation threshold Deond is defined to be the largest average degree which the planted CSP ensemble and the following simple null model become mutually contiguous and thus statistically indistinguishable N -> co. The null model is defined as follows: there is no single unknown assignment, but instead for every ordered subset S of k variables, Nature independently samples an unknown local assignment 0g E m,.. and the observation is included with probability k if g(os) = 1.
For Doond < kP/N < DKs, there exists some other fixed point of the BP operator whose marginals, once rounded to an assignment, achieves strictly higher overlap than does BP with messages initialized randomly. The prediction is that in this regime, no efficient algorithm can achieve optimal recovery (Krzakala & Zdeborová, 2009).
17

---

Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffusions
Conjecture B.13 (IRSB cavity prediction). Suppose k; m, 9 satisfy Assumption B. 1, and let Dgs and Deaud denote the associated Kesten-Stigum and condensation thresholds for the average degree. Then for all P for which D kPPN < DKs, the best overlap achieved by computationally efficient algorithm for recovering is strictly less than the best overlap achievable.
Proof of Proposition 3.3. At masking fraction a satisfying the bounds in the Proposition, with probability at least a 2 DKs/NN-1)> 2 (1(1) we have that all tokens corresponding to latents T=(s) get masked. Independently of this, the number of unmasked tokens among the observation tokens Og is distributed as (in -1) .... (N - + 1 1-a), so by standard binomial tail bounds, with constant probability (depending on the gap between Drond and Dks) this lies between r Deoud N/k and 7 N/k. Furthermore, of these unmasked tokens in expectation y fraction of them correspond to observations for which the associated predicate evaluates to 1. Conditioned on the above events, the masking problem thus reduces exactly to inference for planted constraint satisfaction problem at average degree Deond < D < Dks, from which the Proposition follows.
C. Experimental details in Section 3
C.1. Experimental details in Section 3.2
r-learner configurations. We consider two distributions of x that interpolate between Unif (Sz. where Sr. denote the uniform distribution over all permutations of indices {0, 1,......,4 L 1} and the point mass at the identical distribution: (Closer) and (Much-closer). To construct those distributions, we start from the identity permutation and perform a certain number of random swapping operations. Since L Log(L) number of swaps results in a distribution that is very close to Unif (Sz) (Bormashenko, 20,1), we use L/10 10 and L swaps to construct the (Closer) and (Much-closer) distributions, respectively, For consistency, we repeat this sampling process three times.
Model and training configurations. As explained in Section 3.2, to evaluate the scaling law of the n-learner, we can simply adapt the autoregressive training setup (a transformer with causal attention) by modifying the input to (xx)) and using a learnable positional embedding layer instead of RoPE. We borrow the training configurations from (Nie et al., 2024), which are also consistent with the TinyLlama (Zhang et al., 2024) configurations. In particular, we use AdamW optimizer (Loshchilov & Hutter, 2017), setting 81 0.9, 92 0.95, and a weight decay of 0.1 and L 2048. A cosine learning rate schedule is applied, with maximum learning rate of 4 x 10 and minimum learning rate of 4 x 10-5, We also note that unless otherwise specified, we maintain the same training configuration throughout the paper.
Examining scaling laws. We conduct IsoFLOP analysis (Hoffmann et al., 2022). For given number of FLOPs C, by varying the number of non-embedding parameters of transformers, we set the iteration numbers so that the total number of tokens observed by the model during training equals C/6NN, following prior studies (Hoffmann et al., 2022: Kaplan et al., 2020). We then select the smallest validation loss and set it as a data point.
C.2. Experimental details in Section 3.3
C.2.1. EXPERIMENT ON L&O-NAE-S SAT DISTRIBUTION
We consider the L&O-NAE-SAT distribution with (N, P) = (20, 280). For each example sequence from L&O-NAB-SAT; we pad the last 212 tokens with an additional token value of 2. We employ 19M MDM with RoPE and maximum sequence length of 512. Then, this MDM is trained for x 1 iterations. To attain a proxy MDM for the Bayes optimal
predictor, we further train it for x 10a iterations.
To measure the error across different tasks, we consider the following setup. For each l E [1, - 1]. we randomly mask é tokens in the latent positions and x (P/N) tokens in the observed positions. Across all masked prediction positions, 6(1 + P/N), we measure the error for each position. For certainty, repeat this process 1000 times. The result in Figure 2 corresponds to the case when L 11, and we observe the same tendency for other values of f.
C.2.2. EXPERIMENT ON TEXT DATA
We take a 170M MDM pretrained with text data for a baseline model. To measure the performance imbalance between likelihood modeling tasks
18

---

Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Difffusions
Bzg Páata Σ logpe pa (1) zo politii................... 181
As done in the experiments in Section 3.2, we sample лs from three different distributions: Un,) (Closer), the point mass of identical distribution. For each case, we calculate the expectation over 1024 samples of xo ~ Pdata-
D. Experimental details in Section
D.I. Experimental details in Section 4.2
D.1.1. EXPERIMENT ON L.OO-NAE-SN DISTRIBUTION
We consider five instances of L&O-NAE-SAT: (N, P) (25, 275), (30, 270), (40, 260), (50, 250), (100, 200). For each distribution, we train a 19M MDM and measure the accuracy difference between vanilla inference and adaptive inference using Top-K probability margin.
D.1.2. EXPERIMENT ON TEXT DATA
Top-K probability margin sampler with temperature. To modify our inference for text data modeling, which does not have a determined answer, we found that adding a certain level of temperature to the oracle is useful. This is because Top- K probability margin or Top-K often leads to greedy sampling, which harms the diversity (entropy) of the generated samples. Therefore, consider variant of the oracle as follows, incorporating noise term
F(0, mt) = Top K (Ipa (x = =1(=)) - po(x' = j2|wx)] +c).
Note that this approach has also been employed for unconditional sampling (Wang et al., 2024; Zheng et al., 2023).
Generative perplexity and entropy. We employ a I.IB MDM pretrained on text data as a baseline. For each sampling step, we unconditionally generate samples using both vanilla and adaptive inference. Next, we calculate the likelihood using LLama2-75 baseline large language model. Moreover, we denote the entropy of generated sample as Σ ps log Pt. where P;- (x IL.L.
D.2. Experimental details on Sudoku and Zebra puzzles
the dataset provided in Shah et al. (2024) train our model. To evaluate our model the same difficulty tasks, we the test dataset proposed in Shah al. (2024). This dataset created by filtering the puzzles from (Radcliffe, 2020) that can be solved using fixed list of strategies. To create hard dataset to evaluate easy-to-hard generalization, we use the remaining puzzles from (Radcliffe, 2020) as they either require new strategy unseen during the training and/or require backtracking. The hard dataset contains around 1M Sudoku puzzles.
Dataset. For both Sudoku and Zebra puzzles,
Model, training, and inference. For the Sudoku dataset, we use 6M GPT-2 model and for the Zebra dataset, we use 19M model but instead of causal attention, we use complete bidirectional attention. We set the learning rate to 0.001 with batch size 128 to train the model for 300 epochs. For the inference, we use 50 reverse sampling steps using the appropriate strategy. Additionally, add Gumbel noise with coefficient of 0.5 to the MDM inference oracle F.
E. Omitted proofs
Proof of Proposition 2.1. We first re-state the Proposition 3.1 from (Zheng et al., 2024). To clarify, (Zheng et al., 2024) generally considers the case beyond the time-embedding denoising network pa-
Proposition E.1 (Proposition 3.1 of (Zheng et al., 2024)). For clean data Vo, let q(x(1)) xo) be the discrete forward process that randomly and uniformly masks lokens of xq. Suppose on O and on = 1. Then the MDM training loss (1) can be reformulated as
Cg
E (x(n)~aq(: mo)
=
1
Σ
eag log po =(())
(6)
n=1
Eiz(n)
19

---

Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Difffusions
To obtain alternative formulation of (6), expand the expectation æ<n) 7( 2o). Since there total L positions of To. have the probability assigned for each srn) equals 1/(5). Therefore,
La
(x(n)~((1z))
egg log po [x(n))
n=
Eur(n))==
E
M o log po [x[M)
(1M1)
MELL]AEM
Σ
x
log (a1 (x[M)
(MM1)
ME[L]iEM
E
log (al x[M]).
MELLIIEM T
((M
Reformulating the MDM loss with x-learner s. this paragraph, we provide the proof of
Σ E
-
E
Po [M]] 20
log po (1g(0) zo (i, IL. (t....,.--111))
L-1 (MMI-))
iE
20 Po--Paato
Alternatiely,, will demonstrate that
Σ MEM (
log po (=] -nr-)nif(B))
Σ
log pe ( 7(1)
zolmfi, L 1f1)
1=0
holds for every 2o. Note that
(nw
log
Fo
7(1)
zo (i.......... L 131)
1=0
x())
x (i.(T(i...... 111)
logpo po
To
xE3,, 0=0
Next, by regarding m(i...,, L- 1} { (j); 7L 1)} = M S [4] and (() i in the equation (1), we count the
number of Sr, that induces specific term log Po (r) | (M). For given M [L] and E M, π must satisfy
(ū) =i { (L 1)} = M.
The number of that satisfies above is (L |M))! x (IM) 1).. Finally, the following calculation concludes the proof.
z1)()) x0 p (i, L-
E Ent-Unit(Sa)
log pa
log pa
po:
TES 1=0
Σ
log 1ogpo(z6(2o/11) x (L - - |M)) x (IM) 1]!
L!
[M(E[L],EM
1
Σ
x (og
-
L-1
iMELLL,
(M--1)
20


Reviews:
[{"review_id": "qwultkwYsc", "reviewer": "Reviewer_xjpE", "summary": "This work presents two contributions to a emerging discrete diffusion model called masked diffusion models.\n\n* The first contribution is a theoretic construction showing the hardness of prediction subtasks within masked diffusion, motivating an inference time solution to sidestep these challenging subtasks.\n\n* The authors then proposed a new criteria to adaptively choose the decoding order in inference time based on probability margins and show that leads to significant performance boosts of masked diffusion models in planning problems.", "questions": "", "limitations": "", "rating": 5, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We greatly appreciate the reviewer's positive evaluation and the insightful comments and questions. Below, we respond to the reviewer’s suggestions and questions.\n\n## (1) Further experiments on text data--Top-k margin outperforms Top-k on challenging code and math tasks\n\nTo examine the effect of different inference strategies on text evaluation tasks, we adapted LLaDA, the 8B MDM model from [1]. We compare three strategies: **Vanilla, Topk, Topk prob -margin**.  The results are presented below.\n\n| **Sampler**      | HumanEval-Single | HumanEval-Multi | HumanEval-Split | Math        | MMLU-Pro | ROCStories |\n|------------------|-------------------------|------------------------|------------------------|--------|----------|-------------|\n| **Vanilla**      | 31.8%                   | 16.5%                  | 14.2%                  | 28.5%  | 33.2%    | 21.23%      |\n| **Top-k**        | 32.9%                   | 20.8%                  | 18.4%                  | 31.3%  | **36.5%** | 21.10%      |\n| **Top-k Margin** | **33.5%**               | **25.4%**              | **22.3%**              | **34.3%** | 35.4%    | **21.41%**   |\n\n\nAs shown in the table, both Top-k and Top-k Prob. Margin consistently outperform vanilla MDM inference, underscoring the importance of adaptively selecting the decoding order to avoid harder problem instances. Notably, in relatively challenging tasks, such as HumanEval-Multiline, HumanEval-Split Line, and Math, **Top-k Margin shows a clear advantage over Top-k**. \nThis is because the Top-k prob margin offers a more reliable estimate of uncertainty when multiple tokens have similar probabilities—a common scenario in these challenging tasks. In addition, particularly in coding and math problems where often a fixed answer exists, selecting the correct intermediate token during inference is critical. (Hence, a wrong token selection can directly lead to the incorrect answer) These results reinforce our claim in Section 4.1: Top-k Margin serves as a better proxy for positional uncertainty than Top-k in such cases. **These results also further highlight the potential of the Top-k Margin strategy for challenging infilling tasks.**\n\n ## (2) On the reason why MDMs outperform ARMs\nMDM is trained across all possible orderings and uses an adaptive inference strategy. This flexibility allows it to discover more efficient reasoning orders tailored to the task or dataset, which can generalize better to unseen data. In contrast, ARMs trained with a fixed order may fail to generalize to unseen (or harder) data (please refer to Table 4 in our paper). Additionally, the harder MDM training (i.e., training in more than one token generation order) might be more (sample) efficient than ARM training that focuses on learning in only one order. Moreover, the ordering used for ARM training—predetermined by humans—may be suboptimal. In contrast, MDM may discover more effective decoding orders by systematically leveraging information from the logits, often outperforming human-specified orderings.\n\n## (3) Implications\nThese hint at the strong potential of MDMs, which we also highlighted in Section 4.3: Since MDMs are trained on all possible masked subproblems, their adaptive inference allows them to discover *good reasoning paths*, potentially leading to better performance than fixed orderings predetermined by humans.\n\n[1] Large language diffusion models. Nie et al. 2025", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "Yes. I have followed the presented theoretical results, including proposition 2.1 and the tractability of subtasks in Example 3.2.", "experimental_designs_or_analyses": "Yes. The experiment design is thoughtful and closely tracks the main claims.", "relation_to_broader_scientific_literature": "This paper improves the understanding of masked diffusion models - presenting theoretical analysis of the hardness of the prediction problems, which motivates inference-time improvements that selects \"easier\" decoding passes.", "essential_references_not_discussed": "Relevant work in literature is well cited.", "other_strengths_and_weaknesses": "## Strengths\n\n* The theoretical results on provable hardness of mask prediction problems in some orders is original and improves our understanding of the intrinsic difficulty of training such models. It also provides sufficient motivation for the inference-time strategy that followed later.\n\n* The results on planning tasks are pretty strong. The proposed probablity margin strategy significantly improves performance on hard Sudoku tasks. Notably, it even outperforms an AR model that is informed of the optimal order to solve the problem.  \n\n## Weaknesses\n\n* The decoding-order selection, while novel, is still heuristic. The experimental evidence is in a narrow domain (e.g., sudoku puzzles) and the generality needs to be tested further, e.g., in text and image experiments.\n\n* It's surprising and unclear why the masked diffusion model that uniformly optimizes predictions over all ordering can ourperform AR models informed of the optimal ordering. Could the authors elaborate on the possible reasons and implications?", "other_comments_or_suggestions": "Can you test your decoding strategy on common text or image tasks often used to evaluate diffusion models?", "questions_for_authors": "N/A", "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": "q4UWZ0cBqs", "reviewer": "Reviewer_XERQ", "summary": "The main contribution of the paper is the use of theoretical arguments and carefully designed experiments to show the following:\n\n1. The complexity of training Masked Diffusion Models (MDMs) is higher than Auto-regressive Models (ARMs). \n\n2. The flexibility of any-order decoding offered by MDMs helps it to perform better than ARMs on specific kinds of data distributions, especially the ones where some (data/instance dependent) positions in the sequence contain harder sub-problems than other positions. \n\nThe paper also introduces a new decoding strategy called the \"top-k probability margin\"-based strategy, which picks the next token to decode based on the margin between top-2 vocabulary items at any specific position.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We greatly appreciate the reviewer's overall positive evaluation and comments. We will make sure to include a paragraph of related work in the main body and fix the typo mentioned. Below, we respond to the reviewer’s main concerns:\n\n## (1) Scope of our contributions\nThe reviewer stated “The contribution of the paper is incremental. It combines claims from existing papers (Nie et al. (2024) and Zheng et al. (2024))” and  “Nie et. al. (2024) and Zheng et. al. (2024b) together have already demonstrated the main claim of this paper to a great extent.”  \n\nWe respectfully disagree with the reviewer on this point. As stated in our introduction, our goal is to **understand** the benefits and drawbacks of training and inference of MDMs over ARMs.\n\n- Even though Nie et al. (2024) show that the autoregressive models outperform MDMs in scaling, they don’t explain the reason behind it. In this work, we give extensive empirical and theoretical evidence that this is due to the **heterogeneity of complexity across masking tasks at training time** (Section 3). **While empirically it has been observed (even well before the work of Nie et al.) that MDMs are more difficult to scale, our paper is the first to provide rigorous insight into why this is the case.**\n- While Zheng et al. (2024) propose a different ordering of the sampling, we view our most important contributions in this direction not to be about proposing new heuristics per se, but about explaining the reason/motivations behind the improvement achieved by these heuristics. Indeed, in our work we **provide principled justification for such adaptive inference schemes**, e.g. by showing that *any-order inference in perfectly trained MDM results in the same true distribution* (line 296, right column), and **disentangle the extent to which different “confidence-based” decoding strategies are actually planning based on uncertainty.**\n- Additionally, both of these works fail to explain that the benefit of MDM (especially with adaptive inference) over ARMs is most dramatic on tasks where the left-to-right token ordering structure *doesn’t hold*. This also explains the reason behind very drastic improvements in tasks like math or coding (e.g., see Table 1 in [1]) where left-to-right ordering doesn’t hold.\n\n[1] Large language diffusion models. Nie et al. 2025.\n\n\n## (2) Top-k margin outperforms Top-k on challenging code and math tasks\nOn text MDMs, for challenging math and coding tasks, we found that Top-k margin outperforms Top-k. For comparison, we adapted LLaDA [1], 8B MDM. For the result, please refer to our response to the reviewer xjPE. Notably, in more challenging tasks, such as HumanEval-Multiline, HumanEval-Split Line, and Math, **Top-k Margin shows a clear advantage over Top-k**. \n\nThis is because the Top-k margin offers a more reliable estimate of uncertainty when multiple tokens have similar probabilities (our claim in Section 4.1)—a common scenario in the challenging tasks in the coding and math domains. **These results also further highlight the potential of the Top-k Margin strategy for challenging infilling tasks.**\n\nTo understand the difference between Top-K and Top-K margin strategy, we consider the following problem. \n\nThe problem prompt given to an MDM is: [If $\\sqrt{5x}\\cdot\\sqrt{10x}\\cdot\\sqrt{18x}=30$, find $x$.] \nThe model’s output using Top-K strategy is:\n\n\n… \n*So the equation becomes:\n\\[\n\\sqrt{900x^3} = 30\n\\]*\n\n*Square both sides to eliminate the square root:\n\\[\n(900x^3)^2 = 30^2\n\\]*\n….\nThe model was wrong by decoding an incorrect sentence (900x^3)^2 = 30^2. At the moment just before decoding ^ (following 900x^3), the model faces multiple plausible options: (1) adding ^, or (2) adding =. This ambiguity arises from the token \"square\", which confuses the model. The Top-k strategy selects ^, as it has the highest probability. **This exemplifies a situation where the model assigns comparable probabilities to multiple tokens at a single location**. In contrast, the probability margin between ^ and = was small, indicating high uncertainty, so Top-k margin shifted focus to a different position where it had greater confidence, leading to the correct statement, *900x^3 = 900*.\n\n## (3) Examples of decoding trajectories for Sudoku:\nFor the following partial Sudoku board, Vanilla MDM inference decodes a cell at random—for example, the cell in the 7th row and 9th column. In contrast, adaptive MDM inference prioritizes cells in the 1st and 2nd rows, which are objectively easier to fill in at earlier.\n\n|   |   |   |   |   |   |   |   |   |\n|---|---|---|---|---|---|---|---|---|\n| 9 | 8 | 3 | 7 | 5 | . | 4 | 1 | 2 |\n| 2 | 4 | 5 | . | 9 | 1 | 3 | . | 6 |\n| 7 | 1 | 6 | 2 | 3 | . | . | . | . |\n| . | 2 | 1 | . | . | 8 | . | . | . |\n| 3 | 7 | . | . | 1 | . | 2 | 6 | . |\n| 6 | 9 | . | . | 2 | . | 8 | . | 1 |\n| 8 | . | 2 | . | . | . | 1 | . | . |\n| 1 | 5 | 7 | . | 8 | 3 | 6 | . | . |\n| . | 6 | . | 1 | 7 | . | 5 | 3 | 8 |", "claims_and_evidence": "## 1. The inference flexibility provided by the any-order decoding in MDMs overweights the drawbacks introduced by training complexity.\n\nNie et. al. (2024) already demonstrated through scaling law curves that the complexity of training MDMs is higher than ARMs. The paper provides some theoretical insight into the phenomenon. Their empirical experiments on text (Figure 2) re-confirm the observations made in Nie et. al. (2024).  Zheng et. al. (2024b) demonstrated that the logits produced by MDMs have useful information for selecting the positions to unmask during inference. In summary, Nie et. al. and Zheng et. al. together have already demonstrated the main claim of this paper. Therefore, in my view, the main contribution of this paper is the use of theoretical arguments and carefully designed synthetic experiments to drive the point home, which the paper does well. \n\n## 2. The proposed top-k probability margin strategy for sampling performs better than the top-k strategy proposed in Zhen et. al. 2024b. \n\nThe proposed top-k prob. margin-based sampling strategy is only demonstrated to be better than top-k (from Zhen et. al. 2024b) on Sudoku puzzles. Both top-k and top-k prob. margin work similarly for Zebra puzzles (Table 3). Moreover, it is not clear if there is any advantage to using the top-k prob. margin-based strategy on real data like the text data used in the paper (Figure 4 does not compare the two adaptive sampling strategies). Therefore, I find the use of top-k prob. margin-based strategy to be not well justified and a possible area to improve in the paper.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "The proof for Proposition 2.1 is correct. The claim in proposition 3.3 looks reasonable; however, I was unable to check the complete proof.", "experimental_designs_or_analyses": "All the experimental settings look sound.", "relation_to_broader_scientific_literature": "As mentioned above, Nie et. al. (2024) and Zheng et. al. (2024b) together have already demonstrated the main claim of this paper to a great extent. Nie et. al., demonstrated empirically on text data that the complexity of training MDMs is higher than ARMs. That said, the scope of Nie et. al. was quite broad and was focused more on the scaling aspect of MDMs. Zheng et. al. introduced the top-k sampling strategy for MDMs and demonstrated that it works much better than random unmasking. However, Zheng et. al. do not discuss the learning aspect of MDMs. This paper is much narrower in scope and tries to tease out the essence of adaptive decoding for MDMs through theoretical arguments and carefully selected experiments.", "essential_references_not_discussed": "The paper includes exhaustive references; however, the related work section is in the Appendix. Since the paper re-states some of the claims in existing papers, it would be better to include at least one paragraph of related work in the main paper.", "other_strengths_and_weaknesses": "The contribution of the paper is incremental. It combines claims from existing papers (Nie et al. (2024) and Zheng et al. (2024)). That said, since the paper only focuses on one claim, it is easy to read and follow.", "other_comments_or_suggestions": "1. Line 873: The expression for $p_i$ does not make sense. It should be $p_i = \\sum_{j=1}^L \\delta(x^j = i) / L$.\n\n2. It might be good to show some decoding trajectories on the Sudoku or Zebra puzzles where vanilla unmasking makes mistakes whereas the adaptive strategy circumvents them?", "questions_for_authors": "1. Do you observe any difference between top-k and top-k margin-based samping on text MDMs? Can you provide some examples of the generated text for various sampling strategies?", "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": "scL1euH8l9", "reviewer": "Reviewer_qGCY", "summary": "The paper takes a close look at training and inference of masked diffusion models (MDMs), which are a type of discrete diffusion models where the noising process consists of randomly “masking” tokens until all tokens are masked, and training a model to reverse this degradation process.\nThe paper claims that this order-agnostic training is inherently harder than left-to-right next-token prediction, that this difference can at least partly explain the gap in performance between AR models and MDMs, and that we can leverage the trained MDM to avoid difficult denoising orders. To this end, a novel sampling adapter for diffusion models is introduced, which significantly improves performance both on generative PPL for language modeling, as well as solving logical puzzles like Sudoku or Zebra puzzles.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their insightful review and address the comments below.\n\n## Soundness of theoretical claims\nThere are several misunderstandings, so we'd like to clarify them. The statement \"There exist distributions where for some noise levels (namely, when all tokens are masked) solving the corresponding subproblem is computationally hard and for some noise levels it is easy (namely, when all tokens are unmasked).”  **is incorrect**. The subproblem we investigate is to estimate the *coordinate-wise marginals* of the posterior distribution, not **full posterior sampling.** If all tokens are masked, full posterior can indeed be hard: take any hard-to-sample distribution. However, estimating marginals isn't necessarily difficult.\n\nFor example, take a hard-to-sample Ising model with density $\\propto e^{-x^TAx},x\\in \\\\{-1,1\\\\}^n$. Thanks to sign-symmetry, marginals are unbiased, so estimating them is trivial—even if full sampling is hard! Our theory emphasizes scenarios where some intermediate masking fractions are computationally harder than either extreme (fully masked or unmasked).\n\nIn vanilla MDM inference, at each step, a random subset of positions is selected to be unmasked. Consequently, the masking patterns encountered correspond to randomly sampled mask indices---precisely the setting considered in Proposition 3.3! In contrast, decoding in a fixed left-to-right order (as in ARM) leads to encountering only left-to-right subproblems. **These hopefully address the confusion on why our results imply the hardness of sampling under certain token orderings.**\n\nRegarding 1RSB, it is a widely accepted conjecture from statistical physics, with extensive literature support. For an introduction, see “Notes on computational-to-statistical gaps” by Bandeira, Perry, and Wein.\n\n## Comparison to AR baseline\nWe ran generative perplexity experiments using an ARM baseline. An 1.1B ARM achieved perplexity 11.745, lower than the 1.1B MDM’s 13.396 with adaptive inference. We acknowledge that our phrasing may have given the impression that adaptive MDM inference outperforms ARM. **However, our claim is more nuanced: adaptive MDM inference helps avoid hard problem instances.** Absolute ARM performance isn't directly relevant. To clarify, for Sudoku puzzles, we included ARM to demonstrate adaptive MDM’s advantage through flexible reasoning orders.\n\n## Other comments\n- *On decoding in any order in MDM*: For “MDMs can actually decode in any order”, we only meant that theoretically, when all the infilling problems are perfectly solved, any-order decoding matches the true likelihood. We'll update the PDF to make it clear.\n- *On the title*: For “train for the worst”, we never claimed optimality of training over all permutations, only the benefits of training in fixed order. For “plan for the best,” the reviewer is conflating “sampling adapters, the likes of which are ubiquitous for auto-regressive models” with what we do. For AR, the adapter has nothing to do with decoding *order*, which is left-to-right by default. The reason we call it planning is that MDMs can decide (plan) which token position to decode at each step.\n- *Entropy dropping*: While there is indeed an entropy decrease on adaptive MDM inference, this drop is negligible. To contextualize this, we measured entropy using the SlimPajama dataset: *average (5.10) and 0.45 quantile (4.85)*. These demonstrate minimal entropy reduction during adaptive inference, thus insignificantly impacting text generation performance.\n\n## Questions\n1) **Learnability of hash functions**: This nuanced issue is covered in “Cryptography in NC0” by Applebaum et al., demonstrating cryptographic primitives implemented via constant-depth circuits can be polynomially learnable.\n2) **$\\pi$-learner**: Each $\\pi$-learner learns sequences according to $\\pi$ and is modeled via causal Transformers trained on permuted data. A higher likelihood indicates easier learning. Average-case hardness, applied to MDM, involves uniformly sampling permutations $\\pi$, resulting in a lower likelihood compared to fixed left-to-right ordering (Fig 2, left, green line). Due to character limits, we kindly ask the reviewer to refer to the 'experimental setup' paragraph in Section 3.2 for more details. We are happy to clarify further on the discussion page.\n3) **Performance on left-to-right sampling**: We observed catastrophic collapse during **left-to-right MDM sampling**, with entropy dropping (~0.28). This is because MDM wasn't explicitly trained for left-to-right order. This result underscores the importance of adaptive inference strategies, where the model selects unmasking positions based on logit-derived uncertainty rather than on a prefixed order.\n4) **Per-token entropy strategy**: While a more natural measure would be per-token entropy, the only reason we went with top-k margin was its efficiency. In preliminary experiments, we also tried per-token entropy, but the performance difference was negligible.", "claims_and_evidence": "The question of whether or not “the benefits of inference flexibility for MDMs is enough to outweigh the drawbacks of training complexity” is claimed to be answered “in the affirmative.” While this does seem to be the case for solving Sudoku puzzles, the answer is not as clear for language modeling. Unfortunately, there is no comparison between AR sampling from an AR model (and/or left-to-right sampling from a MDM as a non-adaptive baseline) and adaptive sampling with the proposed sampling algorithm from a MDM. If AR models still have better sample quality than MDMs with adaptive sampling, which is rather plausible, then the answer to the original question would unfortunately change to “it depends.”\n\nL35 ff., Col. 2: I cannot verify the claim that “MDMs can actually be used to decode in any order” based on the provided reference. As far as I can tell, the cited paper does not conduct any experiments about the order-sensitivity of MDMs. Of course, in theory MDMs can decode in any order, but whether or not this is true in practice is a different question.", "methods_and_evaluation_criteria": "Besides what is mentioned above, the methods and evaluation criteria seem sound.", "theoretical_claims": "The paper makes the following two core theoretical claims.\n\nClaim 1:\n> we provide theoretical [...] evidence that the overhead imposed by training complexity quantifiably impacts MDMs’ performance.\n\nClaim 2:\n> We prove that even for simple, benign models of data, there are noise levels at which a large fraction, but not all, of the corresponding subproblems that MDMs must solve are computationally intractable.\n\nWhile the proposed L&O distributions seem very contrived (and calling it a “benign” model of data is, IMO, a stretch when it is known that these problems can be computationally hard), Proposition 3.3 does seem correct. The fact that it does require assuming the “1RSP cavity prediction” conjecture to be true may need to be highlighted more prominently (although I’m personally not familiar with this conjecture and cannot judge whether it is generally accepted to make this assumption). However, the original statement (claim 2) seems somewhat trivially true: Of course there exist distributions where for some noise levels (namely, when all tokens are masked) solving the corresponding subproblem is computationally hard and for some noise levels it is easy (namely, when all tokens are unmasked). A more useful proposition may be that some _orders_ of filling in the missing tokens are harder than others. How might Proposition 3.3 imply that some infilling orders are computationally hard?\n\nSimilarly, and from what I can tell, the statement(s) proved in Appendix B.2 (Proposition B.5) does/do not necessarily imply that “order-aware training is tractable yet order-agnostic training is hard”. As far as I understand, what is shown is that for some noise level, finding the solution is computationally hard. If this indeed implies that some _orders_ are more difficult than others, further explanation and/or proof is warranted.\n\nAll in all, the original claim (claim 1) of “providing theoretical evidence that [...] training complexity quantifiably impacts MDMs’ performance” is either overstated or stands on shaky ground. The provided hardness proofs seem correct upon skimming, but from what I can tell they do not lead to a conclusion this strong. It is possible that I'm simply not seeing the final logical step, in which case it should be a simple fix of adding a more detailed explanation that leads to the final conclusion.", "experimental_designs_or_analyses": "Sampling adapters often decrease the diversity of generated samples, which in extreme cases can lead to a collapse of the distribution. For the proposed adaptive inference there does indeed seem to be somewhat of a decrease in entropy, which warrants providing some qualitative examples in the appendix to prove that no catastrophic collapse is occurring.", "relation_to_broader_scientific_literature": "The paper shines a light on the discrepancy between AR models and discrete (masked) diffusion models, which has been observed many times in the literature. It also proposes a novel adaptive sampling technique that drastically improves performance both on text and logic puzzles. Both of these are valuable contributions to the literature of masked diffusion models.", "essential_references_not_discussed": "It is known that auto-regressive models trained with teacher forcing face some fundamental challenges [1], which may (at least partly) be to blame for their poor performance on Sudoku. There has also been work on the directionality of AR language models, finding that there seems to be a slight but consistent left-to-right bias in human language. Finally, a recent study (concurrent work) has applied image-based diffusion models to solving Sudoku [3].\n\nAll of these are not strictly necessary to cite, but may help tie the results into the broader literature. It is for the authors to decide whether or not to include them.\n\n- [1] Bachmann & Nagarajan, 2024. https://arxiv.org/abs/2403.06963\n- [2] Papadopoulos et al., 2024. https://arxiv.org/abs/2401.17505\n- [3] Wewer et al., 2025. https://arxiv.org/abs/2502.21075", "other_strengths_and_weaknesses": "The paper makes an important observation on how the infilling order in masked diffusion can have a major effect on both upstream and downstream performance. While the theoretical part seems a bit shaky, the empirical evidence is convincing. \n\nBesides the theoretical part, the main weakness of the paper lies in lacking scientific rigor and a tendency to overstate or misrepresent the actual results. For example, the phrase “train for the worst” seems to imply that it is optimal to train on all possible permutations jointly, but this claim is not tested in the paper. Similarly, “planning for the best” implies that there is some sort of planning involved, which there isn’t (the term “planning”, in the context of Machine Learning, generally refers to the act of looking ahead of and beyond the immediate next step). In actuality, the paper proposes a sampling adapter, the likes of which are ubiquitous for autoregressive models. Applying this idea to discrete diffusion models is a novel and valuable contribution, and obfuscating it through a misleading title is not necessary.\n\nDespite the concerns about soundness and in light of the strong empirical results, I am inclined to recommend an accepting decision and will be happy to update my score if these concerns can be addressed.", "other_comments_or_suggestions": "Nits:\n- L104: $e_{x_0^i}$ should be bold.\n- Figure 1 (bottom): 2nd line, 2nd step; mask tokens should presumably have a black background.\n- Definition 3.1: As stated, the vocabulary size is $q+1$, and $p_{data}$ should presumably be $\\{1, \\dots, q\\}^L$. Also, lowercase $n$ is not defined and presumably refers to uppercase $N$.\n- L216: I think it should be $g(x | S)$, not $g(x |_S)$.\n- Figure 3 is not referenced in the text.\n- Conjecture B.13 does not have a citation.\n- Def. 3.1: Overloaded notation: $\\pi$ is used for both permutation and latent distribution.", "questions_for_authors": "1. L189 ff.: If The observations are cryptographic hash functions, is it not true that the observations themselves are also not efficiently learnable? Or, put differently, are cryptographic hash functions efficiently learnable? It seems to me like the answer would be no, in which case this example is not only worst-case, but also violates our assumption.\n2. How is the “hardness” of a $\\pi$-learner measured, and what is the exact hardness of the quoted “$\\pi$-learner-much_closer”, “$\\pi$-learner-closer”, and “$\\pi$-learner-unif”? How does this compare to the average-case hardness, which would apply for MDM?\n3. If different orderings have different inherent difficulty, and MDM is trained on all of them jointly, how does MDM perform on inherently easy (e.g. left-to-right) orders? According to the claims in the paper, we would expect this to be quite close to AR models and would be a nice addition to Figure 2.\n4. As an alternative to the top-k probability margin: Could the per-token entropy be a better proxy for uncertainty?", "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": "DAxT7rpaoZ", "reviewer": "Reviewer_oBo5", "summary": "The authors first show masked diffusion models (MDMs) indeed train on computationally intractable subproblems compared to their autoregressive counterparts. Then an adaptive Top-K probability margin inference strategy is proposed to sidestep hard subproblems that are not properly learned in the training time. The proposed inference strategy has proven to be effective in the Sudoku puzzle task.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer’s valuable questions and comments. Below, we address the main concerns.\n\n## (1) Further experiments on text data\nIn response to the reviewer’s comments, we ran additional experiments and found that **Top-k margin indeed outperforms Top-k on challenging code and math tasks**. \n\nSpecifically, to examine the effect of different inference strategies on text evaluation tasks, we adapted LLaDA, the 8B MDM model from [1]. We compare three strategies: **Vanilla, Topk, Topk-margin**.  The results are presented below.\n\n| **Sampler**      | HumanEval-Single | HumanEval-Multi | HumanEval-Split | Math        | MMLU-Pro | ROCStories |\n|------------------|-------------------------|------------------------|------------------------|--------|----------|-------------|\n| **Vanilla**      | 31.8%                   | 16.5%                  | 14.2%                  | 28.5%  | 33.2%    | 21.23%      |\n| **Top-k**        | 32.9%                   | 20.8%                  | 18.4%                  | 31.3%  | **36.5%** | 21.10%      |\n| **Top-k Margin** | **33.5%**               | **25.4%**              | **22.3%**              | **34.3%** | 35.4%    | **21.41%**   |\n\n\nAs shown in the table, both Top-k and Top-k Margin consistently outperform the Vanilla MDM inference, underscoring the importance of adaptively selecting the decoding order to avoid harder problem instances. Notably, in more challenging tasks, such as HumanEval-Multiline, HumanEval-Split Line, and Math, Top-k Margin shows a clear advantage over Top-k. \n\nThis is because, particularly in coding and math problems where the fixed answer exists, selecting the correct intermediate token during inference is critical. Moreover, the Top-k margin offers a more reliable estimate of uncertainty when multiple tokens have similar probabilities—a common scenario in these challenging tasks. These results reinforce our claim in Section 4.1: Top-k Margin serves as a better proxy for positional uncertainty than Top-k in such cases. **These results also further highlight the potential of the Top-k Margin strategy for challenging infilling tasks.** \n\n**We also emphasize that our main contribution lies in the fundamental understanding of token ordering in MDM, rather than in proposing a superior inference strategy.** Nevertheless, as demonstrated in the experimental results on Sudoku puzzles and math/coding tasks, our proposed Top-$k$ Margin inference shows promising potential compared to Top-$k$ in certain scenarios, highlighting practical implications of our work. **Given that the main weakness raised was that we did not demonstrate effectiveness for our inference strategy on text data, we hope the reviewer will consider raising their score.**\n\n## (2) Clarification about Top-K probability margin\nThe Top-K probability margin strategy is not only useful when there are multiple possible correct values for a given position, but can also be effective when there is a single correct value. The reason is that the strategy is used for the **distribution $p_\\theta$ estimated by the model** rather than for the **true distribution $p_{data}$.**\n\nTo understand this distinction, consider the case of Sudoku. Recall that $x_t$ denotes the sequence at decoding time $t$ – in the case of Sudoku, this is the partially filled Sudoku puzzle. The posterior data distribution at the $i$-th location, given the partially filled puzzle, is denoted by $p_{data}(x^i | x_t)$. Since we are dealing with puzzles that have unique solutions, the reviewer is correct that $p_{data}(x^i = v | x_t) = 1$ at the correct value $v$ and 0 otherwise.\n\nHowever, the Top-K probability margin strategy doesn’t rely on $p_{data}(x^i | x_t)$ but instead uses $p_{\\theta}(x^i | x_t)$ during the adaptive inference. As explained in Section 3, the learned posterior can be quite different from the true posterior. Intuitively, $p_{\\theta}(x^i | x_t)$ reflects the model’s uncertainty about the correct value at the $i$-th location. When the model is unsure, it assigns high probabilities to a few candidate values which we refer to as the possible values at that position.\n\nThe Top-K probability margin is effective in Sudoku because situations often arise where the model is uncertain between two or more possible values — say, $v_1$ and $v_2$ — assigning high probabilities to both. In these cases, the top probability $\\max( p_{\\theta}(x^i = v_1 | x_t), p_{\\theta}(x^i = v_2 | x_t) )$ does not provide a reliable estimate of whether the model knows the correct value at the ith location. However, the Top-K probability margin $| p_{\\theta}(x^i = v_1 | x_t) - p_{\\theta}(x^i = v_2 | x_t) |$ serves as a more effective measure of the model’s uncertainty.", "claims_and_evidence": "The claim of \"Complexity at training time\" of MDMs is well shown both theoretically and empirically on text data.\nThe claim of \"Planning for the best\" of MDMs is less persuasive. For instance, the Top-K probability margin inference strategy proves effective for Sudoku puzzles but does not work for the Zebra puzzle, and its performance on text data remains unknown.", "methods_and_evaluation_criteria": "The evaluation of the imbalanced subproblems during the training time of MDMs and the proposed adaptive inference strategy makes sense in general.", "theoretical_claims": "I reviewed the proofs in Section 2 and they seemed correct to me.", "experimental_designs_or_analyses": "The experimental designs in Sections 3 and 4 are mostly correct.", "relation_to_broader_scientific_literature": "The findings in the paper may be useful in building the next-generation LLMs based on the diffusion model.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "Strengths:\n- The analysis of why \"MDMs train on hard problems\" is sufficient and helps those unfamiliar with MDMs understand their characteristics.\n- The Top-K probability margin inference strategy is intuitive and effective on Sudoku.\n\nWeaknesses:\n- The proposed Top-K probability margin inference strategy proves effective only for Sudoku puzzles but not for text data, which raises concerns about the technical contribution to the general text domain and diffusion LLMs.", "other_comments_or_suggestions": "N.A.", "questions_for_authors": "Regarding adaptive Top-K probability margins: As it is stated \"When multiple values have similar probabilities at a position, Top-K probability margin will provide a better estimate of the uncertainty of a position\". Does this mean the Top-K probability margin strategy is only useful when there are multiple possible values for a given position? If so, why is the Top-K probability margin so effective for Sudoku as each position has only one solution in a standard Sudoku puzzle. I would appreciate further clarification on this.", "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": {"ethics_expertise_needed": ["Other expertise"]}}]

Author rebuttal:

Meta-review: The paper theoretically argues that masked diffusion models effectively solve a harder problem than autoregressive models. However, they show that this can be circumvented if one were to adaptively order the token decoding. This is emprically demonstrated on logic (Sudoku) puzzles (and later in the authors' response, Math and HumanEval-coding tasks).


All reviewers agree that the results here improve our understanding of masked diffusion models significantly (`xjpE`, `qGCY`), the theory is original (`xjpE`) and overall findings maybe useful in building future diffusion models (`oBo5`). The reviewers agree that the experiments are  well-designed (`xjpE`, `XERQ`) and sound (`qGCY`) and support the claims (`xjpE`). Reviewers (`XERQ`, `oBo5`) say the proofs look correct or reasonable.  

Reviewers `xjpE` and `oBo5` also say that the results are strong as far as Sudoku is concerned; however multiple reviewers raised the concern that there were no gains in Zebra, and no results were presented in text/image tasks (`xjpE`, `qGCY`, `oBo5`). The authors have responded with new affirmative results on math/coding tasks. I believe this addresses the reviewers concerns adequately.

As Reviewer `qGCY` suggests in detail, in future versions of the paper, I hope that the authors will take care to make their claims more precise.
Decision: Accept (oral)