[論文レビュー] Fast and Uncertainty-Aware Directional Message Passing for Non-Equilibrium Molecules

Johannes Gasteiger, Shankari Giri+2 more|arXiv (Cornell University)|Nov 28, 2020

Machine Learning in Materials Science被引用数 160

ひとこと要約

DimeNet++は非平衡分子に対する方向性メッセージパッシングを高速化し、精度を向上させ、COLLと共に非平衡研究を可能にし、不確実性定量化分析を可能にします。

ABSTRACT

Many important tasks in chemistry revolve around molecules during reactions. This requires predictions far from the equilibrium, while most recent work in machine learning for molecules has been focused on equilibrium or near-equilibrium states. In this paper we aim to extend this scope in three ways. First, we propose the DimeNet++ model, which is 8x faster and 10% more accurate than the original DimeNet on the QM9 benchmark of equilibrium molecules. Second, we validate DimeNet++ on highly reactive molecules by developing the challenging COLL dataset, which contains distorted configurations of small molecules during collisions. Finally, we investigate ensembling and mean-variance estimation for uncertainty quantification with the goal of accelerating the exploration of the vast space of non-equilibrium structures. Our DimeNet++ implementation as well as the COLL dataset are available online.

研究の動機と目的

• 機械学習による分子モデリングを非平衡反応および歪んだ配置へ拡張する。
• エネルギーと原子力学的力のための高速で正確かつ汎用性の高いGNNを開発する。
• 反応性分子配置の非平衡データセットCOLLを作成する。
• 非平衡領域におけるエネルギーと力の予測のための不確実性定量化戦略を調査する。

提案手法

• DimeNet の高価な二項方向相互作用をHadamard積に置換し、基底表現には2層MLPで補償する。
• 埋め込みサイズを削減し、速度と表現力を向上させるために、ダウンプロジェクションおよびアッププロジェクション層を含む埋め込み階層を実装する。
• 可能な箇所でネットワークの深さを6層から4層に削減し、埋め込み/テンソル表現を調整して精度を維持する。
• 高速でスケーラブルな方法で結合距離と角度をモデル化するため、更新されたアーキテクチャで方向性メッセージパッシングを維持する。
• QM9とCOLLデータセットで評価し、非平衡配置に対する速度・精度・頑健性を評価する。
• エンサンブル法と平均-分散推定などの不確実性定量化アプローチと、それらのエネルギーおよび力の予測における限界を検討する。

実験結果

リサーチクエスチョン

• RQ1 QM9 のような平衡ベンチマークで、DimeNet++ はより高速な実行時間を達成しつつ精度を向上または維持できるか？
• RQ2 COLL のような高度に非平衡で反応性のある配置に対して、DimeNet++ はどれくらい一般化できるか？
• RQ3 標準的不確実性定量化手法（エンサンブル法、平均-分散推定）は、非平衡領域におけるエネルギーと力の不確実性予測に信頼できるか？
• RQ4 MLモデルのトレーニングと不確実性推定に影響を与える非平衡分子データの特徴と課題は何か？

主な発見

• DimeNet++ は元の DimeNet より8倍速く、平均的に約10%の精度向上をもたらし、QM9 のエネルギーで約20%の改善をもたらす。
• COLLデータセットでは、DimeNet++ はSchNetを大幅に上回る（DimeNet++ の MAE E = 0.047 eV、MAE F = 0.040 eV、SchNetの方が高い誤差）。
• COLLデータセット（140k configurations）は、歪んだ結合・角度を含むより広い非平衡領域を包含し、QM9よりも強い課題を課す。
• エンサンブリングと平均-分散推定による不確実性定量化は難航する：エネルギーの不確実性は力の不確実性の信頼できる代理指標ではなく、エンサンブルは顕著な計算オーバーヘッドを伴う。
• エンサンブリングは力の不確実性推定をより良くするがコストが高い。一方、MVE は力の不確実性に対して有用性が限定的である。