Source-linked AI summary
MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields
Ilyes Batatia, Dávid Péter Kovács, Gregor N. C. Simm, Christoph Ortner, Gábor Csányi
TL;DR
Force-field models must balance accuracy with computational cost and scalability, while two-body MPNN messages limit expressivity. MACE addresses this gap with efficient equivariant higher-body-order messages, reducing the required iterations to two and achieving strong accuracy across benchmarks. The approach also improves learning-curve steepness, though its training-data generation remains computationally intensive.
Problem
Most MPNNs use 2-body messages, limiting universality, while accurate equivariant models can require high computational cost and poor scalability.
Method
MACE extends equivariant MPNNs with efficiently constructed higher-body-order messages using tensor products and symmetrization.
Results
MACE reaches state-of-the-art accuracy on the rMD17, 3BPA, and AcAc benchmarks while requiring only two message-passing iterations and training faster than previous accurate models.
Takeaways & Limitations
Higher-order messages produce efficient, parallelizable models and improve the steepness of learning curves.
Takeaways & Limitations
Machine-learned force fields require ab initio training data, leading to substantial computation and energy consumption.
Abstract
from arXiv · showhide
Creating fast and accurate force fields is a long-standing challenge in computational chemistry and materials science. Recently, several equivariant message passing neural networks (MPNNs) have been shown to outperform models built using other approaches in terms of accuracy. However, most MPNNs suffer from high computational cost and poor scalability. We propose that these limitations arise because MPNNs only pass two-body messages leading to a direct relationship between 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 order messages. In particular, we show that using four-body messages reduces the required number of message passing iterations to just two, resulting in a fast and highly parallelizable model, reaching or exceeding state-of-the-art accuracy on the rMD17, 3BPA, and AcAc benchmark tasks. We also demonstrate that using higher order messages leads to an improved steepness of the learning curves.
1 Introduction
MACE addresses the limited accuracy, efficiency, and scalability of earlier force-field models by combining equivariant message passing with higher-body-order messages. The architecture uses only two message-passing iterations while achieving strong benchmark performance and improved learning behavior.
- Earlier graph neural-network force fields generally approached but did not improve upon atom-centered descriptor methods in computational efficiency.
- Most MPNN interatomic potentials use 2-body invariant messages, making them non-universal approximators.
- Two message-passing iterations suffice for high accuracy, unlike the typical five or six iterations of MPNNs, improving scalability and parallelizability.
- MACE combines equivariant message passing with efficient many-body messages and achieves state-of-the-art performance on challenging benchmarks.
- MACE also shows greater extrapolation generalization, altered learning-curve power laws, and faster training and evaluation than previous accurate models.
2 Background
MPNNs represent atomistic structures as spatial graphs and iteratively update node features through neighbor messages. Equivariant variants preserve structured transformations of internal features under rotations and reflections.
- An atomistic MPNN represents atoms as graph nodes and connects pairs whose separation is within a specified distance.
- A forward pass repeatedly constructs neighbor-pooled messages, updates node features, and reads out target quantities such as atomic site energies.
- Equivariant GNN features transform according to a specified group action rather than remaining unchanged under geometric transformations.
- For atomic potential energies, the relevant symmetry group is O(3), covering rotations and reflections of particle positions.
- Features with L = 0 are invariant scalars, whereas L > 0 features are equivariant vectors, matrices, or higher-order tensors.
3 Related Work
Prior work established both systematic higher-body-order descriptors and equivariant MPNNs, but accurate equivariant models often required many layers and high computational cost. MACE combines higher body order with message passing to occupy an intermediate design point.
- Atomic Cluster Expansion provides systematic complete polynomial basis functions with high body order at constant cost per basis function.
- Equivariant MPNNs improved performance over earlier local-descriptor models but often used L = 3 messages and 4 to 6 iterations, increasing computational cost.
- Three-body invariant messages add angular information and improve accuracy over 2-body messages, but explicit triplet or quadruplet sums increase computational cost.
- Multi-ACE characterizes models by layers, message body order, equivariance or invariance, and features per layer.
- MACE combines high body order with message passing rather than relying exclusively on many layers or a single high-body-order layer.
4 The MACE Architecture
MACE constructs equivariant higher-order features efficiently by tensorizing and symmetrizing pooled neighbor features, avoiding explicit multidimensional sums. Its basis can represent many-body equivariant mappings while keeping computation near pairwise cost.
- MACE introduces a new message-construction mechanism within the general MPNN framework.
- Neighbor pooling produces permutation-invariant 2-body features while retaining directional information about the atomic environment.
- The first-layer operation is computationally cheap because initial node features correspond to invariant chemical elements.
- Higher-order features are formed by tensor products followed by symmetrization of the lower-order features.
- The tensor-product order ν corresponds to the order of the many-body expansion terms in the message.
- MACE avoids explicit sums over triplets and quadruplets by absorbing them into its efficient feature construction.
- Under mild conditions, the higher-order features form a complete many-body basis computable at a cost comparable to pairwise interactions.
- Invariant node features are mapped to site-energy contributions, preserving invariance in the readout.
5 Results
MACE uses higher-order messages to reach accurate predictions with only two message-passing layers, improving computational efficiency and parallelization. Higher body order also steepens learning curves and supports strong accuracy, extrapolation, and speed across molecular benchmarks.
- Effect of Higher Order Messages: Two message-passing layers suffice for MACE to converge in accuracy, whereas two-body approaches require up to five iterations.This reduces the number of layers used in subsequent experiments and supports more parallel computation.
- Learning Curves: Higher-order invariant messages steepen learning curves and substantially improve accuracy, with ν = 3 approaching highly accurate two-body equivariant models.At fixed ν = 3, increasing equivariance shifts curves without materially changing their slope.
- Chemical Elements: MACE avoids explicit species-dependent scaling by using continuous species embeddings and excluding the species dimension from higher-order tensor products.The resulting model scales as O(1) with the number of chemical elements S.
- Receptive Field and Computational Cost: MACE combines reduced receptive-field growth with node-based tensor products, addressing computational and multi-GPU parallelization costs of conventional MPNNs.Traditional MPNNs can reach receptive fields of up to 30 Å, whereas MACE builds correlations through node operations after evaluating the expensive equivariant tensor product once.
- Computational Cost: MACE is substantially faster than competing models while maintaining strong accuracy: L = 0 is nearly 10 times faster, and L = 2 is around four times faster than other equivariant MPNNs.The reported MACE training run reached converged BOTNet accuracy in approximately 30 minutes, compared with more than a day for BOTNet on 3BPA.
- Benchmark Accuracy: MACE outperforms other models across the reported 3BPA tasks, including about 30% lower errors than NequIP and Allegro at 1200 K and 40% lower dihedral-slice energy errors than BOTNet.The invariant L = 0 model often nearly matches or exceeds competitive equivariant models.
6 Discussions
MACE extends equivariant message passing from two-body to many-body messages, reducing the number of iterations while retaining high accuracy and extrapolation capability. The resulting models are efficient and parallelizable, though larger-system evaluation remains future work.
- MACE extends traditional equivariant MPNNs from 2-body to many-body message passing.
- Two message-passing iterations reduce computational requirements and produce efficient, parallelizable models.
- MACE reaches state-of-the-art accuracy on the rMD17, 3BPA, and AcAc benchmarks.
- Future work should test MACE on larger systems, including condensed phases and solids.
7 Reproducibility Statements
The paper reports reproducibility measures including error bars, ablation studies, documented hyperparameters and data, and publicly available source code.
- Error bars from different seeds and various ablation studies are included where necessary and appropriate.
- Appendix A.5 provides the hyperparameters and data description.
- Source code is available through the ACEsuit MACE GitHub repository.
8 Ethical Statements
The paper notes potential societal benefits from better force fields but emphasizes that machine-learning force fields require computationally intensive ab initio training data.
- Better force fields may speed drug discovery and help understand, control, and design new materials.
- Machine-learning force fields rely on ab initio training data, causing heavy computation and large energy consumption.
- Machine-learned force fields significantly reduce molecular-modeling costs compared with using solely ab initio methods.
Checklist
The checklist indicates that the paper reports its contributions, scope, limitations, societal impacts, reproducibility materials, training details, error bars, compute resources, and asset citations.
- Author and ethics review: The authors report that they accurately described the paper’s contributions and scope and discussed its limitations.They also report discussing potential negative societal impacts and following ethics review guidelines.
- Reproducibility: The authors report that theoretical-results requirements were not applicable, while experimental reproducibility materials and training details were included.They marked code, data, instructions, training details, data splits, and hyperparameters as provided.
- Experimental reporting: The authors report running experiments with error bars and documenting the total compute and resource types used.Both experimental uncertainty reporting and computational-resource disclosure were marked yes.
- Existing and new assets: The authors report citing creators of existing assets, while asset licensing and release of new assets were marked not applicable.The checklist separately records citation, licensing, and new-asset availability.
- Human-subjects review: Human-subjects, consent, personally identifiable information, offensive-content, participant-instruction, risk, and compensation disclosures were marked not applicable.The checklist marks these requirements as N/A across the relevant items.
A.1 Acetylacetone Dataset: Additional Experiments
The acetylacetone experiments test MACE and competing models on dihedral rotation and intramolecular hydrogen transfer, including extrapolation beyond the training distribution. For hydrogen transfer, all models reproduce the barrier shape, while MPNNs closely match its height.
- The acetylacetone experiments compare BOTNet, NequIP, MACE, and linear ACE on two potential-energy-surface trajectories.The trajectories cover a dihedral rotation and a minimum-energy path for intramolecular hydrogen transfer.
- Training contains dihedral angles below 30°, so accurate predictions up to 180° require substantial extrapolation.The dihedral experiment evaluates rotation around an O-C-C-C angle beyond the training range.
- The hydrogen-transfer task probes bond-breaking behavior absent from the training data, although the reaction region remains relatively near the training distribution.The training-data histogram indicates that this path is not too far from the sampled region.
- All models reproduce the hydrogen-transfer barrier shape, and the MPNN models closely match the barrier height.