[論文レビュー] 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

ひとこと要約

論文は深層学習力場におけるエネルギーと力のアレイオトリック不確実性を correlated ノイズ過程をモデリングして結びつけ、力の分散表現を導出し、等変性メッセージパッシングニューラルネットワークのエンサンブルを用いて分子データセット上で一貫した不確実性を検証します。

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.

研究の動機と目的

• 機械学習力場におけるエネルギーと力の一貫した不確実性推定の必要性を動機づける。
• エネルギーのアレオリック不確実性を力の不確実性と空間的に相関した過程を通じて結びつけるノイズモデルを提案する。
• カラー付きノイズの下でエネルギーと力の分散の閉形式表現を導出する。
• ディープエンサンブルとベイズ的解釈によって認識的不確実性を示す。
• 等変 PaiNN ベースのエンサンブルを ANI-1x および Transition1x の非平衡分子データセットで評価する。

提案手法

• エネルギー観測を E_obs = E_theta + rho_theta * eta としてモデリングし、eta は微分可能な自己相関を持つ。
• 力の分散を導出: Var(-∂E_obs/∂r) = gamma_hat * rho_theta^2 + (∂rho_theta/∂r)^2 (式 (Eq. 8))。
• gamma_hat をノイズカーネルの二乗逆長スケールとして解釈し、それを学習するかハイパーパラメータとして設定する。
• 分散導出を単純化するために wide-sense stationary ノイズを用いる（付録 A）。
• ベイズ的ディープエンサンブルの枠組みを用いて事後サンプリングから認識的不確実性を得る（付録 B）。
• 厳密性と適切な aleatoric/epistemic 成分をキャリブレーションした 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 提案手法は vanilla および white-noise ベースラインと比較して、非平衡分子データ（ANI-1x および Transition1x）でどう機能するか？
• RQ5 この一貫した不確実性アプローチの実用的なキャリブレーションと予測性能の利点は何か？

主な発見

• カラー付きノイズエネルギーモデルは、力の分散をエネルギー不確実性と結ぶ閉形式を gamma_hat (式 (Eq. 8)) により提供する。
• 力の分散は分散出力の勾配の二乗項と rho_theta^2 に比例する項を含み、一貫したエネルギー/力の不確実性推定を可能にする。
• ベイズ的ディープエンサンブルから認識的不確実性を得られ、アレオトリック成分と認識的不確実性の寄与を分離できる（付録 B）。
• ANI-1x および Transition1x で、カラー付き / ノイズベースのエンサンブルは vanilla および white-noise ベースラインと比較して MAE/RMSE で競争力を持ち、NLL、ENCE、RZV、CV などの不確実性指標が改善またはキャリブレーションされている（表 1）。
• キャリブレーションされた不確実性は、予測精度を犠牲にすることなく検証データセットのキャリブレーションを通じて改善可能。

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