---
title: 'MatRIS-MoE: Invariant Graph MoE for uMLIPs'
url: https://www.emergentmind.com/topics/matris-moe
type: topic
---

# MatRIS-MoE: Invariant Graph MoE for uMLIPs

Searching arXiv for the specified paper and closely related context.
MatRIS-MoE is an invariant graph-attention Mixture-of-Experts model for universal interatomic potentials (uMLIPs), introduced together with Janus, a high-dimensional distributed training framework designed to break the billion-parameter barrier on Exascale systems [2604.15821]. In this formulation, uMLIPs are treated as foundation models for quantum-accurate physical simulations across inorganic materials and organic molecules spanning the periodic table. The central contribution is the coupling of an invariant architecture that preserves physically required symmetries with a hybrid-parallel training system that supports second-order differentiation and large-scale sparse expert execution, thereby compressing the training of billion-parameter uMLIPs from weeks to hours [2604.15821].

## 1. Architectural definition and scope

MatRIS-MoE is presented as a billion-parameter Mixture-of-Experts model built upon invariant architecture for universal machine learning interatomic potentials [2604.15821]. Its purpose is to support multi-task learning across heterogeneous chemical domains while retaining the geometric and symmetry constraints expected of interatomic potential models. The model inherits the geometric graph construction of MatRIS: each atomic configuration under periodic boundary conditions is first tiled into a supercell and then converted into a spatial graph whose nodes are atom indices $i$ with atomic number $Z_i$ and Cartesian coordinate $X_i$, whose edges are neighbor pairs $(i,j)$ within a cutoff $R_c$ with scalar distance $r_{ij}=\|X_i-X_j\|$, and whose triplets encode angles $\theta_{ijk}=\angle(X_iX_j,X_iX_k)$ for $i$–$j$–$k$ interactions [2604.15821].

These raw geometric quantities are embedded as
$$
v_i^0 = \mathrm{Emb}_Z(Z_i), \qquad
e_{ij}^0 = \mathrm{Emb}_r(r_{ij}), \qquad
a_{ijk}^0 = \mathrm{Emb}_\theta(\theta_{ijk}).
$$
The network core consists of $T=6$ Interaction Blocks. Within each block, the tuple $\{v_i^t,e_{ij}^t,a_{ijk}^t\}$ is updated through a triangular-update MLP, a message-update MoE, multi-head self-attention over node features, a feature-update MoE, and a final refinement MLP [2604.15821]. This organization makes MatRIS-MoE simultaneously a geometric graph model, an attention model, and a sparse expert model.

The paper positions the model in the context of uMLIPs that require second-order derivatives during training, which creates an immediate systems-level challenge: training infrastructure must support automatic differentiation not only for energies but also for force matching and stress prediction [2604.15821]. A plausible implication is that MatRIS-MoE should be understood not merely as a neural architecture, but as a model-system co-design in which architectural sparsity and distributed execution are inseparable.

## 2. Invariant representation and interaction pipeline

The invariant character of MatRIS-MoE is rooted in its atomic-environment representation and its energy readout. For each interaction block $t$, the model applies a sequence of updates:

1. a triangular-update MLP (“Tri”) that refines edge and angle features;
2. a message-update MoE (“MoE$_0$”) specialized on constructing messages;
3. multi-head self-attention (MHSA) on the node features;
4. a feature-update MoE (“MoE$_1$”) to refine post-attention features;
5. a small MLP “Refine” that produces $v_i^{t+1}, e_{ij}^{t+1}, a_{ijk}^{t+1}$ [2604.15821].

The MHSA sublayer is defined per head $h$ by
$$
Q_{i,h}=W_h^Q v_i', \qquad K_{j,h}=W_h^K v_j',
$$
$$
\alpha_{ij,h}=\mathrm{softmax}_{j\in\mathcal N(i)}(Q_{i,h}\cdot K_{j,h}/\sqrt d),
$$
$$
m_i^h=\sum_j \alpha_{ij,h}\, W_h^V v_j',
\qquad
v_i''=\mathrm{Concat}_h[m_i^h]W^O.
$$
This makes attention local to the graph neighborhood structure $\mathcal N(i)$ rather than globally dense [2604.15821].

The readout stage predicts per-atom energies $E_i$ and, optionally, magnetic moments $m_i$. Total energy is then formed additively:
$$
E_{\rm total}=\sum_{i=1}^N E_i(v_i^T).
$$
The paper states that this construction ensures size extensivity and strict rotational/translational invariance [2604.15821]. Forces and stress are obtained by automatic differentiation:
$$
F_i=-\frac{\partial E_{\rm total}}{\partial X_i}, \qquad
\sigma=\frac{1}{\mathcal V}\frac{\partial E_{\rm total}}{\partial \varepsilon}.
$$

This representation is notable because invariance is not imposed only at the output level; it is baked into the graph construction, the use of scalar distances and angular triplets, and the additive energy decomposition. This suggests that MatRIS-MoE follows the established physical design principle that symmetry compliance should be enforced structurally rather than recovered statistically.

## 3. Mixture-of-Experts mechanism

The Mixture-of-Experts design is explicitly motivated by multi-task learning across heterogeneous chemical domains [2604.15821]. Each MoE layer is implemented as a two-stage sparse router with element-dependent routing. The reported expert counts are:

| Model scale | Number of experts $E$ | Top-$K$ routing |
|---|---:|---:|
| Small | 72 | 4 |
| Medium | 40 | 8 |
| Large | 72 | 16 |

The corresponding parameter counts are 1.09 B for the small model, 2.47 B for the medium model, and 11.5 B for the large model [2604.15821].

Routing is defined per atom and per layer. For atom $i$ at layer $t$, with feature $h_i=v_i^t$, a linear gater computes logits
$$
g_k(h_i)=w_k^T h_i+b_k,\qquad k=1\dots E,
$$
followed by softmax probabilities
$$
p_k(h_i)=\frac{\exp(g_k(h_i))}{\sum_{j=1}^E \exp(g_j(h_i))}.
$$
Only the top-$K$ experts by $p_k$ are activated, and expert $k$ is an MLP $f_k(\cdot)$ of hidden size $d$, equal to the interaction-block width. The resulting MoE output is
$$
o_i^{\rm MoE}=\sum_{k\in \mathcal E_i} p_k(h_i)\, f_k(h_i),
$$
where $\mathcal E_i$ denotes the activated expert set [2604.15821].

A distinctive feature is that routing is performed per element type: each atom’s element ID determines its subset of experts, and this mapping is fixed, or time-independent, during dynamics [2604.15821]. The paper further states that no auxiliary “importance” or “load-balance” regularizer is added to the loss; instead, workload balance is achieved at runtime by Janus through just-in-time sparse expert planning [2604.15821]. This departs from many generic MoE formulations in which balancing is enforced by auxiliary terms. A plausible implication is that the model treats chemical identity as a strong prior for specialization, reducing the need for learned balancing heuristics.

The paper also distinguishes between the active and stored parameter complexity: parameter count per MoE layer grows as $O(K\cdot d^2)$ in the activated subnetwork, while the full model stores $E\cdot d^2$ parameters [2604.15821]. This formalizes the central efficiency premise of sparse activation.

## 4. Objective functions and second-order training

MatRIS-MoE is trained with a per-structure base loss that combines energy, force, stress, and magnetic-moment terms:
$$
L_i = w_E |E_i - E_i^{\rm ref}|
+ w_F \|F_i-F_i^{\rm ref}\|_2^2
+ w_S \|\sigma_i-\sigma_i^{\rm ref}\|_F^2
+ w_M |M_i-M_i^{\rm ref}|.
$$
This is then embedded in a robust multi-task loss over a batch of $N$ structures:
$$
\mu=\tfrac1N\sum_i L_i,\qquad
\sigma=\sqrt{\tfrac1N\sum_i(L_i-\mu)^2},\qquad
z_i=\tfrac{L_i-\mu}{\sigma},
$$
$$
\mathcal W(z_i,\tau)=
\begin{cases}
1 & z_i\le\tau,\\
\frac{\tau}{z_i} & z_i>\tau,
\end{cases}
\qquad
\mathcal L_{\rm robust}
=\frac1N\sum_{i=1}^N \mathcal W(z_i,\tau)^2\,L_i,\qquad \tau=2.0.
$$
The same section of the paper also restates the MoE aggregation as
$$
p_k(x)=\frac{\exp(g_k(x))}{\sum_j\exp(g_j(x))},\qquad
o_k(x)=f_k(x),\qquad
\bar o(x)=\sum_{k\in \mathrm{top}\mbox{-}K} p_k(x)\,o_k(x)
$$
[2604.15821].

Training is characterized by the need for second-order derivatives. Janus implements a double-backward execution for force matching, with each iteration comprising four phases: forward for energy, first backward for force computation, double backward for force-loss propagation, and final backward for energy-loss propagation [2604.15821]. The paper identifies this requirement as one of the central obstacles in scaling uMLIPs, because second-order force-matching substantially increases memory and communication pressure [2604.15821].

This training formulation is significant because it integrates robust multi-task weighting with a systems design specifically intended for second-order automatic differentiation. The paper’s discussion suggests that the training difficulty is not just model size, but the interaction between large sparse models and higher-order derivative computation.

## 5. Janus and FS-3D distributed execution

Janus is described as the first high-dimensional hybrid-parallel training framework that breaks through the billion-parameter barrier for uMLIPs on two Exascale supercomputers [2604.15821]. Its basic execution unit is FS-3D, short for “Fully Sharded 3-Dim,” which combines three forms of sharding:

| FS-3D component | Function |
|---|---|
| FSDP | shards all non-MoE parameters, gradients, and optimizer state across GPUs |
| FSGP | slices each mini-batch’s atomic graph over multiple ranks to reduce activation memory |
| FSEP | shards the $E$ experts across ranks, prefetching only the top-$K$ experts per atom via all-to-all each step |

At the beginning of each step, Janus performs just-in-time sparse expert planning: it routes all MoE layers, counts activated tokens per expert, and then deterministically assigns experts to ranks in a greedy least-loaded fashion [2604.15821]. The paper attributes three direct effects to this procedure: expert sparsification, in which only active expert parameters are communicated; neighbor dispatch, in which atom features are routed to the owning ranks by an all-to-all; and sparse gradient return, in which expert gradients are all-to-all’d back, with missing-expert gradients zero-filled before final gradient synchronization [2604.15821].

FS-3D restores and re-shards parameters only on demand, defers gradient synchronization until the final phase, and reuses recorded operator dependencies to prefetch parameters and overlap communication [2604.15821]. This is paired with hardware-aware optimizations that differ by system. On GPU-based CNIS, the paper reports pipelined gradient synchronization across gradient buckets, atom-type aware FP16 compression for MoE all-to-all with 50% smaller volume and negligible loss, and hand-tuned kernels for neighbor gather, edge aggregation, MHSA, and MoE dispatch/combine to maximize Tensor-Core throughput [2604.15821]. On ARMv9-based LineShine, the reported optimizations are software-defined asynchronous MPI streams pinned to dedicated communication cores and SDMA-driven double-buffered HBM↔DDR data movement for large GEMM and operator inputs and weights, hiding memory latency and achieving 1.4× higher effective bandwidth [2604.15821].

The paper states that these system–algorithm co-designs maintain over 90% parallel efficiency at 20K+ ranks while supporting second-order differentiation natively [2604.15821]. In this framing, Janus is not an auxiliary training script but a full distributed runtime tailored to the numerical structure of uMLIP training.

## 6. Exascale performance and stated implications

The reported evaluation is conducted on two Exascale supercomputers. CNIS is described as 5,632 nodes with 8 GPGPUs per node, 65.5 TFLOPS FP32 per GPU, and a theoretical 2.952 EFLOPS FP32. LineShine is described as 20,480 nodes with 304 ARMv9 cores and 32 GB HBM per die, with theoretical 5.0 EFLOPS FP32 [2604.15821].

For weak scaling to the full machine, using the large 11.5 B-parameter model and constant per-unit graph size, the paper reports:

| System | Weak-scaling result |
|---|---|
| CNIS | 152 PFLOPS $\rightarrow$ 1,048 PFLOPS (93.8% eff.), sustained 762 PFLOPS (25.8% of 2.952 EF) |
| LineShine | 133 PFLOPS $\rightarrow$ 1,200 PFLOPS (90.3% eff.), sustained 1,033 PFLOPS (21.0% of 5 EF) |

For strong scaling from one-eighth to full scale at fixed global problem size, again for the large model, the reported results are 0.80 M $\rightarrow$ 1.63 M edges/sec with 53.9% efficiency and peak 602 PFLOPS on CNIS, and 0.90 M $\rightarrow$ 1.77 M edges/sec with 50.6% efficiency and peak 672 PFLOPS on LineShine [2604.15821].

The abstract separately reports peak performance of 1.2/1.0 EFLOPS in single precision at over 90% parallel efficiency, corresponding to 24%/35.5% of theoretical peak, on the two Exascale systems [2604.15821]. End-to-end, the paper states that MatRIS-MoE compresses billion-parameter uMLIP training from weeks—illustrated by “21 days on 256 H200 GPUs for UMA”—to hours on Exascale, with up to 3,200× normalized throughput and 2796× for the large model versus the UMA baseline [2604.15821].

The paper identifies three core barriers that the combined MatRIS-MoE plus Janus design addresses: second-order force-matching at scale, billion-parameter model parallelism, and Exascale communication bottlenecks [2604.15821]. Specifically, it states that Janus’s double-backward lifecycle and FS-3D sharding reduce memory footprint by 4–8× versus naive automatic differentiation; that hybrid FSDP + FSGP + FSEP shards both model weights and activation graphs, eliminating the need for pure data parallelism; and that hardware-tailored collectives, FP16 MoE compression, and SDMA/DDReactor pipelines deliver $>90\%$ weak-scaling efficiency [2604.15821].

The broader interpretation offered in the paper is methodological. It lists as lessons for other universal MLIPs: embedding sparsity early, combining data-, model-, and graph-parallelism in a unified sharding framework, co-designing kernels and collectives with target hardware, and exploiting problem structure such as element IDs and fixed routing patterns for just-in-time planning instead of heavy regularization [2604.15821]. This suggests that MatRIS-MoE should be viewed both as a concrete uMLIP instantiation and as an architectural template for future billion-parameter scientific models trained at Exascale.

Source: https://www.emergentmind.com/topics/matris-moe