Scaling Machine Learning Interatomic Potentials with Mixtures of Experts

This paper introduces Mixture-of-Experts (MoE) and Mixture-of-Linear-Experts (MoLE) architectures for Machine Learning Interatomic Potentials, demonstrating that element-wise routing with shared nonlinear experts achieves state-of-the-art accuracy across multiple benchmarks while revealing chemically interpretable specialization aligned with periodic-table trends.

The Gist

Imagine you are trying to teach a computer to understand how atoms stick together to form everything in the universe—from water molecules to the steel in a bridge. This is the job of Machine Learning Interatomic Potentials (MLIPs).

Think of an MLIP as a super-smart "chemistry translator." It takes the positions of atoms and predicts how much energy they have and how they will move. The more accurate this translator is, the better we can simulate drug discovery, new materials, or chemical reactions.

However, there's a problem: To make these translators smarter, we usually just make the computer "brain" bigger and denser. But making a brain that is huge and dense is like trying to carry a library in your backpack—it's too heavy, too slow, and eventually, it gets too messy to learn anything new.

This paper introduces a clever new way to build these chemistry brains using something called Mixture of Experts (MoE). Here is the simple breakdown of what they did and why it works.

1. The Old Way vs. The New Way

• The Old Way (Dense Model): Imagine a single, giant chef who tries to cook every dish in the world. As the menu gets bigger, the chef gets overwhelmed, makes mistakes, and cooks very slowly.
• The New Way (Mixture of Experts): Instead of one giant chef, imagine a kitchen with a team of specialized sous-chefs.
  • There is a Head Chef (The Router) who looks at the order.
  • If the order is for a steak, the Head Chef calls the Grill Expert.
  • If the order is for a salad, the Head Chef calls the Veggie Expert.
  • Crucially, only a few experts work on any single order, even though the whole team is huge. This makes the kitchen fast and efficient.

2. The Two Big Challenges They Solved

The authors realized that simply copying the "team of chefs" idea from language models (like the AI that writes this text) wouldn't work for atoms. They had to fix two specific problems:

Challenge A: The "Smoothness" Problem
In a language model, words are distinct (a "cat" is not a "dog"). But in chemistry, atoms move in a smooth, continuous flow. If the Head Chef suddenly switches from the "Grill Expert" to the "Veggie Expert" just because an atom moved a tiny bit, the energy prediction might jump wildly. That's physically impossible (energy can't teleport!).

• The Fix: They created a "Shared Expert" (a generalist chef who is always working) and made sure the switching between specialists happens very smoothly, so the energy curve never breaks.

Challenge B: The "Who is Who?" Problem
In a language model, the router decides which expert to use based on the sentence. In chemistry, the router needs to decide based on the type of atom (e.g., is this a Carbon atom or an Oxygen atom?).

• The Fix: They built a system where the router looks at the specific chemical identity of each atom. A Carbon atom gets routed to Carbon-specialist experts, while an Oxygen atom goes to Oxygen-specialist experts. This is called Element-wise Routing.

3. The "Secret Sauce" Findings

The paper tested many variations of this kitchen team and found some surprising rules:

• The "Generalist" is Vital: The best teams aren't just a bunch of specialists. They need a few Shared Experts who are always active. These experts learn the "common sense" of chemistry that applies to everything (like how atoms generally repel or attract). The paper found that having about half your active team be "specialists" and the other half be "generalists" is the sweet spot.
• Non-Linearity Matters: The authors found that letting the specialists do their own complex thinking before combining their answers (Non-linear) works much better than just averaging their answers first. It's like letting a master painter finish a masterpiece before showing it to the group, rather than just mixing their paint buckets together.
• Global vs. Local Routing: If you make the Head Chef decide based on the whole molecule (Global), the system crashes and becomes unstable. But if the Head Chef decides based on each individual atom (Local/Element-wise), the system is rock-solid and incredibly accurate.

4. The Result: A Chemically Intuitive AI

When they tested this new "MoE" model on massive datasets of molecules and materials, it beat all previous records.

But the coolest part? They looked under the hood to see how the experts were thinking. They used a technique called PCA (a way to map complex data) and found something magical: The experts naturally organized themselves according to the Periodic Table.

• The experts handling "Alkali Metals" (like Sodium) clustered together in one corner of the map.
• The experts handling "Transition Metals" clustered in the center.
• The experts for "Noble Gases" had their own spot.

It turns out the AI didn't just memorize numbers; it rediscovered the logic of chemistry. It learned that elements in the same column of the Periodic Table behave similarly, and it assigned specific "experts" to handle those specific chemical families.

Summary

This paper is like upgrading a chemistry simulator from a single, tired genius to a highly organized, specialized hospital.

• It uses a team of specialists (Experts) who only work when needed.
• It keeps general doctors (Shared Experts) on duty to handle common cases.
• It assigns patients (Atoms) to the right specialist based on their specific condition (Chemical Element).

The result is a system that is faster, cheaper to run, and significantly smarter, capable of predicting how the building blocks of our universe will behave with unprecedented accuracy.

Technical Summary

Here is a detailed technical summary of the paper "Scaling Machine Learning Interatomic Potentials with Mixtures of Experts."

1. Problem Statement

Machine Learning Interatomic Potentials (MLIPs) bridge the gap between the high accuracy of Quantum Mechanical (QM) methods and the computational efficiency of classical force fields. However, improving their expressive capacity through traditional scaling (increasing model depth or width) faces two major bottlenecks:

1. Computational Inefficiency: Dense architectures suffer from all-to-all computational dependencies, limiting parallel efficiency.
2. Optimization Instability: As models grow larger, the optimization landscape becomes complex, leading to diminishing returns and training instability.

While Mixture-of-Experts (MoE) architectures have successfully addressed these issues in Large Language Models (LLMs) by decoupling model capacity from computational cost via sparse activation, their direct application to MLIPs is hindered by two fundamental challenges:

• Equivariance Compatibility: Standard MoE designs often conflict with the equivariant representations required by state-of-the-art MLIPs (e.g., MACE, SevenNet).
• Physical Continuity: MLIPs model continuous potential energy surfaces (PES). Abrupt switching between experts (common in discrete token-based MoE) can introduce numerical instabilities or non-physical discontinuities, violating energy conservation laws.

Previous attempts, such as the Mixture-of-Linear-Experts (MoLE) framework (used in UMA models), addressed these by using linear combinations and global (configuration-level) gating. However, MoLE limits expressive capacity by lacking nonlinear activation on expert outputs and fails to capture element-specific nuances due to its global gating assumption.

2. Methodology

The authors propose a novel MoE-integrated DPA3 architecture. DPA3 is chosen as the baseline because it uses exclusively invariant node and edge features, allowing for valid nonlinear operations on expert outputs while remaining compatible with equivariant GNNs.

Key methodological innovations include:

• Element-wise Routing (MoE-E): Unlike previous global routing strategies that average features across the entire configuration, this approach encodes the chemical identity of each atom (via atomic number $Z_i$) into a latent representation. The router then selects experts based on the specific atom type, enabling element-dependent expert specialization.
• Nonlinear Expert Aggregation: The architecture applies a nonlinear activation function ($\sigma$) within each expert before aggregating their outputs. This contrasts with MoLE, which aggregates linearly and applies activation only at the end.
• Shared Experts Mechanism: To ensure stability and capture universal chemical knowledge, a subset of experts is always activated with fixed weights, independent of the input atom type. The remaining experts are sparsely activated based on the routing mechanism.
• Mathematical Formulation:
  • MoE-E: $y_i = \sum \alpha_{ij} \sigma(W_j x_i + b_j) + \sum \sigma(W_j x_i + b_j)$ (Shared experts).
  • MoLE-E: $y_i = \sigma(\sum \alpha_{ij} (W_j x_i + b_j) + \dots)$ (Linear aggregation before activation).

The study systematically compares four variants: MoE-E (Element-wise, Nonlinear), MoE-G (Global, Nonlinear), MoLE-E (Element-wise, Linear), and MoLE-G (Global, Linear).

3. Key Contributions

1. Architectural Innovation: The first successful integration of standard nonlinear MoE mechanisms into an equivariant-compatible MLIP framework (DPA3), overcoming the limitations of linear-only MoLE approaches.
2. Element-wise Routing Strategy: Demonstrating that routing based on specific atomic species (element-wise) is superior to global configuration-level routing for capturing chemical diversity and ensuring training stability.
3. Shared Expert Optimization: Establishing that allocating approximately 50% of activated experts as "shared experts" provides the optimal trade-off between capturing universal chemical laws and enabling element-specific specialization.
4. Interpretability: Providing evidence that the learned routing weights spontaneously organize according to periodic table trends, offering a chemically interpretable mechanism for the model's success.

4. Results

The proposed MoE-E architecture was benchmarked against standard DPA3 and widened dense baselines across three major datasets: OMol25 (organic molecules), OMat24 (solid-state materials), and OC20M (catalytic surfaces).

• Performance Gains:
  • MoE-E achieved State-of-the-Art (SOTA) accuracy, outperforming both the baseline DPA3 and a "6× Params" dense baseline (which has 6 times the parameters) across all datasets.
  • On OMol25, MoE-E reduced normalized Energy MAE by 0.10 and Force MAE by 0.10 compared to MoLE-E.
  • It demonstrated superior parameter efficiency, achieving higher accuracy than dense models with significantly more parameters under the same computational budget.
• Ablation Studies:
  • Sparse Activation + Shared Experts: Performance scales monotonically with the number of activated experts ($K$) only when shared experts are present. Without shared experts, performance saturates or degrades.
  • Nonlinearity vs. Linearity: MoE-E (nonlinear) significantly outperformed MoLE-E (linear), especially when shared experts were introduced. This highlights the necessity of nonlinear specialization for complex PES modeling.
  • Routing Strategy: MoE-G (Global routing) suffered from catastrophic training failure and numerical instability. MoE-E (Element-wise) was the only stable and high-performing configuration.
• Interpretability (PCA Analysis):
  • Principal Component Analysis (PCA) of the expert routing weights revealed that the model implicitly learned the periodic table.
  • Elements clustered by chemical group (e.g., alkali metals, transition metals, lanthanides) and followed periodic trends (atomic number progression) in the latent space. This confirms the model effectively captures element-specific chemical characteristics.

5. Significance

This work establishes MoE-based conditional computation as a principled and scalable alternative to dense parameter scaling for next-generation atomistic foundation models.

• Scalability: It proves that MLIPs can scale to massive parameter counts without the computational prohibitions of dense models, provided sparse activation and shared experts are used.
• Physical Grounding: By enforcing element-wise routing and shared experts, the model maintains physical consistency (smoothness/differentiability) while gaining the ability to model diverse chemical environments.
• Chemical Intuition: The discovery that the model spontaneously learns periodic trends suggests that MoE architectures can serve as a bridge between data-driven learning and established chemical intuition, leading to more transferable and interpretable potentials.

Future Directions: The authors note that while the architecture is theoretically scalable, future work must focus on developing distributed training and inference frameworks specifically optimized for sparse expert parallelism to fully unlock the computational efficiency of these models.