[论文解读] Coherent energy and force uncertainty in deep learning force fields

Peter Bjørn Jørgensen, Jonas Busk+2 more|arXiv (Cornell University)|Dec 7, 2023

Machine Learning in Materials Science参考文献 16被引用 8

一句话总结

本论文通过建模相关噪声过程，将深度学习力场中的能量与力的本征不确定性联系起来，推导力的方差表达，并在等变信息传递神经网络集合上对分子数据集进行一致性不确定性验证。

ABSTRACT

In machine learning energy potentials for atomic systems, forces are commonly obtained as the negative derivative of the energy function with respect to atomic positions. To quantify aleatoric uncertainty in the predicted energies, a widely used modeling approach involves predicting both a mean and variance for each energy value. However, this model is not differentiable under the usual white noise assumption, so energy uncertainty does not naturally translate to force uncertainty. In this work we propose a machine learning potential energy model in which energy and force aleatoric uncertainty are linked through a spatially correlated noise process. We demonstrate our approach on an equivariant messages passing neural network potential trained on energies and forces on two out-of-equilibrium molecular datasets. Furthermore, we also show how to obtain epistemic uncertainties in this setting based on a Bayesian interpretation of deep ensemble models.

研究动机与目标

• 说明在机器学习力场中对能量和力的一致性不确定性估计的必要性。
• 提出一个噪声模型，通过空间相关过程将能量本征不确定性与力不确定性联系起来。
• 推导在有色噪声下能量与力方差的闭式表达。
• 通过深度集合实现本体不确定性（epistemic uncertainty），并给出贝叶斯解释。
• 在非平衡分子数据集（ANI-1x 和 Transition1x）上，使用等变 PaiNN 基于的集合对方法进行评估。

提出的方法

• 将能量观测建模为 E_obs = E_theta + rho_theta * eta，其中 eta 具有可微分的自相关性。
• 推导力的方差：Var(-∂E_obs/∂r) = gamma_hat * rho_theta^2 + (∂rho_theta/∂r)^2 (式(8))。
• 将 gamma_hat 解释为噪声核的平方倒长尺度，并将其作为超参数进行学习或设定。
• 使用广义平稳噪声简化方差推导（附录 A）。
• 采用贝叶斯深度集合框架通过后验采样获得本体不确定性（附录 B）。
• 训练等变的基于 PaiNN 的集合，具备经过标定的本体/本体不确定性分量，并在能量和力的 MAE/RMSE、NLL、以及校准指标（表1）上进行评估。

Figure 1: Ensemble model (with 10 instances) fitted to energy (left) and force (right) observations with correlated noise.

实验结果

研究问题

• RQ1 如何在 ML 力场中将能量与力的本征不确定性一致地连接起来？
• RQ2 对不确定性估计的影响：将相关（有色）噪声建模引入能量和力的不确定性会产生哪些影响？
• RQ3 通过这种一致性框架能否在深度集合中捕捉本体不确定性？
• RQ4 与原生和白噪声基线相比，该方法在非平衡分子数据（ANI-1x 和 Transition1x）上的表现如何？
• RQ5 在实际校准和预测性能方面，这种一致性不确定性方法有哪些好处？

主要发现

• 有色噪声的能量模型给出一个与能量不确定性耦合的力方差的闭式表达，耦合通过超参数 gamma_hat (式(8)) 实现。
• 力的方差包含一个与方差输出梯度的平方成分和一个与 rho_theta^2 成比例的成分，因此能够实现能量/力不确定性的统一估计。
• 本体不确定性可以通过贝叶斯深度集合获得，从而将本体与本体不确定性分离（附录 B）。
• 在 ANI-1x 和 Transition1x 上，有色/有噪声的集合在 MAE/RMSE 方面具有竞争力，并在不确定性指标（NLL、ENCE、RZV、CV）方面优于原生和白噪声基线（表1）。
• 经过验证集校准可以在不牺牲预测准确性的前提下提升校准的不确定性表现。

Figure 2: Reliability diagrams for models trained on the ANI-1x dataset.