Generalizing Denoising to Non-Equilibrium Structures Improves Equivariant Force Fields

TL;DR

This work tackles the data scarcity challenge in ab initio atomistic modeling by introducing Denoising Non-Equilibrium Structures (DeNS), a force-conditioned denoising auxiliary task that extends denoising to non-equilibrium configurations. By encoding original forces into the denoising target, DeNS resolves the ill-posedness of denoising non-equilibrium structures and synergizes with SE(3)/E(3)-equivariant GNNs like EquiformerV2. Extensive experiments on OC20, OC22, and MD17 show that DeNS delivers new state-of-the-art results and significantly improves training efficiency, with modest overhead and broad applicability to other architectures. The approach offers a practical path toward more data-efficient, accurate surrogate models for energy and force predictions in atomistic simulations, catalysis design, and related domains.

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.

Figures (5)

• Figure 1: Illustration of denoising equilibrium and non-equilibrium structures. As shown in (a), we relax a non-equilibrium structure $S_{\text{non-eq}}$ and obtain the final relaxed, equilibrium structure $S_{\text{eq}}$. The path between $S_{\text{non-eq}}$ and $S_{\text{eq}}$ forms a relaxation trajectory. All structures $S$ along the trajectory except $S_{\text{eq}}$ are non-equilibrium and have non-zero forces $F(S)$. For denoising structures $S$, we add noise to their 3D atomic coordinates to obtain corrupted structures $\tilde{S}$ and predict the corresponding noise given $\tilde{S}$. We compare denoising non-equilibrium and equilibrium structures in (b) and (c), respectively, and show that denoising non-equilibrium structures can be ill-posed in (b). The issue in (b) can be addressed with force encoding, where we take forces as additional inputs, as in (d) and (e).
• Figure 2: Training process when incorporating DeNS as an auxiliary task. The pseudocode for training with DeNS can be found in Section \ref{['appendix:sec:pseudocode_dens']}. "Equivariant GNN", "energy head", "force head" and "noise head" are shared across (a), (b) and (c). For each batch of structures, we use the original task (a) for some structures and DeNS ((b) or (c)) for the others. Using partially corrupted structures as in (c) is empirically better than (b). We note that the force label $F(S_{\text{non-eq}})$ and energy label $E(S_{\text{non-eq}})$ used in (b) and (c) are the same as those in (a) and that training with DeNS does not require any additional data.
• Figure 3: Visualization of corrupted structures in OC20 dataset. We add noise of different scales to original structures (column 1). For each row, we sample $\bm{\epsilon}_{i} \sim \mathcal{N}(0, I_3)$, multiply $\bm{\epsilon}_{i}$ with $\sigma = 0.1 \text{ (column 2)}, 0.3\text{ (column 3)} \text{ and } 0.5\text{ (column 4)}$, and add the scaled noise to the original structures. For columns 2, 3 and 4, the lighter colors denote the atomic positions of the original structures. Here we add noise to all the atoms in a structure for better visual effects.
• Figure 4: Visualization of corrupted structures in OC22 dataset. We add noise of different scales to original structures (column 1). For each row, we sample $\bm{\epsilon}_{i} \sim \mathcal{N}(0, I_3)$, multiply $\bm{\epsilon}_{i}$ with $\sigma = 0.1 \text{ (column 2)}, 0.3\text{ (column 3)} \text{ and } 0.5\text{ (column 4)}$, and add the scaled noise to the original structures. For columns 2, 3 and 4, the lighter colors denote the atomic positions of the original structures. Here we add noise to all the atoms in a structure for better visual effects.
• Figure 5: Visualization of corrupted structures in MD17 dataset. We add noise of different scales to original structures (column 1). For each row, we sample $\bm{\epsilon}_{i} \sim \mathcal{N}(0, I_3)$, multiply $\bm{\epsilon}_{i}$ with $\sigma = 0.01 \text{ (column 2)}, 0.03\text{ (column 3)}, 0.05\text{ (column 4)} \text{ and } 0.07\text{ (column 5)}$, and add the scaled noise to the original structures. For columns 2, 3, 4 and 5, the lighter colors denote the atomic positions of the original structures. Here we add noise to all the atoms in a structure for better visual effects.