MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields
s
AffiliationAddress
Abstract
Creating fast and accurate force fields is a long-standing challenge in computationalchemistry and materials science.
Recently, several equivariant message passingneural networks (MPNNs) have been shown to outperform models built usingother approaches in terms of accuracy.
However, most MPNNs suffer from highcomputational cost and poor scalability.
We propose that these limitations arisebecause MPNNs only pass two-body messages leading to a direct relationshipbetween the number of layers and the expressivity of the network.
In this work,we introduce MACE, a new equivariant MPNN model that uses higher body ordermessages.
In particular, we show that using four-body messages reduces therequired number of message passing iterations to just two , resulting in a fast andhighly parallelizable model, reaching or exceeding state-of-the-art accuracy on therMD17, 3BPA, and AcAc benchmark tasks.
We also demonstrate that using higherorder messages leads to an improved steepness of the learning curves.
Introduction
The earliest approaches for creating force fields (interatomic potentials) using machine learning techniques were using local atom-centered symmetric descriptors and feed-forward neural networks [6],
Gaussian Process regression[2] or linear regression [44, 47].
The first attempts to use graph neural networks to model the potential energy of atomistic systems had only limited success.
The DTNN [42],
SchNet [41]
HIP-NN [35]
PhysNet [48] or DimeNet [20, 29] approaches could only come closeto but not improve upon the atomic descriptor-based methods in terms of computational efficiencyand accuracy on public benchmarks.
Furthermore, most MPNN interatomic potentials use 2-bodyinvariant messages, making them non-universal approximators [38].
The MACE architecture presented here allows for the efficient computation of equivariant messageswith high body order.
As a result of the increased body order of the messages, only two messagepassing iterations are necessary to achieve high accuracy - unlike the typical five or six iterationsof MPNNs, making it scalable and parallelizable.
Finally, our implementation has remarkablecomputational efficiency, reaching state-of-the-art results on the 3BPA benchmark after 30 mins oftraining on NVIDIA A100 GPUs.
We summarise our main contributions as follows
• We introduce MACE, a novel architecture combining equivariant message passing withefficient many-body messages.
The MACE models achieve state-of-the-art performance
on challenging benchmark tests.
They also display greater generalization capabilities overother approaches on extrapolation benchmarks.
• We demonstrate that many-body messages change the power of the empirical power-law ofthe learning curves.
Furthermore, we show experimentally that the addition of equivariantmessages only shifts the learning curves but does not change the power law when higherorder messages are used.•
We show that MACE does not only outperform previous approaches in terms of accuracybut also does so while being significantly faster to train and evaluate than the previous mostaccurate models.
Background2.1 MPNN Interatomic Potentials
MPNNs [22, 9] are a type of graph neural network (GNN, [40, 4, 27, 51]) that parametrises amapping from a labeled graph to a target space, either a graph or a vector space.
When applied toparameterise properties of atomistic structures (materials or molecules), the graph is embedded in3-dimensional (3D)
Euclidean space, where each node represents an atom, and edges connect nodesif the corresponding atoms are within a given distance of each other.
We represent the state of eachnode i in layer t of the MPNN by a tuple
[Equation]
where r i ∈ R 3 is the position of atom i , z i the chemical element, and h ( t ) i are its learnable features.
A forward pass of the network consists of multiple message construction , update , and readout steps.
During message construction, a message m ( t ) i is created for each node by pooling over its neighbors:
[Equation]
(2) where M t is a learnable message function and (cid:76) j ∈N ( i ) is a learnable, permutation invariant poolingoperation over the neighbors of atom i (e.g., a sum).
In the update step, the message m ( t ) i istransformed into new features
[Equation]
3) where U t is a learnable update function.
After T message construction and update steps, the learnablereadout functions R t map the node states σ ( t ) i to the target, in this case the site energy of atom i ,
E
[Equation]
2.2 Equivariant Graph Neural Networks
In equivariant GNNs, internal features h ( t ) i transform in a specified way under some group action [1,12, 32, 46, 49].
When modelling the potential energy of an atomic structure, the group of interest is
O (3) , specifying rotations and reflections of the particles.
1 We call a GNN O (3) equivariant if it hasinternal features that transform under the rotation Q ∈ O (3) as
[Equation]
where Q · ( r 1 , ..., r N ) denotes the action of the rotation on the set of atomic positions and D ( Q ) is amatrix representing the rotation Q , acting on message h ( t ) i .
In general, elements of the feature vectorcan be labeled according to the irreducible representation they transform with.
We will write h ( t ) i,kLMto indicate a collection of features on atom i , indexed by k , that transform according to
[Equation]
i
[Equation]
where D L ( Q ) ∈ R (2 L +1)
× (2 L +1) is a Wigner D-matrix of order L .
A feature labelled with L = 0describes an invariant scalar.
Features labeled with L > 0 , describe equivariant features, formallycorresponding to equivariant vectors, matrices or higher order tensors.
The features of invariantmodels, such as SchNet[41] and DimeNet[29], transform according to D ( Q ) = 1 , the identity matrix.
Models such as NequIP [5], equivariant transformer [45], PaiNN [43], or SEGNNs [8], in addition toinvariant scalars, employ equivariant internal features that transform like vectors or tensors.
Related WorkACE - Higher Order Local Descriptors
In the last few years, there have been two significantbreakthroughs in machine learning force fields.
First, the Atomic Cluster Expansion (ACE) [16]provided a systematic framework for constructing high body order complete polynomial basisfunctions (features) at a constant cost per basis function, independent of body order [17].
It has alsobeen shown that ACE includes many previously developed atomic environment representations asspecial cases, including Atom Centred Symmetry Functions [6], the Smooth Overlap of Atomic
Positions (SOAP) descriptor [2], Moment Tensor Potential basis functions [44], and the hypersphericalbispectrum descriptor [2] used by the SNAP model [47].
These local models are limited by theircutoff distance and their relatively rigid architecture compared to the overparametrised MPNNs,leading to somewhat lower accuracy, in particular, for molecular force fields.
Equivariant MPNNs
The second breakthrough was using equivariant internal features in MPNNs.
These equivariant MPNNs, such as Cormorant [1], Tensor Field Networks [46], EGNN [39],
PaiNN [43], Equivariant Transformers [45], SEGNN [8], NewtonNet [23], and NequIP [5] were ableto achieve higher performance than previous local descriptor-based models.
However, they sufferfrom two significant limitations: first, the most accurate models used L = 3 spherical tensors asmessages and 4 to 6 message passing iterations [5], which resulted in a relatively high computationalcost.
Second, using this many iterations significantly increased the receptive field of the network,making them difficult to parallelise across multiple GPUs [36].
Higher Order Message Passing
Most MPNNs use a message passing scheme based on two-bodymessages, meaning they simultaneously depend on the states of two atoms.
It has been recognisedthat it can be beneficial to include angles into the features, effectively creating 3-body invariantmessages [29].
This idea has also been exploited in other invariant MPNNs, in particular, by
SphereNet [34] and GemNet [30].
Even though these models improved the accuracy compared tothe 2-body message passing, they were limited by the computational cost associated with explicitlysumming over all triplets or quadruplets to compute the higher order features.
Multi-ACE Framework
Recently, multi-ACE has been proposed as a unifying framework of
E (3) -equivariant atom-centered interatomic potentials, extending the ACE framework to includemethods built on equivariant MPNNs [3].
A similar unifying theories were also put forward by[37] and [7].
The idea is to parameterise the message m ( t ) i in terms of invariant or equivariant
ACE models.
This framework sets out a design space in which each model can be characterisedin terms of: (1) the number of layers, (2) the body order of the messages, (3) the equivariance (orinvariance) of the messages, and (4) the number of features in each layer.
The framework highlightsthe relationship between the overall body order of the models and message passing, also previouslydiscussed in Kondor [31].
Most previously published models achieved high accuracy by either using4 to 6 layers [5, 43] or increasing the local body order with a single layer [33, 36].
With our model,we fall in between these two extremes by combining high body order with message passing.
The MACE Architecture
Our MACE model follows the general framework of MPNNs outlined in Section 2.
Our keyinnovation is a new message construction mechanism.
We expand the messages m ( t ) i in a hierarchical body order expansion,
[Equation]
where the u functions are learnable, the sums run over the neighbors of i , and ν is a hyper-parametercorresponding to the maximum correlation order, the body order minus 1, of the message functionwith respect to the states.
Even though we refer to the message as ( ν + 1) -body with respect to thestates, the overall body order with respect to the positions can be larger depending on the body order ofthe states themselves.
Crucially, by writing (cid:80) j 1 ,...,j ν , which includes self-interaction (e.g., j 1 = j 2 ),we will later obtain a tensor product structure with a computationally efficient parameterisation,that allows us to circumvent the seemingly exponential scaling of the computational cost with thecorrelation order ν .
This contrasts with previous models, such as DimeNet [28, 29], that compute3-body features via the more standard many-body expansion (cid:80) j 1 < ··· <j ν .
Below, we describe the
MACE architecture in detail.
Message Construction
At each iteration, we first embed the edges using a learnable radial basis
R ( t ) kl 1 l 2 l 3 , a set of spherical harmonics
Y m 1 l 1 , and a learnable embedding of the previous node featuresh ( t ) j,l 2 m 2 using weights W ( t ) kl 2 .
The A ( t ) i -features are obtained by pooling over the neighbours N ( i ) toobtain permutation invariant 2-body features whilst, crucially, retaining full directional information,and thus, full information about the atomic environment:
[Equation]
where C l 3 m 3 l 1 m 1 ,l 2 m 2 are the standard Clebsch-Gordan coefficients ensuring that A ( t ) i,kl 3 m 3 maintain thecorrect equivariance, r ji is the (scalar) interatomic distance, and ˆ r ji is the corresponding unit vector.
R ( t ) kl 1 l 2 l 3 is obtained by feeding a set of radial features that embed the radial distance r ji using Besselfunctions multiplied by a smooth polynomial cutoff (cf.
Ref.
[29]) to a multi-layer perceptron (MLP).
See Section A.4 for details.
In the first layer, the node features h ( t ) j correspond to the (invariant)chemical element z j .
Therefore, (8) can be further simplified:
[Equation]
[Equation]
This simplified operation is much cheaper, making the computational cost of the first layer low.
The key operation of MACE is the efficient construction of higher order features from the A ( t ) ifeatures.
This is achieved by first forming tensor products of the features, and then symmetrising:
[Equation]
where the coupling coefficients C LMη, lm corresponding to the generalised Clebsch-Gordan coefficients(details in A.3) ensuring that B ( t ) i,ηkLM are L -equivariant, the weights w ( t )
k ˜ kl ξ are mixing the channels( k ) of A ( t ) i , and ν is a given correlation order.
C LMη, lm is very sparse and can be pre-computed such that(10) can be evaluated efficiently (see Appendix A.3.3).
The additional index η simply enumerates allpossible couplings of l 1 , . . .
, l ν features that yield the selected equivariance specified by the L index.
The B ( t ) i -features are constructed up to some maximum ν .
This variable in (10) is the order of thetensor product, and hence, can be identified as the order of the many-body expansion terms in (7).
The computationally expensive multi-dimensional sums over all triplets, quadruplets, etc., are thuscircumvented and absorbed into (9) and (8).
The message m ( t ) i can now be written as a linear expansion
[Equation]
where W ( t ) z i kL,η is a learnable weight matrix that depends on the chemical element
z i of the receivingatom and message symmetry L .
Thus, we implicitly construct each term u in (7) by a linearcombination of B ( t ) i,ηkLM features of the corresponding body order.
Under mild conditions on the two-body bases A ( t ) i , the higher order features B ( t ) i,ηkLM can beinterpreted as a complete basis of many-body interactions [17], which can be computed at a costcomparable to pairwise interactions.
Because of this, the expansion (11) is systematic .
It can inprinciple be converged to represent any smooth ( ν + 1) -body equivariant mapping in the limit ofinfinitely many features (proof in [17]).
Update
In MACE, the update is a linear function of the message and the residual connection [25]:
[Equation]
Readout
In the readout phase, the invariant part of the node features is mapped to a hierarchicaldecomposition of site energies via readout functions:
[Equation]
where
[Equation]
The readout only depends on the invariant features h ( t ) i,k 00 to ensure that the site energy contributions
E ( t ) i are invariant as well.
To maintain body ordering, we use linear readout functions for all layersexcept the last, where we use a one-layer MLP.
Results5.1 Effect of Higher Order MessagesNumber of Layers
In this section, we investigate the effect of using higher order messages.
Many MPNN architectures [41, 48] exclusively pass two-body invariant messages resulting in anincomplete representation of the local environment [38].
Equivariant message-passing schemes [5,43 , 8] lift the degeneracy of most structures by containing directional information in the messages.
MPNNs that only employ two-body messages at each layer can increase the body order either bystacking layers [31] which simultaneously increases the model’s receptive field or by using non-linearactivation functions, generate only a subset of all possible higher order features.
By constructinghigher order messages using the MACE architecture, we disentangle the increase in body order fromthe increase of the receptive field.
In Figure 1, we show the accuracy of MACE, NequIP, and BOTNet [3] on the 3BPA benchmark [33]as a function of the number of message passing layers.
Approaches employing 2-body messagepassing require up to five iterations for their accuracy to converge.
By constructing many bodymessages, the number of required layers to converge in accuracy reduces to just two.
In all subsequentexperiments, we use two-layer MACE models.
Furthermore, we compare BOTNet, which does not use any non-linearities in the update step to
NequIP, which does.
Otherwise, the two models are very similar.
We observe that the increase in bodyorder through non-linearities within the update provides only marginal improvement, highlightingthe difference between an increase in body order through non-linearities (NequIP) and higher ordersymmetric messages (MACE).
Consequently, higher order message passing allows one to reduce thenumber of layers, thereby increasing speed and ease of parallelization over multiple GPUs.
We notethat MACE does not improve after two layers as the diameter of the 3BPA molecule is about 9 Å andradial cutoff in each layer is 5 Å.
[Figure]
Learning Curves
We study how higher order message passing affects the learning curves.
Arecent study of the NequIP model [5] showed that the inclusion of equivariant features results inenhanced data efficiency, increasing the slope of the log-log plot of predictive error as a function ofthe dataset size.
They showed that adding equivariance not only shifts the learning curves, but alsochanges the powers in the empirical power law of the learning curves, which is usually constant for agiven dataset [26].
On the left panel of Figure 2, we replicate the experiments of [5] by training a series of invariant
MACE models with increasing correlation order ν on the aspirin molecule from the rMD17 dataset.
We observe that adding higher order messages changes the steepness of the learning curves, evenwithout equivariant messages.
The model with correlation order ν = 1 corresponds to a two-layer2-body invariant model, similar to SchNet.
This model is the least accurate due to the incompletenature of 2-body invariant representations of the local environment [38].
The invariant messages withν = 2 are akin to those in DimeNet, which explicitly puts angular information into the messages.
Wesee that including higher order information significantly improves the model’s accuracy.
Finally, bygoing beyond any current message passing potential by setting ν = 3 , we achieve similar performanceto a highly-accurate 2-body, equivariant MPNN while only using higher order invariant messages.
On the middle panel of Figure 2, we keep the correlation order fixed at ν = 3 and gradually increasethe symmetry order L of the messages.
While the slope remains nearly unchanged, the curves areshifted.
In the right panel of Figure 2, we keep the correlation order fixed at ν = 1 and graduallyincrease the symmetry order L of the messages.
We see only a marginal slope change when addingequivariant features, which could be attributable to the relatively low expressiveness of a two-layer
MACE restricted to correlation order ν = 1 .
These results suggest two routes to improve invariant2-body MPNN models: creating higher correlation order messages or incorporating equivariantmessages.
By exploiting both of these options, the MACE model achieves state-of-the-art accuracy.
5.2 Scaling and Computational CostChemical Elements
A significant limitation of existing atomic environment representations is thattheir size grows with the number of chemical elements S and correlation order ν as S ν .
Data-drivencompression schemes have been proposed [50] to solve this issue, and MPNNs incorporate similarembeddings of the chemical elements into a fixed-size vector space.
MACE uses a continuous speciesembedding and when constructing the higher order features in (10), it does not include the speciesdimension k in the tensor product resulting in O (1) scaling of the model with the number of chemicalelements S . [Figure]
Receptive Field
A severe limitation of many previously published MPNNs was their largereceptive field, making it difficult to parallelize the evaluation across multiple GPUs.
In traditional
MPNNs, the total receptive field of each node, which grows with each message passing iteration,can be up to 30 Å.
This scaling results in the number of neighbours being in the thousands in acondensed phase simulation, preventing any efficient parallelization [36].
By decoupling the increasein correlation order of the messages from the number of message passing iterations, MACE onlyrequires two layers resulting in a much smaller receptive field.
With a local radial cutoff of 4 to 5 Å,the overall receptive field remains small, making the model more parallelisable.
Computational Cost
The computational bottleneck of equivariant MPNNs is the equivarianttensor product (8).
This tensor product is evaluated on edges.
In MACE, we only evaluate thisexpensive tensor product once, within the second layer, and build up correlations through the tensorproduct of (10).
Importantly, this operation is carried out on nodes.
Typically the number of nodes isorders of magnitudes smaller than the number of edges resulting in a computational advantage.
Inaddition, we developed a loop tensor contraction algorithm for the efficient implementation of (10)and (11) detailed in Section A.3.
We report evaluation times for BOTNet, NequIP, and multiple versions of MACE in Table 2.
Weobserve that the invariant MACE ( L = 0 ) is close to 10 times faster than BOTNet and NequIPwhile achieving similar accuracy at high temperatures.
MACE with L = 1 and L = 2 is 5 and 4times faster than BOTNet and NequIP, respectively, while outperforming them at every temperature.
We acknowledge that accurate speed comparisons between codes are hard to obtain, and furtherinvestigations need to be carried out.
It is also essential to consider training times.
In order to do a faircomparison, all the timings were realised using the mace code that implements all the above models.
Models that are significantly faster to train are better suited for applications of active learning, whichis typically how databases for materials science applications are built [13–15].
The MACE modelreported in Table 2 takes approximately 30 mins to reach the accuracy of a converged BOTNet model,taking more than a day to be trained on the 3BPA dataset using NVIDIA A100 GPUs.
5.3 Benchmark Results
5.3.1 rMD17: Molecular Dynamics Trajectory
The revised MD17 (rMD17) dataset contains train test splits randomly selected from a long moleculardynamics trajectory of ten small organic molecules [11].
For each molecule, the splits consist of 1000training and test configurations.
Table 1 shows that MACE achieves excellent accuracy, improvingthe state of the art for some molecules, particularly those with the highest errors.
As several methodsachieve similar accuracy on the standard task of predicting energies and forces based on the whole
[Table]
2.7 Ethanol
E 0.
6.7 F 2.
32.6 Malonaldehyde E 0.
10.0 F 4.
43.3 Naphthalene E 0.
2.1 F 1.
9.2 Paracetamol E 1.
9.7 F 4.
31.5 Salicylicacid E 0.
6.5 F 3.
28.4 Toluene
E 0.
3.1 F 1.
12.1 Uracil
E 0.
4.4 F 2.
training set, we also trained MACE and NequIP, another accurate model, on just 50 configurations toincrease the difficulty of the benchmark.
In this case, we found that MACE outperformed NequIP formost molecules.
5.3.2 3BPA: Extrapolation to Out-of-domain Data
The 3BPA dataset introduced in [33] tests a model’s extrapolation capabilities.
Its training set contains500 geometries sampled from 300 K molecular dynamics simulation of the large and flexible drug-likemolecule 3-(benzyloxy)pyridin-2-amine.
The three test sets contain geometries sampled at 300 K,600 K, and 1200 K to assess in- and out-of-domain accuracy.
A fourth test set consists of optimizedgeometries, where two of the molecule’s dihedral angles are fixed, and a third is varied between 0and 360 degrees resulting in so-called dihedral slices through regions of the PES far away from thetraining data.
The root-mean-squared errors (RMSE) on energies and forces for several models are shown in Table 2.
It can be seen that MACE outperforms the other models on all tasks.
In particular, when extrapolatingto 1200 K data, MACE with L = 2 outperforms NequIP and Allegro models by about 30% .
Further,
MACE with L = 2 outperforms the next best model, BOTNet, by 40% on energies for the dihedralslices.
Finally, the MACE model with invariant messages ( L = 0 ) often nearly matches or exceedsthe performance of competitive equivariant models.
MACE shows excellent results while also featuring low computational cost compared to many othermodels.
The L = 0 model, which approaches previous models in terms of accuracy, outpaces them order to implemented the NequIP and BOTNet architectures in the we for which is published together with this
For the precise specification of our NequIP implementa tion see the Appendix A PyTorch timings were realised on an NVIDIA
[Table]
by nearly a factor of 10, whereas the L = 2 model achieves state-of-the-art accuracy and is aroundfour times faster than other equivariant MPNN models.
In the table, we characterise the evaluationspeed of the models by reporting the “latency” which is defined as the time it takes to compute forceson a structure, which is typically independent of the number of atoms until GPU threads are filled(typically 10,000 atoms for these models on an Nvidia A100 80GB GPU).
In Figure 3, we compare the BOTNet, NequIP, and MACE ( L = 2 ) by inspecting their energy profilefor three dihedral slices.
Overall, it can be seen that all models produce smooth energy profiles andthat, in general, MACE comes closest to the ground truth.
The fact that MACE outperforms the
[Figure]
279 other methods in the middle panel, which contains geometries furthest from the training dataset [3],suggests superior extrapolation capabilities.
5.3.3 AcAc: Flexibility and Reactivity
A similar benchmark dataset assessing a model’s extrapolation capabilities to higher temperatures,bond breaking, and bond torsions of the acetylacetone molecule was proposed in [3].
In Table 3,we show that MACE achieves state-of-the-art results on this dataset as well.
For details, see Appendix 5.3.3.
[Table]
286Discussions
With MACE, we extend traditional (equivariant) MPNNs from 2-body to many-body message passingin a computationally efficient manner.
Our experiments show that the approach reduces the requirednumber of message passing, leading to efficient and parallelizable models.
Furthermore, we havedemonstrated the high accuracy and good extrapolation capabilities of MACE, reaching state-of-theart accuracy on the rMD17, 3BPA, and AcAc benchmarks.
Future development should concentrateon testing MACE on larger systems, including condensed phases and solids.
Reproducibility Statements
We have included error bars via different seeds and various ablation studies wherever necessary andappropriate.
We have stated all hyper-parameters and data description in the Appendix A.4.
Codewill be made available online.
Ethical Statements
The societal impact of MACE is challenging to predict.
However, better force fields have a positiveimpact on society by speeding up drug discovery and through helping to understand, control, anddesign new materials.
However, machine learning force fields rely on generating ab initio trainingdata leading to heavy computation and large energy consumption.
Machine learned force fields doalleviate the costs of doing molecular modelling significantly when compared with using solely abinitio methods.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes]
(c) Did you discuss any potential negative societal impacts of your work?
[Yes]
(d) Have you read the ethics review guidelines and ensured that your paper conforms tothem?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[N/A]
(b) Did you include complete proofs of all theoretical results?
[N/A]
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as a URL)?
[Yes]
(b) Did you specify all the training details (e
g., data splits, hyper-parameters, how theywere chosen)?
[Yes]
(c) Did you report error bars (e
g., with respect to the random seed after running experiments multiple times)?
[Yes] (d) Did you include the total amount of compute and the type of resources used (e
g., typeof GPUs, internal cluster, or cloud provider)?
[Yes]
If you are using existing assets (e
g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes]
(b) Did you mention the license of the assets?
[N/A]
(c) Did you include any new assets either in the supplemental material or as a URL?
[N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains
personally identifiable information or offensive content?
N/ A
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots , if
applicable
[N/A
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A
(c) Did you include the estimated hourly wage paid to participants and the total amount
spent on participant compensation?
[N/A