[論文レビュー] So3krates: Equivariant attention for interactions on arbitrary length-scales in molecular systems

J. Thorben Frank, Oliver T. Unke+1 more|arXiv (Cornell University)|May 28, 2022

Machine Learning in Materials Science被引用数 15

ひとこと要約

So3kratesは、動的で注意誘導の下での球面調和座標（SPHCs）を学習することで、分子系における非局所的電子的効果をモデル化する画期的な等長的アテンション機構を導入する。幾何的情報を原子特徴から分離し、SPHC空間における球面自己アテンションを適用することで、0.25–0.4倍のパラメータでSOTAの精度を達成するとともに、学習および推論が2–14倍高速化され、任意の長さスケールにわたる長距離相関の学習が効率的に行える。

ABSTRACT

The application of machine learning methods in quantum chemistry has enabled the study of numerous chemical phenomena, which are computationally intractable with traditional ab-initio methods. However, some quantum mechanical properties of molecules and materials depend on non-local electronic effects, which are often neglected due to the difficulty of modeling them efficiently. This work proposes a modified attention mechanism adapted to the underlying physics, which allows to recover the relevant non-local effects. Namely, we introduce spherical harmonic coordinates (SPHCs) to reflect higher-order geometric information for each atom in a molecule, enabling a non-local formulation of attention in the SPHC space. Our proposed model So3krates - a self-attention based message passing neural network - uncouples geometric information from atomic features, making them independently amenable to attention mechanisms. Thereby we construct spherical filters, which extend the concept of continuous filters in Euclidean space to SPHC space and serve as foundation for a spherical self-attention mechanism. We show that in contrast to other published methods, So3krates is able to describe non-local quantum mechanical effects over arbitrary length scales. Further, we find evidence that the inclusion of higher-order geometric correlations increases data efficiency and improves generalization. So3krates matches or exceeds state-of-the-art performance on popular benchmarks, notably, requiring a significantly lower number of parameters (0.25 - 0.4x) while at the same time giving a substantial speedup (6 - 14x for training and 2 - 11x for inference) compared to other models.

研究の動機と目的

• 従来の機械学習力場が計算的・アーキテクチャ的制限のためしばしば無視される、分子系における非局所的電子的効果のモデル化という課題に取り組むこと。
• 局所的近傍を越える長距離相互作用を捉える、高次幾何的相関を表現する微分可能で学習可能な表現を開発すること。
• 3次元分子系のSO(3)対称性に適応した自己アテンション機構を用いて、任意の長さスケールにわたる分子相互作用の効率的でスケーラブルなモデル化を可能にすること。
• 球面調和座標（SPHCs）を学習可能な特徴として明示的に幾何的相関をモデル化することで、データ効率と一般化性能を向上させること。
• モデルの複雑さと推論コストを低減しつつ、分子性質予測においてSOTAの性能を維持または上回ること。

提案手法

• 原子の幾何的情報を動的で学習可能な表現として球面調和座標（SPHCs）を導入し、固定または手作業で作成された幾何的特徴に置き換える。
• 不変な原子特徴と等長的なSPHCsを分離することで、それぞれに独立したアテンション機構を適用可能とし、SPHC空間における局所的および非局所的アテンションを実現する。
• 球面フィルタを用いた球面自己アテンション機構を構築し、ユークリッド空間における連続的フィルタを球面調和空間へ一般化することで、SO(3)等長性を保持する。
• SPHCsを用いたアテンションによるメッセージパッシングを採用し、長距離にわたり情報を伝搬させることで、非局所的電子的効果を捉える。
• エネルギーと力の両方の損失関数を組み合わせて訓練し、エネルギーの積分定数は、力のみが監視対象の場合、最小二乗法によるフィッティングで算出する。
• 訓練中にSPHCsを更新することで、分子構造内の高次相関を反映する最適な幾何的表現をモデルが学習可能となる。

実験結果

リサーチクエスチョン

• RQ1 学習可能な注目ベースの球面調和座標表現が、分子系における非局所的電子的効果のモデル化を改善できるか？
• RQ2 SPHCsによる幾何的情報と原子特徴の分離が、機械学習力場におけるデータ効率と一般化性能を向上させるか？
• RQ3 球面自己アテンション機構が、任意の長さスケールにわたる分子系における長距離相関を効果的に捉えることができるか？
• RQ4 高次幾何的相関をSPHCsで組み込むことで、モデルの複雑さを低減しつつ予測精度を維持または向上させられるか？
• RQ5 So3kratesは、既存のSOTAモデルと比較して、パラメータ効率、学習速度、推論速度のどの面でも優れているか？

主な発見

• So3kratesはQM7-X250およびMD17を含む複数のベンチマークでSOTA性能を達成し、先行SOTAモデルと比較してモデルパラメータが0.25–0.4倍に削減された。
• 他のアテンションベースのモデルと比較して、学習時で6–14倍、推論時で2–11倍の高速化を達成したが、依然として高い精度を維持している。
• So3kratesは、キュムレンの水素回転子における長距離相関を的確に捉えており、局所モデルとは異なり、ab-initioの真値を忠実に再現している。
• SPHCsによる高次幾何的相関の組み込みにより、限られたデータでも優れた性能を示すなど、データ効率と一般化性能が向上した。
• アーキテクチャにSPHCsの学習を組み込むことで、固定または手作業で作成された幾何的特徴を用いるモデルよりも、効果的な非局所的アテンションが実現され、性能が優れていることが示された。
• 力のみで訓練された場合、エネルギーの積分定数は最小二乗法によるフィッティングで算出され、追加の監視なしに正確なエネルギー予測が保証される。