具有原始-对偶推断的约束扩散模型

# 约束扩散模型与原对偶推理

来源: https://arxiv.org/html/2606.17192
Samar Hadou Yiğit Berkay Uslu Alejandro Ribeiro
电气与系统工程系, 宾夕法尼亚大学
\{selaraby, ybuslu, aribeiro\}@seas.upenn.edu

###### 摘要

本文开发了带有原对偶推理 (PDI) 的约束扩散模型，用于从具有平均约束的熵正则化优化问题的最优分布中进行采样。我们将约束采样形式化于拉格朗日对偶域，其中最优分布呈现为以最优对偶变量为索引的吉布斯分布。与在采样前估计此对偶乘子并在生成过程中冻结它不同，PDI 联合推断最优原分布及其参数化对偶变量。每个反向扩散步骤使用与当前乘子关联的得分场进行去噪，然后利用去噪样本的估计约束违反量通过对偶上升来更新乘子。为了实现这种条件得分场，我们在推理过程中遇到的对偶变量所诱导的吉布斯分布族上训练一个单一的对偶条件得分网络。我们证明了沿推理轨迹生成的对偶变量的时间平均收敛到对偶最优值的一个邻域，并通过与时间安排相关的稳定性因子来界定残余对偶失配对终端分布的影响。我们在从高斯混合、无线资源分配和投资组合管理中进行的约束采样上评估了 PDI。

## 1 引言

扩散模型通过学会逆转加噪过程，已成为从复杂高维分布中采样的主导框架 (Sohl-Dickstein et al., 2015; Ho et al., 2020; Song et al., 2021b)。虽然其最显著的成就在于无约束的感知任务，如图像 (Rombach et al., 2022)、音频 (Kong et al., 2021) 和视频合成 (Ho et al., 2022)，但相同的分布建模能力对于解决一大类约束优化问题也很有价值，这些问题的解是概率性的而非确定性的。例如，无线资源分配采用随机的时间共享策略，在满足每个用户服务保证的前提下最大化网络效用 (Uslu et al., 2025a; Babazadeh Darabi and Coleri, 2025)。类似地，投资组合管理旨在分散投资，在预期风险约束下最大化收益 (Tiwari, 2026; He et al., 2025a)。然而，标准的扩散模型缺乏在采样过程中强制执行此类统计约束的可靠机制。本文开发了具有原对偶推理的扩散模型来弥合这一差距。

对具有平均约束的分布进行优化的熵正则化形式，在拉格朗日对偶域中具有简洁的形式。对于固定的对偶变量，拉格朗日函数的分布极小化器是一个吉布斯分布，其能量关于对偶变量是线性的。这将对偶采样与大量基于扩散的、从非归一化吉布斯目标中采样的文献联系起来 (Zhang and Chen, 2022; Vargas et al., 2023; Berner et al., 2024; Richter and Berner, 2024)。原则上，可以求解出最优对偶变量，然后训练一个扩散模型从相应的吉布斯分布中采样。然而，在我们的设定中，采样器必须处理一系列约束优化实例，每个实例都可能产生不同的最优乘子。在采样前估计这些乘子是脆弱且昂贵的。最优乘子是通过吉布斯分布下的难以处理的期望来定义的，因此每次对偶更新都需要估计正在学习的分布的统计量。之前的对偶训练 (DT) 方法 (Khalafi et al., 2024, 2025) 通过在训练中交替进行对偶更新和得分模型训练来解决这个问题。由于每次对偶迭代中得分更新次数有限，学习到的得分模型可能反映这些目标的累积历史，而不是由最终乘子索引的吉布斯得分场。一旦训练结束，生成的采样器就固定了，无法在生成过程中对约束违反做出响应。

本文提出了不同的观点。我们不将对偶变量视为采样前必须估计的量，而是将其作为采样状态的一部分。我们提出了原对偶推理 (PDI)，它将对偶上升与反向扩散过程交错进行。在每个去噪步骤，PDI 使用与当前乘子关联的得分场来更新样本，基于 Tweedie 估计的去噪样本估计约束违反量 (Efron, 2011)，并在下一步反向步骤之前更新乘子。对偶变量成为采样器推理时的一个状态，将轨迹引向可行性，并且采样器遵循一个路径依赖的吉布斯得分场序列。这种动态耦合在整个推理过程中使得分场与当前乘子保持一致，允许早期的噪声步骤吸收对偶失配，而后续步骤在乘子已接近最优乘子区域时细化样本。我们的分析通过证明时间平均对偶变量的收敛性以及界定残余对偶失配如何通过反向过程传播来捕捉这种效应。

为了实现这种原对偶耦合，我们训练一个以对偶变量为条件的单一得分网络，使得模型代表一个吉布斯目标族，而不是单个约束分布。这为 PDI 提供了一种直接机制来适应转移的约束和分布外 (OOD) 实例，只要诱导的对偶轨迹保持在训练期间覆盖的乘子范围内。实验表明，PDI 优于 DT，并且在转移约束下保持更强的鲁棒性。当通过冻结候选最终乘子并继续对该固定吉布斯目标进行得分训练来增强 DT 时，它变得与 PDI 有竞争力，但这需要额外的训练阶段。这表明 PDI 的优势不仅仅是估计一个好的乘子，而是在推理期间保持乘子更新与去噪轨迹的耦合。

我们做出以下贡献。

1. (C1) 我们将具有平均约束的约束采样形式化为拉格朗日对偶域中的鞍点问题，其中最优分布是以最优对偶变量为索引的吉布斯分布。
2. (C2) 我们引入了 PDI 算法，一种反向扩散采样器，其中原始去噪和对偶提升共同演化。在每一步，采样器使用与当前乘子关联的得分场，然后使用约束违反量更新乘子。
3. (C3) 我们训练一个单一的对偶变量条件得分网络，该网络在推理过程中遇到的对偶变量索引的吉布斯分布族上泛化。
4. (C4) 我们建立了时间平均对偶迭代收敛到最优乘子邻域。我们还通过时间安排相关的稳定性因子来界定残余对偶失配对终端分布的影响。
5. (C5) 我们在受约束的高斯混合、具有每个用户速率要求的无线功率控制问题以及具有预期风险约束的投资组合构建任务上实证验证了 PDI。

图 1 具体展示了推理时对偶更新的效果。我们展示了一个无线优化问题，其中用户共享一个信道，每个用户都需要一个最小平均服务速率。无约束采样器忽略需求，导致许多用户服务不足；DT 在采样前对对偶变量进行去噪，因此无法在生成过程中纠正违反；而 PDI 通过在去噪时更新对偶变量，将样本导向（平均）可行性，并将服务最差的用户提升到更接近其目标。

![图 1: 在遍历最小速率 r_min 约束下的无线功率分配。节点代表用户，颜色表示每个用户的速率需求满足程度，红色标记未达到要求的用户。这里，平均约束是有意义的，因为相互干扰阻止任何单一解同时满足所有要求。因此最优解必须随时间交替样本（即功率分配和/或信道接入），以便一些样本补偿其他样本。无约束采样器使许多用户低于目标速率（红节点），而 DT 通过训练时的对偶更新部分提高了可行性。PDI 通过在去噪期间更新对偶变量，进一步提高了低速率用户的性能。](caption)

## 2 相关工作

##### 带约束的扩散模型采样。引导技术通过修改 Tweedie 去噪公式中的得分 (Efron, 2011)，将反向扩散导向奖励或约束对齐，如分类器引导 (Dhariwal and Nichol, 2021)、无分类器引导 (Ho and Salimans, 2022) 和基于 CLIP 的变体 (Nichol et al., 2022)。另一条工作线通过边界反射 (Lou and Ermon, 2023)、镜像映射 (Liu et al., 2023)、黎曼形式 (De Bortoli et al., 2022)、对数障碍结构 (Fishman et al., 2023)、可行性投影 (Christopher et al., 2024) 和约束反向步骤 (Huang et al., 2024) 直接在扩散过程上强制执行域可行性约束。最关键的是，这些方法处理的是逐点约束，而不是分布（平均）约束。

##### 带拉格朗日对偶形式的约束扩散建模。与我们的工作密切相关的是 Khalafi et al. (2024, 2025)；Chamon et al. (2024)，它们采用原对偶形式处理平均约束。具体来说，Khalafi et al. (2024, 2025) 将外循环对偶上升与内循环得分匹配配对，但每个对偶迭代都重新训练得分网络，而原对偶 Langevin Monte Carlo (Chamon et al., 2024) 在采样时共同演化样本和乘子，但缺乏学习到的生成模型。我们的方法与这些工作共享拉格朗日鞍点结构，但在两个关键方面不同。首先，对偶变量与去噪轨迹共同演化，而不是在单独的训练循环或遍历 Langevin 链中。其次，单个以对偶变量为条件的得分网络是在乘子的分布上训练的，而不是为每次迭代重新拟合。扩展相关工作见附录 A (https://arxiv.org/html/2606.17192#A1)。

## 3 约束扩散模型

考虑一个目标分布 μ*，它在期望中权衡了变分目标函数的最小化与满足一组约束。更具体地说，μ* 是以下问题的最小化器

P* = min_{μ ∈ P_2}  E_μ[ f_0(x) ] - β H(μ),     s.t.   E_μ[ f(x) ] ⪯ 0,      (1)

其中 f_0: X → R, f: X → R^d, X ⊆ R^p 是可行决策域，P_2(X) 表示所有支撑在 X 上且具有有限二阶矩的概率分布空间。我们在目标中添加了一个熵正则化项，以惩罚退化解并确保样本多样性。这种形式在多玩家合作博弈和资源分配问题中尤其相关，其中单个确定性策略可能无法同时满足所有竞争性需求。在这样的设定中，从 μ* 中采样的时间共享策略可以在平均意义上满足约束，同时改善最优性和可行性之间的权衡。

问题 (3) 在 μ 上是 β 强凸的，具有零对偶间隙，可以通过构造拉格朗日函数在对偶域中处理：

L(μ, λ) = E_{x ~ μ}[ f_0(x) + λ^T f(x) ] - β H(μ),     (2)

其中 λ ⪰ 0 包含拉格朗日乘子。对偶问题定义为：

D* = max_{λ ⪰ 0} min_{μ ∈ P_2} E_{x ~ μ}[ f_0(x) + λ^T f(x) ] - β H(μ).     (3)

这两个问题通过以下命题等价。

###### 假设 1 (存在严格可行解)。存在 μ' ∈ P_2 使得 E_μ'[ f_0(x) ] - β H(μ') < ∞ 且 E_μ'[ f(x) ] ≺ 0。

**命题 1** (强对偶与最优解的特征)。在假设 1 下，对于 β > 0，我们有 P* = D*，并且最优对 (μ*, λ*) 是拉格朗日函数的鞍点，即对于任意 μ ∈ P_2(X) 和 λ ∈ R_+^d，有 L(μ*, λ) ≤ L(μ*, λ*) ≤ L(μ, λ*)。同时 λ* 是有限的，即 λ* ∈ Λ ⊂ R_+^d。此外，μ* 服从吉布斯分布律：

μ*(x) = μ_{λ*}(x) = (1 / Z(λ*)) ⋅ exp( - (1/β) (f_0(x) + λ*^T f(x)) ).