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

TL;DR

This paper proposes DeNS (Denoising Non-Equilibrium Structures) as an auxiliary self-supervised learning task to improve equivariant force fields in atomistic machine learning. By incorporating forces as conditioning signals during denoising of non-equilibrium atomic structures, DeNS transforms an ill-posed problem into a well-defined one, enabling significant performance gains on OC20, OC22, and MD17 with minimal training overhead.

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.

Motivation & Objective

• Address the limited scale of ab initio atomistic datasets by leveraging underutilized non-equilibrium structures in existing trajectories.
• Overcome the ill-posed nature of denoising non-equilibrium structures, where multiple target configurations exist due to non-zero forces.
• Improve the generalization and efficiency of equivariant neural networks for energy and force prediction in molecular systems.
• Demonstrate that auxiliary training with DeNS achieves performance gains comparable to pre-training, but with significantly lower computational cost.

Proposed method

• Corrupt 3D atomic coordinates of non-equilibrium structures by adding noise to create corrupted inputs.
• Condition the denoising task on the original forces of the uncorrupted structure to uniquely define the target reconstruction.
• Train a model to predict the noise vector that restores the original non-equilibrium structure given the corrupted structure and original forces.
• Integrate DeNS as an auxiliary loss during training of force field models, alongside primary tasks like energy and force prediction.
• Use equivariant message passing networks (e.g., Equiformer) that naturally handle forces as vector-valued features through irreducible representations (irreps).
• Apply force encoding via spherical harmonics to project forces into the network’s irreducible representation space, ensuring equivariance.

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.

Experimental results

Research questions

• RQ1 Can denoising non-equilibrium atomic structures be made tractable and effective as a self-supervised pretraining task?
• RQ2 Does incorporating original forces as conditioning signals resolve the ill-posedness of denoising non-equilibrium structures?
• RQ3 To what extent does DeNS improve the performance of equivariant force fields on standard benchmarks like OC20, OC22, and MD17?
• RQ4 How does DeNS compare to pre-training or architectural scaling in terms of performance gain and training efficiency?
• RQ5 Is DeNS generalizable across different equivariant network architectures and hyperparameters?

Key findings

• DeNS achieves state-of-the-art performance on both OC20 and OC22 benchmarks, outperforming prior methods without requiring additional datasets.
• On the MD17 dataset, DeNS reduces training time by 3.1× compared to a deeper model (L_max=3) while achieving better performance than the baseline (L_max=2).
• Training with DeNS yields a 20.8% improvement in force MAE on MD17, surpassing the 17.2% gain from pre-training on a much larger dataset (PCQM4Mv2), with only a 10.5% increase in training cost.
• Force encoding is critical: DeNS without force encoding shows inconsistent gains and fails to improve performance on several molecules, while DeNS with force encoding consistently improves results across all molecules.
• The method generalizes across architectures: DeNS improves performance on variants of Equiformer, including those without attention or layer normalization, indicating broad applicability to equivariant models.
• DeNS provides performance gains comparable to pre-training but with substantially lower computational cost, making it a highly efficient alternative for data-scarce atomistic ML.

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