Contrastive Distribution Matching for Amortized Sequential Monte Carlo in Discrete Diffusion

TL;DR CDM 通过对比分布匹配学习 twist 函数摊销 SMC 推理，在离散扩散模型中实现高效奖励对齐采样，仅需不到5%额外计算开销，性能超越现有方法。

问题

离散扩散模型在文本、DNA 序列、蛋白质结构等结构化分类数据生成领域展现出巨大潜力，但如何从经 reward 函数倾斜的分布中高效采样仍是一个核心瓶颈。

现有方法局限

当前主流的 Twisted SMC 方法通过引入 twist 函数修正提议分布，理论上可从 reward-tilted 分布中渐近精确采样。然而，最优 twist 函数在离散状态空间下没有解析形式，必须借助 蒙特卡洛近似：在推理时为每个状态运行大量前向模拟来估计未来期望奖励。这使得每次采样都需要数十甚至上百次模型前向传播，计算开销是基础推理的数十倍，严重限制其在实时或大规模应用中的部署。其他替代方案如基于微调的提议分布，则容易陷入 模式坍塌，牺牲生成多样性。

问题难点与重要性

• 技术挑战：离散状态的组合爆炸导致 twist 估计的方差极大，所需 MC 样本量随序列长度指数增长；而高维空间中的近似误差会破坏 SMC 理论的渐近精确性保证。
• 业界关注：可控生成是当前 AI 落地的关键需求，例如生成无毒性文本、设计高可药性蛋白、对齐大语言模型偏好等。若每次采样都需大量计算，则无法用于交互式应用或高吞吐管线。因此，摊销推理成本、在几乎不增加延迟的前提下逼近最优 twist，成为亟待解决的问题。

类似在 LLM 对齐中，通过 reward model 指导解码时，如何避免每个 token 都重跑一次昂贵评估，是所有在线可控生成方法的共同追求。

核心洞察

• CDM 通过对比学习将推理阶段的 SMC 扭转函数计算代价摊销至训练过程中，从而避免了传统方法在推理时反复进行昂贵的蒙特卡罗估计。与此前直接微调生成模型或在线学习提案分布的方式不同，CDM 利用正负样本对比匹配直接拟合最优扭转函数，实现了推理效率与采样质量的解耦，推理开销降低至单次前传的 5% 以下。
• 利用离散扩散模型前向核的闭式解，CDM 重新推导了梯度估计器，使得扭转函数的训练无需通过反向传播穿越采样过程，大幅降低了训练计算量。这一设计不仅让学习到的扭转函数在推理时几乎零额外开销，而且保证了梯度估计的低方差，从而在多种复杂离散生成任务（如 DNA 序列设计、蛋白质设计）中稳定收敛，相比同类方法在相同墙钟时间内表现更优。

方法

核心思路：摊销式扭曲函数学习

CDM 将 Twisted SMC 中推理时的昂贵 twist 估计转化为离线训练任务，通过对比学习得到参数化扭曲函数 $\psi_\theta$，使其能够近似最优扭曲，从而在采样时以极低开销获得接近精确的 reward-tilted 分布。

输入与数据流

• 输入：离散扩散模型的前向转移核 $q(x_t|x_{t-1})$（具有闭式表达），以及目标奖励函数 $r(x)$。
• 训练数据：正样本从当前扭曲策略诱导的分布中采样（或使用经验回放缓冲），负样本从基础扩散分布 $p(x)$ 中采集，形成对比对。

关键模块：对比分布匹配

1. 扭曲函数参数化：$\psi_\theta(x_t, t)$ 输出一个标量，用于修正 SMC 提议分布，使其偏向高奖励区域。通常设计为轻量 head（如 MLP）附着在基础模型表示上，评估成本 <5% 单次前向传播。
2. 训练目标：最大化扭曲后分布与真实 reward-tilted 后验之间的相似度。CDM 采用对比损失，鼓励正样本对（高奖励轨迹）的扭曲值高于负样本，从而隐式匹配分布。
3. 高效梯度估计：利用离散扩散模型前向核的闭式形式，将原本需要 Monte Carlo 估计的梯度重新表达为可微运算，避免高方差采样，显著加速训练。
4. 负采样策略：通过维持一个负样本缓冲池，动态更新以稳定对比学习过程。

输出与推理

训练完成后，在 SMC 重采样和提议步骤中直接调用 $\psi_\theta$，无需任何 MC 模拟即可获得扭曲权重。这使得推理阶段的计算复杂度与普通 SMC 相当，但混合质量接近精确方法。

与同类方法的差异

不同于微调提议分布（易模式坍塌）或推理时 MC 估计 twist（计算量数倍于基础模型），CDM 首次在离散扩散中通过摊销对比学习实现高效且精确的 reward-aligned 采样，将计算从推理转移到训练，同时保持生成多样性与准确性。

实验

实验设计

CDM 在四个离散扩散生成任务上验证：有毒文本生成（控制毒性分数）、调控 DNA 序列设计（优化转录因子结合）、蛋白质可设计性（最大化结构合理性）和扩散 LLM 对齐（奖励优化）。基线包括微调提议分布（fine-tuned proposals）、标准 Twisted SMC、Best-of-N 等。所有比较在**匹配挂钟时间（matched wall-clock time）**下进行，确保计算效率公平。

关键发现

• CDM 在所有任务中一致优于基线，学到的扭曲函数额外计算开销低于基模型单次前向的 5%。
• 在蛋白质设计中，CDM 生成的序列可设计性分数显著更高；在 LLM 对齐中，产生更符合人类偏好的文本。
• CDM 有效避免微调提议的模式坍塌，生成样本多样性明显优于 fine-tuned proposals。

对比分析

CDM 与现有方法的本质区别在于摊销推断：通过对比分布匹配预训练扭曲函数，消除了 SMC 每步重估扭曲的高昂开销。相比标准 Twisted SMC，CDM 推理计算量大幅降低，实现可扩展部署；相较于微调提议，CDM 保留 SMC 的无偏性，不会因提议的过度优化导致分布简化。这表明对比学习驱动的摊销 SMC 是离散扩散模型中奖励对齐的一种高效、鲁棒的新范式。

行业影响

落地场景

CDM 为离散扩散模型提供了一种低成本的奖励偏向采样能力，适用于任何需要生成符合特定偏好或约束的结构化离散数据的场景。典型落地产品包括：

• 内容安全与可控文本生成：社交媒体、客服对话系统、营销文案生成中，确保输出无害、符合品牌调性。
• 药物发现与合成生物学：蛋白质序列设计、调控 DNA 序列优化，帮助生物技术公司快速筛选高活性候选分子。
• 大语言模型对齐：在扩散语言模型推理时注入安全、有用性奖励，抑制毒性或不当回复。
• 推荐与个性化：电商或内容平台中，生成符合用户偏好或高转化率的物品描述、标签组合等。

商业价值

CDM 的核心商业贡献在于 大幅降低推理计算开销。传统 Twisted SMC 需要昂贵的蒙特卡洛近似来估计最优 twist 函数，而 CDM 通过离线训练一个参数化 twist 头，将推理开销压缩至基模型单次前向的 5% 以内，近乎无额外成本。这直接带来：

• 降本：减少 GPU 资源消耗，支持更大规模实时服务。
• 增收：在高价值场景（如药物开发）中加快迭代速度，缩短研发周期。
• 体验提升：生成结果更贴合目标需求，减少后处理或人工审核成本。

与现有工作流的接口

CDM 的集成极为轻量：

• 模型侧：仅在基础扩散模型上附加一个 twist 预测头，可用轻量网络实现。无需修改原模型参数，可与已有 checkpoints 兼容。
• 推理侧：将训练好的 twist 函数作为 SMC 的 proposal 调节项，无缝嵌入现有采样循环。
• 训练数据：只需收集正负样本对（高/低奖励样本），支持在线或离线训练，适配已有 RLHF 或偏好数据集。

具体落地 Use Case

1. 电商平台评论/描述生成：利用 CDM 对齐扩散文本模型，使生成的商品描述更吸引点击，同时避免夸大或违规内容。例如，某跨境电商平台需要为千万级 SKU 生成多语言描述，CDM 可在极小推理成本下让输出同时满足高转化率和政策合规。

2. 制药企业蛋白质设计管线：使用基于离散扩散的蛋白序列生成模型，CDM 引入可设计性、稳定性等奖励，直接筛选高潜力抗体或酶变体。学习到的 twist 函数可一次性部署，在每次采样中持续引导，加速湿实验前的计算机筛选，降低湿实验成本。

局限

• **扭曲函数对奖励变化的适应性受限** CDM 通过对比学习从**特定奖励函数**的正负样本中学习扭曲函数。当奖励函数频繁变化或需快速适配新奖励时，必须重新收集样本并训练，导致**训练成本重复**。这与纯推理时优化（如在线 Monte Carlo 估计）形成对比，后者虽推理开销大但无需前置训练，在奖励函数动态变化的环境中可能更灵活。此问题在论文中未被明确讨论，但限制了该方法在需要持续适配不同奖励信号的应用中的实用价值。
• **依赖离散扩散模型的闭式前向核** CDM 训练梯度估计利用了**离散状态空间扩散模型**的前向核闭式形式，从而避免昂贵的 Monte Carlo 梯度估计。然而，这一优势也将其适用范围**局限在具有此类闭式解的模型族**（如使用均匀转移或离散匹配的扩散模型）。对于连续状态扩散模型、自回归模型或其他生成框架，若不重新推导梯度估计，该方法无法直接迁移，限制了其作为通用 SMC 加速方案的可能性。
• **扭曲函数泛化能力未充分验证** 论文在多个任务（毒性文本生成、DNA 序列设计、蛋白质设计等）上展示了有效性，但所有任务均使用**训练时见过的奖励函数**进行评估。对于**分布外奖励函数**或**未见过的核结构**，学习到的扭曲函数的泛化性能缺乏系统研究。在实际部署中，若测试奖励与训练分布有偏移，扭曲函数可能失效，导致 SMC 采样效率退化回原始水平。