[論文レビュー] Point Edge Transformer model trained on highly distorted methane configurations

Sergey N. Pozdnyakov, Michele Ceriotti|arXiv (Cornell University)|May 30, 2023

Image Processing and 3D Reconstruction被引用数 15

ひとこと要約

この論文は、非等変バックボーンの上に正確な回転等変性を課す Equivariant Coordinate System Ensemble (ECSE) を備えた Point Edge Transformer (PET) を提案し、複数の原子レベルベンチマークで最先端の結果を達成します。

ABSTRACT

Point clouds are versatile representations of 3D objects and have found widespread application in science and engineering. Many successful deep-learning models have been proposed that use them as input. The domain of chemical and materials modeling is especially challenging because exact compliance with physical constraints is highly desirable for a model to be usable in practice. These constraints include smoothness and invariance with respect to translations, rotations, and permutations of identical atoms. If these requirements are not rigorously fulfilled, atomistic simulations might lead to absurd outcomes even if the model has excellent accuracy. Consequently, dedicated architectures, which achieve invariance by restricting their design space, have been developed. General-purpose point-cloud models are more varied but often disregard rotational symmetry. We propose a general symmetrization method that adds rotational equivariance to any given model while preserving all the other requirements. Our approach simplifies the development of better atomic-scale machine-learning schemes by relaxing the constraints on the design space and making it possible to incorporate ideas that proved effective in other domains. We demonstrate this idea by introducing the Point Edge Transformer (PET) architecture, which is not intrinsically equivariant but achieves state-of-the-art performance on several benchmark datasets of molecules and solids. A-posteriori application of our general protocol makes PET exactly equivariant, with minimal changes to its accuracy.

研究の動機と目的

• 原子スケールの機械学習における正確な回転対称性の必要性を動機づけ、シミュレーションのアーティファクトを回避する。
• どのバックボーンアーキテクチャにも事後的に正確な回転等変性を付与する一般的な対称化プロトコル（ECSE）を提案し、他の対称性を犠牲にしない。
• エッジ中心のトランスフォーマー ベースのメッセージパッシングを活用して表現力を高めつつ、局所性と効率性を維持する Point Edge Transformer (PET) を開発する。
• 歪んだメタン CH4 の配置、低分子、周期的結晶など、多様な原子種を含むデータセットに対して PET の最先端性能を示す。

提案手法

• ECSE（Equivariant Coordinate System Ensemble）を導入し、カットオフ内の隣接ペアによって定義されるすべての座標系で予測を平均化し、滑らかさを確保する重みを付与する。
• 原子環境ごとに適応的な内部カットオフ R_in(A_i) を定義して座標系の数を制限しつつ、少なくとも1つの適切な基準系を保証する。
• 等変性を持たないバックボーンに事後的に ECSE を適用して、他の不変性（置換、並進、滑らかさ）を保ちながら正確な回転等変性を得る。
• Point Edge Transformer (PET) は中心原子と隣接原子のトークンを処理し、並置変換に不変な自己注意を用いたエッジ機能中心のトランスフォーマー型メッセージパッシングアーキテクチャを提示する。
• 注意機構を滑らかなカットオフ関数 f_c で実現し、近傍が現れたり消えたりしても微分可能性と連続性を保つ。
• PET を ECSE の有無でベンチマークしたデータセットには、液体水、COLL、MnO、HME21、HEA、QM9 dipole を含む。

実験結果

リサーチクエスチョン

• RQ1 任意のバックボーンモデルに対して正確な回転等変性を適用できる一般的な対称化プロトコルは、他の対称性を劣化させずに適用可能か。
• RQ2 ECSE によって強化された PET アーキテクチャは、液体・分子・固体の多様なデータセットを含む多様な原子系データにおいて最先端の精度を達成するか。
• RQ3 モデル深さ/ブロックを増やすことによる性能向上はどの程度か、ECSE は推論コストと精度にどう影響するか。
• RQ4 PET は標準的なベンチマークや歪んだ構成に対して、NEQUIP、MACE などの先行の等変モデルと比べてどうか。

主な発見

データセット	MAE f (COLL)	MAE E /at. (COLL)	RMSE f (MnO)	RMSE E /at. (MnO)	MAE f (HME21)	MAE E /at. (HME21)	MAE f (HEA)	MAE E /at. (HEA)
SOTA	26.4 [17]	47 [86]	125 [88]	1.11 [88]	138 [89]	15.7 [89]	190 [90]	10 [90]
PET ( y0 )	23.1	12.0	22.7	0.312	140.5 ± 2.0	17.8 ± 0.1	60.2	1.87
PET ( yS )	23.1	11.9	22.7	0.304	141.6 ± 1.9	17.8 ± 0.1	60.1	1.87
PET (ens.)				128.5	16.8

• PET は液体水の関連ベンチマークにおいて、状態最先端の等変モデル NEQUIP に対する誤差を約 30% 減少させる。
• 半径 4.25 Å のカットオフで、液体水様データ中の力の MAE が 14.4 meV/Å となり、長距離相互作用の扱いが強いことを示す。
• COLL データセットでは、PET は力の MAE を約 13% 改善し、結合エネルギーの誤差を以前の最良法と比べて 4 倍低減。
• MnO、HME21、HEA データセットでは、y0/yS バリアントの PET は競合的または最先端に近い性能を達成し、報告された指標の多くで PET yS は PET y0 とほぼ同等の精度を示す。
• PET モデルのアンサンブル（ens.）は誤差をさらに低減し、単一モデルよりもいくつかの指標で顕著な改善を示す（例：アンサンブル行の報告値 16.8 など）。
• ECSE は正確な回転等変性（共変出力）を最小限の影響で実現でき、厳密な対称性要件を持つ原子スケールシミュレーションに PET を適用可能にする。