Mixture of Experts Framework in Machine Learning Interatomic Potentials for Atomistic Simulations
This work addresses the inference cost bottleneck of MLIPs for large-scale atomistic simulations, offering a practical speedup for materials science applications.
The authors introduce a multifidelity Mixture-of-Experts framework for machine learning interatomic potentials that spatially partitions simulation domains, assigning high-fidelity models to chemically complex regions and low-fidelity models to simple regions. They address mechanical mismatch at interfaces with a co-training strategy, achieving predictive accuracy comparable to full high-fidelity simulations at more than twice the computational speed on a Pt+CO catalytic system.
First-principles atomistic simulations are essential for understanding complex material phenomena but are fundamentally limited by their computational cost. While Machine Learning Interatomic Potentials (MLIPs) have drastically improved cost for a given accuracy, their inference cost remains a bottleneck for massive systems or long timescales. To address this, we introduce a multifidelity "Mixture-of-Experts" framework based on the E(3)-equivariant Allegro architecture. Our method spatially partitions the simulation domain into a chemically complex region (e.g., reactive interfaces) and a simple region (e.g., bulk lattice), assigning models of varying capacity to each. Among the challenges in such static domain decomposition, the mechanical mismatch between models at the interface is particularly critical, as it can generate artificial stress fields and instability. We address this challenge with a co-training strategy in which the loss function includes agreement constraints -- penalties on per-atom energy and force discrepancies between models evaluated on shared bulk environments -- forcing the independent models to learn a consistent physical description of the bulk material. We validate this approach on a realistic Pt+CO catalytic system, demonstrating that the co-trained models maintain exact energy conservation, align their bulk mechanical response (e.g., equation of state and bulk modulus), and achieve predictive accuracy comparable to a full high-fidelity simulation at more than twice the computational speed.