Learning Energy-Based Models from Stochastic Interpolants using Spatiotemporal Differences¶

会议: NeurIPS 2026
arXiv: 2605.26850
领域: 学习理论
关键词: 能量模型, 随机插值, 噪声对比估计, 时空差异, 扩散生成模型

一句话总结¶

针对能量模型（EBM）中纯时间差异易受支撑集不匹配破坏、纯空间差异在多模态下严重失效的瓶颈，提出时空噪声对比估计（stNCE）统一二分类框架，通过联合建模数据空间与时间的流动差异，实现了高效且具备非参数一致性的单步似然评估。

研究背景与动机¶

学习未知数据分布的底层无标度概率密度并构建能量模型（EBM），在分子动力学玻尔兹曼采样纠偏、组合生成模型（如密度乘积与退火微调）、以及大模型偏好对齐（如对奖励倾斜分布的加权采样）等领域中扮演着核心角色。然而，传统能量模型由于配分函数不可积，推理成本通常极高；近期的研究趋势致力于构建支持单步评估（Single NFE）的似然估计模型。当前占主导地位的技术方案，例如去噪得分匹配（DSM）和时间得分匹配（TSM），大多基于随机插值（Stochastic Interpolants）构建加噪轨迹，将干净数据逐渐混入参考高斯分布，从而在连续时间变量上诱导出一个数据与噪声联合分布。

然而，现有方法在模型训练的差异分解机制上存在严重的理论与实践两极分化。一类方法依赖于空间差异（如条件噪声对比估计 CNCE、得分匹配及能量扩散模型），另一类方法依赖于时间差异（如标准 NCE、时间得分匹配 TSM、伸缩密度比估计 TRE）。过去研究普遍直觉地认为，只要引入多尺度噪声加噪（Amortization over multiple times），就能缓和单一层级的估计难度。但实质上，这一策略并未消除其本质缺陷：纯空间差异在多模态分布中计算模式权重时，积分轨迹必须横跨几乎无样本落入的低密度“真空区域”，导致空间梯度的积分产生严重漂移；而纯时间差异在数据分布与参考高斯分布的支撑集缺乏重叠（Support Mismatch）时，同样会在低密度过渡带发生评估崩溃。

为了打破空间与时间割裂导致的病态估计，本文从微积分路径积分视角切入：只要在数据空间与时间维度上同时离散并沿高密度流形转移，就能始终在概率质量集中的“安全走廊”内估计密度差。本文的核心 idea 是将能量模型的训练统一构建为数据空间与时间联合扰动的时空噪声对比估计（stNCE）二分类任务，并引入沿概率流形扩散-逆向传播的正反向扰动核（forward-reverse kernel），彻底规避多模态与支撑集错配的灾难性误差。

方法详解¶

整体框架¶

stNCE 的核心目标是学习含时能量函数 $E_\theta(x, t)$ 以及对数配分函数曲线 $\log Z(t)$，使得联合概率密度建模为 $p_\theta(x, t) = p(t) \exp(-E_\theta(x, t) - \log Z_t)$。输入为由随机插值轨迹 $x_t = (1-t)x_0 + t x_1$ 生成的数据样本及对应时间步（$t=0$ 为标准高斯基准分布，$t=1$ 为真实数据），输出为模型对任意单步点所评估的未归一化负能量与规整化因子。整个学习管线构造为一个扩展空间 $(x, t)$ 上的二分类判别问题：给定来自联合真实分布的锚点，利用预先设计的扰动核采样对比点，通过优化逻辑斯蒂损失迫使分类器输出恰好拟合真实时空对数密度差。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["随机插值联合采样<br/>(x, t) ~ p_d(x, t)"] --> B["时空扰动核转移<br/>(x', t') ~ p_n(x', t' | x, t)"]
    B --> C["联合分类对数几率计算<br/>F_theta(x, t, x', t')"]
    C --> D["时空逻辑斯蒂损失极小化<br/>min L_stNCE(F_theta)"]
    D --> E["能量模型单步推断<br/>E_theta(x, 1) & log Z_1"]

关键设计¶

1. 联合时空二分类框架与非参数一致性保证

传统对比估计仅在单一变量上设置负类扰动，stNCE 将样本增广至时空联合状态对 $((x, t), (x', t'))$。数据类 $A$ 的联合分布定义为 $p_A(x, t, x', t') = p_d(x, t) p_n(x', t' \mid x, t)$，而负类 $B$ 则将参考点与扰动点对调，即 $p_B(x, t, x', t') = p_d(x', t') p_n(x, t \mid x', t')$。二分类器的判别对数几率为： $$F_\theta(x, t, x', t') = \log p_\theta(x \mid t) - \log p_\theta(x' \mid t') + \log p(t) - \log p(t') + \log p_n(x', t' \mid x, t) - \log p_n(x, t \mid x', t')$$ 通过对二元交叉熵损失做泛函变分，分类器极小值严格对应于真值对数密度之差。理论证明表明，只要基准分布在 $t=0$ 处满足边界针定 $p_\theta(x \mid 0) = p_0(x)$，该极小化过程即可在非参数意义下保证 $p_{\hat{\theta}} \to p_d$ 的点态强一致性；更重要的是，其渐近协方差矩阵满足以样本量 $N_d$ 线性衰减的方差收敛界。

2. 保持数据流形约束的正反向扰动核（Forward-Reverse Kernel）

如果采用朴素的各向同性高斯白噪声扰动空间坐标（如 stNCE-w），在真实数据处于高维低维流形的场景下，扰动后的样本点会瞬间跌出数据支撑集，落入极低密度区域，导致反向梯度极不稳定。为此，本文设计了沿概率流扩散轨迹演化的正反向核：先依时间建议转移 $t \to t'$，随后根据 $t'$ 与 $t$ 的相对大小分别施加扩散前向加噪或基于指数积分器的逆向去噪： $$p_n(x' \mid x, t, t') = \begin{cases} \mathcal{N}\left(x' \;\middle|\; \frac{x}{t}, \, \left(1 - \frac{(1-t)^2}{t^2}\right) I\right), & t' < t \\ \mathcal{N}\left(x' \;\middle|\; \frac{t'}{t}x + \frac{2(1-t)(t'-t)}{t}\nabla_x \log p_t(x), \, \frac{(t'-t)(t+t'-2tt')}{t^2}I\right), & t' > t \end{cases}$$ 在逆向去噪分支中，漂移项依赖于空间得分函数 $\nabla_x \log p_t(x)$。本文提出了自引导估计（stNCE-s），即直接复用模型自身的空间梯度预测值作为得分代理。由于满足贝叶斯时空可逆恒等式 $p_d(x, t) p_n(x', t' \mid x, t) = p_d(x', t') p_n(x, t \mid x', t')$，对数几率公式中的先验漂移项得以完全抵消，迫使采样子步始终贴紧概率密度流形的中心骨架。

3. 对偶得分匹配与联合训练加速拓展（stNCE-s++DSM）

在理论连续极限下，当时间增量 $\Delta t \to 0$ 时，正反向 SDE 扰动下的 stNCE 损失收敛为最大似然加权的空间得分匹配损失；而当扰动核退化为各向同性微元高斯时，stNCE 则严格恢复为联合拟合空间得分 $\nabla_x \log p_d$ 与时间得分 $\partial_t \log p_d$ 的对偶得分匹配（Dual SM）。为了兼顾宏观时空跃迁与局部微分精度，算法构建了混合变体 stNCE-s++DSM：利用时空差分捕捉跨模态的相对比重，同时叠加瞬时去噪得分匹配（DSM）损失以锐化局部的几何曲率，且通过参数重用机制（Reuse Scheme）让两次前向计算共享特征，避免了二阶自动微分的繁重开销。

损失函数 / 训练策略¶

训练时从均匀分布采样时间标量 $t \in [0, 1]$，并通过正反射机制（Folding Mechanism）保证生成的时间对 $(t, t')$ 均严格服从均匀边际分布。损失采用对称化单项形式： $$\mathcal{L}_{\mathrm{stNCE}}(\theta) = -2\,\mathbb{E}_{p_d(x, t)p_n(x', t' \mid x, t)}\left[\log \sigma\left(F_\theta(x, t, x', t')\right)\right]$$ 在与 DSM 联合优化时，总目标函数设为 $\mathcal{L}_{\text{total}} = \mathcal{L}_{\mathrm{stNCE}\text{-s}} + \mathcal{L}_{\mathrm{DSM}}$。为了消除时间端点 $t=1$ 处的网络输出发散奇点，设计了基于标准差调制的 EDM 预条件变换，将时变阻尼系数因子从 $1-t$ 替换为 $1-0.75t$ 截断平滑，利用 AdamW 优化器进行全量参数更新。

实验关键数据¶

主实验¶

在 784 维 MNIST 高斯混合数据集（100 个高维稀疏模式，强多模态且与高斯基准无重叠）、MNIST 图像无条件密度估计以及 ImageNet $64\times64$ 生成建模任务中，系统比较了纯时间差异（tNCE、CTSM）、纯空间差异（tCNCE、DSM）以及各时空联合变体的性能。

数据集 / 设定	指标	本文 (stNCE-s)	联合增强 (stNCE-s++DSM)	基线最优 (Dual SM / Glow)	传统空间方法 (DSM)	传统时间方法 (tNCE)
MNIST GMM (784维)	密度均方误差 (MSE) $\downarrow$	121.6 $\pm$ 11.77	—	42.16 $\pm$ 7.11 (Dual SM)	$1\times 10^7$	$2\times 10^5$
MNIST GMM (784维)	空间差异比率误差 (Ratio) $\downarrow$	28.31 $\pm$ 1.93	—	17.82 $\pm$ 1.02 (Dual SM)	$1\times 10^4$	958.11
MNIST 图像测试集	负对数似然 (NLL bpd) $\downarrow$	1.61 (单步NFE)	1.00 (单步NFE)	1.05 (Glow, 归一化流)	—	7.01
ImageNet 64 测试集	负对数似然 (NLL bpd) $\downarrow$	—	2.94 (单步NFE)	3.36 (Dual SM) / 3.31 (FM)	—	—

注：在 MNIST GMM 上，stNCE-s 的单步训练耗时仅为 0.36s，相较需要二阶高阶导数的 Dual SM（0.70s）提速近一倍；在 ImageNet 64 任务中，stNCE-s++DSM 的 2.94 bpd 显著领先于现有连续能量扩散模型与主流 Flow Matching（3.31 bpd）。

消融实验¶

针对扰动核类型（混合核 stNCE-m、白噪声核 stNCE-w、自引导正反核 stNCE-s）、时间得分引导质量（真实得分 stNCE-o vs 自测得分 stNCE-s）以及多尺度时空训练步长进行消融评估。

配置	MNIST GMM Ratio $\downarrow$	单步运行时间 (s) $\downarrow$	说明
stNCE-s (正反向核+自引导)	28.31	0.36	流形内采样，权衡精度与计算开销的最佳组合
stNCE-o (正反向核+Oracle得分)	40.25	0.32	采用解析精确得分作为逆向引导
stNCE-w (高斯白噪声扰动核)	412.73	0.40	有限非零高斯噪声使样本脱离流形，误差骤增
stNCE-m (时间/空间交替混合核)	600.77	0.45*	朴素相加 tNCE 与 tCNCE，未打通联合时空流动通道
RNE-s (自引导重加权正则化)	3000+	0.44	缺乏非参数一致性，纯自回归引导出现严重退化

关键发现¶

流形保持是时空采样的决定性边界：有限扰动下，stNCE-w 的误差比 stNCE-s 高出逾一个数量级，这证明在高维空间盲目添加各向同性高斯球噪声会触发“维度灾难”，迫使判别器在极其稀疏的外围空间拟合伪能量；唯有依赖概率流驱动的正反向核能稳居高密度支撑带。
纯空间与纯时间维度的缺陷互补：纯时间方法 tNCE 因支撑集缺乏交集，MSE 高达 $2\times 10^5$；纯空间方法 DSM 因孤立模态间屏障而发散至 $1\times 10^7$。stNCE 将二者编织进同一轨迹，将 MSE 直降至百数量级。
计算开销断崖式下降：在分子系统（Chignolin）折叠模拟实验中，stNCE-s++DSM 达到与偏微分物理方程约束法（FPE）相近的 PMF（0.059 vs 0.037），但训练耗时从 49.57 GPU 小时骤减至 18.85 GPU 小时，提速达 2.6 倍。

亮点与洞察¶

将多类能量模型训练归一为经典二分类问题：本文以极其简练的数学视角，证明了标准 NCE、条件 NCE、能量扩散模型、时间得分匹配以及对偶得分匹配，均只是 stNCE 在特定扰动核与时间先验下的极端退化特例。
避免高阶自动微分的单步似然评估：传统规范化能量模型常需要借助散度求解（Hutchinson 迹估计）或二阶导数正则，计算极其笨重；stNCE-s 完全由一阶前向分类对数差驱动，兼顾极高运算吞吐与单步快速求值。
可迁移的流形正反向采样范式：通过扩散正向过程与基于指数积分器的逆向恢复核配对，为后续各类基于得分模型的对比学习提供了普适的流形保真采样模板。

局限与展望¶

作者承认的局限：纯 NCE 分类损失所学习的能量曲面在能量幅值上极其精确，但在未引入 DSM 联合约束时，其直接求导得到的空间梯度向量（得分函数）光滑度不如专业得分扩散模型，导致依赖梯度郎之万采样的图像生成质量受限。
实验与机制局限：自引导核（stNCE-s）在训练初始阶段模型梯度极不准确，此时逆向积分步引入的有偏伪样本可能在初期的数个 epoch 内对分类器形成短时扰动，虽然最终能够收敛，但延缓了早期的训练稳定性。
改进思路：后续可探索将无导数自适应马尔可夫采样器（如近端采样或哈密顿蒙特卡洛调整）融入 stNCE 评估环路，进一步解放对空间得分导数的依赖。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 建立了时空联合差异的统一分类理论，彻底厘清了传统时间/空间差异的失效机理。
实验充分度: ⭐⭐⭐⭐⭐ 覆盖一维概念模拟、784维混合高斯、MNIST、ImageNet64以及高难度分子动力学力场。
写作质量: ⭐⭐⭐⭐⭐ 数学推导严密自洽，几何物理图像生动清晰，分类谱系组织极具洞察力。
价值: ⭐⭐⭐⭐⭐ 为单步单次前向（Single NFE）的高效能量模型估计树立了极具竞争力的全新范式。

数据集 / 设定	指标	本文 (stNCE-s)	联合增强 (stNCE-s++DSM)	基线最优 (Dual SM / Glow)	传统空间方法 (DSM)	传统时间方法 (tNCE)
MNIST GMM (784维)	密度均方误差 (MSE) \(\downarrow\)	121.6 \(\pm\) 11.77	—	42.16 \(\pm\) 7.11 (Dual SM)	\(1\times 10^7\)	\(2\times 10^5\)
MNIST GMM (784维)	空间差异比率误差 (Ratio) \(\downarrow\)	28.31 \(\pm\) 1.93	—	17.82 \(\pm\) 1.02 (Dual SM)	\(1\times 10^4\)	958.11
MNIST 图像测试集	负对数似然 (NLL bpd) \(\downarrow\)	1.61 (单步NFE)	1.00 (单步NFE)	1.05 (Glow, 归一化流)	—	7.01
ImageNet 64 测试集	负对数似然 (NLL bpd) \(\downarrow\)	—	2.94 (单步NFE)	3.36 (Dual SM) / 3.31 (FM)	—	—