MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields (2206.07697v2)

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

• The paper introduces a novel architecture that uses four-body messages to reduce iterations while achieving state-of-the-art performance on benchmark datasets.
• It emphasizes equivariant transformations that accurately capture rotational symmetries, enhancing model expressivity and data efficiency.
• MACE significantly cuts computational costs and training times, enabling rapid force field predictions for molecular simulations.

MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields

The paper introduces MACE, an advanced neural network architecture for developing force fields, focusing on efficiency and precision within the domain of computational chemistry and materials science. Recognizing that traditional message passing neural networks (MPNNs) suffer from scalability issues due to their reliance on predominantly two-body message interactions, this work seeks to circumvent these limitations by integrating higher-order message passing mechanisms.

Overview and Contributions

1. Higher Order Message Passing: MACE introduces a new architecture that utilizes four-body messages rather than traditional two-body messages. This reduces the number of required message passing iterations to merely two, enhancing both the speed and parallelizability of the network. This change permits the model to achieve state-of-the-art performance on well-known benchmark datasets, such as rMD17 and 3BPA, with a significantly reduced computational cost.
2. Equivariant Graph Neural Networks: The authors underline the importance of equivariant features in MPNNs to ensure that the networks can adequately represent rotational symmetries present in the physical properties of molecular systems. The MACE model incorporates this by employing an equivarient transformation within its layered architecture.
3. Learning Curves: The paper evidences that the incorporation of higher-order messages alters the empirical power law of the learning curves, highlighting their impact on generalization and data efficiency. This result demonstrates that the blanket increase in message body order leads to better model expressivity without necessitating additional layers.
4. Practical Benefits: Not only does MACE achieve better accuracy in predicting molecular properties, but it also does so more rapidly than previous models. The paper reports a significant reduction in training time compared to traditional models, especially when leveraging modern computing hardware like NVIDIA A100 GPUs.

Implications and Future Directions

The introduction of MACE substantially contributes to the field of machine learning-assisted computational chemistry by reducing the complexity and computational overhead of modeling interatomic interactions. The ability to achieve state-of-the-art results while maintaining fast training times opens the possibility for broader applications, particularly in drug discovery and materials science, where rapid iterative cycles are beneficial.

Theoretically, MACE's approach of disentangling message complexity from iterative deepening suggests new directions for neural network architecture design. It challenges the status quo of deep networks and encourages finding optimal balances between depth, complexity, and computational cost.

Future research may delve into extending these methods to more extensive systems, including macromolecules and crystalline materials, where the computational economies of scale achieved here could be even more pronounced. Additionally, exploring the application of higher-order message passing in other domains could reveal versatile utilities across traditional boundaries of neural network applications.

In conclusion, MACE serves as a significant advancement in modeling molecular dynamics by integrating higher-order equivariant message passing, presenting transformative potential for the computational simulation community.