[論文レビュー] Equivariant message passing for the prediction of tensorial properties and molecular spectra

Kristof T. Schütt, Oliver T. Unke+1 more|arXiv (Cornell University)|Feb 5, 2021

Machine Learning in Materials Science被引用数 265

ひとこと要約

本論文は、分子グラフの回転的等価性を持つメッセージ伝搬ニューラルネットワークである PaiNN を提案し、スカラーおよびテンソル的特性を予測し、分子スペクトルの効率的なシミュレーションを可能にします。小型モデルで最先端の性能を達成し、推論を高速化します。

ABSTRACT

Message passing neural networks have become a method of choice for learning on graphs, in particular the prediction of chemical properties and the acceleration of molecular dynamics studies. While they readily scale to large training data sets, previous approaches have proven to be less data efficient than kernel methods. We identify limitations of invariant representations as a major reason and extend the message passing formulation to rotationally equivariant representations. On this basis, we propose the polarizable atom interaction neural network (PaiNN) and improve on common molecule benchmarks over previous networks, while reducing model size and inference time. We leverage the equivariant atomwise representations obtained by PaiNN for the prediction of tensorial properties. Finally, we apply this to the simulation of molecular spectra, achieving speedups of 4-5 orders of magnitude compared to the electronic structure reference.

研究の動機と目的

• 分子グラフにおける方向情報を伝搬させる際の回転不変表現の限界を動機づける。
• データ効率と表現力を向上させるため、回転等価なメッセージ伝搬と PaiNN アーキテクチャを提案する。
• 等価表現を用いてテンソル性質（例：偶極子モーメント、分極率）の予測を可能にする。
• 電子構造計算に比べて大幅な速度向上を示す RPMD シミュレーションを通じた分子スペクトルへの応用を示す。

提案手法

• 回転的に等価な関数でスカラー特徴とベクトル特徴が更新される等価メッセージ伝搬を開発する。
• 不変なスカラー信号には半径基底関数フィルターを用いた連続フィルタ畳み込みを、等価信号にはベクトル対応の畳み込みを用いる。
• 方向情報を伝搬するため、残差アーキテクチャを介してスカラー更新とベクトル更新を結びつける。
• 原子ごとの潜在特徴と局所ベクトル特徴を組み合わせたランク-1のテンソル分解を通じてテンソル性質を予測する。
• エネルギー、力、偶極子モーメント、分極率を対象に、重み減衰と学習率スケジューリングを用いた Adam でエンドツーエンドに学習する。

実験結果

リサーチクエスチョン

• RQ1 回転的に等価なメッセージ伝搬は、標準的な分子特性ベンチマークにおいて、さまざまなデータレジームで不変 GNN を上回ることができるのか？
• RQ2 等価表現は、過度なモデルの複雑さなしに、偶極子や分極率といったテンソル性質を正確に予測することを可能にするか？
• RQ3 PaiNN は、第一原理法と比較して高速で正確な軌道を用いたIR/ラマン分光の分子スペクトルシミュレーションを加速できるか？

主な発見

対象	単位	SchNet	PhysNet	DimeNet++	Cormorant	L1Net	PaiNN
μ	D	0.033	0.053	0.030	0.038	0.043	0.012
α	a0^3	0.235	0.062	0.044	0.085	0.088	0.045
ε_HOMO	meV	41	32.9	24.6	34	46.0	27.6
ε_LUMO	meV	34	24.7	19.5	38	34.6	20.4
Δε	meV	63	42.5	32.6	38	67.5	45.7
<R^2>	a0^2	0.073	0.765	0.331	0.961	0.354	0.066
ZPVE	meV	1.7	1.39	1.21	2.03	1.56	1.28
U_0	meV	14	8.15	6.32	22	13.46	5.85
U	meV	19	8.34	6.28	21	13.83	5.83
H	meV	14	8.42	6.53	21	14.36	5.98
G	meV	14	9.40	7.56	20	13.99	7.35
c_v	cal/mol/K	0.033	0.028	0.023	0.026	0.031	0.024

• PaiNN は QM9 における複数のスカラー特性で最先端または競合的な結果を達成し、パラメータ数は約0.6Mと、DimeNet++（約1.8M）などのベースラインより小さい。
• PaiNN は QM9 で DimeNet++ に対して推論時間を70%以上削減（50分子のバッチで13 ms 対 45 ms）。
• アブレーション研究は、等価アーキテクチャの全要素が精度に寄与し、ベクトル特徴畳込みとスカラー–ベクトル結合が顕著な向上をもたらすことを示している。
• 等価表現は、原子中心の分 polarization 領分と全体的な幾何学を活用して、偶極子モーメントや分極率などのテンソル性質を正確に予測することを可能にする。
• PaiNN は赤外・ラマンスペクトルの RPMD ベースの高速シミュレーションを実現し、電子構造参照よりも大幅に速度を向上させ、エタノール（RPMD）とアスピリンデータセットのスペクトルを実験と緊密に一致させている。
• 等価特徴は小さなカットオフでも方向情報を捉える能力を向上させ、局所近傍を超えた効率的な伝搬を可能にする。