[論文レビュー] Uncertainty Estimation for Molecules: Desiderata and Methods

Tom Wollschläger, Nicholas Gao+3 more|arXiv (Cornell University)|Jun 20, 2023

Machine Learning in Materials Science被引用数 5

ひとこと要約

本論文は、グラフニューラルネットワーク（GNN）ベースの分子力場における不確実性推定手法として、六つの主要な望ましい特性（物理的インスパイラションの3つと応用指向の3つ）を満たす、ガウス過程に基づく局所化ニューラルカーネル（LNK）を提案する。LNKは、高い予測精度を維持しながら、単一のフォワードパスで信頼性の高い不確実性評価を可能にし、分布外（OOD）予測における不均衡検出において、モンテカルロドロップアウトや証拠的回帰よりも最大2.5倍および2.1倍低いAUC-ROC誤差を達成する。

ABSTRACT

Graph Neural Networks (GNNs) are promising surrogates for quantum mechanical calculations as they establish unprecedented low errors on collections of molecular dynamics (MD) trajectories. Thanks to their fast inference times they promise to accelerate computational chemistry applications. Unfortunately, despite low in-distribution (ID) errors, such GNNs might be horribly wrong for out-of-distribution (OOD) samples. Uncertainty estimation (UE) may aid in such situations by communicating the model's certainty about its prediction. Here, we take a closer look at the problem and identify six key desiderata for UE in molecular force fields, three 'physics-informed' and three 'application-focused' ones. To overview the field, we survey existing methods from the field of UE and analyze how they fit to the set desiderata. By our analysis, we conclude that none of the previous works satisfies all criteria. To fill this gap, we propose Localized Neural Kernel (LNK) a Gaussian Process (GP)-based extension to existing GNNs satisfying the desiderata. In our extensive experimental evaluation, we test four different UE with three different backbones and two datasets. In out-of-equilibrium detection, we find LNK yielding up to 2.5 and 2.1 times lower errors in terms of AUC-ROC score than dropout or evidential regression-based methods while maintaining high predictive performance.

研究の動機と目的

• 分布外（OOD）予測において標準的なGNNが失敗する状況において、分子力場における信頼性の高い不確実性推定の重要性に対応する。
• 有効な不確実性推定のための六つの主要な望ましい特性（物理的インスパイラションの3つ：対称性、エネルギー保存、局所性、および応用指向の3つ：高速性、解釈可能性、OOD検出能力）を特定する。
• 分子機械学習分野における既存の不確実性推定手法を調査し、それらが六つの望ましい特性をすべて満たさないことを示す。
• GNNバックボーンとガウス過程回帰を統合することで、すべての望ましい特性を満たす、新規の単一フォワードパス不確実性推定手法LNKを提案する。
• 複数のバックボーンと不確実性推定手法を用いて、二つのデータセット上でLNKを実証的に評価し、OOD検出と予測性能に注目する。

提案手法

• 分子力場における不確実性推定のための六つの望ましい特性を定義する：物理的インスパイラションの3つ（ユークリッドおよび置換対称性の不変性、カurlフリーな力によるエネルギー保存）と応用指向の3つ（高速性、解釈可能性、OOD検出能力）。
• 分子グラフ上に局所化されたカーネル関数を学習することで不確実性をモデル化する、微分可能でGPベースのGNN拡張である「局所化ニューラルカーネル（LNK）」を提案する。
• 既存のGNNバックボーン（例：DimeNet++）にLNKを統合することで、バックボーンの再トレーニングなしに不確実性推定を可能にし、単一のフォワードパスで実現する。
• 物理的対称性を尊重し、力場が構造的に保存的であることを保証する微分可能カーネルメカニズムを採用する。
• 変動に基づくアプローチにより、エネルギー（u_E）および力（u_F）の両方の不確実性評価を可能にし、u_Fはカーネルのトレースから計算する。
• エネルギー/力の予測と不確実性のキャリブレーションを組み合わせた確率的損失関数を用いて、エンドツーエンドでモデルを訓練し、良好にキャリブレートされた不確実性推定を実現する。

実験結果

リサーチクエスチョン

• RQ1 分子力場における不確実性推定に求められる物理的および応用指向の基本的要件は何か？
• RQ2 分子機械学習分野における既存の不確実性推定手法は、これらの望ましい特性をどの程度満たしているか？
• RQ3 分子の対称性および保存則を尊重しつつ、高速で単一フォワードパスによる不確実性推定を可能にするGPベースの手法を設計できるか？
• RQ4 提案されたLNK手法は、モンテカルロドロップアウトや証拠的回帰などの既存の不確実性推定技術と比較して、OOD検出および予測性能においてどのように差をつけるか？
• RQ5 GNNバックボーンにLNKを統合することで、高い予測精度を維持しながら不確実性のキャリブレーションを著しく改善できるか？

主な発見

• 分子機械学習分野における既存の不確実性推定手法は、六つの提案された望ましい特性をすべて満たさず、特に物理的整合性または推論速度に欠ける。
• QM7Xデータセットにおける非平衡状態検出において、LNKはモンテカルロドロップアウトおよび証拠的回帰よりも最大2.5倍および2.1倍低いAUC-ROC誤差を達成する。
• LNKは高い予測性能を維持し、1％のドロップアウト率においてエネルギーのMAEが0.059 kcal/mol、力のMAEが0.007 kcal/molにとどまり、最先端モデルと同等の性能を示す。
• 分散ベースの手法（MCD やアンサンブル）と比較して、LNKは4.6倍以上高速であり、QM7Xデータセットでの推論時間は1サンプルあたり59.3±2.68 msで、証拠的回帰と同等の速度を達成する。
• 推論時にドロップアウトを用いるMCD-LNKの変種は、ドロップアウト率の増加に伴って安定した性能を示し、標準的なMCDを上回るOOD検出性能と予測精度を達成する。
• アンサンブルは力のAUC-ROCが最高の0.996を記録するが、トレーニング時間が5倍に増加し、推論時間も274±8.25 msと著しく長くなるため、アクティブラーニングやリアルタイム応用には現実的ではない。