[論文レビュー] TensorNet: Cartesian Tensor Representations for Efficient Learning of Molecular Potentials

Guillem Simeon, Gianni De Fabritiis|arXiv (Cornell University)|Jun 10, 2023

Machine Learning in Materials Science被引用数 17

ひとこと要約

TensorNetは、Cartesian rank-2テンソルを用いたO(3)-等価性ニューラルネットワークを導入し、分子ポテンシャルを効率的に学習する。これにより、いくつかの球テンソルモデルよりもはるかに少ないパラメータと低い計算量で最先端の精度を達成するとともに、エネルギーと力だけでなくベクトルおよびテンソルの分子量を予測することを可能にする。

ABSTRACT

The development of efficient machine learning models for molecular systems representation is becoming crucial in scientific research. We introduce TensorNet, an innovative O(3)-equivariant message-passing neural network architecture that leverages Cartesian tensor representations. By using Cartesian tensor atomic embeddings, feature mixing is simplified through matrix product operations. Furthermore, the cost-effective decomposition of these tensors into rotation group irreducible representations allows for the separate processing of scalars, vectors, and tensors when necessary. Compared to higher-rank spherical tensor models, TensorNet demonstrates state-of-the-art performance with significantly fewer parameters. For small molecule potential energies, this can be achieved even with a single interaction layer. As a result of all these properties, the model's computational cost is substantially decreased. Moreover, the accurate prediction of vector and tensor molecular quantities on top of potential energies and forces is possible. In summary, TensorNet's framework opens up a new space for the design of state-of-the-art equivariant models.

研究の動機と目的

• 分子シミュレーションと創薬における効率的で高精度な原子間ポテンシャルの必要性を動機づける。
• Cartesian rank-2テンソルを Atomic 埋め込みに用いたO(3)-等価アーキテクチャを提案する。
• 高階の球面モデルと比較した計算効率とパラメータ効率を示す。
• エネルギー、力、テンソル量を予測できるモデルの能力を示す。
• より少ないパラメータで最先端の性能を確立するため、TensorNetを多様なデータセットで評価する。

提案手法

• 原子を全ての3x3 Cartesianテンソルとして表現し、それを标量、ベクトル、テンソル成分（I, A, S）に分解する。
• 相対位置と原子番号から得られるエッジベースの初期化を用いて、エッジごとのテンソル特徴を形成する。
• 不変ノルムを計算し、それを用いてパリティを保つ行列演算による原子ごとの更新を生み出す。
• 層を通じてO(3)等価性とパリティを保持する行列積風の相互作用を適用する。
• 最終テンソルノルムから導出された原子ごとの出力の和として全体エネルギーを予測する。
• テンソル成分から抽出することでベクトル出力とテンソル出力を有効にし、反対称成分を含む。

Figure 1: Key steps, from top to bottom, in the embedding and interaction modules for some central atom $i$ and neighbors $j$ and $k$ found within the cutoff radius. a) Relative position vectors are used to initialize edge-wise tensor components, modified using edge-wise invariant functions, and sum

実験結果

リサーチクエスチョン

• RQ1 Cartesian rank-2テンソル表現は、高階の球面モデルよりも少ないパラメータで分子ポテンシャルの最先端精度を達成できるか。
• RQ2 単純な行列積と不可約分解を用いるO(3)等価テンソル演算は、エネルギー・力・テンソル特性の正確な予測に十分か。
• RQ3 TensorNetは多様な分子データセット（QM9, rMD17, SPICE/ANI1x/COMP6）において既存のベースラインと比較してどう機能するか。
• RQ4 予測精度に対する等価性、相互作用設計、カットオフの影響はどうか。
• RQ5 TensorNetはスカラー、ベクトル、テンソルの分子特性を同時に信頼して予測できるか。

主な発見

Property	DimeNet++	ET	PaiNN	Allegro	MACE	TensorNet
U0	6.3	6.3	5.9	4.7	4.1	3.9(1)
U	6.3	6.3	5.7	4.4	4.1	3.9(1)
H	6.5	6.5	6.0	4.4	4.7	4.0(1)
G	7.6	7.6	7.4	5.7	5.5	5.7(1)

• TensorNetはQM9のU0, U, HでAllegroとMACEを上回り、3.9(1) meVスタイルの参照値でありながらAllegroのパラメータの23%しか使用していない状態で最先端の精度を達成した。
• rMD17では、TensorNetは1–2層の相互作用で競争力のあるエネルギー・力のMAEを達成し、他の球面モデルよりはるかに少ないパラメータ（0.535–0.770M程度まで）である。
• TensorNetは軽量な1LモデルでANI1x/COMP6の一般化に強さを示し、複数のベンチマークでいくつかのベースラインより低いエネルギーと力のMAEを提供。
• エタノールの真空条件下で、TensorNetはFieldSchNetとPaiNNよりエネルギーと力のMAEを約2倍改善し、双極子と極化率予測も正確。
• アブレーション研究は、完全なO(3)等価性と相互作用積が精度を劇的に改善することを示し、0L対1L/2Lのトレードオフはカットオフと分子サイズに依存する。
• TensorNetはGPU上で推論と訓練の速度が良好で、数百原子クラスの分子にも対応し、より大きな系では低次の等価性トランスフォーマーと競合する効率を維持する。

(b) Interaction and node update