[論文レビュー] TorchMD-NET: Equivariant Transformers for Neural Network based Molecular Potentials

Philipp Thölke, Gianni De Fabritiis|arXiv (Cornell University)|Feb 5, 2022

Machine Learning in Materials Science被引用数 58

ひとこと要約

TorchMD-NET は分子の量子力学的特性を予測するための等変換Transformerを導入し、QM9、MD17、ANI-1 のベンチマークで最先端の精度と効率を達成し、アテンション分析による洞察を提供する。

ABSTRACT

The prediction of quantum mechanical properties is historically plagued by a trade-off between accuracy and speed. Machine learning potentials have previously shown great success in this domain, reaching increasingly better accuracy while maintaining computational efficiency comparable with classical force fields. In this work we propose TorchMD-NET, a novel equivariant transformer (ET) architecture, outperforming state-of-the-art on MD17, ANI-1, and many QM9 targets in both accuracy and computational efficiency. Through an extensive attention weight analysis, we gain valuable insights into the black box predictor and show differences in the learned representation of conformers versus conformations sampled from molecular dynamics or normal modes. Furthermore, we highlight the importance of datasets including off-equilibrium conformations for the evaluation of molecular potentials.

研究の動機と目的

• 分子のエネルギーと力を予測する回転に対して等変な Transformer（ET）アーキテクチャを開発する。
• 標準的な量子化学ベンチマーク（QM9、MD17、ANI-1）で高い精度と計算効率を達成する。
• アテンション重みの分析およびデータレジーム間の表現比較を通じてモデルの解釈性を提供する。
• ニューラルネットワークポテンシャルを評価する際の非平衡構造の重要性を示す。

提案手法

• 原子種と局所近傍を各原子につき2つの学習済み埋め込みで埋め込む。
• RBF展開されたエッジデータを介して原子間距離を組み込む修正済みのマルチヘッド注意機構を用いる。
• 注意値経路に連続フィルターグラフ畳込みを組み込み、特徴に距離情報を注入する。
• ゲート付き等変ブロックを適用してスカラーおよびベクトルの更新を計算し、座標に対するエネルギーの微分による力の予測を可能にする。
• Adamを用いた平均二乗誤差で学習し、学習率ウォームアップとプラトーに基づく減衰を適用する；エネルギーと力の予測時には指数平滑化を適用する。

実験結果

リサーチクエスチョン

• RQ1 等変換トランスフォーマーは QM9、MD17、および ANI-1 のベンチマークで既存のニューラルネットワークポテンシャルを上回ることができるのか？
• RQ2 注意機構にエッジ（距離）情報を組み込むことは予測精度と効率にどのような影響を与えるか？
• RQ3 性能に最も影響を与える構造要素は何か（例: 隣接埋め込み、更新層、ベクトル特徴量など）？
• RQ4 注意重みによって静的および非平衡構造を含む意味のある化学的・動的構造が明らかになるか？
• RQ5 エネルギーと力の予測における訓練データの多様性（非平衡構造）の重要性はどの程度か？

主な発見

• 等変換トランスフォーマー（ET）は QM9 のターゲットに対して不変・共変アーキテクチャと比較して最先端の MAE を達成し、特化した等変モデルと対等またはそれを上回る。
• MD17 では、多くの分子においてエネルギーと力の精度でトップを獲得し、従来法を上回ることが多い。アブレーションでは隣接埋め込みと更新層が性能に大きく寄与することを示している。
• ANI-1 では大規模な非平衡データセットで強い性能を示し、テスト MAE で複数のベースラインを上回る。
• アテンション分析は、訓練データのタイプ（静的 QM9 対 dynamic ANI-1/MD17）がアテンションの焦点を形作ることを示し、動的データセットでは水素原子がより注目され、静的データセットではコア原子が強調される。
• 非平衡構造は正確なエネルギー予測にとって重要である。アテンションパターンとアブレーション結果から示唆される。
• 小型の ET バリアントも最先端モデルと競争力を保ち、パラメータ数に比例したデータ効率とスケーラビリティの有利性を示している。