Breaking the Training Barrier of Billion-Parameter Universal Machine Learning Interatomic Potentials

2026/04/17 by Yuanchang Zhou, Hongyu Wang, Yiming Du +12 · 2 voices
#cs.DC #cs.LG

paper · pdf

Abstract

Universal Machine Learning Interatomic Potentials (uMLIPs), pre-trained on massively diverse datasets encompassing inorganic materials and organic molecules across the entire periodic table, serve as foundational models for quantum-accurate physical simulations. However, uMLIP training requires second-order derivatives, which lack corresponding parallel training frameworks; moreover, scaling to the billion-parameter regime causes explosive growth in computation and communication overhead, making its training a tremendous challenge. We introduce MatRIS-MoE, a billion-parameter Mixture-of-Experts model built upon invariant architecture, and Janus, a pioneering high-dimensional distributed training framework for uMLIPs with hardware-aware optimizations. Deployed across two Exascale supercomputers, our code attains a peak performance of 1.2/1.0 EFLOPS (24%/35.5% of theoretical peak) in single precision at over 90% parallel efficiency, compressing the training of billion-parameter uMLIPs from weeks to hours. This work establishes a new high-water mark for AI-for-Science (AI4S) foundation models at Exascale and provides essential infrastructure for rapid scientific discovery.

Citations

• The Open Polymers 2026 (OPoly26) Dataset and Evaluations
• Cross Learning between Electronic Structure Theories for Unifying Molecular, Surface, and Inorganic Crystal Foundation Force Fields
• Optimizing Cross-Domain Transfer for Universal Machine Learning Interatomic Potentials
• The Open Catalyst 2025 (OC25) Dataset and Models for Solid-Liquid Interfaces
• The Open DAC 2025 Dataset for Sorbent Discovery in Direct Air Capture
• MOFSimBench: Evaluating Universal Machine Learning Interatomic Potentials In Metal--Organic Framework Molecular Modeling
• UMA: A Family of Universal Models for Atoms
• A Graph Neural Network for the Era of Large Atomistic Models
• The Open Molecules 2025 (OMol25) Dataset, Evaluations, and Models
• Orb-v3: atomistic simulation at scale
• A Foundational Potential Energy Surface Dataset for Materials
• Open Materials 2024 (OMat24) Inorganic Materials Dataset and Models
• CatTSunami: Accelerating Transition State Energy Calculations with Pre-trained Graph Neural Networks
• EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree Representations
• DMC-ICE13: ambient and high pressure polymorphs of ice from Diffusion Monte Carlo and Density Functional Theory
• MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields
• Towards Training Billion Parameter Graph Neural Networks for Atomic Simulations
• Pushing the limit of molecular dynamics with ab initio accuracy to 100 million atoms with machine learning

Discussions

• So the CPU is only properly described in 2 ML papers (arxiv.org/pdf/2604.158..., arxiv.org/pdf/2605.086...) I found no official specs Its an ARMv9 based custom CPU design with SVE/SME support. 2 socke [bsky, 7 points, 1 comments]
• => arXiv, Apr 17, 2026 arxiv.org/abs/2604.15821 🇨🇳 Two distinct Exascale China New-gen Intelligent Supercomputer (CNIS) CISC+GPU & LineShine (灵晟) 20,480x 2x [2x (152x ARMv9 cores+128 GB DDR), 8x HBM [bsky, 0 points, 0 comments]

Related

• Global universal scaling and ultrasmall parameterization in machine-learning interatomic potentials with superlinearity
• Global Universal Scaling and Ultra-Small Parameterization in Machine Learning Interatomic Potentials with Super-Linearity
• Automated Machine Learning Pipeline: Large Language Models-Assisted Automated Dataset Generation for Training Machine-Learned Interatomic Potentials
• Taming Multi-Domain, -Fidelity Data: Towards Foundation Models for Atomistic Scale Simulations
• AI2Pot: A scalable and unified framework for machine-learning interatomic potential development and large-scale molecular dynamic simulations