Contrastive Distribution Matching for Amortized Sequential Monte Carlo in Discrete Diffusion¶

会议: NeurIPS 2026
arXiv: 2605.23346
OpenReview: https://openreview.net/forum?id=fGFv81kGOt
代码: https://cdm-smc.github.io/
领域: 优化/理论
关键词: 离散扩散模型, 奖励对齐, 序贯蒙特卡洛, 扭曲函数摊销, 对比学习

一句话总结¶

针对离散扩散模型在奖励对齐中评估最优扭曲函数（Twist Function）高度依赖昂贵蒙特卡罗采样的推理瓶颈，本文提出基于前向 KL 散度的对比分布匹配（CDM）与基于扩散前向核的高效缓冲区训练方案，将扭曲函数评估摊销为单次网络前向传播（额外开销低于 5%），在生物序列设计与扩散大模型对齐等多项任务上实现显著超越基线的帕累托缩放边界。

研究背景与动机¶

离散扩散模型（如掩码扩散模型 MDM）在语言建模、基因调控 DNA 序列设计以及从头蛋白质合成等结构化离散数据生成中取得了重要突破。在这类模型中，下游任务的核心需求往往是对齐特定标量奖励（如人类偏好得分、细胞系活性或蛋白质可设计性），即从由奖励调优的目标分布中进行精确采样。在连续扩散模型中，序贯蒙特卡洛（SMC）因渐近无偏性而被广泛用于奖励对齐，并依赖分数梯度或 Tweedie 公式对中间状态的真实样本均值给出极低计算代价的解析近似。

然而，将扭曲序贯蒙特卡洛（Twisted SMC）迁移至离散状态空间时遇到了根本性的计算瓶颈。由于离散空间天然不可导，基于梯度的局部最优提议需要依赖 Gumbel-Softmax 松弛，极易引入严重的梯度偏差与优化不稳定性，且无法支持不可导的黑盒或 API 奖励。更为致命的是，离散扩散模型缺失类似 Tweedie 公式的解析关系，使得传统 Twisted SMC 必须在逆向去噪的每一步中，从当前状态向终点采样大量蒙特卡罗（MC）干净样本以估计扭曲函数。当奖励模型自身计算代价高昂（如 AlphaFold/ESMFold 类结构折叠模型或超大模型 API 评估）时，这种推理期嵌套采样会带来不可承受的计算开销。

此前基于软 Q 学习（Soft Value）的回归式扭曲函数训练方法尝试在预训练阶段拟合该价值函数，但其仅在奖励无关的基线分布样本上进行均方误差回归，导致训练与推理分布严重失配，在真正关键的高奖励目标区域预测极其不准。本文的核心 idea 是：将 Twisted SMC 扭曲函数的学习建模为对齐目标中间分布与参数化分布的前向 KL 散度极小化，推导出兼具正样本提权与负样本抑制的对比分布匹配（CDM）梯度，并利用离散扩散模型的封闭形式前向加噪核实现单批干净样本跨时间步重复利用的高效免采样训练。

方法详解¶

整体框架¶

CDM 的核心目标是消除离散扩散逆向采样过程中估计扭曲函数 \(\psi_t^*\) 的在线蒙特卡罗展开瓶颈，将其摊销为一个轻量级标量预测头 \(\psi_t^\phi\) 的单次前向推理。方法整体分为两个关键阶段：训练阶段通过维护高奖励正样本经验池，结合闭式前向加噪核与当前参数重要性采样的负样本，执行对比分布匹配更新；推理阶段则利用紧密集成在扩散骨干网络最后一层的轻量级扭曲头直接输出扭曲比率，驱动 Twisted SMC 高效粒子重采样。

%%{init: {'theme': 'base', 'themeVariables': {'primaryColor': '#EEF2FF', 'primaryBorderColor': '#6366F1', 'lineColor': '#4F46E5', 'fontSize': '13px'}, 'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入：离散扩散模型骨干与标量奖励函数"] --> B["前向核分解正样本池构建<br/>从目标分布采样干净样本并前向加噪"]
    B --> C["对比梯度估计与参数更新<br/>正项提升目标似然 + 负项抑制当前分布偏移"]
    C --> D["双头共享骨干轻量级参数化<br/>共享主干特征并行输出 Logits 与 Twist 标量"]
    D --> E["摊销 Twisted SMC 粒子推断<br/>常数时间计算重要性权重完成渐近无偏采样"]

关键设计¶

1. 对比分布匹配目标：解决训练与测试分布失配

传统的回归型扭曲学习（Soft Value）在无偏基线模型分布 \(p_t^{\text{base}}\) 上最小化均方误差，无法准确聚焦于经过奖励指数调控的高概率目标区域。CDM 转而在每个离散时间步对齐调制分布 \(p_t^\phi(\mathbf{x}_t) = \frac{1}{\mathcal{Z}_t^\phi} p_t^{\text{base}}(\mathbf{x}_t) \psi_t^\phi(\mathbf{x}_t)\) 与理论最优目标分布 \(p_t^*(\mathbf{x}_t) \propto p_t^{\text{base}}(\mathbf{x}_t) \psi_t^*(\mathbf{x}_t)\) 之间的前向 KL 散度：

\[\mathcal{L}_{\text{CDM}}(\phi) = \mathbb{E}_t \left[ \mathcal{D}_{\text{KL}}(p_t^*(\mathbf{x}_t) \parallel p_t^\phi(\mathbf{x}_t)) \right]\]

通过对参数 \(\phi\) 求导并应用对数导数技巧，归一化常数 \(\mathcal{Z}_t^\phi\) 的梯度自然化简为在参数化分布 \(p_t^\phi\) 下的期望，从而推导出具有清晰对比结构的梯度形式：

\[-\nabla_\phi \mathcal{L}_{\text{CDM}}(\phi) = \mathbb{E}_t \left[ \mathbb{E}_{p_t^*(\mathbf{x}_t)} \left[ \nabla_\phi \log \psi_t^\phi(\mathbf{x}_t) \right] - \mathbb{E}_{p_t^\phi(\mathbf{x}_t)} \left[ \nabla_\phi \log \psi_t^\phi(\mathbf{x}_t) \right] \right]\]

其中正项促使扭曲函数在目标高奖励样本上增大输出，直接纠正训练-推理分布偏差；负项则在当前模型生成的样本上抑制低质量区域，保证了整体概率密度的校准，避免了回归目标在分布外区域的盲目发散。

2. 基于扩散前向核的高效样本复用梯度估计：破除重复奖励评估壁垒

若直接对目标分布 \(p_t^*\) 在每个时间步独立执行 SMC 采样，会导致高昂的奖励评估开销反向吞噬训练效率。本文证明了离散扩散模型轨迹级最优目标分布具备独特的马尔可夫结构分解：任意中间时间步的目标边际分布精确等于最终干净目标分布经过基线前向转移核的加噪扩散，即 \(p_t^*(\mathbf{x}_t) = \sum_{\mathbf{x}_0} p_0^*(\mathbf{x}_0) p^{\text{base}}(\mathbf{x}_t \mid \mathbf{x}_0)\)。这一性质在自回归生成模型中并不存在。

基于该闭式分解定理，正项梯度可以严格重写为对无噪声最终状态 \(p_0^*(\mathbf{x}_0)\) 与一步离散前向加噪转移 \(p^{\text{base}}(\mathbf{x}_t \mid \mathbf{x}_0)\) 的期望。在算法实现上，CDM 维护容量为 \(B_{\text{buffer}}\) 的干净正样本池 \(\mathcal{B}\)，仅每隔 \(n_{\text{update}}\) 步更新一次，并在每次参数更新时通过掩码加噪核采样任意时间步的带噪状态 \(\mathbf{x}_t\)。这意味着单次高成本的奖励评估生成的干净样本能够跨时间步、跨梯度迭代成百上千次复用，使得训练高代价奖励模型（如结构折叠与大模型 API）在计算上完全可行。对于负样本，CDM 采用单步重要性采样（IS）与指数移动平均（EMA）软目标参数更新，兼顾了采样效率与训练稳定性。

3. 双头共享特征的极轻量扭曲函数参数化：将推理开销压至极致

以往为扩散模型学习价值函数或引导网络的工作通常从零训练一个独立的神经网络，这不仅放弃了预训练去噪网络中蕴含的丰富离散语义表示，还在推理计算 SMC 粒子重采样权重时导致计算量翻倍。

CDM 将扭曲函数 \(\psi_t^\phi(\mathbf{x}_t)\) 直接参数化为附加在去噪骨干网络最后一层特征表示上的轻量级标量头，与原始去噪 Logits 头并行计算。针对不同数据模态，设计了三类轻量结构：标准均值池化 MLP、结合正弦位置编码后逐 Token 映射的 MLP+PE（针对 CNN 等无内置位置编码的骨干）、以及轻量级 2 层 Transformer（针对长文本高阶依赖）。在单次前向计算中，去噪骨干网络特征被两头复用，评估扭曲函数的额外计算耗时仅占骨干网络前向传播的 0.4% 至 4.6%，将每个时间步基于扭曲比率更新粒子权重的代价降低至常数级别。

损失函数 / 训练策略¶

训练完整算法结合了目标缓冲池、闭式加噪以及基于 EMA 的软更新： 1. 缓冲区 \(\mathcal{B}\) 维护来自目标分布的带权样本 \(\{(\mathbf{x}_0^{*(k)}, \tilde{W}_0^{*(k)})\}_{k=1}^{B_{\text{buffer}}}\)，通过运行带扭曲估计的 SMC 生成，每 \(n_{\text{update}}\)（文本任务取 4，蛋白质折叠取 10）步周期刷新； 2. 每次迭代随机抽取离散时间步 \(t \sim \mathcal{U}(\{1, \dots, T\})\)，从池中取出样本通过掩码离散转移概率 \(p(\mathbf{x}_t \mid \mathbf{x}_0) = \text{Cat}(\mathbf{x}_t; \alpha_t \mathbf{x}_0 + (1-\alpha_t)\mathbf{m})\) 瞬时加噪； 3. 从带有 EMA 扭曲权重 \(\psi^{\phi_{\text{EMA}}}\) 的提议中抽取负样本，计算对比梯度并使用 AdamW 优化器（学习率 \(1\times 10^{-4}\) 至 \(3\times 10^{-4}\)，权重衰减 0.01）更新扭曲头参数； 4. 扭曲网络参数平滑同步至 EMA 副本，确保负采样分布轨迹连续平稳。

实验关键数据¶

主实验¶

论文在四种异构的离散扩散对齐任务上进行了系统评测：DiT-MDM 文本毒性生成（OpenWebText）、CNN-MDM 增强子调控 DNA 序列设计（HepG2 细胞系活性）、DPLM-2 蛋白质共生成序列与结构可设计性（基于 ESMFold 折叠的自洽 scRMSD）、以及 8B 规模离散扩散语言模型 LLaDA-8B-Instruct 的偏好对齐（Skywork/ArmoRM）。在固定推理 Wall-Clock 时间下，CDM 建立了全新的帕累托最优边界。

下表汇总了提议微调模型（d1 与 DRAKES）与 CDM 结合以及独立运行在毒性文本生成与蛋白质可设计性上的量化表现：

任务 / 方案	评价指标 (主目标)	多样性指标 1	多样性/质量指标 2	结果表现说明
文本生成: d1 [94]	奖励: -0.933	Self-BLEU: 0.051	PPL: 332.322	GRPO 微调，出现明显的模式坍塌与困惑度恶化
文本生成: DRAKES [82]	奖励: -1.051	Self-BLEU: 0.027	PPL: 322.424	反向传播松弛微调，奖励较低且语言困惑度高
文本生成: CDM (Ours)	奖励: -0.845	Self-BLEU: 0.015	PPL: 124.500	在取得更高对齐奖励的同时，保持了最低的重复率与最优困惑度
蛋白质生成: d1 [94]	奖励: -1.876	聚类数 (Clusters): 14	内部 TM-score: 0.851	出现严重模式坍塌，折叠结构过于相似（高 TM-score）
蛋白质生成: CDM (Ours)	奖励: -1.723	聚类数 (Clusters): 17	内部 TM-score: 0.499	在更高设计自洽性下，蛋白质折叠结构多样性显著更高

各任务中主要构件与奖励模型的运行延迟分析如下表所示，充分印证了扭曲头轻量化对破除 SMC 推理瓶颈的关键作用：

实验领域	扩散骨干单步延迟 \(\mathbf{x}_\theta\) (ms)	CDM 扭曲头延迟 \(\psi^\phi\) (ms)	奖励模型评估延迟 \(r\) (ms)	扭曲头/骨干延迟占比 (%)
毒性文本生成 (DiT)	14.227	0.101	4.414	0.710%
调控 DNA 设计 (CNN)	9.829	0.289	11.438	2.940%
蛋白质可设计性 (DPLM-2)	33.901	0.144	1356.0	0.425%
扩散大模型对齐 (LLaDA-8B)	28.619	1.327	24.268	4.637%

消融实验¶

消融实验深入验证了训练期蒙特卡罗样本数 \(M\)、正样本采样算法（IS vs SMC）以及正样本池刷新频率 \(n_{\text{update}}\) 对最终收敛表现的影响：

训练配置 / 消融项	毒性生成最终奖励 (↑)	蛋白质生成最终奖励 (↑)	机制说明与结论
Soft Value (\(M=1\))	-1.8060	-1.8237	回归基准在单一 MC 样本下受方差干扰表现严重欠拟合
Soft Value (\(M=4\))	-0.9030	-1.6895	增加 MC 样本数改善了价值目标，但训练耗时倍增且性能迅速饱和
CDM (\(M=1\))	-0.4614	-1.6950	对比学习显著缓解对 \(M\) 的依赖，极小 MC 算力下依然超越强回归基线
CDM (\(M=4\))	-0.5079	-1.5725	在复杂折叠任务上进一步取得最优得分，性能大幅领先 Soft Value
CDM 采用重要性采样 (IS)	-1.0367	-1.8049	IS 在正样本采样中存在严重权重退化，导致正梯度估计有偏失真
CDM 采用序贯蒙特卡罗 (SMC)	-0.4614	-1.5725	依赖粒子重采样机制维持了高质量正样本，大幅增强了下游表现

关键发现¶

打破算力与精度的矛盾：在蛋白质设计等奖励极其沉重（ESMFold 折叠单次需 1356 ms）的场景下，传统 SMC（\(M=4\)）在墙钟时间内完全无法实用，而 CDM 扭曲头前向仅需 0.144 ms，不仅让 SMC 大规模粒子采样变为现实，更在保留测试集（scTM）上验证了未见指标的泛化有效性。
正负样本协同优于单向拟合：Soft Value 回归在训练初期迅速陷入平台期，而 CDM 的对比梯度结构由于正项拉升目标态分布、负项惩罚偏离态分布，在相同墙钟训练时间内收敛速度显著加快，且收敛上限显著更高。
正样本缓冲区机制的高效性：即使将正样本更新间隔 \(n_{\text{update}}\) 从 1 扩大至 10，CDM 的对齐指标依然稳定保持在高位，证明了扩散前向扩散核在离散状态下复用样本的极其强大的统计稳健性。

亮点与洞察¶

扩散前向核闭式分解的巧妙应用：巧妙利用了离散扩散边际分布的马尔可夫分解定理 \(p_t^*(\mathbf{x}_t) = \sum_{\mathbf{x}_0} p_0^*(\mathbf{x}_0) p^{\text{base}}(\mathbf{x}_t \mid \mathbf{x}_0)\)，绕过了每个逆向时间步均需调用奖励函数的限制，将极其昂贵的采样转化为廉价的闭式加噪，此技术在自回归生成模型中天然不可用。
与提议微调（Fine-Tuned Proposals）无缝正交兼容：CDM 是基于引导的摊销框架，不绑定特定的转移核；直接将其嫁接在已经通过 RL/GRPO（如 d1）或梯度松弛（如 DRAKES）微调后的模型上，依然能进一步带来显著的复合性能增益。
克服后训练的模式坍塌：微调方法常因强化学习偏置丢失生成多样性（如 d1 和 DRAKES 的 Self-BLEU 与 PPL 严重退化），而 CDM 借助保留基线生成先验并在逆向过程中通过扭曲粒子重采样调节分布，完美兼顾了高奖励与生成多样性。

局限与展望¶

复杂奖励的表征容量瓶颈：作者指出当前轻量级扭曲头（MLP 或小型 Transformer）在面对极度复杂、高维多模态或强组合特性的奖励表面时，标量近似精度可能面临容量上限。
极度稀疏奖励下的冷启动难题：当初始基准模型在目标奖励空间中几乎无法采样出有效正样本时，正样本缓冲池将面临严重的冷启动方差。未来工作需结合主动探索机制或离线经验回放。
离散/类别奖励的泛化扩展：目前实验主要验证了连续标量偏好与生物适应度指标，如何原生拓展至离散分类判别奖励或多约束硬性组合优化规则仍有待探索。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次系统性推导了离散扩散模型下 Twisted SMC 的对比分布匹配目标，并从理论上证明并利用了前向核闭式加噪分解。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖离散文本、基因调控 DNA、蛋白质多模态共设计、以及 8B 扩散大模型对齐四大前沿任务，包含详尽的墙钟时间对齐与真实泛化评测。
写作质量: ⭐⭐⭐⭐⭐ 数学推导严密自洽，附录算法与理论证明完备，图表组织逻辑清晰。
价值: ⭐⭐⭐⭐⭐ 彻底解决了离散扩散模型应用 Twisted SMC 时被昂贵奖励卡住的致命难题，为扩散大模型对齐与 AI4Science 提供了极具实用价值的推理加速工具。