[論文レビュー] Spherical Channels for Modeling Atomic Interactions

C. Lawrence Zitnick, Abhishek Das+6 more|arXiv (Cornell University)|Jun 29, 2022

Machine Learning in Materials Science被引用数 36

ひとこと要約

球面チャネルネットワーク（SCN）を導入したGNNで、球面調和基底のチャネルを用いて原子のエネルギーと力を回転に対して操作可能な埋め込みと可変等変性を持つ表現でモデリングし、エネルギーと力の予測においてOC20で最先端の結果を達成します。

ABSTRACT

Modeling the energy and forces of atomic systems is a fundamental problem in computational chemistry with the potential to help address many of the world's most pressing problems, including those related to energy scarcity and climate change. These calculations are traditionally performed using Density Functional Theory, which is computationally very expensive. Machine learning has the potential to dramatically improve the efficiency of these calculations from days or hours to seconds. We propose the Spherical Channel Network (SCN) to model atomic energies and forces. The SCN is a graph neural network where nodes represent atoms and edges their neighboring atoms. The atom embeddings are a set of spherical functions, called spherical channels, represented using spherical harmonics. We demonstrate, that by rotating the embeddings based on the 3D edge orientation, more information may be utilized while maintaining the rotational equivariance of the messages. While equivariance is a desirable property, we find that by relaxing this constraint in both message passing and aggregation, improved accuracy may be achieved. We demonstrate state-of-the-art results on the large-scale Open Catalyst dataset in both energy and force prediction for numerous tasks and metrics.

研究の動機と目的

• DFTによる原子エネルギーと力の高い計算コストを動機づけ、対処する。
• 球面チャネルを介して角情報をエンコードするグラフニューラルネットワークを開発し、予測精度を向上させる。
• 回転等変性の役割と、それを緩和することでモデルの表現力と性能を向上させる可能性を探る。
• エネルギーと力タスクにおいてOpen Catalyst 2020（OC20）データセットで最先端の性能を示す。
• これまでの最先端モデルと比較してサンプル効率とスケーラビリティを評価する。

提案手法

• 各原子をノード、近傍をエッジとしてグラフベースのモデルで表現する。
• ノード埋め込みを球面上の球面チャネルの集合としてエンコードし、度数Lまでの球面調和関数で展開する。
• メッセージ伝搬中に角情報を伝えるため、エッジ方向に合わせてエッジ指向の埋め込みを回転させる。
• 回転をグローバルフレームへ戻すニューラルネットワークを用いてエッジメッセージを計算し、場合によっては等変性を奨励するために複数のz軸回転（タップ）を使用する。
• 球面グリッド上の点ごとに制約なしの非線形性を適用してメッセージを集約し、更新のために球面調和関数へ戻す。
• orientations に対する神経網路を用いて球面積分として各原子エネルギーを推定し、エネルギー勾配または球面方向上での直接的な力推定のいずれかで力を計算する。
• AdamWで学習し、データ並列化と混合精度を活用し、適用可能な場合には複数の分解能（バンド）の球面調和関数を用いてメモリと精度のバランスを取る。

実験結果

リサーチクエスチョン

• RQ1 SCNの球面チャネル埋め込みはOC20上で既存のGNNよりエネルギーと力の予測を改善するか。
• RQ2 厳密な回転等変性を緩和してm=0を超える場合、モデルの精度と訓練効率にどのような影響があるか。
• RQ3 回転させた埋め込みと複数のタップを力の予測に用いることは角情報にどのような影響を与えるか。
• RQ4 OC20タスクにおけるSCNのサンプル効率は従来の最先端モデルと比べてどうか。

主な発見

• SCNはS2EFタスクで力のMAEと力コサインの点で、複数のベースラインと比較して最先端の改善を達成。
• OC20ではmを[-1,1]または[-2,2]とするSCN変種が厳密なm=0よりも有利であり、厳密な等変性を緩和することの利点を示す。
• より大きなLと深いKへスケーリングしたSCNは強力な性能を示し、IS2REエネルギー予測と全体的な力の精度で顕著な改善を得る。
• F_e計算時に複数回の回転（タップ）をサンプリングすることで力の予測が大幅に改善される一方、スループットにコストがかかる場合がある。一方で2つのバンド（L, L-1）を使用すると、スループットへの影響は限定的でエネルギーと力の予測が改善される。
• 直接力予測（勾配ベースのエネルギー力ではなく）を用いる設定で、速度と精度のトレードオフが良好になる場面が多く、なおかつ競争力のある精度を維持する。
• SCNは複数の強力なベースラインよりサンプル効率が高く、訓練例が少ない条件でも同等の精度を達成できることを示す。