Bayesian Neural Networks versus deep ensembles for uncertainty quantification in machine learning interatomic potentials

技術要約：機械学習原子間ポテンシャルにおける不確実性定量化のためのベイズニューラルネットワークとディープアンサンブルの比較

問題提起
機械学習原子間ポテンシャル（MLIP）は、密度汎関数理論（DFT）のような第一原理手法に代わる計算効率の高い選択肢として、原子スケールのモデリングにおいて不可欠なものとなっている。しかし、標準的なディープラーニングアーキテクチャは決定論的であり、予測の不確実性を本質的に定量化するメカニズムを欠いている。この限界は、モデルが分布外（out-of-distribution）のデータに適用されることが多い材料科学において極めて重要である。信頼できる不確実性の推定がなければ、シミュレーションは誤解を招く結果を生むリスクがあり、モデルの洗練のための能動学習戦略も指針を欠くことになる。Deep Ensembles (DE) と Bayesian Neural Networks (BNNs) は、この問題に対処するための確立された戦略であるが、MLIPの特定の文脈、特にエネルギーと力の同時予測における、異なるデータレジームを横断した比較性能については、体系的な評価が必要である。

手法
著者らは、ænet-PyTorch フレームワーク内で2つの主要な不確実性定量化アプローチを実装し、比較を行った：

1. Deep Ensembles (DE): Lakshminarayananらによる定式化に従い、異なる乱数シードで初期化された5つの独立したニューラルネットワークを訓練した。不確実性は、アンサンブルの予測値の標本標準偏差として定量化した。
2. Variational Bayesian Neural Networks (VBNN): 著者らは、Pyro確率プログラミングフレームワークを用いた変分推論（VI）を用いてBNNを実装した。2つの特定の変分ガイドがテストされた：
   • Local Reparametrization Trick (LRT): 勾配の分散を抑えるためにBayes by BackpropアルゴリズムとLRTを用いた、完全な因子分解型ガウス事後分布による平均場近似を利用。
   • Radial Guide (RAD): 重みの間のグローバルな相関を捉えるために放射状変換を導入した非平均場アプローチであり、証拠下限（ELBO）の標準的なモンテカルロサンプリングを介して最適化された。

実験設定
本研究は、以下の2つの異なるデータセットを用いてこれらの手法を評価した：

• TiO2 データセット: チタン酸化物の様々な相を表す7,815構造。モデルは、データ効率を評価するために、フルデータセット（High-Dataレジーム）および20%のサブセット（Low-Dataレジーム）の両方で訓練された。エネルギーのみの訓練と、エネルギー・力同時訓練の両方が行われた。
• QM7 データセット: 化学的に多様なシステムにおける汎用性をテストするために使用される、7,165個の小さな有機分子。モデルは、原子化エネルギーのみを用いてHigh-Dataレジームで訓練された。

評価指標
性能は、包括的な一連の指標によって評価された：

• 精度 (Accuracy): 平均絶対誤差 (MAE)、平方根平均二乗誤差 (RMSE)、および負の対数尤度 (NLL)。
• 不確実性の質 (Uncertainty Quality): 平方根平均較正誤差 (RMSCE)、鋭さ (Sharpness: 平均的な予測不確実性)、不確実性と誤差の間のSpearman順位相関係数 ($r_s$)、および高不確実性の予測と高誤差領域の整合性を測定する「Overlap Score」。

主な結果

1. 予測精度:

   • High-Data レジーム (TiO2) では、DEがVBNNを一貫して上回り、最小のMAE (6 meV/atom) および RMSE (9 meV/atom)、そして最も負に大きいNLL (-4.2) を達成し、優れた精度と較正性を示した。
   • Low-Data レジーム においても、DEは精度（MAE 35 meV/atom）において優位性を維持した。しかし、LRTは競争力のある精度（MAE 54 meV/atom）とDEよりも低いRMSE（259 meV/atom）を示し、外れ値に対する堅牢性が高いことが示唆された。RADは低データ設定では性能が悪かった。
   • QM7 データセット では、逆転現象が起きた。BNNs（LRTおよびRAD）はDE（15.2 meV/atom）よりも低いエネルギーRMSE（13.0–13.3 meV/atom）を達成しており、これは変分事前分布が化学的に多様な空間に対して有益な正則化を提供していることを示唆している。

2. 不確実性の定量化:

   • 較正性 (Calibration): DEは一般に、最も優れた較正された不確かな予測（最低のRMSCE）を提供した。Low-Dataレジームにおいて、RADは驚くべきことに優れた較正性（RMSCE 0.05）を維持し、High-Dataでの性能と同等であったが、LRTとDEは較正性が低下した。
   • 鋭さ (Sharpness): DEは最も鋭い（自信のある）予測を生み出した。LRTは、特に低データレジームにおいて、過度に広い不確実性区間（高いSharpness値）を生み出す傾向があった。
   • 相関とオーバーラップ: DEは、低データレジームにおいて予測不確実性と実際の誤差の間の強い単調相関 ($r_s = 0.45$) を示し、LRT (0.28) やRAD (0.19) を上回った。また、DEは最高のOverlap Score (55%) を達成しており、高誤差領域を最も効果的に特定できることを示した。

3. 力（Force）の予測:

   • Low-Dataレジームにおいて、VBNNs（LRTおよびRAD）はDEよりも有意に低い力のRMSEを達成した。これは、事前分布による正則化によるものと考えられる。しかし、力の較正性はすべてのモデルにおいて低く、DEのみが低データレジームにおいて合理的な較正性を示した。

4. 計算効率:

   • 訓練: DEは1エポックあたりのコストは低いが、5つの独立したモデルを訓練する必要がある。DEの総訓練時間はHigh-DataレジームではLRTと同等であったが、Low-Dataレジームではわずかに高くなった。
   • 推論: DEは推論において大幅に高速である (0.03 s) のに対し、VBNNs (0.64–0.71 s) は遅い。これは、BNNが予測ごとに自動微分を含む20個のモンテカルロサンプルを必要とするためである。

意義と結論
本論文は、Deep Ensemblesは理論的にはBNNほど根拠が強固ではないものの、精度、較正性、および推論効率に関するほとんどの指標において、経験的に変分BNNを上回っていると結論付けている。特に、化学空間が比較的狭い場合（TiO2）において顕著である。信頼できる不確実性定量化が優先され、かつ実装の単純さが求められる場合は、DEが推奨される。

しかし、本研究は、VBNN（特にLRTのバリアント）が競争力のあるトレードオフを提供し、データが乏しい、あるいは化学的に多様なレジーム（QM7に見られるような）において、変分事前分布が効果的な正則化として機能するため、より優れている可能性があることを強調している。著者らは、BNNはランダムな初期化に敏感であり、堅牢性を確保するために複数の訓練実行を必要とすると述べている。本研究は、不確実性を考慮したMLIPを選択するための実践的な指針を提供しており、DEとVBNNの選択は、データのレジームおよび対象となるシステムの化学的多様性に依存すべきであることを示唆している。

貢献

• エネルギー・力同時訓練をサポートする、ænet-PyTorch フレームワーク内でのGPU加速されたVBNN実装の導入。
• 無機（TiO2）および有機（QM7）のデータセットを用い、High-DataおよびLow-DataレジームにわたるDE、LRT、RADの体系的な比較。
• DEが一般的に優れている一方で、化学的に多様なデータセットではBNNの相対的な性能が逆転することを示し、あらゆる文脈におけるアンサンブルの優位性の普遍性に疑問を投げかけたこと。
• 能動学習シナリオにおける不確実性推定の有用性を評価するための、新しい指標である「Overlap Score」の提供。