Orb-v3: atomistic simulation at scale

TL;DR

Contrary to recent literature, it is found that non-equivariant, non-conservative architectures can accurately model physical properties, including those which require higher-order derivatives of the potential energy surface.

Abstract

We introduce Orb-v3, the next generation of the Orb family of universal interatomic potentials. Models in this family expand the performance-speed-memory Pareto frontier, offering near SoTA performance across a range of evaluations with a >10x reduction in latency and > 8x reduction in memory. Our experiments systematically traverse this frontier, charting the trade-off induced by roto-equivariance, conservatism and graph sparsity. Contrary to recent literature, we find that non-equivariant, non-conservative architectures can accurately model physical properties, including those which require higher-order derivatives of the potential energy surface. This model release is guided by the principle that the most valuable foundation models for atomic simulation will excel on all fronts: accuracy, latency and system size scalability. The reward for doing so is a new era of computational chemistry driven by high-throughput and mesoscale all-atom simulations.

Figures (10)

• Figure 1: The Pareto frontier for a range of universal Machine Learning Interatomic Potentials. The $K_{SRME}$ metric assesses a model's ability to predict thermal conductivity via the Wigner formulation of heat transport pota2024thermal and requires accurate geometry optimizations as well as second and third order derivatives of the PES (computed via finite differences). The y-axis measure a model's forward passes per second on a dense periodic system of 1000 atoms, disregarding graph construction time, measured on a NVIDIA H200. Point sizes represent max GPU memory usage. Y-axis jitter (+/- 5 steps/second) has been applied to allow visualization of overlapping points. Model families include a range of specific models with broadly the same architecture, but may be different sizes or trained on different datasets. More details are provided in Appendix \ref{['sec:model_families']}.
• Figure 2: Speed + max GPU memory allocated on an NVIDIA H200 for the computation of energies, forces and stress. The batch size is fixed to 1, but we vary the number of atoms across the subplots. Relative times are computed with respect to the fastest model: orb-v3 Direct (20 neighbors). Times include both model inference and graph construction, with the latter marked by hatched lines. The graph construction method for Orb is a function of the number of atoms, as described in Appendix \ref{['sec: efficient graph construction']}. A key takeaway from this figure is that extreme scalability requires a confluence of i) efficient graph construction ii) Finite max neighbors iii) Non-conservative direct predictions. For the baselines, we use mace-medium-mpa-0 (v0.3.10, cuequivariance-torch v0.1.0), mattersim-v1.0.0-5m (v1.1.2), 7net-mf-ompa (v0.11.0). All models are benchmarked using PyTorch v2.6.0+cu124. Alternative libraries, like JAX, may yield further improvements for some models, but is out of scope for this work.
• Figure 3: (left) Scatter plot comparing the measured invariance (the standard deviation of the energy prediction over a randomized set of rotations) to the norm of the rotational gradient $||\Delta_\text{rot}||$, for all 103 structures in the thermal conductivity benchmark. Gray dots are obtained using Orb-v3 trained on OMat24 with the default loss function; red dots are obtained using Orb-v3 trained with equigrad regularization. (right) Thermal conductivity benchmark performance for two different methods in Phonopy; auto exploits the crystal symmetry to reduce the number of displacements to consider. For non-invariant models, this reduction is invalid, but models trained with equigrad regularization partially alleviate this difference due to increased invariance under rotation.
• Figure 4: Binned confidence predictions from Orb-v3's confidence head on on a random sample of systems from 3 datasets. MP Traj systems are sampled from the validation set; Small Molecules are systems randomly sampled from optimization trajectories of 162 commmon organic molecules from g2_ase_mols (the g2 subset, made available in ASE), and IZA are 233 relaxed zeolite structures, all optimized with VASP at the PBE level of theory. Even for out of distribution datasets, confidence bin predictions correlate well with Force MAE at the atom level.
• Figure 5: Stable simulation of the Carbonic Anhydrase enzyme II system using orb-v3-direct-inf-omat for over 700 ps. The enzyme is depicted as its amino acid representation for visual clarity, but simulations use the full all-atom representation.