[論文レビュー] Generalizing Denoising to Non-Equilibrium Structures Improves Equivariant Force Fields

Yi-Lun Liao, Tess Smidt+2 more|arXiv (Cornell University)|Mar 14, 2024

Computer Graphics and Visualization Techniques被引用数 4

ひとこと要約

本稿では、原子的機械学習における等変換力場の性能向上を目的として、DeNS（ノイズ除去非平衡構造）を補助的自己教師あり学習タスクとして提案する。非平衡原子構造のノイズ除去時に力を条件付け信号として組み込むことで、元来不適切な定式化である問題を適切に定式化し、OC20、OC22、MD17の各ベンチマークで最小限の学習コストで顕著な性能向上を達成できる。

ABSTRACT

Understanding the interactions of atoms such as forces in 3D atomistic systems is fundamental to many applications like molecular dynamics and catalyst design. However, simulating these interactions requires compute-intensive ab initio calculations and thus results in limited data for training neural networks. In this paper, we propose to use denoising non-equilibrium structures (DeNS) as an auxiliary task to better leverage training data and improve performance. For training with DeNS, we first corrupt a 3D structure by adding noise to its 3D coordinates and then predict the noise. Different from previous works on denoising, which are limited to equilibrium structures, the proposed method generalizes denoising to a much larger set of non-equilibrium structures. The main difference is that a non-equilibrium structure does not correspond to local energy minima and has non-zero forces, and therefore it can have many possible atomic positions compared to an equilibrium structure. This makes denoising non-equilibrium structures an ill-posed problem since the target of denoising is not uniquely defined. Our key insight is to additionally encode the forces of the original non-equilibrium structure to specify which non-equilibrium structure we are denoising. Concretely, given a corrupted non-equilibrium structure and the forces of the original one, we predict the non-equilibrium structure satisfying the input forces instead of any arbitrary structures. Since DeNS requires encoding forces, DeNS favors equivariant networks, which can easily incorporate forces and other higher-order tensors in node embeddings. We study the effectiveness of training equivariant networks with DeNS on OC20, OC22 and MD17 datasets and demonstrate that DeNS can achieve new state-of-the-art results on OC20 and OC22 and significantly improve training efficiency on MD17.

研究の動機と目的

• 既存の軌道に含まれる未利用の非平衡構造を活用することで、第一原理的原子的データセットのスケール制限を緩和する。
• 非平衡構造のノイズ除去が不適切な定式化である問題を克服する。非ゼロの力が存在するため、複数の目的構造が存在する。
• 分子系におけるエネルギーおよび力の予測のための等変換ニューラルネットワークの一般化性能と効率性を向上させる。
• DeNSによる補助的学習が、事前学習と同等の性能向上を達成するが、はるかに低い計算コストで実現できることを示す。

提案手法

• 非平衡構造の3次元原子座標にノイズを加えることで、汚れた入力を生成する。
• 元の構造の元の力を条件付け信号として用いることで、再構成の目的を一意に定義する。
• 汚れた構造と元の力を入力として与えられたモデルが、元の非平衡構造を回復させるために必要なノイズベクトルを予測する。
• 力場モデルの学習中に、エネルギー・力予測という主タスクに加えて、DeNSを補助損失として統合する。
• 不変表現（irreps）を介してベクトル値特徴（力）を自然に処理できる等変換メッセージパッシングネットワーク（例：Equiformer）を用いる。
• 球面調和関数による力符号化を適用し、力をネットワークの不変表現空間に射影することで、等変換性を保証する。

Figure 1 : Illustration of denoising equilibrium and non-equilibrium structures. In this figure, we relax a non-equilibrium structure ( red point) and form a local relaxation trajectory (black dotted arrow). All the points along the trajectory except the blue point are non-equilibrium structures.

実験結果

リサーチクエスチョン

• RQ1 非平衡原子構造のノイズ除去を、自己教師あり事前学習タスクとして実行可能かつ有効なものとすることができるか？
• RQ2 元の力を条件付け信号として組み込むことで、非平衡構造のノイズ除去における不適切な定式化が解消されるか？
• RQ3 DeNSは、OC20、OC22、MD17といった標準ベンチマークにおいて、等変換力場の性能をどの程度向上させるか？
• RQ4 事前学習やアーキテクチャスケーリングと比較して、DeNSの性能向上と学習効率はどのように差異を示すか？
• RQ5 DeNSは、異なる等変換ネットワークアーキテクチャおよびハイパーパrameterに一般化可能か？

主な発見

• DeNSは、追加データセットを必要とせず、OC20およびOC22ベンチマークで最先端の性能を達成した。
• MD17データセットでは、DeNSを用いることで、深さがL_max=3のモデルよりも3.1倍の高速化を達成したが、性能はベースライン（L_max=2）を上回った。
• DeNSを用いた学習では、MD17における力のMAEが20.8%改善され、非常に大きなデータセット（PCQM4Mv2）で事前学習した場合の17.2%の向上を上回ったが、学習コストはわずか10.5%増加にとどまった。
• 力符号化は極めて重要である：力符号化なしのDeNSでは一貫性のない向上が観察され、いくつかの分子では性能が向上しなかったが、力符号化ありのDeNSでは、すべての分子で一貫して性能向上が得られた。
• アーキテクチャに一般化可能である：DeNSは、アテンションやレイヤーノーマライゼーションを備えないEquiformerの変種に対しても性能向上を示し、等変換モデルへの広範な適用可能性を示した。
• DeNSは、事前学習と同等の性能向上を達成するが、はるかに低い計算コストで実現できるため、データが限られる原子的機械学習の文脈において、極めて効率的な代替手法である。

Figure 2 : Training process when incorporating DeNS as an auxiliary task. The upper blue block corresponds to the original task of energy and force predictions (Equation 1 ), and the lower red block corresponds to training DeNS (Equation 6 ). “Equivariant GNN” and “energy head” are shared across the