Atomistic Foundation Models

• Atomistic foundation models are large, pre-trained geometric ML models that learn transferable representations of atomic structures and potential energy surfaces.
• They leverage advanced architectures such as graph neural networks, equivariant networks, and transformers to predict energies, forces, and other observables accurately.
• AFMs enable efficient transfer learning by fine-tuning on small, task-specific datasets, reducing data requirements and enhancing simulation fidelity across various applications.

Atomistic foundation models (AFMs) are large, pre-trained geometric machine-learning models and machine-learned interatomic potentials trained on diverse, large-scale atomistic datasets to learn general, transferable representations of atomic structures and of the Born–Oppenheimer potential energy surface, together with derivatives such as forces and stresses. In current literature, the term covers graph neural networks, equivariant neural networks, and transformer-like architectures that are intended to be reused across downstream tasks through fine-tuning, new output heads, or distillation, rather than retrained from scratch for each material, molecule, or simulation regime (Kong et al., 14 Apr 2025, Yuan et al., 13 Mar 2025).

1. Definition and conceptual scope

AFMs are defined in one widely used formulation as models “pre-trained on diverse, large-scale atomistic datasets” that learn “general, transferable representations of atomic structures” and can then be fine-tuned on “small, task-specific datasets” for downstream materials problems (Kong et al., 14 Apr 2025). Closely related terms in the literature include foundation machine-learning potentials, universal machine learning potentials, and universal MLIPs; in all cases the central idea is a reusable atomistic backbone trained once at large scale and adapted many times (Póta et al., 2024, Anam et al., 3 Sep 2025).

This distinguishes AFMs from conventional task-specific ML potentials. Traditional models such as SchNet- or CGCNN-style systems are typically trained from scratch on one dataset or one material family, and their performance degrades sharply when training data are scarce. AFMs, by contrast, are meant to carry priors learned from millions to hundreds of millions of structures, spanning many elements, structures, and thermodynamic conditions, so that downstream adaptation does not need to relearn basic chemistry or local geometry (Kong et al., 14 Apr 2025, Yuan et al., 13 Mar 2025).

The scope of the term is still under active negotiation. One perspective distinguishes current “universal PBE potentials” from fully realized foundation models, reserving the latter for models that also exhibit scaling laws, multi-task or multi-domain pretraining, and broad fine-tuning behavior across many downstream observables (Yuan et al., 13 Mar 2025). The acronym is also domain-sensitive: in smart agriculture literature, “AFMs” denotes Agriculture Foundation Models, not atomistic foundation models, so disambiguation is required in cross-disciplinary contexts (Li et al., 2023).

2. Geometric and architectural foundations

AFMs arise from geometric deep learning. Their inputs are atomic identities and structures of molecules or crystals; these are encoded as node and edge attributes of a graph, processed by message passing or transformer-like attention, and mapped to latent atom-level and system-level embeddings that support predictions of energies, forces, stresses, band gaps, elastic properties, and related observables (Kong et al., 14 Apr 2025). A standard energy-based formulation writes the learned potential as $E(\{\mathbf{R}_i\})$, with forces given by

$$\mathbf{F}_i = -\nabla_{\mathbf{R}_i} E(\{\mathbf{R}_j\}).$$

Physical symmetry is a defining architectural constraint. AFMs are trained to respect translational invariance and rotational equivariance or invariance. For a rotationally equivariant model predicting vector quantities $f(\mathbf{X})$, one expects

$f(R\mathbf{X} + \mathbf{t}) = R\,f(\mathbf{X}),$

while scalar outputs such as energy satisfy $f(R\mathbf{X} + \mathbf{t}) = f(\mathbf{X})$ (Kong et al., 14 Apr 2025). In practice, the model families used under the AFM label include message-passing GNNs with many-body interactions, E(3)-equivariant neural networks, and graph transformers with higher-order irreducible-representation channels (Yuan et al., 13 Mar 2025).

Different architectural traditions implement these constraints differently. Tensor-product equivariant models such as MACE organize intermediate features into irreducible representation channels indexed by angular momentum $L$; a later probing study found that dipole information concentrates in $L=1$ channels while HOMO–LUMO-gap information concentrates in $L=0$ channels, indicating property-specific routing through the equivariant basis (Steier, 3 Mar 2026). AlphaNet follows a different route, constructing equivariant local frames $\mathcal{F}_{ij}$ for atom pairs, scalarizing geometric information in those frames, and tensorizing it back into global coordinates, thereby enforcing equivariance without explicit spherical-tensor message passing (Yin et al., 13 Jan 2025).

3. Pretraining regimes and representative model families

AFM pretraining is usually supervised on quantum-mechanical targets such as energies $E$, forces $$\mathbf{F}_i = -\nabla_{\mathbf{R}_i} E(\{\mathbf{R}_j\}).$$0, and stresses $$\mathbf{F}_i = -\nabla_{\mathbf{R}_i} E(\{\mathbf{R}_j\}).$$1, often with multi-task losses of the form

$$\mathbf{F}_i = -\nabla_{\mathbf{R}_i} E(\{\mathbf{R}_j\}).$$2

or, more generally, $$\mathbf{F}_i = -\nabla_{\mathbf{R}_i} E(\{\mathbf{R}_j\}).$$3 (Kong et al., 14 Apr 2025). The literature also includes self-supervised or unsupervised objectives such as denoising and contrastive learning, as well as synthetic-data and multi-fidelity strategies intended to expand coverage without recomputing all data at one level of theory (Kong et al., 14 Apr 2025, Yuan et al., 13 Mar 2025).

Representative AFMs already span several architectural families and data regimes. MatterTune currently supports JMP, MatterSim-v1, ORB, and EquiformerV2. Within that set, JMP-S has about 30M parameters, JMP-L about 235M, and the family is reported as trained on about 120M structures across molecules and materials; MatterSim-v1 is trained on about 17M structures; ORB-v1 is listed at about 25.2M parameters and about 32.1M structures with denoising plus energy, force, and stress objectives; and EquiformerV2 appears in variants ranging from about 31.2M parameters on about 1.58M structures to about 86.6M parameters on about 102M structures (Kong et al., 14 Apr 2025). A thermal-conductivity benchmark treats M3GNet, CHGNet, MACE-MP-0, SevenNet, and ORB-MPtraj as representative “mp-fMLPs” trained on Materials Project-scale data (Póta et al., 2024).

A parallel line of work addresses data integration. TEA, or Total Energy Alignment, was introduced to combine heterogeneous molecular and crystalline datasets computed with different functionals, basis sets, and treatments of core electrons; the resulting MACE-Osaka24 model is described as the first open-source neural network potential trained on a unified dataset covering both molecular and crystalline systems (Shiota et al., 2024). Allegro-FM extends the same multi-domain logic to an E(3)-equivariant, strictly local architecture trained on TEA-merged SPICE and MPtrj data and reported to cover 89 elements (Nomura et al., 9 Feb 2025).

4. Transfer learning and scientific applications

The immediate motivation for AFMs is data scarcity. One paper states that “a growing numbers of studies have shown atomistic FMs can improve accuracies of GNNs significantly over models trained from scratch … as well as reduce data requirements by an order of magnitude or more” (Kong et al., 14 Apr 2025). MatterTune provides a concrete downstream picture: in few-shot liquid-water fine-tuning on a 1000-structure QMD dataset, MatterSim-V1-1M achieved energy MAE of 1.21 meV/atom and force MAE of 38.37 meV/Å with 900 samples, and 1.20 meV/atom and 40.65 meV/Å with only 30 unique training structures; in subsequent MD, only MatterSim and EqV2 yielded RDFs consistent with experimental data, showing both the promise of few-shot adaptation and the inadequacy of force MAE alone as a simulation criterion. On MatBench property prediction, the best AFM fine-tuned with uniform hyperparameters outperformed the best models trained from scratch; for MP Gap, the best leaderboard MAE was 0.156 eV, while ORB-V2, EqV2-31M-mp, and JMP-S achieved 0.093, 0.098, and 0.119 eV on fold0, respectively (Kong et al., 14 Apr 2025).

AFMs also support more physics-heavy workflows. A conductivity framework combines AFMs with the Wigner Transport Equation, using the AFM as a drop-in replacement for DFT in the calculation of harmonic and anharmonic force constants. In a benchmark over 103 binary crystals, the best model, MACE-MP-0, achieved zero-shot $$\mathbf{F}_i = -\nabla_{\mathbf{R}_i} E(\{\mathbf{R}_j\}).$$4 within a factor of 2 of DFT for about 69% of materials, and for LiBr fine-tuning on only 3 DFT frames reduced $$\mathbf{F}_i = -\nabla_{\mathbf{R}_i} E(\{\mathbf{R}_j\}).$$5 from about 47% to about 2% while bringing phonon bands, linewidths, and conductivity into close agreement with DFT and experiment (Póta et al., 2024).

The representational use of AFMs extends beyond classical materials informatics. In protein modeling, intermediate features from MACE, Egret, OrbNet-v2, and AIMNet2 were used as descriptors of local protein environments. Those AFM-derived embeddings were reported to capture secondary motifs, amino-acid identity, and protonation state; to support pK$$\mathbf{F}_i = -\nabla_{\mathbf{R}_i} E(\{\mathbf{R}_j\}).$$6 prediction; and to enable a physics-informed NMR chemical-shift predictor with state-of-the-art accuracy for most backbone and side-chain atom types (Bojan et al., 29 May 2025).

5. Representation structure, uncertainty, and distillation

AFMs are increasingly studied as representation learners in their own right. Composition Projection Decomposition (CPD) removes linear composition signal from learned representations and probes the residual geometry subspace. Across eight models on QM9 and Materials Project, this produced a marked disentanglement gradient: for HOMO–LUMO gap, MACE achieved $$\mathbf{F}_i = -\nabla_{\mathbf{R}_i} E(\{\mathbf{R}_j\}).$$7 after composition removal, while ANI-2x yielded $$\mathbf{F}_i = -\nabla_{\mathbf{R}_i} E(\{\mathbf{R}_j\}).$$8 under Ridge but recovered to $$\mathbf{F}_i = -\nabla_{\mathbf{R}_i} E(\{\mathbf{R}_j\}).$$9 under an MLP probe. The same study showed that MACE routes dipole information through $f(\mathbf{X})$0 channels and HOMO–LUMO-gap information through $f(\mathbf{X})$1, whereas the same pattern was not observed for ViSNet’s vector-scalar architecture (Steier, 3 Mar 2026).

Uncertainty estimation has become a second major theme. A multi-head committee built on MACE attaches several lightweight output heads to a fixed AFM backbone and uses their disagreement as an uncertainty proxy. Applied to MACE-MP-0b, that strategy trained only about 15,000 new parameters, or about 0.2% of the model size; condensed the foundation-model training set to about 5% of its original size; and provided force-uncertainty estimates that correlated strongly with true errors across in-distribution and out-of-distribution datasets (Beck et al., 13 Aug 2025). A complementary heterogeneous-ensemble method defines a universal uncertainty metric $f(\mathbf{X})$2 from pretrained uMLIPs; on the OMat24 test set it reached Spearman correlation $f(\mathbf{X})$3 with true configuration-level force errors, and the authors proposed $f(\mathbf{X})$4 eV/Å as a practical cutoff below which AFM predictions are generally safe for deployment (Liu et al., 28 Jul 2025).

Compression follows naturally from those developments. Distillation of AFMs through synthetic data was shown to transfer knowledge from large graph-network teachers into smaller graph models and into ACE, yielding speedups of more than $f(\mathbf{X})$5 and more than $f(\mathbf{X})$6, respectively, across domains ranging from liquid water and metallic hydrogen to porous silica, hybrid halide perovskites, and an explicit-solvent organic S$f(\mathbf{X})$7 reaction (Gardner et al., 12 Jun 2025).

6. Infrastructure, scaling, and open problems

The AFM ecosystem is now supported by dedicated software and scaling strategies. MatterTune provides modular abstractions for backbones, properties, data, trainers, and applications; it integrates with PyTorch Lightning for distributed and customizable fine-tuning and with ASE through MatterTuneCalculator and MatterTunePropertyPredictor for MD, geometry optimization, and batch property screening (Kong et al., 14 Apr 2025). At the large-scale end, Allegro-FM reports parallel efficiency of 0.964 on the Aurora exaflop/s supercomputer and scaling to multi-billion-atom systems, illustrating one route to exascale AFM deployment (Nomura et al., 9 Feb 2025). A complementary route is to prune pretrained AFM backbones by removing low-contribution message-passing layers and then distribute the resulting shallow graph models with halo-based partitioning; that workflow, implemented in MatterTune, was reported to support million-atom simulations, nanosecond timescales, and single-GPU simulations of up to 5.18 million atoms while retaining AFM-level accuracy after fine-tuning (Kong et al., 25 Sep 2025).

The remaining difficulties are not merely engineering details. Fine-tuning itself is optimizer-sensitive: across molecular, crystalline, and liquid regimes, AdamW and ScheduleFree were found to deliver better curvature conditioning and force accuracy than Adam or SGD, and a short second-order L-BFGS refinement stage improved phonon spectra and interfacial dynamics without increasing inference cost (Liu et al., 5 Dec 2025). Physical consistency also remains uneven across model families. MatterTune’s water benchmark shows that low force MAE does not guarantee acceptable RDFs when force prediction is not energy conserving (Kong et al., 14 Apr 2025). In phonon benchmarking over 2,429 crystalline materials, MACE and CHGNet achieved force-prediction accuracy comparable to EquiformerV2, yet their IFC fitting errors produced poor lattice-thermal-conductivity predictions; by contrast, fine-tuned EquiformerV2 consistently led on second-order IFCs, LTC, and related phonon properties, demonstrating that force accuracy alone is not a sufficient AFM evaluation metric for vibrational physics (Anam et al., 3 Sep 2025).

Current research therefore treats AFMs simultaneously as reusable atomistic backbones, simulation engines, and representation models, while leaving several questions open: how to incorporate long-range physics beyond finite cutoffs, how to benchmark out-of-distribution behavior without relying on error cancellation, how to automate fine-tuning without eroding physical fidelity, and how to decide when a universal potential has become a fully realized foundation model rather than a strong domain-general precursor (Yuan et al., 13 Mar 2025).