[Paper Review] 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 Science4 citations

TL;DR

This paper proposes ET-OREO, a force-centric pre-training framework for 3D molecular conformations that unifies equilibrium and off-equilibrium data by learning atomic forces directly. By leveraging forces as a unified objective—using denoising and zero-force regularization for equilibrium data and direct force supervision for off-equilibrium data—it achieves 3× higher force accuracy than untrained Equivariant Transformers and 2.45× faster inference than NequIP in molecular dynamics simulations, while matching state-of-the-art performance on QM9 property prediction.

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.

Motivation & Objective

• To address the limitation of existing 3D molecular pre-training models that focus only on equilibrium conformations and fail on off-equilibrium data.
• To unify pre-training across both equilibrium and off-equilibrium molecular conformations using a single, physically grounded objective.
• To improve the accuracy and stability of molecular dynamics simulations by learning atomic forces directly from diverse conformational data.
• To enhance the generalization and transferability of 3D molecular representations for downstream tasks like property prediction and simulation.
• To reduce the cost of data generation by enabling effective pre-training on small, inexpensive polymers while extrapolating to large, unseen systems.

Proposed method

• Propose a force-centric pre-training objective that treats atomic forces as the primary signal, leveraging their physical invariance and transferability across molecules.
• For off-equilibrium conformations, directly supervise the model to predict atomic forces obtained from ab initio methods like DFT.
• For equilibrium conformations, apply zero-force regularization and force-based denoising to approximate near-equilibrium forces and stabilize training.
• Train the Equivariant Transformer model on a unified dataset of ~15 million conformations from multiple sources, including both equilibrium and off-equilibrium states.
• Use the pre-trained model as a strong encoder for downstream tasks, including molecular dynamics simulations and quantum property prediction on QM9.
• Incorporate a de-noising objective on equilibrium data that does not rely on energy but instead on force consistency, improving robustness.

(b) Potential energy curve during simulations.

Experimental results

Research questions

• RQ1 Can a unified pre-training objective based on atomic forces effectively represent both equilibrium and off-equilibrium 3D molecular conformations?
• RQ2 How does force-centric pre-training improve the accuracy and stability of molecular dynamics simulations compared to energy-based or denoising-based methods?
• RQ3 To what extent does incorporating off-equilibrium conformations enhance the generalization of 3D molecular representations for property prediction tasks?
• RQ4 Can a model pre-trained on small molecules and low-cost DFT data successfully extrapolate to large, unseen polymers in simulation?
• RQ5 How does the force-based objective compare to traditional energy-based or denoising objectives in terms of data efficiency and transfer performance?

Key findings

• ET-OREO achieves approximately 3 times higher force accuracy compared to the un-pretrained Equivariant Transformer model.
• The model enables stable molecular dynamics simulations with 2.45 times faster inference speed than NequIP, while maintaining high correlation with DFT reference forces.
• ET-OREO successfully simulates large, unseen polymers with up to 360 atoms, demonstrating strong generalization and conformational sampling ability.
• The model achieves state-of-the-art performance on QM9 HOMO-LUMO property prediction, matching NoisyNode and outperforming models pre-trained from scratch.
• Despite access to rich off-equilibrium data, no significant accuracy gain is observed in equilibrium property prediction, suggesting such tasks may not benefit from off-equilibrium inclusions.
• The zero-force regularization and force-based denoising effectively stabilize training on equilibrium data, resolving instability issues seen in vanilla Equivariant Transformers.