置信捷径：掩码扩散模型的一种推理失效模式

# 自信捷径：掩码扩散模型的推理失败模式
来源：https://arxiv.org/html/2605.29123

###### 摘要

掩码扩散语言模型 (MDMs) 独特地支持任意顺序生成，而基于置信度的解码目前是事实上的标准推理策略。为了优化这一策略，最近的训练方案试图将训练掩码模式直接与生成过程中观察到的模式对齐。然而，我们认为基于置信度的解码本质上与复杂推理所需的逻辑流轨迹不一致，而置信度对齐训练则进一步固化了这种错位。我们使用多位数加法来具体说明这一点：解码策略在解决其长距离依赖之前就过早地预测了局部简单的数字，从而在具有挑战性的输入上产生高置信度错误。虽然传统的随机掩码在这个具有挑战性的长尾上保持了较低的失败率，但置信度对齐训练将错误率放大了一个数量级。在五个不同的推理任务中，同样的模式以任务相关的严重程度出现：基于置信度的解码在高度复杂的输入上引发失败，而置信度对齐训练则加剧了这种情况。相比之下，随机掩码——尽管被认为效率较低——却稳健地保留了解决挑战性长尾所必需的推理轨迹条件。  
²²脚注：通讯作者：Albert No.

## 1 引言

离散扩散模型 (Austin 等, 2021 (https://arxiv.org/html/2605.29123#bib.bib19); Lou 等, 2023 (https://arxiv.org/html/2605.29123#bib.bib10); Campbell 等, 2022 (https://arxiv.org/html/2605.29123#bib.bib20)) 已成为自回归语言建模的替代方案。在离散扩散变体中，掩码扩散模型 (MDMs) (Sahoo 等, 2024 (https://arxiv.org/html/2605.29123#bib.bib12); Shi 等, 2024 (https://arxiv.org/html/2605.29123#bib.bib11); Ou 等, 2024 (https://arxiv.org/html/2605.29123#bib.bib13)) 尤为突出：标记在吸收态 [MASK] 和原始文本之间转换，模型被训练从部分观察到的上下文中重建被掩码的位置。最近的工作表明，MDMs 可以扩展到大规模语言建模 (Nie 等, 2025 (https://arxiv.org/html/2605.29123#bib.bib5); Ye 等, 2025b (https://arxiv.org/html/2605.29123#bib.bib8); Gong 等, 2025 (https://arxiv.org/html/2605.29123#bib.bib21))，并在推理、规划和代码生成基准上取得有竞争力的性能 (Ye 等, 2025a (https://arxiv.org/html/2605.29123#bib.bib7); Zhao 等, 2025 (https://arxiv.org/html/2605.29123#bib.bib22))。MDMs 的一个核心吸引力在于其解码灵活性。与自回归模型不同（它们固守固定的从左到右分解），MDM 可以支持任意的生成顺序：不同位置可以在不同时间被取消掩码，每种解码顺序都会导致对同一序列分布的不同分解 (Chang 等, 2022 (https://arxiv.org/html/2605.29123#bib.bib15); Kim 等, 2025 (https://arxiv.org/html/2605.29123#bib.bib3); Ye 等, 2025b (https://arxiv.org/html/2605.29123#bib.bib8))。主要的推理启发式方法是**基于置信度的解码**，它揭示出具有最大 top-1 概率、间隔或负熵分数的位置。这种启发式方法在生成时直观上很吸引人，因为它首先确定在当前上下文下看起来最容易的标记。它也激发了最近的置信度对齐训练方案。PAPL (Peng 等, 2026 (https://arxiv.org/html/2605.29123#bib.bib17)) 对每个标记的损失进行加权，使其偏向于模型已经自信预测的位置，而 PUMA (Kim 等, 2026b (https://arxiv.org/html/2605.29123#bib.bib4)) 则修改了掩码过程，使得训练状态类似于基于置信度的推理轨迹。

然而，对于推理任务，生成顺序不仅仅是呈现的问题。推理问题通常有一个**逻辑流顺序**：在这个顺序中，中间事实变得合理，而后来的事实变得确定。在多位数加法中，稳定的顺序是从最小有效位数开始，因为每个进位都必须从较低位传播，然后较高位才能完全确定。更一般地，有用的生成顺序是熟练的求解者建立解决方案的顺序。如果以不同的顺序对 MDM 进行解码，那么它必须在后验事实的前提条件仍然被掩码的情况下预测它们，从而迫使模型对未解决的推理状态进行边缘化。因此，基于置信度的解码可能会偏离推理顺序：它偏好局部容易的标记，而不一定是其依赖关系已解决的标记。这种区别在推理分布的长尾部分最为重要。依赖顺序长或严格的输入通常很少见，但它们并非病态的异常值。它们是推理模型最有价值的那些情况的受控版本：长进位链、狭窄的迷宫走廊、深度嵌套的表达式，或者最终难以解决的数学和科学问题——它们的依赖结构无法通过局部启发式方法缩短。一个通过遵循置信度捷径而在常见实例上表现良好的模型，恰恰可能在那些将复杂推理与单纯插值区分开来的输入上失败。

我们使用多位数加法作为最清晰的设置来揭示这种偏差。加法足够简单，正确的依赖结构是确切已知的：每个数字的值取决于从低位传播来的进位，因此唯一的推理顺序是从最小有效位数开始。同时，它允许一个强大的分布捷径。在典型的数字采样下，长进位链很少见，并且高位数字通常可以从一个短的局部窗口预测出来，而无需遍历整个链条。这两个要素——一个已知的推理顺序和一个在大多数输入上有效的可用捷径——使我们能够直接测量解码行为。具体来说，我们可以确定基于置信度的解码是遵循逻辑推理顺序还是捷径，并观察当这两条路径偏离时模型的行为。

我们的实证研究将均匀随机掩码与两种置信度对齐训练方案 (PAPL 和 PUMA) 在五个推理任务（加法、迷宫、ListOps、Countdown 和数独）上进行了比较。我们在每个领域内固定架构和计算量；只有训练干预措施不同。我们评估基于置信度的解码、随机解码以及可用的任务特定逻辑流或求解器推导顺序，并按结构难度对结果进行分层。结果表明，置信度对齐可以放大局部容易的预测与真实推理顺序依赖关系之间的不匹配。根据任务的不同，这种对齐甚至可能导致关键性的整体失败。我们的贡献是：

- • 提出了 MDM 解码的推理顺序视角，并解释了为什么当置信度与逻辑流依赖顺序不同时，基于置信度的解码可能是次优的。
- • 对多位数加法上的置信度捷径给出了具体分析，描述了置信度对齐训练方案如何放大失败。
- • 提供了一个受控的五任务实证研究，表明置信度对齐训练可以以跨任务定性不同的方式放大推理顺序覆盖缺口。

## 2 预备知识

#### 符号。
令 \(\mathbf{x} = (x_1, \dots, x_L) \in \mathcal{V}^L\) 表示一个干净序列，词汇表为 \(\mathcal{V}\)，并附加一个特殊掩码标记 \(\mathtt{M}\)。对于子集 \(\mathbf{M} \subseteq [L] = \{1, \dots, L\}\)，我们记 \(\overline{\mathbf{M}} = [L] \setminus \mathbf{M}\) 为补集，\(\mathbf{x}_{\overline{\mathbf{M}}}\) 为 \(\mathbf{x}\) 在 \(\overline{\mathbf{M}}\) 位置上的子序列。我们将 \(\mathbf{M}\) 称为掩码索引，\(\overline{\mathbf{M}}\) 称为可见索引。

### 2.1 掩码扩散模型

掩码扩散模型 (MDM) 学习从部分掩码的上下文中重建标记。给定 \(\mathbf{x} \sim p_{\text{data}}\)，采样一个掩码率 \(\lambda \sim \mathrm{Unif}(0,1)\)，并且 \([L]\) 中的每个位置独立地以概率 \(\lambda\) 被掩码。令 \(\mathbf{M} \subseteq [L]\) 为结果掩码集。模型观察到 \(\mathbf{x}_{\overline{\mathbf{M}}}\)，并为每个掩码位置分配一个分类分布：
\[ p_{\theta}(x_i \mid \mathbf{x}_{\overline{\mathbf{M}}}) \in \Delta(\mathcal{V}), \quad i \in \mathbf{M}. \]
标准的 MDM 去噪目标为：
\[ \mathcal{L}(\theta) = -\,\mathbb{E}_{\mathbf{x},\,\lambda,\,\mathbf{M}}\left[\frac{1}{\lambda}\sum_{i\in\mathbf{M}}\log p_{\theta}(x_i \mid \mathbf{x}_{\overline{\mathbf{M}}})\right], \]
其中 \(1/\lambda\) 对期望的掩码标记比例进行归一化。

#### 与顺序无关的解释。
MDM 目标等价于对生成顺序的均匀期望 (Kim 等, 2025 (https://arxiv.org/html/2605.29123#bib.bib3)):
\[ \mathcal{L}(\theta) \;\propto\; -\,\mathbb{E}_{\pi \sim \mathrm{Unif}(\mathbb{S}_L)}\left[\sum_{j=1}^L \log p_{\theta}\bigl(x_{\pi(j)}\,\big\|\, \mathbf{x}_{\pi(:\,j)}\bigr)\right], \]
其中 \(\mathbb{S}_L\) 表示 \([L]\) 上排列 \(\pi\) 的对称群，\(\pi(:\,j) = \{\pi(1), \dots, \pi(j-1)\}\) 是在步骤 \(j\) 之前未掩码的前缀。每个生成顺序都被均匀训练；推理时使用的顺序由解码策略决定。

### 2.2 解码策略

MDM 推理从完全掩码的序列开始，并迭代地取消掩码标记。每一步，令 \(\mathbf{M} \subseteq [L]\) 为当前掩码集；模型为每个 \(i \in \mathbf{M}\) 产生一个分布 \(p_{\theta}(\cdot \mid \mathbf{x}_{\overline{\mathbf{M}}})\)。解码策略选择一个揭示集 \(R \subseteq \mathbf{M}\) 并填充这些位置，通常通过贪心预测：
\[ x_i \leftarrow \arg\max_{v\in\mathcal{V}} p_{\theta}(v \mid \mathbf{x}_{\overline{\mathbf{M}}}), \quad i \in R. \]

#### 基于置信度的解码。
最常见的解码策略选择当前模型最自信的位置。对于掩码位置 \(i \in \mathbf{M}\)，定义 top-1 置信度：
\[ c_{\theta}^i = \max_{v\in\mathcal{V}} p_{\theta}(v \mid \mathbf{x}_{\overline{\mathbf{M}}}). \]
基于置信度的解码根据 \(c_{\theta}^i\) 选择得分最高的位置并首先解码它们。变体使用相关的不确定性度量，例如 top-2 概率之间的间隔或负预测熵 (Chang 等, 2022 (https://arxiv.org/html/2605.29123#bib.bib15); Kim 等, 2025 (https://arxiv.org/html/2605.29123#bib.bib3); Ye 等, 2025b (https://arxiv.org/html/2605.29123#bib.bib8))。

#### 置信度对齐训练。
最近的工作修改了 MDM 训练，使得训练分布更好地匹配基于置信度的推理。PAPL (规划感知路径学习) (Peng 等, 2026 (https://arxiv.org/html/2605.29123#bib.bib17)) 保持 i.i.d. 随机掩码过程，但重新加权了掩码位置上的每个标记损失：模型已经自信预测的位置获得更大的损失权重。PUMA (渐进式取消掩码) (Kim 等, 2026b (https://arxiv.org/html/2605.29123#bib.bib4)) 替换了掩码过程本身：从完全掩码的序列开始，它根据模型置信度分数选择的顺序迭代地取消掩码真实标记，并使用得到的中间状态作为去噪上下文。这两种方案都将训练分布转向推理时的置信度轨迹，其基本原理是对齐两者应该提高生成质量。每种方法的细节请参考附录 D (https://arxiv.org/html/2605.29123#A4)。

## 3 加法作为推理顺序的受控透镜

加法作为掩码扩散模型中推理的一个异常清晰的诊断工具。该任务是简单的，其确切的依赖结构是已知的，并且真实推理与高精度捷径之间的差距可以用封闭形式描述。这使得加法成为一个理想的动机案例，然后再检验具有更复杂和模糊依赖关系的任务。

### 3.1 任务设置

我们考虑 32 位数加法。每个示例包含两个操作数和一个和："a+b=c"。提示包含两个操作数和等号，答案区域包含和的 33 个输出数字，包括可能的进位输出数字。在生成过程中，答案区域最初被掩码，模型必须填充所有输出数字。我们评估三种训练方案——标准均匀随机掩码、PAPL 和 PUMA——并结合两种解码策略。第一种是基于置信度的解码，每一步取消掩码置信度最高的数字。第二种是加法特定的、从最小有效位数开始的策略，严格遵守算术依赖顺序。我们根据测试实例的最长进位传播链长度对其进行分层。

### 3.2 最优推理顺序是 LSB 优先

尽管输入和输出字符串通常按最高有效位优先顺序书写，但算术计算是反向流动的。从最小有效位开始索引，令 \(a_0, b_0\) 和 \(c_0\) 为最小有效位，\(c_{32}\) 表示最终进位输出。令 \(r_i\) 表示进入位置 \(i\) 的进位。算术计算如下：
\[ r_0 = 0, \quad c_i = (a_i + b_i + r_i) \bmod 10, \quad r_{i+1} = \mathbf{1}\{a_i + b_i + r_i \geq 10\} \quad (i = 0, \dots, 31), \quad c_{32} = r_{32}. \]
这里，\(c_i\) 在 \(r_i\) 已知后完全确定，而 \(r_i\) 只有在解决较低阶位置后才建立。因此，从 \(c_0\) 到 \(c_{32}\) 取消掩码答案永远不会迫使模型在其算术前提条件可用之前预测一个数字。这使得最小有效位优先 (LSB-first) 解码成为加法的最优逻辑流顺序。

### 3.3 诱人的捷径

加法也在相同的机制中隐藏着一个强大的捷径。定义逐位和 \(s_i = a_i + b_i\)。如果 \(s_i \leq 8\)，则位置 \(i\) 的进位输出被**杀死**：无论 \(r_i\) 如何，\(r_{i+1} = 0\)。如果 \(s_i \geq 10\)，则进位被**生成**：无论 \(r_i\) 如何，\(r_{i+1} = 1\)。只有当 \(s_i = 9\) 时，该位置会**传播**传入的进位，得到 \(r_{i+1} = r_i\)。我们将这些分别分类为 \(k\)、\(g\) 和 \(p\) 单元。一个最大连续 \(p\) 单元序列形成一个**进位链**，该链严格由最近的低阶 \(g\) 或 \(k\) 单元（终止它的单元）界定。为了确定高阶位置的进位，通常只需要向下查看短距离直到遇到一个非传播数字 (\(g\) 或 \(k\)) 即可。