MARA: Continuous SE(3)-Equivariant Attention for Molecular Force Fields

TL;DR

MARA introduces a continuous spherical attention module that operates on angular and radial coordinates to approximate SE(3)-equivariance in molecular force fields. The module is pluggable into existing SE(3)-equivariant backbones (e.g., MACE) and uses a discretized attention-on-sphere mechanism with a radial scalar field to weight local environments by both direction and distance. Across benchmarks, MARA improves energy and force predictions and reduces tail-risk, while ablations show its effectiveness across angular resolutions and grid densities, complemented by interpretable attention maps. The approach enhances expressiveness, stability, and reliability of atomistic models and offers a flexible framework for integrating continuous geometric attention into SE(3)-equivariant MLFFs with modest computational overhead.

Abstract

Machine learning force fields (MLFFs) have become essential for accurate and efficient atomistic modeling. Despite their high accuracy, most existing approaches rely on fixed angular expansions, limiting flexibility in weighting local geometric interactions. We introduce Modular Angular-Radial Attention (MARA), a module that extends spherical attention -- originally developed for SO(3) tasks -- to the molecular domain and SE(3), providing an efficient approximation of equivariant interactions. MARA operates directly on the angular and radial coordinates of neighboring atoms, enabling flexible, geometrically informed, and modular weighting of local environments. Unlike existing attention mechanisms in SE(3)-equivariant architectures, MARA can be integrated in a plug-and-play manner into models such as MACE without architectural modifications. Across molecular benchmarks, MARA improves energy and force predictions, reduces high-error events, and enhances robustness. These results demonstrate that continuous spherical attention is an effective and generalizable geometric operator that increases the expressiveness, stability, and reliability of atomistic models.

Figures (10)

• Figure 1: Illustration of a Machine Learning Force Field (MLFF) modeled with Equivariant Graph Neural Networks (EGNNs). Given a molecular configuration at a specific temperature, the model predicts the forces acting on each atom, which are then used to estimate the molecule’s subsequent configuration. Molecular configurations that occur more frequently are captured more accurately by the model, while less frequent configurations, which can be visualized through the interatomic distribution, are inherently more uncertain.
• Figure 2: The molecule is represented as a 3D molecular graph, in which the nodes are characterized by scalar features and spatial coordinates. Different equivariant graph neural networks process this information using different message passing mechanisms: Equivariant Graph Neural Networks (EGNN), SE(3)-equivariant networks (SE(3)-GNN), and SE(3)-GNN enriched with MARA.
• Figure 3: Violin plots of force error distributions for rMD17 molecules. Blue violins represent the MACE baseline, and orange represent MACE-MARA. Each molecule shows separate distributions for the two models, illustrating the overall error spread.
• Figure 4: Interpretability of MARA. (a) Radial slice of the attention map showing attention from heavy atoms ($O_1$, $O_2$, $C_1$, $C_2$, $C_3$) to hydrogen ($H_1$) during intramolecular proton transfer. (b) Angular slice of the attention map showing attention between $C_3$ and $H_3$ as $H_3$ is rotated around $C_3$ at a fixed distance of $1,\text{\AA}$
• Figure 5: Module operating diagram