技术摘要：用于扩散模型推理时分布控制的平均场框架
问题陈述
扩散模型正越来越多地被用作可控采样器，即在推理阶段通过引导生成过程来满足特定目标。虽然现有方法能有效处理逐点奖励 （定义在单个样本上，r : R d → R r: \mathbb{R}^d \to \mathbb{R} r : R d → R 分布奖励 （定义在生成的测度上，R : P ( R d ) → R R: \mathcal{P}(\mathbb{R}^d) \to \mathbb{R} R : P ( R d ) → R
目前的分布引导方法通常将奖励的梯度项纳入推理动力学中（例如，排斥势或集成似然）。然而，与早期的逐点引导方法类似，这些启发式方法缺乏关于采样分布的理论保证。具体而言，仅仅添加梯度项并不能确保粒子追踪到期望的倾斜测度 μ ∗ \mu^* μ ∗ arg max μ { R ( μ ) − KL ( μ ∥ p 1 ) } \arg \max_{\mu} \left\{ R(\mu) - \text{KL}(\mu \| p_1) \right\} arg μ max { R ( μ ) − KL ( μ ∥ p 1 ) } μ ∗ ∝ p 1 e r \mu^* \propto p_1 e^r μ ∗ ∝ p 1 e r 隐式 目标，其倾斜势依赖于测度本身。目前尚不存在能够纠正动力学和权重以追踪该隐式分布的有理论依据的框架。
方法论
作者提出了一个平均场框架 来形式化并解决推理时的分布控制问题。
1. 理论公式化
该问题被形式化为在每个时间步 t ∈ [ 0 , 1 ] t \in [0,1] t ∈ [ 0 , 1 ] μ t ∗ \mu^*_t μ t ∗ μ t ∗ ( d x ) = 1 Z t e Ψ t ( x , μ t ∗ ) p t ( x ) d x \mu^*_t(dx) = \frac{1}{Z_t} e^{\Psi_t(x, \mu^*_t)} p_t(x) dx μ t ∗ ( d x ) = Z t 1 e Ψ t ( x , μ t ∗ ) p t ( x ) d x Ψ t ( x , μ ) = δ R t δ μ ( x , μ ) \Psi_t(x, \mu) = \frac{\delta R_t}{\delta \mu}(x, \mu) Ψ t ( x , μ ) = δ μ δ R t ( x , μ )
2. 加权 McKean-Vlasov 动力学
为了追踪这一路径，作者推导了一个加权 McKean-Vlasov 系统 。与标准的梯度引导不同，引导项 ∇ x Ψ t ( X t ; μ t ) \nabla_x \Psi_t(X_t; \mu_t) ∇ x Ψ t ( X t ; μ t ) μ t \mu_t μ t
粒子 X t X_t X t A t A_t A t d X t = [ b t ( X t ) + ϵ t s t ( X t ) + ϵ t ∇ x Ψ t ( X t ; μ t ) ] d t + 2 ϵ t d B t dX_t = \left[ b_t(X_t) + \epsilon_t s_t(X_t) + \epsilon_t \nabla_x \Psi_t(X_t; \mu_t) \right] dt + \sqrt{2\epsilon_t} dB_t d X t = [ b t ( X t ) + ϵ t s t ( X t ) + ϵ t ∇ x Ψ t ( X t ; μ t ) ] d t + 2 ϵ t d B t d A t = [ b t ( X t ) ⋅ ∇ x Ψ t ( X t ; μ t ) + Ψ ˙ t ( X t ) ] d t dA_t = \left[ b_t(X_t) \cdot \nabla_x \Psi_t(X_t; \mu_t) + \dot{\Psi}_t(X_t) \right] dt d A t = [ b t ( X t ) ⋅ ∇ x Ψ t ( X t ; μ t ) + Ψ ˙ t ( X t ) ] d t Ψ ˙ t ( x ) = d d t Ψ t ( x , μ t ) \dot{\Psi}_t(x) = \frac{d}{dt} \Psi_t(x, \mu_t) Ψ ˙ t ( x ) = d t d Ψ t ( x , μ t ) Ψ ˙ t \dot{\Psi}_t Ψ ˙ t ∂ t r \partial_t r ∂ t r μ t \mu_t μ t Ψ ˙ t \dot{\Psi}_t Ψ ˙ t
3. 有限粒子算法
为了实现这一点，作者提出了一个加权相互作用粒子方案 （算法 1）：
位置更新： 使用基础漂移、分数（score）以及相对于当前经验测度 μ ^ t \hat{\mu}_t μ ^ t 对数权重更新： 更新权重以纠正引导动力学与目标倾斜路径之间的差异。这需要计算 Ψ ˙ t \dot{\Psi}_t Ψ ˙ t 求解 Ψ ˙ t \dot{\Psi}_t Ψ ˙ t 论文提供了两种计算此隐式时间导数项的方法：
隐式线性系统： 通过求解由积分方程导出的线性系统（命题 3.3）。不动点迭代： 一种轻量级的 Picard 迭代方案，用于近似权重更新。
重采样： 采用标准的序列蒙特卡洛（SMC）重采样以防止权重退化。
该框架将标准的逐点奖励引导作为一种特例恢复。
主要贡献
分布引导的形式化： 作者将分布级引导形式化为追踪一个倾斜测度，其中倾斜依赖于测度本身，建立了分布层面的目标与逐点奖励倾斜目标的类比。推导加权 McKean-Vlasov 动力学： 他们推导出了一个具有纠正性 Feynman-Kac 重加权的原则性动力系统，确保加权相互作用粒子能够追踪所需的倾斜路径 μ t ∗ \mu^*_t μ t ∗ 理论恢复与基础： 该框架恢复了标准的逐点引导作为其特例，并为现有的启发式批次级引导方法（如多样性采样、实验引导）提供了理论基础。收敛保证： 论文证明了当粒子数 N → ∞ N \to \infty N → ∞
实验结果
作者在低维和高维设置下验证了该框架：
低维（1D 混合分布）：
任务： 使用平方最大均差（MMD）奖励，将双峰高斯混合分布向具有不同模态权重的目标分布进行倾斜。发现： 标准梯度引导会将粒子推向目标，但无法精确匹配倾斜分布 μ ∗ \mu^* μ ∗ μ ∗ \mu^* μ ∗ λ \lambda λ J ( μ ) − J ( μ ∗ ) J(\mu) - J(\mu^*) J ( μ ) − J ( μ ∗ )
高维（蛋白质构象）：
任务 1 (HIV-1 Protease)： 引导 Boltz-2 生成的系综以匹配实验观测到的残基距离分布（DEER 数据）。基础模型偏向于“关闭”状态；目标是恢复“半开放”状态的比例。
结果： 平均场引导准确匹配了经验残基距离分布，而仅使用梯度的引导产生了形状略有不同且峰值更强的分布。
任务 2 (Adenylate Kinase)： 向源自分子动力学模拟的二维键角分布进行引导。
结果： 与仅梯度方法相比，平均场引导产生了更强的向目标分布偏移的效果。
任务 3 (X-ray 晶体学)： 引导 Protenix 生成结果以拟合表现出异构构象的电子密度图（4OLE 结构）。
结果： 在相同的引导强度下，平均场引导实现了更高的与目标密度的余弦对齐，并且相比于仅梯度引导，它在替代构象（构象 B）中放置了更多的样本，这表明仅梯度方法相对于真实的逆问题目标存在引导不足的问题。
意义与主张
本文声称通过为后者提供原则性的理论基础 ，弥合了逐点控制与分布控制之间的鸿沟。
正确性： 主要主张是，所提出的程序在无限粒子极限下精确追踪 预设的奖励倾斜分布 μ ∗ \mu^* μ ∗ 可解释性： 该框架阐明了现有的批次级引导方法实际上是在执行梯度步骤，但缺乏实现采样来自真实后验/倾斜测度所需的纠正性重加权。实用性： 虽然该方法引入了额外的计算成本（求解 Ψ ˙ t \dot{\Psi}_t Ψ ˙ t 局限性： 作者承认，依赖重采样对于某些模态（如图像）可能是一个限制（可能导致重复样本），并且该方法要求奖励函数是可微的。他们建议结合方差缩减技术（如 Poisson 方程求解器）可以缓解重采样问题。
这项工作并非旨在取代所有现有方法，而是提供了一个严谨的理论基础，解释了为什么某些启发式方法有效，并提供了一种确保其目标指向正确分布的机制。