[論文レビュー] MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields

Ilyes Batatia, Dávid Péter Kovács+3 more|arXiv (Cornell University)|Jun 15, 2022

Machine Learning in Materials Science被引用数 269

ひとこと要約

MACE は高次元（多体）の等変性メッセージ伝搬を原子間ポテンシャルに導入し、2層ネットワークで最先端の精度と benchmarks での高速学習を実現します。 rMD17, 3BPA, AcAc のようなベンチマークで。

ABSTRACT

Creating fast and accurate force fields is a long-standing challenge in computational chemistry and materials science. Recently, several equivariant message passing neural networks (MPNNs) have been shown to outperform models built using other approaches in terms of accuracy. However, most MPNNs suffer from high computational cost and poor scalability. We propose that these limitations arise because MPNNs only pass two-body messages leading to a direct relationship between the number of layers and the expressivity of the network. In this work, we introduce MACE, a new equivariant MPNN model that uses higher body order messages. In particular, we show that using four-body messages reduces the required number of message passing iterations to just two, resulting in a fast and highly parallelizable model, reaching or exceeding state-of-the-art accuracy on the rMD17, 3BPA, and AcAc benchmark tasks. We also demonstrate that using higher order messages leads to an improved steepness of the learning curves.

研究の動機と目的

• 化学と材料科学のための高速で正確な力場を動機づける。
• 二体 MP NN の制限を克服するために高体数のメッセージを導入する。
• 効率的なテンソル演算を備えたスケーラブルで並列可能なアーキテクチャを開発する。
• 挑戦的なベンチマーク（rMD17, 3BPA, AcAc）で最先端の精度を示す。
• 体の順序と等変性が学習曲線と汎化にどう影響するかを調査する。

提案手法

• nu-body までの相互作用を含む階層的な多体メッセージ構築を導入する。
• 近傍集合の和としてメッセージを表現し、体の順序を増加させる（Equation 7）。
• テンソル積と Clebsch–Gordan 結合によって等変性を強制する高次特徴を構築する（Equations 8–10）。
• 残差接続を用いた線形更新でメッセージを伝播する（Equation 12）。
• 特徴の不変部分を読み出して、層ごとの寄与の和としてサイトエネルギーを得る（Equation 13）。
• 高次メッシングを計算的に実現するために効率的なループテンソル収縮を実装する。

実験結果

リサーチクエスチョン

• RQ1 等変 MPNN における高次（多 body）メッセージは、精度を損なうことなく必要なメッセージ伝播層の数を減らすことができるか？
• RQ2 体の順序と等変性を高めると、力場予測の学習曲線とデータ効率にどのように影響するか？
• RQ3 高次で等変性のあるメッセージは、分域外分子配置の外挿と汎化に優れた性能を提供するか？
• RQ4 標準ベンチマーク上で既存の最先端の等変MPNNと比較して、MACE は速度とスケーラビリティの点でどのように性能を発揮するか？

主な発見

• 高次メッセージは必要な層数を二層に削減しつつ、ベンチマークタスクで最先端の精度を達成または超える。
• 十分な体順序を用いた不変メッセージは等変モデルの性能に近づく一方で、高次の等変メッセージは最良の結果をもたらす。
• L=2（二層、2L+M 等変メッセージ）を用いた MACE は 3BPA 外挿テストで競合モデルを上回り、ベンチマーク全体でエネルギーと力の予測において優れている。
• MACE は学習と評価の速度を大幅に向上させ、前述のトップモデルと比べて精度を維持または向上させつつ高速。
• 学習曲線は、高次の体情報を組み込むとデータ効率の勾配がシフト・急峻になり、学習ダイナミクスの改善を確認できる。