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

In plain words

A machine-learning interatomic potential is a neural network that predicts how much energy a set of atoms has and which way each atom is pushed. Simulations can then run on that prediction instead of an expensive quantum calculation. Most of these networks only let atoms see their close neighbours. That works for small molecules, but big molecules and salts also have effects that span long distances, and a purely local model misses them.

This work takes a plain attention network and adds one global step. Every atom can look at every other atom in the same system, so distant atoms can influence each other in a single step. The model also has a local step that handles the close neighbours in detail. The authors then test two hand-built physics features. One encodes directions between atoms, the other encodes distances. Those features help small models trained on little data. As the data and the model grow, their benefit fades, while the global step keeps helping. The result is a simple model that scales up well on a large molecule dataset, and its simulated liquids come closer to measured densities and vaporisation enthalpies than the baselines it is compared with.

TL;DR

• AllScAIP is an attention-based interatomic potential that predicts energy and forces with an energy-conserving option §3.3 §3.4.
• Each block combines local neighbourhood attention with all-to-all attention over the atoms of one system §3.1.
• Two optional encodings, LAE (angular) and ERoPE (radial), help small, low-data models but their benefit shrinks with scale §4 §6.
• Author-reported: lowest overall OMol25 validation energy error for the 102M-sample medium model, competitive force error §5.3 tab. 4.
• Cost grows from O(Nk) to O(N²) with system size, so the model targets roughly 10³–10⁵ atoms §5.1 fig. 4.

Problem

Local message passing with distance cutoffs cannot carry information across long ranges in one step. Explicit long-range physics terms (electrostatics, dispersion, charge models) handle this but add modelling and engineering burden §1. The paper asks whether an attention model can learn long-range behaviour from data at scale, and which hand-built geometric inductive biases still matter once data and parameters grow §1 §4.

Method

Two attention stages per block §3.1. Local neighbourhood attention runs over a fixed stencil per atom: the atom, a self token and up to k neighbours (shape nodes × (k+1) × d_model). It runs two directional passes, centre to neighbours and neighbours back to the centre, and a smooth distance envelope softly fades far pairs. Global node attention packs all atoms of one system into a single sequence and applies standard multi-head self-attention. The local stage costs O(Nk); the global stage costs O(N²) §3.1 fig. 1.

Geometric encodings, optional §3.2. Both encodings add geometric information to attention logits.

• LAE (neighbourhood attention only): real spherical harmonics of the unit edge direction are computed for degrees l = 0…L. Each degree block is repeated r_l times, with the constraint that the r_l sum to the per-head dimension d_h. Queries and keys are expanded to the same layout and multiplied element-wise by this vector before the usual scaled dot product §3.2 fig. 3.
• ERoPE (node attention only): a bank of M log-spaced frequencies gives sinc(ω·r) features of the pairwise distance. Each head mixes them with its own weight vector, and the result is added to the attention logit as a bias that depends only on distance §3.2 fig. 3. The paper calls this Euclidean rotary encoding, but the operation is an additive bias, not a rotation of queries and keys §3.2 node_attention.py.

Block structure §3.3 fig. 2. The model embeds atomic numbers, a radial basis expansion of edge distances and node/edge direction features, plus optional charge and spin embeddings. Each block then runs, in order: neighbourhood attention; RMSNorm with a residual path; an edge feed-forward network; the self token is taken as the atom's node feature; all-to-all node attention with residual, RMSNorm; a node feed-forward network. The node feature is put back into slot 0 and concatenated with the edge features for the next block. Once all blocks have run, the node features go to an energy output. Forces can be predicted by their own head, or taken as the negative gradient of the energy so they stay consistent with it. A stress head based on that gradient is added when needed. required §3.3.

Inductive biases §3.4 tab. 1. Translation invariance, permutation equivariance, extensivity, locality, and energy conservation are enforced. Rotational equivariance and long-range effects are left to be learned. LAE and ERoPE are optional soft priors. Energy conservation follows a differentiable kNN construction from Liu et al. (2026) §3.4.

# One AllScAIP block  [§3.3]
# [code: models/allscaip/modules/graph_attention_block.py#L77-L107]
# Input: neighbor_reps (N_atoms, 1 + k, d)

def allscaip_block(data, neighbor_reps):
    # 1. Local attention over each atom's stencil, two passes  [§3.1]
    neighbor_reps = neighborhood_attention(data, neighbor_reps)
    # 2. The self token becomes the atom's node feature  [§3.3]
    node_reps = neighbor_reps[:, 0]
    # 3. Feed-forward on the neighbour slots  [§3.3]
    edge_reps = edge_ffn(neighbor_reps[:, 1:])
    # 4. Global attention over all atoms of each system  [§3.1] [§3.3]
    node_reps = node_attention(data, node_reps)
    # 5. Feed-forward on the node features  [§3.3]
    node_reps = node_ffn(node_reps)
    # 6. Put node features back into slot 0 and rejoin the edge features  [§3.3]
    return concatenate([node_reps[:, None], edge_reps], axis=1)

# Legendre angular encoding (LAE) for one directed edge i -> j  [§3.2]
# [code: src/fairchem/core/models/allscaip/utils/data_preprocess.py#L117-L129]

def lae_vector(unit_direction, rep_counts):
    # Real spherical harmonics Y_l for l = 0..L, each scaled by 1/sqrt(2l+1)
    harmonics = real_spherical_harmonics(unit_direction, L)
    # Repeat degree block l exactly rep_counts[l] times
    blocks = [repeat(harmonics[l], rep_counts[l]) for l in range(L + 1)]
    # Length is sum_l rep_counts[l] * (2l+1) = per-head size after expansion
    return concatenate(blocks)                      # (d_lae,)

# Per head: expand query and key the same way, then modulate  [§3.2]
# [code: models/allscaip/modules/neighborhood_attention.py#L192-L277]

def lae_modulate(query, key, gamma):
    # Repeat each query/key dimension 2l+1 times within its degree block
    query_expanded = expand_by_degree(query, rep_counts)
    key_expanded = expand_by_degree(key, rep_counts)
    # Element-wise modulation with the edge code gamma
    return query_expanded * gamma, key_expanded * gamma

# Distance bias for node attention (ERoPE)  [§3.2]
# [code: models/allscaip/modules/node_attention.py#L91-L137]
# [code: models/allscaip/utils/data_preprocess.py#L207-L271]

def radial_bias(distances, valid_pairs, radial_weight):
    # M log-spaced frequencies omega_k (code bounds: pi/(4*30) to pi/0.25)
    omegas = log_spaced(pi / (4 * 30.0), pi / 0.25, num=M)   # (M,)
    # Sinc features of each pair distance  [§3.2]
    features = sinc(distances[..., None] * omegas)           # (N, N, M)
    # Per-head mix, zero for invalid pairs, softplus for positivity
    mixed = softplus(features @ radial_weight.T)              # (N, N, H)
    mixed = where(valid_pairs, mixed, 0) + 1e-6
    # Code rescales so the mean over all valid pairs in the batch is 1
    mixed = mixed * count(valid_pairs) / sum(mixed, over=valid_pairs)
    # Additive attention-logit bias (paper: linear mix; code: log of mix)
    return log(mixed)                                         # (N, N, H)

# Energy and forces  [§3.3] [§3.4]
# [code: src/fairchem/core/models/allscaip/AllScAIP.py#L371-L389, #L410-L484]

def predict_energy_and_forces(atoms):
    # Per-atom energy from the final node features, masked for padding  [§3.3]
    per_atom_energy = energy_ffn(final_node_features(atoms))
    # Sum the per-atom values within each system (code default: sum)
    energy = scatter_sum(per_atom_energy, by=system_index)
    # Energy-conserving variant: forces are the negative gradient  [§3.4]
    forces = -gradient(energy.sum(), positions)
    return energy, forces

Hyperparameters

Model sizes tab. 10 (paper). The per-head size is hidden_size / heads in every row.

Shared settings from tab. 10 (paper): activation GELU; RMSNorm; node direction expansion 10; edge direction expansion 6; edge distance expansion 512; 32 radial frequencies (attn_num_freq); FFN multiplier 2; output multiplier 2; max elements 110; max batch size 96; kNN soft with sigmoid scale 0.2 and LSE scale 0.1; Gaussian distance function; envelope on; dropout 0.0. Attention kernel: memory-efficient md_NeNo.yaml.

Training settings (OMol25; direct-force, medium):

Medium model used in the headline results: the released config disables LAE and ERoPE and keeps neighbourhood and node attention [code: configs/allscaip/backbone/md_NeNo.yaml#L13-L15]. This matches the NeNo row of tab. 2 (see Results).

Results

Author-reported. Setup: OMol25 (Open Molecules 2025), validation split by chemical domain §5.3 tab. 4. Units: energy MAE in meV/atom, force MAE in meV/Å.

Component ablation §4 tab. 2. Medium 85M and small 35M models, OMol25 4M (80 epochs) and 102M (10 epochs). Totals across the four domains, converted from the source macros (×1000 to meV). Throughput is ns/day on one H200 with 1000 atoms, graph generation off tab. 2 caption.

Throughput is the same number in both data blocks for the same architecture tab. 2.

Symmetry and conservation checks §5.2 tab. 3. PBC and non-PBC energy additivity errors, force cosine under random rotation, and NVE energy drift on MD22 (meV/atom/ps).

OMol25 validation by domain §5.3 tab. 4 (selected rows, E / F).

Bold marks the best total energy in each block. The paper states the direct force medium model has the best energy overall and remains competitive in force §5.3. The 4M conservative fine-tune is the best overall energy on that split §5.3.

OMol25 evaluation suite §5.3 tab. 5. Energy MAE (meV) per category, models sorted by average rank (lower is better). Selected rows:

Column key: Pocket = ligand pocket interaction energy; Strain = ligand strain energy; Conf. = conformers; Prot. = protonation; Dist. = distance scaling (LR test); IE/EA; Spin = spin gap. All values are energy MAE in meV.

Distance-scaling (LR) test: the author-reported claim is about 90% lower error than the second-best model §5.3. See Paper vs code for why that number is hard to reproduce from this table.

Other OMol25 results. Long-range binning by system diameter and charge shows AllScAIP's error stays low while eSEN and UMA grow app. A fig. 10. Density and enthalpy of vaporisation from NPT and NVT runs on 39 liquids are plotted, not tabulated; the text reports lower MAE and higher R² than MACELES and eSEN §5.3 fig. 7.

Efficiency §5.1 fig. 4 fig. 8. Speed falls roughly as 1/N once O(N²) node attention dominates. Profiling puts the crossover earlier for larger models app. A. The authors target 10³–10⁵ atoms.

OC20 and OMat24 app. B tab. 8. Total-energy MAE on OC20 validation (ID / OOD-Both) and OMat24 per-atom energy. Units are not given in the table.

SPICE app. C tab. 9. Small conservative AllScAIP (34M per the appendix) trained on SPICE with MACE-OFF splits. Energy / force MAE, meV and meV/Å, on seven test subsets. Selected:

Solvated amino acids forces are better for eSEN (3.68 vs 4.86); the table does not match a blanket "lowest across splits" claim for forces. See Paper vs code.

Model configs app. E tab. 10: see Hyperparameters.

Paper vs code

Code: official repository linked on the paper's cover at facebookresearch/fairchem, commit d9716429092b4339dded8fd8e64fd3fdec261204, MIT licence. Everything below was checked against the medium and small AllScAIP files under src/fairchem/core/models/allscaip/ and configs/allscaip/.

• ERoPE form. Paper: a per-head linear mix of sinc features, added to the logit §3.2. Code: a softplus of the mix, masked, normalised over all valid pairs in the batch, then a log before adding to the logit [code: src/fairchem/core/models/allscaip/modules/node_attention.py#L114-L131].
• ERoPE frequency range. Paper: M log-spaced frequencies in an unspecified range §3.2. Code: 32 frequencies between π/(4·30) and π/0.25 [code: src/fairchem/core/models/allscaip/utils/data_preprocess.py#L250-L256].
• ERoPE initial weights. Code: 1/M per head [code: src/fairchem/core/models/allscaip/modules/node_attention.py#L46-L50].
• LAE in both passes. Paper describes LAE on neighbourhood attention in general §3.2. Code applies it to both the source and destination passes neighborhood_attention.py.
• LAE scaling. Code divides the spherical harmonics by sqrt(2l+1) before expansion; the extra sqrt rescaling of q and k is commented out [code: src/fairchem/core/models/allscaip/utils/data_preprocess.py#L129] neighborhood_attention.py.
• Frequency list. Paper: constraint only §3.2. Code default [20, 10, 4, 10, 20] with head-dim check configs.py.
• Headline medium config. Released medium config disables LAE and ERoPE [code: configs/allscaip/backbone/md_NeNo.yaml#L14-L15]. The paper's NeNo ablation row gives the same totals as the headline 102M model in tab. 4 (0.641 vs 0.64 meV/atom; 5.237 vs 5.24 meV/Å). This match is an inference from the numbers, not a statement by the paper.
• Energy reduction. Code default is sum over atoms [code: src/fairchem/core/models/allscaip/configs.py#L71]. Paper does not state it.
• Epoch count, 102M. Figure 5's caption says 12 epochs fig. 5; the text and the ablation caption say 10 §5.3 tab. 2; the code config uses 10 [code: configs/allscaip/omol_all.yml#L35].
• Precision. Paper: full fp32 on V100 and H200 §5.3 app. C. Code config: TF32 on omol_all.yml.
• Loss and optimiser. Not in the paper. Code: per-atom MAE energy loss with weight 10, L2-norm force loss with weight 10, AdamW lr 8e-4, wd 1e-3, cosine schedule omol_direct.yaml omol_all.yml.
• Energy-conservation kNN. Paper cites a differentiable kNN construction §3.4. Code: soft kNN ranks with sigmoid scale 0.2 and LSE scale 0.1 [code: configs/allscaip/backbone/md_NeNo.yaml#L24-L26], built in src/fairchem/core/models/allscaip/utils/allscaip_radius_graph.py.
• SPICE and forces. Paper text: lowest errors "across splits" app. C. The table has eSEN better on solvated amino-acid forces tab. 9.
• Distance-scaling claim. Paper: about 90% reduction vs second-best §5.3. In tab. 5 the second-best distance-scaling entry is AllScAIP-sm-ft-cons. at 15.473 meV, a reduction of about 35%. The 90% figure matches the eSEN-md-d. entry (109.767) instead. Inference from the table, not a claim by the paper.
• Per-atom metric bug. The paper footnotes a known FAIRChem per-atom MAE bug and gives total-energy tables §5.1 footnote tab. 6.
• No discrepancies found: none of the layer-count, width or head settings in tab. 10 differ from the configs.

Pitfalls

• ERoPE naming. Despite the name, ERoPE does not rotate queries and keys. It adds a distance-only bias to each head's logits. Implement it as a bias, not as RoPE.
• Batch-dependent ERoPE (inferred from code). The rescaling in ERoPE divides by the sum over all valid pairs in the whole padded batch, not per system. A system's logits can therefore change with the other systems in its batch. The encodings are only used by the small configs, not the medium default node_attention.py.
• Log-space bias. The code takes the log of a softplus of the mix. Use the code's form if you want to reproduce released behaviour, not the paper's linear form §3.2.
• Unstated ranges. The ERoPE frequency range, the LAE rep counts and the neighbour count k are not in the paper. The code values above are defaults or config values, not paper values.
• Quadratic cost. Global attention costs O(N²) per system. Memory and time grow with the largest system, not just the total atom count §5.1 fig. 4. Plan batching around that.
• Compilation needs fixed shapes. The code disables compilation when max_atoms is unset, and the model's padding depends on it [code: src/fairchem/core/models/allscaip/AllScAIP.py#L131-L155].
• Conservative training is slower. The paper notes the gradient-based energy path is slower than direct force because the differentiable kNN is not yet optimised fig. 5.
• Hidden precision choice. The training config enables TF32 while the paper says fp32. Results may differ slightly; check which you run.
• Table and figure numbers here were counted from the LaTeX source in order, not checked against the rendered PDF.
• Blank table cells. The OC20 numbers for the 4M medium fine-tuned conservative row are blank in tab. 8.
• Untuned OC20 and OMat24 runs. The authors say they did not tune these runs much app. B, so the comparison with UMA is not a tuned result.

Lineage

• Builds on: EScAIP for the neighbourhood attention stage §3.1; Euclidean rotary-style radial encoding from Frank et al. (2024) §3.2; differentiable kNN from Liu et al. (2026) for energy conservation §3.4.
• Compared against: UMA (Wood et al., 2025) §5.2 tab. 4 tab. 5; eSEN (Fu et al., 2025) tab. 4 tab. 5; MACE and MACELES (Kim et al., 2025) app. C fig. 7; GemNet-OC tab. 4 tab. 5.
• Superseded by: not known from the content of this wiki.