[論文レビュー] The Design Space of E(3)-Equivariant Atom-Centered Interatomic Potentials

Ilyes Batatia, Simon Batzner+7 more|arXiv (Cornell University)|May 13, 2022

Machine Learning in Materials Science被引用数 98

ひとこと要約

本論文は ACE と NequIP を Multi-ACE フレームワーク内で統合し、設計上の重要な選択肢を特定するためのアブレーションを実施し、解釈可能でありながら高精度なモデルとして BOTNet を導入します。

ABSTRACT

The rapid progress of machine learning interatomic potentials over the past couple of years produced a number of new architectures. Particularly notable among these are the Atomic Cluster Expansion (ACE), which unified many of the earlier ideas around atom density-based descriptors, and Neural Equivariant Interatomic Potentials (NequIP), a message passing neural network with equivariant features that showed state of the art accuracy. In this work, we construct a mathematical framework that unifies these models: ACE is generalised so that it can be recast as one layer of a multi-layer architecture. From another point of view, the linearised version of NequIP is understood as a particular sparsification of a much larger polynomial model. Our framework also provides a practical tool for systematically probing different choices in the unified design space. We demonstrate this by an ablation study of NequIP via a set of experiments looking at in- and out-of-domain accuracy and smooth extrapolation very far from the training data, and shed some light on which design choices are critical for achieving high accuracy. Finally, we present BOTNet (Body-Ordered-Tensor-Network), a much-simplified version of NequIP, which has an interpretable architecture and maintains accuracy on benchmark datasets.

研究の動機と目的

• ACE をどのようにして E(3) 対称性に拡張し、多層 MPNN フレームワークに組み込むかを明確にする。
• ディスクリプタベースの ACE と等変な MPNN を Multi-ACE 設計空間の下で統一する。
• 統一されたフレームワーク内で NequIP のアブレーション研究を系統的に実施して設計選択を検討する。
• 高精度に不可欠な要素を特定し、性能を維持したままより簡潔で解釈可能なモデル（BOTNet）を導入する。
• 設計空間内での迅速な実験を可能にする実践的な洞察とモジュール型実装を提供する。

提案手法

• 等変な ACE 層をメッセージパッシング方式として積み重ねる Multi-ACE フレームワークを定義する。
• 連続的な元素埋め込みを用いた等変な一粒子基底を組み込んで ACE を拡張する。
• ACE 密度の工夫と積基底を用いて高次のボディオーダー特徴を形成する。
• Wigner-D 機構を用いて O(3) 回転に対して基底関数を対称化し、不変/等変特徴を得る。
• 線形/非線形変換および層ごとの状態更新を介して MPNN の更新とリードアウトを ACE 層に結びつける。
• 非本質的な非線形活性化を伴わず精度を保つ、簡略化され解釈可能な BOTNet 変種を導入する。

実験結果

リサーチクエスチョン

• RQ1 Multi-ACE フレームワークのどの設計要素がベンチマークデータセットで高い精度を達成するうえで最も重要ですか？
• RQ2 NequIP の等変/非等変成分のアブレーションによって、インドメイン内およびドメイン外の性能はどのように変化しますか？
• RQ3 BOTNet のような簡略化モデルは解釈性とモジュラリティを改善しつつ NequIP レベルの精度を維持できますか？
• RQ4 正規化と元素埋め込みは等変な分子間ポテンシャルの安定性と外挿挙動においてどのような役割を果たしますか？

主な発見

• 統一された Multi-ACE フレームワークはディスクリプタベースの ACE と等変性 MPNN を結びつけ、設計選択の系統的な探索を可能にする。
• アブレーション研究は NequIP のどの革新がピーク性能に不可欠であり、どの点で単純化が妥当かを浮き彫りにする。
• 正規化と連続的元素埋め込みはモデルの安定性と外挿挙動に重要な役割を果たす。
• BOTNet はベンチマークデータセット全体で精度を維持しつつ、はるかに単純な解釈可能なアーキテクチャとして浮上する。
• このフレームワークは、最も複雑な非線形活性化のいくつかを用いなくても高い精度を達成できることを示している。