Energy Loss Functions for Physical Systems¶
会议: NeurIPS 2025
arXiv: 2511.02087
代码: 有
领域: 图像生成
关键词: 能量损失函数, 物理先验, 扩散模型, 对称不变性, 分子生成
一句话总结¶
提出基于物理能量的损失函数框架，通过反向 KL 散度与玻尔兹曼分布推导出以成对距离为基础的能量差损失，天然满足 SE(d) 不变性，在分子生成和自旋基态预测中显著优于 MSE 和交叉熵损失。
研究背景与动机¶
将机器学习应用于物理系统时，数据通常稀缺且昂贵，而物理先验知识的利用大多聚焦于架构层面（如等变网络）。损失函数这一互补方向被严重忽视。
核心问题：通用损失函数（MSE、交叉熵）为何不适合物理系统？
- MSE 对应非物理能量：MSE 等价于以数据为中心的各向同性谐振势 \(E = \|\hat{y} - y\|^2\)——这描述的是外力将粒子拉回固定点，而非粒子间相互作用
- 对称性破坏：物理配置通常在刚体运动 SE(d) 下等价，但 MSE 会惩罚通过旋转/平移与目标等价的正确预测
- MSE 最小化器偏差：当存在多个对称等价目标时，MSE 的最小化器是所有目标的均值——这不是任何合法物理配置
方法详解¶
整体框架¶
损失函数设计的核心思路：
- 将每个数据样本的不确定性建模为以该数据为中心的玻尔兹曼分布
- 使用反向 KL 散度（而非前向 KL / MLE）作为优化目标
- 最终损失等价于数据与预测之间的能量差
反向 KL 的两大优势：(a) 配分函数 \(Z(y, T)\) 不依赖模型参数 \(\theta\)，无需计算难以处理的归一化常数；(b) 只需定义数据点附近的势函数（预测可能很差，在预测处定义能量无意义）。
关键设计¶
1. 从 Reverse KL 到能量损失¶
对于玻尔兹曼分布 \(p(\hat{y}|y) = \exp(-E(\hat{y},y)/T) / Z(y,T)\)，反向 KL 损失为：
\(\log Z\) 项不含 \(\theta\)，可忽略。模型因预测偏离数据的能量增量而受惩罚。
MSE 是能量损失的特例：取 \(E(\hat{y},y) = \|\hat{y}-y\|^2\), \(T=2\sigma^2\)，得到高斯条件分布。但这对应的是非物理的外部谐振势。
2. 原子系统的二次对势能¶
对 \(n\) 个原子在 \(d\) 维空间的系统，定义成对距离的二次势能：
物理含义：粒子间相互作用势能的二阶 Taylor 近似（弹簧模型）。弹簧系数 \(k_{ij}(y)\) 的选择：
| 系数类型 | 来源 | 适用场景 |
|---|---|---|
| 常数 \(k\) | Morse 势 Taylor 近似 | 理论分析 |
| 指数衰减 | 经验最优 | EDM/GDM 模型 |
| 逆距离平方 | Lennard-Jones 势 Taylor 近似 | 物理对应 |
| 逆距离 | 经验 | JODO 模型 |
3. 对称不变性——免费获得¶
能量损失仅依赖不变的成对距离，天然满足 \(E(d)\) 不变性。进一步，损失对距离矩阵的自同构群也不变：
推论：全局最小化器集合为 \(\{g \cdot y \mid g \in G\}\)——模型可回归到任何与数据对称等价的配置，无需 Kabsch 对齐（后者需要 \(O(n^3)\)）。
4. 扩散模型中的应用¶
将能量损失直接替换扩散训练中的 MSE。关键理论结果：
Proposition 4.4：对 SE(d) 不变密度，在小 \(\sigma_t\) 下，最小范数最小化器：
即正确恢复了 score 函数。
Proposition 4.5：能量损失的 score 估计器方差不高于 MSE：
5. 线性缩放：刚性理论¶
成对距离有 \(O(n^2)\) 项。利用 Laman 定理：仅需 \(O(n)\) 条边构成的刚性图即可从距离恢复坐标。稀疏刚性图将计算量降至线性，且不影响全局最优解。
6. 离散系统：自旋能量损失¶
对自旋系统 \(\hat{y}, y \in \{1,-1\}^\Lambda\)，定义局部场能量：
局部场捕获翻转自旋的能量变化。目标变为变分自由能最小化。局部场损失是凸的，以数据为唯一最小化器（\(h^0 > 4\) 时）。
损失函数 / 训练策略¶
能量损失直接替换 MSE/交叉熵，无需修改其他训练流程。关键调优：\(k_{ij}\) 函数形式和位置/原子类型损失权重比。所有对比实验均进行了充分的学习率和权重扫描。
实验关键数据¶
主实验¶
表1：GEOM-Drugs 分子生成（GDM-aug）
| 损失 | 分子稳定性 (%) | 原子稳定性 (%) | 有效性 (%) | 唯一性 (%) |
|---|---|---|---|---|
| MSE | 0.8 | 85.6 | 94.8 | 100 |
| Energy | 24.6 | 96.0 | 89.7 | 100 |
分子稳定性从 0.8% 跃升至 24.6%，改善约 30 倍。
表2：QM9 分子生成（GDM-aug）
| 损失 | 分子稳定性 (%) | 原子稳定性 (%) | 有效性 (%) |
|---|---|---|---|
| MSE | 83.7±2.3 | 98.3 | 93.6 |
| Kabsch align | 82.3±0.5 | 97.8 | 90.8 |
| Energy | 89.8±2.8 | 99.3 | 97.7 |
| Energy (sparse) | 89.1±0.9 | 99.0 | 97.4 |
表3：JODO 3D 与对齐指标
| 模型 | 分子稳定性 | 有效性 | Bond↓ | Angle↓ |
|---|---|---|---|---|
| JODO (原始) | 93.4 | — | 0.1475 | 0.0121 |
| JODO+Energy(Inv.) | 94.3 | 97.1 | 0.1125 | 0.0046 |
能量损失在近 SOTA 模型 JODO 上仍带来显著改善，对齐指标大幅优化。
消融实验¶
弹簧系数 \(k_{ij}\) 消融（QM9，GDM-aug）
| 系数类型 | 分子稳定性 (%) | 有效性 (%) |
|---|---|---|
| 指数衰减 | 89.8 | 97.7 |
| 逆距离平方 | 84.6 | 96.6 |
| 逆距离 | 84.5 | 95.0 |
| 常数 | 83.6 | 93.6 |
自旋基态预测
| 损失 | 测试能量 |
|---|---|
| 交叉熵 | 58.8±0.8 |
| Margin | 49.87±1.5 |
| 局部场能量 | 45.6±1.6 |
| 真实能量 | 14.6±0.3 |
关键发现¶
- 能量损失 + 非等变架构 > 等变架构 + MSE，计算代价可忽略
- 数据效率极高：50% 训练数据，能量损失产生 >75% 稳定分子（MSE 远不及）
- 稀疏刚性图将 \(O(n^2)\) 降至 \(O(n)\)，几乎无性能损失
- 不同模型最优 \(k_{ij}\) 不同（GDM 用指数衰减，JODO 用逆距离），需按任务调优
- 正则形状实验直观展示：能量损失对旋转增强完全免疫，MSE 在 \(\theta_{aug}=\pi\) 时崩溃
亮点与洞察¶
- 反向 KL + 玻尔兹曼 = 能量差损失：推导优雅，物理直觉清晰
- 不变性免费获得：距离基损失天然 SE(d) 不变，无需 Kabsch 的 \(O(n^3)\) 对齐
- 架构无关、即插即用：直接替换任何需要 MSE 的训练目标
- MSE 的物理病理性被形式化：Figure 1 的三幅对比图极其直观——旋转后正确配置被 MSE 错误惩罚
- 刚性理论的巧妙运用：利用图刚性实现 \(O(n^2) \to O(n)\) 缩放
局限性 / 可改进方向¶
- Score 恢复理论仅在小噪声下精确，高噪声需显式修正
- 弹簧系数选择仍需启发式调优，缺乏自动选择机制
- 未考虑扭转角信息，更丰富的替代能量（如包含角度项）值得探索
- 仅在小分子和 2D 自旋系统验证，蛋白质/晶体等大系统需进一步测试
- 非物理系数（指数衰减）优于理论上更正确的常数系数——这一张力需深入理解
相关工作与启发¶
- AlphaFold 3 (Abramson 2024): 已启发式使用距离损失作为正则项，本文为其提供了第一性原理的理论基础
- Kabsch 对齐损失 (Klein 2023): 另一种 SE(3) 不变损失，但需 \(O(n^3)\) 对齐
- JODO (Huang 2023): 近 SOTA 联合 2D/3D 分子扩散模型，能量损失进一步提升其指标
- 启发：损失函数作为注入物理先验的通道被严重低估——架构等变性与损失不变性互补而非替代
评分¶
- 新颖性：⭐⭐⭐⭐⭐ — 从第一性原理推导损失函数，物理视角独到
- 技术深度：⭐⭐⭐⭐⭐ — 连续/离散统一框架，扩散 score 恢复证明，刚性理论应用
- 实验充分度：⭐⭐⭐⭐ — 分子+自旋+形状三类实验，但分子规模偏小
- 实用价值：⭐⭐⭐⭐⭐ — 即插即用、计算高效、理论有保障
选择各向同性谐波势 \(E(\hat{\mathbf{y}}, \mathbf{y}) = \|\hat{\mathbf{y}} - \mathbf{y}\|^2\) 恰好回退到高斯分布和 MSE 损失。但这种"能量"描述的是外力将粒子拉回参考位置，物理上不合理。
3. 原子系统的二次对势能量¶
对于 \(n\) 个原子的位置 \(\mathbf{y} \in \mathbb{R}^{n \times d}\)，提出基于粒子间距离的二次对势：
这是一般交互势在成对距离上的二阶 Taylor 近似。系数 \(k_{ij}(\mathbf{y})\) 有多种选择：常数（Morse 势近似）、逆距离、逆距离平方（LJ 势近似）、指数衰减等。
4. 对称不变性¶
命题：当 \(k_{ij}(\mathbf{y})\) 不变时，能量损失对 \(G = E(d) \times (\text{Aut}(k(\mathbf{y})) \cap \text{Aut}(\Delta y))\) 群不变——即自然尊重欧氏变换和粒子置换对称性。
推论：损失的全局极小值族恰好是 \(\{g \cdot \mathbf{y} | g \in G\}\)，模型可以回归到任何与数据对称等价的构型。
相比 Kabsch 对齐后再算 MSE 的方法，能量损失无需对齐操作，计算更高效。
5. 用于扩散模型¶
将能量损失直接替代扩散模型训练中的 MSE。理论分析表明，在低噪声水平下，常系数能量损失的最优预测器近似为正确的 score function 加上一个刚体运动方向的自由度：
最小范数极小值恰好是正确的 score。此外证明了能量损失估计器的方差不超过 MSE 估计器。
6. 自旋系统的离散能量损失¶
对于 Ising 型哈密顿量 \(E(\mathbf{y}) = -\frac{1}{2}\sum_{ij} J_{ij} \mathbf{y}_i \mathbf{y}_j\)，定义局部场能量：
损失变为变分自由能：\(\mathcal{J}(\theta) = \frac{1}{T}\left[\mathbb{E}_q[E] - T S[q] + T \log Z\right]\)。
实验关键数据¶
分子生成 — QM9（GDM-aug 模型）¶
| 损失 | 分子稳定率(%) | 原子稳定率(%) | 有效率(%) | 唯一性(%) |
|---|---|---|---|---|
| MSE | 83.7 ± 2.3 | 98.3 | 93.6 | 100.0 |
| Kabsch 对齐 | 82.3 ± 0.5 | 97.8 | 90.8 | 100.0 |
| 能量损失 | 89.8 ± 2.8 | 99.3 | 97.7 | 99.9 |
| 能量损失(稀疏) | 89.1 ± 0.9 | 99.0 | 97.4 | 100.0 |
分子稳定率提升 +6.1%，有效率提升 +4.1%。
GEOM-Drugs 大分子¶
| 损失 | 分子稳定率(%) | 原子稳定率(%) | 有效率(%) |
|---|---|---|---|
| MSE | 0.8 | 85.6 | 94.8 |
| 能量损失 | 24.6 | 96.0 | 89.7 |
分子稳定率从 0.8% 飙升至 24.6%——在大分子场景下改进极其显著。
JODO 模型（近 SOTA）¶
能量损失在 atom stability (+0.4%)、mol stability (+3.8%)、validity (+2.8%) 上全面提升；bond length MMD 从 0.1218 降至 0.0928，表明生成的键长更精确。
自旋基态预测（16×16 格子）¶
| 损失 | 测试能量 |
|---|---|
| 交叉熵 | 58.8 ± 0.8 |
| Margin Loss | 49.87 ± 1.5 |
| 局部场能量损失 | 45.6 ± 1.6 |
| 真实能量（非分类） | 14.6 ± 0.3 |
能量损失相比交叉熵大幅降低预测构型的能量。
低数据量实验¶
仅用 50% 训练数据（50K 样本），能量损失仍生成超过 75% 稳定分子；MSE 在相同数据量下表现大幅恶化——数据效率提升显著。
亮点¶
- 优雅的理论框架：从 reverse KL + Boltzmann 分布出发，统一了 MSE / 交叉熵 / 能量损失，给出了损失函数的物理解释
- 自然的对称不变性：无需 Kabsch 对齐就能尊重 SE(d) 对称性，且计算开销可忽略
- 与架构正交：能量损失是架构无关的即插即用模块，与等变架构（如 EGNN）的改进是互补的
- 理论保证：证明了扩散模型中能量损失能正确估计 score function，且方差更低
- 大分子场景（GEOM-Drugs）的改进从 0.8% → 24.6%，定量上令人印象深刻
局限性 / 可改进方向¶
- 用于扩散模型时，正确 score 的恢复仅在低噪声水平下精确；高噪声下需要显式校正
- 能量函数的选择仍有一定的 ad hoc 性，系数 \(k_{ij}\) 需要在新任务上消融
- 标准版本计算复杂度为 \(O(N^2)\)（粒子对），虽然稀疏刚性图可降至 \(O(N)\)，但稀疏图的构建需要额外处理
- 未探索扭转角等更丰富的几何描述符，也未在蛋白质或晶体材料上验证
与相关工作的对比¶
| 维度 | MSE | Kabsch + MSE | 能量损失 |
|---|---|---|---|
| 对称不变 | ✗ | ✓（需对齐） | ✓（天然） |
| 物理先验 | 无 | 无 | 粒子间交互 |
| 计算开销 | O(N) | O(N³)对齐 | O(N²) 或 O(N) |
| Score 恢复 | 精确 | 精确 | 近似精确+低方差 |
| 与等变架构互补 | 部分 | 部分 | 强互补 |
启发与关联¶
- "损失函数也是一种先验"：损失函数隐含了对误差的分布假设，择定更符合物理的分布（Boltzmann）比简单高斯更有意义
- 对称性在损失层面解决比在架构层面更轻量：等变架构设计复杂且增加参数量，而不变损失函数仅改变目标计算
- 能量损失的梯度方向指向最近的合法构型（而非数据本身），这对对称性破缺的预测任务尤其有价值
- 框架可直接推广到任何处于热平衡的物理系统——晶体材料、蛋白质、软物质等
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ — 从第一性原理推导损失函数，理论深度出色
- 实验充分度: ⭐⭐⭐⭐ — 分子+自旋两类系统，多模型验证，消融充分
- 写作质量: ⭐⭐⭐⭐⭐ — 理论推导严谨优雅，直觉解释到位，图示启发性强
- 价值: ⭐⭐⭐⭐⭐ — 提供了一个通用且实用的框架，可直接改善现有分子生成流程