Graph Energy Matching: Transport-Aligned Energy-Based Modeling for Graph Generation¶

会议: NeurIPS 2026
arXiv: 2603.23398
论文: Project Page
领域: 图学习
关键词: 图生成、能量模型、最优传输、JKO 方案、分子生成

一句话总结¶

本文提出离散图能量生成框架 Graph Energy Matching (GEM)，借鉴 JKO 最优传输映射视角学习置换不变标量势能函数，并通过能量阈值门控机制无缝衔接确定性贪心传输与局部平衡 Metropolis-Hastings 混合采样，在分子图无条件生成、属性约束优化及图测地线插值任务上全面达到或超越主流离散扩散与流匹配模型。

研究背景与动机¶

分子发现与材料设计等科学任务高度依赖离散图结构的生成建模，不仅要求生成高化学有效性与多样性的分子图，还需要在推理阶段灵活融入理化属性约束或先验偏好。当前主流的无条件图生成方法（如 DiGress、DeFoG 和 Variational Flow Matching）大多依赖离散时间去噪或连续时间马尔可夫链（CTMC）将简单噪声先验逐步去噪为目标图。这类方法虽然在无条件生成质量上表现优异，但它们仅通过带噪中间状态的转移概率隐式定义数据分布，缺乏直接在干净图流形上刻画相对似然的显式概率模型；当在推理阶段施加用户指定的属性引导时，预测器必须在化学键和原子价态严重受损的流形外带噪图上评估性质，导致引导信号失真与有效率下降。

相比之下，基于能量的模型（EBM）通过标量能量函数直接建模干净图的相对对数似然，天然支持在推理时零样本组合外部约束与属性目标而无需重新训练主干生成模型。然而，以 GraphEBM 为代表的传统离散图能量模型长期受困于高维组合空间中的采样效率低下问题：流形外的低密度区域遍布虚假局部极小值，纯靠局部随机游走或连续嵌入空间 Langevin 动力学极易陷入次优模态，不仅导致对比散度训练频繁崩溃，也使生成样本的保真度大幅落后于离散扩散模型。连续空间中的 Energy Matching 虽然通过 Jordan–Kinderlehrer–Otto (JKO) 变分格式统一了传输引导与平衡态微调，但其依赖连续时间常微分/随机微分方程（ODE/SDE）轨迹积分，无法直接推广到依赖局部组合状态跳变的离散图空间。

本文的切入角度是：虽然图的采样移动发生在离散的节点与边类别空间，但只要将离散图嵌入到实值独热连续表示空间并在线性化局部邻域内评估势能梯度，就能把连续 JKO 传输映射的一阶最优性条件转化为离散局部图编辑的评分准则，从而用同一个标量能量场同时驱动全局流形外传输与局部流形内探索。核心 idea：提出基于 JKO 传输映射视角的离散图能量匹配框架（GEM），通过近似 Fused Gromov–Wasserstein 最优传输耦合在离散插值路径上对齐能量负梯度与传输位移，并在推理时采用无显式时间步的能量门控机制，将零温极限下的确定性贪心降能传输与正温下的局部平衡 Metropolis-Hastings 混合采样有机统一。

方法详解¶

整体框架¶

GEM 的核心目标是仅学习一个置换不变的标量势能网络 \(V_\theta(x)\)，使其既能为远离数据流形的噪声图提供指向高概率区域的确定性传输方向，又能在数据流形内部精确刻画玻尔兹曼分布 \(\pi_\theta(x) \propto \exp(-\beta_{\mathrm{mh}} V_\theta(x))\)。给定最多含 \(n\) 个节点的无向图 \(x = (x_{\mathcal{V}}, x_{\mathcal{E}})\)，方法首先将离散节点类别与无向边类别映射为实值独热拼接向量 \(\hat{x} \in \mathbb{R}^d\)，送入基于 DeFoG 的置换等变 Graph Transformer 主干（融合相对随机游走概率 RRWP 特征），最后对高层节点特征、边特征做置换不变求和池化并与全局图特征拼接，经双层 MLP 头部输出标量能量 \(V_\theta(x)\)。在训练阶段，模型通过小批量近似最优传输配对噪声图与真实图，沿着类别插值分布最小化梯度流匹配损失与 MCMC 对比学习损失；在采样阶段，模型从均匀噪声图或真实训练图出发，根据当前样本能量是否达到训练集平均能量水平自动在贪心传输提议与带 Metropolis-Hastings (MH) 接受准则的离散 Langevin 混合提议之间切换，最终输出高有效性与高新颖性的目标分子图。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入：离散图独热嵌入<br/>噪声图或数据图对"] --> B["置换不变能量势能与<br/>传输对齐训练<br/>学习标量势函数与负梯度场"]
    B --> C["能量阈值门控的<br/>离散贪心传输提议<br/>沿负梯度确定性局部编辑"]
    C -->|达到目标能量或陷入局部极值| D["局部平衡离散混合提议<br/>与 MH 修正<br/>温度退火与细致平衡采样"]
    D --> E["输出：高似然生成图<br/>或属性约束/测地线插值图"]

关键设计¶

1. 置换不变能量势能与传输对齐训练：利用小批量最优传输耦合将连续 JKO 一阶最优性条件转化为离散插值路径上的梯度匹配与对比学习

在离散图空间中直接训练能量模型极易因为负样本采样器无法到达数据流形而发散。为解决这一冷启动难题，GEM 从零温极限（\(\epsilon \to 0\)）下的单步 JKO 变分问题出发：将从源分布到目标分布的演化写成确定性传输映射 \(T_k(x) \in \arg\min_y \{ \frac{1}{2\eta} c(x, y) + V_\theta(y) \}\)。在独热嵌入空间 \(\hat{x}\) 中取单位步长 \(\eta = 1\) 并利用一阶最优性条件，可以得到势能函数关于嵌入向量的梯度应当与指向目标数据图的负位移向量对齐，即 \(\nabla_{\hat{x}} \tilde{V}_\theta(\hat{x}) = -(\hat{x}^{\mathrm{data}} - \hat{x}^0)\)。为了在保持图置换不变性的前提下构造噪声图 \(x^0 \sim \pi_0\) 与真实数据图 \(x^{\mathrm{data}} \sim \pi_{\mathrm{data}}\) 的小批量配对，本文定义了加权节点与边独热欧氏距离的局部代价 \(c_{\mathrm{loc}}(x,y)\)，并在节点重标号群 \(S_n\) 上取极小化得到硬置换 Fused Gromov–Wasserstein (FGW) 代价 \(c_{\mathrm{FGW}}(x,y)\)。由于精确求解 FGW 等价于二次指派问题（QAP），训练时采用高效的节点类型、边类型及无序节点对边类型归一化直方图签名进行线性指派近似配对，随后在配对图之间按时间 \(t \sim \mathcal{U}(0,1)\) 独立对每个节点和边进行分类分布插值采样得到离散中间图 \(x^t\)，从而在纯离散图输入上计算流匹配回归目标。

仅靠流匹配只能塑造从噪声指向流形的宏观梯度场，尚不足以雕刻流形内部的精细能量盆地。因此，GEM 在完成 \(N_{\mathrm{warmup}}\) 步纯流匹配预热后，引入对比散度损失 \(\mathcal{L}_{\mathrm{CL}}(\theta)\)，拉低真实数据图 \(x^+\) 的能量并推高由当前模型通过贪心传输加混合采样生成的负样本 \(x^-\) 的能量（对采样转移过程做梯度截断）。这种先预热势能场、冻结网络校准混合采样器超参数、再联合微调对比损失的策略，确保了 MCMC 负样本在引入对比学习之初就能稳定落入高概率数据流形附近，彻底消除了传统离散 EBM 训练早期的模式崩溃。

2. 能量阈值门控的离散贪心传输提议：在零温极限下通过一阶泰勒线性化执行降能局部图编辑，无需显式时间步索引即可快速跨越低密度区

离散扩散与流匹配模型在推理时必须严格按照预设的时间步 \(t\) 调度去噪率矩阵，而能量模型本身不依赖显式时间索引。为了从远离流形的均匀随机噪声图快速移动到有效分子区域，GEM 在传输阶段（\(s=0\)）对势能函数 \(\tilde{V}_\theta(\hat{y})\) 在当前图嵌入 \(\hat{x}\) 处做一阶泰勒展开，在恰好发生 \(N\) 处离散修改（改变原子类型或增删改化学键）的候选邻域集合 \(C_N(x)\) 中，确定性地挑选使线性化能量变化与局部传输代价之和最小的候选图：

\[ y^*(x) = \arg\min_{y \in C_N(x)} \left( (\hat{y} - \hat{x})^\top \nabla_{\hat{x}} V_\theta(\hat{x}) + \lambda_{\mathcal{V}} \|\hat{x}_{\mathcal{V}} - \hat{y}_{\mathcal{V}}\|_2^2 + \lambda_{\mathcal{E}} \|\hat{x}_{\mathcal{E}} - \hat{y}_{\mathcal{E}}\|_2^2 \right) \]

在实现中，贪心传输仅接受投影在严格负梯度方向（即能使能量下降）的局部编辑。与需要昂贵连续积分的采样器不同，这种贪心跳跃充当了极其高效的确定性预烧（burn-in）机制。更重要的是，GEM 设计了自适应的能量阈值门控来决定何时结束传输阶段：一旦马尔可夫链当前样本的能量下降到训练集样本的平均能量参考值 \(V_{\mathrm{th}} = \mathbb{E}_{\pi_{\mathrm{data}}}[V_\theta]\) 以下，或者所有候选编辑均无法继续降低能量（陷入局部极小值），采样器便自动将状态指示器切换至 \(s=1\)，无缝转入随机混合阶段，既避免了在流形外盲目随机游走，又防止了贪心搜索在局部极值处死锁。

3. 局部平衡离散混合提议与 MH 修正：结合嵌入空间距离惩罚与能量梯度构造离散 Langevin 提议，并通过 Metropolis-Hastings 接受率保证平稳分布收敛

当样本进入高概率数据区域后，继续使用确定性贪心更新会导致样本坍缩到孤立能量极小点而丧失多样性与理论分布保证。为此，GEM 在混合阶段（\(s=1\)）将连续嵌入空间中的过阻尼 Langevin 动力学离散化，利用能量梯度 \(\nabla_{\hat{x}} V_\theta(\hat{x})\) 与局部图编辑距离惩罚构造多站点独立分解的分类提议分布 \(q_{\mathrm{mixing}}(x \to y)\)：

\[ q_{\mathrm{mixing}}(x \to y) \propto \exp\bigl( -\lambda_{\mathcal{V}}^L \|\hat{y}_{\mathcal{V}} - \hat{x}_{\mathcal{V}}\|_2^2 - \lambda_{\mathcal{E}}^L \|\hat{y}_{\mathcal{E}} - \hat{x}_{\mathcal{E}}\|_2^2 - \beta^L \nabla_{\hat{x}} V_\theta(\hat{x})^\top (\hat{y} - \hat{x}) \bigr) \]

具体实现时，每个节点 \(i\) 变为类别 \(c_i\) 的对数几率由当前类别与目标类别的梯度差 \(\beta^L (g_{i,x_i}^{\mathcal{V}} - g_{i,c_i}^{\mathcal{V}})\) 减去非自环修改惩罚 \(\lambda_{\mathcal{V}}^L \mathbb{I}_{\{c_i \neq x_i\}}\) 给出（无向边同理），从而可以通过单次前向加反向传播并行采样所有节点和边的局部跳变；若独立采样恰好停留在原图（无任何修改），则按比例因子 \(\rho \in (0,1)\) 软化温度与修改惩罚进行有界次数的重采样以强制促进局部探索。随后，采样器引入目标逆温度为 \(\beta_{\mathrm{mh}}\) 的 Metropolis-Hastings 接受概率 \(\alpha(x,y) = \min\bigl(1, \frac{\pi_\theta(y) q_{\mathrm{mixing}}(y \to x)}{\pi_\theta(x) q_{\mathrm{mixing}}(x \to y)}\bigr)\)，严格满足细致平衡条件，从理论上保证马尔可夫链以 \(\pi_\theta(x) \propto \exp(-\beta_{\mathrm{mh}} V_\theta(x))\) 为平稳分布。此外，当直接从真实训练集分子出发（Data Initialization）以生成新颖变体时，模型跳过贪心传输阶段，直接采用双温度模拟退火策略——在固定提议温度 \(\beta^L\) 的同时，让 MH 目标逆温度 \(\beta_{\mathrm{mh}}\) 从较低初值（高接受率以跳出原训练样本邻域）在 \(S_{\mathrm{anneal}}\) 步内线性升温至终值，从而在保持高化学有效性与分布保真度的同时高效获得高新颖度分子。

损失函数 / 训练策略¶

总训练目标为流匹配损失与对比散度损失的加权和：

\[ \min_\theta \; \mathcal{L}_{\mathrm{Flow}}(\theta) + \lambda_{\mathrm{CL}} \mathcal{L}_{\mathrm{CL}}(\theta) = \mathbb{E}_{(x^0, x^{\mathrm{data}})\sim\tilde{\gamma},\, t\sim\mathcal{U}[0,1]} \left[ \|\nabla_{\hat{x}_t} \tilde{V}_\theta(\hat{x}^t) + (\hat{x}^{\mathrm{data}} - \hat{x}^0)\|_2^2 \right] + \lambda_{\mathrm{CL}} \left( \mathbb{E}_{x^+ \sim \pi_{\mathrm{data}}}[\tilde{V}_\theta(\hat{x}^+)] - \mathbb{E}_{x^- \sim \mathrm{sg}(\pi_\theta)}[\tilde{V}_\theta(\hat{x}^-)] \right) \]

在 MOSES 数据集上，训练分两阶段进行：第一阶段设置 \(\lambda_{\mathrm{CL}} = 0.0\)、学习率 \(10^{-4}\)、批大小 128，训练 330k 步预热势能场；随后冻结网络，在 100 次独立运行上通过最小化生成样本能量自动校准混合采样器超参数（得到 \(\beta^L = \beta_{\mathrm{mh}} = 9.55, \lambda_{\mathcal{V}}^L = 0.23, \lambda_{\mathcal{E}}^L = 1.88\)）；第二阶段开启 \(\lambda_{\mathrm{CL}} = 0.1\)、学习率 \(10^{-5}\)、采样链长 \(N_{\mathrm{CL}} = 500\)，继续微调 1000 步（通用模型负样本采用 50% 噪声初始化与 50% 数据初始化，数据初始化专用模型采用 100% 数据初始化）。在推理阶段，这一能量参数化还可直接拓展至两类下游任务：① 属性约束条件生成：无需重训生成主干，直接将条件势能定义为 \(V_\theta^{\mathrm{cond}}(x) = V_\theta(x) + \lambda_{\mathrm{prop}} \|f_\phi(x, t) - \zeta\|_2^2\)（其中 \(f_\phi\) 为噪声条件属性回归器，伪时间 \(t\) 由当前能量在噪声与数据能量带之间的线性插值代理给出并在到达数据分布后截断为 \(t=1\)）；② 图间测地线插值：在节点与边概率单纯形上用三次 B 样条参数化连续路径 \(\gamma(t)\)，最小化能量加权路径长度 \(\int_0^1 \exp(\beta V_\theta(\hat{\gamma}(t))) \|\dot{\hat{\gamma}}(t)\|_{\mathrm{loc}} \mathrm{d}t\) 与平均能量正则项，从而找出沿着高概率有效化学空间演化的分子过渡路径。

实验关键数据¶

主实验¶

表 1 展示了在 MOSES（约 150 万药效分子，最多 27 个重原子，生成 25k 样本，1000 步）和 QM9（约 15 万小有机分子，最多 9 个重原子，生成 10k 样本，500 步）上的无条件生成结果，评价指标为有效-唯一-新颖率（V.U.N. \(\uparrow\)）和 Fréchet ChemNet 距离（FCD \(\downarrow\)）：

数据集与初始化方式	方法	V.U.N. \(\uparrow\)	FCD \(\downarrow\)	说明
MOSES (Noise Init.)	GraphEBM*	0.081	9.83	传统图能量基线（替换同款主干）
MOSES (Noise Init.)	VFM*	0.814	2.71	变分流匹配基线
MOSES (Noise Init.)	DeFoG*	0.822	1.95	之前离散流匹配 SOTA
MOSES (Noise Init.)	GEM (本文)	0.856	1.51	V.U.N. 提升 3.4 个百分点，FCD 降低 22.6%
MOSES (Data Init.)	GraphEBM*	0.322	7.31	数据初始化的传统 EBM
MOSES (Data Init.)	GEM (本文)	0.898	0.76	达到最高有效唯一新颖率与极低分布距离
QM9 (Noise Init.)	GraphEBM*	0.079	7.16	传统图能量基线
QM9 (Noise Init.)	VFM	0.489	0.44	新颖率较低（0.49），故 FCD 偏低
QM9 (Noise Init.)	DeFoG	0.348	0.81	离散流匹配基线
QM9 (Noise Init.)	GEM (本文)	0.612	0.92	V.U.N. 比 VFM 大幅提升 12.3 个百分点
QM9 (Data Init.)	GraphEBM*	0.103	5.15	数据初始化的传统 EBM
QM9 (Data Init.)	GEM (本文)	0.502	0.31	同时超越噪声初始化基线的 FCD

表 2 展示了在 MOSES 上针对四种理化属性约束（每个约束在无条件样本中满足率约为 20%）的条件生成结果（5k 样本，1000 步，由 RDKit 精确评估 C.V.U.N. \(\uparrow\) 与 C.V.U. \(\uparrow\)）：

方法	\(\mathrm{logS} \ge -2.25\) (C.V.U.N. / C.V.U.)	\(\mathrm{QED} \ge 0.9\) (C.V.U.N. / C.V.U.)	\(\mathrm{TPSA} \le 50\) (C.V.U.N. / C.V.U.)	\(\mathrm{logP} \le 1.5\) (C.V.U.N. / C.V.U.)
G-VFM*	0.72 / 0.75	0.26 / 0.28	0.72 / 0.80	0.70 / 0.74
DeFoG*	0.74 / 0.79	0.26 / 0.31	0.77 / 0.82	0.78 / 0.82
GEM (本文)	0.80 / 0.81	0.40 / 0.40	0.80 / 0.81	0.85 / 0.85

消融实验¶

表 3 汇总了在 MOSES 噪声初始化下（1000 步）对传输与混合两阶段采样机制的消融实验（原表 5），以及在属性条件生成中对干净回归器与噪声条件回归器的消融实验（原表 7）：

实验组	配置变体	核心指标表现	说明
采样阶段消融 (Unconditional)	Full GEM (Transport + Mixing)	V.U.N. = 0.856	完整双阶段能量门控采样
采样阶段消融 (Unconditional)	Transport-only (同权重)	V.U.N. = 0.712	仅用贪心传输，缺少局部混合微调掉点 0.144
采样阶段消融 (Unconditional)	Mixing-only (同权重)	V.U.N. = 0.485	无贪心传输预烧，直接混合采样大幅退化
采样阶段消融 (Unconditional)	Mixing-only (重新训练)	V.U.N. = 0.210	缺乏传输预烧导致对比学习约 60 步即崩溃
属性回归器消融 (Conditional)	Clean Regressor (仅干净分子训练)	C.V.U.N. = 0.531 / 0.154 / 0.552	对应 logS / QED / logP，在带噪中间态严重失真
属性回归器消融 (Conditional)	Noise-conditioned Regressor	C.V.U.N. = 0.836 / 0.422 / 0.846	结合能量代理伪时间 \(t\)，三项任务全面跃升

关键发现¶

贪心传输预烧是离散 EBM 稳定训练与高质量采样的核心支柱：消融表明，即便使用相同的预训练势能权重，去掉传输阶段仅靠混合采样会使 MOSES V.U.N. 从 0.856 暴跌至 0.485；若在对比学习训练中也去掉贪心传输预烧，模型在约 60 个迭代步后即发生早期模式崩溃（V.U.N. 仅为 0.21）。同时，单独使用确定性贪心传输也只能达到 0.712，说明进入流形后的随机 MH 混合对于修补局部化学价态和探索多模态至关重要。
在充足推理预算下显著反超离散流匹配：NVIDIA H100（batch size 128）真实挂钟时间测试显示，DeFoG 在前 2 秒内快速收敛至 V.U.N. 0.823 / FCD 1.956 后即陷入平台期；而 GEM 凭借持续的能量降维与 MH 细致平衡修正，在 6~15 秒内将噪声初始化性能稳步推升至 V.U.N. 0.860 / FCD 1.714，数据初始化更达到 V.U.N. 0.891 / FCD 0.655。
能量加权测地线有效捕获化学流形几何：在 MOSES 分子对插值实验中，仅最小化图编辑代价的插值路径在跨越较大 FGW 距离（如 \([15, 20)\) 区间）时会穿过大量低概率非化学区域，导致中间分子有效率骤降；而 GEM 的能量加权测地线成功绕开高能垒无效结构，在整个距离区间内始终维持显著更高的中间分子化学有效率。

亮点与洞察¶

无时间步索引的能量门控双模切换机制：不同于连续 Energy Matching 依赖僵硬的预设时间阈值 \(t=1\) 切换 ODE 与 SDE，GEM 直接利用学习到的势能函数自身作为“流形距离计”，当样本能量降至训练集平均能量 \(V_{\mathrm{th}}\) 或贪心编辑停滞时自适应触发混合采样。这一设计让采样器在无需显式时间输入的情况下兼具流匹配的快速传输能力与 MCMC 的渐近平稳性保证。
用能量高度巧妙代理扩散时间以实现抗噪引导：由于干净属性回归器在带噪图上误差极大，而纯 EBM 又没有扩散时间步 \(t\)，作者利用当前图能量在初始噪声平均能量与真实数据平均能量之间的相对位置线性映射出伪时间 \(t \in [0, 1]\) 喂给噪声条件回归器。这一精巧的桥接技巧使无时间步的 EBM 也能无缝享受时间条件引导器的抗噪优势，在困难的 \(\mathrm{QED} \ge 0.9\) 任务上将 C.V.U.N. 从 DeFoG 的 0.26 提升至 0.40。
可迁移至通用离散组合优化与结构生成的范式：本文展示的“独热嵌入一阶泰勒线性化 + 局部编辑代价惩罚 + 停滞重采样”提议核，不依赖特定分子领域先验，可直接迁移到蛋白质序列协同设计、神经架构搜索或芯片布局图生成等任何具有局部邻域编辑结构的离散组合生成任务中。

局限与展望¶

梯度回传带来的单步推理计算开销：由于 GEM 只参数化标量势能 \(V_\theta(x)\)，每次生成提议都需要对节点和边的连续独热嵌入求一阶梯度 \(\nabla_{\hat{x}} V_\theta(\hat{x})\)，相比直接前向输出转移速率矩阵的离散扩散或流匹配模型（如 DeFoG），单步推理存在约 \(2\sim 3\times\) 的理论计算开销，在极低延迟预算（如 \(<2\) 秒）下不如 DeFoG 快。
JKO 最优传输推导与实际离散采样器之间存在理论近似间隙：正文中的连续 JKO 传输映射推导仅作为损失设计与提议构造的启发式动机；实际运行中由于采用了直方图线性指派近似 FGW、局部一阶泰勒线性化以及离散贪心截断，严格的马尔可夫链平稳分布保证仅适用于固定温度的 MH 混合阶段，而不覆盖传输阶段与温度退火过程。
节点规模固定与更大规模非分子图验证缺失：当前框架要求先从经验分布采样节点数并对小图做零填充掩码，仅支持同尺寸图之间的局部节点/边类型修改，无法在单条采样轨迹中动态增删节点数量；未来可结合可变维度跳转 MCMC（如生灭过程）并在更大规模的通用社交图、代码图或材料晶格图上验证传输对齐能量模型的扩展性。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 将连续 JKO 能量匹配优雅地推广到置换不变的离散图空间，用能量门控机制统一了确定性贪心传输与满足细致平衡的离散 MH 混合采样。
实验充分度: ⭐⭐⭐⭐⭐ 在 QM9 与 MOSES 上覆盖了噪声/数据双初始化无条件生成、四项理化属性条件优化、图测地线有效性分析以及多随机种子、组件消融与 H100 挂钟时间评测。
写作质量: ⭐⭐⭐⭐⭐ 从连续 JKO 变分原理到离散泰勒线性化提议核的数学推导层层递进，附录对近似匹配、退火调度、理论保证边界与超参标定的交代极为坦诚详尽。
价值: ⭐⭐⭐⭐⭐ 成功填平了离散能量模型与离散扩散模型之间的生成质量鸿沟，为需要零样本约束组合与似然评估的科学离散结构生成提供了强有力的新范式。