A recipe for scalable attention-based MLIPs: unlocking long-range accuracy with all-to-all node attention
| Authors | Eric Qu et al. |
| Year | 2026 |
| Field | Machine Learning |
| arXiv | 2603.06567 |
| Download | |
| Categories | cs.LG, cs.CE |
Abstract
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.
Engineering Breakdown
Plain English
This paper presents AllScAIP, a machine-learning model designed to predict atomic forces and energies in molecular systems by learning from simulations rather than physics equations. The core problem is that existing models fail at capturing long-range interactions (forces between distant atoms) when scaled to large, complex systems like proteins and electrolyte solutions, forcing engineers to hand-code additional physics rules. AllScAIP solves this using an attention-based architecture that learns long-range dependencies directly from data, and demonstrates it can train on 100 million samples efficiently. The key finding is that while explicit physics biases help when data is scarce, they become unnecessary and potentially limiting once you have enough data and model capacity.
Core Technical Contribution
The authors introduce a data-driven attention mechanism (all-to-all node attention) that replaces hand-crafted physics terms for capturing long-range interactions in molecular systems. Rather than explicitly encoding distance cutoffs or physics equations, AllScAIP learns which atoms should influence which other atoms through standard transformer-style attention, allowing the model to discover long-range patterns automatically. This is paired with energy conservation guarantees (the model respects fundamental physics laws about energy stability), which prior attention-based approaches lacked. The core novelty is proving that pure learned attention can match or exceed physics-augmented models at scale, challenging the conventional wisdom that neural potentials need explicit inductive biases.
How It Works
AllScAIP takes as input the 3D positions and types of atoms in a molecular system, then processes them through an all-to-all attention layer where each atom queries every other atom to determine long-range interactions—this is the key difference from prior work that uses fixed distance cutoffs. The attention weights learned by the model implicitly capture which pairwise interactions matter, without requiring engineers to specify cutoff distances or hand-designed basis functions. The architecture enforces energy conservation by computing forces as gradients of a scalar energy function, guaranteeing that the model respects fundamental physics constraints. The output is per-atom forces and total system energy, which are differentiable and can be used to drive molecular dynamics simulations or optimize structures. Training scales to 100M samples using standard GPU/TPU infrastructure by batching many independent molecular snapshots together.
Production Impact
For materials science and drug discovery teams, AllScAIP could replace or dramatically reduce reliance on expensive quantum mechanical simulations (DFT) by providing accurate force predictions in milliseconds instead of hours. If you're building a molecular dynamics pipeline today, swapping in AllScAIP means you can simulate larger biomolecules (proteins in solvent) or longer timescales (microseconds instead of nanoseconds) while maintaining accuracy. The attention mechanism makes the model interpretable—you can visualize which atoms the model considers important for predictions, helping chemists understand failure modes. The main trade-off is that you need 100M+ training samples (generated from simulations), which is expensive upfront but amortized across many downstream predictions; for systems with <1M available training frames, traditional physics-biased models remain more sample-efficient. Latency is low (suitable for online prediction during molecular dynamics), and the model produces differentiable outputs, enabling gradient-based optimization of molecular properties.
Limitations and When Not to Use This
The paper doesn't address how the model generalizes to chemical systems outside its training distribution—if you train on organic molecules and apply it to inorganic crystals, performance likely degrades severely, and the paper provides no transfer learning strategy. Energy conservation is enforced through architectural constraints, but this guarantee only holds during inference; it's unclear whether the model maintains accuracy at the extreme configurations (high temperatures, stretched bonds) that appear in realistic MD simulations. The all-to-all attention scales quadratically with system size (O(N²) memory and compute), making it impractical for large-scale systems like nanoparticles with >10,000 atoms; the paper doesn't discuss approximations or sparse attention variants to address this. The work assumes access to high-quality reference data (DFT or CCSD(T) calculations) to train on, but doesn't explore robustness to label noise—real-world data pipelines often have systematic errors or inconsistencies that could corrupt learning.
Research Context
This work continues the trend of scaling neural network-based interatomic potentials (MLIPs), building on earlier successes with graph neural networks and equivariant architectures but pushing the frontier toward learning complex many-body physics without explicit inductive biases. It directly challenges the premise of recent work like Allegro and NequIP, which achieved state-of-the-art results by carefully engineering symmetries and cutoff functions; AllScAIP shows that with sufficient data, learned attention may outperform hand-crafted design. The paper contributes to broader ML research questions about when inductive biases help versus hurt—finding empirically that they matter in low-data regimes but become a bottleneck at scale, which has implications for other domains (vision, NLP). The work opens the door to purely data-driven force field development and may accelerate adoption of ML potentials in industrial computational chemistry, where the ability to simulate realistic solvation and electrolyte effects has been a major gap.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::