基于注意力折扣的自适应采样器用于掩码扩散语言模型

# 基于注意力折扣的自适应采样器用于掩码扩散语言模型
来源：https://arxiv.org/html/2606.10829
尤素夫·萨欣 伯尔尼大学 瑞士伯尔尼 & 艾哈迈德·罗基·赛基亚 洛桑联邦理工学院 瑞士洛桑 & 沃尔坎·切维尔 洛桑联邦理工学院 瑞士洛桑 & 保罗·法瓦罗 伯尔尼大学 瑞士伯尔尼

###### 摘要

掩码扩散语言模型可以通过每次去噪迭代揭示多个标记来减少推理步骤，但这种并行性很脆弱：当位置预测相互耦合时，即使单个位置很有信心，同时提交它们也可能不安全。现有免训练采样器（如 Top-k、Fast-dLLM 和 EB-Sampler）主要控制要揭示的标记数量，同时通常根据逐标记分数对候选进行排序，忽略了所选集合内的交互作用。我们提出了 ADAS，一种用于并行掩码扩散解码的免训练重排序规则。ADAS 保持基础采样器的停止规则不变，仅修改子集构建方式：它贪婪地折扣那些对已选定但预测仍不确定的位置具有强注意力的候选。与将注意力转化为硬兼容性约束的图约束方法不同，ADAS 保持注意力的连续性并将其用作软边际惩罚。在 LLaDA-8B-Base 和 Dream-7B-Base 上对 GSM8K、MATH500、HumanEval 和 MBPP 的评估中，将 ADAS 嵌入 Top-k、Fast-dLLM 和 EB-Sampler，在匹配的去噪器评估次数下，平均分别提高了 9.11 和 10.46 个百分点，每次前向传播的运行时开销仅为 3.1%。这些结果表明，软注意力折扣重排序是一种简单且模块化的方式，可提高掩码扩散语言模型高度并行解码的质量。

## 1 引言

掩码扩散语言模型 (MDLMs) (Sahoo et al., 2024 (https://arxiv.org/html/2606.10829#bib.bib7)) 提供了与因果自回归模型不同的速度-质量权衡：每次去噪步骤可以揭示一个位置子集，而不是单个下一个标记。这种并行性对于推理和代码生成任务尤其有吸引力，因为长输出会使推理延迟成本高昂。因此，核心解码问题不仅是*什么*标记放在每个掩码位置，而且还要考虑*哪些位置*可以在同一步骤中安全提交。

第二个问题正是激进并行解码变得脆弱的地方。许多实用的采样器通过逐标记量（如置信度、边际或熵）对候选位置进行排序，然后使用停止规则或预算来决定揭示多少个位置。当一次只解开少量标记时，此类规则有效，但隐式地认为高得分候选是兼容的。在我们的受控诊断中，强制模型同时预测两个耦合的字面量会将准确率从 71% 降至 30%，这表明联合提交可能比单个预测所暗示的要困难得多。

最近的研究使这种失败模式越来越明显。一些方法将扩散解码视为规划或搜索 (Peng et al., 2025 (https://arxiv.org/html/2606.10829#bib.bib6); Lee et al., 2025 (https://arxiv.org/html/2606.10829#bib.bib12))；其他方法通过调度 (Luxembourg et al., 2025 (https://arxiv.org/html/2606.10829#bib.bib4); Israel et al., 2025 (https://arxiv.org/html/2606.10829#bib.bib5)) 或学习到的解掩码策略 (Jazbec et al., 2025 (https://arxiv.org/html/2606.10829#bib.bib14); Hong et al., 2025 (https://arxiv.org/html/2606.10829#bib.bib13); Bao et al., 2025 (https://arxiv.org/html/2606.10829#bib.bib11)) 来控制并行量。Ben-Hamou et al. (2025 (https://arxiv.org/html/2606.10829#bib.bib9)) 的理论视角还将*模型误差*与*联合依赖误差*分开，阐明当多个位置同时提交时，良好的边际预测并不足够。

我们将这一观点再推进一步。广泛使用的解码器通常在施加于所选子集的*约束*上有所不同，但仍依赖于逐标记目标来构建该子集。Top-k 固定基数，Fast-dLLM 使用置信度阈值条件，EB-Sampler 强制使用熵预算。这些机制决定了*何时停止添加标记*；它们本身并不使候选的边际价值依赖于已选择的其他位置。

我们的提议是让子集构建具有依赖感知性，同时保持停止规则不变。我们引入了 ADAS，一个免训练的基于注意力折扣的自适应采样器，它贪婪地构建要解掩码的位置集。从标准标记置信度分数开始，当某个剩余候选对其已选定但不确定的位置有强注意力时，ADAS 会对其进行折扣。等价地，ADAS 在非加性子集效用下执行贪婪构建：置信度提供个体边际价值代理，而注意力提供成对兼容性代理。

这种设计有意地将排序与预算控制分离。ADAS 不训练规划器、不添加验证器、不执行前瞻搜索，也不重新设计采样器的停止准则。相反，它在现有解码器内部升级贪婪排序步骤，使得相同的注意力折扣更新可以用于固定基数、置信度阈值或熵预算规则。这使得 ADAS 与学习到的解掩码策略 (Jazbec et al., 2025 (https://arxiv.org/html/2606.10829#bib.bib14); Hong et al., 2025 (https://arxiv.org/html/2606.10829#bib.bib13); Asano et al., 2026 (https://arxiv.org/html/2606.10829#bib.bib3))、前瞻方法 (Lee et al., 2025 (https://arxiv.org/html/2606.10829#bib.bib12)) 以及基于调度或验证器的方法 (Luxembourg et al., 2025 (https://arxiv.org/html/2606.10829#bib.bib4); Israel et al., 2025 (https://arxiv.org/html/2606.10829#bib.bib5)) 处于互补位置。

经验上，收益恰好出现在动机所预测的地方：低 NFE、高度并行的区域，其中逐标记排序最为脆弱。在数学推理和代码生成基准测试中，在匹配的去噪器评估预算下，添加 ADAS 分数更新持续改进 Top-k、EB-Sampler 和 Fast-dLLM。平均而言，在 LLaDA-8B-Base 上，ADAS 将匹配 NFE 性能提高了 9.11 点，在 Dream-7B-Base 上提高了 10.46 点，在 90 个匹配操作点中的 80 个上获得了正向增益。

我们的贡献如下。

- • 我们诊断了并行 MDLM 解码中的联合依赖：强制耦合位置一起预测将受控任务准确率从 71% 降至 30%，并且掩码标记自注意力能够区分依赖对和非依赖对。
- • 我们引入了 ADAS，一个免训练的注意力折扣贪婪选择器，它将*排序*与*停止*解耦，可嵌入 Top-k、Fast-dLLM 和 EB-Sampler，每次模型前向传播仅增加 3.1% 的运行时开销。
- • 在两个 MDLM 和四个推理/代码基准测试中，ADAS 将匹配 NFE 性能平均提高了 9.11 和 10.46 点。

## 2 相关工作

关于掩码扩散解码的最新研究表明，生成质量在很大程度上取决于每一步解掩码位置的顺序和分组。因此，一些方法将解码视为显式的规划或策略学习问题。Peng et al. (2025 (https://arxiv.org/html/2606.10829#bib.bib6)) 将采样形式化为规划器引导的解码，Asano et al. (2026 (https://arxiv.org/html/2606.10829#bib.bib3)) 将解掩码位置与解掩码内容分离并学习监督规划器，Lee et al. (2025 (https://arxiv.org/html/2606.10829#bib.bib12)) 使用对候选解码轨迹的前瞻搜索，而 Jazbec et al. (2025 (https://arxiv.org/html/2606.10829#bib.bib14)); Hong et al. (2025 (https://arxiv.org/html/2606.10829#bib.bib13)); Bao et al. (2025 (https://arxiv.org/html/2606.10829#bib.bib11)) 学习解掩码策略或过滤器，以改进手工设计的启发式方法。这些方法表明，逐标记置信度顺序通常不够，尤其是在弱并行区域之外。

另一条互补的工作线通过更好的推理调度或辅助信号来改进扩散解码的速度-质量权衡。Luxembourg et al. (2025 (https://arxiv.org/html/2606.10829#bib.bib4)) 提出了一种结构调度器，通过扩张分组减少有害的交互作用，而 Israel et al. (2025 (https://arxiv.org/html/2606.10829#bib.bib5)) 使用辅助自回归模型在解码过程中调整并行量。Ben-Hamou et al. (2025 (https://arxiv.org/html/2606.10829#bib.bib9)) 通过将解码误差分解为模型误差和联合依赖误差，提供了一个有用的理论框架，并启发了诸如 EB-Sampler 等自适应多标记解掩码规则。Kim et al. (2026b (https://arxiv.org/html/2606.10829#bib.bib10)) 提出了 KLASS，一个免训练采样器，它将标记置信度与连续去噪分布之间的时间 KL 散度相结合，以选择用于解掩码的稳定标记。

最接近的相关方法是 DAPD (Kim et al., 2026a (https://arxiv.org/html/2606.10829#bib.bib2))，它也使用自注意力来估计掩码位置之间的依赖关系。DAPD 通过对称化成对注意力分数、将其阈值化为二值边，并选择一个独立的掩码位置集进行并行解码，从而构建一个注意力诱导的依赖图。因此，其解码规则是图约束的：一个注意力边充当硬冲突，选定的并行批次必须满足近似独立条件。

ADAS 使用相同的广泛信息来源——自注意力——但实例化了一种不同的解码原则。它不构建二值图，不对注意力进行阈值化，也不施加独立集约束。相反，ADAS 保持注意力的连续性，并将其用作贪婪子集构建内部的边际惩罚。因此，ADAS 不声明两个位置不兼容。当一个强耦合候选的置信度增益超过估计的提交风险时，它仍然可以被选择。这种软公式很重要，因为注意力是有害依赖的不完美代理：并非每个高注意力对都应被禁止，也并非每个低注意力对都能保证独立。

这两种方法在集成点上也不同。DAPD 是一个独立的依赖感知解码器，通过图着色定义自己的并行批次。ADAS 是一个与采样器无关的重排序模块：它保持 Top-k、Fast-dLLM 或 EB-Sampler 的停止规则、预算和可接受性条件不变，只改变候选的提出顺序。简而言之，DAPD 使用注意力构建一个禁止同时更新的图；ADAS 使用注意力在现有采样器中计算软边际折扣。

## 3 预备知识

### 3.1 符号

设 V = {1, ..., K} ∪ {m} 表示一个有限词汇表，附加了一个特殊的掩码标记 m。一个序列表示为 x ∈ V^d，其中 d 是序列长度。如果 x_i = m，则位置 i 是*掩码的*，否则是*非掩码的*。对于部分掩码的序列，设 M ⊆ {1, ..., d} 为掩码位置集，M̄ 表示其补集。我们用 x^M̄ 表示已观测到的标记。我们用 q(· | x^M̄) 表示给定当前揭示上下文时掩码位置的真实条件分布，用 p_θ(· | x^M̄) 表示模型的条件分布。

### 3.2 掩码扩散语言模型

掩码扩散语言模型 (Sahoo et al., 2024 (https://arxiv.org/html/2606.10829#bib.bib7)) 通过迭代去噪生成文本。从高度掩码的序列开始，模型反复预测掩码位置的标记分布，并逐步解掩码其中一部分，直到没有掩码剩余。因此，一步解码涉及两个耦合的决策：在选定位置放置*什么*标记，以及下一步在*哪里*解掩码。

在解码步骤中，参数为 θ 的 MDLM 为每个掩码位置 i ∈ M 预测一个条件分布 p_θ(x^i | x^M̄)。许多现有的解码策略随后根据此分布为每个掩码位置分配一个逐标记分数，并根据采样器特定的停止规则构建要解掩码的子集 S ⊆ M。一个常见选择是置信度分数

c_i = max_{x^i} p_θ(x^i | x^M̄),   (1)

即位置 i 最可能标记值的概率。不同的采样器主要在如何约束子集 S 上有所不同，例如，通过固定其大小、阈值化置信度或强制熵预算。

### 3.3 模型误差与联合依赖误差

并行解码的一个核心困难在于，模型为掩码位置提供逐标记条件边缘分布，而在同一步骤中解码多个位置需要推理它们的*联合*行为。正如 Ben-Hamou et al. (2025 (https://arxiv.org/html/2606.10829#bib.bib9)) 所强调的，解掩码子集 S ⊆ M 所产生的误差可以分解为*模型误差*项和*联合依赖误差*项：

∑_{i∈S} D_KL( q(x^i | x^M̄), p_θ(x^i | x^M̄) ) ⏟ 模型误差 + D_KL( q(x^S | x^M̄), ∏_{i∈S} q(x^i | x^M̄) ) ⏟ 联合依赖误差

模型误差捕捉模型逐标记预测的不准确性，而联合依赖误差衡量将选定位置视为条件独立所引入的差异。后者是真实条件分布 q 在 S 上的多信息：它与模型无关，表征了选择用于同时提交的位置之间的内在耦合。

由于 q 在解码时不可观测，ADAS 使用去噪器的自注意力 A_ij 作为耦合强度的模型端代理。当此代理信息量大时，折扣那些对已选定且不确定的位置有强注意力的候选，应能降低下一次并行提交的预期依赖成本。第 4 节 (https://arxiv.org/html/2606.10829#S4) 测试了这一代理假设。

## 4 受控依赖诊断

在介绍基于注意力折扣的自适应采样器 (ADAS) 之前，我们测试注意力引导并行解掩码背后的两个假设：联合解掩码依赖位置比联合解掩码独立位置更困难，并且自注意力为此类依赖提供了有用的代理。

我们使用具有已知依赖结构的合成算术谓词。完整的谓词集为

E = { A+B+C=D, A·B·C=D, min{A,B,C}=D, max{A,B,C}=D }.
请参阅图注
图 1：依赖对和非依赖对的成对注意力值分布