Machine-learning interatomic potential

Overview

A machine-learning interatomic potential (MLIP) takes the atoms of a system and returns their energy and the force on each atom. A simulation can then advance in time using the learned model instead of a quantum calculation at every step 2603.06567 §1. The hard part is range. A model where each atom sees only neighbours inside a distance cutoff cannot carry information across long distances in one step. Explicit long-range physics terms can fill that gap, but they add modelling and engineering work 2603.06567 §1.

The one paper covered here tests a different route: a global attention step over all atoms of a system, combined with a local neighbourhood step, with an option to make forces consistent with energy 2603.06567 §3.1 [arxiv:2603.06567 §3.3]. Per-atom energies are reduced to one energy per system; the paper does not state the reduction. Forces come either from a separate output head or from the negative gradient of the energy 2603.06567 §3.3 [arxiv:2603.06567 §3.4].

Variants

Local attention plus global all-to-all attention

Each block runs two attention steps. Local attention runs over each atom's fixed stencil, which holds the atom itself and up to k neighbours, in two directions. A smooth distance envelope fades far pairs. Global node attention packs all atoms of one system into one sequence and applies standard multi-head self-attention, so any two atoms can interact in one step. The local step costs O(Nk) and the global step O(N²) per system 2603.06567 §3.1. Card: AllScAIP.

Optional geometric encodings: LAE and ERoPE

Two optional encodings bias attention with geometry. LAE applies to local attention: real spherical harmonics of the unit edge direction, expanded by degree, modulate queries and keys before the dot product 2603.06567 §3.2. ERoPE applies to global attention: a bank of log-spaced sinc features of the pair distance, mixed per head, is added to the attention logit as a bias 2603.06567 §3.2. Despite the name, ERoPE does not rotate queries and keys. The paper's own ablations show the encodings cost throughput 2603.06567 tab. 2, and the released code differs from the written form of ERoPE in several ways (see Paper vs code below).

Direct forces or energy-conserving forces

A direct model predicts forces with its own head. A conservative model takes the negative gradient of the predicted energy, so forces and energy agree by construction 2603.06567 §3.3 2603.06567 §3.4. The paper reports both on OMol25. On the 102M split the direct medium model gets 0.64 meV/atom energy and 5.24 meV/Å force, and the conservative one gets 0.67 and 5.78 2603.06567 tab. 4. Conservative training is slower, because the gradient path through the differentiable neighbour construction is not yet optimised 2603.06567 fig. 5.

When to use which

Implementation sketch

The blocks and the output head, in the order the paper describes them. The encodings are left out for clarity.

# One block of a local plus global attention MLIP  [arxiv:2603.06567 §3.3]

def block(nodes, edges, system_index):
    # Each atom attends over its stencil: itself plus up to k neighbours  [§3.1]
    stencil = local_attention(nodes, edges)
    # The self slot becomes the atom's node feature  [§3.3]
    node = stencil[:, 0]
    # Feed-forward on the neighbour slots gives the edge features  [§3.3]
    edge = edge_ffn(stencil[:, 1:])
    # Every atom attends to every other atom of its own system  [§3.1]
    node = global_attention(node, system_index)
    node = node_ffn(node)
    # Put node features back in slot 0 and join the edge features  [§3.3]
    return concat(node[:, None], edge)

def energy_and_forces(atoms, conservative):
    # Run the blocks, then read a per-atom energy  [§3.3]
    nodes, edges = embed(atoms)
    for _ in range(num_blocks):
        stencil = block(nodes, edges, atoms.system_index)
        nodes, edges = stencil[:, 0], stencil[:, 1:]
    per_atom = energy_head(nodes)
    # Sum per system; the paper does not state the reduction  [§3.3]
    energy = scatter_sum(per_atom, by=atoms.system_index)
    if conservative:
        # Forces are the negative gradient of energy  [§3.4]
        return energy, -grad(energy.sum(), atoms.positions)
    # Otherwise a separate head predicts forces  [§3.3]
    return energy, force_head(nodes)

Contradictions

• Encodings and scale. The paper's text says the benefit of LAE and ERoPE shrinks as data and model size grow 2603.06567 §4. The ablation is mixed. At 35M on 102M samples, removing both raises energy error from 0.849 to 0.971 meV/atom 2603.06567 tab. 2. At 85M on 102M samples, removing both lowers energy error 2603.06567 tab. 2. Both sides are cited; this page does not resolve them.
• Distance-scaling claim. The paper says its model has about 90% lower error than the second-best model on the distance-scaling test [arxiv:2603.06567 §5.3]. The table entries behind that test give a smaller reduction, and the 90% figure matches a different row 2603.06567 tab. 5. Read the table before quoting the claim.
• Accuracy claims. The paper says the direct medium model has the best overall energy and is competitive on force 2603.06567 §5.3. On the same split, the eSEN model has a lower total force error, 4.76 against 5.24 meV/Å 2603.06567 tab. 4. On SPICE, the paper says its model is lowest across splits 2603.06567 app. C, but eSEN has lower force error on the solvated amino-acid split 2603.06567 tab. 9.
• Epoch count. The 102M run is 12 epochs in one figure caption and 10 in the text and in the ablation caption 2603.06567 fig. 5 [arxiv:2603.06567 §5.3] 2603.06567 tab. 2.

Pitfalls

• Precision. The paper says training runs in full fp32 [arxiv:2603.06567 §5.3]. The released training config enables TF32, so results may differ slightly from the paper's numbers.
• Batch-dependent ERoPE. The released code's normalisation of the ERoPE bias runs over all valid pairs in a padded batch, not per system. One system's logits can then change with the other systems in its batch. The paper does not describe this 2603.06567 §3.2.
• Unstated settings. The paper leaves out the ERoPE frequency range, the LAE repetition counts and the neighbour count k. Values from the released config are defaults, not paper values 2603.06567 §3.2.
• Quadratic global cost. Plan batches around the largest system, because global attention cost follows the largest system, not the total atom count 2603.06567 §5.1 2603.06567 fig. 4.

Paper vs code

• ERoPE form. The paper gives a linear mix of sinc features added to the logit 2603.06567 §3.2. The released code applies softplus, masking and a batch-wide normalisation, then a log before adding the bias.
• LAE in both passes. The paper describes LAE on local attention in general 2603.06567 §3.2. The released code applies it to both the source and destination passes.
• Headline model. The released medium config turns LAE and ERoPE off. Its numbers match the no-encoding row of the ablation 2603.06567 tab. 2 2603.06567 tab. 4. This match is inferred from the numbers, not stated by the paper.

Papers

• AllScAIP: local plus global all-to-all attention, optional LAE and ERoPE encodings, energy-conserving forces, and OMol25 results.