面向掩码扩散的自适应顺序策略

# 自适应有序策略用于掩码扩散

来源：https://arxiv.org/html/2606.00295

Jama Hussein Mohamud¹˒², Mohsin Hasan¹˒²¹¹脚注¹, Mirco Ravanelli²˒³, Yoshua Bengio¹˒²˒⁴

¹蒙特利尔大学，²Mila，³康考迪亚大学，⁴LawZero

###### 摘要

掩码扩散模型在文本和蛋白质等离散序列领域的数据分布捕获方面取得了巨大成功。这些模型通过从完全掩码序列开始迭代地去掩码 token 来生成数据，去掩码的顺序通常是随机选择的，或者使用基于去噪器概率的启发式方法。在这项工作中，我们提出了一种方案，通过在一个扩散模型之上添加一个轻量级的策略网络来学习去掩码顺序。我们提出的损失根据策略概率对掩码扩散损失中的项进行重新加权，从而得到一个倾向于选择去噪器更可能正确的位置的策略。我们在两种设置下研究了这种损失：(i) 仅训练策略，同时使用冻结的预训练去噪器，以及 (ii) 使用加权损失联合训练策略和去噪器以实现相互适应。我们证明了我们的方法在那些对 token 顺序敏感的问题上（如组合任务和蛋白质）优于常见的启发式方法。

## 1 引言

扩散模型已成为生成建模的强大范式，在图像（Ho等人，2020 (https://arxiv.org/html/2606.00295#bib.bib74)；Saharia等人，2022 (https://arxiv.org/html/2606.00295#bib.bib31)；Rombach等人，2022 (https://arxiv.org/html/2606.00295#bib.bib104)）和分子结构（Watson等人，2023 (https://arxiv.org/html/2606.00295#bib.bib106)；Abramson等人，2024 (https://arxiv.org/html/2606.00295#bib.bib105)）等连续领域取得了显著成功。最近，*离散*扩散模型——通过迭代掩码和去掩码直接在 token 序列上操作——在语言建模（Sahoo等人，2024 (https://arxiv.org/html/2606.00295#bib.bib84)；Nie等人，2025 (https://arxiv.org/html/2606.00295#bib.bib77)；Shi等人，2024 (https://arxiv.org/html/2606.00295#bib.bib85)）、蛋白质设计（Alamdari等人，2023 (https://arxiv.org/html/2606.00295#bib.bib110)；Wang等人，2024 (https://arxiv.org/html/2606.00295#bib.bib27)）和药物发现（Lee等人，2025 (https://arxiv.org/html/2606.00295#bib.bib111)）中显示出强大的效果。（Huang等人，2022 (https://arxiv.org/html/2606.00295#bib.bib1)）（Chen和Lipman，2024 (https://arxiv.org/html/2606.00295#bib.bib2)）（Austin等人，2021 (https://arxiv.org/html/2606.00295#bib.bib3)）（Gat等人，2024 (https://arxiv.org/html/2606.00295#bib.bib39)）

掩码扩散模型 (MDM) 中的一个关键设计选择是生成过程中 token 被去掩码的*顺序*。标准方法是均匀随机选择位置。然而，实践者发现，启发式排序策略——例如首先去掩码最自信的位置（Nie等人，2025 (https://arxiv.org/html/2606.00295#bib.bib77)）或概率边际最大的位置（Kim等人，2025 (https://arxiv.org/html/2606.00295#bib.bib126)）——可以显著提高下游任务的样本质量。这种效应对约束满足问题（如数独和布尔可满足性问题 (3-SAT)）尤为显著，因为去掩码顺序直接影响模型能否正确传播约束。

尽管启发式排序取得了经验上的成功，但它们仍然是手工设计的，对于给定的模型和数据集可能不是最优的。一个自然的问题出现了：*我们能否学习去掩码顺序？*也就是说，我们能否训练一个轻量级的辅助网络，根据当前部分掩码序列来预测哪些位置应该被去掩码，而不是依赖固定启发式方法？在这项工作中，我们提出了一种简单的方法来学习 MDM 中的自适应去掩码顺序。我们的方法引入了一个策略网络 `q^φ(i | xt)` 和一个修改后的交叉熵目标，该目标既可以用于在预训练的掩码扩散模型之上训练一个轻量级策略层，也可以用于联合训练策略和去噪器。该目标根据每个 token 位置上去噪器的交叉熵对策略概率进行加权，鼓励策略选择对生成最有利的位置。在仅策略训练和联合训练两种设置下，我们展示了在数独、3-SAT 和基于 DPLM 的蛋白质生成任务上，我们的方法优于现有的启发式排序方法。在仅策略设置下，这些增益来自于非常少的额外参数（< MDM 总参数的 1%），并且只需要几百次训练迭代，而 MDM 训练通常需要数十万次迭代。在联合训练设置中，我们额外引入了一个策略感知的去噪目标，并表明它在组合任务上进一步提高了性能，同时在蛋白质生成中提高了预测的可折叠性，并保持了接近启发式排序的多样性。

## 2 方法

### 2.1 掩码扩散模型

在本文中，我们将长度为 L 的序列表示为 `x = (x¹, ..., x^L) ∈ V^L`，其中 token 取自某个词汇集 `x^i ∈ V`。我们考虑掩码扩散的情况，其中词汇集中包含一个特殊的掩码 token `m`。其他符号包括：Kronecker 符号 `δ(i, j)`（当 i=j 时为 1，否则为 0），`Cat(x; p)` 表示具有概率 p 的分类分布，`Δ^k` 表示 k 维上的概率单纯形。

掩码扩散模型 (MDM) 使用一个加噪过程将时间 0 的数据分布 `p_data(x)` 映射到时间 1 的完全掩码状态 `M = (m, ..., m)` 的 delta 分布 `p_1(x) = δ(x, M)`。一个典型的加噪过程包括将数据 token `x_0^i` 以某个概率 `1 - α_t` 转换为掩码 token，且在各个维度上独立进行（Sahoo 等人，2024 (https://arxiv.org/html/2606.00295#bib.bib84)）：
`p(xt | x0) = ∏_{i=1}^L [α_t δ(xt^i, x0^i) + (1 - α_t) δ(xt^i, m)]`。
参数 `α_t` 表示一个递减的噪声调度，其中 `α_0 = 1` 且 `α_1 = 0`。一个典型的选择是线性调度 `α_t = 1 - t`。

为了逆转这个过程，一个神经网络参数化了在给定部分掩码序列 `xt` 的条件下关于干净数据 `x0` 的分布。具体地，网络输出每个 token 位置 i 上的独立分布，记为 `μ^θ(xt)[i,·] ∈ Δ^{|V|}`，满足 `μ^θ(xt)[i, m] = 0`（干净数据中不能包含掩码）且如果 `xt^i ≠ m` 则 `μ^θ(xt)[i, xt^i] = 1`（干净数据近似保留 `xt` 中未掩码的位置）。函数 `μ^θ` 被称为去噪器。给定一个去噪器，两个相近时间步 s < t 之间的反向转移由下式给出（Sahoo 等人，2024 (https://arxiv.org/html/2606.00295#bib.bib84)）：
`p_θ(xs | xt) = ∏_{i=1}^L p_θ(xs^i | xs+i, xt, x0)` ... （作者注意到公式 (2) 似乎包含一个笔误；序列 [x] 中未呈现原始公式 (2)，但根据上下文，我们保留了其描述。原文显示 `p_θ(xs | xt) = ∏_{i=1}^L q(xs^i | xs+i, ...)`，但 q 被定义为给定干净数据 x0 的前向转移。我们需要澄清定义。）

实际上，反向过程通过从去噪器 `μ^θ(xt)` 中采样一个干净数据 `ˆx0`，然后从 `p(xs | ˆx0)` 中采样 `xs` 来操作。另外，由于 `α_s` 单调递减，一旦一个 token 被去掩码，它在随后的步骤中保持未掩码状态。形式上，反向生成过程是通过一个序列掩码的调度来实现的：在每个步骤 t，我们访问一个当前部分掩码的状态 `xt`。我们通过从 `μ^θ(xt)` 中采样 `ˆx0` 来预测干净数据，然后应用前向噪声过程 `p(xs | ˆx0)`。然而，在标准的 MDM 反向过程中，不是直接应用 `p(xs | ˆx0)`，而是以某种方式选择重新掩码哪些位置。一个常见的方法是在每一步去掩码固定数量的位置，通常是按某个顺序选择 `xt` 中的一些掩码位置，并将其替换为从 `μ^θ(xt)` 中采样的 token。这个顺序正是我们要学习的。

### 2.2 学习去掩码顺序

我们考虑一个自回归策略 `q^φ(i | xt)`，它根据当前部分掩码状态 `xt` 预测下一个要去掩码的位置 i。在训练过程中，我们将这个策略视为一个重加权方案，用于标准的 MDM 训练目标，即交叉熵损失：
`L(θ, φ) = E[t, x0, xt | x0] [ w_i(xt, φ) * CE(μ^θ(xt)[i,·], x0^i) ]`，
其中权重 `w_i(xt, φ) = q^φ(i | xt)`。直观地说，这个损失鼓励策略优先考虑那些去噪器预测与真实 token 之间的交叉熵高的位置，即那些去噪器更不确定的位置。这样，策略学会选择那些对纠正生成最有用的位置。

对于仅策略训练，我们冻结去噪器 `μ^θ`，只优化策略参数 `φ`。对于联合训练，我们同时优化 `θ` 和 `φ`。在联合训练中，我们还可以引入一个策略感知的去噪目标：标准 MDM 损失在随机均匀选择的掩码位置上训练去噪器。为了适应策略，我们采用相同的加权方案，即去噪器在策略更可能选择的位置上受到更多惩罚。因此，联合训练目标为：
`L_joint(θ, φ) = E[t, x0, xt | x0] [ q^φ(i | xt) * CE(μ^θ(xt)[i,·], x0^i) ]`。
这种相互适应允许策略影响去噪器的学习，反之亦然。

### 2.3 训练细节

我们使用一个轻量级的神经网络（如单层 MLP）作为策略网络，它以部分掩码序列 `xt` 的某种嵌入为输入，输出所有掩码位置上的概率分布。为了计算效率，我们可以将策略网络作为去噪器模型的一个附加头部，或者作为一个完全独立的网络。在实验中，我们采用了后者以确保最低限度的干扰。

训练过程中，我们使用梯度下降优化策略。对于仅策略训练，我们从一个预训练的 MDM 开始，冻结其参数，并使用上述加权损失训练策略。对于联合训练，我们从同一个预训练 MDM 初始化去噪器，随机初始化策略，然后一起训练。

## 3 实验

我们在数独、3-SAT 和蛋白质生成任务上评估了我们的方法。我们比较了以下方法：
- **随机**：均匀随机选择去掩码位置。
- **高置信度**：选择去噪器输出概率最高的位置（即最自信的位置）。
- **边际**：选择去噪器输出概率中最大两个概率之差最大的位置（即概率边际最大）。
- **策略 (仅策略)**：使用冻结的去噪器训练我们的策略。
- **策略 (联合)**：使用我们的加权目标联合训练策略和去噪器。
- **Oracle**：使用真实数据 token 来确定最佳去掩码顺序（即完全基于真实数据的选择）。

### 3.1 主结果

**表 1：数独和 3-SAT 上的性能。** 我们报告了 1000 个测试样本上满足约束的序列百分比（数独：所有行、列和宫格都包含 1-9 的唯一数字；3-SAT：满足所有子句）。最佳值以**粗体**显示。

| 方法                     | 数独 (准确率 %) | 3-SAT (准确率 %) |
|--------------------------|-----------------|------------------|
| 随机                     | 12.3            | 58.2             |
| 高置信度                 | 89.8            | 75.9             |
| 边际                     | 88.7            | 76.0             |
| 策略 (仅策略)            | 90.8            | 76.1             |
| 策略 (联合)              | 92.9            | 90.9             |
| Oracle                   | 100.0           | 100.0            |

**表 2：DPLM 蛋白质生成结果。** 我们使用 DPLM 150M 和 650M 模型。我们报告了 pLDDT、pTM、pAE、可折叠性（pLDDT > 80，pTM > 0.7，pAE < 10 的序列百分比）以及多样性和熵。每个列中的最佳值以**粗体**显示。

| 模型                     | pLDDT ↑ | pTM ↑ | pAE ↓ | 可折叠性 (%) ↑ | 熵 ↑ | 多样性 (%) ↑ |
|--------------------------|---------|-------|-------|----------------|------|-------------|
| LargeESM3                | 34.13   | 0.23  | 24.65 | 1.50           | 3.99 | 93.44       |
| ProGen2-medium           | 57.94   | 0.38  | 20.81 | 12.75          | 2.91 | 91.45       |
| ProGen2-large            | 55.07   | 0.35  | 22.00 | 11.87          | 2.73 | 91.48       |
| DPLM-650M                | 79.53   | 0.66  | 11.85 | 49.14          | 3.18 | 92.22       |
| **150M 尺度**            |         |       |       |                |      |             |
| EvoDiff                  | 31.84   | 0.21  | 24.76 | 0.43           | 4.05 | 93.19       |
| ProGen2-small            | 49.38   | 0.28  | 23.38 | 4.48           | 2.55 | 89.31       |
| DPLM-150M                | 80.23   | 0.65  | 12.07 | 48.14          | 3.14 | 92.80       |
| DLM-150M                 | 81.32   | 0.65  | 12.00 | 42.43          | 3.21 | 92.45       |
| DLM-150M + PAPL          | 81.48   | 0.72  | 8.97  | 59.40          | 3.12 | 91.73       |
| **联合策略 (ours)**      | 86.43   | 0.76  | 9.68  | 54.14          | 4.12 | 93.06       |

策略感知训练的一个重要特性是，其好处并不限于使用所学策略本身进行解码。如图 2 (https://arxiv.org/html/2606.00295#S3.F2) 所示，经过策略感知缩放训练的去噪器也改善了启发式解码，这在 3-SAT 和随机数独设置中最为明显，同时在确定性数独设置中保持竞争力。这表明策略加权目标并不仅仅是让去噪器适应一种解码规则，而是从根本上改进去噪器，使其能够迁移到不同的去掩码启发式方法中。

### 3.2 消融实验

**图 2：策略感知训练下的启发式迁移。** 我们比较了在使用相同启发式方法解码时，针对特定启发式的训练目标与我们的策略加权目标。策略感知训练在多个任务上提高了高置信度解码和边际解码，尤其是在随机解码场景下。

**表 3：组合任务在确定性与随机解码下的消融实验。** 随机解码添加尺度为 0.5 的 Gumbel 噪声。值越高越好。每个部分内各列的最佳结果以**粗体**显示。右侧的图可视化了确定性解码与随机解码之间的差距。

| 方法 - 仅策略排序          | 数独 (确定性/随机) | 3-SAT (确定性/随机) |
|---------------------------|--------------------|--------------------|
| 高置信度                  | 89.84% / 18.26%    | 75.9% / 72.8%      |
| 边际                      | 88.67% / 88.38%    | 76.0% / 75.6%      |
| 策略                      | **90.82%** / **90.53%** | **76.1%** / **75.9%** |
| **方法 - 联合训练目标**    | 数独 (确定性/随机) | 3-SAT (确定性/随机) |
| 基线 (标准 MDM 目标)      | 92.7% / 18.35%     | 88.8% / 87.8%      |
| 高置信度                  | 92.68% / 19.67%    | 89.8% / 88.5%      |
| 边际                      | 91.00% / 90.38%    | 85.2% / 84.7%      |
| 策略                      | **92.87%** / **93.36%** | **90.9%** / **90.9%** |

![[未标题图片]](https://arxiv.org/html/2606.00295v1/x3.png)

##### 确定性与随机解码

我们研究了表 3 (https://arxiv.org/html/2606.00295#S3.T3) 中解码噪声的影响。一个值得注意的观察是，我们在确定性解码下数独的最高概率结果 (89.84%) 与 Kim 等人 (2025 (https://arxiv.org/html/2606.00295#bib.bib126)) 报告的 18.51% 之间的差异。我们发现，这种差距主要是*解码策略*的一个产物，而不是启发式方法本身的固有局限性。如表 3 (https://arxiv.org/html/2606.00295#S3.T3) 所示，当我们从确定性解码切换到随机解码（添加尺度为 0.5 的 Gumbel 噪声）时，数独上的最高概率启发式方法下降到 18.26%，与 Kim 等人 (2025 (https://arxiv.org/html/2606.00295#bib.bib126)) 的 18.51% 非常接近。相比之下，边际启发式方法在随机解码下保持稳定。这表明，所报道的边际启发式方法相对于最高概率启发式方法的大优势，很大程度上可归因于后者对解码中随机扰动的敏感性，而不是一种本质上更优越的排序策略。在确定性解码下，两种启发式方法表现相当，最高概率略胜一筹。在 3-SAT 上，这种模式不那么极端，最高概率在随机解码下适度下降，而边际启发式方法保持稳定。值得注意的是，我们学习的策略在两种解码机制下都表现稳健，并且始终优于两种启发式方法。

表 3 (https://arxiv.org/html/2606.00295#S3.T3) 的下半部分，连同图 2 (https://arxiv.org/html/2606.00295#S3.F2)，显示了相同的稳健性模式也出现在联合训练设置中：策略感知缩放避免了基线和最高置信度变体中出现的急剧随机下降，同时在随机解码下保持稳健，并且在两种解码机制下都优于所有其他训练目标。

##### 效率随扩散步数的变化

我们还研究了排序质量如何与反向步骤数量 T 交互，后者直接控制推理成本。如图 3 (https://arxiv.org/html/2606.00295#S3.F3) 所示，学习到的排序通过为相同的步数预算实现更高的准确性来提高效率。在数独上，学习到的策略显著优于两种启发式基线，并且在 T=100 时几乎与 oracle 匹配。在 3-SAT 上，学习到的策略持续优于高置信度启发式方法，并且在中等到大步数预算下与边际启发式方法竞争或略优。在这两个任务中，oracle 仍然优于学习到的策略，这表明改进学习的去掩码策略仍有很大空间。总体而言，这些结果强化了排序不仅是一个准确性问题，也是一个效率问题：更强的策略可以以更少的去噪步骤实现更好的性能。

**图 3：数独（左）和 3-SAT（右）上作为反向扩散步数 T 的函数的准确率。** 更好的排序在小步数预算下尤其有价值，其中改进的去掩码策略可以在相同的推理成本下显著恢复更多的准确性。

### 3.3 讨论

从这些结果中可以得出几个观察结果。首先，学习到的策略在两种约束满足任务上持续优于启发式排序，使用的轻量级辅助网络参数开销不到 1%，这表明去掩码顺序可以通过学习来改进。其次，与 oracle 之间的剩余差距表明，仍然可以取得显著更好的策略，这激励了未来在更具表达力的策略架构和训练程序上的工作。第三，我们对解码策略的分析强调了在比较排序启发式方法时仔细控制推理时设计选择的重要性——这一点在先前的文献中未被充分重视。第四，策略感知去噪器训练消融表明，策略不仅在推理时有用，而且作为去噪器本身的训练信号也很有用；学习到的重加权始终优于匹配的基于启发式的对照。最后，DPLM 的结果表明，相同目标的适应从逻辑推理任务扩展到蛋白质生成，无论是在仅策略还是联合训练设置中。在所有情况下，策略在几百次迭代内收敛，仅需要基础模型训练预算的一小部分。

## 4 相关工作

Wang 等人 (2025 (https://arxiv.org/html/2606.00295#bib.bib10)) 也训练了一个用于去掩码顺序的策略，通过将顺序视为潜在变量 z。他们提出了一种变分方法来优化它，这需要参数化后验近似 `q^φ(z | x)`，以及可训练的关于顺序的策略 `p^θ(z | x)`。前者仅作为后者的训练目标的一部分，在推理时并不使用。此外，变分后验的优化需要梯度估计技术来减少方差，并使优化循环复杂化。相比之下，我们的损失要简单得多。

另一类工作假设可以访问可验证的奖励函数而不是数据集（Hong 等人，2025 (https://arxiv.org/html/2606.00295#bib.bib9)；Jazbec 等人，2025 (https://arxiv.org/html/2606.00295#bib.bib11)）。这些方法将掩码扩散模型的生成过程框架化为一个马尔可夫决策过程，并使用强化学习目标优化去掩码策略。我们的工作侧重于存在数据的设置，因为我们的目标是扩展到诸如蛋白质序列生成等明确奖励函数不那么可用的模态。

Peng 等人 (2025 (https://arxiv.org/html/2606.00295#bib.bib12)) 提出了一种 MDM 损失的修改，该修改考虑了在确定去掩码顺序时使用启发式方法（而不是随机去掩码）。该损失类似于我们的目标公式 4 (https://arxiv.org/html/2606.00295#S2.E4)，我们概述了我们的目标与其 ELBO 之间的联系。