[論文レビュー] FAENet: Frame Averaging Equivariant GNN for Materials Modeling

Alexandre Duval, Victor Schmidt+5 more|arXiv (Cornell University)|Apr 28, 2023

Machine Learning in Materials Science被引用数 7

ひとこと要約

FAENetは、幾何学的データの確率的フレーム平均化（SFA）を通じて、構造的制約なしにE(3)同変性を達成する、新しい軽量グラフニューラルネットワークを導入する。この手法により、OC20およびQM9データセットで最先端の精度と計算効率を達成し、スケーラビリティと対称性の保持の両面で、先行する同変モデルを上回る。

ABSTRACT

Applications of machine learning techniques for materials modeling typically involve functions known to be equivariant or invariant to specific symmetries. While graph neural networks (GNNs) have proven successful in such tasks, they enforce symmetries via the model architecture, which often reduces their expressivity, scalability and comprehensibility. In this paper, we introduce (1) a flexible framework relying on stochastic frame-averaging (SFA) to make any model E(3)-equivariant or invariant through data transformations. (2) FAENet: a simple, fast and expressive GNN, optimized for SFA, that processes geometric information without any symmetrypreserving design constraints. We prove the validity of our method theoretically and empirically demonstrate its superior accuracy and computational scalability in materials modeling on the OC20 dataset (S2EF, IS2RE) as well as common molecular modeling tasks (QM9, QM7-X). A package implementation is available at https://faenet.readthedocs.io.

研究の動機と目的

• 複雑なアーキテクチャ設計によって対称性を強制する既存の同変GNNが抱える表現力とスケーラビリティの制限を解消すること。
• モデルアーキテクチャではなく、データ変換によってE(3)対称性を保持する柔軟なフレームワークの開発。
• 幾何情報の自由な処理を可能にしつつ、理論的および実験的対称性不変性を維持する、軽量で高性能なGNN—FAENetの設計。
• SFAに基づく対称性強制が、データ増強や対称性強制なしの手法と比較して、材料物性予測において優れていることを実証的に検証すること。

提案手法

• E(3)同変性を強制するための対称性保持型データ拡張技術として、確率的フレーム平均化（SFA）を提案。原子配置を標準フレームに投影することで、同変性を実現する。
• 対称性保持アーキテクチャの制約なしに、相対的な原子位置を処理するGNNアーキテクチャであるFAENetを導入。同変性はSFAに依存する。
• フレーム平均化を用いて入力グラフを標準座標系に変換し、任意のモデルが回転、反転、平行移動に対して同変になるようにする。
• 一般化と対称性のロバストネスを向上させるために、SFAで微分可能なフレーム選択メカニズムを採用し、複数の基準フレームを確率的にサンプリングする。
• 相対位置をエッジ特徴として用いる幾何的メッセージパッシングを活用し、表現力を最大化するとともに、データ前処理による対称性の保持を実現する。
• 回転不変性、反転不変性、力の同変性などの指標を用いて、複数のデータセットで対称性保持度を評価する。

Figure 1: Overview of FAENet architecture making use of the Stochastic Frame Averaging framework. FAENet takes an input sample $D$ and makes an invariant $\Phi^{\text{{invar}}}_{{\mathcal{F}}}$ (e.g. energy) and/or equivariant prediction $\Phi^{\text{{equiv}}}_{{\mathcal{F}}}$ (e.g. atomic forces) b

実験結果

リサーチクエスチョン

• RQ1 アーキテクチャの制約なしに、データ変換によってGNNにおけるE(3)同変性を効果的に強制できるか？
• RQ2 SFAに基づく対称性強制は、データ増強や対称性強制なしの手法と比較して、モデルの精度と対称性保持度においてどのように異なるか？
• RQ3 FAENetのような軽量GNNは、複雑な対称性制約付きモデルと比較して、材料物性予測タスクにおいてより高い精度とスケーラビリティを達成できるか？
• RQ4 完全なフレーム平均化や標準的なデータ増強と比較して、SFAは対称性一般化をどの程度向上させるか？
• RQ5 FAENetは、固体結晶や分子を含む多様な材料データセットに対して、最小限のアーキテクチャ変更で効果的に一般化できるか？

主な発見

• S2EFデータセットにおいて、FAENet＋SFAはテストMAEが464 meVに達し、SchNet（918 meV）とForceNet（655 meV）を上回り、優れた精度とスケーラビリティを示した。
• SFAはS2EFで回転不変性誤差を6.94 meV、反転不変性誤差を6.96 meVにまで低減。データ増強（8.88および9.12 meV）やNo-FA（26.81および27.61 meV）を顕著に上回った。
• IS2REでは、FAENet＋SFAがMAE 569 meVを達成し、ForceNet比で12%、SchNet比で13%の改善を示した。対称性保持度も高く維持された。
• SE(3)-SFAバージョンは完全なFAよりわずかに優れた指標を示し、少ないフレームからの学習が一般化を向上させつつ対称性を損なわないことを示唆した。
• No-FAおよびNo-Trainでは、回転不変性誤差が51.1 meVにまで上昇し、対称性強制が強固なパフォーマンスに不可欠であることを実証した。
• SFAは、特に3次元対称性下で、データ増強を上回る対称性保持と予測精度を達成。S2EFでは回転同変性誤差が60%低下した。

Figure 2: OC-20 Direct IS2RE performance / scalability trade-off. FAENet achieves the best MAE across methods while being much faster than the majority of baselines.