Ab Initio Accuracy Beyond the Cutoff: Distributed Message-Passing Equivariant Attention at Tens-of-Millions-of-Atom Scale
Abstract
Molecular dynamics predicts molecular motion from atomic forces and trajectory integration. Density functional theory (DFT) provides first-principles forces, but its O(N3) cost limits simulations to thousands of atoms; empirical force fields scale to millions but miss bond breaking and charge transfer. Equivariant message-passing Transformers learn the ab initio potential at linear cost. Stacked layers propagate information beyond a single cutoff, giving four layers a receptive field up to 23Å. At this frontier, scale is limited by communication: layers exchange boundary features across devices, force backpropagation repeats the exchange, and moving atoms change neighbor graphs and tensor shapes. We present Tessera, a distributed runtime that fixes workload shapes between rebuilds. It sorts atoms along a space-filling curve into equal-count shards, ensuring identical atom counts and tensor shapes per rank. Peer sets scale as K0.5, not K. A conservative neighbor skin preserves topology across updates, enabling each rebuild interval to replay as one CUDA Graph. Fused kernels, three-level neighbor search, and overlapped All-to-All complete it. On 512 GPUs, Tessera sustains 30.7M atoms at 89.7% weak-scaling efficiency (77--90% across three per-GPU loads); per-step time rises 1.11--1.29x as machine and problem grow 16x. Graph reuse cuts full-step time 1.87×. Across an 8× device increase, neighbor-search cost grows 1.7×, versus 3.0× for Morton order. In a live $10.24$M-atom salt solution, NVE energy drift is 0.42 μeV/atom/ps over $10$ ps, and structure matches a serial reference. A 10.06M-atom Norovirus capsid sustains 100 ps of NVT at ≈ 8 ps/day on 256 A100 GPUs. Over 200,000 steps, it has no broken bonds or anomalously close pairs and uses three compiled variants, with no graph breaks or eager fallbacks.
Related
- Alignment-Free Optical Combining with Scalar-Message Expectation Propagation Detection for Distributed Multi-Aperture Coherent Reception
- Train for Accuracy, Execute at Scale: Architecture-Preserving Inference for Equivariant Atomistic Foundation Models
- Data analysis of millions of Type IIB flux vacua
- CTMRG as a tensor-network message-passing algorithm
- Why Multi-Layer Message Passing Works: Completeness Theory for Graph Neural Network Interatomic Potentials
BibTeX
@misc{kan2026initio,
title = {{Ab Initio Accuracy Beyond the Cutoff: Distributed Message-Passing Equivariant Attention at Tens-of-Millions-of-Atom Scale}},
author = {Allex Kan and Ji Zhang and Hao Yin and Hao Li and Owen Fu and Xunxing Chen and Lin Huang and Zion Wang and Lucas Li and Jiajun Cheng and Arthur Jiang and Junping Zhao and Jia Zhang},
howpublished = {arXiv},
year = {2026},
month = nov,
doi = {10.48550/arxiv.2611.06122},
eprint = {2611.06122},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2611.06122},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - Ab Initio Accuracy Beyond the Cutoff: Distributed Message-Passing Equivariant Attention at Tens-of-Millions-of-Atom Scale
AU - Allex Kan
AU - Ji Zhang
AU - Hao Yin
AU - Hao Li
AU - Owen Fu
AU - Xunxing Chen
AU - Lin Huang
AU - Zion Wang
AU - Lucas Li
AU - Jiajun Cheng
AU - Arthur Jiang
AU - Junping Zhao
AU - Jia Zhang
PY - 2026
DA - 2026/11/04
PB - arXiv
DO - 10.48550/arxiv.2611.06122
UR - https://arxiv.org/abs/2611.06122
AB - Molecular dynamics predicts molecular motion from atomic forces and trajectory integration. Density functional theory (DFT) provides first-principles forces, but its O(N^3) cost limits simulations to thousands of atoms; empirical force fields scale to millions but miss bond breaking and charge transfer. Equivariant message-passing Transformers learn the ab initio potential at linear cost. Stacked layers propagate information beyond a single cutoff, giving four layers a receptive field up to 23Å. At this frontier, scale is limited by communication: layers exchange boundary features across devices, force backpropagation repeats the exchange, and moving atoms change neighbor graphs and tensor shapes. We present Tessera, a distributed runtime that fixes workload shapes between rebuilds. It sorts atoms along a space-filling curve into equal-count shards, ensuring identical atom counts and tensor shapes per rank. Peer sets scale as K^0.5, not K. A conservative neighbor skin preserves topology across updates, enabling each rebuild interval to replay as one CUDA Graph. Fused kernels, three-level neighbor search, and overlapped All-to-All complete it. On 512 GPUs, Tessera sustains 30.7M atoms at 89.7% weak-scaling efficiency (77--90% across three per-GPU loads); per-step time rises 1.11--1.29x as machine and problem grow 16x. Graph reuse cuts full-step time 1.87×. Across an 8× device increase, neighbor-search cost grows 1.7×, versus 3.0× for Morton order. In a live $10.24$M-atom salt solution, NVE energy drift is 0.42 μeV/atom/ps over $10$ ps, and structure matches a serial reference. A 10.06M-atom Norovirus capsid sustains 100 ps of NVT at ≈ 8 ps/day on 256 A100 GPUs. Over 200,000 steps, it has no broken bonds or anomalously close pairs and uses three compiled variants, with no graph breaks or eager fallbacks.
ER -