技术摘要：用于学习基于能量的生成模型的扩散分类损失

1. 问题陈述

基于能量的模型（EBMs）为概率建模提供了一个灵活的框架，将数据密度定义为可学习能量函数的指数。与学习对数密度梯度（即得分）的基于得分的生成模型（例如扩散模型）不同，EBMs 显式地对能量景观进行建模。这种能力对于需要直接访问能量值的下游任务至关重要，例如构建玻尔兹曼生成器、组合采样（结合多个模型）以及估计统计物理中的自由能差异。

然而，训练 EBMs 仍然是一个重大挑战。由于涉及难以处理的归一化常数（需要昂贵的嵌套采样），直接的最大似然（ML）估计在计算上是不可行的。因此，研究人员通常依赖去噪得分匹配（DSM），该方法能够高效地学习得分函数（能量的梯度）。虽然基于得分的方法在生成方面有效，但它们存在“模式盲区”。由于得分仅依赖于局部梯度，它无法捕捉不相交的高密度区域之间的相对比例（混合权重）。因此，基于得分的目标无法可靠地恢复真实的能量景观或断开模式之间的正确权重，限制了其在需要精确密度估计的任务中的效用。

2. 方法论：扩散分类（DiffCLF）

为了克服得分匹配的局限性同时保持计算效率，作者提出了扩散分类（DiffCLF）。该方法将学习随时间变化的能量函数的问题重新表述为跨不同噪声水平的监督分类任务。

核心公式

该方法在一个通用的随机过程框架 $Y_t = X_t + \gamma(t)Z$ 下运行，其中 $X_t$ 是源过程，$Z$ 是高斯噪声。目标是利用参数化为 $p_\theta^t(y) = \exp(-U_\theta^t(y) + F_\theta^t)$ 的 EBM 来估计边际密度 $p_t$（直至归一化常数），其中 $U_\theta^t$ 是能量，$F_\theta^t$ 是可学习的对数归一化常数（偏置）。

DiffCLF 目标 $L_{clf}$ 通过采样 $N$ 个时间步 $t_1, \dots, t_N$ 来定义，并将这些时间的边际分布视为不同的类别。给定从特定时间 $t_i$ 采样的样本 $y$，模型被训练以分类是哪个时间步生成了它。损失函数是分类交叉熵：
$$L_{clf}(\theta; t_{1:N}) = -\frac{1}{N} \sum_{i=1}^N \mathbb{E}_{p_{t_i}} \left[ \log \frac{p_\theta^{t_i}(Y_{t_i})}{\sum_{j=1}^N p_\theta^{t_j}(Y_{t_i})} \right]$$
该公式有效地学习了不同噪声水平之间的密度比。给定数据 $y$ 属于类别 $i$ 的后验概率被建模为 EBM 上的 softmax：
$$p_\theta(c=i|y) = \frac{p_\theta^{t_i}(y)}{\sum_{j=1}^N p_\theta^{t_j}(y)}$$

联合优化

为了确保学习到的能量不仅与分类任务一致，而且在梯度上也是准确的，DiffCLF 与标准的去噪得分匹配（DSM）目标相结合。联合目标为：
$$L_{joint} = L_{DSM} + L_{clf}$$
作者证明，虽然仅靠 $L_{clf}$ 具有非唯一的极小值点（允许密度的任意缩放），但与 $L_{DSM}$ 的组合确保了能够恢复真实地面真值边际 $p_t$ 的唯一极小值点。

计算效率

DiffCLF 的计算成本与标准 DSM 相当。虽然 DSM 每个时间样本需要两次神经网络评估（用于梯度的前向和反向），但多类 DiffCLF 目标需要 $N$ 次评估。当结合使用时，每个时间步的总成本为 $N+1$ 次评估。在二分类情况（$N=2$）下，与 DSM 相比仅增加了一次额外评估，使得该方法具有高度的可扩展性。

3. 主要贡献

• 新颖目标：引入 DiffCLF，这是一种轻量级且灵活的目标，将跨噪声水平的对数密度估计重新表述为多类分类问题。
• 理论保证：证明了联合目标（$L_{DSM} + L_{clf}$）能够唯一地恢复真实分布，并且不像纯得分匹配方法那样存在模式盲区。作者还建立了估计量的一致性和渐近正态性。
• 与先前工作的联系：该论文建立了 DiffCLF 与现有方法之间的理论联系，表明在连续时间极限下，二分类损失收敛于时间得分匹配（tSM），并将该方法与噪声对比估计（NCE）和多态 Bennett 接受率（MBAR）联系起来。
• 实证验证：展示了该方法在各种随机过程（扩散模型和随机插值）及维度上的有效性。

4. 实验结果

作者通过在合成数据集和真实世界数据集上的广泛实验验证了 DiffCLF：

• 合成高斯混合模型：在高维高斯混合模型（高达 128 维）上，与 DSM 和条件时间得分匹配（CtSM）相比，DiffCLF 实现了显著更低的分类损失，并且在恢复混合权重方面具有更好的一致性。虽然所有方法都实现了可比的 Fisher 散度（得分准确性）和最大均值差异（样本质量），但只有 DiffCLF 正确捕捉了不相交模式的相对比例。
• 分子系统：在分子势能（Müller-Brown、丙氨酸二肽、Chignolin）上的实验表明，DiffCLF 在 $t=0$ 时学习了更准确的能量函数，且未降低扩散采样的质量。其性能与 Fokker-Planck 正则化方法相当，但训练时间显著更快（加速 1.5–2.5 倍）。
• 下游应用：
  • 模型组合：在"AND"和"OR"组合任务中，使用 DiffCLF 训练的模型成功保留了组合分布的正确比例，而基于 DSM 的模型由于模式盲区未能维持这些权重。
  • 玻尔兹曼生成器：当用于通过序贯蒙特卡洛（SMC）构建玻尔兹曼生成器时，与 DSM 和 CtSM 基线相比，DiffCLF 训练的能量产生了更准确的样本，并具有更低的 Wasserstein/KS 距离。
  • 自由能估计：在用于估计溶剂化自由能的热力学积分（TI）任务中，DiffCLF 提供了比基线方法更准确的估计。

5. 意义与主张

该论文声称，DiffCLF 为训练基于能量的模型这一长期存在的问题提供了一个简单、高效且理论健全的方案，无需嵌套采样的计算负担，也克服了模式盲度的准确性限制。

作者强调，他们的方法将基于得分的生成模型的效用扩展到了单纯的样本生成之外。通过实现准确的能量估计，DiffCLF 促进了科学计算和概率建模中的一系列下游应用，包括：

• 玻尔兹曼生成器：为复杂的玻尔兹曼分布创建无偏采样器。
• 组合生成：在不重新训练的情况下，结合多个生成模型以形成新的分布（乘积或混合）。
• 自由能估计：准确计算统计物理中的自由能差异。

该工作将 DiffCLF 定位为一种通用的训练原则，能够与现有框架（如 DSM）无缝集成，并广泛适用于不同的随机过程，包括用于离散扩散的连续时间马尔可夫链。作者总结道，虽然他们目前的实验规模有限（主要集中在合成和分子系统上），但该方法为未来的大规模应用（如图像建模和文本生成）提供了一个有前景的方向。