vskip=0pt

# Generalizing Denoising to Non-Equilibrium Structures
Improves Equivariant Force Fields

Yi-Lun Liao<EMAIL_ADDRESS>
Massachusetts Institute of Technology
Work partially done during an internship at FAIR, Meta Tess Smidt
<EMAIL_ADDRESS>
Massachusetts Institute of Technology Abhishek Das<EMAIL_ADDRESS>
FAIR, Meta

###### 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.

##  1 Introduction

Graph neural networks (GNNs) have made remarkable progress in approximating
high-fidelity, compute-intensive quantum mechanical calculations like density
functional theory (DFT) for atomistic systems (Gilmer et al., 2017; Zhang et
al., 2018; Unke et al., 2021; Batzner et al., 2022; Rackers et al., 2023; Lan
et al., 2022), enabling new insights in applications such as molecular
dynamics simulations (Musaelian et al., 2023) and catalyst design (Chanussot*
et al., 2021; Lan et al., 2022). However, unlike other domains such as natural
language processing (NLP) and computer vision (CV), the scale of atomistic
data is quite limited since generating data requires compute-intensive ab
initio calculations. For example, the largest atomistic dataset, OC20
(Chanussot* et al., 2021), contains about $138$M examples while GPT-3 (Brown
et al., 2020) is trained on hundreds of billions of words and ViT-22B
(Dehghani et al., 2023) is trained on around $4$B images.

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.

To start addressing this gap, we take inspiration from self-supervised
learning methods in NLP and CV and explore how we can adapt them to learn
better atomistic representations from existing labeled data. Specifically, one
of the most popular self-supervised learning methods in NLP (Devlin et al.,
2019) and CV (He et al., 2022) is training a denoising autoencoder (Vincent et
al., 2008), where the idea is to mask or corrupt a part of the input data and
learn to reconstruct or denoise the original, uncorrupted data. Denoising
assumes we know a unique target structure such as a sentence and an image in
the case of NLP and CV. Indeed, this is the case for equilibrium structures
(e.g., $S_{\text{eq}}$ at a local energy minimum in Figure 1b) as has been
demonstrated by previous works leveraging denoising for pretraining on
atomistic data (Jiao et al., 2022; Zaidi et al., 2023; Liu et al., 2023; Wang
et al., 2023; Feng et al., 2023a). However, most previous works are limited to
equilibrium structures, and equilibrium structures constitute only a small
portion of available data since structures along a trajectory to get to a
local minimum are all non-equilibrium. Hence, it is important to generalize
denoising to leverage the much larger set of non-equilibrium structures.

Since a non-equilibrium structure has non-zero atomic forces and atoms are not
confined to local energy minima, it can have more possible atomic positions
than a structure at equilibrium. As shown in Figure 1a, this can make
denoising a non-equilibrium structure an ill-posed problem since there are
many possible target structures. To address the issue, we propose to take the
forces of the original non-equilibrium structures as inputs when denoising
non-equilibrium structures. Intuitively, the forces constrain the atomic
positions of a non-equilibrium structure. With the additional information, we
are able to predict the original non-equilibrium structure satisfying the
input forces instead of predicting any arbitrary structures as shown in Figure
1c. Previous works on denoising equilibrium structures (Jiao et al., 2022;
Zaidi et al., 2023; Liu et al., 2023; Feng et al., 2023b; a) end up being a
special case where the forces of original structures are close to zero.

Based on the insight, in this paper, we propose to use denoising non-
equilibrium structures (DeNS) as an auxiliary task to better leverage
atomistic data. For training DeNS, we first corrupt a structure by adding
noise to its 3D atomic coordinates and then reconstruct the original
uncorrupted structure by predicting the noise. For noise predictions, a model
is given the forces of the original uncorrupted structure as inputs to make
the transformation from a corrupted non-equilibrium structure to an
uncorrupted non-equilibrium structure tractable. When used along with the
original tasks like predicting the energy and forces of non-equilibrium
structures, DeNS improves the performance of the original tasks with a
marginal increase in training cost. We further discuss how DeNS can better
leverage training data to improve performance and show the connection to self-
supervised learning methods in other domains.

Because DeNS requires encoding forces, it favors equivariant networks, which
build up equivariant features at each node with vector spaces of irreducible
representations (irreps) and have interactions or message passing between
nodes with equivariant operations like tensor products. Since forces can be
projected to vector spaces of irreps with spherical harmonics, equivariant
networks can easily incorporate forces in node embeddings. Moreover, with the
reduced complexity of equivariant operations (Passaro & Zitnick, 2023) and
incorporation of the Transformer network design (Liao & Smidt, 2023; Liao et
al., 2023) from NLP (Vaswani et al., 2017) and CV (Dosovitskiy et al., 2021),
equivariant networks have become the state-of-the-art methods on large-scale
atomistic datasets.

We mainly focus on how DeNS can improve equivariant networks and conduct
extensive experiments on OC20 (Chanussot* et al., 2021), OC22 (Tran* et al.,
2022) and MD17 (Chmiela et al., 2017; Schütt et al., 2017; Chmiela et al.,
2018) datasets. EquiformerV2 (Liao et al., 2023) trained with DeNS achieves
better S2EF results and comparable IS2RE results on OC20 (Section 4.1).
EquiformerV2 trained with DeNS sets new state-of-the-art results on OC22,
improving energy by ${\sim}13\%$, forces by ${\sim}7\%$, and IS2RE by
${\sim}12\%$ (Section 4.2). On MD17, Equiformer ($L_{max}=2$) (Liao & Smidt,
2023) trained with DeNS achieves better results and saves $3.1\times$ training
time compared to Equiformer ($L_{max}=3$) without DeNS (Section 4.3), where
$L_{max}$ denotes the maximum degree. Additionally, DeNS can improve other
equivariant networks like eSCN (Passaro & Zitnick, 2023) on OC20 and SEGNN-
like networks (Brandstetter et al., 2022) on MD17.

##  2 Related Works

##### Denoising 3D Atomistic Structures.

Denoising structures have been used to boost the performance of GNNs on 3D
atomistic datasets (Godwin et al., 2022; Jiao et al., 2022; Zaidi et al.,
2023; Liu et al., 2023; Feng et al., 2023b; Wang et al., 2023; Feng et al.,
2023a). The approach is to first corrupt data by adding noise and then train a
denoising autoencoder to reconstruct the original data by predicting the
noise, and the motivation is that learning to reconstruct data enables
learning generalizable representations (Devlin et al., 2019; He et al., 2022;
Godwin et al., 2022; Zaidi et al., 2023). Since denoising equilibrium
structures do not require labels and is self-supervised, similar to BERT
(Devlin et al., 2019) and MAE (He et al., 2022), it is common to pre-train via
denoising on a large dataset of equilibrium structures like PCQM4Mv2 (Nakata &
Shimazaki, 2017) and then fine-tune with supervised learning on smaller
downstream datasets. Besides, Noisy Nodes (Godwin et al., 2022) use denoising
equilibrium structures as an auxiliary task along with original tasks without
pre-training on another larger dataset. However, most of the previous works
are limited to equilibrium structures, which occupy a much smaller amount of
data than non-equilibrium ones. In contrast, the proposed DeNS generalizes
denoising to non-equilibrium structures with force encoding so that we can
improve the performance on the larger set of non-equilibrium structures. We
provide a detailed comparison to previous works on denoising in Section A.2.

#####  SE(3)/E(3)-Equivariant Networks.

Refer to Section A.1 for discussion on equivariant networks.

##  3 Method

###  3.1 Problem Setup

Calculating quantum mechanical properties like energy and forces of 3D
atomistic systems is fundamental to many applications. An atomistic system can
be one or more molecules, a crystalline material and so on. Specifically, each
system $S$ is an example in a dataset and can be described as
$S=\left\\{(z_{i},\mathbf{p}_{i})\mid i\in\\{1,...,|S|\\}\right\\}$, where
$z_{i}\in\mathbb{N}$ denotes the atomic number of the $i$-th atom and
$\mathbf{p}_{i}\in\mathbb{R}^{3}$ denotes the 3D atomic position. The energy
of $S$ is denoted as $E(S)\in\mathbb{R}$, and the atom-wise forces are denoted
as $F(S)=\left\\{\mathbf{f}_{i}\in\mathbb{R}^{3}\mid
i\in\\{1,...,|S|\\}\right\\}$, where $\mathbf{f}_{i}$ is the force acting on
the $i$-th atom. In this paper, we define a system to be an equilibrium
structure if all of its atom-wise forces are close to zero. Otherwise, we
refer to it as a non-equilibrium structure. Since non-equilibrium structures
have non-zero atomic forces and thus are not at an energy minimum, they have
more degrees of freedom and constitute a much larger set of possible
structures than those at equilibrium.

In this work, we focus on the task of predicting energy and forces given non-
equilibrium structures. Specifically, given a non-equilibrium structure
$S_{\text{non-eq}}$, GNNs predict energy $\hat{E}(S_{\text{non-eq}})$ and
atom-wise forces $\hat{F}(S_{\text{non-
eq}})=\left\\{\hat{\mathbf{f}}_{i}(S_{\text{non-eq}})\in\mathbb{R}^{3}\mid
i\in\\{1,...,|S_{\text{non-eq}}|\\}\right\\}$ and minimize the following loss
function:

$\lambda_{E}\cdot\mathcal{L}_{E}+\lambda_{F}\cdot\mathcal{L}_{F}=\lambda_{E}\cdot|E^{\prime}(S_{\text{non-
eq}})-\hat{E}(S_{\text{non-eq}})|+\lambda_{F}\cdot\frac{1}{|S_{\text{non-
eq}}|}\sum_{i=1}^{|S_{\text{non-eq}}|}|\mathbf{f}_{i}^{\prime}(S_{\text{non-
eq}})-\hat{\mathbf{f}}_{i}(S_{\text{non-eq}})|^{2}$ (1)

where $\lambda_{E}$ and $\lambda_{F}$ are energy and force coefficients
controlling the relative importance between energy and force predictions.
$E^{\prime}(S_{\text{non-eq}})=\frac{E(S_{\text{non-
eq}})-\mu_{E}}{\sigma_{E}}$ is the normalized ground truth energy obtained by
first subtracting the original energy $E(S_{\text{non-eq}})$ by the mean of
energy labels in the training set $\mu_{E}$ and then dividing by the standard
deviation of energy labels $\sigma_{E}$. Similarly,
$\mathbf{f}_{i}^{\prime}=\frac{\mathbf{f}_{i}}{\sigma_{F}}$ is the normalized
atom-wise force. For force prediction, we can either directly predict them
from latent representations like node embeddings as commonly used for OC20 and
OC22 datasets or take the negative gradients of predicted energy with respect
to atomic positions for datasets like MD17.

###  3.2 Denoising Non-Equilibrium Structures (DeNS)

####  3.2.1 Formulation of Denoising

Denoising structures have been used to improve the performance of GNNs on 3D
atomistic datasets. They first corrupt data by adding noise and then train a
denoising autoencoder to reconstruct the original data by predicting the
noise. Specifically, given a 3D atomistic system
$S=\left\\{(z_{i},\mathbf{p}_{i})\mid i\in\\{1,...,|S|\\}\right\\}$, we create
a corrupted structure $\tilde{S}$ by adding Gaussian noise with a tuneable
standard deviation $\sigma$ to the atomic positions $\mathbf{p}_{i}$ of the
original structure $S$:

$\tilde{S}=\left\\{(z_{i},\tilde{\mathbf{p}_{i}})\mid
i\in\\{1,...,|S|\\}\right\\}\text{,}\quad\text{where}\quad\tilde{\mathbf{p}_{i}}=\mathbf{p}_{i}+\bm{\epsilon}_{i}\quad\text{and}\quad\bm{\epsilon}_{i}\sim\mathcal{N}(0,\sigma
I_{3})$ (2)

We denote the set of noise added to $S$ as
$\text{Noise}(S,\tilde{S})=\left\\{\bm{\epsilon_{i}}\in\mathbb{R}^{3}\mid
i\in\\{1,...,|S|\\}\right\\}$. When training a denoising autoencoder, GNNs
take $\tilde{S}$ as inputs, output atom-wise noise predictions
$\hat{\bm{\epsilon}}(\tilde{S})_{i}$ and minimize the L2 difference between
normalized noise $\frac{\bm{\epsilon}_{i}}{\sigma}$ and noise predictions
$\hat{\bm{\epsilon}}(\tilde{S})_{i}$:

$\mathbb{E}_{p(S,\tilde{S})}\left[\frac{1}{|S|}\sum_{i=1}^{|S|}\left|\frac{\bm{\epsilon}_{i}}{\sigma}-\hat{\bm{\epsilon}}(\tilde{S})_{i}\right|^{2}\right]$
(3)

where $p(S,\tilde{S})$ denotes the probability of obtaining the corrupted
structure $\tilde{S}$ from the original structure $S$. We divide the noise
$\bm{\epsilon}_{i}$ by the standard deviation $\sigma$ to normalize the
outputs of noise predictions.

When the original structure $S$ is an equilibrium structure, denoising is to
find the structure corresponding to the nearest energy local minima given a
high-energy corrupted structure. This makes denoising equilibrium structures a
many-to-one mapping and therefore a well-defined problem. However, when the
original structure $S$ is a non-equilibrium structure, denoising, the
transformation from a corrupted non-equilibrium structure to the original non-
equilibrium one, can be an ill-posed problem since there are many possible
target structures as shown in Figure 1a.

####  3.2.2 Force Encoding

The reason that denoising non-equilibrium structures can be ill-posed is
because we do not provide sufficient information to specify the properties of
the target structures. Concretely, given an original non-equilibrium structure
$S_{\text{non-eq}}$ and its corrupted counterpart $\tilde{S}_{\text{non-eq}}$,
some structures interpolated between $S_{\text{non-eq}}$ and
$\tilde{S}_{\text{non-eq}}$ could be in the same data distribution and
therefore be the potential target structures of denoising. In contrast, when
denoising equilibrium structures as shown in Figure 1b, we implicitly provide
the extra information that the target structure should be at equilibrium with
near-zero forces, and this therefore limits the possibility of the target of
denoising. Motivated by the assumption that the forces of the original
structures are close to zeros when denoising equilibrium ones, we propose to
encode the forces of original non-equilibrium structures when denoising non-
equilibrium ones as illustrated in Figure 1c. Specifically, when training
denoising non-equilibrium structures (DeNS), GNNs take both a corrupted non-
equilibrium structure $\tilde{S}_{\text{non-eq}}$ and the forces
$F(S_{\text{non-eq}})$ of the original non-equilibrium structure
$S_{\text{non-eq}}$ as inputs, output atom-wise noise predictions
$\hat{\bm{\epsilon}}\left(\tilde{S}_{\text{non-eq}},F(S_{\text{non-
eq}})\right)_{i}$ and minimize the L2 difference between normalized noise
$\frac{\bm{\epsilon}_{i}}{\sigma}$ and noise predictions
$\hat{\bm{\epsilon}}\left(\tilde{S}_{\text{non-eq}},F(S_{\text{non-
eq}})\right)_{i}$:

$\mathcal{L}_{\text{DeNS}}=\mathbb{E}_{p(S_{\text{non-
eq}},\tilde{S}_{\text{non-eq}})}\left[\frac{1}{|S_{\text{non-
eq}}|}\sum_{i=1}^{|S_{\text{non-
eq}}|}\left|\frac{\bm{\epsilon}_{i}}{\sigma}-\hat{\bm{\epsilon}}\left(\tilde{S}_{\text{non-
eq}},F(S_{\text{non-eq}})\right)_{i}\right|^{2}\right]$ (4)

Equation 4 is more general and reduces to Equation 3 when the target of
denoising becomes equilibrium structures with near-zero forces. Since we train
GNNs with $\tilde{S}_{\text{non-eq}}$ and $F(S_{\text{non-eq}})$ as inputs and
$\text{Noise}(S_{\text{non-eq}},\tilde{S}_{\text{non-eq}})$ as outputs, they
implicitly learn to leverage $F(S_{\text{non-eq}})$ to reconstruct
$S_{\text{non-eq}}$ instead of predicting any arbitrary non-equilibrium
structures. Empirically, we find that force encoding is critical to the
effectiveness of DeNS on OC20 S2EF-2M dataset (Index 1 and Index 2 in Table
LABEL:tab:design_choices) and MD17 dataset (Index 2 and Index 3 in Table 6).

Since DeNS requires encoding forces, DeNS favors equivariant networks, which
can easily incorporate forces as well as other higher-degree tensors like
stress into their node embeddings. Specifically, the node embeddings of
equivariant networks are equivariant irreps features built from vectors spaces
of irreducible representations (irreps) and contain $C_{L}$ channels of
type-$L$ vectors with degree $L$ being in the range from $0$ to maximum degree
$L_{max}$. $C_{L}$ and $L_{max}$ are architectural hyper-parameters of
equivariant networks. To obtain the force embedding $x_{\mathbf{f}}$ from the
input force $\mathbf{f}$, we first project $\mathbf{f}$ into type-$L$ vectors
with spherical harmonics
$Y^{(L)}\left(\frac{\mathbf{f}}{||\mathbf{f}||}\right)$, where $0\leq L\leq
L_{max}$, and then expand the number of channels from $1$ to $C_{L}$ with an
$SO(3)$ linear layer (Geiger et al., 2022; Geiger & Smidt, 2022; Liao & Smidt,
2023):

$x_{\mathbf{f}}^{(L)}=\text{SO3\\_Linear}^{(L)}\left(||\mathbf{f}||\cdot
Y^{(L)}\left(\frac{\mathbf{f}}{||\mathbf{f}||}\right)\right)$ (5)

$x_{\mathbf{f}}^{(L)}$ denotes the channels of type-$L$ vectors in force
embedding $x_{\mathbf{f}}$, and $\text{SO3\\_Linear}^{(L)}$ denotes the
$SO(3)$ linear operation on type-$L$ vectors. Since we normalize the input
force when using spherical harmonics, we multiply
$Y^{(L)}\left(\frac{\mathbf{f}}{||\mathbf{f}||}\right)$ with the norm of input
force $||\mathbf{f}||$ to recover the information of force magnitude. After
computing force embeddings for all atom-wise forces, we simply add the force
embeddings to initial node embeddings to encode forces in equivariant
networks.

On the other hand, it might not be that intuitive to encode forces in
invariant networks since their internal latent representations such as node
embeddings and edge embeddings are scalars instead of type-$L$ vectors or
geometric tensors. One potential manner of encoding forces in latent
representations is to project them into edge embeddings by taking inner
products between forces and edge vectors of relative positions. This process
is the inverse of how GemNet-OC (Gasteiger et al., 2022) decodes forces from
latent representations. Since equivariant networks have been shown to
outperform invariant networks on atomistic datasets and are a more natural fit
to encoding forces, we mainly focus on equivariant networks and how DeNS can
further advance their performance.

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 two
tasks. For each batch of structures, we use the original task for some
structures and DeNS for the others.

####  3.2.3 Training with DeNS

##### Auxiliary Task.

We propose to use DeNS as an auxiliary task along with the original task of
predicting energy and forces and summarize the training process in Figure 2.
Specifically, given a batch of structures, for each structure, we decide
whether we optimize the objective of DeNS (Equation 6) or the objective of the
original task (Equation 1). This introduces an additional hyper-parameter
$p_{\text{DeNS}}$, the probability of optimizing DeNS. We use an additional
noise head for noise prediction, which slightly increases training time.
Additionally, when training DeNS, similar to Noisy Nodes (Godwin et al.,
2022), we also leverage energy labels and predict the energy of original
structures. Therefore, given an original non-equilibrium structure
$S_{\text{non-eq}}$ and the corrupted counterpart $\tilde{S}_{\text{non-eq}}$,
training DeNS corresponds to minimizing the following loss function:

$\lambda_{E}\cdot\mathcal{L}_{E}+\lambda_{\text{DeNS}}\cdot\mathcal{L}_{\text{DeNS}}=\lambda_{E}\cdot\left|E^{\prime}(S_{\text{non-
eq}})-\hat{E}\left(\tilde{S}_{\text{non-eq}},F(S_{\text{non-
eq}})\right)\right|+\lambda_{\text{DeNS}}\cdot\mathcal{L}_{\text{DeNS}}$ (6)

where $\mathcal{L}_{\text{DeNS}}$ denotes the loss function of denoising as
defined in Equation 4. We note that we also encode forces as discussed in
Section 3.2.2 to predict the energy of $S_{\text{non-eq}}$, and we share the
energy prediction head across Equation 1 and Equation 6. The loss function
introduces another hyper-parameter $\lambda_{\text{DeNS}}$, DeNS coefficient,
controlling the relative importance of the auxiliary task. Besides, the
process of corrupting structures also results in another hyper-parameter
$\sigma$ as shown in Equation 2. We provide the pseudocode in Section E. We
note that we only train DeNS with force encoding on the training set without
using any force labels on the validation and testing sets.

##### Partially Corrupted Structures.

Empirically, we find that adding noise to all atoms in a structure can lead to
limited performance gain of DeNS on OC22 and MD17 datasets, and we surmise
that there are still several structures satisfying input forces if we add
noise to all atoms, making denoising too difficult a task and potentially not
well-defined. To address this issue, we use partially corrupted structures,
where we only add noise to and denoise a subset of atoms. Specifically, for
corrupted atoms, we encode their atom-wise forces and predict the noise. For
other uncorrupted atoms, we do not encode forces and train for the original
task of predicting forces. We also predict the energy of the original
structures given partially corrupted structures. This introduces an additional
hyper-parameter, corruption ratio $r_{\text{DeNS}}$, which is the ratio of the
number of corrupted atoms to that of all atoms.

##### Multi-Scale Noise.

Inspired by prior works on denoising score matching (Song & Ermon, 2019;
2020), we empirically find that incorporating multiple noise scales together
for denoising improves energy and force predictions on OC20 and OC22 datasets.
Specifically, we choose the standard deviations $\\{\sigma_{k}\\}_{k=1}^{T}$
to be a geometric progression that satisfy
$\frac{\sigma_{T}}{\sigma_{T-1}}=...=\frac{\sigma_{3}}{\sigma_{2}}=\frac{\sigma_{2}}{\sigma_{1}}>1$:

$\sigma_{k}=\text{exp}\left(\log_{e}\sigma_{\text{low}}+\frac{k-1}{T-1}\cdot\left(\log_{e}\sigma_{\text{high}}-\log_{e}\sigma_{\text{low}}\right)\right)\quad\text{where}\quad
k=1,...,T$ (7)

Here we use $\sigma_{1}=\sigma_{\text{low}}=0.01$ and $T=50$ and only tune
$\sigma_{T}=\sigma_{\text{high}}$ when multi-scale noise is adopted. When
training with DeNS, for each structure, we first sample a single noise
standard deviation $\sigma$ from the $T$ values, and then follow Equations 2
and 4. We surmise that multi-scale noises are more likely to span the
distribution of meaningful non-equilibrium structures across a diverse range
of geometries compared to a fixed noise standard deviation.

####  3.2.4 Discussion

##### Why Does DeNS Help.

DeNS can better leverage training data to improve the performance in the
following two ways. First, DeNS adds noise to non-equilibrium structures to
generate structures with new geometries and therefore naturally achieves data
augmentation (Godwin et al., 2022). Second, training DeNS enourages learning a
different yet highly correlated interaction. Since we encode forces as inputs
and predict the original structures in terms of noise corrections, DeNS
enables learning the interaction of transforming forces into structures, which
is the inverse of force predictions. As demonstrated in the works of Noisy
Nodes (Godwin et al., 2022) and UL2 (Tay et al., 2023), training a single
model with multiple correlated objectives to learn different interactions can
help the performance on the original task.

##### Connection to Self-Supervised Learning.

DeNS shares similar intuitions to self-supervised learning methods like BERT
(Devlin et al., 2019) and MAE (He et al., 2022) and other denoising methods
(Vincent et al., 2008; 2010; Godwin et al., 2022; Zaidi et al., 2023) – they
remove or corrupt a portion of input data and then learn to predict the
original data. Learning to reconstruct data can help learning generalizable
representations, and therefore these methods can use the task of
reconstruction to improve the performance of downstream tasks. However, unlike
those self-supervised learning methods (Devlin et al., 2019; He et al., 2022;
Zaidi et al., 2023), DeNS requires force labels for encoding, and therefore,
we propose to use DeNS as an auxiliary task optimized along with original
tasks and do not follow the previous practice of first pre-training and then
fine-tuning. Additionally, we note that before obtaining a single equilibrium
structure, we need to run relaxations and generate many intermediate non-
equilibrium ones, which is the labeling process as well. We hope that the
ability to leverage more from non-equilibrium structures as proposed in this
work can encourage researchers to release data containing intermediate non-
equilibrium structures in addition to final equilibrium ones. Moreover, we
note that DeNS can also be used in fine-tuning. For example, we can first pre-
train models on PCQM4Mv2 dataset and then fine-tune them on the smaller MD17
dataset with both the original task and DeNS.

##### Marginal Increase in Training Time.

Since we use an additional noise head for denoising, training with DeNS
marginally increases the time of each training iteration. We optimize DeNS for
some structures and the original task for the others for each training
iteration, and we demonstrate that DeNS can improve the performance given the
same amount of training iterations. Therefore, training with DeNS only
marginally increase the overall training time.

##  4 Experiments

###  4.1 OC20 Dataset

##### Dataset and Tasks.

We start with experiments on the large and diverse Open Catalyst 2020 dataset
(OC20) (Chanussot* et al., 2021), which consists of about $1.2$M Density
Functional Theory (DFT) relaxation trajectories. Each DFT trajectory in OC20
starts from an initial structure of an adsorbate molecule placed on a catalyst
surface, which is then relaxed with the revised Perdew-Burke-Ernzerhof (RPBE)
functional (Hammer et al., 1999) calculations to a local energy minimum.
Relevant to DeNS, all the intermediate structures from a relaxation
trajectory, except the relaxed structure, are considered non-equilibrium
structures. The relaxed or equilibrium structure has forces close to zero.

The primary task in OC20 is Structure to Energy and Forces (S2EF), which is to
predict the energy and per-atom forces given an equilibrium or non-equilibrium
structure from any point in the trajectory. These predictions are evaluated on
energy and force mean absolute error (MAE). Once a model is trained for S2EF,
it is used to run structural relaxations from an initial structure using the
predicted forces till a local energy minimum is found. The energy predictions
of these relaxed structures are evaluated on the Initial Structure to Relaxed
Energy (IS2RE) task.

##### Training Details.

Please refer to Section B.1 for details on DeNS, architectures, hyper-
parameters and training time.

Epochs |  $\sigma_{\text{high}}$ |  forces |  energy
---|---|---|---
$12$ |  $0.125$ |  19.10 |  284
$12$ |  $0.250$ |  19.11 |  276
$12$ |  $0.500$ |  baselinecolor19.31 |  baselinecolor273
$12$ |  $0.750$ |  19.43 |  273
$12$ |  $1.000$ |  19.52 |  274
$20$ |  $0.250$ |  18.41 |  271
$20$ |  $0.500$ |  baselinecolor18.48 |  baselinecolor263
$20$ |  $0.750$ |  18.66 |  263
$20$ |  $1.000$ |  18.66 |  261
$30$ |  $0.250$ |  17.85 |  265
$30$ |  $0.500$ |  baselinecolor17.96 |  baselinecolor255
$30$ |  $0.750$ |  18.06 |  255
$30$ |  $1.000$ |  18.06 |  256

(a)

Epochs |  $p_{\text{DeNS\leavevmode\nobreak\ }}$ |  forces |  energy
---|---|---|---
$12$ |  $0.125$ |  19.64 |  276
$12$ |  $0.250$ |  19.31 |  273
$12$ |  $0.500$ |  baselinecolor19.32 |  baselinecolor271
$12$ |  $0.750$ |  19.90 |  281
$20$ |  $0.250$ |  18.66 |  263
$20$ |  $0.500$ |  baselinecolor18.49 |  baselinecolor262
$20$ |  $0.750$ |  18.92 |  269
$30$ |  $0.250$ |  18.06 |  255
$30$ |  $0.500$ |  baselinecolor17.83 |  baselinecolor255
$30$ |  $0.750$ |  18.24 |  262

(b)

Epochs |  $\lambda_{\text{DeNS\leavevmode\nobreak\ }}$ |  forces |  energy
---|---|---|---
$12$ |  $5$ |  19.64 |  277
$12$ |  $10$ |  baselinecolor19.31 |  baselinecolor273
$12$ |  $15$ |  19.25 |  275
$20$ |  $5$ |  19.02 |  262
$20$ |  $10$ |  18.66 |  263
$20$ |  $15$ |  baselinecolor18.56 |  baselinecolor263
$30$ |  $5$ |  18.40 |  260
$30$ |  $10$ |  18.06 |  255
$30$ |  $15$ |  baselinecolor17.94 |  baselinecolor258

(c)

| EquiformerV2 | EquiformerV2 + DeNS
---|---|---
Epochs |  forces |  energy | # params | training time |  forces |  energy | # params | training time
$12$ |  20.46 |  285 |  $83$M | 1398 |  19.32 |  271 |  $89$M | 1501
$20$ |  19.78 |  280 |  $83$M | 2330 |  18.49 |  262 |  $89$M | 2501
$30$ |  19.42 |  278 |  $83$M | 3495 |  17.83 |  255 |  $89$M | 3752
| eSCN | eSCN + DeNS
$20$ |  20.61 |  290 |  $52$M | 1802 |  19.07 |  279 |  $52$M | 1829

(d)

Index |  |  forces |  energy
---|---|---|---
1 | DeNS |  19.32 |  271
2 | Without force encoding |  20.16 |  271
3 |  $\lambda_{E}=0$ in Eq. 6 |  19.57 |  281
4 |  Fixed $\sigma=0.1$ |  19.58 |  283
5 |  With $\sigma$ encoding |  19.66 |  279

(e)

Table 1:  Ablation results of EquifomerV2 trained with DeNS on the $2$M subset
of the OC20 S2EF dataset. We report mean absolute errors for forces in meV/Å
and energy in meV, and lower is better. Errors are averaged over the four
validation sub-splits of OC20. The training time is in GPU-hours and measured
on V100 GPUs. (a)-(c) We use $\sigma_{\text{high}}=0.5$ for $12$ epochs,
$\sigma_{\text{high}}=0.75$ for $20$ and $30$ epochs, $p_{\text{DeNS}}=0.25$,
and $\lambda_{\text{DeNS}}=10$ as the default DeNS-related hyper-parameters
and vary them to study how they affect the performance. The settings with the
best energy-force trade-offs are marked in gray. (d) We train EquiformerV2 and
eSCN and compare the results of training with and without DeNS. The results of
EquiformerV2 trained with DeNS are from (b). (e) We train EquiformerV2 for
$12$ epochs with the best setting in (b) to verify the design choices of DeNS.

####  4.1.1 Ablation studies

We use EquiformerV2 (Liao et al., 2023) and S2EF-$2$M split of OC20 to
investigate how DeNS-related hyper-parameters affect the performance, compare
the results of training with and without DeNS and verify some design choices
of DeNS.

##### Hyperparameters.

In Tables LABEL:tab:sigma, LABEL:tab:prob, LABEL:tab:lambda, we vary the upper
bound on standard deviations of Gaussian noise $\sigma_{\text{high}}$, the
probability of optimizing DeNS $p_{\text{DeNS\leavevmode\nobreak\ }}$, and the
loss coefficient for DeNS $\lambda_{\text{DeNS\leavevmode\nobreak\ }}$ to
study how the hyper-parameters of DeNS affect performance. We find that the
optimal settings are similar when training for different epochs and have the
following observations. First, as we increase $\sigma_{\text{high}}$, force
predictions become worse monotonically, while energy predictions improve but
saturate at $\sigma_{\text{high}}=0.5$. Second, $p_{\text{DeNS}}=0.5$ and
$\lambda_{\text{DeNS}}=10\text{ and }15$ work better than any other values
across the three different epochs.

##### Comparison between Training with and without DeNS.

Table LABEL:tab:comparison_with_without_dens summarizes the results of
training with and without DeNS. For EquiformerV2, incorporating DeNS as an
auxiliary task boosts the performance of energy and force predictions and only
increases training time by $7.4\%$ and the number of parameters from $83$M to
$89$M. Particularly, EquiformerV2 trained with DeNS for $12$ epochs
outperforms EquiformerV2 trained without DeNS for $30$ epochs, saving
$2.3\times$ training time. Additionally, we also show that DeNS can be
applicable to other equivariant networks like eSCN (Passaro & Zitnick, 2023),
and training eSCN with DeNS improves both energy and force MAE while slightly
increasing training time by $1.5\%$. The different amounts of increase in
training time between EquiformerV2 and eSCN are because they use different
modules for noise prediction.

##### Design Choices.

We conduct experiments to verify the design choices of DeNS and summarize the
results in Table LABEL:tab:design_choices. All the models follow the best
setting of training with DeNS for $12$ epochs in Table LABEL:tab:prob.
Comparing Index 1 and Index 2, we show that encoding forces $F(S_{\text{non-
eq}})$ in Equations 4 and 6 enables denoising non-equilibrium structures to
further improve performance. DeNS without force encoding only results in
slightly better force MAE than training without DeNS as in Table
LABEL:tab:comparison_with_without_dens. We also compare DeNS with and without
force encoding on MD17 in Section 4.3 and find that force encoding is
critical. Comparing Index 1 and Index 3, we demonstrate that predicting energy
of original structures given corrupted ones can be helpful to the original
task. Additionally, we compare the performance of using a fixed $\sigma$ and
multi-scale noise, and the comparison between Index 1 and Index 4 shows that
multi-scale noise improves both energy and force predictions. Since we sample
standard deviation $\sigma$ when using multi-scale noise, we investigate
whether we need to encode $\sigma$. The comparison between Index 1 and Index 5
shows that DeNS without $\sigma$ encoding works better, and thus we can use
the same approach when we use either a fixed $\sigma$ or multi-scale noise.

| Throughput | S2EF validation | S2EF test | IS2RE test
---|---|---|---|---
| Samples / | Energy MAE | Force MAE | Energy MAE | Force MAE | Energy MAE
Model | GPU sec. $\uparrow$ |  $\mathrm{(}\mathrm{m}\mathrm{eV}\mathrm{)}$ $\downarrow$ |  $\mathrm{(}\mathrm{m}\mathrm{eV}\mathrm{/}\mathrm{\text{Å}}\mathrm{)}$ $\downarrow$ |  $\mathrm{(}\mathrm{m}\mathrm{eV}\mathrm{)}$ $\downarrow$ |  $\mathrm{(}\mathrm{m}\mathrm{eV}\mathrm{/}\mathrm{\text{Å}}\mathrm{)}$ $\downarrow$ |  $\mathrm{(}\mathrm{m}\mathrm{eV}\mathrm{)}$ $\downarrow$
GemNet-OC-L-E (Gasteiger et al., 2022) | $7.534\,085$ | $238.930\,850\,000\,000\,02$ | $22.068\,625$ | $230.174\,375$ | $20.9918$ | -
GemNet-OC-L-F (Gasteiger et al., 2022) | $3.178\,525$ | $252.255\,025$ | $19.9955$ | $241.134\,825$ | $19.012\,375$ | -
GemNet-OC-L-F+E (Gasteiger et al., 2022) | - | - | - | - | - | $347.634\,875$
SCN L=6 K=16 (4-tap 2-band) (Zitnick et al., 2022) | - | - | - | $228$ | $17.8$ | $328$
SCN L=8 K=20 (Zitnick et al., 2022) | - | - | - | $237$ | $17.2$ | $321$
eSCN L=6 K=20 (Passaro & Zitnick, 2023) | $2.9$ | $242.955\,993\,6$ | $17.090\,926\,48$ | $236$ | $16.2$ | $323$
EquiformerV2 ($\lambda_{E}=4$, $31$M) (Liao et al., 2023) | $7.1$ | $232$ | $16.26$ | $228$ | $15.51$ | $315$
EquiformerV2 ($\lambda_{E}=2$, $153$M) (Liao et al., 2023) | $1.8$ | $230$ | $14.6$ | $227.0$ | $13.8$ | $311$
EquiformerV2 ($\lambda_{E}=4$, $153$M) (Liao et al., 2023) | $1.8$ | $227$ | $15.04$ | $219$ | $14.20$ | $309$
EquiformerV2 + DeNS ($\lambda_{E}=4$, $160$M) | $1.8$ | 221 | 14.2 | 216 | 13.4 | 308

Table 2:  OC20 results on S2EF validation and test splits and IS2RE test split
when trained on OC20 S2EF-All+MD split. Throughput is reported as the number
of structures processed per GPU-second during training and measured on V100
GPUs.

####  4.1.2 Main results

We train EquiformerV2 ($160$M) with DeNS on the S2EF-All+MD split of OC20. The
model follows the same configuration as EquiformerV2 ($153$M) trained without
DeNS, and the additional parameters are due to force encoding and one
additional block of equivariant graph attention for noise prediction. We
report results in Table 2. All test results are computed via the EvalAI
evaluation server111eval.ai/web/challenges/challenge-page/712. EquiformerV2
trained with DeNS achieves better S2EF results and comparable IS2RE results,
setting the new state-of-the-art results. The improvement is not as
significant as that on OC20 S2EF-2M, OC22 (Section 4.2) and MD17 (Section 4.3)
datasets since the OC20 S2EF-All+MD training set contains much more structures
along relaxation trajectories, making new 3D geometries generated by DeNS less
helpful. However, DeNS is valuable because most datasets are not as large as
OC20 S2EF-All+MD dataset (about $172$M structures in the training set) but
have sizes closer to OC20 S2EF-2M ($2$M structures), OC22 ($8.2$M structures),
and MD17 ($950$ structures) datasets.

###  4.2 OC22 Dataset

##### Dataset and Tasks.

The Open Catalyst 2022 (OC22) dataset (Tran* et al., 2022) focuses on oxide
electrocatalysis and consists of about $62$k DFT relaxations obtained with
Perdew-Burke-Ernzerhof (PBE) functional calculations. One crucial difference
between OC22 and OC20 is that the energy targets in OC22 are DFT total
energies. DFT total energies are harder to predict but are the most general
and closest to a DFT surrogate, offering the flexibility to study property
prediction beyond adsorption energies. Analogous to the task definitions in
OC20, the primary tasks in OC22 are S2EF-Total and IS2RE-Total. We train
models on the OC22 S2EF-Total dataset, which has $8.2$M structures, and
evaluate them on energy and force MAE on the S2EF-total validation and test
splits. Then, we use these models to perform relaxations starting from initial
structures in the IS2RE-Total test split and evaluate the predicted relaxed
energies on energy MAE.

##### Training Details.

Please refer to Section C.1 for details on architectures, hyper-parameters and
training time.

##### Results.

We use EquiformerV2 with energy coefficient $\lambda_{E}=4$ and force
coefficient $\lambda_{F}=100$ to demonstrate how DeNS can further improve
state-of-the-art results and summarize the comparison in Table 3. Compared to
EquiformerV2 ($\lambda_{E}=4$, $\lambda_{F}=100$) trained without DeNS,
EquiformerV2 trained with DeNS consistently achieves better energy and force
MAE across all the S2EF-Total validation and test splits, with $9.6\%$ to
$16.7\%$ improvement in energy MAE and $5.1\%$ to $10.4\%$ improvement in
force MAE. For IS2RE-Total, EquiformerV2 trained with DeNS achieves the best
energy MAE results. The improvement on IS2RE-Total from training with DeNS on
only OC22 is greater than that of training on the much larger OC20 and OC22
datasets in previous works. Specifically, training GemNet-OC on OC20 and OC22
datasets (about $134$M + $8.2$M structures) improves IS2RE-Total energy MAE ID
by $129$meV and OOD by $50$meV compared to training GemNet-OC on only OC22
dataset ($8.2$M structures). Compared to training without DeNS, training
EquiformerV2 with DeNS improves ID by $151$meV and OOD by $167$meV. Thus,
training with DeNS clearly improves sample efficiency and performance on OC22.

|  | S2EF-Total validation | S2EF-Total test | IS2RE-Total test
---|---|---|---|---
| Number of | Energy MAE (meV) $\downarrow$ | Force MAE (meV/Å) $\downarrow$ | Energy MAE (meV) $\downarrow$ | Force MAE (meV/Å) $\downarrow$ | Energy MAE (meV) $\downarrow$
Model | parameters | ID | OOD | ID | OOD | ID | OOD | ID | OOD | ID | OOD
GemNet-OC (Gasteiger et al., 2022) |  $39$M | 545 | 1011 | 30 | 40 | 374 | 829 | 29.4 | 39.6 | 1329 | 1584
GemNet-OC (trained on OC20 + OC22) (Gasteiger et al., 2022) |  $39$M | 464 | 859 | 27 | 34 | 311 | 689 | 26.9 | 34.2 | 1200 | 1534
eSCN (Passaro & Zitnick, 2023) |  $200$M | - | - | - | - | 350 | 789 | 23.8 | 35.7 | - | -
EquiformerV2 ($\lambda_{E}=1,\lambda_{F}=1$) (Liao et al., 2023) |  $122$M | 343 | 580 | 24.42 | 34.31 | 182.8 | 677.4 | 22.98 | 35.57 | 1084 | 1444
EquiformerV2 ($\lambda_{E}=4,\lambda_{F}=100$) (Liao et al., 2023) |  $122$M | 433 | 629 | 22.88 | 30.70 | 263.7 | 659.8 | 21.58 | 32.65 | 1119 | 1440
EquiformerV2 + DeNS ($\lambda_{E}=4,\lambda_{F}=100$) |  $127$M | 391 | 524 | 21.13 | 27.82 | 233.9 | 589.9 | 20.47 | 30.06 | 968 | 1273

Table 3: OC22 results on S2EF-Total validation and test splits and IS2RE-Total
test split. Models are trained on the OC22 S2EF-Total training split unless
otherwise noted.

| Aspirin | Benzene | Ethanol | Malonaldehyde | Naphthalene | Salicylic acid | Toluene | Uracil | Training time | Number of
---|---|---|---|---|---|---|---|---|---|---
Model | energy | forces | energy | forces | energy | forces | energy | forces | energy | forces | energy | forces | energy | forces | energy | forces | (GPU-hours)$\downarrow$ | parameters
SchNet (Schütt et al., 2017) | 16.0 | 58.5 | 3.5 | 13.4 | 3.5 | 16.9 | 5.6 | 28.6 | 6.9 | 25.2 | 8.7 | 36.9 | 5.2 | 24.7 | 6.1 | 24.3 | - | -
DimeNet (Gasteiger et al., 2020) | 8.8 | 21.6 | 3.4 | 8.1 | 2.8 | 10.0 | 4.5 | 16.6 | 5.3 | 9.3 | 5.8 | 16.2 | 4.4 | 9.4 | 5.0 | 13.1 | - | -
PaiNN (Schütt et al., 2021) | 6.9 | 14.7 | - | - | 2.7 | 9.7 | 3.9 | 13.8 | 5.0 | 3.3 | 4.9 | 8.5 | 4.1 | 4.1 | 4.5 | 6.0 | - | -
TorchMD-NET (Thölke & Fabritiis, 2022) | 5.3 | 11.0 | 2.5 | 8.5 | 2.3 | 4.7 | 3.3 | 7.3 | 3.7 | 2.6 | 4.0 | 5.6 | 3.2 | 2.9 | 4.1 | 4.1 | - | -
NequIP ($L_{max}=3$) (Batzner et al., 2022) | 5.7 | 8.0 | - | - | 2.2 | 3.1 | 3.3 | 5.6 | 4.9 | 1.7 | 4.6 | 3.9 | 4.0 | 2.0 | 4.5 | 3.3 | - | -
Equiformer ($L_{max}=2$) | 5.3 | 7.2 | 2.2 | 6.6 | 2.2 | 3.1 | 3.3 | 5.8 | 3.7 | 2.1 | 4.5 | 4.1 | 3.8 | 2.1 | 4.3 | 3.3 | 17 |  $3.50$M
Equiformer ($L_{max}=3$) | 5.3 | 6.6 | 2.5 | 8.1 | 2.2 | 2.9 | 3.2 | 5.4 | 4.4 | 2.0 | 4.3 | 3.9 | 3.7 | 2.1 | 4.3 | 3.4 | 59 |  $5.50$M
Equiformer ($L_{max}=2$) + DeNS | 5.1 | 5.7 | 2.3 | 6.1 | 2.2 | 2.6 | 3.2 | 4.4 | 3.7 | 1.7 | 4.3 | 3.9 | 3.5 | 1.9 | 4.2 | 3.3 | 19 |  $4.00$M
Equiformer ($L_{max}=3$) + DeNS | 5.0 | 5.2 | 2.3 | 6.1 | 2.2 | 2.4 | 3.2 | 4.1 | 3.7 | 1.6 | 4.2 | 3.2 | 3.4 | 1.8 | 4.1 | 2.9 | 61 |  $6.35$M

Table 4: Mean absolute error results on the MD17 testing set. Energy and force
are in units of meV and meV/Å. We additionally report the time of training
different Equiformer models averaged over all molecules and the numbers of
parameters to show that the proposed DeNS can improve performance with minimal
overhead.

| Method |  | Aspirin | Benzene | Ethanol | Malonaldehyde | Naphthalene | Salicylic acid | Toluene | Uracil
---|---|---|---|---|---|---|---|---|---|---
Index | Attention | Layer normalization | DeNS |  | energy | forces | energy | forces | energy | forces | energy | forces | energy | forces | energy | forces | energy | forces | energy | forces
1 | ✓ | ✓ |  |  | 5.3 | 7.2 | 2.2 | 6.6 | 2.2 | 3.1 | 3.3 | 5.8 | 3.7 | 2.1 | 4.5 | 4.1 | 3.8 | 2.1 | 4.3 | 3.3
2 | ✓ | ✓ | ✓ |  | 5.1 | 5.7 | 2.3 | 6.1 | 2.2 | 2.6 | 3.2 | 4.4 | 3.7 | 1.7 | 4.3 | 3.9 | 3.5 | 1.9 | 4.2 | 3.3
3 |  | ✓ |  |  | 5.2 | 7.7 | 2.4 | 6.2 | 2.3 | 3.9 | 3.3 | 6.2 | 3.8 | 2.2 | 4.1 | 4.7 | 3.3 | 2.4 | 4.2 | 4.4
4 |  | ✓ | ✓ |  | 5.2 | 6.1 | 2.4 | 6.1 | 2.2 | 2.9 | 3.2 | 5.1 | 3.7 | 1.7 | 4.2 | 3.9 | 3.4 | 2.0 | 4.2 | 3.4
5 |  |  |  |  | 5.3 | 9.3 | 2.4 | 9.2 | 2.3 | 4.5 | 3.4 | 8.2 | 3.7 | 2.4 | 4.2 | 5.5 | 3.3 | 2.9 | 4.2 | 4.8
6 |  |  | ✓ |  | 5.2 | 7.3 | 2.4 | 8.1 | 2.2 | 3.4 | 3.4 | 6.7 | 3.7 | 1.9 | 4.2 | 4.4 | 3.4 | 2.2 | 4.2 | 3.8

Table 5: Mean absolute error results of variants of Equiformer ($L_{max}=2$)
without attention and layer normalization on the MD17 testing set. Energy and
force are in units of meV and meV/Å. Index 1 and Index 2 correspond to
“Equiformer ($L_{max}=2$)” and “Equiformer ($L_{max}=2$) + DeNS” in Table 4.

###  4.3 MD17 Dataset

##### Dataset.

The MD17 dataset (Chmiela et al., 2017; Schütt et al., 2017; Chmiela et al.,
2018) consists of molecular dynamics simulations of small organic molecules.
The task is to predict the energy and forces of these non-equilibrium
molecules. We use $950$ and $50$ different configurations for training and
validation sets and the rest for the testing set.

##### Training Details.

Please refer to Section D.2 for additinoal implementation details of DeNS,
hyper-parameters and training time.

##### Results.

We train Equiformer ($L_{max}=2$) (Liao & Smidt, 2023) with DeNS based on
their implementation and summarize the results in Table 4. DeNS improves the
results on all molecules, and Equiformer ($L_{max}=3$) trained with DeNS
achieves overall best results. Particularly, Equiformer ($L_{max}=2$) trained
with DeNS acheives better results on all the tasks and requires $3.1\times$
less training time than Equiformer ($L_{max}=3$) trained without DeNS. This
demonstrates that for this small dataset, training an auxiliary task and using
data augmentation are more efficient and result in larger performance gain
than increasing $L_{max}$ from $2$ to $3$. Additionally, we find that the
gains from training DeNS as an auxiliary task are comparable to pre-training.
Zaidi et al. (2023) uses TorchMD-NET (Thölke & Fabritiis, 2022) pre-trained on
the PCQM4Mv2 dataset and report results on Aspirin. Their improvement on force
MAE is about $17.2\%$ (Table 3 in Zaidi et al. (2023)). Training Equiformer
with DeNS results in $20.8\%$ improvement on force MAE without relying on
another dataset. Note that we only increase training time by $10.5\%$ while
their method takes much more time since PCQM4Mv2 dataset is more than
$3000\times$ larger than MD17. Moreover, we also train variants of Equiformer
($L_{max}=2$) by removing attention and layer normalization to investigate the
performance gain of DeNS on different network archtiectures. The results are
summarized in Table 5, and DeNS improves all the models. We note that
Equiformer without attention and layer normalization reduces to SEGNN
(Brandstetter et al., 2022) but with a better radial basis function. Since the
models cover many variants of equivariant networks, this suggests that DeNS is
general and can be helpful to many equivariant networks.

##### Comparison between DeNS with and without Force Encoding.

We use Equiformer ($L_{max}=2$) to compare DeNS with and without force
encoding and demonstrate that force encoding is critical. The results are
summarized in Table 6. DeNS with force encoding (Index 3) achieves the best
results across all molecules. Compared to training without denoising (Index
1), DeNS without force encoding (Index 2) only achieves slightly better
results on some molecules (i.e., Benzene, Malondaldehyde, and Toluene) and
much worse results on others. For molecules on which DeNS without force
encoding is helpful, adding force encoding can achieve even better results.
For others, force encoding is indispensable for DeNS to be effective.

|  |  | Aspirin | Benzene | Ethanol | Malonaldehyde | Naphthalene | Salicylic acid | Toluene | Uracil
---|---|---|---|---|---|---|---|---|---|---
Index | Method |  | energy | forces | energy | forces | energy | forces | energy | forces | energy | forces | energy | forces | energy | forces | energy | forces
1 | Equiformer ($L_{max}=2$) |  | 5.3 | 7.2 | 2.2 | 6.6 | 2.2 | 3.1 | 3.3 | 5.8 | 3.7 | 2.1 | 4.5 | 4.1 | 3.8 | 2.1 | 4.3 | 3.3
2 | Equiformer ($L_{max}=2$) + DeNS without force encoding |  | 8.6 | 9.1 | 2.3 | 6.3 | 2.3 | 3.3 | 3.2 | 5.8 | 7.7 | 6.1 | 5.2 | 10.6 | 3.7 | 2.0 | 5.5 | 6.5
3 | Equiformer ($L_{max}=2$) + DeNS with force encoding |  | 5.1 | 5.7 | 2.3 | 6.1 | 2.2 | 2.6 | 3.2 | 4.4 | 3.7 | 1.7 | 4.3 | 3.9 | 3.5 | 1.9 | 4.2 | 3.3

Table 6:  Comparison between DeNS with and without force encoding on MD17
dataset. Mean absolute error results are evaluated on the MD17 testing set.
Energy and force are in units of meV and meV/Å. Index 1 and Index 3 correspond
to “Equiformer ($L_{max}=2$)” and “Equiformer ($L_{max}=2$) + DeNS” in Table
4.

##  5 Conclusion

In this paper, we propose to use denoising non-equilibrium structures (DeNS)
as an auxiliary task to better leverage training data and improve performance
of original tasks of energy and force predictions. Denoising non-equilibrium
structures can be an ill-posed problem since there are many possible target
structures. To address the issue, we propose to take the forces of original
structures as inputs to specify which non-equilibrium structures we are
denoising. With force encoding, DeNS successfully improves the performance of
original tasks when it is used as an auxiliary task. We conduct extensive
experiments on OC20, OC22 and MD17 datasets to demonstrate that DeNS can boost
the performance of energy and force predictions with minimal increase in
training cost and is applicable to many equivariant networks.

## Acknowledgement

We thank Larry Zitnick, Xinlei Chen, Zachary Ulissi, Muhammed Shuaibi, Saro
Passaro, Anuroop Sriram, and Brandon Wood for helpful discussions. We
acknowledge the MIT SuperCloud and Lincoln Laboratory Supercomputing Center
(Reuther et al., 2018) for providing high performance computing and
consultation resources that have contributed to the research results reported
within this paper.

Yi-Lun Liao and Tess Smidt were supported by DOE ICDI grant DE-SC0022215.

## References

  * Batatia et al. (2022) Ilyes Batatia, David Peter Kovacs, Gregor N. C. Simm, Christoph Ortner, and Gabor Csanyi.  MACE: Higher order equivariant message passing neural networks for fast and accurate force fields.  In _Advances in Neural Information Processing Systems (NeurIPS)_ , 2022.
  * Batzner et al. (2022) Simon Batzner, Albert Musaelian, Lixin Sun, Mario Geiger, Jonathan P. Mailoa, Mordechai Kornbluth, Nicola Molinari, Tess E. Smidt, and Boris Kozinsky.  E(3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials.  _Nature Communications_ , 13(1), May 2022.  doi: 10.1038/s41467-022-29939-5.  URL https://doi.org/10.1038/s41467-022-29939-5.
  * Brandstetter et al. (2022) Johannes Brandstetter, Rob Hesselink, Elise van der Pol, Erik J Bekkers, and Max Welling.  Geometric and physical quantities improve e(3) equivariant message passing.  In _International Conference on Learning Representations (ICLR)_ , 2022.  URL https://openreview.net/forum?id=_xwr8gOBeV1.
  * Brown et al. (2020) Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei.  Language models are few-shot learners.  In _Advances in Neural Information Processing Systems (NeurIPS)_ , 2020.
  * Chanussot* et al. (2021) Lowik Chanussot*, Abhishek Das*, Siddharth Goyal*, Thibaut Lavril*, Muhammed Shuaibi*, Morgane Riviere, Kevin Tran, Javier Heras-Domingo, Caleb Ho, Weihua Hu, Aini Palizhati, Anuroop Sriram, Brandon Wood, Junwoong Yoon, Devi Parikh, C. Lawrence Zitnick, and Zachary Ulissi.  Open catalyst 2020 (oc20) dataset and community challenges.  _ACS Catalysis_ , 2021.  doi: 10.1021/acscatal.0c04525.
  * Chmiela et al. (2017) Stefan Chmiela, Alexandre Tkatchenko, Huziel E. Sauceda, Igor Poltavsky, Kristof T. Schütt, and Klaus-Robert Müller.  Machine learning of accurate energy-conserving molecular force fields.  _Science Advances_ , 3(5):e1603015, 2017.  doi: 10.1126/sciadv.1603015.  URL https://www.science.org/doi/abs/10.1126/sciadv.1603015.
  * Chmiela et al. (2018) Stefan Chmiela, Huziel E. Sauceda, Klaus-Robert Müller, and Alexandre Tkatchenko.  Towards exact molecular dynamics simulations with machine-learned force fields.  _Nature Communications_ , 9(1), sep 2018.  doi: 10.1038/s41467-018-06169-2.
  * Dehghani et al. (2023) Mostafa Dehghani, Josip Djolonga, Basil Mustafa, Piotr Padlewski, Jonathan Heek, Justin Gilmer, Andreas Steiner, Mathilde Caron, Robert Geirhos, Ibrahim Alabdulmohsin, Rodolphe Jenatton, Lucas Beyer, Michael Tschannen, Anurag Arnab, Xiao Wang, Carlos Riquelme, Matthias Minderer, Joan Puigcerver, Utku Evci, Manoj Kumar, Sjoerd van Steenkiste, Gamaleldin F. Elsayed, Aravindh Mahendran, Fisher Yu, Avital Oliver, Fantine Huot, Jasmijn Bastings, Mark Patrick Collier, Alexey Gritsenko, Vighnesh Birodkar, Cristina Vasconcelos, Yi Tay, Thomas Mensink, Alexander Kolesnikov, Filip Pavetić, Dustin Tran, Thomas Kipf, Mario Lučić, Xiaohua Zhai, Daniel Keysers, Jeremiah Harmsen, and Neil Houlsby.  Scaling vision transformers to 22 billion parameters.  _arXiv preprint arXiv:2302.05442_ , 2023.
  * Devlin et al. (2019) Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.  BERT: Pre-training of deep bidirectional transformers for language understanding.  _arxiv preprint arxiv:1810.04805_ , 2019.
  * Dosovitskiy et al. (2021) Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby.  An image is worth 16x16 words: Transformers for image recognition at scale.  In _International Conference on Learning Representations (ICLR)_ , 2021.  URL https://openreview.net/forum?id=YicbFdNTTy.
  * Feng et al. (2023a) Rui Feng, Qi Zhu, Huan Tran, Binghong Chen, Aubrey Toland, Rampi Ramprasad, and Chao Zhang.  May the force be with you: Unified force-centric pre-training for 3d molecular conformations.  _arXiv preprint arXiv:2308.14759_ , 2023a.
  * Feng et al. (2023b) Shikun Feng, Yuyan Ni, Yanyan Lan, Zhi-Ming Ma, and Wei-Ying Ma.  Fractional denoising for 3D molecular pre-training.  In _International Conference on Machine Learning (ICML)_ , 2023b.
  * Fuchs et al. (2020) Fabian Fuchs, Daniel E. Worrall, Volker Fischer, and Max Welling.  Se(3)-transformers: 3d roto-translation equivariant attention networks.  In _Advances in Neural Information Processing Systems (NeurIPS)_ , 2020.
  * Gasteiger et al. (2020) Johannes Gasteiger, Janek Groß, and Stephan Günnemann.  Directional message passing for molecular graphs.  In _International Conference on Learning Representations (ICLR)_ , 2020.
  * Gasteiger et al. (2022) Johannes Gasteiger, Muhammed Shuaibi, Anuroop Sriram, Stephan Günnemann, Zachary Ulissi, C Lawrence Zitnick, and Abhishek Das.  GemNet-OC: Developing Graph Neural Networks for Large and Diverse Molecular Simulation Datasets.  _Transactions on Machine Learning Research (TMLR)_ , 2022.
  * Geiger & Smidt (2022) Mario Geiger and Tess Smidt.  e3nn: Euclidean neural networks.  _arXiv preprint arXiv:2207.09453_ , 2022.
  * Geiger et al. (2022) Mario Geiger, Tess Smidt, Alby M., Benjamin Kurt Miller, Wouter Boomsma, Bradley Dice, Kostiantyn Lapchevskyi, Maurice Weiler, Michał Tyszkiewicz, Simon Batzner, Dylan Madisetti, Martin Uhrin, Jes Frellsen, Nuri Jung, Sophia Sanborn, Mingjian Wen, Josh Rackers, Marcel Rød, and Michael Bailey.  e3nn/e3nn: 2022-04-13, April 2022.  URL https://doi.org/10.5281/zenodo.6459381.
  * Gilmer et al. (2017) Justin Gilmer, Samuel S. Schoenholz, Patrick F. Riley, Oriol Vinyals, and George E. Dahl.  Neural message passing for quantum chemistry.  In _International Conference on Machine Learning (ICML)_ , 2017.
  * Godwin et al. (2022) Jonathan Godwin, Michael Schaarschmidt, Alexander L Gaunt, Alvaro Sanchez-Gonzalez, Yulia Rubanova, Petar Veličković, James Kirkpatrick, and Peter Battaglia.  Simple GNN regularisation for 3d molecular property prediction and beyond.  In _International Conference on Learning Representations (ICLR)_ , 2022.  URL https://openreview.net/forum?id=1wVvweK3oIb.
  * Hammer et al. (1999) B. Hammer, L. B. Hansen, and J. K. Nørskov.  Improved adsorption energetics within density-functional theory using revised perdew-burke-ernzerhof functionals.  _Phys. Rev. B_ , 1999.
  * He et al. (2022) Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick.  Masked autoencoders are scalable vision learners.  In _IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_ , 2022.
  * Huang et al. (2016) Gao Huang, Yu Sun, Zhuang Liu, Daniel Sedra, and Kilian Q. Weinberger.  Deep networks with stochastic depth.  In _European Conference on Computer Vision (ECCV)_ , 2016.
  * Jiao et al. (2022) Rui Jiao, Jiaqi Han, Wenbing Huang, Yu Rong, and Yang Liu.  Energy-motivated equivariant pretraining for 3d molecular graphs.  _arXiv preprint arXiv:2207.08824_ , 2022.
  * Jing et al. (2021) Bowen Jing, Stephan Eismann, Patricia Suriana, Raphael John Lamarre Townshend, and Ron Dror.  Learning from protein structure with geometric vector perceptrons.  In _International Conference on Learning Representations (ICLR)_ , 2021.  URL https://openreview.net/forum?id=1YLJDvSx6J4.
  * Kondor et al. (2018) Risi Kondor, Zhen Lin, and Shubhendu Trivedi.  Clebsch–gordan nets: a fully fourier space spherical convolutional neural network.  In _Advances in Neural Information Processing Systems 32_ , pp. 10117–10126, 2018.
  * Lan et al. (2022) Janice Lan, Aini Palizhati, Muhammed Shuaibi, Brandon M Wood, Brook Wander, Abhishek Das, Matt Uyttendaele, C Lawrence Zitnick, and Zachary W Ulissi.  AdsorbML: Accelerating adsorption energy calculations with machine learning.  _arXiv preprint arXiv:2211.16486_ , 2022.
  * Le et al. (2022) Tuan Le, Frank Noé, and Djork-Arné Clevert.  Equivariant graph attention networks for molecular property prediction.  _arXiv preprint arXiv:2202.09891_ , 2022.
  * Liao & Smidt (2023) Yi-Lun Liao and Tess Smidt.  Equiformer: Equivariant graph attention transformer for 3d atomistic graphs.  In _International Conference on Learning Representations (ICLR)_ , 2023.  URL https://openreview.net/forum?id=KwmPfARgOTD.
  * Liao et al. (2023) Yi-Lun Liao, Brandon Wood, Abhishek Das*, and Tess Smidt*.  EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree Representations.  _arxiv preprint arxiv:2306.12059_ , 2023.
  * Liu et al. (2023) Shengchao Liu, Hongyu Guo, and Jian Tang.  Molecular geometry pretraining with SE(3)-invariant denoising distance matching.  In _International Conference on Learning Representations (ICLR)_ , 2023.  URL https://openreview.net/forum?id=CjTHVo1dvR.
  * Miller et al. (2020) Benjamin Kurt Miller, Mario Geiger, Tess E. Smidt, and Frank Noé.  Relevance of rotationally equivariant convolutions for predicting molecular properties.  _arxiv preprint arxiv:2008.08461_ , 2020.
  * Musaelian et al. (2022) Albert Musaelian, Simon Batzner, Anders Johansson, Lixin Sun, Cameron J. Owen, Mordechai Kornbluth, and Boris Kozinsky.  Learning local equivariant representations for large-scale atomistic dynamics.  _arxiv preprint arxiv:2204.05249_ , 2022.
  * Musaelian et al. (2023) Albert Musaelian, Anders Johansson, Simon Batzner, and Boris Kozinsky.  Scaling the leading accuracy of deep equivariant models to biomolecular simulations of realistic size.  _arXiv preprint arXiv:2304.10061_ , 2023.
  * Nakata & Shimazaki (2017) Maho Nakata and Tomomi Shimazaki.  Pubchemqc project: A large-scale first-principles electronic structure database for data-driven chemistry.  _Journal of chemical information and modeling_ , 57 6:1300–1308, 2017.
  * Passaro & Zitnick (2023) Saro Passaro and C Lawrence Zitnick.  Reducing SO(3) Convolutions to SO(2) for Efficient Equivariant GNNs.  In _International Conference on Machine Learning (ICML)_ , 2023.
  * Rackers et al. (2023) Joshua A Rackers, Lucas Tecot, Mario Geiger, and Tess E Smidt.  A recipe for cracking the quantum scaling limit with machine learned electron densities.  _Machine Learning: Science and Technology_ , 4(1):015027, feb 2023.  doi: 10.1088/2632-2153/acb314.  URL https://dx.doi.org/10.1088/2632-2153/acb314.
  * Ramakrishnan et al. (2014) Raghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp, and O Anatole von Lilienfeld.  Quantum chemistry structures and properties of 134 kilo molecules.  _Scientific Data_ , 1, 2014.
  * Reuther et al. (2018) Albert Reuther, Jeremy Kepner, Chansup Byun, Siddharth Samsi, William Arcand, David Bestor, Bill Bergeron, Vijay Gadepally, Michael Houle, Matthew Hubbell, Michael Jones, Anna Klein, Lauren Milechin, Julia Mullen, Andrew Prout, Antonio Rosa, Charles Yee, and Peter Michaleas.  Interactive supercomputing on 40,000 cores for machine learning and data analysis.  In _2018 IEEE High Performance extreme Computing Conference (HPEC)_ , pp. 1–6. IEEE, 2018.
  * Ruddigkeit et al. (2012) Lars Ruddigkeit, Ruud van Deursen, Lorenz C. Blum, and Jean-Louis Reymond.  Enumeration of 166 billion organic small molecules in the chemical universe database gdb-17.  _Journal of Chemical Information and Modeling_ , 52(11):2864–2875, 2012.  doi: 10.1021/ci300415d.  URL https://doi.org/10.1021/ci300415d.  PMID: 23088335.
  * Sanchez-Gonzalez et al. (2020) Alvaro Sanchez-Gonzalez, Jonathan Godwin, Tobias Pfaff, Rex Ying, Jure Leskovec, and Peter W. Battaglia.  Learning to simulate complex physics with graph networks.  In _International Conference on Machine Learning (ICML)_ , 2020.
  * Satorras et al. (2021) Víctor Garcia Satorras, Emiel Hoogeboom, and Max Welling.  E(n) equivariant graph neural networks.  In _International Conference on Machine Learning (ICML)_ , 2021.
  * Schütt et al. (2017) K. T. Schütt, P.-J. Kindermans, H. E. Sauceda, S. Chmiela, A. Tkatchenko, and K.-R. Müller.  Schnet: A continuous-filter convolutional neural network for modeling quantum interactions.  In _Advances in Neural Information Processing Systems (NeurIPS)_ , 2017.
  * Schütt et al. (2017) Kristof T. Schütt, Farhad Arbabzadah, Stefan Chmiela, Klaus R. Müller, and Alexandre Tkatchenko.  Quantum-chemical insights from deep tensor neural networks.  _Nature Communications_ , 8(1), jan 2017.  doi: 10.1038/ncomms13890.
  * Schütt et al. (2021) Kristof T. Schütt, Oliver T. Unke, and Michael Gastegger.  Equivariant message passing for the prediction of tensorial properties and molecular spectra.  In _International Conference on Machine Learning (ICML)_ , 2021.
  * Smith et al. (2017) Justin S. Smith, Olexandr Isayev, and Adrian E. Roitberg.  Ani-1: A data set of 20m off-equilibrium dft calculations for organic molecules.  _arXiv preprint arXiv:1708.04987_ , 2017.
  * Smith et al. (2018) Justin S. Smith, Benjamin Tyler Nebgen, Nicholas Lubbers, Olexandr Isayev, and Adrian E. Roitberg.  Less is more: sampling chemical space with active learning.  _The Journal of chemical physics_ , 2018.
  * Song & Ermon (2019) Yang Song and Stefano Ermon.  Generative modeling by estimating gradients of the data distribution.  _Advances in Neural Information Processing Systems (NeurIPS)_ , 2019.
  * Song & Ermon (2020) Yang Song and Stefano Ermon.  Improved techniques for training score-based generative models.  _Advances in Neural Information Processing Systems (NeurIPS)_ , 2020.
  * Srivastava et al. (2014) Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov.  Dropout: A simple way to prevent neural networks from overfitting.  _Journal of Machine Learning Research_ , 15(56):1929–1958, 2014.  URL http://jmlr.org/papers/v15/srivastava14a.html.
  * Tay et al. (2023) Yi Tay, Mostafa Dehghani, Vinh Q. Tran, Xavier Garcia, Jason Wei, Xuezhi Wang, Hyung Won Chung, Dara Bahri, Tal Schuster, Steven Zheng, Denny Zhou, Neil Houlsby, and Donald Metzler.  UL2: Unifying language learning paradigms.  In _International Conference on Learning Representations (ICLR)_ , 2023.  URL https://openreview.net/forum?id=6ruVLB727MC.
  * Thölke & Fabritiis (2022) Philipp Thölke and Gianni De Fabritiis.  Equivariant transformers for neural network based molecular potentials.  In _International Conference on Learning Representations (ICLR)_ , 2022.  URL https://openreview.net/forum?id=zNHzqZ9wrRB.
  * Thomas et al. (2018) Nathaniel Thomas, Tess E. Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley.  Tensor field networks: Rotation- and translation-equivariant neural networks for 3d point clouds.  _arxiv preprint arXiv:1802.08219_ , 2018.
  * Townshend et al. (2020) Raphael J. L. Townshend, Brent Townshend, Stephan Eismann, and Ron O. Dror.  Geometric prediction: Moving beyond scalars.  _arXiv preprint arXiv:2006.14163_ , 2020.
  * Tran* et al. (2022) Richard Tran*, Janice Lan*, Muhammed Shuaibi*, Brandon Wood*, Siddharth Goyal*, Abhishek Das, Javier Heras-Domingo, Adeesh Kolluru, Ammar Rizvi, Nima Shoghi, Anuroop Sriram, Zachary Ulissi, and C. Lawrence Zitnick.  The open catalyst 2022 (oc22) dataset and challenges for oxide electrocatalysis.  _arXiv preprint arXiv:2206.08917_ , 2022.
  * Unke et al. (2021) Oliver Thorsten Unke, Mihail Bogojeski, Michael Gastegger, Mario Geiger, Tess Smidt, and Klaus Robert Muller.  SE(3)-equivariant prediction of molecular wavefunctions and electronic densities.  In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), _Advances in Neural Information Processing Systems (NeurIPS)_ , 2021.  URL https://openreview.net/forum?id=auGY2UQfhSu.
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  In _Advances in Neural Information Processing Systems (NeurIPS)_ , 2017.
  * Vincent et al. (2008) Pascal Vincent, H. Larochelle, Yoshua Bengio, and Pierre-Antoine Manzagol.  Extracting and composing robust features with denoising autoencoders.  In _International Conference on Machine Learning (ICML)_ , 2008.
  * Vincent et al. (2010) Pascal Vincent, H. Larochelle, Isabelle Lajoie, Yoshua Bengio, and Pierre-Antoine Manzagol.  Stacked denoising autoencoders: Learning useful representations in a deep network with a local denoising criterion.  _Journal of Machine Learning Research_ , 2010.
  * Wang et al. (2023) Yuyang Wang, Chang Xu, Zijie Li, and Amir Barati Farimani.  Denoise pretraining on nonequilibrium molecules for accurate and transferable neural potentials.  _Journal of Chemical Theory and Computation_ , 2023.
  * Weiler et al. (2018) Maurice Weiler, Mario Geiger, Max Welling, Wouter Boomsma, and Taco Cohen.  3D Steerable CNNs: Learning Rotationally Equivariant Features in Volumetric Data.  In _Advances in Neural Information Processing Systems 32_ , pp. 10402–10413, 2018.
  * Zaidi et al. (2023) Sheheryar Zaidi, Michael Schaarschmidt, James Martens, Hyunjik Kim, Yee Whye Teh, Alvaro Sanchez-Gonzalez, Peter Battaglia, Razvan Pascanu, and Jonathan Godwin.  Pre-training via denoising for molecular property prediction.  In _International Conference on Learning Representations (ICLR)_ , 2023.
  * Zhang et al. (2018) Linfeng Zhang, Jiequn Han, Han Wang, Roberto Car, and Weinan E.  Deep potential molecular dynamics: A scalable model with the accuracy of quantum mechanics.  _Phys. Rev. Lett._ , 120:143001, Apr 2018.  doi: 10.1103/PhysRevLett.120.143001.  URL https://link.aps.org/doi/10.1103/PhysRevLett.120.143001.
  * Zitnick et al. (2022) Larry Zitnick, Abhishek Das, Adeesh Kolluru, Janice Lan, Muhammed Shuaibi, Anuroop Sriram, Zachary Ulissi, and Brandon Wood.  Spherical channels for modeling atomic interactions.  In _Advances in Neural Information Processing Systems (NeurIPS)_ , 2022.

## Appendix

  * A Related Works

    1. A.1 $SE(3)/E(3)$-Equivariant Networks

    2. A.2 Comparison to previous works on denoising

  * B Details of experiments on OC20

    1. B.1 Training details

  * C Details of experiments on OC22

    1. C.1 Training details

  * D Details of experiments on MD17

    1. D.1 Additional details of DeNS

    2. D.2 Training details

  * E Pseudocode for training with DeNS

  * F Visualization of corrupted structures

##  Appendix A Related Works

###  A.1 SE(3)/E(3)-Equivariant Networks

Equivariant neural networks (Thomas et al., 2018; Kondor et al., 2018; Weiler
et al., 2018; Fuchs et al., 2020; Miller et al., 2020; Townshend et al., 2020;
Batzner et al., 2022; Jing et al., 2021; Schütt et al., 2021; Satorras et al.,
2021; Brandstetter et al., 2022; Thölke & Fabritiis, 2022; Le et al., 2022;
Musaelian et al., 2022; Batatia et al., 2022; Liao & Smidt, 2023; Passaro &
Zitnick, 2023; Liao et al., 2023) use vector spaces of irreducible
representations (irreps) for equivariant irreps features and adopt equivariant
operations such as tensor products to achieve equivariance to 3D rotation.
Previous works differ in which equivariant operations are used and the
combination of those operations. TFN (Thomas et al., 2018) and NequIP (Batzner
et al., 2022) use tensor products for equivariant graph convolution with
linear messages, with the latter utilizing extra gate activation (Weiler et
al., 2018). SEGNN (Brandstetter et al., 2022) applies gate activation to
messages passing for non-linear messages (Gilmer et al., 2017; Sanchez-
Gonzalez et al., 2020). SE(3)-Transformer (Fuchs et al., 2020) adopts
equivariant dot product attention with linear messages. Equiformer (Liao &
Smidt, 2023) improves upon previous models by combining MLP attention and non-
linear messages and additionally introducing equivariant layer normalization
and regularizations like dropout (Srivastava et al., 2014) and stochastic
depth (Huang et al., 2016). However, these networks rely on compute-intensive
$SO(3)$ convolutions built from tensor products, and therefore they can only
use small values for maximum degrees $L_{max}$ of irreps features. eSCN
(Passaro & Zitnick, 2023) significantly reduces the complexity of $SO(3)$
convolutions by first rotating irreps features based on relative positions and
then applying $SO(2)$ linear layers, enabling higher values of $L_{max}$.
EquiformerV2 (Liao et al., 2023) adopts eSCN convolutions and proposes an
improved version of Equiformer to better leverage the power of higher
$L_{max}$, achieving the current state-of-the-art results on OC20 (Chanussot*
et al., 2021) and OC22 (Tran* et al., 2022) datasets.

We refer readers to the works (Liao & Smidt, 2023; Liao et al., 2023) for
detailed background on equivariant networks.

###  A.2 Comparison to previous works on denoising

We discuss the comparisons between previous works on denoising (Godwin et al.,
2022; Zaidi et al., 2023; Feng et al., 2023b; Wang et al., 2023) and this work
in chronological order as below.

Godwin et al. (2022) first proposes the idea of adding noise to 3D coordinates
and then using denoising as an auxiliary task. The auxiliary task is trained
along with the original task without relying on another large dataset. Their
approach requires known equilibrium structures and therefore is limited to QM9
(Ramakrishnan et al., 2014; Ruddigkeit et al., 2012) and OC20 IS2RE datasets
and can not be applied to force prediction such as OC20 S2EF dataset. For QM9,
all the structures are at equilibrium, and for OC20 IS2RE, the target of
denoising is the relaxed, equilibrium structure. Denoising without force
encoding is well-defined on both QM9 and OC20 IS2RE. In contrast, this work
proposes using force encoding to generalize their approach to non-equilibrium
structures, which have much larger datasets than equilibrium ones. Force
encoding can achieve better results on OC20 S2EF-2M dataset without any
overhead (Index 1 and Index 2 in Table 1(e)) and is indispensable on MD17
dataset (Section 4.3).

Zaidi et al. (2023) adopts the denoising approach proposed by Godwin et al.
(2022) as a pre-training method and therefore requires another large dataset
containing unlabelled equilibrium structures for pre-training. On the other
hand, Godwin et al. (2022) and this work use denoising along with the original
task and do not use any additional unlabeled data.

Feng et al. (2023b) follows the same practice of pre-training via denoising
(Zaidi et al., 2023) and proposes a different manner of adding noise.
Specifically, they separate noise into dihedral angle noise and coordinate
noise and only learn to predict coordinate noise. However, adding noise to
dihedral angles requires tools like RDKit to obtain rotatable bonds and cannot
be applied to other datasets like OC20 and OC22.

Although Zaidi et al. (2023) and Feng et al. (2023b) report results of force
prediction on MD17 dataset, they first pre-train models on PCQM4Mv2 dataset
(Nakata & Shimazaki, 2017) and then fine-tune the pre-trained models on MD17
dataset. We note that their setting is different from ours since we do not use
any dataset for pre-training. As for fine-tuning on MD17 dataset, Zaidi et al.
(2023) simply follows the same practice in standard supervised training. Feng
et al. (2023b) explores fine-tuning with objectives similar to Noisy Nodes
(Godwin et al., 2022), but the performance gain is much smaller than ours.
Concretely, in Table 5 in Feng et al. (2023b), the improvement on force
prediction on Aspirin is about $2.6\%$ while we improve force MAE by $20.8\%$.

Wang et al. (2023) uses the same pre-training method as Zaidi et al. (2023)
but applies it to ANI-1 (Smith et al., 2017) and ANI-1x (Smith et al., 2018)
datasets, which contain non-equilibrium structures. However, Wang et al.
(2023) does not encode forces, and we show in Section 4.3 that denoising non-
equilibrium structures without force encoding can sometimes lead to worse
results compared to training without denoising.

##  Appendix B Details of Experiments on OC20

###  B.1 Training Details

Since each structure in OC20 S2EF dataset has a pre-defined set of fixed and
free atoms, when training DeNS, we only add noise to free atoms and denoise
corrupted free atoms. Noise is added to all the free atoms, and therefore the
corruption ratio $r_{\text{DeNS\leavevmode\nobreak\ }}$ is $1$. Additionally,
we find that multi-scale noise works better than using a fixed $\sigma$ and
thus corrupt structures with multi-scale noise when training DeNS.

We add an additional block of equivariant graph attention to EquiformerV2 for
noise prediction. We mainly follow the hyper-parameters of training
EquiformerV2 without DeNS on OC20 S2EF-2M and S2EF-All+MD datasets. For
training EquiformerV2 on OC20 S2EF-All+MD dataset, we increase the number of
epochs from $1$ to $2$ for better performance. This results in higher training
time than other methods. However, we note that we already demonstrate training
with DeNS can achieve better results given the same amount of training time in
Table LABEL:tab:comparison_with_without_dens. Table 7 summarizes the hyper-
parameters of training EquiformerV2 with DeNS for the ablation studies on OC20
S2EF-2M dataset in Section 4.1.1 and for the main results on OC20 S2EF-All+MD
dataset in Section 4.1.2. Please refer to the work of EquiformerV2 (Liao et
al., 2023) for details of the architecture.

V100 GPUs with 32GB are used to train models. We use $16$ GPUs for training on
OC20 S2EF-2M dataset and $128$ GPUs for OC20 S2EF-All+MD dataset. The training
time and the numbers of parameters of different models on OC20 S2EF-2M dataset
can be found in Table LABEL:tab:comparison_with_without_dens. For OC20 S2EF-
All+MD dataset, the training time is $65465$ GPU-hours and the number of
parameters is $160$M.

##  Appendix C Details of Experiments on OC22

###  C.1 Training Details

Different from OC20, all the atoms in a structure in OC22 are free, and
therefore we cannot limit denoising to a pre-defined set of atoms. Instead, we
use partially corrupted structures as discussed in Section 3.2.3 and add noise
to and denoise random subsets of atoms. We follow the practice on OC20 dataset
and use multi-scale noise to corrupt structures.

We add an additional block of equivariant graph attention to EquiformerV2 for
noise prediction. Table 8 summarizes the hyper-parameters for the results on
OC22 in Table 3.

We use $32$ V100 GPUs (32GB) for training. The training time is $5082$ GPU-
hours, and the number of parameters $127$M.

| EquiformerV2 ($89$M) on | EquiformerV2 ($160$M) on
---|---|---
Hyper-parameters | S2EF-2M | S2EF-All+MD
Optimizer | AdamW | AdamW
Learning rate scheduling | Cosine learning rate with linear warmup | Cosine learning rate with linear warmup
Warmup epochs | $0.1$ | $0.01$
Maximum learning rate |  $2\times 10^{-4}$ for $12$ epochs | $4\times 10^{-4}$
|  $4\times 10^{-4}$ for $20,30$ epochs |
Batch size | $64$ | $512$
Number of epochs | $12,20,30$ | $2$
Weight decay | $1\times 10^{-3}$ | $1\times 10^{-3}$
Dropout rate | $0.1$ | $0.1$
Stochastic depth | $0.05$ | $0.1$
Energy coefficient $\lambda_{E}$ | $2$ | $4$
Force coefficient $\lambda_{F}$ | $100$ | $100$
Gradient clipping norm threshold | $100$ | $100$
Model EMA decay | $0.999$ | $0.999$
Cutoff radius (Å) | $12$ | $12$
Maximum number of neighbors | $20$ | $20$
Number of radial bases | $600$ | $600$
Dimension of hidden scalar features in radial functions $d_{edge}$ | $(0,128)$ | $(0,128)$
Maximum degree $L_{max}$ | $6$ | $6$
Maximum order $M_{max}$ | $2$ | $3$
Number of Transformer blocks | $12$ | $20$
Embedding dimension $d_{embed}$ | $(6,128)$ | $(6,128)$
$f_{ij}^{(L)}$ dimension $d_{attn\\_hidden}$ | $(6,64)$ | $(6,64)$
Number of attention heads $h$ | $8$ | $8$
$f_{ij}^{(0)}$ dimension $d_{attn\\_alpha}$ | $(0,64)$ | $(0,64)$
Value dimension $d_{attn\\_value}$ | $(6,16)$ | $(6,16)$
Hidden dimension in feed forward networks $d_{ffn}$ | $(6,128)$ | $(6,128)$
Resolution of point samples $R$ | $18$ | $18$
Upper bound on standard deviations of Gaussian noise $\sigma_{\text{high}}$ | - | $0.75$
Probability of optimizing DeNS $p_{\text{DeNS\leavevmode\nobreak\ }}$ | - | $0.25$
Loss coefficient for DeNS $\lambda_{\text{DeNS\leavevmode\nobreak\ }}$ | - | $10$
Corruption ratio $r_{\text{DeNS}}$ | $1.0$ | $1.0$

Table 7:  Hyper-parameters of training EquiformerV2 with DeNS on OC20 S2EF-2M
dataset and OC20 S2EF-All+MD dataset. For OC20 S2EF-2M dataset, the DeNS-
related hyper-parameters, $\sigma_{\text{high}}$, $p_{\text{DeNS}}$ and
$\lambda_{\text{DeNS}}$, can be found in Table 1.

Hyper-parameters | Value or description
---|---
Optimizer | AdamW
Learning rate scheduling | Cosine learning rate with linear warmup
Warmup epochs | $0.1$
Maximum learning rate | $2\times 10^{-4}$
Batch size | $128$
Number of epochs | $6$
Weight decay | $1\times 10^{-3}$
Dropout rate | $0.1$
Stochastic depth | $0.1$
Energy coefficient $\lambda_{E}$ | $4$
Force coefficient $\lambda_{F}$ | $100$
Gradient clipping norm threshold | $50$
Model EMA decay | $0.999$
Cutoff radius (Å) | $12$
Maximum number of neighbors | $20$
Number of radial bases | $600$
Dimension of hidden scalar features in radial functions $d_{edge}$ | $(0,128)$
Maximum degree $L_{max}$ | $6$
Maximum order $M_{max}$ | $2$
Number of Transformer blocks | $18$
Embedding dimension $d_{embed}$ | $(6,128)$
$f_{ij}^{(L)}$ dimension $d_{attn\\_hidden}$ | $(6,64)$
Number of attention heads $h$ | $8$
$f_{ij}^{(0)}$ dimension $d_{attn\\_alpha}$ | $(0,64)$
Value dimension $d_{attn\\_value}$ | $(6,16)$
Hidden dimension in feed forward networks $d_{ffn}$ | $(6,128)$
Resolution of point samples $R$ | $18$
Upper bound on standard deviations of Gaussian noise $\sigma_{\text{high}}$ | $0.5$
Probability of optimizing DeNS $p_{\text{DeNS\leavevmode\nobreak\ }}$ | $0.5$
Loss coefficient for DeNS $\lambda_{\text{DeNS\leavevmode\nobreak\ }}$ | $100$
Corruption ratio $r_{\text{DeNS}}$ | $0.5$

Table 8: Hyper-parameters for OC22 dataset.

##  Appendix D Details of Experiments on MD17

###  D.1 Additional Details of DeNS

It is necessary that gradients consider both the original task and DeNS when
updating learnable parameters, and this affects how we sample structures for
DeNS when only a single GPU is used for training models on the MD17 dataset.
We zero out forces corresponding to structures used for the original task so
that a single forward-backward propagation can consider both DeNS and the
original task. In contrast, if we switch between DeNS and the original task
for different iterations, gradients only consider either DeNS or the original
task, and we find that this does not result in better performance on the MD17
dataset than training without DeNS.

###  D.2 Training Details

We use the same codebase as Equiformer (Liao & Smidt, 2023) for experiments on
the MD17 dataset and follow most of the original hyper-parameters for training
with DeNS. For training DeNS, we use an additional block of equivariant graph
attention for noise prediction, which slightly increases training time and the
number of parameters. We use single-scale noise with a fixed standard
deviation $\sigma$ when corrupting structures. The hyper-parameters introduced
by training DeNS and the values of energy coefficient $\lambda_{E}$ and force
coefficient $\lambda_{F}$ on different molecules can be found in Table 9.
Empirically, we find that linearly decaying DeNS coefficient
$\lambda_{\text{DeNS}}$ to $0$ thoughout the training can result in better
performance. For the Equiformer variant without attention and layer
normalization, we find that using normal distributions to initialize weights
can result in training divergence and therefore we use uniform distributions.
For some molecules, we find training Equiformer variant without attention and
layer normalization with DeNS is unstable and therefore reduce the learning
rate to $3\times 10^{-4}$.

We use one A5000 GPU with 24GB to train different models for each molecule.
The training time and the numbers of parameters can be found in Table 4.

Hyper-parameter |  | Aspirin | Benzene | Ethanol | Malonaldehyde | Naphthalene | Salicylic acid | Toluene | Uracil
---|---|---|---|---|---|---|---|---|---
Energy coefficient $\lambda_{E}$ |  | 1 | 1 | 1 | 1 | 2 | 1 | 1 | 1
Force coefficient $\lambda_{F}$ |  | 80 | 80 | 80 | 100 | 20 | 80 | 80 | 20
Probability of optimizing DeNS $p_{\text{DeNS}}$ |  | 0.25 | 0.25 | 0.25 | 0.25 | 0.25 | 0.125 | 0.125 | 0.25
Standard deviation of Gaussian noises $\sigma$ |  | 0.05 | 0.05 | 0.05 | 0.05 | 0.05 | 0.025 | 0.025 | 0.05
DeNS coefficient $\lambda_{\text{DeNS}}$ |  | 5 | 5 | 5 | 5 | 5 | 5 | 5 | 5
Corruption ratio $r_{\text{DeNS}}$ |  | 0.25 | 0.25 | 0.25 | 0.25 | 0.25 | 0.25 | 0.25 | 0.25

Table 9:  Hyper-parameters of training with DeNS on the MD17 dataset. The
other hyper-parameters not listed here are the same as the original Equiformer
($L_{max}=2$) trained without DeNS.

##  Appendix E Pseudocode for training with DeNS

We provide the pseudocode for training with DeNS in Algorithm 1 and note that
Line 5 can be parallelized. For denoising partially corrupted structures
discussed in Section 3.2.3, we only add noise to a subset of atoms (Line 16 –
19) and predict the corresponding noise (Line 30).

Algorithm 1 Training with DeNS

1:Input:

2:$p_{\text{DeNS}}$: probability of optimizing DeNS

3:$\lambda_{\text{DeNS}}$: DeNS coefficient

4:$\sigma$: standard deviation of Gaussian noise if multi-scale noise is not
used

5:$\sigma_{\text{high}}$: upper bound on standard deviations of Gaussian noise
if multi-scale noise is used

6:$\sigma_{\text{low}}$: lower bound on standard deviations of Gaussian noise
if multi-scale noise is used

7:$r_{\text{DeNS}}$: corruption ratio

8:$\lambda_{E}$: energy coefficient

9:$\lambda_{F}$: force coefficient

10:GNN: graph neural network for predicting energy, forces and noise

11:while training do

12: $\mathcal{L}_{\text{total}}=0$

13: Sample a batch of $B$ structures $\\{\left(S_{\text{non-
eq}}\right)^{j}\mid j\in\\{1,...,B\\}\\}$ from the training set

14: for $j=1$ to $B$ do $\triangleright$ This for loop can be parallelized

15: Let $\left(S_{\text{non-eq}}\right)^{j}=\left\\{(z_{i},\mathbf{p}_{i})\mid
i\in\\{1,...,|\left(S_{\text{non-eq}}\right)^{j}|\\}\right\\}$

16: Sample $p$ from a uniform distribution ${\mathbf{U}}(0,1)$ to determine
whether to optimize DeNS

17: if $p<p_{\text{DeNS}}$ then $\triangleright$ Optimize DeNS based on
Equation 6

18: if multi-scale noise is used then

19: Sample $\sigma_{\text{sample}}$ from $\\{\sigma_{k}\\}_{k=1}^{T}$

20: else

21: $\sigma_{\text{sample}}=\sigma$

22: end if

23: for $i=1$ to $|\left(S_{\text{non-eq}}\right)^{j}|$ do

24: $q_{i}\sim{\mathbf{U}}(0,1)$

25: if $q_{i}<r_{\text{DeNS}}$ then $\triangleright$ Add noise to and denoise
the atom

26: $\bm{\epsilon}_{i}\sim\mathcal{N}(0,\sigma_{\text{sample}}I_{3})$

27: $\tilde{\mathbf{p}_{i}}=\mathbf{p}_{i}+\bm{\epsilon}_{i}$

28: $m_{i}=1$ $\triangleright$ Denoise the atom when calculating
$\mathcal{L}_{\text{DeNS}}$

29: else

30: $\tilde{\mathbf{p}_{i}}=\mathbf{p}_{i}$

31: $m_{i}=0$

32: end if

33: $\tilde{\mathbf{f}}_{i}^{\prime}=\mathbf{f}_{i}^{\prime}\cdot m_{i}$
$\triangleright$ Encode the atomic force if the atom is corrupted

34: end for

35: Let $(\tilde{S}_{\text{non-
eq}})^{j}=\left\\{(z_{i},\tilde{\mathbf{p}_{i}})\mid
i\in\\{1,...,|\left(S_{\text{non-eq}}\right)^{j}|\\}\right\\}$

36: Let $\tilde{F}\left((S_{\text{non-
eq}})^{j}\right)=\left\\{\tilde{\mathbf{f}}_{i}^{\prime}\mid
i\in\\{1,...,|\left(S_{\text{non-eq}}\right)^{j}|\\}\right\\}$

37:
$\hat{E},\hat{F},\hat{\bm{\epsilon}}\leftarrow\text{GNN}\left((\tilde{S}_{\text{non-
eq}})^{j},\tilde{F}\left((S_{\text{non-eq}})^{j}\right)\right)$

38: $\mathcal{L}_{E}=\left|E^{\prime}\left((S_{\text{non-
eq}})^{j}\right)-\hat{E}\right|$

39: $\mathcal{L}_{\text{DeNS}}=\frac{1}{|(S_{\text{non-
eq}})^{j}|}\sum_{i=1}^{|(S_{\text{non-
eq}})^{j}|}m_{i}\cdot\left|\frac{\bm{\epsilon}_{i}}{\sigma_{\text{sample}}}-\hat{\bm{\epsilon}}_{i}\right|^{2}$
$\triangleright$ Calculate $\mathcal{L}_{\text{DeNS}}$ based on Equation 4

40: $\mathcal{L}_{F}=\frac{1}{|(S_{\text{non-
eq}})^{j}|}\sum_{i=1}^{|(S_{\text{non-
eq}})^{j}|}(1-m_{i})\cdot|\mathbf{f}_{i}^{\prime}\left((S_{\text{non-
eq}})^{j}\right)-\hat{\mathbf{f}}_{i}|^{2}$ $\triangleright$ Predict forces
for uncorrupted atoms

41:
$\mathcal{L}_{\text{total}}=\mathcal{L}_{\text{total}}+\lambda_{E}\cdot\mathcal{L}_{E}+\lambda_{\text{DeNS}}\cdot\mathcal{L}_{\text{DeNS}}+\lambda_{F}\cdot\mathcal{L}_{F}$

42: else$\triangleright$ Optimize the original task based on Equation 1

43: $\hat{E},\hat{F},\\_\leftarrow\text{GNN}\left((S_{\text{non-
eq}})^{j}\right)$

44: $\mathcal{L}_{E}=\left|E^{\prime}\left((S_{\text{non-
eq}})^{j}\right)-\hat{E}\right|$

45: $\mathcal{L}_{F}=\frac{1}{|(S_{\text{non-
eq}})^{j}|}\sum_{i=1}^{|(S_{\text{non-
eq}})^{j}|}|\mathbf{f}_{i}^{\prime}\left((S_{\text{non-
eq}})^{j}\right)-\hat{\mathbf{f}}_{i}|^{2}$

46:
$\mathcal{L}_{\text{total}}=\mathcal{L}_{\text{total}}+\lambda_{E}\cdot\mathcal{L}_{E}+\lambda_{F}\cdot\mathcal{L}_{F}$

47: end if

48: end for

49: $\mathcal{L}_{\text{total}}=\frac{\mathcal{L}_{\text{total}}}{B}$

50: Optimize GNN based on $\mathcal{L}_{\text{total}}$

51:end while

##  Appendix F Visualization of corrupted structures

We visualize how adding noise of different scales affect structures in OC20,
OC22 and MD17 datasets in Figure 3, Figure 4 and Figure 5, respectively.

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.  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.  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.