Equilibrium Matching: Generative Modeling with Implicit Energy-Based Models¶

会议: NeurIPS 2026
arXiv: 2510.02300
论文: https://arxiv.org/abs/2510.02300
代码: https://raywang4.github.io/equilibrium_matching/
领域: 图像生成
关键词: 平衡态匹配, 隐式能量模型, 优化采样, 流匹配, 图像生成

一句话总结¶

Equilibrium Matching (EqM) 摒弃了扩散与流匹配模型对时间/噪声条件的依赖及非平衡动力学，通过构造在真实数据流形处梯度消失的目标场来学习隐式能量景观的平衡态梯度，使推理生成转化为自适应步长的梯度下降优化，在 ImageNet 256×256 上取得 1.90 FID 并原生支持部分加噪去噪与无监督 OOD 检测。

研究背景与动机¶

近年来扩散模型与流匹配（Flow Matching, FM）在连续生成建模中取得了巨大成功，它们通过预定义前向加噪过程并训练神经网络学习沿时间反向去噪的速度场或得分函数。然而，这些主流方法在本质上都属于非平衡动力学（non-equilibrium dynamics）：模型必须显式接收连续时间步 $t$ 或噪声水平作为输入条件，在不同噪声尺度下学习完全不同的局部演化速度，并且推理采样被严格绑定在从 $t=1$ 到 $t=0$ 的固定时间积分区间与刚性步长网格上。一旦脱离预设的初始纯高斯噪声分布或调整时间积分路径，采样质量便会急剧劣化。

直觉上，构建一种不依赖时间条件、处处时间不变（time-invariant）的平衡态动力学不仅能摆脱繁琐的时间调度，还能赋予生成过程更灵活的推理时优化能力。然而，直接将扩散或流匹配模型的时间条件去除（如 Sun 等人在 2025 年的探索）会导致生成质量显著滑坡；而经典的能量模型（EBMs）虽然天生具备时间不变的平衡态能量景观，却长期受困于对比散度（Contrastive Divergence）等 MCMC 采样带来的训练极不稳定、难以扩展至高维复杂图像的问题。后续的两阶段混合方法（如 Energy Matching）又需要分段训练且在大规模数据上落后于纯流匹配基线。这种质量退化背后的本质症结在于：传统流匹配所拟合的瞬时速度场与平衡态能量景观之间存在根本性的数学不兼容，真实数据点并未被塑造成能量极小值。

针对这一核心矛盾，本文的切入角度是重新从隐式能量模型的视角审视生成速度场，设计与平衡态动力学严格自洽的目标梯度。核心 idea：通过构建在真实数据流形处梯度归零、向高噪声区域单调递增的目标梯度场，使模型在无需任何时间/噪声条件的前提下学习隐式能量景观的稳态梯度，将生成采样重构为基于凸优化技术的梯度下降寻优过程。

方法详解¶

整体框架¶

Equilibrium Matching (EqM) 旨在摆脱扩散与流匹配中的非平衡时间条件演化，建立一个以真实数据为局部极小值的平滑能量曲面。在训练阶段，算法在真实样本 $x$ 与高斯噪声 $\epsilon$ 之间采样隐式插值状态 $x_\gamma = \gamma x + (1-\gamma)\epsilon$（其中插值因子 $\gamma \sim \mathcal{U}(0,1)$ 对网络完全不可见），并驱使无时间条件的网络 $f(x_\gamma)$ 拟合一个从噪声指向数据的下降方向 $(\epsilon - x)$ 乘以随接近数据流形而衰减的标量系数 $c(\gamma)$。在推理阶段，生成过程完全转化为在所学能量景观上的梯度下降优化：从随机初始状态出发，利用预测梯度以及 Nesterov 加速梯度（NAG）等一阶优化算法迭代寻优，甚至可依据梯度范数动态提前终止，实现自适应计算分配。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["训练输入：真实图像 x 与高斯噪声 ε<br/>随机隐式插值 x_γ = γ x + (1-γ) ε"] --> B["目标梯度场构造<br/>方向 (ε - x) 乘截断衰减权重 c(γ)"]
    B --> C["平衡态网络 f(x_γ)<br/>无时间/噪声输入条件，直接拟合稳态梯度"]
    C --> D["显式能量函数建模（可选 EqM-E）<br/>点积标量 g(x_γ) 梯度的反向传播对齐"]
    D --> E["优化式采样与自适应计算<br/>从噪声出发执行 NAG-GD 寻优，梯度低于阈值自适应停机"]
    E --> F["输出：收敛至能量极小值的高保真生成图像"]

关键设计¶

1. 目标梯度场构造：与隐式能量函数对齐的零驻点约束

直接剔除流匹配中的时间条件之所以导致性能退化，是因为原版流匹配目标 $x - \epsilon$ 在数据端点处依然保持非零速度，强行拟合无法使数据点成为动力系统的稳定不动点。EqM 从能量模型的视角出发，将目标向量反向为由噪声指向数据的下降方向 $\epsilon - x$，并引入幅度衰减函数 $c(\gamma)$，将训练损失定义为： $$\mathcal{L}_{\text{EqM}} = \mathbb{E}_{x,\epsilon,\gamma} \left\| f(x_\gamma) - (\epsilon - x)c(\gamma) \right\|^2$$ 为了确保真实数据点成为能量曲面的驻点与局部极小值，函数 $c(\gamma)$ 必须严格满足流形零梯度边界条件 $c(1) = 0$。经过对线性衰减 $c_{\text{linear}}(\gamma)=1-\gamma$、分段衰减等多种形式的系统探索，论文发现截断衰减（Truncated Decay）表现最优：当远离数据流形（$\gamma \le a$）时维持恒定梯度以提供足够的收敛推力，而在临近数据流形（$\gamma > a$）时线性收缩至零： $$c_{\text{trunc}}(\gamma) = \begin{cases} 1, & \gamma \le a \\ \frac{1-\gamma}{1-a}, & \gamma > a \end{cases}$$ 在实验中选定截断拐点 $a=0.8$ 并配合整体缩放系数 $\lambda=4.0$，模型在理论与高维几何性质下保证了在真实样本集 $\mathcal{X}$ 上 $\|f(x^{(i)})\|_2 \approx 0$，且任意驻点 $\hat{x}$ 均以极高概率属于数据分布 $P(\hat{x} \in \mathcal{X}) \approx 1$。

2. 显式能量函数建模：从输出梯度反推标量能量

为了使 EqM 不仅能作为隐式梯度场生成样本，还能像传统 EBM 一样直接提供能量标量用于分布外检测（OOD Detection）和模型组合，作者进一步提出了显式能量匹配变体（EqM-E）。该变体直接指定标量函数 $g(x_\gamma)$，并通过自动微分将其输入梯度 $\nabla_{x_\gamma} g(x_\gamma)$ 与目标场对齐： $$\mathcal{L}_{\text{EqM-E}} = \mathbb{E}_{x,\epsilon,\gamma} \left\| \nabla g(x_\gamma) - (\epsilon - x)c(\gamma) \right\|^2$$ 在无需额外增加模型参数的前提下，论文评估了两种标量化形式：点积能量函数 $g(x_\gamma) = x_\gamma \cdot f(x_\gamma)$（对应输入梯度 $\nabla g(x_\gamma) = f(x_\gamma) + x_\gamma^T \nabla f(x_\gamma)$）与平方 $L_2$ 范数能量函数 $g(x_\gamma) = -\frac{1}{2}\|f(x_\gamma)\|_2^2$。实验证明，点积形式的参数曲面更加平滑易优化，不仅避免了二阶项带来的训练不稳定性，而且赋予了网络直接评估未归一化对数密度的能力。

3. 优化式采样与自适应计算：从固定积分迈向灵活梯度下降

传统扩散/流模型将采样框定为常微分方程（ODE）或随机微分方程（SDE）在固定区间的数值积分，步长与时间网格强耦合；而 EqM 学习的是时间不变的稳态保守场，因此推理过程被自然转化为无约束优化问题。最基本的更新为标准梯度下降（GD）：$x_{k+1} \leftarrow x_k - \eta f(x_k)$。当能量函数 $L$-光滑且有下界时，梯度下降在 $K$ 步内保证以 $\mathcal{O}(1/K)$ 的收敛率达到临界驻点。

在此基础上，经典的优化加速算法可被无缝迁移到采样管线中。引入 Nesterov 加速梯度（NAG-GD）： $$x_{k+1} \leftarrow x_k - \eta f\big(x_k + \mu(x_k - x_{k-1})\big)$$ 通过在前瞻点评估梯度，NAG 显著加速了跨越平坦区域与震荡峡谷的进程，尤其在少步数采样下（如 20-50 步）相比标准 GD 展现出巨大的保真度优势。此外，优化视角天然催生了自适应计算预算分配（Adaptive Compute）：通过监控梯度范数 $\|f(x_k)\|_2 < g_{\text{min}}$ 决定动态停机，模型能针对简单样本提前收敛，平均节省高达 60% 的函数评估次数（NFE）。

一个完整示例¶

以无条件生成一张 $256 \times 256$ 图像为例，EqM 的前向优化流程如下： 1. 初始化：从标准高斯先验采样起始潜在张量 $x_0 \sim \mathcal{N}(0, I)$，设定前瞻动量状态 $x_{-1} = x_0$，选定步长 $\eta = 0.0017$、动量因子 $\mu = 0.3$ 以及停机阈值 $g_{\text{min}} = 10$。 2. 前瞻与梯度评估：在第 $k$ 步，计算前瞻状态 $x_k^{\text{look}} = x_k + \mu (x_k - x_{k-1})$，送入主干网络获得稳态梯度向量 $g_k = f(x_k^{\text{look}})$。 3. 极小值寻优更新：记录历史状态并沿负梯度方向更新 $x_{k+1} = x_k - \eta g_k$。 4. 动态停机判定：计算梯度欧式范数 $\|g_k\|_2$。若其低于 $g_{\text{min}}$，表明当前状态已落入数据流形能量池的不动点区域，提前终止循环输出 $x_{k+1}$；否则持续迭代至最大步数（例如 250 步）。整个过程无需感知时间标签，如同在复杂能量曲面上执行高效的球体滚落寻优。

损失函数 / 训练策略¶

EqM 采用标准 MSE 损失优化网络权重，骨干网络采用基于 Transformer 的 SiT/DiT 架构（涵盖 S/2, B/2, L/2 至 XL/2）。为了完全剥离时间条件，模型直接将时间输入固定为 0，不再向自适应归一化层传入时间嵌入。优化器采用 AdamW，批大小为 256，恒定学习率 $1 \times 10^{-4}$，无学习率预热和权重衰减。梯度场超参数统一固定为截断衰减 $a=0.8$、整体放大因子 $\lambda=4.0$。

实验关键数据¶

主实验¶

在 ImageNet 256×256 类条件图像生成基准上，EqM-XL/2 在 1400 轮训练下取得了 1.90 的 FID，全面超越了代表性 GAN、扩散模型与流匹配（SiT）基线。

模型	方法类别	FID ↓	sFID ↓	IS ↑
StyleGAN-XL	GAN	2.30	4.02	265.1
VDM++	扩散模型 (Diffusion)	2.12	-	267.7
DiT-XL/2	扩散模型 (Diffusion)	2.27	4.60	278.2
SiT-XL/2	流匹配 (Flow Matching)	2.06	4.49	277.5
EqM-XL/2 (本文)	平衡态匹配 (EqM)	1.90	4.54	275.7

在不同采样算法的直接对比中（ImageNet 256×256, XL/2 架构），EqM 无论是使用 ODE 数值积分还是优化式梯度下降，均稳定领先于 SiT-XL/2：

模型	采样器类型	步长 $\eta$	NAG 动量 $\mu$	FID ↓
SiT-XL/2	Euler (ODE 积分)	0.0040	-	2.10
SiT-XL/2	Heun (SDE 积分)	0.0040	-	2.06
EqM-XL/2	Euler (ODE 积分)	0.0017	-	1.93
EqM-XL/2	GD (标准梯度下降)	0.0017	-	1.93
EqM-XL/2	NAG-GD (加速梯度下降)	0.0017	0.3	1.90

消融实验¶

在 EqM-B/2（训练 80 轮）上系统评估目标场函数 $c(\gamma)$ 形式与超参数选择对生成性能的影响：

目标函数形式 $c(\gamma)$	截断点 $a$	起始值 $b$	FID ↓	说明
constant (常数)	-	-	40.81	等价于无时间条件的负流匹配目标
linear (线性衰减)	-	-	50.47	全程线性下降至 0，初始推力不足
truncated (截断衰减)	0.5	-	38.98	前期恒定，后期衰减
truncated (截断衰减)	0.8	-	38.34	最优场形态（配乘 $\lambda=4$ 后降至 32.85）
truncated (截断衰减)	0.9	-	41.22	衰减区间过窄，边界震荡
piecewise (分段线性)	0.8	0.8	38.84	初始幅度低于 1
piecewise (分段线性)	0.8	1.4	38.75	初始幅度高于 1

在分布外检测（OOD Detection）能力上，基于 CIFAR-10 训练的显式能量模型 EqM-E (B/4 点积变体) 与典型似然/能量基线的 AUROC 对比如下：

模型架构	SVHN	Textures	Constant	平均 AUROC
PixelCNN++	0.32	0.33	0.71	0.45
GLOW	0.24	0.27	-	0.26
IGEBM	0.63	0.48	0.39	0.50
EqM-E (本文)	0.55	0.49	1.00	0.68

关键发现¶

截断衰减与零边界是稳态生成的核心：对比常数与纯线性目标，截断衰减（$a=0.8$）使得模型在高噪声区域能获得充足且稳定的指向性梯度，而在近数据流形区域迅速平滑衰减为零，严格满足驻点要求；配以梯度倍率 $\lambda=4.0$ 时性能达到 32.85 FID。
时间条件在平衡态场中呈现相反效果：当强制加入时间条件时，截断衰减反而使 FID 从 36.68 恶化至 41.89；而在无时间条件下，截断衰减将 FID 从 40.81 大幅优化至 32.85，证实了能量不动点理论与时间不变假设的内在契合。
灵活步长与自适应计算：Flow Matching 强依赖固定积分步长 $\eta = 1/N = 0.004$，步长微调即导致质量崩溃；而 EqM 在极宽的步长范围内保持高保真生成，且自适应停机机制仅用平均 40% 的计算量（节省 60% NFE）即可达到 33.79 FID（固定 250 步为 32.85 FID）。
部分加噪图像去噪与多概念组合：EqM 能直接输入任意加噪程度的图像进行去噪优化，噪声越小 FID 越好；而流匹配在未告知精确时间步时处理中间加噪样本性能严重崩溃。同时，两个类别的 EqM 梯度场可直接相加（$\nabla E_1 + \nabla E_2$）实现高质量概念组合生成（如“熊猫+山谷”）。

亮点与洞察¶

将生成动力学范式从数值积分重铸为凸/非凸优化：以往生成模型受制于时间维度的刚性步进，EqM 将其升华为时间无关的能量景观优化，使得成熟优化器（如 Nesterov 动量、自适应学习率、线搜索）能直接转化为推理加速武器。
优雅打通 Flow Matching 与 EBM 的鸿沟：克服了传统 EBM 训练中 MCMC 采样发散的死结，借用流匹配插值路径构建确定性负梯度监督目标，兼具流匹配的高维训练稳定性与能量模型的自然属性（OOD 检测、能量组合）。
完全解耦推理算力与时间调度：利用梯度范数直接作为流形收敛判定标准，使生成模型具备了类似 LLM 思维链长短可调节的“自适应推理时间算力”特性。

局限与展望¶

显式能量变体与隐式梯度场的性能折损：虽然 EqM-E 能显式输出标量能量并在 OOD 检测中夺冠，但由于引入了输入梯度的二阶反传项，其图像生成 FID 相比隐式 EqM 仍有可感知的差距（B/4 架构下隐式为 57.54，显式点积为 73.40）。
高度调优基准上的超参适配：在卷积架构（如 CIFAR-10 U-Net）上，受制于基线代码库对标准流匹配时间调度的极端过拟合，EqM 虽胜过无时间条件流匹配（3.36 vs 3.96），但略逊于深度调优的标准流匹配（2.09）。
未来方向：探索更强的一阶与拟牛顿优化器在潜在空间采样中的应用；研究如何将 EqM 拓展至视频生成与离散模态建模中。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 颠覆了扩散/流模型依赖时间调度的固有范式，以简洁优雅的目标设计统一了生成流动与能量稳态。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖 ImageNet 缩放定律、多采样器对比、超参消融、自适应计算及 OOD 与组合生成等多维度验证。
写作质量: ⭐⭐⭐⭐⭐ 数学推导严密自洽，直觉清晰，实验分析具有深刻的洞察力。
价值: ⭐⭐⭐⭐⭐ 为生成式 AI 的推理时优化与自适应计算开辟了全新的方法论路径。