[論文レビュー] Equiformer: Equivariant Graph Attention Transformer for 3D Atomistic Graphs

Yi-Lun Liao, Tess Smidt|arXiv (Cornell University)|Jun 23, 2022

Machine Learning in Materials Science被引用数 63

ひとこと要約

Equiformerは、SE(3)/E(3)-等変グラフトランスフォーマーで、不可約表現と等変グラフ注意機構を統合し、3D原子グラフ上の量子特性を予測する。競争力のある精度と向上した効率性を実現。

ABSTRACT

Despite their widespread success in various domains, Transformer networks have yet to perform well across datasets in the domain of 3D atomistic graphs such as molecules even when 3D-related inductive biases like translational invariance and rotational equivariance are considered. In this paper, we demonstrate that Transformers can generalize well to 3D atomistic graphs and present Equiformer, a graph neural network leveraging the strength of Transformer architectures and incorporating SE(3)/E(3)-equivariant features based on irreducible representations (irreps). First, we propose a simple and effective architecture by only replacing original operations in Transformers with their equivariant counterparts and including tensor products. Using equivariant operations enables encoding equivariant information in channels of irreps features without complicating graph structures. With minimal modifications to Transformers, this architecture has already achieved strong empirical results. Second, we propose a novel attention mechanism called equivariant graph attention, which improves upon typical attention in Transformers through replacing dot product attention with multi-layer perceptron attention and including non-linear message passing. With these two innovations, Equiformer achieves competitive results to previous models on QM9, MD17 and OC20 datasets.

研究の動機と目的

• 3D幾何学的帰納バイアスを取り入れることにより、Transformersを3D原子グラフへ一般化できるよう動機づけ、実現する。
• Transformerの操作を等変対応のものに置換し、テンソル積を追加することでEquiformerを開発する。
• 内容情報と幾何情報を組み合わせる等変グラフ注意機構を導入する。
• QM9、MD17、OC20で評価し、最新モデルに対する競争力のある性能を示す。

提案手法

• 不可約表示特徴上のテンソル積を含むSE(3)/E(3)-等変操作を用いてTransformerの構成要素を置換する。
• L次元ベクトルの組み合わせを効率的に行い、パラメータ増加を制御する深さ方向テンソル積を用いる。
• スカラー（タイプ-0）特徴にMLPLアテンションを用い、より高次の特徴には非線形メッセージ伝搬を用いる新規の等変グラフ注意機構を採用する。
• 入力を原子・エッジ次数の埋め込みで埋め込み、距離条件付きの重みのためにラジアル基底・ラジアル関数を組み込む。
• 非スカラーの不可約表示にはゲート型の非線形活性化を適用して等変性を保持し、計算複雑性を管理するために深さ方向の手法を用いる。

実験結果

リサーチクエスチョン

• RQ1 SE(3)/E(3)-等変特徴を備えたとき、Transformersは3D原子グラフへ効果的に一般化できるか？
• RQ2 提案された等変グラフ注意は、3Dグラフにおいて標準のドット積注意より効果的か？
• RQ3 等変操作を用いた最小限の改変を加えたTransformersは、QM9、MD17、OC20で従来のSE(3)/E(3)-等変アーキテクチャと同等またはそれを上回るか？
• RQ4 Equiformerで高次の不可約表示（Lmax）を使用する場合の訓練時間とスケーラビリティのトレードオフは何か？

主な発見

• Lmax=2-3のEquiformerは、QM9、MD17、およびOC20データセット全体で従来モデルと競合する結果を達成する。
• QM9では、EquiformerはNequIPやTorchMD-NETを含むいくつかのベースラインよりも全体的に12の回帰タスクで良好な結果を示す。
• MD17では、Lmaxが高いEquiformerは複数の分子でエネルギーおよび力のMAEを改善し、TorchMD-NETやNequIPなどの代替手法より優れている。
• OC20では、IS2REデータで訓練した場合、最先端SEGNNやGraphormerよりIS2RE性能を向上させ、従来モデルと比較して訓練時間を2.3倍から15.5倍削減できる。
• Equiformerの等変グラフ注意（MLPベースの注意と非線形メッセージ）はドット積注意より表現力が高く、等変トランスフォーマーの性能を向上させる。
• Lmaxを増やすとOC20とMD17の結果が一般に改善されるが、いくつかのケースでは過学習を引き起こす可能性がある。