[論文レビュー] Spatial Attention Kinetic Networks with E(n)-Equivariance

Yuanqing Wang, John D. Chodera|arXiv (Cornell University)|Jan 21, 2023

Machine Learning in Materials Science被引用数 11

ひとこと要約

SAKEは、E(n)-不変/共変な単純な空間アテンション形を導入し、それを用いて多体系の不変および共変タスクに対して高速で正確なSAKEモデルを構築します。いくつかのベースラインを上回りつつ、はるかに高速です。

ABSTRACT

Neural networks that are equivariant to rotations, translations, reflections, and permutations on n-dimensional geometric space have shown promise in physical modeling for tasks such as accurately but inexpensively modeling complex potential energy surfaces to guiding the sampling of complex dynamical systems or forecasting their time evolution. Current state-of-the-art methods employ spherical harmonics to encode higher-order interactions among particles, which are computationally expensive. In this paper, we propose a simple alternative functional form that uses neurally parametrized linear combinations of edge vectors to achieve equivariance while still universally approximating node environments. Incorporating this insight, we design spatial attention kinetic networks with E(n)-equivariance, or SAKE, which are competitive in many-body system modeling tasks while being significantly faster.

研究の動機と目的

• 物理系のニューラルモデルに幾何学を意識した帰納バイアスを組み込む必要性を動機づける。
• E(n)-不変性/共変性を実現する、単純でニューラルにパラメータ化された空間的アテンション機構を提案する。
• n体ダイナミクスにおける速度と位置の更新に空間アテンションを活用するSAKEアーキテクチャを開発する。
• 標準ベンチマーク全体で、不変タスク（ポテンシャルエネルギー）と共変タスク（ダイナミクス）において、競争力のある性能と向上した速度を示す。

提案手法

• 空間アテンションを、エッジベクトルのニューラルにパラメータ化された線形結合を集約し、それらのノルムを取って不変な記述子を形成する関数として定義する。
• エッジ埋め込み上の関数 lambda_i を用いたマルチヘッド風の設定と、エッジベクトル上の共変関数 f に続き、結合されたノルムを処理する mu を用いる。
• ノード埋め込みと距離情報をラジアル基底関数で結合した SchNet風エッジ埋め込みを用いて対相互作用をエンコードする。
• 虚偽の速度更新とオイラー風の位置更新を組み込み、ほぼハミルトニアン積分器を作成し、ダイナミクスの E(n)-共変性を保証する。
• 相互作用の各異方性を促進するために、セマンティックな注意と距離遮断ベースの注意を含める。
• スペーシャルアテンションが性能の主な寄与であることを示すアブレーションに基づく正当化を提供する。

実験結果

リサーチクエスチョン

• RQ1 局所ノード環境に対して、単純なニューラルにパラメータ化された空間アテンション機構はE(n)-不変性/共変性を達成できるか？
• RQ2 SAKEは不変タスク（例：ポテンシャルエネルギー予測）と共変タスク（例：ダイナミクス予測）で、精度と速度の観点から最先端モデルと比べてどうか？
• RQ3 空間アテンション、セマンティックアテンション、速度/位置更新がモデル性能に与える影響は何か？

主な発見

• SAKEはMD17、ISO17、QM9、および合成N体動力学ベンチマーク全体で、不変タスクおよび共変タスクにおいて競争力のあるまたは優れた性能を達成する。
• SAKEは一般に、いくつかの最先端共変モデルよりも大幅に少ないトレーニングおよび推論時間を必要とする。
• アブレーション研究は、空間アテンションがエネルギーと力の予測の改善に寄与し、追加の成分がさらなる改善を提供することを示す。
• SAKEは未見の化学空間（ISO17の未知分子）への強い外挿性能と有利な軌道予測結果を示す。
• 球面調和函数を用いる手法と比較して、SAKEはさまざまなタスクで競争力のある精度を保ちつつ、より高速な代替手段を提供する。