一句话总结

来自哈佛大学、麻省理工学院、加州大学伯克利分校和得克萨斯大学奥斯汀分校的研究人员证明，掩码扩散模型（MDM）相对于自回归模型会学习到计算上难以处理的填充子问题，然而自适应地选择token解码顺序可以绕过这些子问题，将预训练MDM在数独上的准确率从不足7%提升到约90%，并超越了参数规模大7倍的教师强制自回归模型。

核心贡献

• 本文提供了理论和实证证据，表明掩码扩散模型（MDM）相较于自回归模型面临计算上难以处理的训练子问题，这些发现通过基于稀疏谓词观测和空洞预测的新平均情况困难性构造加以形式化。

• 这项工作为MDM推理引入了一种自适应token解码顺序策略，通过动态选择token生成顺序，使模型能够绕过训练中遇到的困难子问题。

• 结果表明，这种自适应推理方法将数独谜题的求解准确率从不足7%提高到约90%，超越了通过教师强制显式训练学习固定解码顺序、参数规模大七倍的自回归模型。

引言

扩散模型主导了图像和音频合成等连续生成任务，但将其适配到文本和蛋白质等离散领域仍在发展之中。掩码扩散模型（MDM）学习逆转随机token掩码过程，是一种有前景的方法，但在语言建模中往往不如自回归模型（ARM）表现出色。虽然MDM在推理时可以通过任意顺序解码提供更大的灵活性，但它们面临更困难的训练目标：从任意未掩码上下文中预测token。作者通过证明MDM即使在良性数据分布下也可能遇到计算上难以处理的掩码子问题来研究这种权衡。为克服这一问题，他们提出了自适应推理策略，智能地选择下一个要解掩码的token，有效绕过这些困难子问题。他们在逻辑谜题上的实证结果显示了显著提升，通过简单的自适应策略将数独准确率从不足7%提升到近90%，甚至超过了通过教师强制训练学习最优解码顺序的ARM。

方法

掩码扩散模型（MDM）为序列生成提供了连续扩散模型的离散状态推广。作者同时构造了前向掩码过程和学习的逆向过程，然后将训练目标重新解释为与顺序无关的学习。他们进一步表明，这种与顺序无关的公式化使模型暴露于计算上困难的填充问题，并提出了一种自适应推理策略，基于模型置信度选择要解掩码的位置，从而无需任何额外训练即可绕过这些困难子问题。

前向过程在词汇表$\{1, \ldots, m\}${1,…,m}（附加掩码token $0$0）上破坏干净序列$x_0 \sim p_{\mathrm{data}}$x0​∼pdata​。在噪声水平$t \in [0,1]$t∈[0,1]下，每个坐标独立地以概率$\alpha_t$αt​保持不变，或以概率$1 - \alpha_t$1−αt​被替换为掩码token，其中调度满足$\alpha_0 \approx 1$α0​≈1和$\alpha_1 \approx 0$α1​≈0。逆向过程通过条件分布$q_{s|t}(x_s | x_t, x_0)$qs∣t​(xs​∣xt​,x0​)（其中$s < t$s<t）定义：未被掩码的坐标保持固定，而被掩码的坐标则作为掩码token与真实token $x_0^i$x0i​之间的类别混合被重新采样。去噪网络$p_\theta(x_t, t)$pθ​(xt​,t)被训练来预测$x_0$x0​上的边际分布，逆向转移通过将网络输出替代$x_0$x0​来近似。训练目标为基于ELBO的损失

$$\mathcal{L}_{\theta} = \int_0^1 \frac{\alpha_t'}{1 - \alpha_t} \mathbb{E}_{\substack{x_0 \sim p_{\mathrm{data}} \\ x_t \sim q_{t|0}(\cdot|x_0)}} \sum_{i: x_t^i = 0} -\log p_\theta(x_0^i | x_t, t) \, dt.$$Lθ​=∫01​1−αt​αt′​​Ex0​∼pdata​xt​∼qt∣0​(⋅∣x0​)​​i:xti​=0∑​−logpθ​(x0i​∣xt​,t)dt.

在实践中，作者指出无时间嵌入的架构通常就足够了，因为被掩码token的数量隐式编码了噪声水平。采样从完全掩码的序列开始，通过从学习到的逆向转移中采样来迭代地解掩码位置。

作者随后重新表述了MDM训练损失，以揭示其与掩码语言建模的联系。设$x_0[M]$x0​[M]表示通过掩码集合$M \subseteq [L]$M⊆[L]中所有索引而从$x_0$x0​获得的部分掩码序列。在温和假设下，损失分解为

$$\mathcal{L}_{\theta} = - \sum_{M \subseteq [L], i \in M} \frac{1}{|M|} \frac{1}{\binom{L}{|M|}} \mathbb{E}_{x_0 \sim p_{\mathrm{data}}} \left[ \log p_\theta(x_0^i | x_0[M]) \right].$$Lθ​=−M⊆[L],i∈M∑​∣M∣1​(∣M∣L​)1​Ex0​∼pdata​​[logpθ​(x0i​∣x0​[M])].

该表达式表明，MDM目标是在所有可能的填充掩码上的损失的线性组合：最优预测器必须解决每个掩码问题，即基于任意可见上下文子集预测每个token的后验边际。作者将此称为与顺序无关的训练，因为模型必须在每种掩码模式下具有同等能力。相比之下，自回归模型（ARM）仅学习一组受限的$L$L个填充问题，每个问题在给定所有前序token的情况下预测下一个token，作者将其称为顺序感知训练。这一区别对后续分析至关重要：与顺序无关的训练严格要求更高，作者表明MDM训练中遇到的许多掩码子问题在计算上是难以处理的。

为使这种困难性精确化，作者引入了一族称为潜变量-观测（L&O）分布的数据分布。在这种分布中，序列包含从先验中采样的$N$N个潜在token，以及$P$P个观测token，每个观测token通过观测函数$\mathcal{O}_1, \ldots, \mathcal{O}_P$O1​,…,OP​作为潜在变量子集的（可能随机化的）函数生成。顺序感知训练对于这些分布在计算上是可处理的：通过重排序列使潜在token出现在前面，每个token可以在给定其前序token的情况下被高效预测。然而，与顺序无关的训练会遇到困难的掩码问题。具体而言，在稀疏谓词观测的示例中，每个观测token是将固定谓词$g$g应用于潜在变量的$k$k元素子集后的输出，作者证明在一定范围的掩码概率下，基于可见上下文预测任何被掩码的潜在token对于多项式时间算法而言在计算上是难以处理的。这种困难性在平均情况下成立，即使算法知道所有分布参数也依然如此，这证明困难性固有于掩码问题本身，而非特定训练过程。因此作者认为，在与顺序无关的训练中，很大一部分子问题在计算上是困难的，并实证验证了在此类数据上训练的MDM表现出性能不均衡：在简单子问题上误差较小，但在困难子问题上误差较大。

鉴于这种困难性，一个自然的问题是MDM在采样过程中是否仍能保持有效。作者表明答案是肯定的：MDM的逆向过程允许多条有效的采样路径。普通推理在每一步随机选择要解掩码的位置，而作者观察到，对于能完美解决所有掩码问题的理想模型，任何解掩码顺序都会产生相同的最终分布，因为链式法则无论顺序如何都会以相同方式分解似然。在实践中，由于模型并非在所有子问题上都同样准确，不同的采样路径会导致不同的似然估计。这推动了自适应MDM推理的提出：作者建议使用预言机$\mathcal{F}(\theta, x_t)$F(θ,xt​)策略性地选择接下来要解码的位置，从而避免困难的掩码问题，而非随机选择要解掩码的位置集合。

作者引入了两种具体的预言机设计。第一种是Top-K概率，按词汇表上的最大预测概率对掩码位置进行排序，即$\max_{j} p_\theta(x^i = j | x_t)$maxj​pθ​(xi=j∣xt​)，并解掩码置信度最高的$K$K个位置。这是置信度的简单有效代理，但当模型在两个或多个高概率值之间不确定时可能会产生误导。为解决此问题，作者提出了Top-K概率边际，该策略按每个位置两个最可能值之间的差值$|p_\theta(x^i = j_1 | x_t) - p_\theta(x^i = j_2 | x_t)|$∣pθ​(xi=j1​∣xt​)−pθ​(xi=j2​∣xt​)∣对位置进行排序。当模型将相似概率分配给多个token时，这种基于边际的得分能提供更可靠的不确定性估计，而当单一值占主导时则退化为Top-K准则。两种预言机都是纯推理时修改：它们不需要额外训练，可应用于任何预训练MDM。作者证明，将预训练MDM与有效的排序预言机配对使用，可以使模型绕过先前识别的困难子问题，尽管模型在困难的掩码任务上表现不完美，仍能大幅改善似然匹配。

实验

实验首先确立了一个事实：掩码扩散模型（MDM）在其训练的众多掩码子问题上存在困难性不均衡：与顺序无关的训练产生的预测任务在计算上比从左到右的自回归训练更困难，且模型的最终误差集中在这些更困难的子问题上。为解决这一问题，作者引入了自适应MDM推理，在生成时选择更简单的子问题。该策略在L&O-NAE-SAT、文本基准测试（包括LLaDA 8B）和逻辑谜题上始终优于普通推理，其中Top-K概率边际在数独上的求解准确率达到89.49%，而普通推理仅为6.88%。此外，自适应MDM推理优于自回归模型，即使这些模型使用显式排序信息训练也是如此，并且在泛化到更困难谜题时表现出更强的鲁棒性，这表明MDM的与顺序无关训练结合自适应推理为序列生成提供了一种更灵活、更有效的方法。

在所有测试的(N, P)配置中，自适应MDM推理相比普通MDM推理显著提高了观测token的预测准确率。普通推理的表现接近或低于75%的朴素猜测基线，而自适应推理的准确率始终保持在88%以上，最佳结果超过93%。对于每对(N, P)，自适应推理都以显著优势超越普通推理，差距约为15到26个百分点。普通推理的得分在62.84%到78.06%之间，有时低于75%的朴素猜测基线，而自适应推理始终保持在88%以上。随着问题规模增大，准确率差距也在扩大，在(100, 200)处观察到最大的相对改进，自适应推理接近89%，而普通推理降至63%以下。

在数独谜题上，自适应掩码扩散模型推理，特别是Top-K概率边际策略，取得了最高准确率，超越了使用或不使用排序信息训练的自回归模型以及普通掩码扩散推理。鉴于掩码扩散模型使用的参数远少于自回归基线，这一改进尤为显著。Top-K概率边际自适应推理达到89.49%的准确率，超过了使用排序信息训练的最强自回归基线（87.18%）。自适应Top-K概率边际方法大幅优于普通掩码扩散推理（6.88%）和Top-K概率采样（18.51%）。最佳掩码扩散模型仅使用6M参数，而自回归模型使用42M参数，同时仍实现了更高的准确率。

在斑马谜题上，自适应MDM推理方法显著优于普通MDM和ARM，以更少的参数实现了超过98%的准确率。Top-K和Top-K概率边际策略均展现出显著提升，其中Top-K略占优势，而即使带有排序信息的更强ARM基线仍低于这些自适应MDM方法。自适应MDM策略的准确率超过98%，远超普通MDM的76.9%以及ARM（即使使用排序信息训练）。最佳自适应MDM方法使用Top-K概率，而Top-K概率边际实现了几乎相同的准确率，两者均仅使用19M参数。ARM需要42M参数，但仍然不如自适应MDM，包括使用监督排序信息的变体。

针对LLaDA 8B模型的自适应推理策略在编码和数学基准测试上始终优于普通推理。Top-K概率边际策略展现出最强的提升，尤其是在更困难的代码补全和数学推理任务上，而两种自适应方法在通用知识和故事推理上均保持竞争力。两种自适应策略在所有评估任务上都优于普通推理，其中边际变体在多行和分行代码生成上取得了最大改进。Top-K概率边际在数学推理和分行代码等具有挑战性的任务上优于普通Top-K，但在MMLU上并非如此，这表明其能更好地处理模糊的token概率。

在需要未见策略或回溯的困难数独谜题上，基于Top-K概率边际的自适应MDM推理策略取得了最高准确率，大幅超越使用排序信息训练的ARM。采用自适应推理的MDM在向更困难谜题分布偏移时也表现出更强的鲁棒性，尽管使用的参数远少。Top-K概率边际自适应MDM推理在困难数独上以显著优势超越带有排序信息的ARM。随机MDM推理表现不佳，这证明了自适应推理策略对MDM的关键作用。最佳MDM在仅使用约七分之一参数的情况下实现了比ARM更高的准确率。

在包括观测token预测、数独、斑马谜题以及编码/数学任务在内的多个基准测试中，自适应掩码扩散模型（MDM）推理始终优于普通MDM，并且通常优于自回归模型（ARM），其中Top-K概率边际策略带来了最强的提升。自适应推理即使在参数显著减少的情况下也能实现高准确率，展现出对更困难问题实例的鲁棒性以及对模糊token概率的更好处理。这些改进凸显了自适应推理策略对MDM的关键作用。