Mixture of experts architectures for machine learning interatomic potentials

Yuzhi Liu; Duo Zhang; Anyang Peng; Weinan E; Linfeng Zhang; Han Wang

physics.chem-phMar 9, 2026

Mixture of experts architectures for machine learning interatomic potentials

Authors: Yuzhi Liu, Duo Zhang, Anyang Peng, Weinan E, Linfeng Zhang, Han Wang

Organizations: AI for Science Institute, Beijing 100080, P. R. China. · DP Technology, Beijing 100080, P. R. China. · Academy for Advanced Interdisciplinary Studies, Peking University, Beijing 100871, P. R. China. · School of Mathematical Sciences, Peking University, Beijing 100871, P. R. China . ·

Abstract

Machine Learning Interatomic Potentials (MLIPs) enable accurate large-scale atomistic simulations, yet improving their expressive capacity efficiently remains challenging. Here we systematically investigate Mixture-of-Experts (MoE) and Mixture-of-Linear-Experts (MoLE) architectures within the DPA3 framework for MLIPs and analyze the effects of routing strategies and expert designs. We show that sparse activation combined with shared experts yields substantial performance gains, and that nonlinear MoE formulations outperform MoLE when shared experts are present, underscoring the importance of nonlinear expert specialization. Furthermore, element-wise routing consistently surpasses configuration-level routing, while global MoE routing often leads to numerical instability. The resulting element-wise MoE model consistently outperforms all DPA3-based baselines across the OMol25, OMat24, and OC20M benchmarks. Analysis of routing patterns reveals chemically interpretable expert specialization aligned with periodic-table trends, indicating that the model effectively captures element-specific chemical characteristics for precise interatomic modeling.

Figures & tables

Figure 1: Schematic illustration of the model architectures. (a) The standard MLP framework. (b) The MoE-E framework, featuring a router that dynamically selects 3 active experts from the pool, complemented by 1 fixed shared expert to capture universal features. (c) The MoLE-E framework, where 4 expert weights are linearly combined via the router to modulate the main network weights, including 1 shared expert component.

Figure 2: Performance benchmarks on the OMol25 dataset for the MoE-E model with a total pool of

N=64

experts. (a) Impact of shared expert allocation. Normalized energy and force MAE are shown as a function of the shared expert ratio. The x-axis denotes the proportion of shared experts (e.g.,

(2/6)

indicates 2 shared experts among 6 total activated experts). Background shading demarcates different

K

activation settings for MoE-E. The horizontal dashed line provides a performance baseline corresponding to a standard 6-layer DPA3 model (approximately 1.3 M parameters). See Supplementary Table 2 for detailed values. (b) Effect of total expert count: Comparative analysis of normalized energy (left) and force (right) MAE as a function of the total number of experts. Two baselines are provided: the standard 6-layer DPA3 model (approximately 1.3 M parameters) and a “

4\times

Params” variant (approximately 5.2 M parameters) achieved by doubling the hidden dimension of the 6-layer DPA3 model. See Supplementary Table 3 for detailed values.

Model	Routing	Normalized Energy MAE			Normalized Force MAE
Model	Routing	4 Experts	16 Experts	64 Experts	4 Experts	16 Experts	64 Experts
MoE	Element-wise (E)	0.814	0.780	0.771	0.809	0.793	0.779
MoE	Global (G)	N/A	N/A	N/A	N/A	N/A	N/A
MoLE	Element-wise (E)	0.913	0.847	0.866	0.914	0.857	0.847
MoLE	Global (G)	0.981	0.946	0.952	0.965	0.936	0.950

Table 1: Normalized Energy and Force MAE values. The results compare the performance of Global (G) vs. Element-wise (E) routing across different expert counts. All MAEs are normalized against a 6-layer DPA3 baseline; thus, values represent relative error ratios where a lower score signifies superior performance relative to the baseline.

Figure 3: Performance comparison across three benchmark datasets: (a) OC20M, (b) OMol25, and (c) OMat24. For each dataset, energy and force prediction errors are reported in terms of MAE, with additional RMSE shown for OC20M. We compare baseline DPA3 models with 6-layer and 12-layer configurations against three enhanced variants: a widened dense model (“6

\times

Params”), MoLE-E, and MoE-E. The “6

\times

Params” model increases the total parameter count by a factor of six through wider hidden layers and is applied only to the 6-layer configuration, enabling comparison with MoE models under a comparable activated parameter budget (with

K=6

active experts). Both MoLE-E and MoE-E employ 64 experts with 3 shared experts, while MoE-E activates

K=6

experts per input. The 6-layer baseline, 12-layer baseline, 6

\times

Params model, MoLE-E, and MoE-E contain approximately 1.3, 2.6, 7.8, 83, and 83 million parameters, respectively; for MoE-E, approximately 7.8 million parameters are activated per forward pass.

Figure 4: Distribution of energy and force prediction errors across three benchmark datasets. Top row: energy MAE per atom. Bottom row: force MAE. Results are shown for (a) OC20M, (b) OMol25, and (c) OMat24. Compared to the dense baseline, both MoE-E and MoLE-E shift the error distributions toward lower values and significantly suppress high-error tails. MoE-E further produces more concentrated distributions, indicating improved stability and more consistent performance across diverse chemical environments. All expert-based models use 64 experts with 3 shared experts; for MoE-E,

K=6

experts are activated per input. All results correspond to 6-layer models.

Figure 5: Expert weighting distribution analysis via PCA for the MoE-E model (64 total experts,

K=6

activated experts, 3 shared experts) trained on the OMat24 dataset. (a) Overview of the general distribution for all elements. (b) Distribution patterns for alkali metals, alkaline earth metals, and noble gases. (c) Regional distribution of p-block groups except noble gas, specifically Boron (Group 13) through Halogen (Group 17) elements.(d) Spatial clustering of transition metal elements.

Machine Learning Interatomic Potentials

Machine learning interatomic potentials predict atomic energies and forces for fast simulations of molecules and materials.

See all 91 papers

Mixture of Experts Framework in Machine Learning Interatomic Potentials for Atomistic Simulations
Apr 28, 2026Gabriel de Miranda Nascimento, Marc L. Descoteaux, Laura Zichi +9Machine Learning Interatomic Potentials Multifidelity Modeling
Department of Materials Science and Engineering, Massachusetts Institute of Technology, Cambridge, MA, USA · John A. Paulson School of Engineering and Applied Sciences, Harvard University, Cambridge, MA, USA · Robert Bosch LLC Research and Technology Center, Watertown, MA, USA
Beyond Adam: SOAP and Muon for Faster, Label-Efficient Training of Machine Learning Interatomic Potentials
Jul 2, 2026Gil Harari, Yoel Zimmermann, Ola Tangen Kulseng +4Machine Learning Interatomic Potentials Neural Network Optimization
John A. Paulson School of Engineering and Applied Sciences, Harvard University, Cambridge, MA, USA · Robert Bosch LLC Research and Technology Center, Watertown, MA, USA
Breaking the Training Barrier of Billion-Parameter Universal Machine Learning Interatomic Potentials
Apr 17, 2026Yuanchang Zhou, Hongyu Wang, Yiming Du +12Machine Learning Interatomic Potentials Communication-Efficient Distributed Training
State Key Lab of Processors, Institute of Computing Technology, Chinese Academy of Sciences · Independent Researcher · Sun Yat-Sen University

Apr 28, 2026J/K move · Enter open · S save

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

Gabriel de Miranda Nascimento, Marc L. Descoteaux, Laura Zichi, Chuin Wei Tan+8

Department of Materials Science and Engineering, Massachusetts Institute of Technology, Cambridge, MA, USA · John A. Paulson School of Engineering and Applied Sciences, Harvard University, Cambridge, MA, USA · Robert Bosch LLC Research and Technology Center, Watertown, MA, USA

Machine Learning Interatomic Potentials Multifidelity Modeling Mixture-of-Experts Models Computational Materials Science

Open paper AskPDF arXiv

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.

Abstract

Machine Learning Interatomic Potentials

Explore similar work

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

Beyond Adam: SOAP and Muon for Faster, Label-Efficient Training of Machine Learning Interatomic Potentials

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