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

为最坏情况训练，为最佳路径规划：理解掩码扩散中的 Token 顺序

Jaeyeon Kim Affiliation:  Harvard University    Kulin Shah Affiliation:  University of Texas Austin Correspondence to: kulinshah@utexas.edu    Vasilis Kontonis Affiliation:  University of Texas Austin    Sham Kakade Affiliation:  Harvard University    Sitan Chen Affiliation:  Harvard University

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 with $7\times$ as many parameters and that were explicitly trained via teacher forcing to learn the right order of decoding. This shows that MDMs without knowledge of the correct token generation order during training and inference can outperform ARMs trained with knowledge of the correct token generation order. We also show the effectiveness of adaptive MDM inference on reasoning tasks such as coding and math on the 8B large language diffusion model (LLaDa 8B).

近年来，掩码扩散模型（masked diffusion models，MDMs）成为离散域生成建模的一条有前景的替代路线。与自回归模型（autoregressive models，ARMs）相比，MDMs 用训练时更高的复杂度换取推理时更大的灵活性：训练时须学会指数级多种填空问题，推理时却几乎能按任意顺序解码 token。本文仔细审视这两种此消彼长的效应。在训练侧，我们从理论与实验两方面说明：相对自回归对照，MDMs 确实在训练计算上不可行的子问题上用力。在推理侧，我们表明：若能恰当、自适应地选择 token 解码顺序，MDMs 的能力会显著增强，从而绕开困难子问题。在数独这类逻辑谜题上，自适应推理能把预训练 MDM 的求解准确率从不足 7% 提升到约 90%，甚至超过参数量为其 7 倍、且经教师强制显式学习正确解码顺序的 ARM。这说明：训练与推理时都不知道正确 token 生成顺序的 MDM，可以胜过带着正确顺序知识训练的 ARM。我们还在 80 亿参数的大语言扩散模型（LLaDa 8B）上，展示了自适应 MDM 推理在编程与数学等推理任务上的效果。

1 Introduction

1 引言

Keywords:

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 et 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 a leading variant, distinguished by a simple and principled objective: to generate samples, learn to reverse a noise process which independently and randomly masks tokens.

关键词：

尽管扩散模型（Ho et al., 2020；Song et al., 2021）已在图像、视频与音频等连续域生成建模中占据主导，把同一思路推广到文本与蛋白质等离散域（Austin et al., 2021；Lou et al., 2024；Hoogeboom et al., 2021b）仍处早期。众多方案里，掩码扩散模型（MDMs）（Lou et al., 2024；Sahoo et al., 2025；Shi et al., 2024）成为主流变体，其目标简洁且有原则：为生成样本，学习反转一个独立、随机地掩码 token 的噪声过程。

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 a 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 on what type of tasks can MDMs be scaled to challenge the dominance of ARMs in discrete generative modeling?

在语言建模等许多应用中，MDMs 仍逊于自回归模型（ARMs）（Nie et al., 2024；Zheng et al., 2024）；后者学习反转的是从左到右依次揭开 token 的噪声过程。不过近期研究提示：在推理（Nie et al., 2024；Kitouni et al., 2025）、规划（Ye et al., 2024）与填空（Gong et al., 2024）等 ARM 短板场景，MDMs 可能更有优势。关键问题随之而来：相对 ARMs，MDMs 的长短处是什么？在哪些任务上把 MDM 做大，才有望挑战离散生成建模中 ARM 的主导地位？

To understand these questions, we turn a microscope to two key competing factors when weighing the merits of MDMs over ARMs:

为回答这些问题，我们把显微镜对准权衡 MDMs 与 ARMs 时的两个关键竞争因素：

• • Complexity at training time: MDMs face a more challenging training task by design. While ARMs predict the next token given an unmasked prefix, MDMs predict a token conditioned on a set of unmasked tokens in arbitrary positions. This inherently increases their training complexity.

• • 训练时的复杂度：MDMs 的训练任务本身就更难。ARMs 在未掩码前缀上预测下一 token；MDMs 则在任意位置的未掩码集合上条件预测某个 token。训练复杂度因此固有地升高。

• • Flexibility at inference time: On the other hand, the sampling paths taken by an MDM are less rigid. Unlike the fixed left-to-right decoding of ARMs, MDMs decode tokens in random order at inference. Even more is possible: MDMs can be used to decode in any order (including left-to-right).

• • 推理时的灵活性：另一方面，MDM 的采样路径更不僵硬。不像 ARM 固定的从左到右解码，MDMs 在推理时按随机顺序解码 token；甚至还能按任意顺序解码（也包括从左到右）。

Therefore, we ask:

因此我们追问：

Are the benefits of inference flexibility for MDMs enough to outweigh the drawbacks of training complexity?

MDMs 在推理灵活性上的收益，是否足以抵消训练复杂度上的代价？

In this work, we provide dual perspectives on this question.

本文对这一问题给出双重视角。

(1) Training for the worst.  First, we provide theoretical and empirical evidence that the overhead imposed by training complexity quantifiably impacts MDMs’ performance.

（1）为最坏情况训练。 首先，我们给出理论与经验证据：训练复杂度带来的额外负担，会量化地影响 MDMs 的表现。

Theoretically, we show examples of simple data distributions with a natural left-to-right order, where ARMs can provably generate samples efficiently. In contrast, there are noise levels at which a large fraction of the corresponding subproblems solved by MDMs for these distributions are provably computationally intractable. Empirically, we validate this claim on real-world text data, known to have left-to-right order and show that the imbalance in training complexity across subproblems persists even in real-world text data (Fig. 2, left).

理论上，我们构造简单且具有自然从左到右顺序的数据分布：ARMs 可被证明能高效生成样本；相对地，在某些噪声水平下，MDMs 为这些分布所解的大量子问题在计算上可证不可行。经验上，我们在已知具有从左到右顺序的真实文本上验证：跨子问题的训练复杂度失衡在真实文本中同样存在（图 2，左）。

(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 building upon the observation (Chang et al., 2022; Zheng et al., 2023) that MDMs which can perfectly solve all masking subproblems can be used to decode in any order.

（2）为最佳路径规划。

上述观察对 MDMs 看似坏消息；但在论文后半部分，我们基于 Chang et al.（2022）、Zheng et al.（2023）的观察——若 MDM 能完美解所有掩码子问题，就能按任意顺序解码——对引导问题给出肯定回答。

In first part of the paper, we show that the imbalance in complexity across subproblems during the training of MDMs results in some of the subproblems being poorly trained and the vanilla MDM inference that unmasks tokens in random order results in evaluating the poorly trained marginals. Therefore, in place of vanilla MDM inference, we consider adaptive strategies that carefully select which token to unmask next. Our key insight is that the adaptive strategies makes it possible to sidestep the hard subproblems from training (Fig. 1). In particular, 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. We show the effectiveness of the adaptive inference in solving logic puzzles, coding, math and infilling tasks. For example, on Sudoku puzzles, a simple adaptive strategy (Section 4.1) improves the accuracy of MDMs from $<7$% to almost 90%.

论文前半表明：MDM 训练中子问题复杂度失衡，会使部分子问题学得很差；而按随机顺序揭开 token 的朴素 MDM 推理，恰恰会去评估这些学得差的边缘分布。因此，我们不以朴素推理为终点，而考虑自适应策略：谨慎选择下一步揭开哪个 token。核心洞见是：自适应策略有可能绕开训练中的困难子问题（图 1）。尤其我们发现：即便不改 MDM 的训练方式，模型的 logits 已含有足够信息来决定正确的揭开顺序。我们在逻辑谜题、编程、数学与填空任务上展示自适应推理的效果；例如在数独上，简单自适应策略（第 4.1 节）把 MDM 准确率从不足 7% 提到近 90%。

Advantage of MDMs over ARMs.

We show that the main effectiveness of MDMs lies in tasks that do not have the same natural token generation order across all sequences (e.g., logic puzzles and reasoning tasks like coding and math). By carefully designing experiments on logic puzzles, we show that MDMs without the knowledge of the correct token generation order during training and inference can outperform ARMs trained with the knowledge of the correct token generation order. In particular, we show that MDMs that decide the correct token generation order during inference via adaptive strategies can outperform ARMs that are trained to learn the right token generation order via supervised teacher forcing (Shah et al., 2024; Lehnert et al., 2024).

MDM 相对 ARM 的优势。

我们表明：MDMs 的主要效力，体现在并非所有序列都共享同一自然 token 生成顺序的任务上（例如逻辑谜题，以及编程、数学等推理任务）。通过对逻辑谜题的细致实验设计，我们说明：训练与推理时都不知道正确 token 生成顺序的 MDMs，可以胜过带着正确顺序知识训练的 ARMs。具体而言，推理时靠自适应策略决定正确生成顺序的 MDMs，可胜过经监督教师强制学习正确生成顺序的 ARMs（Shah et al., 2024；Lehnert et al., 2024）。

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.

结构安排。

第 2 节给出 MDMs 预备知识与记号；第 3 节审视 MDM 训练，展示跨子问题的计算不可行性失衡；第 4 节讨论 MDMs 的自适应推理，并考察其在多种任务上对似然建模的影响。

2 Masked Diffusion Models (MDM)

2 掩码扩散模型（MDM）

In this section, we explain the framework of Masked Diffusion Models (Shi et al., 2024; Sahoo et al., 2025) and highlight its interpretation as an order-agnostic learner. MDMs gradually add masking noise to the true discrete data and learn the marginal distribution of the induced reverse process. We formally define both the forward and reverse processes for MDMs below.

本节说明掩码扩散模型框架（Shi et al., 2024；Sahoo et al., 2025），并强调其作为与顺序无关的学习者的解释。MDMs 向真实离散数据逐步加入掩码噪声，并学习诱导的反向过程边缘分布。下面形式化定义 MDM 的前向与反向过程。

Let the distribution $p_{\rm{data}}$ on $\{1,\ldots,m\}^{L}$ be the data distribution over sequences of length $L$ and with vocabulary $\{1,\ldots,m\}$. We use $0$ to denote the “mask” token.

设数据分布 $p_{\rm{data}}$ 定义在 \{1,\ldots,m}^{L} 上，即长度为 $L$、词表为 \{1,\ldots,m} 的序列分布。用 $0$ 表示「掩码」token。

Forward process.

For a given $x_{0}\sim p_{\rm{data}}$ and a noise level $t\in[0,1]$, the forward process $x_{t}\sim q_{t|0}(\cdot\,|\,x_{0})$ is a coordinate-independent masking process via $q_{t|0}(x_{t}|x_{0})=\prod_{i=0}^{L-1}q_{t|0}(x_{t}^{i}|x_{0}^{i})$, where

前向过程。

给定 $x_{0}\sim p_{\rm{data}}$ 与噪声水平 $t\in[0,1]$，前向过程 $x_{t}\sim q_{t|0}(\cdot\,|\,x_{0})$ 是坐标独立的掩码过程：$q_{t|0}(x_{t}|x_{0})=\prod_{i=0}^{L-1}q_{t|0}(x_{t}^{i}|x_{0}^{i})$，其中

$$q_{t|0}(x_{t}^{i}\mid x_{0}^{i})=\mathrm{Cat}\bigl(\alpha_{t}\mathbf{e}_{x_{0}^{i}}+(1-\alpha_{t})\mathbf{e}_{0}\bigr).$$

Here, $\alpha_{t}$ is a predefined noise schedule satisfying $\alpha_{0}\approx 1,\alpha_{1}\approx 0$ and $\mathbf{e}_{x_{0}^{i}}\in\mathbb{R}^{m+1}$ is a one-hot vector corresponding to the value of token $x_{0}^{i}$. $\mathrm{Cat}(\pi)$ denotes the categorical distribution given by $\pi\in\Delta^{m}$. In other words, for each $i$-th coordinate, $x_{t}^{i}$ is masked to the mask token $0$ with probability $1-\alpha_{t}$ and remains unchanged otherwise.

这里 $\alpha_{t}$ 是预定噪声日程，满足 $\alpha_{0}\approx 1,\alpha_{1}\approx 0$；$\mathbf{e}_{x_{0}^{i}}\in\mathbb{R}^{m+1}$ 是对应 token $x_{0}^{i}$ 取值的 one-hot 向量。$\mathrm{Cat}(\pi)$ 表示由 $\pi\in\Delta^{m}$ 给出的类别分布。换言之，对每个第 $i$ 坐标，$x_{t}^{i}$ 以概率 $1-\alpha_{t}$ 被掩成掩码 token $0$，否则保持不变。

Figure 1: (Top) MDM training can be seen as learning multiple 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).

图 1：（上）MDM 训练可视为学习多个掩码预测问题，其中一些更难学，造成性能失衡（第 3 节）。（下）推理时，自适应 MDM 可避开困难问题实例，从而提升性能（第 4 节）。

Reverse process.

The reverse process of the above forward process is denoted by $q_{s|t}(x_{s}|x_{t},x_{0})$ and is given by $q_{s|t}(x_{s}|x_{t},x_{0})=\prod_{i=0}^{L-1}q_{s|t}(x_{s}^{i}|x_{t},x_{0})$ for any $s<t$, where

反向过程。

上述前向过程的反向记为 $q_{s|t}(x_{s}|x_{t},x_{0})$，对任意 $s<t$ 有 $q_{s|t}(x_{s}|x_{t},x_{0})=\prod_{i=0}^{L-1}q_{s|t}(x_{s}^{i}|x_{t},x_{0})$，其中

$$q_{s|t}(x_{s}^{i},\lvert,x_{t},x_{0})=\begin{cases}\mathrm{Cat}(\mathbf{e}{x{t}^{i}})\quad&x_{t}^{i}\neq 0\
\mathrm{Cat}\left(\frac{1-\alpha_{s}}{1-\alpha_{t}}\mathbf{e}{0}+\frac{\alpha{s}-\alpha_{t}}{1-\alpha_{t}}\mathbf{e}{x{0}^{i}}\right)\quad&x_{t}^{i}=0,.\end{cases}$$

The reverse transition probability $q_{s|t}(x_{s}^{i}|x_{t},x_{0})$ is approximated using $g_{\theta}(x_{s}^{i}|x_{t})\triangleq q_{s|t}(x_{s}^{i}\,\lvert\,x_{t},x_{0}\leftarrow p_{\theta}(\cdot|x_{t},t))$ where $p_{\theta}(\cdot|x_{t},t)$ is a denoising network trained to predict the marginal distribution on $x_{0}^{i}$ via an ELBO-based loss for all masked tokens at noise scale $t$ (i.e., for all $i$ such that $x_{t}^{i}=0$). To be precise, $q_{s|t}\left(x_{s}^{i}\mid x_{t},x_{0}\leftarrow p_{\theta}(\cdot|x_{t},t)\right)$ indicates the conditional probability where $p_{\theta}(\cdot|x_{t},t)$ is placed in the position of $e_{x_{0}^{i}}$ within $q_{s|t}(x_{s}^{i}\mid x_{t},x_{0})$. The denoising network is trained to minimize the following loss derived from the score-entropy (Lou et al., 2024; Sahoo et al., 2025; Shi et al., 2024; Ou et al., 2024):

反向转移概率 $q_{s|t}(x_{s}^{i}|x_{t},x_{0})$ 用 $g_{\theta}(x_{s}^{i}|x_{t})\triangleq q_{s|t}(x_{s}^{i}\,\lvert\,x_{t},x_{0}\leftarrow p_{\theta}(\cdot|x_{t},t))$ 近似，其中 $p_{\theta}(\cdot|x_{t},t)$ 是去噪网络，经基于 ELBO 的损失训练，以在噪声尺度 $t$ 上预测所有被掩码 token 上的 $x_{0}^{i}$ 边缘分布（即所有使 $x_{t}^{i}=0$ 的 $i$）。精确地说，$q_{s|t}\left(x_{s}^{i}\mid x_{t},x_{0}\leftarrow p_{\theta}(\cdot|x_{t},t)\right)$ 表示把 $p_{\theta}(\cdot|x_{t},t)$ 放进 $q_{s|t}(x_{s}^{i}\mid x_{t},x_{0})$ 中 $e_{x_{0}^{i}}$ 位置后的条件概率。去噪网络最小化由分数熵导出的如下损失（Lou et al., 2024；Sahoo et al., 2025；Shi et al., 2024；Ou et al., 2024）：

$$\mathcal{L}{\theta}=\int{0}^{1}\frac{\alpha_{t}^{\prime}}{1-\alpha_{t}}\displaystyle\mathop{\mathbb{E}}{\begin{subarray}{c}x{0}\sim p_{\rm data}\
x_{t}\sim q_{t|0}(\cdot|x_{0})\end{subarray}}\sum_{i:x_{t}^{i}=0}-\log p_{\theta}(x_{0}^{i}|x_{t},t)dt,$$

where $\alpha_{t}^{\prime}=\frac{d\alpha_{t}}{dt}$ and the summation is computed over masked tokens (i.e., all $i$ such that $x_{t}^{i}=0$). In practice, a time-embedding-free architecture for the denoising network, i.e., $p_{\theta}(\cdot|x_{t},t)=p_{\theta}(\cdot|x_{t})$ is generally used as $x_{t}$ implicitly contains information about $t$ via the number of masked tokens.

其中 $\alpha_{t}^{\prime}=\frac{d\alpha_{t}}{dt}$，求和只对掩码 token（即所有使 $x_{t}^{i}=0$ 的 $i$）。实践中去噪网络常采用无时间嵌入结构，即 $p_{\theta}(\cdot|x_{t},t)=p_{\theta}(\cdot|x_{t})$，因为 $x_{t}$ 已通过掩码 token 个数隐式携带 $t$ 的信息。

The reverse sampling process starts from the fully masked sentence $x_{1}=(0,\ldots,0)$. Suppose we have a partially \fully masked sequence $x_{t}$ at a given noise level $t\in(0,1]$. Then, to obtain $x_{s}$ for a predetermined noise level $s<t$, we sample $x_{s}^{i}\sim g_{\theta}(\cdot|x_{t})$ for all $i$. This process is repeated recursively from $t=1$ to $t=0$.

反向采样从全掩码句 $x_{1}=(0,\ldots,0)$ 开始。设在噪声水平 $t\in(0,1]$ 已有部分或全掩码序列 $x_{t}$；对预定的 $s<t$，对所有 $i$ 采样 $x_{s}^{i}\sim g_{\theta}(\cdot|x_{t})$ 得到 $x_{s}$。该过程从 $t=1$ 递归做到 $t=0$。

2.1 Reformulating the training and inference of MDMs

2.1 重述 MDM 的训练与推理

In this section, we first discuss training of MDMs and compare it with “left-to-right” order training of autoregressive models in Section 2.1.1. Then, we reformulate vanilla MDM inference in Section 2.1.2 to set the stage for the upcoming discussion.

本节先在第 2.1.1 小节讨论 MDM 训练，并与自回归模型的「从左到右」顺序训练对比；再在第 2.1.2 小节重述朴素 MDM 推理，为后文铺垫。

2.1.1 Order-agnostic training of MDMs

Recent works (Zheng et al., 2024; Ou et al., 2024) have observed that the learning problem of MDM is equivalent to a masked language model. Building upon their analysis, we reformulate the loss $\mathcal{L}_{\theta}$ to show that $\mathcal{L}_{\theta}$ is a linear combination of the loss for all possible infilling masks. We first define $x_{0}[M]$ as a masked sequence, obtained from original sequence $x_{0}$ where indices in the mask set $M$ (a subset of $[L]\triangleq\{1,2,\ldots,L\}$) are replaced with mask token $0$.

2.1.1 MDM 的与顺序无关训练

近期工作（Zheng et al., 2024；Ou et al., 2024）观察到：MDM 的学习问题等价于掩码语言模型。沿其分析，我们重写损失 $\mathcal{L}_{\theta}$，表明它是所有可能填空掩码上损失的线性组合。先定义 $x_{0}[M]$：由原序列 $x_{0}$ 把掩码集合 $M$（$[L]\triangleq\{1,2,\ldots,L\}$ 的子集）中的下标换成掩码 token $0$ 得到的掩码序列。

Proposition 2.1.

Assume $\alpha_{0}=1$, $\alpha_{1}=0$ and denoising network $p_{\theta}$ is time-embedding free. Then $\mathcal{L}_{\theta}\leq-\mathbb{E}_{x_{0}\sim p_{\rm data}}[\log p_{\theta}(x_{0})]$ and

命题 2.1。

假设 $\alpha_{0}=1$、$\alpha_{1}=0$，且去噪网络 $p_{\theta}$ 无时间嵌入。则 $\mathcal{L}_{\theta}\leq-\mathbb{E}_{x_{0}\sim p_{\rm data}}[\log p_{\theta}(x_{0})]$，并且

$$\mathcal{L}_{\theta}=-\sum_{M\subseteq[L],i\in M}\frac{1}{|M|}\frac{1}{\binom{L}{|M|}}\displaystyle\mathop{\mathbb{E}}_{x_{0}\sim p_{\rm data}}[\log p_{\theta}(x^{i}_{0}|x_{0}[M])], \tag{1}$$

where $|M|$ is the size of the set $M$ and $p_{\theta}(x_{i}\mid x_{0}[M])$ indicates the conditional probability of the $i$-th coordinate from $p_{\theta}(x_{t})$.

其中 $|M|$ 是集合 $M$ 的大小，$p_{\theta}(x_{i}\mid x_{0}[M])$ 表示 $p_{\theta}(x_{t})$ 在第 $i$ 坐标上的条件概率。

The proof of the above proposition is given in Appendix E. As the MDM loss is a linear combination of the loss for all possible infilling mask $M$, the minimizer of the loss $\mathcal{L}_{\theta}$ learns to solve every masking problem. In other words, the optimal predictor $p_{\theta}$ is the posterior marginal of the $i$-th token, conditioned on $x_{0}[M]$ for all masks $M$.

命题证明见附录 E。由于 MDM 损失是所有填空掩码 $M$ 上损失的线性组合，损失 $\mathcal{L}_{\theta}$ 的极小化者须学会解每一个掩码问题。换言之，最优预测器 $p_{\theta}$ 是在所有掩码 $M$ 下，第 $i$ 个 token 关于 $x_{0}[M]$ 的后验边缘。

On the other hand, Autoregressive Models (ARMs) learn to predict $i^{\textrm{th}}$ token $x^{i}$ based on all preceding tokens, from $x^{0}$ to $x^{i-1}$. This is equivalent to predicting $x^{i}$ by masking positions from $i$ to $L-1$. Therefore, the training objective for ARMs can be expressed as:

另一方面，自回归模型（ARMs）学习在前缀 token $x^{0}$ 到 $x^{i-1}$ 上预测第 $i$ 个 token $x^{i}$。这等价于把位置 $i$ 到 $L-1$ 掩掉再预测 $x^{i}$。因此 ARM 训练目标可写成：

$$\log p_{\theta}(x_{0})=\sum_{i=0}^{L-1}\log p_{\theta}(x_{0}^{i}|x_{0}[\{i,\ldots,L-1\}]). \tag{2}$$

Typically, ARMs are trained to predict tokens sequentially from left to right. We refer to this as left-to-right training. However, it’s also possible to train these models to predict tokens sequentially based on a fixed, known permutation of the sequence. We refer to this general approach as order-aware training.

通常 ARMs 按从左到右顺序预测 token，我们称之为从左到右训练。也可按序列的某个固定、已知置换顺序依次预测，我们把这类做法统称为感知顺序的训练。

To understand the comparison between the training objective of MDMs and ARMs, we want to highlight the equivalence between any-order autoregressive loss and MDM loss (Hoogeboom et al., 2021a; Ou et al., 2024). In particular, under conditions of Proposition 2.1, MDM loss is equal to

为理解 MDM 与 ARM 训练目标的对比，我们强调任意顺序自回归损失与 MDM 损失的等价性（Hoogeboom et al., 2021a；Ou et al., 2024）。在命题 2.1 的条件下，MDM 损失等于

$$\displaystyle\mathcal{L}{\theta}=-\mathop{\mathbb{E}}{\begin{subarray}{c}x_{0}\sim p_{\textrm{data}}\
\pi\sim\textrm{Unif}(\mathbb{S}{L})\end{subarray}}\left[\sum{i=0}^{L-1}\log p_{\theta}\left(x_{0}^{\pi(i)}\Big|x_{0}[\pi{i,\ldots,L-1}]\right)\right],$$

where $\textrm{Unif}(\mathbb{S}_{L})$ is a uniform distribution over all the permutations of length $L$ (See Appendix E.1 for the proof). Observe that if the expectation is only with respect to the identity permutation, then the loss becomes an autoregressive loss. This shows that MDM loss solves exponentially more subproblems than ARM loss. In contrast to ARM loss, MDM does not prefer any particular (e.g., left-to-right) order during the training; therefore, we call its training order-agnostic training.

其中 $\textrm{Unif}(\mathbb{S}_{L})$ 是长度 $L$ 的全部置换上的均匀分布（证明见附录 E.1）。若期望只对恒等置换取，损失就退化为自回归损失。由此可见，MDM 损失比 ARM 损失多解指数级个子问题。与 ARM 损失不同，MDM 训练不偏好任何特定顺序（例如从左到右）；因此我们称其训练为与顺序无关的训练。

2.1.2 Order-agnostic inference of MDMs

The MDM inference can be decomposed into two steps: (a) randomly selecting a set of positions to unmask and (b) assigning token values to each position via the denoising network $p_{\theta}$. More precisely, we can reformulate the reverse process $x_{s}\sim g_{\theta}(\cdot|x_{t})$ as follows.

2.1.2 MDM 的与顺序无关推理

MDM 推理可拆成两步：（a）随机选取要揭开的位置集合；（b）用去噪网络 $p_{\theta}$ 为各位置赋 token 值。更精确地，可将反向过程 $x_{s}\sim g_{\theta}(\cdot|x_{t})$ 重述如下。

Therefore, the inference in MDM is implemented by randomly selecting $S$ and then filling each token value according to the posterior probability $p_{\theta}(x_{s}^{i}|x_{t})$.

因此，MDM 推理的实现是：随机选取 $S$，再按后验概率 $p_{\theta}(x_{s}^{i}|x_{t})$ 填入各 token 值。

On the other hand, ARMs are trained to predict tokens sequentially from left to right and therefore, generate tokens also in left-to-right order. In contrast, vanilla MDM inference generates the tokens in a random order.

另一方面，ARMs 按从左到右训练，因而也按从左到右生成；朴素 MDM 推理则按随机顺序生成 token。

3 MDMs train on hard problems

3 MDMs 在困难问题上训练

Figure 2: Left: MDMs train on hard problems (Section 3.2). x-axis and y-axis correspond to $\log(\text{FLOPs})$ and $-\log p_{\theta}(x)$, respectively. 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-likelihoods. Right: Task error imbalance (Section 3.3). MDM’s performance varies across different tasks. For text data (top right), this is indicated by validation loss. For L&O-NAE-SAT (bottom right), MDM performs well on the masking problems for observation positions (light region) but struggles with latent positions (dark region).

图 2：左：MDMs 在困难问题上训练（第 3.2 节）。横轴与纵轴分别为 $\log(\text{FLOPs})$ 与 $-\log p_{\theta}(x)$。MDM（蓝）在似然建模上逊于 ARM（橙）。MDM 所训的多数掩码问题（其他曲线）比 ARM 遇到的更难，表现为更小的对数似然。右：任务误差失衡（第 3.3 节）。MDM 在不同任务上表现不一。文本数据（右上）用验证损失体现；对 L&O-NAE-SAT（右下），MDM 在观测位置的掩码问题上表现好（浅色区），在潜变量位置上则吃力（深色区）。

In this section, we provide theoretical and empirical evidence that when the data distribution has left-to-right order (or any fixed known order) then autoregressive training in left-to-right order (or in the known order) is more tractable than MDMs. In particular, for such distributions with fixed order, we show that ARMs can efficiently sample from the distributions but for MDMs, we theoretically and empirically demonstrate that a large portion of masking subproblems $p_{\theta}(x^{i}_{0}\mid x_{0}[M])$ can be difficult to learn.

本节给出理论与经验证据：当数据分布具有从左到右顺序（或任一固定已知顺序）时，按该顺序做自回归训练，比 MDMs 更易处理。尤其对这类有固定顺序的分布，我们说明 ARMs 能高效采样；而对 MDMs，我们从理论与实验两方面表明：大量掩码子问题 $p_{\theta}(x^{i}_{0}\mid x_{0}[M])$ 可能很难学。

In Section 3.1, we show several examples of simple, non-pathological distributions for which: (1) the masking problems encountered during order-aware training (such as in ARMs) are computationally tractable, yet (2) many of the ones encountered during order-agnostic training (such as in MDMs) are computationally intractable. In Section 3.2, we empirically show that text data also exhibits this gap between the computational complexity of order-aware and order-agnostic training and therefore, MDMs train on subproblems of wide variety of complexity (depending on the order/masks). In Section 3.3, we empirically show that the variety in training complexity results in performance imbalance across subproblems: MDMs trained on data from such distributions exhibits small errors on easy subproblems but suffers from large errors on harder ones.

第 3.1 节给出若干简单、非病态分布的例子，满足：（1）感知顺序训练（如 ARMs）遇到的掩码问题在计算上可行；（2）与顺序无关训练（如 MDMs）遇到的许多问题在计算上不可行。第 3.2 节经验上说明文本数据同样存在感知顺序与与顺序无关训练之间的计算复杂度鸿沟，因而 MDMs 在复杂度跨度很大的子问题上训练（取决于顺序/掩码）。第 3.3 节经验上说明：训练复杂度的多样性会带来跨子问题的性能失衡：在这类分布上训练的 MDMs 在容易子问题上误差小，在困难子问题上误差大。

3.1 Benign distributions with hard masking problems

3.1 带困难掩码问题的良性分布

We now describe a simple model of data under which we explore the computational complexity of masking problems and show the contrast between masking problems encountered by MDMs and ARMs.

下面描述一个简单的数据模型，用以考察掩码问题的计算复杂度，并对比 MDMs 与 ARMs 所遇掩码问题的差异。

Definition 3.1.

A latents-and-observations (L&O) distribution is a data distribution $p_{\rm data}$ over sequence of length $L$ with alphabet size $m$ (precisely, $p_{\rm data}$ is over $\{0,\ldots,m\}^{L}$) is specified by a permutation $\pi$ over indices $\{1,2,\ldots,L\}$, number of latent tokens $N$, number of observation tokens $P$ such that $N+P=L$, prior distribution $p_{\textrm{prior}}$ of latent variables over $\{1,\ldots,m\}$ and efficiently learnable observation functions $\mathcal{O}_{1},\ldots,\mathcal{O}_{P}:\{1,\ldots,m\}^{N}\to\Delta(\{0,\ldots,m\})$,[^1]

定义 3.1。

潜变量与观测（latents-and-observations，L&O）分布是长度 $L$、字母表大小 $m$ 上的数据分布 $p_{\rm data}$（精确地说定义在 $\{0,\ldots,m\}^{L}$ 上），由以下要素指定：指标 \{1,2,\ldots,L} 上的置换 $\pi$、潜变量 token 数 $N$、观测 token 数 $P$（$N+P=L$）、潜变量在 $\{1,\ldots,m\}$ 上的先验 $p_{\textrm{prior}}$，以及可高效学习的观测函数 $\mathcal{O}_{1},\ldots,\mathcal{O}_{P}:\{1,\ldots,m\}^{N}\to\Delta(\{0,\ldots,m\})$，[^1]

• • (Latent tokens) For $i=1,\ldots,N$, sample $x^{\pi(i)}$ independently from the prior distribution $p_{\textrm{prior}}$ of the latents.

• • （潜变量 token）对 $i=1,\ldots,N$，独立地从先验 $p_{\textrm{prior}}$ 采样 $x^{\pi(i)}$。

• • (Observation tokens) For $j=1,\ldots,P$, sample $x^{\pi(N+j)}$ independently from $\mathcal{O}_{j}(x^{\pi(1)},\ldots,x^{\pi(N)})$.

• • （观测 token）对 $j=1,\ldots,P$，独立地从 $\mathcal{O}_{j}(x^{\pi(1)},\ldots,x^{\pi(N)})$ 采样 $x^{\pi(N+j)}$。

L&O distributions contain two types of tokens: (1) latent tokens and (2) observation tokens. Intuitively, latent tokens are tokens in the sequence, indexed by $\pi(1),\pi(2),\ldots,\pi(N)$ that serve as “seeds” that provide randomness in the sequence; the remaining tokens, called observation tokens (indexed by $\pi(N+1),\pi(N+2),\ldots,\pi(N+P)$), are determined as (possibly randomized) functions of the latent tokens via $\mathcal{O}_{1},\ldots,\mathcal{O}_{P}$. Observe that L&O distributions specified by a permutation $\pi$ have a natural generation order by permutation $\pi$.

L&O 分布含两类 token：（1）潜变量 token；（2）观测 token。直观上，由 $\pi(1),\ldots,\pi(N)$ 索引的潜变量 token 充当序列中的「种子」，提供随机性；其余由 $\pi(N+1),\ldots,\pi(N+P)$ 索引的观测 token，则经 $\mathcal{O}_{1},\ldots,\mathcal{O}_{P}$ 由潜变量（可能带随机）决定。可见，由置换 $\pi$ 指定的 L&O 分布具有沿 $\pi$ 的自然生成顺序。

Order-aware training

Order-aware training, i.e. by permuting the sequence so that $\pi$ becomes the identity permutation and then performing autoregressive training, is computationally tractable: predicting $x^{\pi(i)}$ given $x^{\pi(1)},\ldots,x^{\pi(i-1)}$ is trivial when $i\leq N$ as the tokens are independent, and computationally tractable when $i>N$ because $x^{\pi(i)}$ only depends on $x^{\pi(1)},\ldots,x^{\pi(N)}$ and is efficiently learnable by assumption. 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.

感知顺序的训练

感知顺序训练——即先把序列置换使 $\pi$ 变为恒等置换，再做自回归训练——在计算上可行：给定 $x^{\pi(1)},\ldots,x^{\pi(i-1)}$ 预测 $x^{\pi(i)}$，当 $i\leq N$ 时因 token 独立而平凡；当 $i>N$ 时，$x^{\pi(i)}$ 只依赖 $x^{\pi(1)},\ldots,x^{\pi(N)}$，且按假设可高效学习。相对地，下文将给出例子：若按 MDM 式做与顺序无关训练，将以高概率撞上困难掩码问题。

Order-agnostic training

We first note that if the observations $(\mathcal{O}_{1},\ldots,\mathcal{O}_{P})$ are given by a cryptographic hash function, then the masking problem of predicting $(x^{\pi(1)},\ldots,x^{\pi(L)})$ given $(x^{\pi(N+1)},\ldots,x^{\pi(N+P)})$ 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 is worst-case in nature – in real-world data, one rarely trains on sequences given by cryptographic hash functions. Furthermore, it only establishes hardness for a specific masking pattern which need not be encountered in the course of running the reverse process.

与顺序无关的训练

首先指出：若观测 $(\mathcal{O}_{1},\ldots,\mathcal{O}_{P})$ 由密码学哈希给出，则在给定 $(x^{\pi(N+1)},\ldots,x^{\pi(N+P)})$ 时预测 $(x^{\pi(1)},\ldots,x^{\pi(L)})$ 的掩码问题，本质上要求求逆哈希，因而在设计上计算不可行。这虽是关于语言建模中 token 顺序角色的常见民间观察，却并不完全令人满意：构造是最坏情形的——真实数据很少由密码学哈希生成序列；而且它只对特定掩码模式建立困难性，该模式在反向过程中未必出现。

We provide several simple instances of L&O distributions that address these issues: instead of leveraging delicate cryptographic constructions, they are average-case in nature and furthermore we can establish hardness for typical masking problems encountered along the reverse process.

我们给出若干简单的 L&O 分布实例来回应这些问题：不依赖精巧的密码学构造，而是平均情形的，并且能对反向过程中典型遇到的掩码问题建立困难性。

In all these examples, the hardness results we establish hold even if the algorithm knows all of the parameters of $p_{\rm data}$ as well as the observation functions $\mathcal{O}_{1},\ldots,\mathcal{O}_{P}$. Due to space constraints, here we focus on the following example, deferring two others to Apps. B.1 and B.2.

在所有这些例子中，即使算法已知 $p_{\rm data}$ 的全部参数以及观测函数 $\mathcal{O}_{1},\ldots,\mathcal{O}_{P}$，困难性结论依然成立。因篇幅所限，这里聚焦下面例子，另两个放到附录 B.1 与 B.2。

Example 3.2(Sparse predicate observations).

Consider the following class of L&O distributions. Given arity $k\geq 2$, fix a predicate function $g:\{1,\ldots,m\}^{k}\to\{0,1\}$. Consider the set of all ordered subsets of $\{1,2,\ldots,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(N-1)\cdots(N-k+1)$). To sample a new sequence, we first sample latent tokens $x^{\pi(1)},\ldots,x^{\pi(N)}$ from the prior distribution $p_{\textrm{prior}}$ and an observation latent corresponding to a $k$-sized subset $S$ is given by $g(\{x^{\pi(i)}\}_{i\in S})$. In other words, each observation latent corresponds to a $k$-sized subset $S$ of $\{1,2,\ldots,N\}$ and the corresponding observation function $\mathcal{O}_{S}(x^{\pi(1)},\ldots,x^{\pi(N)})$ is given by $g(\{x^{\pi(i)}\}_{i\in S})$.

例 3.2（稀疏谓词观测）。

考虑如下 L&O 分布类。给定元数* $k\geq 2$，固定谓词函数 $g:\{1,\ldots,m\}^{k}\to\{0,1\}$。取 $\{1,2,\ldots,N\}$ 上所有大小为 $k$ 的有序子集，令观测潜变量总数 $P$ 等于该集合大小（故 $P=N!/(N-k)!=N(N-1)\cdots(N-k+1)$）。采样新序列时，先从先验 $p_{\textrm{prior}}$ 采样潜变量 token $x^{\pi(1)},\ldots,x^{\pi(N)}$；对应于大小为 $k$ 的子集 $S$ 的观测潜变量由 $g(\{x^{\pi(i)}\}_{i\in S})$ 给出。换言之，每个观测潜变量对应 $\{1,2,\ldots,N\}$ 的一个 $k$ 元子集 $S$，相应观测函数 $\mathcal{O}_{S}(x^{\pi(1)},\ldots,x^{\pi(N)})$ 为 $g(\{x^{\pi(i)}\}_{i\in S})$。*

Proposition 3.3.

Let $x$ be a sample from an L&O distribution $p_{\rm data}$ with sparse predicate observations as defined in Example 3.2, with arity $k$ and predicate $g$ satisfying Assumption B.11, and let $\gamma$ be the probability that $g$ is satisfied by a random assignment from $\{1,\ldots,m\}^{k}$. Let $D_{\rm KS}$ and $D_{\rm cond}$ be some constants associated with the predicate function $g$ (see Definition B.12). Suppose each token in $x$ is independently masked with probability $\alpha$, and $M$ is the set of indices for the masked tokens. If $1-\gamma^{-1}D_{\rm KS}/kN^{k-1}\leq\alpha\leq 1-\gamma^{-1}D_{\rm cond}/kN^{k-1}$, then under the 1RSB cavity prediction (see Conjecture B.13), with probability $\Omega_{k}(1)$ over the randomness of the masking, no polynomial-time algorithm can solve the resulting subproblem of predicting any of the masked tokens among $x^{\pi(1)},\ldots,x^{\pi(N)}$ given $x[M]$.

命题 3.3。

设 $x$ 取自例 3.2 所定义的带稀疏谓词观测的 L&O 分布 $p_{\rm data}$，元数为 $k$，谓词 $g$ 满足假设 B.11；令 $\gamma$ 为随机赋值取自 $\{1,\ldots,m\}^{k}$ 时 $g$ 被满足的概率。设 $D_{\rm KS}$ 与 $D_{\rm cond}$ 为与谓词 $g$ 相关的常数（见定义 B.12）。若每个 token 独立以概率 $\alpha$ 被掩码，$M$ 为被掩码下标集合，且 $1-\gamma^{-1}D_{\rm KS}/kN^{k-1}\leq\alpha\leq 1-\gamma^{-1}D_{\rm cond}/kN^{k-1}$，则在1RSB 腔方法预言*（见猜想 B.13）下，对掩码随机性以概率 $\Omega_{k}(1)$，没有任何多项式时间算法能在给定 $x[M]$ 时，解出预测 $x^{\pi(1)},\ldots,x^{\pi(N)}$ 中任一被掩码 token 的子问题。*

The complete proof of the proposition is given in Section B.4. We also provide a proof outline in Section B.3 for a comprehensive understanding.

命题完整证明见第 B.4 节；第 B.3 节给出证明提纲以便把握全貌。

3.2 Empirical evidence of hardness via likelihoods

3.2 经似然度量的困难性经验证据

In the previous section, we provided theoretical evidence that order-aware training is tractable when data has a natural order but the order-agnostic training is not. In this section, we provide empirical evidence to support this claim, using natural text data. Additionally, 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 performance gap is primarily due to the order-agnostic training of MDMs. Natural text inherently follows a left-to-right token order, and we show that as training deviates from this order, model performance progressively declines.

上一节给出理论证据：数据有自然顺序时，感知顺序训练可行，而与顺序无关训练不可行。本节用自然语言文本提供经验支持。此外，近期研究（Nie et al., 2024；Zheng et al., 2024）表明 MDMs 在自然语言上逊于 ARMs。本节给出证据：这一差距主要来自 MDMs 的与顺序无关训练。自然语言固有地遵循从左到右的 token 顺序；我们说明：训练越偏离该顺序，模型表现越差。

To understand the importance of the order during the training, we use the following setting: Given a permutation $\pi$ of indices $\{0,1,\ldots,L-1\}$, define a $\pi$-learner to be a likelihood model $\log p_{\theta}(x_{0})$ given as follows:

为理解训练顺序的重要性，采用如下设定：给定指标 \{0,1,\ldots,L-1} 的置换 $\pi$，定义 $\pi$-学习器 为如下似然模型 $\log p_{\theta}(x_{0})$：

$$\log p_{\theta}(x_{0})=\sum_{i=0}^{L-1}\log p_{\theta}\bigl(x_{0}^{\pi(i)}\Big|x_{0}[\pi\{i,\ldots,L-1\}]\bigr)\, \tag{3}$$

In other words, the $\pi$-learner predicts the token at position $\pi(i)$ given the clean tokens $x_{0}^{\pi(0)},\ldots,x_{0}^{\pi(i-1)}$ and masked tokens $x_{0}^{\pi(i)},\ldots,x_{0}^{\pi(L-1)}$. If $\pi$ is the identity permutation, this reduces to the standard (left-to-right) autoregressive training. Note that the MDM loss encodes a $\pi$-learner for every permutation $\pi$ because the MDM loss (1) is equivalent to the average loss of those $\pi$-learners over $\pi$ sampled from $\mathrm{Unif}(\mathbb{S}_{L})$:

换言之，$\pi$-学习器在干净 token $x_{0}^{\pi(0)},\ldots,x_{0}^{\pi(i-1)}$ 与掩码 token $x_{0}^{\pi(i)},\ldots,x_{0}^{\pi(L-1)}$ 的条件下，预测位置 $\pi(i)$ 上的 token。若 $\pi$ 为恒等置换，则退化为标准（从左到右）自回归训练。注意 MDM 损失对每个置换 $\pi$ 都编码了一个 $\pi$-学习器，因为 MDM 损失（1）等价于对 $\pi\sim\mathrm{Unif}(\mathbb{S}_{L})$ 取平均的那些 $\pi$-学习器损失：

$$\displaystyle\mathcal{L}{\theta}=-\mathop{\mathbb{E}}{\begin{subarray}{c}x_{0}\sim p_{\textrm{data}}\
\pi\sim\textrm{Unif}(\mathbb{S}{L})\end{subarray}}\left[\sum{i=0}^{L-1}\log p_{\theta}\left(x_{0}^{\pi(i)}\Big|x_{0}[\pi{i,\ldots,L-1}]\right)\right],$$

where $\mathbb{S}_{L}$ denotes the set of all permutations over $\{0,1,\ldots,L-1\}$. The proof of the above equivalence is given in Appendix E. Therefore, by measuring the ‘hardness’ of each $\pi$-learner, we can probe differences in hardness between arbitrary masking problems and left-to-right masking problems.

其中 $\mathbb{S}_{L}$ 表示 \{0,1,\ldots,L-1} 上全部置换的集合。上述等价证明见附录 E。因此，度量每个 $\pi$-学习器的「难度」，就能探测任意掩码问题与从左到右掩码问题之间的难度差异。

Experimental setup.

We use the Slimpajama dataset (Soboleva et al., 2023) to evaluate the performance of training in different orders. To train a $\pi$-learner, we employ a transformer with causal attention and use permuted data $\pi(x_{0})$ as input. By varying $\pi$ 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 $\pi$-learner.

实验设置。

我们用 Slimpajama 数据集（Soboleva et al., 2023）评估不同顺序下的训练表现。训练 $\pi$-学习器时，采用带因果注意力的 Transformer，并以置换后的数据 $\pi(x_{0})$ 为输入。在固定其余训练配置（模型、优化等）下只改变 $\pi$，便可用所得似然（按式 3 计算）作为刻画 $\pi$-学习器所解子问题难度的度量。

In our experiments, the sequence length $L$ is $2048$, so repeating the scaling laws for each $\pi$ is infeasible. Instead, we sample $\pi\sim\mathrm{Unif}(\mathbb{S}_{L})$ and examine the scaling law of the $\pi$-learner’s likelihood. We leverage the codebase from (Nie et 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 a learnable 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 $\pi$ and the identity permutation affects the scaling law, we consider two interpolating distributions over permutations between $\mathrm{Unif}(\mathbb{S}_{L})$ (i.e, MDM training) and the point mass at the identical permutation (i.e, ARM training). We sample three permutations from the interpolating distribution and $\mathrm{Unif}(\mathbb{S}_{L})$ and plot the scaling law for each of the permutation. Due to space constraints, we provide further experimental details in Appendix C.1.

实验中序列长度 $L=2048$，对每个 $\pi$ 重复缩放律不可行。于是我们采样 $\pi\sim\mathrm{Unif}(\mathbb{S}_{L})$，考察 $\pi$-学习器似然的缩放律。我们借用 Nie et al.（2024）的代码库，其中给出了 MDM 与 ARM 的基线缩放律。又因 RoPE 对从左到右顺序有归纳偏置，所有实验改用可学习位置嵌入以纠正；相应地也重跑了原先使用 RoPE 的基线。为考察 $\pi$ 与恒等置换的距离如何影响缩放律，我们考虑在 $\mathrm{Unif}(\mathbb{S}_{L})$（即 MDM 训练）与恒等置换上的点质量（即 ARM 训练）之间的两种插值置换分布。从插值分布与 $\mathrm{Unif}(\mathbb{S}_{L})$ 各采样三个置换，并画出各自的缩放律。更多实验细节因篇幅放在附录 C.1。

Results.

As shown in Fig. 2, the scaling law for a $\pi$-learner with uniformly random $\pi$ is worse than that of an ARM. This elucidates the inherent hardness of masking problems $p_{\theta}(x_{i}\mid x_{0}[M])$ beyond left-to-right prediction and also explains why MDM, which is trained simultaneously on all $\pi\in\mathbb{S}_{L}$, is worse than ARM in likelihood modeling. Additionally, as $\pi$ gets closer to the identity permutation, the scaling laws also get closer to ARM ($\pi$-learner-closer and $\pi$-learner-much-closer in Fig. 2). This also supports the common belief that ARM is a good fit for text data as it inherently follows a left-to-right ordering.

结果。

如图 2 所示，均匀随机 $\pi$ 的 $\pi$-学习器缩放律逊于 ARM。这说明除从左到右预测之外，掩码问题 $p_{\theta}(x_{i}\mid x_{0}[M])$ 固有更难，也解释了为何同时在所有 $\pi\in\mathbb{S}_{L}$ 上训练的 MDM，在似然建模上逊于 ARM。此外，当 $\pi$ 越接近恒等置换，缩放律也越接近 ARM（图 2 中的 $\pi$-learner-closer 与 $\pi$-learner-much-closer）。这也支持常见信念：ARM 适合文本，因为文本固有从左到右顺序。

That said, it should also be noted that even though MDMs are trained on exponentially more masking problems than ARM ($\Theta(L2^{L})$ versus $L$), its performance is not significantly worse than $\pi$-learners. We attribute this to the blessing 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.

不过也应注意：尽管 MDMs 比 ARM 多训指数级掩码问题（$\Theta(L2^{L})$ 对 $L$），其表现并不显著差于单个 $\pi$-学习器。我们将其归因于任务多样性之福：多任务训练可因任务间正向迁移而同时惠及优化动态（Kim et al., 2024）与验证表现（Tripuraneni et al., 2021；Maurer et al., 2016；Ruder, 2017）。

3.3 Error is imbalanced across masking problems

3.3 误差在掩码问题上失衡

In previous sections, we have demonstrated that the hardness of different masking problems $p_{\theta}(x^{i}\mid x_{0}[M])$ can vary significantly, potentially hindering the MDM’s learning. In this section, we provide empirical evidence that the MDM’s final performance exhibits a similar imbalance across subproblems. Details are provided in App. C.2.

前几节已表明不同掩码问题 $p_{\theta}(x^{i}\mid x_{0}[M])$ 的难度可相差很大，可能妨碍 MDM 学习。本节给出经验证据：MDM 的最终表现在子问题上呈现类似失衡。细节见附录 C.2。

L&O-NAE-SAT.

Consider an L&O distribution with $\pi$ given by the identity permutation and where each observation $\mathcal{O}_{j}$ is deterministically given by $\mathrm{NAE}(x_{i_{1}},x_{i_{2}},x_{i_{3}})\triangleq 1-\mathbf{1}[x_{i_{1}}=x_{i_{2}}=x_{i_{3}}]$ for some randomly chosen (prefixed) triples $(i_{1},i_{2},i_{3})\in[N]$. For an MDM trained on this distribution, we measure the error it achieves on each task $\log p_{\theta}(x_{0}|x_{0}[M])$ via $\mathbb{E}_{x_{0}}\Bigl\|\log p_{\theta}(x_{0}|x_{0}[M])-\log p_{\rm data}(x_{0}|x_{0}[M])\Bigr\|^{2}$, where $p_{\rm data}(x_{0}|x_{0}[M])$ denotes the Bayes-optimal predictor. Technically, we do not have access to this, so instead we train another MDM for a much larger number of iterations and use this as a proxy. Fig. 2 reveals that prediction tasks for latent positions (light region) exhibit larger errors compared to those for observation positions (dark region).

L&O-NAE-SAT。

考虑 $\pi$ 为恒等置换的 L&O 分布，每个观测 $\mathcal{O}_{j}$ 由 $\mathrm{NAE}(x_{i_{1}},x_{i_{2}},x_{i_{3}})\triangleq 1-\mathbf{1}[x_{i_{1}}=x_{i_{2}}=x_{i_{3}}]$ 确定性给出，其中三元组 $(i_{1},i_{2},i_{3})\in[N]$ 预先随机选定。对在该分布上训练的 MDM，我们用 $\mathbb{E}_{x_{0}}\Bigl\|\log p_{\theta}(x_{0}|x_{0}[M])-\log p_{\rm data}(x_{0}|x_{0}[M])\Bigr\|^{2}$ 度量各任务 $\log p_{\theta}(x_{0}|x_{0}[M])$ 上的误差，其中 $p_{\rm data}(x_{0}|x_{0}[M])$ 表示贝叶斯最优预测器。技术上我们得不到它，于是另训一个迭代次数远更多的 MDM 作代理。图 2 显示：潜变量位置上的预测任务（浅色区）误差大于观测位置（深色区）。

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 $\mathbb{E}_{x_{0}\sim p_{\rm{data}}}\left[\sum_{i=0}^{L-1}\log p_{\theta}\left(x_{0}^{\pi(i)}\Big|x_{0}[\pi\{i,\ldots,L-1\}]\right)\right]$. This captures the accumulation of error across subproblems $p_{\theta}\left(x_{0}^{\pi(i)}\Big|x_{0}[\pi\{i,\ldots,L-1\}]\right)$, since $p_{\theta}(x_{0}|x_{0}[M])=p_{\rm{data}}(x_{0}|x_{0}[M])$ minimizes this metric. Fig. 2 shows a clear gap between different subproblems.

文本。

这里回顾第 3.2 节的文本实验。因得不到贝叶斯最优预测器，我们采用度量 $\mathbb{E}_{x_{0}\sim p_{\rm{data}}}\left[\sum_{i=0}^{L-1}\log p_{\theta}\left(x_{0}^{\pi(i)}\Big|x_{0}[\pi\{i,\ldots,L-1\}]\right)\right]$。它刻画子问题 $p_{\theta}\left(x_{0}^{\pi(i)}\Big|x_{0}[\pi\{i,\ldots,L-1\}]\right)$ 上误差的累积，因为 $p_{\theta}(x_{0}|x_{0}[M])=p_{\rm{data}}(x_{0}|x_{0}[M])$ 极小化该度量。图 2 显示不同子问题之间存在清晰差距。

The theoretical and empirical evidence demonstrates that MDMs perform better in estimating $p_{\theta}(x_{0}|x_{0}[M])$ for some subproblems $M$ than for others. We therefore want to avoid encountering hard subproblems $M$ at inference time. In the next section, we show that while vanilla MDM inference can run into such subproblems, simple modifications at the inference stage can effectively circumvent these issues, resulting in dramatic, training-free performance improvements.

理论与经验证据表明：对某些子问题 $M$，MDMs 估计 $p_{\theta}(x_{0}|x_{0}[M])$ 更好。因此我们希望在推理时避开困难子问题 $M$。下一节将说明：朴素 MDM 推理可能撞上这类子问题，但推理阶段的简单修改就能有效规避，从而带来显著的无需再训练的性能提升。

4 MDMs can plan around hard problems

4 MDMs 可以绕开困难问题来规划

We previously argued that due to the complex nature of masking subproblems, MDM must perform poorly on certain ones $p_{\theta}(x^{i}|x_{t})$. Therefore, during vanilla MDM inference, MDM inevitably encounters such difficult subproblems 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.

前文论证：因掩码子问题复杂，MDM 必然在某些 $p_{\theta}(x^{i}|x_{t})$ 上表现差。于是朴素 MDM 推理在步骤（b）难免撞上困难子问题。这似乎意味着必须从根本上重审 MDM 的训练方式；但本节将说明：出人意料的是，仅在推理阶段做简单修改——无需再训练——就能绕开这些问题，并带来显著性能提升。

MDM offers multiple sampling paths.

The vanilla MDM inference (Algorithm 2.1.2) aim to align the intermediate distributions with the forward process, as used in continuous 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.

MDM 提供多条采样路径。

朴素 MDM 推理（算法 2.1.2）旨在让中间分布与前向过程对齐，这与连续扩散的用法一致。但与连续扩散不同，MDM 的反向过程允许多条有效采样路径（不同的揭开 token 顺序），它们都与 MDM 前向过程的起始分布匹配。

We first show that when we have an ideal MDM that perfectly solves all masking problems, i.e., $p_{\theta}(x_{0}^{i}|x_{0}[M])=p_{\rm{data}}(x_{0}^{i}|x_{0}[M])$, then using any sampling path (unmasking the tokens in any order) results in the same distribution. Consider the following sampler: For every step, $S$ is a set with one index selected agnostically (without following any distribution). For any clean sample $x_{0}$ generated by this sampler, note that $p_{\theta}(x_{0})=\prod_{i=0}^{L-1}p_{\theta}\left(x_{0}^{\pi(i)}\Big|x_{0}[\pi\{i,\ldots,L-1\}]\right)$ by chain rule, and this is equal to $\prod_{i=0}^{L-1}p_{\rm{data}}\left(x_{0}^{\pi(i)}\Big|x_{0}[\pi\{i,\ldots,L-1\}]\right)=p_{\rm{data}}(x_{0})$. Therefore, other choices of $S$, not necessarily following Algorithm 2.1.2, still capture the true likelihood.

我们首先说明：若有理想 MDM，能完美解所有掩码问题，即 $p_{\theta}(x_{0}^{i}|x_{0}[M])=p_{\rm{data}}(x_{0}^{i}|x_{0}[M])$，则沿任意采样路径（按任意顺序揭开 token）得到的分布相同。考虑如下采样器：每一步 $S$ 是不可知地（不遵循特定分布）选出的单下标集合。对该采样器生成的任一干净样本 $x_{0}$，由链式法则有 $p_{\theta}(x_{0})=\prod_{i=0}^{L-1}p_{\theta}\left(x_{0}^{\pi(i)}\Big|x_{0}[\pi\{i,\ldots,L-1\}]\right)$，且等于 $\prod_{i=0}^{L-1}p_{\rm{data}}\left(x_{0}^{\pi(i)}\Big|x_{0}[\pi\{i,\ldots,L-1\}]\right)=p_{\rm{data}}(x_{0})$。因此，即便 $S$ 的选取不遵循算法 2.1.2，仍能刻画真实似然。

Figure 3: Generative Perplexity. We compare the resulting generative perplexity (GenPPL) of adaptive vs. vanilla MDM inference. We employ a pretrained $170$M MDM and LLaMA-$7$B (Touvron et al., 2023) as inference and evaluation, respectively. Adaptive MDM inference ( Blue) leads to a substantial reduction in generative perplexity, while maintaining the entropy.

图 3：生成困惑度。 比较自适应与朴素 MDM 推理的生成困惑度（GenPPL）。推理用预训练 170M MDM，评估用 LLaMA-7B（Touvron et al., 2023）。自适应 MDM 推理（蓝）在保持熵的同时大幅降低生成困惑度。

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:

实践中，与理想情形不同，MDM 并非在所有子问题上同样出色，如第 3.3 节所示。因而不同采样路径对应不同的似然建模能力。受此启发，我们考虑 MDMs 的自适应推理：

Instead of selecting $S$ randomly, adaptive MDM inference leverages an oracle $\mathcal{F}(\theta,x_{t})$ to select $S$ strategically to avoid hard masking problems. This naturally raises the question of how to design an effective oracle $\mathcal{F}$.

不再随机选取 $S$，自适应 MDM 推理借助 oracle $\mathcal{F}(\theta,x_{t})$ 策略性地选择 $S$，以避开困难掩码问题。这自然引出如何设计有效 oracle $\mathcal{F}$ 的问题。

In the following sections, we demonstrate that adaptive MDM inference with careful choices of $\mathcal{F}$ enhance MDM’s likelihood matching ability. In other words, a pretrained MDM, even if it performs poorly on certain hard subproblems, still contains sufficient information to avoid them when paired with an effective oracle $\mathcal{F}$.

下文各节表明：搭配精心选择的 $\mathcal{F}$，自适应 MDM 推理能增强 MDM 的似然匹配能力。换言之，即便预训练 MDM 在某些困难子问题上表现差，只要配上有效的 oracle $\mathcal{F}$，仍含有足够信息去避开它们。

4.1 Effective design of ordering oracle

4.1 顺序 oracle 的有效设计

We introduce two different oracles, Top probability and Top probability margin. Intuitively, both strategies are based on the idea that $S$ 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.

我们引入两类 oracle：最高概率（Top probability）与最高概率间隔（Top probability margin）。直观上，两者都基于「应按模型对各位的『确定程度』来选 $S$」。需注意：它们不应与 ARMs 中的 nucleus sampling 等概念混淆（Holtzman et al., 2019）；这里的 oracle 用于选择下一步解码的位置而非取值，因而只在 MDMs 语境下有意义。

($N,P$)	Vanilla inference	Adaptive inference
$(25,275)$	78.06%	93.76%
$(30,270)$	75.70%	93.54%
$(40,260)$	74.60%	92.21%
$(50,250)$	67.94%	90.01%
$(100,200)$	62.84%	88.91%

Table 1: L&O-NAE-SAT. Adaptive MDM inference achieves better likelihood matching than vanilla MDM inference. Note that naive guessing leads to $75\%$ accuracy, indicating that vanilla inference performs similarly or worse than naive guessing.

($N,P$)	朴素推理	自适应推理
$(25,275)$	78.06%	93.76%
$(30,270)$	75.70%	93.54%
$(40,260)$	74.60%	92.21%
$(50,250)$	67.94%	90.01%
$(100,200)$	62.84%	88.91%

表 1：L&O-NAE-SAT。自适应 MDM 推理的似然匹配优于朴素 MDM 推理。注意朴素猜测准确率约 75%，说明朴素推理与之相当或更差。

Top probability (Zheng et al., 2023).

Suppose we want to unmask $K$ positions at time step $t$, i.e., select $|S|=K$. In the top probability, the uncertainty of a position is estimated by the maximum probability assigned to any value in the vocabulary. More precisely, the certainty at position $i$ is $\max_{j\in\{0,\ldots,m-1\}}p_{\theta}(x^{i}=j|x_{t})$ and $\mathcal{F}(\theta,x_{t})=\text{Top }K\left(\max p_{\theta}(x^{i}|x_{t})\right)$.

最高概率（Zheng et al., 2023）。

设在时间步 $t$ 要揭开 $K$ 个位置，即选 $|S|=K$。在最高概率策略中，某位置的不确定性由词表上任一取值被赋予的最大概率估计。更精确地，位置 $i$ 的确定性为 $\max_{j\in\{0,\ldots,m-1\}}p_{\theta}(x^{i}=j|x_{t})$，且 $\mathcal{F}(\theta,x_{t})=\text{Top }K\left(\max p_{\theta}(x^{i}|x_{t})\right)$。