Learned Force Fields Are Ready For Ground State Catalyst Discovery

We present evidence that learned density functional theory (``DFT'') force fields are ready for ground state catalyst discovery. Our key finding is that relaxation using forces from a learned potential yields structures with similar or lower energy to those relaxed using the RPBE functional in over 50\% of evaluated systems, despite the fact that the predicted forces differ significantly from the ground truth. This has the surprising implication that learned potentials may be ready for replacing DFT in challenging catalytic systems such as those found in the Open Catalyst 2020 dataset. Furthermore, we show that a force field trained on a locally harmonic energy surface with the same minima as a target DFT energy is also able to find lower or similar energy structures in over 50\% of cases. This ``Easy Potential'' converges in fewer steps than a standard model trained on true energies and forces, which further accelerates calculations. Its success illustrates a key point: learned potentials can locate energy minima even when the model has high force errors. The main requirement for structure optimisation is simply that the learned potential has the correct minima. Since learned potentials are fast and scale linearly with system size, our results open the possibility of quickly finding ground states for large systems.


Introduction
In July 2022 the single biggest workload of the largest UK supercomputer, ARCHER2 [6], was density functional theory [22,26] calculations using the Vienna ab initio Simulation Package ("VASP") [29,28,27]. The hope of such research is that by screening a very large number of compounds for specific properties, we can find answers for the grand scientific challenges of our transition to a green economy: better batteries, solar cells and cheap renewable fuels. This project also has significant expense: ARCHER2 cost £79M to build.
An important use case for DFT is structure optimisation. The goal of structure optimisation is to find the global minimum of the potential energy surface (PES) of a material, a quantity that helps characterize its functional properties. The success of the optimisation depends on the complexity of the PES, the optimisation algorithm and the DFT approximation known as the exchangecorrelation functional.
Here we present the surprising result that machine learned potentials trained on DFT data con-sistently find low energy structures of the target functional outside of the training distribution. This occurs even when force estimates are inaccurate or noisy, and even when the model is trained on a simple harmonic approximation to the energy around ground states in the training set. Our key result is shown in Fig. 1: in over 50% of cases on out of distribution tasks, our learned potential finds structures with similar or lower energies to those optimised with RPBE in VASP using the same initial geometry. The median convergence time is one minute, which is two orders of magnitude faster than traditional DFT codes. This is despite being trained on a dataset where training labels do not represent true minima either.
Our findings have a surprising implication: learned force fields may already be ready for use in ground state discovery for catalytic systems. If DFT optimisation is a useful tool for catalyst discovery, and force field optimised structures have a lower mean energy than DFT geometries, then force fields are useful for catalyst discovery. This is despite the fact that learned models still have non-negligible errors with respect to DFT energies and forces. Our experiments suggest that learned potentials succeed because they have the correct energy minima, even if their force errors are high. To test this hypothesis, we trained a force field on a local harmonic approximation to the DFT energy around optimised structures in the Open Catalyst 2020 ("OC20") dataset [5]. We call this an "Easy Potential" (EP) because it is so simple. Potentials trained on actual DFT data are termed ML-FF, for "machine-learned force fields". Optimising structures with this Easy Potential, and then refining them with an ML-FF, yielded even lower energy structures on average than optimisation with an ML-FF alone. In addition to explaining the success of learned force fields, this experiment indicates that Easy Potentials can accelerate and improve optimisation. This further increases the utility of machine learning in catalyst discovery.
Finally, when combining Easy Potentials with VASP relaxations in a hybrid approach, we find that the structures match or exceed the stability of VASP-optimized geometries in 88% of cases. The mean improvement is 0.46 eV, which is remarkably large. Further, these relaxations require half as many DFT steps as VASP optimization from scratch. This means that even practitioners not ready to fully embrace learned force fields can dramatically improve and accelerate their computations by initialising DFT-based relaxations through Easy Potentials.

Results
We are given a set of atoms specified by their positions ∈ ℝ 3 and atomic numbers ∈ {1, . . . , 118}. Denoting the set of atomic positions as = , . . . , | | , the forces exerted on the atoms are defined as ì = −∇ ( ) ∈ | |×3 , where ( ) is the potential energy for and · is the concatenation operator. The ground state of a PES is the set of atomic positions that minimises the energy, = min ( ( )).

Density Functional Theory
DFT is a method to compute the total energy of a system using a functional of the electron density [26,35,22]. DFT exists along a spectrum of methods to calculate the energy and forces of an atomic configuration that trade off computation time and accuracy. On one end there are traditional force fields, which include simple classical physics-inspired terms like the Lennard-Jones potential [24]. These potentials can be rapidly evaluated, but often lack transferability and/or accuracy. On the other end, there are quantum chemical methods such as CCSD(T) [1] and quantum Monte Carlo [7,33], which are highly accurate but too expensive for all but the simplest systems. Reducing computational cost while maintaining accuracy is an ongoing effort [18].
DFT is the most widely used method for computing energies for systems in condensed matter physics, materials science, and chemistry, with key papers having over 100 000 citations [35]. DFT's success comes from being sufficiently accurate for many applications, particularly in materials science and medicinal chemistry, but with an ( 3 ) run time and moderate prefactor that make it usable in practice. An important use case for DFT is structure optimisation, in which the atomic configuration with the minimum energy is found using a gradient-based optimisation algorithm. This process is also known as relaxation or finding the ground state of a material. The ground state energy of materials can be used to derive many thermodynamic properties, such as adsorption and formation energies [5]. For this reason, structure optimisation using DFT is a workhorse of computational materials science.
Structure optimisation can be sensitive to the choice of functional and the choice of optimisation algorithm [21]. DFT codes such as VASP [29,28,27] and Quantum Espresso [14] use conjugate gradients or quasi-Newton derived methods such as RMM-DIIS [44] and LBFGS [30] for optimisation. These methods are tightly integrated into the software packages.

Learned Interatomic Potentials
Despite its favorable cost-accuracy trade-off relative to other quantum chemical methods, DFT is still rather expensive. Applications such as ab initio molecular dynamics ("MD") [8], which require thousands of energy calculations for each system, are often prohibitively expensive. Further, its cubic scaling makes DFT very challenging for evaluating large systems of many thousands of atoms, though linear scaling algorithms are under constant development [37]. The goal of learned atomic potentials [43] is to replace DFT with a neural network that is as accurate as DFT, but that is much faster to evaluate.
Given a training set of set of atomic structures with associated DFT forces and energies = {( 1 , ( 1 ), ì 1 ), . . . , ( | | , ( | | ), ì | | ))} we train a learned interatomic potential by minimising the approximation error of the energies and forces across a sample of DFT calculations, ( ( ( ), ( )) + ( ì ( ), ì ( ))) where is the expectation. Many learned interatomic potentials compute the forces by differentiating the energy with respect to the atomic positions, while others learn the forces independently [23,17]. We also predict forces directly, since this is faster than differentiating the energy.
A common class of neural networks used for learned atomic potentials are Message Passing Neural Networks [3,15], which represent a system as a spatial graph. Atoms are represented as nodes, with edges between atoms constructed based on inter-atomic distances and using a cutoff radius. Atom nodes and inter-atomic edges are then featurized using learned embeddings of atom types [40] as well as varying representations of positional information. In our work we use the Graph Net Simulator (GNS) [38], a generic architecture suitable for molecular property prediction [17]. Other architectures such as SchNet [40], DimeNet [11,10], GemNet [12,13], and PaiNN [41] can also be used for the tasks in this work.
Our learned potential, hereafter referred to as ML-FF, is trained and evaluated on different subsets of the OC20 dataset [5]. OC20 is the largest dataset of ground state structures generated with the same DFT settings throughout, and is thus the most suitable for analysing ground state discovery. The S2EF ("structures to energy and forces") subset has 134 million point-wise DFT calculations of both energies and forces. These are used to train models that predict and ì for use in relaxation. Four validation datasets with approximately 25k structures each (initial and relaxed) are available to benchmark the model performance. The four validation sets have indistribution catalysts and adsorbates (ID), out of distribution catalysts (OOD_CAT), out of distribution adsorbates (OOD_ADS), and out of distribution catalysts and adsorbates (OOD_BOTH). Test sets are not publicly available and hence were not used in this work. We also note that the labels for the validation and test sets are in practice often wrong and do not correspond to the true minimum. Our main purpose in using the validation set labels is to determine when our model produces higher quality output than OC20 itself. We train our model following the training method described by Zaidi et al. [46] (full details in Methods). The model has a mean absolute error (MAE) of 0.029-0.038 eV/Å for forces and 0.18-0.25 eV for energies, and 0.5 cosine similarity for forces on validation datasets. This error is similar to recent results in the literature; for example, GemNet [12] in 2021 reported 0.6 cosine similarity and 0.29 eV energy MAE on the OC20 test sets. We relax the structures using the Adamw optimiser [45], a weighted adaptive momentum optimiser that we found to best handle noisy force predictions.

Easy Potentials
While structure optimisation is typically guided by forces, one can also find ground states by following the gradient of a surrogate potentialw ith the same global minimum as . Indeed, one can choose a surrogate that is smoother than , leading to faster optimisation and a higher likelihood of escaping local minima.
This point is illustrated in Fig. 2. We start with two PESs, one a complicated energy surface (top left), and the other a parabola (bottom left) that has been centered on the minimum of the first surface. Optimising the parabola is fast and easy, taking only five steps of optimisation. The complicated PES is hard to optimise, taking 23 steps in total. However, constructing the parabola requires already knowing the locations of the local minima, and so it cannot be used for structure optimisation. This is where machine learning can help: when trained to match the parabola in different but related systems, the Easy Potential generalises to minima in the system of interest, while maintaining the same advantageous optimisation qualities (bottom right).
The above example emphasises an important point: the only requirement for a learned potential being useful for material structure optimisation is that its forces are close to zero near the ground states of the target DFT functional. Indeed, our experiments show that Easy Potentials successfully identify low energy configurations outside of their training set even when they exhibit high force prediction error in high energy regions. Counter-intuitively, this also means that ML-FF with high force prediction errors in areas outside a minimum may offer optimisation benefits, as noisy gradients [34] can be a successful strategy to avoid getting stuck in local minima.
For each ground stateˆin the training set, the Easy Potential approximates the energy of nearby states as ( ) = (ˆ) + 1 2 || −ˆ|| 2 2 , where the sum is over each atom in with position , and || · || 2 is the L2 norm. The specific construction of the neural network predicting ( ) in this manner is inspired by Godwin et al's work into regularisation to prevent over-smoothing [17]. The forces are then ì ( ) =ˆ− . Nearby states were sampled from intermediate points on the DFT relaxation trajectory, then augmented with additive Gaussian noise. Easy Potential gives rise to large energy and force errors when relaxed naively as a result of small errors in bond distances. As a result, we adopt a two stage relaxation proce-dure. First, we relax Easy Potential, and then we resolve bond errors with a second stage of relaxation using ML-FF. Adamw is again used as the optimiser. Full details of the relaxation and training procedures can be found in the Methods section.

ML-FF Relaxations
Our first experiments demonstrate the ability of ML-FF relaxations to find minimum energy wells of the target DFT energy. Table 1 compares the energies of relaxed ML-FF structures ("predicted structures") with structures contained in the OC20 datasets. Since we have limited access to VASP, we obtain these results using DFT in Quantum Espresso for each set of structures. We use DFT settings as close as possible to those that generated the OC data. (As discussed below, VASP was also used on a small subset of the data to confirm that our results were not artifacts of Quantum Espresso.) The Quantum Espresso settings yield 93%-96% convergence rates on the validation datasets. Full details of the parameters used, and a comparison between VASP and Quantum Espresso can be found in the methods. We also calculate the root-mean-square deviation of atomic positions (RMSD). Our analysis indicates that structures with lower DFT energies than OC20 correlate with higher RMSDs. This suggests that ML-FF escapes local minima to obtain qualitatively different structures with lower energies, an example of which is shown in Fig. 1.

Easy Potential Relaxations
Does training on synthetic quadratic approximations yield similar results, which tests our hypothesis of primarily needing to train on similar minima? Our experiments for Easy Potential, presented in To confirm that these results were not due to different DFT settings in Quantum Espresso and VASP, we also evaluate a sample of 1000 OOD_BOTH structures using VASP, with the exact same settings used in OC20. To check our settings, we first confirmed that VASP results on the ground truth structures matched the energies provided by the OC20 dataset. We observe a distribution shift in which a higher number of structures are less stable, but still find an improvement of −0.13 eV on average.
The majority of structures predicted by Easy Potential are still equivalent or lower in energy than DFT relaxations. We also carry out a study where we allow for a larger iteration budget in the refinement stage (950 versus 100 iterations in the second stage), which further improves results against VASP. Finally, we use again use the OOD_BOTH sample relaxed via EP and continue relaxations using VASP in a hybrid approach, which results in -0.46 eV average improvements (for a median of 71 additional DFT relaxation steps) against the original VASP OC20 relaxations (median 147 DFT steps).   Table  2. Two expected effects are observed. Structures centered around the DFT margin of error (0.1 eV) have very small average RMSD with few outliers. Mean RMSD then increases for both more stable and less stable structures.

ML-FF Analysis
Here we perform ablation studies to provide further insight into the performance and trade-offs of ML-FF relaxations. In Table 3, we show that increasing the number of optimisation steps in our relaxation yields lower-energy geometries. We observe substantial improvements going from 100 to 1000 iterations, but marginal subsequent gains. Note that these results are obtained using only ML-FF, i.e. not improving results by first using the harmonic approximation for initialisation.
In Table 4, we compare runtimes with CPUs and with GPUs. Our main results were obtained using a maximum number of 1000 iterations, corresponding to a 10 minute median relaxation time on CPU and 46.5 seconds on a V100 GPU. Note that the model and relaxation procedure are deterministic, but we observe small runtimescheduling related non-determinism on GPU platforms. Classical DFT solvers are typically run on CPUs, and CPU relaxation of Easy Potentials is fully deterministic. We expect future work to further refine execution times by additionally optimising for a specific target platform.
In Fig. 4, we show an example ML-FF relaxation trajectory from OOD_BOTH, illustrating how using an adaptive momentum optimiser combined with noisy force predictions allows escaping from local minima several times before converging.  DFT scales cubically in system size, whereas GNS scales linearly for OC20 catalytic systems since the maximum number of interactions is capped by the cutoff radius. In Fig. 5, we measure inference cost to predict energies and forces for structures from the OOD_BOTH validation dataset. Executed on a single V100 GPU, we observe linear scaling against system size. The largest structures in OOD BOTH have less than 250 atoms, and we further create two synthetic systems with 512 and 1024 atoms, and with the theoretical maximum number of edges that our model admits (see Methods). We see that these large systems follow the same linear scaling trend.

Bond Length Analysis
We further validate ML-FF's ability to predict minima by evaluating bond lengths. To do so we predict the potential energy of atom pairs at various bond distances. The bond length is the distance that minimizes the energy. We compare ML-FF predictions with the ground truth computed with Quantum Espresso. The results are illustrated in Fig. 6: the predicted distance is within 10% of the target for all the pairs we sampled. Furthermore, our model and the ground truth tend to show similar behavior near the energy minimum, in many cases displaying non-parabolic behaviour. We observe more difficulties on some pairs (e.g. Cl-Cl, P-P, Se-Se) which is likely due to a limited number of these bonds in the training data.

Discussion
The key finding of our experiments is that learned interatomic potentials are surprisingly accurate at identifying low energy wells of target DFT PESs, even if force predictions are noisy or the model is trained on a simple harmonic potential. The latter are similar in spirit to gradient fields developed for ML conformer generation [42,32], which are trained on minimum-energy structures rather than actual forces. We speculate that DFT packages like VASP struggle to escape local minima due to using classical optimisers and/or nonnoisy ground truth forces. A detailed investigation of the importance of these two factors is left for future work. The main impact of our findings is that learned force fields are demonstrably useful for practical calculations today.
Using ML-FFs and Easy Potentials opens up the possibility of relaxing large systems of tens of thousands of atoms due to their fast execution and linear scaling. We imagine Easy Potentials being used in collaboration with classical methods, for example to obtain an initial fast result within a low energy well, which can then be further refined. As shown when combining EP with VASP in table 2, this combines the guarantees of DFT while reducing the number of steps required while also improving results.
Lastly, our results imply that focusing on force and structure distance errors on Open Catalyst and similar datasets may not accurately capture the progress in learned force fields. Instead, predicted structures should be evaluated with singlepoint DFT calculations. Practitioners investigating other approaches to learned potentials may be similarly surprised to find that, by incorporating harmonic potentials or novel optimisation strategies for their relaxations, their force fields may already be sufficiently accurate for ground state catalyst discovery.

Model details
Our model is a variant of the Graph Net Simulator [38] with the modifications described by Godwin et al. [17] for molecular property prediction. Below, we provide a high level summary of the model and configurations. Please refer to the original papers for comprehensive explanations.
GNS consists of an encoder that processes input data into a latent space, a processor that executes message passing and updates the latent representation, and decoders that interpret the resulting node embeddings to per-node and graph-level predictions. In the encoder, we use the set of atoms to construct a directed graph ( , ) with featurized sets of vertices = { 1 , 2 , . . . , | | } and edges = { , } , where an edge between and is added if their distance is less than a connectivity radius of 6 Å, for a maximum of 20 edges per atom determined via a k-d tree for nearest neighbours. Each edge , is represented with the displacement vector − and the Euclidean norm as features. Distances are then further featurised by computing Gaussian radial basis functions (mean = 0, standard deviation = 0.5 Å). For each atom ( , ) a vertex is represented through a learned embedding of the atom type . For OC20, we also utilise another input feature, a "tag" for each atom that indicates whether it is part of the slab, surface, or catalyst, by creating a separate learned embedding. Embedding sizes were set to 512, and all multi-layer perceptrons (MLPs) had size 1024 and used shifted softplus activations.
Following the input graph construction, the processor iteratively updates node and edge features by applying Interaction Networks [2]. Each application represents one message passing step with a different set of weights. We repeat this for a configurable number of message passing steps, and then apply the entire processor multiple times sequentially while re-using weights. During training we average the output from each application. At inference time we only use the output from the final step. We used 10 repeated applications of 5 message passing layers.
Finally, after message passing, forces are decoded by applying an MLP to each node feature. Energy predictions are decoded by applying an MLP to each node feature, aggregating the nodelevel results using a sum aggregation function, and finally applying another MLP to predict the total potential potential energy of the system. Note that energy predictions are not strictly required for relaxations if not differentiating them. We primarily use them for assessing the model's internal evaluation consistency during relaxation (i.e. for analysis without DFT).
GNS is implemented in Jax [4] using Haiku [19] for layer implementations, Jraph [16] for GNN components, Optax [20] for optimisation, and Automap for distributed training compilation [39]. Harmonic potentials are trained for 250 000 gradient steps on the OC20 ISRS dataset using Gaussian noise around positions which were interpolated between initial and finial relaxed structures ( = 0, = 0.5). Force predictions on S2EF for ML-FF are then trained for up to 2 million gradient steps with early stopping for model selection. The two step procedure follows the methodology described by Zaidi et al. [46].
All training is performed on a cluster of 16 TPU v4 devices with 32 GB of RAM. Graphs are dynam-ically batched with up to 10 graphs, 1024 vertices and 12800 edges in a single batch (i.e. the batch is padded to these dimensions once one of them is reached, since input structures have varying dimensions). Parameters are updated using a mean squared error loss and the Adam optimiser [25] ( 1 = 0.9, 2 = 0.95) with a warm up (10 5 warm up steps, 10 −5 warmup start learning rate and 10 −4 warmup max learning rate), and a cosine decay schedule (5 · 10 6 cosine cycle length with 10 −6 min learning rate). Finally, parameters are smoothed using an exponentially moving average decay with a decay parameter of 0.9999.

Relaxation settings
Relaxations are implemented using the Adamw [31] optimiser from the Optax library [20] (decay parameters 1 = 0.9, 2 = 0.99). A learning rate of 0.1 is used to update structures. Relaxations are run with two stopping conditions (stopped when either is reached): • A maximum number of relaxation steps (ML-FF: 1000, Easy Potential/ML-FF two-stage procedure: 50 Easy Potential, 100 ML-FF). • The Frobenius norm of the predicted forces per-atom below a threshold hyperparameter for all atoms (ML-FF: 0.005, Easy Potential/ML-FF two-stage: 0.03).
As gradients for Adamw, we use the negative direct per-atom force predictions rather than explicitly computing the derivative of the predicted energy with respect to positions. As discussed in the Model Details section, atom graphs are constructed based on radial distance cutoffs. The relaxed atomic positions change inter-atomic distances. We hence recompute edges after updating positions by re-executing a nearest neighbour lookup within the distance cutoff.
Since relaxations are implemented in JAX, we in practice use jit compilation to XLA which requires static shapes. Graphs change shape throughout relaxations, so we compile ahead-oftime to an initial power-of-two for edge counts and pad graphs to this size, and only recompile dynamically as needed to the next power of two to minimize the total number of recompilations throughout a relaxation. For large scale inference, pre-compiling a range of sizes and matching each molecule to the next power-of-two could avoid any runtime compilation and optimize throughput. Relaxations are measured on both CPUs and NVIDIA V100 GPUs with 16 GB RAM.

Dataset preparation
We used the OC20 dataset as provided in https://github.com/Open-Catalyst-Project/ocp. We conducted some data cleaning, removing 580 structures where the surface crossed the unit cell -axis. This was observed after training, and future models could be adjusted to also handle a periodic -axis if necessary.

DFT settings
We faithfully reproduce the settings described in the Open Catalyst paper [5], which utilises VASP, while we use the open source Quantum Espresso package (apart from a small sample of VASP validations). We use the RPBE functional and computed the -points for each structure using the same formula as [5]: Here = 40 and a 1 , a 2 are the dimensions in the first and second axis of the unit cell of the structure. Since a significant part of our structures are metallic we used a Marzari-Vanderbilt smearing with a Gaussian spreading of 0.2. We do not observe a significant energy difference for the insulators of the dataset with or without smearing. Therefore, for practical reasons, we used smearing for all DFT computations. Moreover, due to computational limitations we do not use the precise pseudo-potentials in [36], but rather choose the most efficient ones. Finally, to improve DFT convergence rates, we use a local densitydependent Thomas-Fermi screening for the mixing mode, with beta equal to 0.7. Each DFT computation was launched on a Google Cloud Platform (GCP) machine with 62 CPUs and 248 GB of RAM.
We use VASP for a small sample of structures for validation. Our VASP settings exactly reproduce the OC20 results when using the configurations specified in the paper [5], together with the corresponding preprocessing scripts that were released with the datasets. The jobflow software was used for automating VASP calculations [9].

Data availability
The predicted structures for ML-FF and Easy Potential will be released upon acceptance of this manuscript.

Code availability
Model code, trained checkpoints, as well as the relaxation implementation will be released upon acceptance of this manuscript.

Author contributions
Jonathan Godwin and Michael Schaarschmidt carried out the original research into learned potentials, including model design and implementation. Morgane Rivière designed and implemented the DFT and relaxation evaluation infrastructure. Alex M. Ganose carried out structure validations and advised on experiment design for evaluating relaxed structures. James S. Spencer implemented the first DFT validation prototype and advised on the configuration of Quantum Espresso, as well as contributed to the interpretation of results, writing and analysis. Alexander L. Gaunt and James Kirkpatrick advised on computational chemistry and relaxation evaluation. Simon Axelrod advised on computational chemistry and contributed to writing. Peter W. Battaglia advised on initial model design for learned potentials as well as on the overall project structure and delivery.
DFT experiments on Google Cloud, as well as to Alison Reid for guiding the open sourcing process. A.M.G. was supported by EPSRC Fellowship EP/T033231/1.