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

cs.LGMar 6, 2026

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

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

Organizations: UC Berkeley · FAIR at Meta · LBNL

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.

Figures & tables

Figure 1: AllScAIP model design. The simple backbone design enables efficient scaling.

Figure 2: Illustration of the attention operations used in AllScAIP. (a) Neighbor attention. (b) Node attention.

Figure 3: Illustration of the geometric encoding used in AllScAIP. (a) Legendre Angular Encoding (LAE) (b) Euclidean rotary position encoding.

						OMol 4M (80 Epochs)
		Ablations				Biomol.		Elytes.		Metal Cplx.		Neutral Org.		Total
Size	Thrpt	NeiAtt	LAE	NodeAtt	ERoPE	E $\downarrow$	F $\downarrow$	E $\downarrow$	F $\downarrow$	E $\downarrow$	F $\downarrow$	E $\downarrow$	F $\downarrow$	E $\downarrow$	F $\downarrow$
35M	2.279	✓	✓	✓	✓	$0.39$	$5.58$	$1.27$	$8.92$	$2.85$	$33.53$	$1.9$	$13.52$	$1.47$	$9.36$
	7.623	✓		✓		$0.5$	$6.32$	$1.52$	$9.81$	$2.96$	$34.56$	$2.22$	$14.51$	$1.76$	$10.21$
85M	1.124	✓	✓	✓	✓	$0.23$	$3.51$	$1.05$	$6.38$	$2.45$	$27.84$	$1.2$	$8.93$	$1.16$	$6.89$
	3.333	✓		✓	✓	$0.29$	$4.28$	$1.17$	$7.21$	$2.67$	$29.79$	$1.53$	$10.53$	$1.22$	$7.6$

Table 2: Component ablations under data and model size scaling . We report Energy / Atom MAE (meV) and Force MAE (meV/Å), lower is better. Each size’s with encoding (Nei Att + LAE + Node Att + ERoPE) configuration is the reference; other rows are colorized relative to that reference ( Green = better , Red = worse ). The Throughput column reports inference speed (ns/day on one H200 with 1000 atoms). We report results for the OMol25 4M split (80 epochs) and the full 102M (10 epochs) 2 2 2 We note that there has been an identified bug in the FAIRChem codebase that made the per atom MAE metric appear to be consistently worse for both AllScAIP and all baselines. We have attached the unnormalized total energy MAE in the Appendix A.4 .

Model	Training Set	Energy $\Delta$ [meV] $\downarrow$	Energy $\Delta$ [meV] $\downarrow$	Cos. Sim. $\uparrow$	[meV / atom / ps] $\downarrow$
		PBC	non-PBC	Rand. Rotation	NVE MD Energy Drift
UMA-M-1p1	UMA-459M	$2.41\times 10^{-2}$	$3.83\times 10^{-3}$	0.9999	$1.5\times 10^{-3}$
AllScAIP-md	NA (Random Init)	$3.62\times 10^{-2}$	$2.49\times 10^{-3}$	0.6827	-
AllScAIP-md	OMol-102M	$4.18\times 10^{-2}$	$4.73\times 10^{-3}$	0.9999	$4.3\times 10^{-3}$

Table 3: Symmetry and Conservation Checks. Extensivity errors for periodic (PBC) and non-periodic (non-PBC) systems, (lower is better), rotational equivariance via force cosine similarity under random rotations (higher is better), and NVE molecular dynamic simulation energy drift on MD22 large molecules (lower is better).

Figure 4: Throughput and Memory vs. System size. Left: Atom-time (number of atoms

\times

ns/day, higher is better) vs. system size. The result is measured on a single H200 141G with graph generation off. We report four model sizes (35M/85M/180M/1B) of AllScAIP, and eSEN baselines. The dotted vertical lines indicate approximately when the

\mathcal{O}(N^{2})

of the node attention dominates over the

\mathcal{O}(Nk)

of the neighborhood attention, where

k

is the max number of neighbors. Right: vram usage vs. system size. Dotted horizontal lines indicate the vram size for common GPUs.

		Biomol.		Elytes.		Metal Cplx.		Neutral Org.		Total
Dataset	Model	E $\downarrow$	F $\downarrow$	E $\downarrow$	F $\downarrow$	E $\downarrow$	F $\downarrow$	E $\downarrow$	F $\downarrow$	E $\downarrow$	F $\downarrow$
All	eSEN-sm-d.	0.67	6.30	1.24	9.41	2.53	33.08	1.23	13.84	1.49	9.92
	eSEN-sm-cons.	0.59	4.61	1.01	8.08	2.30	28.86	0.84	11.11	1.27	8.25
	eSEN-md-d.	0.34	2.61	0.69	4.40	1.73	19.99	0.59	5.63	0.84	4.76
	GemNet-OC	0.15	3.88	0.56	5.98	1.83	25.12	0.86	10.38	0.66	6.52
	AllScAIP-sm-ft-cons.	0.22	3.84	0.75	6.01	2.09	26.02	1.00	8.55	0.85	6.61

Table 4: OMol25 validation results. Energy / Atom MAE (meV) and Force MAE (meV/Å) across Biomolecules, Electrolytes, Metal Complexes, and Neutral Organics. AllScAIP variants achieve the lowest overall energy error on both splits, with competitive force errors.

Figure 5: OMol25 energy error vs. efficiency throughput (inference). Energy / Atom MAE (meV,

\downarrow

) vs. throughput (ns/day,

\uparrow

) for 35M/85M models with/without encodings; compared with eSEN ( Fu et al., 2025 ) baselines. Conservative models are labeled with a hollow marker. Left: Models trained on OMol25 4M (80 epochs). Right: OMol25 102M (12 epochs). Our models trace the Pareto front at 4M; at 102M the gap between with/without encodings shrinks or flips, indicating these inductive bias may be unnecessary at scale. The larger speed gap between the direct force and conservative AllScAIP models, compared to eSEN, occurs because the differentiable kNN graph construction used in AllScAIP is a newly introduced operation that has not yet been optimized ( Liu et al., 2026 ) .

	Energy		Avg.	Ligand pocket	Ligand strain	Conformers	Protonation	Dist. scaling	IE/EA	Spin gap
Model	Cons.	Training Set	Rank	Ixn Energy $\downarrow$	Strain Energy $\downarrow$	$\Delta$ Energy $\downarrow$	$\Delta$ Energy $\downarrow$	$\Delta$ Energy $\downarrow$	$\Delta$ Energy $\downarrow$	$\Delta$ Energy $\downarrow$
AllScAIP-md-d.		OMol-102M	3.43	29.496	4.149	4.985	16.104	10.082	143.035	350.825
AllScAIP-md-cons.	✓	OMol-102M	4.14	50.532	4.103	3.874	17.092	36.145	148.976	317.400
eSEN-md-d.		OMol-102M	4.14	64.246	3.113	3.566	14.988	109.767	145.776	303.883
UMA-M-1p1	✓	UMA-459M	4.86	76.778	2.962	2.532	18.457	138.094	136.386	335.058
UMA-S-1p2	✓	UMA-520M	5.00	66.604	4.365	3.738	18.914	32.174	176.939	246.565

Table 5: Evaluation results across OMol25 test evaluations. Results are reported in energy MAE error [meV] (lower is better). Models are sorted by the average ranking of individual categories, with a lower average ranking indicating better overall performance.

Figure 6: OMol25 distance-scaling evaluation. A group of molecules are uniformly compressed and stretched by a scalar "stretch factor" to probe the long-range ability (x-axis;

<1

: compressed,

>1

: stretched) and report energy absolute error per atom (meV;

\downarrow

). AllScAIP-md-cons. stays flat and low across the full range, while eSEN and UMA md-1p1 degrade sharply under stretching, indicating poor long-range capacities. Insets show example geometries at selected factors labeled with dotted lines.

Figure 7: MD Simulations vs. Experimental Density and Enthalpy of Vaporization. NPT MD on 39 different molecular liquids and NVT MD on isolated molecules for 300 ps. Left: predicted vs. experimental densities. Right: predicted vs. experimental enthalpy of vaporization. Top: models trained on OMol25. Bottom: models trained on SPICE (note that SPICE structures, computed at a different DFT level of theory, is also a subset of OMol25). AllScAIP attains lower MAE and higher

R^{2}

than MACELES and eSEN on both properties.

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 8: Runtime breakdown by component vs. system size. Stacked bars show the percentage of wall-clock inference time spent in each module of AllScAIP on a single H200 (graph construction off), for three model sizes (top: 35M, middle: 85M, bottom: 180M). Colors: Neighbor attention (blue,

\mathcal{O}(Nk)

), Neighbor FFN (magenta), Node attention (orange,

\mathcal{O}(N^{2})

), Node FFN (green), and Input/embeddings (red). At small systems the runtime is dominated by the local block (neighbor attention). As the number of atoms grows into the tens of thousands, the cost of the global node attention increases and eventually becomes the majority of runtime; the crossover happens earlier for larger models.

Figure 9: Throughput vs. System size. Throughput (ns/day, higher is better) is measured on a single H200 141G with graph generation off. Lines show four model sizes (35M/85M/180M/1B) of AllScAIP and eSEN baselines.

Figure 10: OMol25 Prediction Error vs. System Diameter (Å) and Total Charge. Energy / Atom MAE (meV,

\downarrow

, top) and Force MAE (meV/Å,

\downarrow

, bottom) binned by (left) system diameter (Å) and (right) total charge on all of OMol25 validation set. Bars compare AllScAIP-md, UMA-md, and eSEN-md. AllScAIP remains lower and more stable across bins, while eSEN and UMA error grows for large-diameter systems and for high total charge.

Figure 11: Interpreting the All-to-all Node Attention. Example protein pocket with 342 atoms in OMol Validation Set (top-left). Top-right: per-head attention weight in the last node-attention layer as a function of pairwise distance; thin lines = individual heads, thick black = head average. The red dashed line marks the local neighbor cutoff (6 Å). Bottom-left: pairwise distance matrix. Bottom-right: head-averaged attention matrix. Several heads allocate noticeable mass to pairs well beyond the local cutoff, and the attention heatmap shows column "hubs" that attract long-range attention across many rows.

						OMol 4M (80 Epochs)
		Ablations				Biomol.		Elytes.		Metal Cplx.		Neutral Org.		Total
Size	Thrpt	NeiAtt	LAE	NodeAtt	ERoPE	E $\downarrow$	F $\downarrow$	E $\downarrow$	F $\downarrow$	E $\downarrow$	F $\downarrow$	E $\downarrow$	F $\downarrow$	E $\downarrow$	F $\downarrow$
35M	2.279	✓	✓	✓	✓	$56.9$	$5.58$	$89.3$	$8.92$	$166.2$	$33.53$	$50.3$	$13.52$	$85.2$	$9.36$
	7.623	✓		✓		$73$	$6.32$	$112.3$	$9.81$	$174$	$34.56$	$58.2$	$14.51$	$102.9$	$10.21$
85M	1.124	✓	✓	✓	✓	$34.1$	$3.51$	$62.2$	$6.38$	$139.6$	$27.84$	$31.4$	$8.93$	$59.7$	$6.89$
	3.333	✓		✓	✓	$43.6$	$4.28$	$70.8$	$7.21$	$153.7$	$29.79$	$40.5$	$10.53$	$68.1$	$7.6$

Appendix

Table 6: Component ablations under data and model size scaling . We report Total Energy MAE (meV) and Force MAE (meV/Å), lower is better. Each size’s with encoding (Nei Att + LAE + Node Att + ERoPE) configuration is the reference; other rows are colorized relative to that reference ( Green = better , Red = worse ). The Throughput column reports inference speed (ns/day on one H200 with 1000 atoms). We report results for the OMol25 4M split (80 epochs) and the full 102M (10 epochs).

		Biomol.		Elytes.		Metal Cplx.		Neutral Org.		Total
Dataset	Model	E $\downarrow$	F $\downarrow$	E $\downarrow$	F $\downarrow$	E $\downarrow$	F $\downarrow$	E $\downarrow$	F $\downarrow$	E $\downarrow$	F $\downarrow$
All	eSEN-sm-d.	96.92	6.30	88.20	9.41	145.45	33.08	34.24	13.84	89.34	9.92
	eSEN-sm-cons.	86.72	4.61	72.51	8.07	132.71	28.86	22.98	11.11	75.98	8.25
	eSEN-md-d.	50.72	2.61	47.37	4.39	97.55	19.99	16.47	5.62	49.21	4.76
	GemNet-OC	23.35	3.88	30.73	5.98	103.10	25.12	21.22	10.38	33.85	6.52
	AllScAIP-sm-ft-cons.	31.94	3.34	39.70	6.06	123.69	25.94	18.77	8.97	42.48	6.41

Appendix

Table 7: OMol25 validation results. Total Energy MAE (meV) and Force MAE (meV/Å) across Biomolecules, Electrolytes, Metal Complexes, and Neutral Organics.

	OC20 (Total Energy)						OMat24
	Val ID			Val OOD-Both			Val
Model	Energy	Force	F. Cos.	Energy	Force	F. Cos.	E./Atom	Forces	Stress	F. Cos.
UMA-S	63.6	24.1	0.63	107.0	29.2	0.65	11.3	57.1	2.9	0.98
UMA-M	43.1	15.8	0.73	70.0	19.2	0.75	10.0	47.3	2.7	0.99
UMA-L	32.6	12.0	0.77	49.8	14.5	0.79	9.7	43.5	2.5	0.99
eqV2-S							11	49.2	2.4	0.985

Appendix

Table 8: OC20 validation (Total Energy) and OMat24 validation results.

	MACE-L		MACE-M		MACELES-M		EScAIP		eSEN		AllScAIP
Dataset	E	F	E	F	E	F	E	F	E	F	E	F
PubChem	0.88	14.75	0.91	20.57	0.71	17.24	0.53	5.86	0.15	4.21	0.13	3.21
DES370K M.	0.59	6.58	0.63	9.36	0.54	7.65	0.41	3.48	0.13	1.24	0.07	0.97
DES370K D.	0.54	6.62	0.58	9.02	0.47	7.48	0.38	2.18	0.15	2.12	0.06	0.95
Dipeptides	0.42	10.19	0.52	14.27	0.52	11.46	0.31	5.12	0.25	3.68	0.06	1.54
Solvated A.A.	0.98	19.43	1.21	23.26	0.82	18.37	0.61	11.52	0.25	3.68	0.13	4.86

Appendix

Table 9: SPICE test results. Energy / Atom MAE (meV,

\downarrow

) and Force MAE (meV/Å,

\downarrow

) across different splits of SPICE. AllScAIP achieve the lowest overall energy and force errors.

Figure 12: Zero-shot MD on MD22: radial distribution

h(r)

. NVT trajectories with OMol-trained AllScAIP on seven MD22 systems. Curves show

h(r)

from the simulations compared with PBE

+

MBD reference runs. Although the reference level differs from OMol, AllScAIP tracks eSEN/UMA in peak locations and overall shape, indicating matched structural statistics.

Model Size	small	medium	large	xlarge
hidden_size	512	640	1024	2048
num_layers	6	10	8	12
atten_num_heads	8	10	16	32
atten_name	memory_efficient
activation	gelu
normalization	rmsnorm

Appendix

Table 10: AllScAIP Model Configurations

Machine Learning Interatomic Potentials

Machine learning interatomic potentials predict atomic energies and forces for fast simulations of molecules and materials.

See all 91 papers

MatRIS: Toward Reliable and Efficient Pretrained Machine Learning Interatomic Potentials
Mar 2, 2026Yuanchang Zhou, Siyu Hu, Xiangyu Zhang +3Machine Learning Interatomic Potentials Computational Materials Science
State Key Lab of Processors, Institute of Computing Technology, Chinese Academy of Sciences · University of Chinese Academy of Sciences · School of Advanced Interdisciplinary Sciences, UCAS
BranchIP: Learning Adaptive Equivariant Computation for Interatomic Potentials
Oct 1, 2026Laura Zichi, Gil Harari, Chuin Wei Tan +6Machine Learning Interatomic Potentials Equivariant Neural Networks
Harvard University · Robert Bosch LLC Research and Technology Center
Scalable long-range machine learning interatomic potentials with electrical response for atomistic material modeling
Oct 16, 2026Yicheng Chen, Xiaoyun Yu, Hyejeong Song +5Machine Learning Interatomic Potentials Computational Materials Science
Department of Materials Science and Engineering, National University of Singapore, Singapore · Department of Physics, National University of Singapore, Singapore

Mar 2, 2026J/K move · Enter open · S save

MatRIS: Toward Reliable and Efficient Pretrained Machine Learning Interatomic Potentials

Yuanchang Zhou, Siyu Hu, Xiangyu Zhang, Hongyu Wang+2

State Key Lab of Processors, Institute of Computing Technology, Chinese Academy of Sciences · University of Chinese Academy of Sciences · School of Advanced Interdisciplinary Sciences, UCAS

Machine Learning Interatomic Potentials Computational Materials Science Invariant Representation Learning Materials Property Prediction

Open paper AskPDF arXiv

Foundation MLIPs demonstrate broad applicability across diverse material systems and have emerged as a powerful and transformative paradigm in chemical and computational materials science. Equivariant MLIPs achieve state-of-the-art accuracy in a wide range of benchmarks by incorporating equivariant inductive bias. However, the reliance on tensor products and high-degree representations makes them computationally costly. This raises a fundamental question: as quantum mechanical-based datasets continue to expand, can we develop a more compact model to thoroughly exploit high-dimensional atomic interactions? In this work, we present MatRIS (Materials Representation and Interaction Simulation), an invariant MLIP that introduces attention-based modeling of three-body interactions. MatRIS leverages a novel separable attention mechanism with linear complexity

O(N)

, enabling both scalability and expressiveness. MatRIS delivers accuracy comparable to that of leading equivariant models on a wide range of popular benchmarks (Matbench-Discovery, MatPES, MDR phonon, Molecular dataset, etc). Taking Matbench-Discovery as an example, MatRIS achieves an F1 score of up to 0.847 and attains comparable accuracy at a lower training cost. The work indicates that our carefully designed invariant models can match or exceed the accuracy of equivariant models at a fraction of the cost, shedding light on the development of accurate and efficient MLIPs.

Abstract

Machine Learning Interatomic Potentials

Explore similar work

MatRIS: Toward Reliable and Efficient Pretrained Machine Learning Interatomic Potentials

BranchIP: Learning Adaptive Equivariant Computation for Interatomic Potentials

Scalable long-range machine learning interatomic potentials with electrical response for atomistic material modeling