Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffusions

We appreciate the reviewer’s valuable questions and comments. Below, we address the main concerns.

(1) Further experiments on text data

In 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.

Specifically, 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.

As 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.

This 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.

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.

(2) Clarification about Top-K probability margin

The 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}$.

To 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.

However, 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.

The 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.

We thank the reviewer for their insightful review and address the comments below.

Soundness of theoretical claims

There 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.

For 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).

In 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.

Regarding 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.

Comparison to AR baseline

We 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.

Other comments

• 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.
• 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.
• 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.

Questions

1. 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.
2. $\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.
3. 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.
4. 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.

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:

(1) Scope of our contributions

The 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.”

We 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.

• 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.
• 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.
• 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.

[1] Large language diffusion models. Nie et al. 2025.

(2) Top-k margin outperforms Top-k on challenging code and math tasks

On 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.

This 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.

To understand the difference between Top-K and Top-K margin strategy, we consider the following problem.

The problem prompt given to an MDM is: [If $\sqrt{5x}\cdot\sqrt{10x}\cdot\sqrt{18x}=30$, find $x$.] The model’s output using Top-K strategy is:

… *So the equation becomes:

$

\sqrt{900x^3} = 30

$ *

*Square both sides to eliminate the square root:

$

(900x^3)^2 = 30^2

$ * …. The 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.

(3) Examples of decoding trajectories for Sudoku:

For 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.

We greatly appreciate the reviewer's positive evaluation and the insightful comments and questions. Below, we respond to the reviewer’s suggestions and questions.

(1) Further experiments on text data--Top-k margin outperforms Top-k on challenging code and math tasks

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 prob -margin. The results are presented below.

As 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. This 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.

(2) On the reason why MDMs outperform ARMs

MDM 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.

(3) Implications

These 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.

[1] Large language diffusion models. Nie et al. 2025