[論文レビュー] Molecular Geometry-aware Transformer for accurate 3D Atomic System modeling

Zheng Yuan, Yaoyun Zhang+4 more|arXiv (Cornell University)|Feb 2, 2023

Machine Learning in Materials ScienceMaterials Science被引用数 3

ひとこと要約

Moleformer は、相対的な距離と角度に基づく回転および並進不変な空間符号化を用いて、原子、共有結合、および非共有原子対を幾何学的に注意を向けた Transformer アーキテクチャを提案する。これは、OC20 におけるエネルギー予測のリラクゼーションにおいて最先端の性能を達成し、QM9 においても競争力ある結果を示し、高次元の幾何的相互作用を明示的にモデル化することで、従来の GNN や Transformer モデルを上回る。

ABSTRACT

Molecular dynamic simulations are important in computational physics, chemistry, material, and biology. Machine learning-based methods have shown strong abilities in predicting molecular energy and properties and are much faster than DFT calculations. Molecular energy is at least related to atoms, bonds, bond angles, torsion angles, and nonbonding atom pairs. Previous Transformer models only use atoms as inputs which lack explicit modeling of the aforementioned factors. To alleviate this limitation, we propose Moleformer, a novel Transformer architecture that takes nodes (atoms) and edges (bonds and nonbonding atom pairs) as inputs and models the interactions among them using rotational and translational invariant geometry-aware spatial encoding. Proposed spatial encoding calculates relative position information including distances and angles among nodes and edges. We benchmark Moleformer on OC20 and QM9 datasets, and our model achieves state-of-the-art on the initial state to relaxed energy prediction of OC20 and is very competitive in QM9 on predicting quantum chemical properties compared to other Transformer and Graph Neural Network (GNN) methods which proves the effectiveness of the proposed geometry-aware spatial encoding in Moleformer.

研究の動機と目的

• 既存の Transformer モデルが 3D 原子系モデリングにおいて、原子のみを入力として扱い、結合、角度、非共有相互作用を明示的にモデル化できないという限界を解決すること。
• 結合角や二面角といった幾何的要因を注意メカニズムに直接組み込むことで、分子エネルギーおよび性質予測の精度を向上させること。
• 周期的境界条件補正とデータ拡張を導入することで、分布外の分子系における一般化性能を向上させること。
• 結合長に対する幾何的監視を提供するエッジレベルの補助損失がエネルギー予測の改善に寄与することを示すこと。

提案手法

• Moleformer は、従来の Transformer が原子のみを入力とするのに対し、原子とエッジ（結合および非共有原子対）を入力として用いる。
• ノードとエッジ間の相対的な距離と角度を計算する幾何学的に注意を向けた空間符号化を導入し、回転および並進不変性を確保する。
• 各層でノードおよびエッジの両方に対して自己注意を適用することで、3体および4体相互作用（例：結合角、二面角）を明示的にモデル化可能となる。
• 周期的構造を有する系、特に触媒シミュレーションにおいて精度を向上させるために、新規の周期的境界条件（PBC）補正を適用する。
• OC20 データセットにおける分布外の吸着種に対して一般化性能を向上させるために、ノイズを含むノードデータ拡張を用いる。
• 予測された結合長に対する補助的なエッジレベル損失を導入し、エネルギー推定性能を向上させる。

Figure 1: The overview of Moleformer. (a) A toy example of a catalyst system where \ch CH3 is the adsorbate and the surface is composed of copper atoms. Moleformer prefers edges between adosrbate-adsorbate atoms and adsorbate-surface atoms as model inputs. (b) Architecture of Moleformer: Nodes and s

実験結果

リサーチクエスチョン

• RQ1 Transformer アーキテクチャにおいて結合および非共有原子対を明示的にモデル化することで、原子のみを入力とするモデルと比較して、3D 原子系のエネルギー予測性能が向上するか？
• RQ2 相対的な距離と角度に基づく幾何学的に注意を向けた空間符号化が、分子エネルギーおよび性質予測タスクの性能に与える影響は何か？
• RQ3 周期的境界条件補正を組み込むことで、周期的構造を有する触媒系におけるエネルギー予測性能がどの程度向上するか？
• RQ4 結合長予測のためのエッジレベル補助損失を追加することで、モデルのエネルギー推定能力が向上するか？
• RQ5 Moleformer は、特に OC20 データセットにおける希少な吸着種を含む分布外の分子に対して、どの程度一般化できるか？

主な発見

• Moleformer は、初期状態からリラックス状態へのエネルギー予測において OC20 データセットで最先端の性能を達成し、ドメイン内スプリットにおいてエネルギー MAE が 413 meV であった。
• 周期的境界条件補正を導入することで、IS2RS における位置 MAE が低下し、IS2RE スプリットにおいてエネルギー MAE が 0.014 eV 減少し、EWT が 0.7% 向上した。
• ノイズを含むノードデータ拡張を導入しないと、分布外の吸着種スプリットにおけるエネルギー MAE が顕著に増加し、データ拡張による一般化性能の向上が示された。
• エッジレベル補助損失を削除すると、ドメイン内および分布外スプリットの両方でエネルギー MAE が上昇し、幾何的監視の有効性が確認された。
• QM9 データセットにおいて、Moleformer は分子軌道エネルギーを含む量子化学的性質の予測において、最近の最先端手法（SEGNN や EQGAT）を上回る競争力のある結果を達成した。
• アブレーションスタディにより、PBC 補正、データ拡張、エッジレベル損失のすべての要素が性能向上に寄与しており、特にエッジレベル損失が最大の向上をもたらしたことが確認された。

Figure 2: PBC Correction. (a)-(c) An example of a catalyst system where \ch CH3 is the adsorbate and the surface is composed of copper atoms viewed along the <110>orientation. The dotted lines show the borders of the unit cell. Grey-colored small atom is a hydrogen atom that goes across the unit cel