Train for Accuracy, Execute at Scale: Architecture-Preserving Inference for Equivariant Atomistic Foundation Models
Abstract
Equivariant atomistic foundation models provide broadly transferable interatomic potentials trained against quantum-mechanical reference data, but their repeated execution at simulation scale remains computationally and memory intensive. We present Symmetrix-XL, an inference engine that scales pretrained MACE checkpoints without retraining, distillation, or modification of their learned weights. It combines streamed-edge execution to avoid graph-wide materialization of expanded edge intermediates, model-specialized code generation to compile checkpoint-specific operators, and tiled execution to reuse a bounded device-memory workspace. On complete LAMMPS-step benchmarks, Symmetrix-XL reduces inference time by 3.1-5.0 times relative to ML-IAP + cuEquivariance across tested A100 and RTX 5090 workloads. On a single A100 80 GB GPU, the tested MACE-OMAT-0 capacity boundary increases from 24,565 atoms to 11.24 million atoms. This substantially lowers the hardware threshold for simulations that would otherwise require spatial decomposition across many GPUs and compute nodes. The same backend weak-scales to 703 million atoms on 64 A800 GPUs at 93.8 percent efficiency. Energy, force, stress, molecular-dynamics stability, and Matbench Discovery evaluations reproduce reference MACE behavior within measured tolerances. Case studies spanning solid-state, interfacial, and reactive systems demonstrate the increased capacity in realistic workflows. These results establish post-training execution as a scaling axis complementary to model redesign, compression, and distributed scale-out, and show that the practical accuracy-cost frontier depends on both model architecture and inference execution.
Related
- Fast and Accurate Equivariant Foundation Models for Atomistic Simulation
- Cartesian tensor equivariant machine-learning force field for spin-dependent atomistic simulations
- NEMORA: Neural Equivariant Multipole Operators for Long-Range Atomistic Learning
- Challenging Atomistic Foundation Models on Disordered Carbons
- Atom-JEPA: Joint-Embedding Predictive Architecture for 3D Atomistic Systems
BibTeX
@misc{fu2026train,
title = {{Train for Accuracy, Execute at Scale: Architecture-Preserving Inference for Equivariant Atomistic Foundation Models}},
author = {Lei Fu and Zihui Feng and Yongheng Li and Hongwei Du and Xin He and Junyi Wu and Kejie Bao and Yueyu Zhang and Zeyu Deng and Ziheng Lu and Bonan Zhu},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2610.01036},
eprint = {2610.01036},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2610.01036},
}
RIS
TY - UNPB
TI - Train for Accuracy, Execute at Scale: Architecture-Preserving Inference for Equivariant Atomistic Foundation Models
AU - Lei Fu
AU - Zihui Feng
AU - Yongheng Li
AU - Hongwei Du
AU - Xin He
AU - Junyi Wu
AU - Kejie Bao
AU - Yueyu Zhang
AU - Zeyu Deng
AU - Ziheng Lu
AU - Bonan Zhu
PY - 2026
DA - 2026/10/01
PB - arXiv
DO - 10.48550/arxiv.2610.01036
UR - https://arxiv.org/abs/2610.01036
AB - Equivariant atomistic foundation models provide broadly transferable interatomic potentials trained against quantum-mechanical reference data, but their repeated execution at simulation scale remains computationally and memory intensive. We present Symmetrix-XL, an inference engine that scales pretrained MACE checkpoints without retraining, distillation, or modification of their learned weights. It combines streamed-edge execution to avoid graph-wide materialization of expanded edge intermediates, model-specialized code generation to compile checkpoint-specific operators, and tiled execution to reuse a bounded device-memory workspace. On complete LAMMPS-step benchmarks, Symmetrix-XL reduces inference time by 3.1-5.0 times relative to ML-IAP + cuEquivariance across tested A100 and RTX 5090 workloads. On a single A100 80 GB GPU, the tested MACE-OMAT-0 capacity boundary increases from 24,565 atoms to 11.24 million atoms. This substantially lowers the hardware threshold for simulations that would otherwise require spatial decomposition across many GPUs and compute nodes. The same backend weak-scales to 703 million atoms on 64 A800 GPUs at 93.8 percent efficiency. Energy, force, stress, molecular-dynamics stability, and Matbench Discovery evaluations reproduce reference MACE behavior within measured tolerances. Case studies spanning solid-state, interfacial, and reactive systems demonstrate the increased capacity in realistic workflows. These results establish post-training execution as a scaling axis complementary to model redesign, compression, and distributed scale-out, and show that the practical accuracy-cost frontier depends on both model architecture and inference execution.
ER -