[論文レビュー] May the Force be with You: Unified Force-Centric Pre-Training for 3D Molecular Conformations

Rui Feng, Qi Zhu+5 more|arXiv (Cornell University)|Aug 24, 2023

Machine Learning in Materials Science被引用数 4

ひとこと要約

本論文では、原子力の直接学習により、平衡状態と非平衡状態のデータを統一する力中心の事前学習フレームワーク、ET-OREOを提案する。力の統一的目標として、平衡状態データにはノイズ除去とゼロ力正則化を、非平衡状態データには直接的な力の監督を用いることで、未事前学習の等変換トランスフォーマーよりも3倍高い力の精度を達成し、NequIPに比べ2.45倍高速な推論速度を実現した。また、QM9の性質予測において最先端の性能を達成した。

ABSTRACT

Recent works have shown the promise of learning pre-trained models for 3D molecular representation. However, existing pre-training models focus predominantly on equilibrium data and largely overlook off-equilibrium conformations. It is challenging to extend these methods to off-equilibrium data because their training objective relies on assumptions of conformations being the local energy minima. We address this gap by proposing a force-centric pretraining model for 3D molecular conformations covering both equilibrium and off-equilibrium data. For off-equilibrium data, our model learns directly from their atomic forces. For equilibrium data, we introduce zero-force regularization and forced-based denoising techniques to approximate near-equilibrium forces. We obtain a unified pre-trained model for 3D molecular representation with over 15 million diverse conformations. Experiments show that, with our pre-training objective, we increase forces accuracy by around 3 times compared to the un-pre-trained Equivariant Transformer model. By incorporating regularizations on equilibrium data, we solved the problem of unstable MD simulations in vanilla Equivariant Transformers, achieving state-of-the-art simulation performance with 2.45 times faster inference time than NequIP. As a powerful molecular encoder, our pre-trained model achieves on-par performance with state-of-the-art property prediction tasks.

研究の動機と目的

• 既存の3次元分子事前学習モデルが平衡状態のコンformationにのみ焦点を当てており、非平衡状態データでは失敗することの制限を解消すること。
• 物理的に根拠のある1つの目的関数を用いて、平衡状態と非平衡状態の両方の分子コンformationにおける前処理を統一すること。
• 多様なコンformationalデータから直接原子力を学習することで、分子動力学的シミュレーションの精度と安定性を向上させること。
• 下流タスク（性質予測やシミュレーションなど）における一般化および転移性を高めること。
• 小規模で安価なポリマーを用いたデータ生成コストを削減し、未観測の大規模システムへの外挿を可能にすること。

提案手法

• 原子力を主な信号として扱い、その物理的不変性と分子間での転移可能性を活用することで、力中心の事前学習目的を提案する。
• 非平衡状態のコンformationについては、DFTなどのab initio手法で得られた原子力を直接モデルに監視させる。
• 平衡状態のコンformationについては、ゼロ力正則化と力に基づくノイズ除去を適用し、近似的に平衡状態に近い力を近似し、学習を安定化させる。
• 複数のソースからの約1,500万個のコンformationを含む統一データセット上で、等変換トランスフォーマーを学習する。
• 事前学習済みモデルを、分子動力学的シミュレーションやQM9における量子的性質予測といった下流タスクの強力なエンコーダーとして活用する。
• エネルギーに依存しない力の一貫性に基づくノイズ除去目的を平衡状態データに組み込み、耐性を向上させる。

(b) Potential energy curve during simulations.

実験結果

リサーチクエスチョン

• RQ1 原子力を根拠とする統一的前処理目的は、平衡状態と非平衡状態の両方の3次元分子コンformationを効果的に表現できるか？
• RQ2 エネルギーベースやノイズ除去ベースの手法と比較して、力中心の事前学習は分子動力学的シミュレーションの精度と安定性をどのように向上させるか？
• RQ3 非平衡状態のコンformationを組み込むことで、性質予測タスクにおける3次元分子表現の一般化能力はどの程度向上するか？
• RQ4 小分子と低コストDFTデータで事前学習したモデルは、未観測の大規模ポリマーにおけるシミュレーションに成功するか？
• RQ5 データ効率性と転移性能の観点から、力ベースの目的関数は従来のエネルギーベースやノイズ除去ベースの目的関数と比較してどのように異なるか？

主な発見

• ET-OREOは、未事前学習の等変換トランスフォーマーと比較して、約3倍高い力の精度を達成した。
• モデルは、NequIPに比べ2.45倍高速な推論速度を実現しながらも、DFT基準力と高い相関を維持し、安定した分子動力学的シミュレーションを可能にした。
• ET-OREOは、最大360原子を有する未観測の大規模ポリマーを効果的にシミュレートし、優れた一般化能力とコンformationサンプリング能力を示した。
• QM9のHOMO-LUMO性質予測において、最先端の性能を達成し、NoisyNodeと同等の性能を示し、スクラッチから事前学習したモデルを上回った。
• 豊富な非平衡状態データにアクセスしても、平衡状態の性質予測において顕著な精度向上は観察されず、このようなタスクは非平衡状態データの組み込みから恩恵を受けない可能性を示唆した。
• ゼロ力正則化と力に基づくノイズ除去は、通常の等変換トランスフォーマーで見られる不安定性を解消する効果を効果的に発揮した。