Mixture of Experts Framework in Machine Learning Interatomic Potentials for Atomistic Simulations

TL;DR

This paper introduces a multifidelity mixture-of-experts framework for machine learning interatomic potentials, partitioning simulation domains to improve efficiency while maintaining accuracy and physical consistency.

Contribution

It proposes a novel co-training strategy with agreement constraints to address interface issues in static domain decomposition for atomistic simulations.

Findings

Maintains energy conservation and mechanical response accuracy.

Achieves over twice the computational speed compared to full high-fidelity models.

Validates effectiveness on a Pt+CO catalytic system.

Abstract

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.