Generalizing Denoising to Non-Equilibrium Structures Improves Equivariant Force Fields
Authors: Yi-Lun Liao
Tess Smidt
Muhammed Shuaibi
Abhishek Das
Venue: Accepted by TMLR ()

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.

Published in Transactions on Machine Learning Research (12/2024)
Generalizing Denoising to Non-Equilibrium Structures Improves Equivariant Force Fields
Yi-Lun Liao
yllioo@mitt.cmm edu
Massachusetts Institute of Technology Work partially done during un internship FAIR, Mela
Tess Smidt
tsmrddt edu
Massachusetts Institute Technology
Muhammed Shuaibi'
mshuaibi@@meta.com
FAIR, Meta
Abhishek Das" Work done at FAIR, Mela
Jbhshkd=@gmail..o
Equal contribution
Reviewed on OpenReview:
Code:
Abstract
Understanding the interactions of atoms such as forces 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, propose use denoising mon-equilibrium structures (DeNS) 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 non-equilibrium structures. The main difference is that non-equilbrium structure does not correspond to local energy minima and has non-zero forces, and therefore can have many possible atomic positions compared to an equilibrium structure. This makes denoising non-equiibrrimm structures 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-cquilibrium structure satisfying the input forces instead 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 OC20, OC22 and MD17 datasets and demonstrate that DeNS can achieve new slate-of-the-ard results on OC20 and OC22 and significantly improve training efficiency on MD17.
auxiliary task to better leverage training data and
ill-posed problem since the
1 Introduction
Graph neural networks (GNNs) have made remarkable progress approximating high-fidelity, compute- intensive quantum mechanical caleulations like density functional theory (DFT) for atomistic systems (Gilmer

---

Published in Transactions on Machine Learning Research (12/2024)
(b) Denoising non-oquilibrium structures is ill- posed since the target of denoising has many possibilities as denoted by "?" below.
(d) In DeNS, we propose to additionalyy use the forces of the original non-equilibriu structure to specily which larget structure we are denoising.
Add noise
GNN denoising
Add noise
Forze encoding
Smase
Sum
Soa
Add noise
noise
Sea
Sia
F(S.)) =0
(a) Local relaxation trajectory on potential energy surface.
(c) Previous works on denoising are limited to
(e) Denoising equilibrium structures hecomes a special case where forces are close to zero.
equilibrium or relaved structures.
Figure| 1: Illustration of denoising equilibrium and non-equilibrium structures. As shown in (a), we relax a non- equilibrium structure Snon-eq and obtain the final relaxed, equilibrium structure Seg. The path between Snon-eq and Seq forms relaxation trajectory. All structures S along the trajectory except Seg are non-equilibrim and have non-zero forces F(S). For denoising structures S, we add noise to their 3D atomic coordinates obtain corrupted structures S and predict the corresponding noise given S. We compare denoising non-equilibrium and equilibrium structures in (b) and (c), respectively, and show that denoising non-oquilibrium structures can be ill- 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).
et al., 2017; Zhang et al., 2018; Unke et al., 2021; Batzner al., 2022; Rackers et al., 2023; Lan et al., 2022), enabling new insights in applications such as molecular dynamics simulations (Musaclian 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* * al., 2021), contains about 138M examples while GPT-3 (Brown al., 2020) is trained hundreds of billions of words and VIT-22B (Dehghani et al., 2023) is trained on around 4B images.
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. Indeod, this is the case for equilibrium structures (c.g., Sen at local energy minimum in Figure 1(c)) 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 relaxation trajectory to get to a local minimum are all non-equilibrium as shown in Figure 1. Hence, is important to generalize denoising to leverage the much larger set of non-equilibrium structures.
Since non-equilibrium structure has non-zero atomic forces and atoms not confined to local energy minima, it can have more possible atomic positions than structure at equilibrium. As shown in Figure 1(b), this can make denoising non-oquilibrium structure an ill-posed problem since there are many possible target structures. To address the issue, we propose force encoding and lake the forces of the original non-equilibrium structures as inputs when denoising non-equilbrrium structures, Intuitively, the forces constrain the atomic positions of non-oquilbrrium 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 (d). Previous works on denoising equilibrium structures (Godwin et al., 2022; Jiao et al., 2022;
2

---

Published in Transactions on Machine Learning Research (12/2024)
Zaidi et al., 2023; Liu et al., 2023; Feng et al., 2023b;a) end up being special case where the forces of original structures are close to xero as in Figure (e)..
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 prediction, a model is given the forces of the original uncorrupted structure as inputs make the transformation from corrupted non-oquilibrimm structure to uncorrupted non-equilirium structure tractable. When used along with the original tasks like predicting the energy and forces of non-equilibrium structures, DeNS improves the performance on 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 supe-vuue learning methods in other domains.
Because DcNS requires encoding forees, it favors cquivariant networks, which build up equivariant features at each node with vector spaces of irredueible represenlations (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. Morcover, with the reduced complexity of equivariant operations (Passaro & Zitmick, 2023) and incorporation of the Transformer network design (Tiao & Smidt, 2023; Liao et al., 2023) from NL.P (Vaswani et al., 2017) and CV (Dosovitakiy et al., 2021), equivariant networks have become the state-of-the-ar 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 (Chmicla al., 2017; Schütt al., 2017; Chiniela et al., 2018) datasets. On OC20 S2EF-2M dalaset, EquiformerV2 (Liao et al., 2023) trained with DeNS can achieve better energy and force results and save 2.3x training time compared to training without DeNS (Section 4.1.1). On OC20 S2ED-AIH+ dataset, EquiformerV? trained with DeNS achieves new state-of-the-art results on Structure to Energy and Forces (S2EF) and Initial Structure to Relaxed Energy (IS2RE) tasks (Section 4.1.2). Similarly, EquiformerV2 trained with DeNS sets new state-of-the-art results on OC22 dataset, improving energy by up to 15%, forces by up to 12%, and IS2RE by up 15% (Section 4.2). On MD17 dataset, Equiformer (Lmex = 2) (Liao & Smidt, 2023) trained with DeNS achieves better results and saves 3.1x x training time compared to Equiformer (Lmar = 3) without DeNS (Section 4.3), where Linax 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 large dataset of equilibrium structures like PCQM4Mv2 (Nakata & Shimazaki, 2017) and then fine-tune with supervised learning on smaller downstream datasets. Besides, the work of Noisy Nodes (Godwin et al., 2022) uses 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-equilbrium ones. In contrast, the proposed DeNS generalizes denoising to non-oquilibrimm structures with force encoding so that we can improve the performance on the larger set of non-equilirrumm structures. We provide a detailed comparison to previous works on denoising in Section A.1.
3

---

Published in Transactions on Machine Learning Research (12/2024)
Networks. Please refer to Section A.2 for discussion on equivariant networks. In addition, since we mainly focus on applying the proposed DeNS to Equiformer series (Liao & Smidt, 2023; Liao et al., 2023), we provide a brief introduction to them in Section A.3.
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 a dataset and can be described as {(2i. P:) e (1. [S]]}, where 2; E N denotes the atomic number of the i-th atom and P; R denotes the 3D atomic position. The energy of S is denoted E(S) R, and the atom-wise forces where fü is the force acting on the i-th atom. Tn 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-equilirrim structure. Since non-equilibrium structures have non-zero alomic 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.
denoted as F(S) { e R i [1,...,||$|]]
In this work, we focus on the task of predicting energy and forces given non-cquliirrum structures, Specifically, given non-equilibrium structure Smon-eq: GNNs predict energy ((Snon-eq) and atom-wise forces ((Snnn-o)) (S R ie I } and minimize the following loss function:
Snon-ca
Ag- LE + dF LF = AE (S'nnnnn)))) E(Son-oo)) I + Xp-
(SSon-e))) - P (Smon-ca)) (1)
S Spon-cq
where 0g and AR energy and force coefficients controlling the relative importance between energy and force predictions. ("(Snon-eq)== F(Snowee)=-s is the normalized ground truth energy obtained by first subtracting the original energy F((uoo--e)) by the mean of energy labels in the training set Lg and then dividing by the standard deviation of energy labels dg. Similarly, f ay prediction, we can either use direct methods, which is to directly predict forces from latent representations like node embeddings, as commonly used for OC20 and OC22 datasets or gradient methods, which is to take the negative gradients of predicted energy with respect to atomic positions, for datascts like MD17.
dg
f. is the normalized atom-wise force. For force
3.2 Denoising Non- Eon-iquiii Structures (DeNS)
3.2.1 Formulation of Denoising
Denoising structures has been used improve the performance of GNNs on 3D atomistic datasets (Godwin et al., 2022; Zaidi et al., 2023; Feng et al., 2023b; Wang al., 2023). We first corrupt data by adding noise and then train denoising autoencoder to reconstruct the original data by predicting the noise. Specifically, given 3D atomistic system S = [(2,,)) P:) lie (1...., (1)),,,)()) we create corrupted structure Š by adding Gaussian noise with a tuneable standard deviation a to the atomic positions p. of the original structure S:
Š = (E., iE (1|)..,|||],. where D: = Pi + E; and E; ~ N((,0la)
(2)
We denote the set of noise added to S as Noise(S, $) (e& R autoencoder, GNNs take 3 as inputs, output stom-wise noise prediction é(S), and minimize the L2 difference between normalized noise e and noise prediction e($);:
When training denoising
S
E
((S.S)
Σ | e,
1
E(S))
(3)
o
where p(S. 5) denotes the probability of obtaining the corrupted structure Š from the original structure S. We divide the noise e, by the standard deviation to normalize the outputs of noise prediction.
4

---

Published in Transactions on Machine Learning Research (12/2024)
When the original structure S is 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-oquilibrim structure, denoising, the transformation from a corrupted non-cquiibrium structure to the original non-equilibriun one, can be an ill-posed problem since there are many possible targel structures as shown in Figure (bb).
3.2.2 Force Encoding
The reason that denoising non-equilibriun structures can be ill-posed is because we do not provide sulficient information to specify the properties of the target structures. Concretely, given an original non-equilibrium structure Snon-eq and its corrupted counterpart Snon-eq: some structures interpolated between Snon-eq and 3 Son-ee could be in the same data distribution and therefore be the potential target structures of denoising. In contrast, when denoising equilibrium structures shown in Figure 1(c), 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-equilbrrium structures when denoising non-equilibrium ones as illustrated in Figure (d). Specifically, when training denoising non-equilibrimm structures (DeNS), GNNs take both a corrupted non-equilibrium structure Snow-eq and the forces F(Snou-eg) of the original non-equilibrimm structure Spon-eq as inputs, output atom-wise noise prediction é (Snon-eq, (S)on-ca)); and minimize the L2 difference between normalized noise e and noise prediction è (Snon-oq: F(Snon-eq));
[Snoneeg
1
--(@@coneep F (no) 2
DDNS
E pp 9 mon
Spon-eq
Equation 4 is more general and reduces to Equation 3 when the target of denoising becomes equilib- rium structures with near-zero forces. Since we train GNNs with Snon-eg and F(Suon-eq) as inputs and Noise(Snon-es) moon-eq) ) as outputs, they implicitly learn to leverage F(Spon-eq) to reconstruct Snou-eq instead of predicting any arbitrary non-equilbrrim structures. Empirically, we find that force encoding critical to the effectiveness of DeNS on OC20 S2EF-2M dataset (Index 2 and Index 3 in Table 1(b)) and MD17 dataset (Index and Index 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 Ct channels of type-T. vectors with degree being in the range from 0 to maximum degree Imaa- CL and Lmaz are architoctural hyper-parametors of equivariant networks. To obtain the force embedding de from the input force f, we first project f into type-L vectors with spherical harmonics (L) ( where 0 & L & Lmaz, and then expand the number of channels from to CL with an SO(3) lincar layer (Geiger et al., 2022; Geiger & Smidt, 2022; Liao & Smidt, 2023):
(D) 3 = S03 Linear
(1E9)) (
(5)
(L) denotes the channels of type-L vectors in force embedding Ir, and SO3 Lincar denotes the SO(3) linear operation on type-L vectors. Since we normalize the input force when using spherical harmonics, we multiply (()) [ with the norm of input force ||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. The pseudocode for encoding forces into node embeddings can be found in Section E.
2f
On the other hand, it might not be that intuitive to encode forces in invariant networks since their internal latent representations such node embeddings and edge embeddings are scalars instead of type-L vectors
5

---

Published in Transactions on Machine Learning Research (12/2024)
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. As for other work (Duval et al., 2023) leveraging frame averaging to achieve equivariance through data transformations, we can follow how unit cell Cartesian coordinates are projected in their framework to encode forces when optimizing for DeNS. Specifically, using only 3D atomic positions and project the input forces to the frame axes. The projected forces remain the same under any E3) transformation and enable using unconstrained functions to encode the input forces.
compute one set of frame axes
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 (Figure 2(b) or Figure 2(c)) or the objective of the original task (Figure 2(a)). This introduces an additional hyper-parameter PDeNs, the probability of optimizing DeNS. We use an additional noise head for noise prodiction, which slightly increases training time. Additionally, when training DeNS, similar to Noisy Nodes (Godwin et al., 2022), also leverage energy labels and predict the energy of original structures. Therefore, given original non-oquilibrium structure Snon-eq and the corrupted counterpart Snon-eqp training DeNS corresponds to minimizing the following loss function:
Ag LE + ADeNS - LDNNS = Ag E' (Snon-eqq) - E (Snon-eq; P(S F(Snou-ea)) ADeNS LDeNs
(6)
where LDNNS 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 Smon-oq: and we share the energy prediction head across Equation 1 and Equation 6. The loss function introduces another syper-parameter ADeNS, DeNS coefficient, controlling the relative importance of the auxiliary task. Besides, the process of corrupting structures also results in another hyper-parameter a as shown in Equation 2. We provide the pseudocode in Section F. We note that the force encoding in DeNS only relies force labels the training set and we do not use any force labels the validation testing sets.
Partially Corrupted Structures. Empirically, find that adding noise to all atoms in structure can sometimes lead to limited performance gain of DeNS. We surmise adding noise to all atoms makes denoising too difficult and potentially not well-defined and that there are still several structures satisfying input forces. To address this issue, we use partially corrupted structures, where we only add noise to and denoise a random subset of atoms as shown in Figure 2(c). 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 rDNss, which is the ratio of the number of corrupted atoms to that of all atoms.
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 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 the inverse of predicting forces. As demonstrated in the works of Noisy Nodes (Godwin al., 2022) and UL2 (Tay et al., 2023), training single model with multiple correlated objectives to learn different interactions can help the performance on original tasks.

---

Published in Transactions on Machine Learning Research (12/2024)
Focces of the orizinal zon- qquilbriaa structure (YSlamm))
lead
Equivariant GNS
Bionny
Cr. Energy of the onignal notrequilbeum stmoner Fi m)
Orginal non-equillbii bant sinadure Sweee
Original task predicting energy and forces.
Forres the original mot- equilbrriu draciome Fl P(
Equivariant GNN
Energy original
bon-eipillbbtemm All
Force enoeding
Original inscpiiliiian stractrre Staee
Compptel
Moise head
Sumpled
(b) DeNS without partially corrupted structures, All the atoms structure are corrupted with Gaussian noise. We encode all the atom-wise forces predict noise.
Forces of the original non equilite fune structure (Saen))))
Fonce heed
P
Cy
Add
Equivariax GNN
Erorgy haad
Force anceding
E
sen-zqquilhiin
Oniginal equutitutruiu
compped Seam
Noiae hoal
Sumpled
(c) DeNS with partially corrupted structures. We add noise to and denoise random subset atoms. corrupted atoms, encode their forces and predict the noise. For other uncorrupted atoms, instead predict their forces. Here only the top two white atoms are corrupted.
Figure Training process when incorporating DeNS as an auxiliary task. The pseudocode training with DeNS can be found in Section F. "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(Suon-eq) and energy label E(Snon-eg) used in (b) and (c) are the same as those (a) and that training with DeNS does nol require any additional dala.
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 el 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 improve the performance on downstream tasks. However, unlike those self-supervised learning methods (Devlin al., 2019; He al., 2022; Zaidi et al., 2023), DeNS requires force labels for encoding, and therefore, propose use DeNS

---

Published in Transactions on Machine Learning Research (12/2024)
Equikemee V2
Equikemne V2 DoNS
Mothod
Epuucho
Sorus
Number of training tine (GPU-bourx)
Serrw
Number of training time paramettns (GPU-hrars)
Index
DeNS
cucoding
Force
Ap
Partially comppted
cuaoge
8931
255
969
18002
19.00
269)
65CN
65CN DeNS
21.32
18.67
278
265
80
20001
240
0258
1882
19.14
208
08M
1829
10.34
371
(a) Comparison of training with and without DeNS.
(b) Design Choices.
Table 1: Ablation results of training with DeNS on OC20 S2EF-2M dataset. We report mean absolute errors for forces in meV/ / and energy meV, and lower better. Errors are averaged over the four validation sub-splits of OC20. (a) We train EquiformerV2 and eSCN and compare the results of training with and without DeNS. The training time is measured on V100 GPUs. (b) We train EquiformerV2 for epochs to verify the design choices of DeNS.
auxiliary task optimized along with original tasks and do not follow the previous practice of first pre- rre-training and then fine-tuning. Additionally, we note that before obtaining a single equílibrium structure, we need to run relaxations and generale many intermediate non-equilibrium ones (Figure (a)),, which is the labeling process as well. We hope that the ability leverage more from non-equilibrium structures as proposed in this work can encourage researchers release data containing intermediate non-equilibrium structures in addition to final equilibrium ones. Moreover, 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 about 1.2M Density Functional Theory (DFT) relaxation trajectorics. Each DFT trajectory in OC20 starts from an initial structure of an adsorbate molccule placed on a catalyst surface, which is then relaxed with the revised Perdew-Burke-Ernnae et al., 1999) calculations to local energy minimum. Relevant to DeNS, all the intermediate structures from relaxation trajectory, except the relaxed structure, are considered non-oquilibrium structures. The relaxed or equilibrium structure has forces close to zero.
(RPBE) functional (Hammer
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-oquilibrium structure from any point the trajectory. These predictions are evaluated on energy and force mean absolute error (MAE). Most of the previous works direct methods to predict forces on OC20, and we follow this practice and investigate how DeNS improve direct methods for force prediction. Once model is trained for S2EF, it is used to run structural relaxations from an initial structure using the predicted forces till 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 Section B.1 for details on DeNS, architectures, hyper- parameear and training time.
8

---

Published in Transactions on Machine Learning Research (12/2024)
4.1.1 Ablation Studies
We use OC20 S2EF-2M dataset lo compare the results of training with and without DeNS and verify the design choices of DeNS.
Comparison between Training with and without DeNS. Table 1(a) summarizes the results of training with and without DeNS. For EquiformerV2, incorporating DeNS as an auxiliary task boosts the performance on energy and force predictions and only increases training time by 7.4% and the number of parameters from 83M to 89M. Particularly, EquiformerV' trained with DeNS for 12 epochs oulperforms EquilormerV? trained without DeNS for 30 epochs and saves 2.3x training time. This suggests that using data augmentation and learning an auxiliary task are more efficient to improve performance compared to simply training for longer. Additionally, we show that DeNS can be applicable to other equivariant networks like eSCN (Passaro & Zitnick, 2023). Training eSCN with DeNS for 20 epochs results in better energy and force MAE compared to EquiformerV2 V2 trained without DeNS for 30 epochs while requiring 1.9x less training time. DeNS slightly increases training time of eSCN by 1.5%, and the different amounts of increase in training time between EquiformerV V2 and eSCN are because they use different modules for noise prediction.
Design Choices. We train EquiformerV2 for 12 epochs with DeNS to verify the design choices of DeNS and summarize the results in Table (bb). Comparing Index 2 and Index 3, we show that encoding forces F(Snon-eq) in Equations and enables denoising non-equilibrium structures to significantly improve both energy and force MAE. Moreover, DeNS without force encoding (Index 3) results in clearly worse force MAE compared to training without DeNS (Index 1). These results verify our claim that denoising non-equilibrium structures naively can be ill-posed and potentially harmful to performance and that force encoding is critical to leverage the gain of denoising. We nole that force encoding only relies force labels in the training set and does not require any additional data. We also compare DeNS with and without force encoding on MD17 in Section 4.3 and have similar observations. Comparing Index 2 and Index 4, we demonstrate that predicting energy of original structures given corrupted ones can improve energy MAE by 5.6% while slightly increasing force MAE by 1.2%. The increase in force MAE is because predicting energy given corrupted structures is equivalent to using Ag # 0 in Equation 6 and implicitly decreases the relative importance of force prediction. Since the decrease in energy MAE is greater than the increase in force MAE, we adopt the practice of predicting energy given corrupted structures. Finally, the comparison between Index 2 and Index 5 shows that partially corrupted structures can further improve the performance gain of DeNS.
4.1.2 Main Results
We train EquiformerV2 (160M) with DeNS on OC20 S2EF-AUI+MD dataset. The model follows the same configuration as EquiformerV2 (153M) trained without DeNS, and the additional paramcters are due to force encoding and one additional block of equivariant graph allention for noise prediction. We report results in Table 2. All test results are computed via the EvalAI evaluation server EquiformerV? trained with DeNS achieves new stato-of-tho-ar results on both S2EF and IS2RE tasks. On the S2EF validation split, EquiformerV2 trained with DeNS improves energy MAE by 5meV and force MAE by 1. 1.0meV/A, /A, which is comparable to the gain brought by increasing the size of EquiformerV2 from 31M to 153M (energy MAE improved by 5meV and force MAE improved by 1.3meV/A) On the S2EF test split, the improvement in energy and force predictions is smaller, which is probably because of different splits. On the IS2RE test split, training Equiformer with DeNS reduces MAF by 16meV, achieving similar performance gain of going from cSCN (Passaro & Zitnick, 2023) to EquiformerV? (Liao et al., 2023) (IS2RE MAE improved by 14mcV) and demonstrating the effectiveness of training with DeNS. We note that the improvement might not be as significant that on OC20 S2EF-2M (Section 4.1.1), OC22 (Section 4.2) and MD17 (Section 4.3) datasets training set contains much more structures along relaxation trajectories,
since the OC20 S2EF- making new 3D geometries generated by DeNS less helpful. However, DeNS is still valuable because most datasets are not as large as OC20 S2EF-A1I+MD dataset (about 172M structures in the training set.) but have sizes closer to OC20 S2EF-2M (2M structures), OC22 (8.2M structures), and MD17 (950 structures) datascts.
eval. eevll. hwwp://aa/.ee.ccllaage///////aae
9

---

Published in Transactions on Machine Learning Research (12/2024)
Throughput
S2EF vulidiaioon
S2EF test
IS2RE test
Samples GPU sec.
Energy MAE (aueV).
Foree MAE (uee/A1 $
Energy MAE
Foree MAE
Energy MAE
Model
(meVV)
(mee)) J
Gem-Net-OC---E (Gasteiger et al., 2022)
7.3
239 232
22.1
230
21,0
GemNet-OO-LLP [Gasteigee el ol., 2022))
3.2
20.0
241
19.0
GemNNet-CC-LFF (Gasteiger et al., 2022)
318
SCN L=6 K=16 (Keap 2.bml)) (Zitnick al., 2022)
228
17.8
328
SCN L=8 K=20 (Zitaick et al, 20122)
237
17.2
321
<SCN L-8 K-20 (Passaro Zitnick, 2025)
2.9
243
17.1
228
15.6
521
Equilormer (Au = 4, 31M) (Lino et ad., 2023)
7.1
232
10.3
228
15.5
315
EquiformerV (Ag 159M) (Liao
2023)
1.8
227
15.0
219
14.2
309
EquilormerV2 + DeNS (Ap 4, 16084)
1.8
222
14.0
214
13.3
293
Table 2: OC20 results S2EF validation and test splits and IS2RE Lest split when trained on OC20 S2EF-All+
dataset. Throughput is reported as the number of structures processed per GPU-second during training and measured on V100 GPUs.
SIFP-Tura viliiation
SIRF-TsEal
183RE-Thaal teel
Number
Faregy MAR (mwh))
MAR (mn)/A]
Energy MAF:
|Prro MAR (mwV/A)
Fowegy MAR (mwV))
Model
porameters
000
GeunNee-OC OC (Gusteige: et 3322)
1384
GomNe-OC (nssati ia DCas ac2o (Gaateige et al., 2022)
SIM
S11
680
31.2
1230
eSCN (Physoro Zitnok, 1053)
200
350
781
15.5
35.7
Equifomme22 (Ag = 4,Ay - 1000 (Liso ol., 2029)
122M
483
620
22.88
50.70
253.7
609.8
21.58
32.65
1119
1440
Bquiforne'2 DeN8 (Ag - 4, = 100)
127M
301.6
533.0
20.08
27.11
236.4
500.7
20.04
29.31
951
1282
Table 3: OC22 results S2EF-Total validation and test splits and IS2RE-Total test split. Models are trained the OC22 S2EF-Total training split unless otherwise noted.
4.2 OC22 Dataset
Datasct and Tasks. The Open Catalyst 2022 (OC22) datasct (Tran* et al., 2022) focuses on oxide electrocatalysis and consists of about 62k DFT relaxations obtained with Perdew-Burke- ofrdee-Burke-Emzerhoo (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-Tote on the OC22 - dataset, which has 8.2M structures, and evaluate them on energy and force MAE on the S2EF-total validation and test splits. Same as OC20, we use direct methods to predict forces here. Then, we use these models to perform relaxations starting from initial structures in the IS2RE-Totar test split and evaluate the predicted relaxed energies energy MAF.
We train models
Training Details. Please refer to Section C.1 for details architectures, hyper parameparseeeers and training time.
Results. We use Equiformer V2 with energy coefficient AÆ = 4 and force coefficient AF = 100 to demonstrate how DeNS can further improve thtt--tf- results and summarize the comparison in Table 3. Compared to EquiformerV2 (>g = 4, A 100) trained without DeNS, EquiformerV? trained with DeNS consistently achieves better energy and force MAE across all the S2EF-Total validation and test splits, with 9.6% 15.3% improvement in energy MAE and 7.1% to 11.7% improvement in force MAE. For IS2RE-Total, EquiformerV2 trained with DeNS achieves the best energy MAE results. The improvement in IS2RE-Tota from training with DcNS on only OC22 is greater than that of training on the much larger OC20 and OC22 datascts reported in previous works. Specifically, training GemNet-OG on OC20 and OC22 datasets (about 134M + 8.2M structures) improves 1S2RE-Total energy MAE ID by 129meV and OOD by 50meV compared to training GemNet-OC on only OC22 dataset (8.2M structures). Compared to training without DeNS, training EquiformerVa V with DeNS improves ID by 168meV and OOD by 158meV. Thus, training with DeNS clearly improves data efficiency and performance on OC22.
10

---

Published in Transactions on Machine Learning Research (12/2024)
Blaund
Sakeslic
Tolacs
Numbe
Modl
bceas
carnyy
frees
freos
eaergy
frees
[aoo
laoss
ursg
Knos
(CPUaom)) pummetes
BeiNaa (Scaaii c. di,, 2017)
16.0
65-8
3.5
13.4
355
19.9
6..6
25.8
6.9
20.2
8.7
91.
6.2
24.7
6.1
26.8
Denalert (Couompee at 2000
21.8
34
6.1
2.8
10:1
455
16.8
5.8
16.2
6.5
9.4
5.0
13.1
PaiNS (Siain 2023)
27
877
5.3
6.1
41
Tom (Tbxike
SepllP Uhe.=
2.2
Equioumax Houu =
5.
7.8
3.8
6.6
9.2
8.1
888
588
8.7
9.3
4.5
4.1
8.8
2.1
4.8
38
17
3/20
Eeuismnre Thoss -50
93
166
23
6.1
233
2.9
3.2
5.4
4.4
20
43
3.9
87
2.8
4.3
34
59
3.502
28
6.1
92
2.6
4.4
8.7
8.6
La
Mosse 3)
Table 4: Mean absolute error results on the MD17 testing set. Energy and force are in units of meV and meV/A. 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
Aspirm
Benaene
Ethano)
Makensldetyde
Nephthalene
Saleyyli acid
Tolume
Uraet]
lodess
Atiatiinn
Layer tormaliaaiio
DoN8
suaeyy
Reus
eaangy
lexas
stangy
lueas
euangy
funow
eusngy
loecea
surrgy
Gueces
souegy
Soecen
souege
Exrt
5.1
722
2.2
6.6
2.2
3.1
3.3
5.8
3.7
2.1
4.5
4.1
3.8
211
J.3
/
6.1
6.7
233
6.1
2.2
2.8
5.2
4.4
3.7
1.7
4.4
8.7
3.6
1.9
4.2
3.3
/
52
24
62
23
39
3.3
6.2
3.
2.2
4.1
4.7
3.3
24
4.2
44
/
/
52
6.1
24
6.1
2.2
29
5.2
6.1
3.7
1.7
4.2
3.9
3.4
2.0
4.2
5.4
53
93
24
92
23
6.
3.4
82
3.7
24
42
5.5
3.3
29
4.2
488
6
32
73
24
81
22
34
34
67
3.7
L.0
4.2
4.4
3.8
22
4.2
38
Table 5: Mean absolute error results of variants of Equiformer (Lmon = 2) without attention and layer normalization the MD17 tesling set. Energy and force are in units of meV and meV/A. Index 1 and Index 2 correspond to "Equiformer (Imar = 2)" and "Equiformer ( (Imax = 2) + DeNS" in Table 4.
Mathod
Auptrin
Beasorre
Fahooel
Nepltiaaien
Salteghe arld
Toimse
Unael
Peete meeding
Ax # In Fej. 6
Partially cerrepte stractuee
Indbee
DWNS
ereegy
fienes
eergy
fieons
earegy
furow
enregy
fores
eneegy
fienes
eneegy
feeess
enregy
forow
earegy
funves
5:3
7:2
2.2
5.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
3.1
8.7
233
6.1
2.2
2.8
3.2
4.4
3.7
1.7
4.4
3.7
3.5
1.81
4.2
3.3
8
8.6
9.1
25
6.5
2.5
3.5
8.2
5.8
6.1
52
108
3.7
2.0
8.5
6.5
6.1
1.8
25
0.1
2.5
2.0
8.2
4.5
8.7
1.7
48
3.0
3.7
1.9
4.2
8.8
5.8
6.9
2.4
8.8
2.8
3.2
3.5
5.4
87
2.0
4.
4.2
4.0
2.0
4.3
3.5
Table 6: Ablation study on the design choices of DeNS using MD17 dataset. Mean absolute error results are evaluated on the testing set. Energy and force are in units of meV and meV/A. Index 1 and Index 2 correspond to "Equiformer (Lmaz= = 2)" and "Equiformer (Lmox = 2) + DeNS" 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. Following previous works, we adopt gradient methods for force prediction. 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.
Main Results. We train Equiformer (Tmar 2) (Liao & Smidt, 2023) and Equiformer (Tmar 3)
with DeNS based on their official implementation, where Lmaz denotes the maximum degree of equivariant representations. As shown in Table 4, DeNS improves the results on all molecules. Along with the results on OC20 and OC22 datasets, DeNS can generally improve the performance on force predictions with both direct (i.e., OC20 and OC22) and gradient (i.e., MD17) methods. Particularly, Equiformer (Tmaa 2) trained with DeNS acheives better results on all the tasks and requires 3.1x less training time than Equiformer (Lmaz = 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 Lmax from 2 to 3. Additionally, we find that the gains from training DeNS as an auxiliary task are comparable to pre-training. For example, Zaidi et al. (2023) uses MD-NET (Thölke & Fabritiis, 2022) pre-trained on the PCQM4M+2 dataset and reports results on Aspirin. Their improvement in force MAE is about 17.2% (Table 3 in Zaidi et al. (2023)). Training Equiformer (Imma = 2) with DeNS results in 20.8% improvement in 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 x larger than the training set of
11

---

Published in Transactions on Machine Learning Research (12/2024)
MD17. Moreover, training with DcNS enables consistent improvement in all molecules when increasing Lmar from 2 to 3, and Equiformer (Tmax 3) trained with DeNS achieves overall best results. In contrast, when DeNS is not used, increasing Lmaz from 2 to 3 can lead to overfitting and worse results on some molecules (i.c., Benzenc and Uracil). Besides, we report simulation- results in Section D.3.
Effect of DeNS on Different Network Architectures. We train variants of Equiformer (Lmaz 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 model variants. 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 general and can be helpful many equivariant networks.
Ablation Study on the Design Choices of DeNS. We use Equiformer (Lmaz 2) to justify the design choices of DeNS. The results are summarized in Table 6 and are similar to those on OC20 S2EP-2M dataset. Comparing Index 2 and Index 3, force encoding consistently results in significant improvement in energy and force MAE. Compared training without denoising (Index 1), DeNS without force encoding (Index 3) 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. Comparing Index 2 and Index 4, predicting energy given corrupted structures results in overall belter performance. Finally, the comparison between Index 2 and Index 5 demonstrates that partially corrupted structures are necessary to achieve better results. For some molecules (i.e., Ethanol and Salicyeli acid), DeNS with noise added to all atoms (Index 5) can be worse than training without DeNS (Index 1).
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 on original tasks of energy and force prodictions. Denoising non-equilibrium structures can be ill-posed problem since there are many possible target structures. To address the issue, propose force encoding and take the forces of original structures as imputs to specify which non-oquilibrium structures we are denoising. With force encoding, DeNS successfully improves the performance on 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 on energy and force predictions across datasets of various scales with minimal increase in training cost and applicable to many equivariant networks. Finally, we note that the proposed DeNS is general and can be directly applied other atomistic datasets containing non-equilibrium structures. Take Materials Project (Jain et al., 2013) for example. Similar to OC20 and OC22 datasots, for each entry in the Materials Project database, the Materials Project Trajectory (MPtrj) dataset (Deng et al., 2023) contains the corresponding relaxation trajectory between initial and relaxed structures. Most structures in MPtrj are non-equiibrrum and have energy and force labels. Therefore, we can apply DeNS to MPtrj dataset in the same way as OC20 and OC22 datasets.
Acknowledgement
We thank Larry Zitnick, Xinlei Chen, Zachary Ulissi, 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 the research results reported within this paper.
Yi-Lun Liao and Tess Smidt were supported by DOE ICDI grant DE-SC0022215.
References
Rdkit: open-source cheminformatics. URL https::/www.rrddee org/
12

---

Published in Transactions on Machine Learning Research (12/2024)
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer normalization. arriv preprint arziv: 0err:1607.060 06450, 2016.
Tles 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 (NewrIPS), 2022.
Simon Batzner, Albert Musaelian, Lixin Sun, Mario Geiger, Jonathan P. Mailoa, Mordechai Kornbluth, Nicola Molinari, Tess E. Smidt, and Boris Kozinsky. E(3)-qquivariant graph neural networks for data-efficient and accurate interatomic potentials. Nature Communications, 13(1), May 2022. doi: 10. 10.:008//4414467-02-----059 URL https: org/ 10. 1030///44660000000 020- 000-144477--222299--555
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://openrevient net/forum?id xwr8g0BeV1.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Veelakantan, 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, Tya Sutskever, and Dario Amodei. Language models are few-shot. learners. In Advances in Neural Information Processing Systems (NewrIPS), 2020.
Noss,,
Lowik Chanussot*, Abhishek Das*, Siddharth Goyal*, Thibaut Lavril Muhammed Shuaibi*, Morgane Ríviere, Kevin Tran, Javier Heras-Doming Caleb Ho, Weihua Hu, Aini Palizhati, Amuroop Sriram, Brandon Wood, Junwoong Yoon, Devi Parikh, C. Lawrence Zitnick, and Zachary Ulissi. Open catalyst 2020 (oc20) dataset and community challengos. ACS Catalysis, 2021. doi: 10.1021 10.102//acscatal..00/l /acscatall.0004555,
Stefan Chmiela, Alexandre Tkatchenko, Huziel F. Sauceda, Igor Poltavsky, Kristof T. Schütt, and Klaus- Robert Müller. Machine learning of accurate energy-conserving molecular force fields. Science Advances, e633:1,,, 2017. doi: 201116//ciadv..16000 URL https://wwet http::///wwww.ieceeeeeemmmm science. www.ooigg/oooooooooooooo..oo.ooo dot/abs//0..100/0 sciadv.1603011
Stefan Chmiela, Huziel E. Sauceda, Klaus-Robert Miller, and Alexandre Tkatchenko. Towards exact molecular dynamics simulations with machine-learned force fields. Nature Communications, 9(1), sep 2018. doi:
00.1088/s11467-018-006688.
Taco S. Cohon, Mario Geiger, Jonas Köhler, and Max Welling. Spherical CNNs. In International Conference on Learning Representations (ICLR), 2018. URL httpae///pppereeeeeeeefe/ofrrrr HddH
Mostafa Dehghani, Josip Djolouga, Basil Mustafa, Piotr Padlewski, Jonathan Heek, Justin Gilmer, Andreas Steiner, Mathilde Caron, Robert Geirhos, Tbrahim 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 Vasconeclos, 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 Xrxi::2302. 05442, 2023.
Bowen Deng, Peichen Zhong, KyuJung Jun, Janosh Riebesell, Kevin Han, Christopher Bartel, and Gerbrand Ceder. Chgnet as a pretrained universal neural network potential for Nature Machine Intelligence, 2023.
atomistic modelling.
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. orziv preprinl arziv:1810.04805, 2019.
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.
13

---

Published in Transactions on Machine Learning Research (12/2024)
An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Represenlations (TCLR), 2021. URL https: / www./eeeeeeie...eoo wtw..e//ffouurtdddtccooooTT
Alexandre Duval, Simon V. Mathis, Chaitanya K. Joshi, Victor Schmidt, Santiago Miret, Fragkiskos D. Malliaros, Taco Cohen, Pietro Liö, Yoshua Bengio, and Michael Bronstein. A hitchhiker's 's guide to geometric gnns for 3d atomic systems. arXiv preprint orXiv:2312.07511 2024.
Alexandre Agm Duval, Victor Schmidt, Alex HernándezGGaria,, Santiago Miret, Fragkiskos D. Malliaros, Yoshua Bengio, and David Rolnick. FAENet: Frame averaging equivariant GNN for materials modeling. In International Conference on Machine Learning (ICML), 2023.
Rui Feng, Qi Zhu, Huan Tran, Binghong Chen, Aubrey Toland, Rampi Ramprasad, and Chao Zhang, May the force be with you: Unified forco-centric pro-training for 3d molecular conformations, arXiv preprint arviv:2308. 14759, 2023a.
Shikun Feng, Yuyan Ni, Yanyan Lan, Zhi-Ming Ma, and Woi-Ying Ma. Fractional denoising for 3D molccular pre-training. International Conference on Machine Learning (ICML), 2023b.
Xiang Fu, Zhenghao Wu, Wujic Wang, Tian Xic, Sinan Keten, Rafacl Gomez-Bombardde Bombarclli, and Tommi Jaakkola. Forces are not enough: Benchmark and critical evaluation for machine learning force fields with molecular simulations. Transactions on Machine Learning Research (TMLR), 2023.
Fabian Fuchs, Daniel E. Worrall, Volker Fischer, and Max Welling, Se(3)-transoormers:: 3d roto-translation equivariant attention networks. In Advances in Neurol Information Processing Systems (NeurIPS), 2020.
Johannes Gasteiger, Janek GroB, and Stephan Günneeaan. Directional message passing for molecular graphs. In International Conference on Learning Representations (ICLR), 2020.
Johannes Gasteiger, Muhammed Shuaibi, Amuroop Sriram, Stephan Günnemann, Zachary Ulissi, C Lawrence Zitnick, and Abhishek Das. GemNet-OC: Developing Graph Neural Networks for Large and Diverse Molecular Simulation Datascts. Transactions Machine Learning Research (TMLR), 2022.
Mario Geiger and Tess Smidt. e3nn: Euclidean neural networks. arXiv preprinl arXiv:2007.0445 2022.
Mario Geiger, Tess Smidt, Alby M., Benjamin Kurt Miller, Wouter Boomsma, Bradley Dice, Kostiantyn Tapchevskyi, Maurice Weiler, Michal Tysskiewicz, 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 htttp::/ddoi.rgeee 1r...o..rr//1...../eeeeeoo 6459381.
Justin Gilmer, Samuel S. Schoenholz, Patrick F. Riley, Oriol Vinyals, and George E. Dahl. Neural message passing for quantum chemistry. In International Conference Machine Learning (ICML), 2017.
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 molccular property prediction and beyond. In International Conference on Learning Representations (ICLR), 2022. URL /ttp::////peeeeeeee/ee
B. Hammer, L. B. Hansen, and J. K. Nørskov. Improved adsorption energetics within density-ffucttional theory using revised perdew-burke--rnzerroo functionals. Phys. Rev. B, 1999.
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.
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.
Anubhav Jain, Shyue Ping Ong, Geoffroy Hautier, Wei Chen, William Davidson Richards, Stephen T. Dacek, Shreyas Cholia, Dan Gunter, David Skinner, Gerbrand Ceder, and Kristin Aslaug Persson. Commentary: The materials project: materials genome approach to accelerating materials innovation. APL Materials, 2013.
14

---

Published in Transactions on Machine Learning Research (12/2024)
Rui Jiao, Jiaqi Han, Wenbing Huang, Yu Rong, and Yang Liu. Energy-motivan equivariant pretraining for 3d molecular graphs. Xrviv preprinl or arXiv:2007.08: 2022.
Bowen Jing, Stephan Fismann, 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://openreeierr.rt wtt.://ww.....tooooor../.oo.....
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.
Janice Lan, Aini Palizhati, Muhammed Shuaibi, Brandon M Wood, Brook Wander, Abhishek Das, Matt Uyttendacle, C Lawrence Zitnick, and Zachary W Ulissi. AdsorbML: Accelerating adsorption energy calculations with machine learning. arXiv preprint ar 1rXi:22211. 16486, 2022.
Tuan Le, Frank Noé, and Djork- Arné Clevert. Equivariant graph attention networks for molecular property prodiction. arXiv preprint ar Xrv::2200..0 Xix:2200.008911 2022.
Yi-Lum Liao and Tess Smidt. Equiformer: Equivariant graph attention transformer for 3d atomistic graphs. In International Conference on Learning Represenlations (ICLR), 2023. URL https:/ httpe////ppeeeeeeeeeeeee forum?id=KumPp Ad-RWWPARRROOD
Yi-Lum Liao, Brandon Wood, Abhishek Das*, and Tess Smidt EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree Representatious. arziv preprint orziv:2306.1200 120599 2023.
Shengchao Liu, Hongyu Guo, and Jian Tang. Molecular geometry pretraining with SE(3) -ivvarant denoising distance malching. In International Conference on Leurning Representations (ICLR), 2023. URL https: ///ppprreveee..mt/ttroor/reeeeee met://orur7dd-CTHHoO...
// //openreview
Benjamin Kurt Miller, Mario Geiger, Tess E. Smidt, and Frank Noé. Relevance rotationally equivariant convolutions for predicting molecular properties. arric preprint arriv:2008.05: 08461, 2020.
Albert Musaelian, Simon Batzner, Anders Johansson, Lixin Sun, Cameron J. Owen, Mordechai Korubluth, and Boris Kozinsky. Learning local equivariant representations for large-scale atomistic dynamics. arriv preprint araiv:2204.0523 2022.
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:2344.10061, 2023.
Maho Nakata and Tomomi Shimazaki. Pubchemqe project: A large-scale first-principles electronic structure database for data-driven chemistry. Journal of chemical information and modeling, 57 6:1300-1308, 2017.
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.
Joshua 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 4(1):015, feb 2023. doi: 10.1088/180 2262-22533,-22 acb3114 URL https://dx: http//xxxxiioio.ooooo/////ooo orgg /10. 10011/222222213///aaom acb314.
Raghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp, and O Anatole von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific Dala, 1, 2014.
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 Michalcas. Interactive supercomputing on 40,000 cores for machine learning and data analysis. In 2018 IEEE High Performance extreme Computing Conference (HPFC), pp. 1-6. IEFE, 2018.
15

---

Published in Transactions on Machine Learning Research (12/2024)
Lars Ruddigkeit, Rund van Deursen, Lorenz C. Blum, and Jean-Louis Reymond. Enumcration 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: /0.1021/ci304155d URL https:///oi..oo / PMID: 23088335.
Alvaro Sanchez- Gonzalez,, Jonathan Godwin, Tobias Pfaff, Rex Ying, Jure Loskovee, and Peter W. Battaglia. Learning to simulate complex physics with graph networks. In International Conference on Machine Learning (ICML), 2020.
Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. B(n) equivariant graph neural networks. In International Conference on Machine Learning (TCML), 2021.
K. T. Schütt, P.-J. Kindermans, E. Sauceda, S. Chmiela, A. Tkatchenko, and K.-R. Müller. Schnet: A contimuous-filter convolutional neural network for modeling quantum interactions. In Advances in Neural Information Processing Systems (NewrIPS), 2017.
Kristof T. Schütt, Farhad Arbabzadah, Stefan Chmniea,, Klaus R. Müller, and Alexandre Tkatchenko. Quantum-chemical insights from deep tensor neural networks. Nature Communications, 8(1), jan 2017. doi: 10. N0.1038//ccomms1333
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.
Justin S. Smith, Olexandr Isayev, and Adrian E. Roitberg. Ani-1: A data set of 20m off-equilibrium dft calculations for organic molecules. arXiv preprinl arXiv: 04988, 2017.
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.
Nitish Srivastava, GeolTrey Hinton, Alex Krizhevsky, Πya 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. hwwp:ll...r..pprrr////eeeeeee
html.
Yi Tay, Mostafa Dehghani, Vinh Q. Tran, Xavier Garcia, Jason Wei, Xuczhi Wang, Hyung Won Chung, Dara Bahri, Tal Schuster, Steven Zheng, Demny Zhou, Neil Houlsby, and Donald Metzler. UL2: Unifying language learning paradigms. In International Conference on Learning Representations (ICLR), 2023. URL https://openreviee pea//foruntdd-GrrWLEEZZZ7
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.o www.:nneeeeeeeeeeeeee www//eeeettaniee hatf/foruu?ddd=lllzzz2zzzz
Nathaniel Thomas, Tess E. Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field works: Rotation- and translation-eqquivvriant neural networks for 3d point clouds. arziv preprint Xie: 08219, 2018.
Raphael J. L. Townshend, Brent Townshend, Stephan Eismann, and Ron O. Dror. Geometric prediction: Moving beyond scalars. arXie preprint wxXi:22006.1100 2020.
Richard Tran*, Janice Lan*, Muhammed Shuaibi*, Brandon Wood*, Siddharth Goyal*, Abhishek Das, Javier Heras-Domingo, Adcesh Kolluru, Ammar Rizvi, Nima Shoghi, Anuroop Sriram, Zachary Ulissi, and C. Lawrence Zitnick. The open catalyst 2022 (oc22) dataset and challenges for oxide electrocatalysis. anXiv preprint ar xrXx:22266.88) 2022.
Oliver Thorsten Unke, Mihail Bogojeski, Michael Gastegger, Mario Geiger, Tess Smidt, and Klaus Robert Muller. -quuvvaqqa prediction of molccular wavefunctions and electronic densitics. In A. Boygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems (NewrIPS), 2021. URL https: https::/ /ttpa//////peeeeeeeeeeeee net wttp://wwww...eeeoee////aaaaaaaa
16

---

Published in Transactions on Machine Learning Research (12/2024)
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 (NewrIPS), 2017.
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.
Pascal Vincent, H. Larochelle, Isabelle Lajoie, Yoshua Bengio, and Pierro-Antoine Manzagol. Stacked denoising autoencoders: Learning useful representations in deep network with local denoising criterion. Journal of Machine Learning Research, 2010.
Yuyang Wang, Changwen Xu, Zijie Li, and Amir Barati Farimani. Denoise pretraining on nonequilibrium molecules for accurate and transfcrable neural potentials. Journal of Chemical Theory and Computation, 2023.
Maurice Weiler, Mario Geiger, Max Welling, Wouter Boomsma, and Taco Cohen. 3D Steerable CNNs: Learning Rotationally Equivariant Features in Volumetric Data. In Advonces in Newral Information Processing Systems 32, pp. 10402-10413, 2018.
Sheheryar Zaidi, Michael Schaarschmidt, James Martens, Hyumjik 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.
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. 1tt.:/eeeeeeet..2.1443111. 1108/PhyyRevLett..... 120. 143001. URL https://link... aps org//oi/1/ 1103//hyysseeeeet.tt.mmm 120. 143001.
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 Newral Information Processing Systems (NeurIPS), 2022.
17

---

Published in Transactions on Machine Learning Research (12/2024)
Appendix
A Related Works
A.1
Comparison to previous works on denoising
A.2
eequvaatt networks
A.3
Equiformer series
B Details of experiments OC20
B.1 Training details
C Details of experiments on OC22
C.1 Training details
D Details of experiments MD17
D.1 Additional details of DeNS
D.2 Training details
D.3 Additional simulation-base results
E Pscudocode for force encoding
Pseudocode for training with DeNŞ
G Visualization of corrupted structures
A Related Works
A.1 Comparison to Previous Works on Denoising
We discuss previous works on denoising (Godwin et al., 2022; Zaidi et al., 2023; Feng et al., 2023b; Wang et al., 2023) in chronological order and compare them with this work 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 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 denoising is the relaxed, equilibrium structure. Denoising without force encoding is well-defined on both QM9 and OC20 IS2RE dalasets. In contrast, this work proposes using force encoding generalize their approach non-equílibrium structures, which have much larger datasets than equilibrium ones. Force encoding achieve better results OC20 S2EF-2M dataset with little overhead (Index and Index in Table 1(b)) and is indispensable 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 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 (rdk) 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 pro-trainod 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 MD17 dataset, Zaidi et al. (2023) simply follows the same practice
18

---

Published in Transactions on Machine Learning Research (12/2024)
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 in 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 to ANT-1 (Smith et al., 2017) and ANI-Ix (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.
A.2 S Networks
We first discuss the concept of equivariance and how equivariant networks achieve equivariance and then compare previous works below. We note that most of the content adapted from Equiformer (Liao & Smidt, 2023) and EquiformerV2 (Liao et al., 2023). We refer readers to the works (Liao & Smidt, 2023; Liao et al., 2023) for more detailed background equivariant networks and the work (Duval et al., 2024) for broader review on geometric GNNs for modeling 3D atomistic systems.
3D atomistic systems are often described in 3D coordinate systems. We have the freedom to choose arbitrary 3D coordinate systems since we can change between different coordinates via the symmetries of 3D space. The relevant 3D symmetris are rotation, translation and inversion. The Euclidean group E(3) consists of 3D rotation, translation and inversion while the special Euclidean group SE(3) is comprised of 3D rotation and translation. The laws of physics remain the same regardless of the coordinate we use, and thus the properties of 3D atomistic systems are equivariant to 3D symmetries. For instance, when 3D atomistic system is rotated, quantities like energy will remain identical while others like forces will rotate accordingly. Mathematically, a function f mapping between vector spaces X and Y is equivariant to a group of transformations G if for any input & E X, output y E Y and group element g E G, we have [(Dx(g)=) = (g) = Dy(g)y, where Dx(g) and (g(g) are transformation matrices or group representations parametrized by g in X and Y. Additionally, f is invariant when Dy (g) is an identity matrix for any g ge G.
Incorporating equivariance into neural networks inductive biases can improve data efficiency and generaliz- ability. 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 el. 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) achiove equivariance to 3D rotation and optionally inversion by using vector spaces of irreducible representations (irreps) as equivariant features. The vector spaces of irreps are (2L + 1)-dimensional, with degree L being a non-negative integer. L can be viewed as the angular frequency of vectors and determines how fast vectors change with respect to rotation of the coordinate system. Vectors of degree L are referred to as type-L L vectors. They are transformed with Wigner-D matrices ((L) when rotating the coordinate system, and (L) of different L acts on independent vector spaces. Euclidean vectors F in R33 such as relative positions and forces can be projected into type-L vectors by using spherical harmonics ((()( [ We concatenate multiple type-T, vectors to build an equivariant feature f. Given the maximum degree Timar of equivariant features, f has Cr type-L vectors, where 0 & L & Limex and Cz, is the number of channels for type-L vectors. Both Lmae and CT are architectural hyper- paper--praeeerrs of equivariant networks. Equivariant operations are applied to equivariant features to preserve equivariance, and two examples are tensor products and SO(3) linear operations. Tensor products are the fundamental operation in equivariant networks for interacting vectors of different L and are used to build convolutions and attention. On the other hand, SO(3) linear operations apply separate linear operations to each group of type-L vectors. Relevant to the proposed method, we encode forces into equivariant features by first projecting forces to type-L vectors with spherical harmonics and then expanding the number of channels to CT with an SO(3) linear operation.
Previous works on equivariant networks mainly differ in which equivariant operations are used and the combination of those operations. TFN (Thomas et al., 2018) and NequlP (Batzner et al., 2022) use tensor products for equivariant graph convolution with lincar messages, with the latter utilizing extra gate activation (Weiler et al., 2018). SEGNN (Brandstetter et al., 2022) applies gate activation to messages passing
19

---

Published in Transactions on Machine Learning Research (12/2024)
for non-linear messages (Gilmer et al., 2017; Manchez-Gonzaler et al., 2020). Transsfraasf (Fuchs et al., 2020) adopts equivariant dot product attention with linear messages. Equiformer (Tiao & 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 compule-intensive SO(3) convolutions built from tensor products, and therefore they can only use small values for maximum degrees Lmaz 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 Imaaz EquiformerV2 (Tiao et al., 2023) adopts eSCN convolutions and proposes an improved version of Equiformer to better leverage the power of higher Lmaz achieving the current state-of-the-art results on OC20 (Chanussot* et. al., 2021) and OC22 (Tran* $ et al., 2022) datasets.
A.3 Equiformer Series
Equiformer (Tiao & Smidt, 2023) is an (()))//(()equuvvariin) graph neural network that combines the inductive biases of 3D-related equivariance with the strength of Transformers (Vaswani et al., 2017). Starting from Transformers, Equiformer introduces three architectural modifications. First, Equiformer adopts equivariant features built from vector spaces of irreps as internal representations to incorporate equivariance. Second, equivariant operations are applied to the equivariant features. These operations include tensor products and the equivariant counterparts of the original operations in Transformers. The latter part consists of equivariant linear operations (i.e., SO(3) linear operations), equivariant layer normalivation (Ba al., 2016) and gate activation (Weiler et al., 2018). Third, Equiformer proposes to apply non-linear functions to both attention weights and message passing, which improves the expressivity of attention in standard Transformers.
Although Equiformer demonstrates that Transformers generalize well to 3D atomistic systems, it is limited to small values of maximum degree Lmaz because of the compute-intensive tensor product operations. Higher degrees can better capture angular resolutions and directional information, and therefore lower I.max can limit the expressivity of Equiformer. To address this limitation, EquiformerV2 (Liao et al., 2023) adopts eSCN (Passaro & Zitnick, 2023) convolutions, which significantly reduce the computational complexity of tensor products, to incorporate higher-degree equivariant representations and proposes three architectural improvements to better leverage the power of higher degrees. First, attention re-normalization is proposed to introduce one additional layer normalization to the non-lincar functions of attention weights. This helps stabilize attention and improves empirical performance. Second, separable 5 activation based on 5 activation (Cohen et al., 2018) is proposed to better mix the information of all degrees and stabilize training. Third, separable layer normalization is proposed to replace the original cquivariant layer normalization order to preserve the relative importance of different degrees.
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 and we only predict forces of free atoms, we only apply DeNS to free atoms. When partially corrupted structures are used, we add noise to and denoise a random subset of free atoms. When training EquiformerV2 on OC20 A2D datasct, we only apply DeNS to structures from the All split.
For force prediction, we adopt direct methods following previous works for a fair comparison. We add an additional block of equivariant graph attention to EquiformerV2 V2 for noise prediction. We mainly follow the of training Equiformer without DeNS on OC20 S2EF-2M and -2EF-AII+MD datascts. For training EquiformerV2 on OC20 S2EF-AIl+M 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 1(a). Table summarizes the ablation studies on OC20 S2FF-2M dataset in Section 4.1.1 and for the main results on OC20 A2D+DAII+
datasel, we increase the number of epochs from 1 to
of training EquiformerV2 V2 with DeNS for the
20

---

Published in Transactions on Machine Learning Research (12/2024)
EquifornerV2 (89M) on
EguiformerV2 (160M) on
Hyper-parameters
OCa1 S2RP-2M dataset
OCau SEEF-AII-MD dataent
Optinizer
AdooW
AdamW
Learning zate schoduling
learning rate Emoar warmap
Cosine learning rase with warmap
Warmup spochs
0.1
0.01
Maxiimmm leazning tate
2 10"1 spochs
10-1
4xx 1001 for 20, 30 spochs
Bateh size
12 epocte
512
128 for 20, opochs
Number of spachts
12, 20,
Weight donay
1 x 10-
10-
Dropout vale
0.1
0.1
Stochassic depth
0.05
0.1
Euergy coeflicient Ag:
2
Force coefficieut
100
100
Guadieal clipping worm threshold
1010
100
Model EMA decay
0,999
0.999
Cutoff radius (A)
12
12
Maxicam nambhr of neighbors
20
Number of radial bases
600
600
Dimension of hirlden sealar festuures in racilal funetions degge
(0,
[0, 128)
Maxinamm degree Lese
Maricam ceder Meer
Number of T'raueforteer blocks
12
201
Emberlinng d'nnension dunded
[6. 128)
(5 dimeresion daata holdow
[6. 64)
Number of alltenion lisedd h
8
(2)) dimuension daten slgin
(0,61)
(0.61)
dimesusion dats edder
Ld)
[6. 16)
Hidden dimerrsion in feed forwaed networks days.
[6. 128)
[6. 128)
Resolation of poins samples
Probability of optimising DeNS PDiNS
0.5
0.125
DeNS coefficient ADuss
Standard deviation of Gaussion noise
0.1 for 12 opochs
0.1
0.15 for 20, 30 epocis 0.5
Commption ratoo rDuas
0,25
Table 7: Hyper-parameters training Equiformer V2 with DeNS on OC20 S2EF-2M dataset and OC20 M2DD-AII+ dataset.
dataset in Section 4.1.2. Please refer to the work of EquiformerV2 (Liao et al., 2023) for details of the architecture. For training eSCN with DeNS, we use the same DeNS-related typer-parameters as those for training EquiformerV2 for 20 epochs and the same module as force prediction to predict noise.
V100 GPUs with 32GB are used to train models. We use 16 GPUs for training EquiformerV2 V2 for 12 epochs and cSCN on OC20 S2EF-2M dataset and use 32 GPUs for training EquiformerV2 for 20 and 30 epochs. We train EquiformerV2 with 128 GPUs on OC20 S2EF-AII+MD dataset. The training time and the numbers of parameters of different models on OC20 S2EF-2M dataset be found in Table (a). For OC20 S2EF-AlI+ dataset, the training time is 88982 GPU-hours and the number of parameters is 160M.
C Details of Experiments on OC22
C.1 Training Details
Different from OC20, all the atoms in structure in OC22 are free, and we apply DeNS to all the free atoms. We add an additional block of equivariant graph atlention to EquiformerV2 for noise prediction. We follow the same practice as on OC20 and use direct methods for force prediction. We follow the same parametrsssses not relevant to DeNS, and Table 8 summarizes the hyper-parmmeters for the results on OC22 in Table 3. We use 32 V100 GPUs (32GB) for training. The training time 5082 GPU-hours, and the number of parameters is 127M.
21

---

Published in Transactions on Machine Learning Research (12/2024)
lyper-pararecter
Value or deseription
Optimizer
AdamAW
Tearning sate schoduling
Cosine loarning rate with linerr warmup
Warmup epochs
Maximum loarning nate
2xx 10
Number of epochs
6
Weight decay
1
Dropout raño
0.1
Stodhnslic depili
Energy confficient, Ag
4
Force coellicieul A
100
Gradient clipping
threshold
50
Model EMA deea
0. ..999
Outoff radius (A)
Maximam number of neighhors
20
Number radial bases
600
Dimensioo of hidden sealar features in radinl fimetions dedge
(0. 128)
Maximam degree laosz
Maximam order Mane
Number Transfotmer blocks
Erabedding dimecnsion dentted
(6. 128)
(i dimerssion dattn hildew
(6. 64)
Number attention heads
10 dimension siyhe
64]
Value dimceusion datta.attee
(6. 16]
dimension
forward neeworks dyts
128)
Resolution of point samples R
18
Probability optimising DeNS powes
DeNS coefficiens Anass
25
Standard
0.15
Compption ratoo "DoN
0.5
Table 8: Hyper-parameters OC22 dataset.
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 sample structures for DeNS when only single GPU is used for training models the MD17 dataset. We zero out forces corresponding to structures used for the original task that single forward-backwarr 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 official implementation of Equiformer (Liao & Smidt, 2023) for experiments on the MD17 dataset. and follow most of the original hyper-parameters for training with DeNS. Gradient methods are used to predict forces. For training DeNS, we use an additional block equivariant graph attention for noise prediction, which slightly increases training time and the number of parameters. The hyper-parameters introduced by training DeNS and the values of energy cocfficient Ag and force cocfficient Ay on different. molecules can be found in Table 9. Empirically, we find that linearly decaying DeNS coefficient ADaNS 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 x 44
We use one A5000 GPU with 24GB to train different models for each molecule. The training time and the numbers of parameters be found in Table 4.
22

---

Published in Transactions on Machine Learning Research (12/2024)
Hyperspramnter
Aspirin
Berwen
Rihonol
Malonaklelyde
Naphtholess
Salleylie
Thllee
Unael
Energy coefleieet Apt
1
1
2
1
1
100
Prohabläty of optimieing DcNS Posse
0.25
0.25
0.25
0.25
0.25
0.25
0.125
0.25
DeNS confliciout Apaas
5
5
5
5
5
6
3
5
Standard devintion Ganssinn
0.005
Corruption natio Fieaa
0.28
0.25
0.25
0.25
0.25
0.25
0.25
0.25
Table 9: Hyper- parameters of training Equiformer (Lmaz 2) and Equiformer (Lmes 3) with DeNS the MD17
dataset. Other
not listed here are the same as the original Equiformer trained without DeNS.
Aepirin
Fthanil
Napketaalene
Saleylle acld
Model
energy
lorees]
stabilly
A69)
energy.
Toceesi
wabillis
eoerge
foree]
stabilty1
eneggy
Torres
nip)
Equitoraee Chass Equiiormee Bhaas
5.5
7:2
5000
0,02
2.2
3.8
260.9
0.00
3.7
211
135.8
0.12
4.5
1.1
500
0,03
1x85
5.1
8.7
500
0.02
2.2
2,8
300
0.00
3.7
1.7
187.2
0.12
4.4
3.7
300
0.05
Table 10: Simulation- based results on MD17 dataset. We report simulation based metrics, stability and distribution
of interatomic distances h(r). Models are the same as in Table 4. Energy and force are in units of meV and meV/A and are evaluated the testing set.
D.3 Additional Simulation-Ban Based Results
Following the work (Fu et al., 2023), we run simulations on the four molecules (i.e., Aspirin, Ethanol, Naphthalene and Salicylic Acid) and compare the two simulation-based metrics, which are stability and distribution interatomic distances h(r). We use the previously trained Equiformer for the simulations. We note that Equiformer models are trained on 950 examples for each molecule instead of 9, 500 as in Fu et al. (2023). The results are summarized Table 10. Training with DeNS helps energy and forces MAE as well stability. The results of l(r) similar. For Naphthalene, training with DeNS improves stability from 183.8/300 to 157.2/300. Both models become unstable when running simulations of Naphthalene, and we surmise that is because we only use 950 examples for training instead of 9, 500. For others, the performance gain in iocllation-baser metrics is not significant since the original Equiformer trained on 950 examples already achieves similar results to other top-performing models trained on 9, 500 examples, suggesting that the potential room for improvement would be quite limited.
E Pseudocode for Force Encoding
We provide the pscudocode for force encoding in Algorithm 1. Here the original node embedding 2; contains only the alom embedding Zt,z- Note that the lype-L vectors in Wx,z are all zeros for since æi,z is obtained by applying an SO(3) linear layer type-0 vectors. We directly add the force embedding oı, the original node embedding encode forces.
Algorithm Force Encoding
1: Fi,= - S03 Linear )ot(2)) t> Extend Equation 5 to all degrees L and apply to the onc-hot.
encoding of atomic numbers zz for each atom
3; Zi,=
Fi,f SO3 Linear (|EE]--
[
4: D2 < I₁ + Zi.F
F Pseudocode for Training with DeNS
We provide the pseudocode for training with DeNS in Algorithm 2 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 random subset of atoms (Line 11 14) and predict the corresponding noise (Line 25).
23

---

Published in Transactions on Machine Learning Research (12/2024)
Algorithm Training with DeNS
1: Input:
PDeNs: probability of optimizing DeNS
ADeNS: DeNS coefficient
a: standard deviation Gaussian noise
"DeNs: corruption ratio
Ag: energy coefficient
Apt force coefficient
GNN: graph neural network for predicting energy, forces and noise
2: while training do
3: Liotal = 0
4: Sample a batch of B structures {(Snon-eq)' 1 [l....BB) from the training set
for 1 do
t- This for loop can be parallelized
Let (Snon-eg)² (xx ie (1....,| (S
7:
Sample p from uniform distribution U(0, 1) to determine whether to optimize DeNS >> Optimize DeNS based on Equation
8:
if < PDeNS then for i to (Snon.eg)' do 9i ~ U(0, 1)
9:
10:
11:
if q. < 7DNNS then
Add noise and denoise the atom
12:
6, oTa)
13:
P; P; + 6; 1
14:
Denoise the alom when calculating DDNSS
15:
else
16:
p; = Pi
17:
mi;
18:
end if
19:
= f. m,
Encode the atomic force the atom is corrupted
20:
end for Tet (Snon-eg)² = (2), po) liE (1,..., (aw.lISSnonnnnnn (noon-cg)
21:
22:
Let F ((Smoureq)))
P.
iE (1, . (Snon-eq) |3
É, S. GNN n(Snon-eg)) 1 ((Swon-eq)*) LE E' ((Snon-cq)
23:
24:
P.
b- Predict energy of the original structure
25:
DDNS [Snoneoy |
(S -==
m, 1/- =11²
Predict noise of corrupted atoms
26:
Lp [(Snomeen)]
(SSonnca))) 2i=1
- mc) 1P, ((Snon-oa) )) - f/2
t> -Pediitt forces of uncorrupted
atoms
27:
Liotal Liotal + AE LE + ADONS LDONS \F LF
28:
else
Optimize the original task based on Equation
E, F. F_ GNN ((Snon-eq))
29:
30:
CE E' ((Snonreq)")
CF = [(Snonen))) 2 (((nomeag)) |P/ ((Snon-ea))) - f
31:
32:
Liotal Lrotal + dg CE + d LE
33:
end if
34:
end for
35:
Lintal Liotal
B
36:
Optimize GNN based on Lrotal
37:
end while
24

---

Published in Transactions on Machine Learning Research (12/2024)
G Visualization of Corrupted Structures
We visualize how adding noise of different scales affects structures in OC20, OC22 and MD17 datascts in Figure 3, Figure and Figure 5, respectively.
25

---

Published in Transactions on Machine Learning Research (12/2024)
Figure 3: Visualization of corrupted structures in OC20 dataset. We add noise of different scales original structures (columm 1). For each row, we sample e (o, Is), multiply with a 0.1 (column 2).00 0.3 (columm 3) and 0.5 (column 4), and add the sealed noise to the original structures. For columus 2, and 4, the lighter colors denote the atomic positions of the original structures. Here we add noise the atoms structure for better visual effects.
26

---

Published in Transactions on Machine Learning Research (12/2024)
Figure Visualization of corrupted structures OC22 dataset. We add noise of different scales to original structures (column 1). For each row, sample (column 2). 0.3 (column 3) and 0.5 (column 4), and add the scaled noise to 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 structure for better visual effects.
N0. Ta), multiply with
27

---

Published in Transactions on Machine Learning Research (12/2024)
Figure 5: Visualization of corrupted structures in MD17 dataset. We add noise of different scales original structures (column 1). For each row, we sample &: (0 Is), multiply with a 0.01 (column 2), 0.03 (column 3), 0.05 (column 4) and 0.07 (column 5), and add the scaled noise to the original struc- tures. columns 2, 3, 1 and the lighter colors denote the atomic positions the original structures. Here we noise to all the atoms in a structure for better visual effects.
28


Reviews:
[{"review_id": "dqCBT8RemI", "reviewer": "Reviewer_q8wm", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- The authors have the great insight of recognizing that the force information is already present in denoising tasks for equilibrium structures, and, leverage the forces of non-equilibrium conformers effectively expand the denoising task to non-equilibrium structures, helping to disentangle the ambiguity associated with non-equilibrium target structures.\n- It is well-written and concise, presenting a clear and focused narrative on the DeNS approach definition, function and resulting improvements.\n- The experiments are thorough, including detailed ablation studies that separate the effects of each design choice, adding to the robustness of the findings.\n- Applying DeNS to the Equiformer model results in improvements across multiple datasets, enhancing both prediction accuracy and training efficiency, which is a great achievement.\n\nWeaknesses:\n- The main weakness of this work is the lack of any benchmarks beyond energy or force errors. It is well known that demonstrating low. errors on energies and forces for test structures does not tell the full picture for models that are to be used for MD simulations or optimization. It would be very important to see whether DeNS improves the performance in any of the MD-related metrics such as stability and distribution of interatomic distances.\n- Other more minor complaints are the very short introduction about equivariant models in general, and the lack of information about the Equiformer model architecture.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This manuscript introduces Denoising Non-Equilibrium Structures (DeNS) as a novel spin on the popular auxiliary task of training models to denoising equilibrium structures to improve atomistic models' performance on tasks like energy and force prediction. The primary contribution is the extension of denoising approaches from equilibrium to non-equilibrium structures, which comprise the majority of atomic data but lack a unique target structure due to the presence of non-zero forces. This approach utilizes the forces from the original, unperturbed structure to guide denoising and reduce the inherent ambiguity in denoising non-equilibrium structures. Due to the requirement of explicitly encoding and representing the forces of the original structure, this approach is particularly well-suited to neural networks that are equivariant to the 3D roto-translational group SO(3).\nThe authors present experiments on a variety of datasets such as OC20, OC22, and MD17, showing that this approach simultaneously improves performance while reducing training time for energy and force prediction for non-equilibrium structures.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- A brief explanation of the concept of equivariance and equivariant models' foundational principles and basic operations would aid accessibility for readers less familiar with the topic. If space is constrained, these details could be included in an appendix.\n- Additional details on the Equiformer model's design would make the paper more self-contained, this information should ideally be included in the appendix.\n- Some additional evaluations based on ML generated molecular dynamics trajectories, such as stability and distribution of interatomic distances, would definitely be useful and add a further use-case relevant demonstration of the untility of the DeNS approach.\n- Most forces prediction models tend to obtain the forces as the analytical derivative of the energies, which guarantees energy conservation. While the alternative approach taken by the authors has become more common as of late, the authors still need to discuss and justify their choice of directly predicting the forces, thus not guaranteeing energy conservation. Commenting on whether this non-conservative approach might extend to models with explicitly energy-conserving forces could also be useful for other readers.", "broader_impact_concerns": "No broader impact concerns to note.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "8yTAvVkvl6", "reviewer": "Reviewer_7SRu", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths\n- The paper is quite clear in its setup and execution.\n- Insights into the challenges of denoising non-equilibrium structures.\n- Improvements in force and energy predictions across different datasets.\n\n### Weaknesses and Questions\n- Given the existence of multiple equilibrium structures (e.g., molecular conformers), can the model handle situations where a corrupted non-equilibrium structure is more closer to a different equilibrium structure, leading to non-consistent denoising?\n- Are the ground truth forces also used during the evaluation time as an encoding for corrupted atoms to predict the ground truth forces and energy?\n- How can the above method be adapted for real-world scenarios where force and energy estimation is required for a given non-equilibrium structure without prior knowledge of its forces, a context in which models like MACE and NequIP typically operate?\n- Do you create one corrupted structure per molecule or many while training?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes DeNS (Denoising Non-equilibrium Structures), an auxiliary task for neural networks in 3D atomistic systems, aimed at better leveraging limited training data. It also offers crucial insights into the challenges of denoising non-equilibrium structures, particularly addressing the ambiguity stemming from multiple possible target structures for a given input and propose a novel solution by incorporating force embeddings\n\nThe efficacy of method is evaluated through comprehensive evaluations, with significant improvements in both force and energy predictions across OC20, OC22, and MD17 datasets.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Clarifying the points in the previous section.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "xCAkjR14N7", "reviewer": "Reviewer_rRiX", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths:**\n* The paper presents a practical auxiliary training method for equivariant models on atomistic modeling. DeNS has the potential to be applied to diverse modeling problems as an auxiliary loss term.\n* The paper describes the method formulation in great detail and provides a detailed discussion of prior work on equivariant networks and denoising approaches.\n* The experimental analysis on Equiformer-V2 is quite detailed for the datasets studied.\n\n**Weaknesses:**\n* The paper could be further strengthened with some additional discussion on how DeNS could be applied to other types of modeling settings. Since the authors mention that partially corrupted structures often train better, it might be good to have a discussion on how DeNS might be practically applied to other situations.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes a new technique for training of equivariant networks called denoising non-equilibrium structures (DeNS). The key differentiator in DeNS is that it enables training of equivariant networks on non-equilibrium atomistic structures by using the force as an additional feature for denoising. Including the force makes the problem no longer ill-posed as the structure to be denoised is better identified. \n\nThe paper starts by providing an introduction and motivation in Section 1, followed by a discussion of related works in Section 2. Both sections focus on prior denoising works and geometric deep learning architectures. Section describes the main method, including the setup of DeNS, the DeNS force encoding and the formulation of DeNS as auxiliary loss for training models on energy and force prediction tasks. Section 3 also includes a discussion on the reasons behind why DeNS appears to be beneficial in training equivariant networks focused on data augmentation, connections to self-supervised learning and marginal training time increases.\n\nSection 4 outlines the main experiments of the paper which focus primarily on training EquiformerV2 with DeNS and understanding the difference in performance on OC20, OC22 and MD17. The results generally show that DeNS is helpful for OC20 and OC22 with some more mixed results on MD17. Section 5 ends the paper with a conclusion.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "**Requested Changes**\n* Provide more detail on how the force of the disturbed structure is calculated and/or provided as input. \n* Discussion of additional related work, including a review of atomistic modeling using GNNs [1]. It would also be interesting to discuss arriving at equivariance via data augmentations [2] as those methods appear to be getting more attention due to their lower computational costs. The paper could also discuss additional datasets [3] and modeling settings [4] where DeNS could be relevant. \n\n\n**Nice to Have**\n* Did you do any ablation on equilibrium-based training compared to DeNS? \n* Do you think DeNS has any effects on the learning manifold? Prior work has tried to visualize energy landscapes to understand this further [5]. Maybe this can provide further details on why partially perturbed structures work better.\n* You briefly mentioned that equivariant models are most suited because forces are an equivariant quantity. What do you think would be needed to extend to other modeling methods? Do you think DeNS could be useful in setting with different materials representations? [6] [7]\n\n\n\n[1] Duval, Alexandre, et al. \"A Hitchhiker's Guide to Geometric GNNs for 3D Atomic Systems.\" arXiv preprint arXiv:2312.07511 (2023).\n\n[2] Duval, Alexandre Agm, et al. \"Faenet: Frame averaging equivariant gnn for materials modeling.\" International Conference on Machine Learning. PMLR, 2023.\n\n[3] Jain, Anubhav, et al. \"Commentary: The Materials Project: A materials genome approach to accelerating materials innovation.\" APL materials 1.1 (2013).\n\n[4] Lee, Kin Long Kelvin, et al. \"Towards foundation models for materials science: The open matsci ml toolkit.\" Proceedings of the SC'23 Workshops of The International Conference on High Performance Computing, Network, Storage, and Analysis. 2023.\n\n[5] Bihani, Vaibhav, et al. \"Low-Dimensional Projections for Visualizing Energy Landscapes of Atomic Systems.\" AI for Accelerated Materials Design-Vienna 2024.\n\n[6] Damewood, James, et al. \"Representations of materials for machine learning.\" Annual Review of Materials Research 53.1 (2023): 399-426.\n\n[7] Alampara, Nawaf, Santiago Miret, and Kevin Maik Jablonka. \"MatText: Do Language Models Need More than Text & Scale for Materials Modeling?.\" AI for Accelerated Materials Design-Vienna 2024.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The reviewers all agree that this work has value, it provides support for the claims and is interesting for the community.
Decision: Accept (poster)