A recipe for scalable attention-based MLIPs: unlocking long-range accuracy with all-to-all node attention

Eric Qu; Brandon M. Wood; Aditi S. Krishnapriyan; Zachary W. Ulissi

Source-linked AI summary

A recipe for scalable attention-based MLIPs: unlocking long-range accuracy with all-to-all node attention

Eric Qu, Brandon M. Wood, Aditi S. Krishnapriyan, Zachary W. Ulissi

arXiv:2603.06567v1Mar 6, 2026cs.LGcond-mat.mtrl-scics.CEphysics.chem-phq-bio.QM

TL;DR

Scalable MLIPs need better treatment of long-range interactions in large, complex systems, where local models often rely on explicit physics-based additions. AllScAIP addresses this with an energy-conserving, attention-based architecture that combines local processing with data-driven all-to-all node attention. It achieves strong molecular accuracy and physics-based performance, while ablations show that geometric priors become less valuable with scale but global attention remains important.

Problem
Large-system MLIPs struggle to capture long-range interactions with scalable local architectures, motivating explicit physics-based additions and leaving broad heterogeneous scaling unresolved.
Method
AllScAIP combines neighborhood self-attention with all-to-all node self-attention in a scalable, energy-conserving, prior-light MLIP.
Results
AllScAIP achieves state-of-the-art molecular energy/force accuracy and physics-based evaluations on OMol25, remains competitive on OMat24 and OC20, and supports stable long-timescale MD recovering experimental observables.
Takeaways & Limitations
Geometric inductive biases improve low-data efficiency, but their benefits diminish or reverse with scale, whereas all-to-all attention remains a durable mechanism for long-range interactions.
Takeaways & Limitations
All-to-all node attention introduces an O(N^2) cost, creating a manageable but explicit efficiency trade-off for long-range accuracy.

Abstract

from arXiv · show

Machine-learning interatomic potentials (MLIPs) have advanced rapidly, with many top models relying on strong physics-based inductive biases. However, as models scale to larger systems like biomolecules and electrolytes, they struggle to accurately capture long-range (LR) interactions, leading current approaches to rely on explicit physics-based terms or components. In this work, we propose AllScAIP, a straightforward, attention-based, and energy-conserving MLIP model that scales to O(100 million) training samples. It addresses the long-range challenge using an all-to-all node attention component that is data-driven. Extensive ablations reveal that in low-data/small-model regimes, inductive biases improve sample efficiency. However, as data and model size scale, these benefits diminish or even reverse, while all-to-all attention remains critical for capturing LR interactions. Our model achieves state-of-the-art energy/force accuracy on molecular systems, as well as a number of physics-based evaluations (OMol25), while being competitive on materials (OMat24) and catalysts (OC20). Furthermore, it enables stable, long-timescale MD simulations that accurately recover experimental observables, including density and heat of vaporization predictions.

1 Introduction

AllScAIP targets long-range accuracy in scalable MLIPs with a data-driven all-to-all node-attention stage and a prior-light design. Ablations indicate that geometric inductive biases help in low-data regimes, whereas global attention remains useful as data and model scale.

Motivation: Long-range interactions matter in larger systems such as biomolecules and electrolytes, but scalable local message-passing MLIPs struggle to capture them.Existing responses commonly add explicit electrostatics, polarization, dispersion, or continuum corrections, but scaling across heterogeneous datasets remains open.
Approach: AllScAIP combines local neighborhood self-attention with all-to-all node self-attention so information can travel globally across the graph.The model also adds optional Legendre Angular Encoding and Euclidean Rotary Position Encoding.
Ablations: In low-data and small-model regimes, geometric encodings and global node attention improve performance, with global attention delivering the largest gains.LAE lowers force error, Euclidean RoPE improves energy, and node attention enables many-hop communication without deep stacks.
Ablations: As data and model capacity increase, geometric encoding benefits diminish or sometimes reverse, while all-to-all node attention remains the most durable source of long-range improvement.This supports the hypothesis that several inductive biases are learnable under scale, but global coupling remains valuable.
Results: AllScAIP reaches state-of-the-art molecular energy/force accuracy and physics-based evaluations on OMol25, while remaining competitive on OMat24 and OC20.OMol25-trained models also recover experimental density and heat-of-vaporization observables in stable, long-timescale MD simulations.
Conclusion: The proposed recipe prioritizes scalable components, keeps priors lightweight, and lets scale learn several otherwise explicit biases.The authors frame this as a prior-light path for scalable MLIPs.

2 Related Works

Related MLIP work spans equivariant architectures and explicit long-range physics, while AllScAIP belongs to the attention-based direction. Existing long-range approaches commonly augment local models with electrostatics, polarization, dispersion, or related physical components.

Machine Learning Interatomic Potentials: Current MLIPs include equivariant node-feature models such as NeuqIP, MACE, SCN, eSCN, Equiformer, EScAIP, MindScAIP, OrbNet, and PET-MAD.These approaches emphasize stronger geometric or symmetry-related inductive biases.
Long-range Interactions in MLIPs: Long-range extensions often predict charges, multipoles, or surrogate charge densities and evaluate electrostatics with Ewald, PME, or FMM.Representative approaches include PhysNet, 4G-HDNNP, AIMNet, BAMBOO, LODE, and DPLR.

3 Methods

AllScAIP uses local neighborhood attention for fine geometry and all-to-all node attention for global communication, with optional angular and radial geometric encodings. Its design enforces selected symmetries, locality, conservation, and extensivity while leaving some structure learnable from scale.

Scaling: The local and global attention stages use O(Nk) and O(N^2) costs, respectively, reflecting the trade-off between local efficiency and global mixing.Both stages use optimized CUDA attention kernels; only token packing differs.
Attention backbone: Neighborhood attention processes each node’s fixed neighbor stencil with standard multi-head self-attention and bidirectional center–neighbor passes.A smooth distance envelope provides padding and masking for far pairs.
Attention backbone: All-to-all node attention applies multi-head self-attention to the full node stream, enabling many-hop coupling in one step.It complements local attention, which handles fine anisotropic interactions.
Geometric encodings: LAE injects compact, rotation-aware angular features into neighborhood-attention logits using spherical-harmonic directional codes.Its multi-order angular structure has linear cost and integrates with standard attention kernels.
Geometric encodings: Euclidean RoPE adds an isotropic radial bias to node-attention logits using sinc features evaluated at multiple frequencies.Learned per-head weights mix the radial basis into an additive attention-logit bias.
Architecture: The architecture initializes atom, distance, and optional charge/spin features, then alternates local attention, normalization and residual operations, edge processing, global attention, and node processing.Updated node features are fed back into the neighbor representation.
Inductive biases: AllScAIP enforces translation and permutation symmetries, locality, energy conservation, and extensivity while optionally supplying light geometric features.Rotation and long-range interactions are treated as structures that large models and datasets may learn.

4 Ablations

Ablations show that geometric inductive biases help most in low-data or small-model settings, while all-to-all node attention remains useful as data and model size increase.

Component ablations: At fixed capacity and 4M training samples, removing LAE worsens force performance, ERoPE primarily harms energies, and removing node attention degrades both.The node-attention gap is especially pronounced for biomolecules, where systems are larger and long-range effects are stronger.
Component ablations: The full component configuration achieves the best results across all evaluated splits in the low-data regime.This supports improved sample efficiency from combining neighborhood attention, geometric encodings, and global node attention.
Model-size scaling: With model scaling from 35M to 85M, absolute errors decrease and the benefit of geometric encodings diminishes, whereas all-to-all attention retains a measurable advantage.At 102M samples, fixed encodings can show parity or slight disadvantages, but global mixing continues to pay off.

5 Results

AllScAIP exhibits predictable system-size scaling, preserves symmetry and conservation properties, and achieves strong molecular accuracy and long-range behavior across evaluations. It also produces realistic molecular-liquid observables and remains competitive on materials and catalyst datasets.

Inference efficiency and system-size scaling: As system size grows, throughput changes from neighborhood-dominated scaling to an O(N^2) node-attention regime with approximately 1/N decline.The transition occurs when all-to-all attention overtakes O(Nk) neighborhood attention; measurements use one H200 141G with graph generation disabled.
Symmetry and conservation checks: The architecture shows near-perfect additivity under periodic supercell doubling and non-periodic vacuum duplication tests.These tests evaluate whether duplicated systems approach twice the single-system energy without introducing unphysical long-range coupling.
Symmetry and conservation checks: AllScAIP trained on OMol-102M achieves 0.9999 force cosine similarity under random rotations, on par with UMA.The test rotates coordinates, rotates the original force prediction, and compares it with the prediction on rotated coordinates.
Open Molecules (OMol25): AllScAIP variants achieve the lowest overall energy error on OMol25 validation with competitive force errors across four chemical categories.The reported categories are Biomolecules, Electrolytes, Metal Complexes, and Neutral Organics.
Open Molecules (OMol25): AllScAIP-md-d. ranks first on the OMol25 test suite and reduces distance-scaling error by about 90% versus the second-best model.Its energy error remains low and flat under molecular compression and stretching, while eSEN and UMA degrade under large stretches.
Molecular dynamics and observables: AllScAIP achieves better MAE and R2 than MACELES for liquid-system evaluations of density and enthalpy of vaporization.The eSEN results systematically over-predict, whereas AllScAIP removes this bias; the broader MD results indicate realistic bulk behavior without explicit long-range terms.
Cross-domain evaluation: AllScAIP achieves competitive performance on the large-scale OMat24 materials and OC20 catalyst datasets.These datasets differ substantially from OMol25 in domain and system type.

6 Discussions

AllScAIP’s long-range capability remains useful across scales, while geometric inductive biases become less beneficial as data and models grow. This accuracy comes with a transition from local O(Nk) computation to globally mixing O(N^2) attention for larger systems.

Scaling and the role of inductive bias: All-to-all node attention remains beneficial at every tested scale, unlike LAE and Euclidean RoPE, whose marginal benefits diminish as data and parameters increase.The ablations support a “scale first, bias second” guideline.
Limitations and opportunities: The O(N^2) cost is presented as a manageable trade-off for long-range accuracy, with hierarchical attention, linear attention, and mixtures-of-experts identified as possible scaling routes.The discussion connects these alternatives to delaying the quadratic regime.
Efficiency and scaling: The long-range module changes the dominant cost from local O(Nk) computation to all-to-all O(N^2) computation beyond a system-size crossover.The crossover occurs earlier for the 180M model and later for the 35M model.
Efficiency and scaling: Throughput is evaluated across four AllScAIP model sizes and eSEN baselines on a single H200, using system size as the independent variable.The throughput plot reports ns/day with graph generation disabled.
Long-range accuracy: AllScAIP maintains lower energy and force MAE across system-diameter and total-charge bins, while eSEN and UMA errors grow for large or highly charged systems.This pattern is reported on the OMol25 validation set.

A.3 Visualization of All-to-all Node Attention

Visualization indicates that all-to-all node attention mixes information over distances beyond the local 6 Å cutoff. Different heads specialize, with some routing information across tens of Å and forming long-range hubs.

Visualization setup: Attention analysis uses pairwise distances and post-softmax per-head matrices from a representative 342-atom protein pocket.Distances are binned in 0.5 Å intervals from 0 to 25 Å.
Distance profiles: Some attention heads place above-baseline mass at large separations, while others remain local, producing a near-uniform head average.The local neighbor cutoff is 6 Å and the baseline is 1/N.
Attention heatmaps: The head-averaged attention matrix reveals long-range hubs: atoms receiving attention from many others, including chemically distinctive charged or polar groups.The distance matrix provides the contact structure used to interpret these bands.
Interpretation: A subset of heads routes information across tens of Å, complementing local neighborhood attention.The visualization is interpreted as evidence of head specialization and long-range mixing.

A.4 Total Energy MAE Results

The reported ablation and validation results use total energy MAE as the evaluation metric.

Metric: The ablation and validation results are presented using total energy MAE.The accompanying tables report energy error as the metric for this analysis.

B OC20 and OMat24 Results

AllScAIP is competitive with UMA on OC20 and OMat24 despite UMA using more data and a larger compute budget. On OC20, its medium model surpasses UMA-S and shows an OOD/ID ratio comparable to UMA-M.

Settings: AllScAIP is trained on OC20 and OMat24 with direct-force training followed by conservative fine-tuning.OC20 uses 3 direct-force epochs and 1 fine-tuning epoch; OMat24 uses 6 and 3, respectively.
Comparison: UMA is trained on 459M samples with a larger compute budget, providing the comparison baseline for OC20 and OMat24.The reported comparison uses the medium AllScAIP models against UMA variants.
OC20: AllScAIP-md-d. surpasses UMA-S on both OC20 Val/ID and Val/OOD-Both.The evaluation metric is total energy MAE.
OMat24: On OMat24, the medium conservative AllScAIP model performs better than UMA-S without substantial tuning effort for either dataset.The source reports this result on the OMat24 validation set.

C SPICE Results

On SPICE, the small energy-conserving AllScAIP variant achieves the lowest overall energy and force errors across the evaluated splits against several long-range MLIP baselines.

Settings: The 34M AllScAIP-sm-cons. model is energy-conserving and trained on SPICE using the MACE-OFF splits.Training used 300 epochs with full fp32 precision on an H200 142GB GPU.
Results: Evaluation covers seven held-out subsets, including PubChem, DES370K monomers/dimers, Dipeptides, Solvated Amino Acids, Water, and QMugs.Performance is reported using Energy / Atom MAE and Force MAE, both with lower values preferred.
Results: AllScAIP-sm-cons. achieves the lowest overall energy and force errors across splits compared with MACE/MACELES, EScAIP, and eSEN.The comparison uses Energy / Atom MAE in meV and Force MAE in meV/Å.

D.1 Results on MD22

Zero-shot NVT simulations on seven MD22 molecules show that OMol-trained AllScAIP reproduces reference structural statistics under distribution shift, despite the reference using a lower theory level.

D.1 Results on MD22: Zero-shot NVT simulations evaluate structural fidelity on seven MD22 molecules using the OMol-trained AllScAIP model.Radial distribution functions are computed from production segments and compared with PBE+MBD reference trajectories.
D.1 Results on MD22: AllScAIP reproduces the reference peak positions and overall radial-distribution shapes across systems.The reference trajectories use PBE+MBD, a lower level of theory than OMol, so exact coincidence is not expected.
D.1 Results on MD22: The matched structural statistics indicate comparable equilibrium structure and no gross biases in intermolecular distances relative to eSEN-sm and UMA-md.The comparison is based on radial distribution functions from the simulated trajectories.