MatRIS-MoE: Invariant Graph MoE for uMLIPs

• The paper introduces a novel invariant graph-attention Mixture-of-Experts model that compresses the training of billion-parameter uMLIPs from weeks to hours.
• It integrates precise atomic geometric representations with symmetric invariance, supporting second-order force matching and robust multi-task loss for accurate physical simulations.
• Janus, a hybrid-parallel training framework, co-designs the model-system with advanced sharding techniques to achieve over 90% weak-scaling efficiency on Exascale supercomputers.

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 (Zhou et al., 17 Apr 2026). 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 (Zhou et al., 17 Apr 2026).

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 (Zhou et al., 17 Apr 2026). 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 (Zhou et al., 17 Apr 2026).

These raw geometric quantities are embedded as

$Z_i$0

The network core consists of $Z_i$1 Interaction Blocks. Within each block, the tuple $Z_i$2 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 (Zhou et al., 17 Apr 2026). 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 (Zhou et al., 17 Apr 2026). 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 $Z_i$3, 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$Z_i$4”) specialized on constructing messages;
3. multi-head self-attention (MHSA) on the node features;
4. a feature-update MoE (“MoE$Z_i$5”) to refine post-attention features;
5. a small MLP “Refine” that produces $Z_i$6 (Zhou et al., 17 Apr 2026).

The MHSA sublayer is defined per head $Z_i$7 by

$Z_i$8

$Z_i$9

$X_i$0

This makes attention local to the graph neighborhood structure $X_i$1 rather than globally dense (Zhou et al., 17 Apr 2026).

The readout stage predicts per-atom energies $X_i$2 and, optionally, magnetic moments $X_i$3. Total energy is then formed additively:

$X_i$4

The paper states that this construction ensures size extensivity and strict rotational/translational invariance (Zhou et al., 17 Apr 2026). Forces and stress are obtained by automatic differentiation:

$X_i$5

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 (Zhou et al., 17 Apr 2026). 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 $X_i$6	Top-$X_i$7 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 (Zhou et al., 17 Apr 2026).

Routing is defined per atom and per layer. For atom $X_i$8 at layer $X_i$9, with feature $(i,j)$0, a linear gater computes logits

$(i,j)$1

followed by softmax probabilities

$(i,j)$2

Only the top-$(i,j)$3 experts by $(i,j)$4 are activated, and expert $(i,j)$5 is an MLP $(i,j)$6 of hidden size $(i,j)$7, equal to the interaction-block width. The resulting MoE output is

$(i,j)$8

where $(i,j)$9 denotes the activated expert set (Zhou et al., 17 Apr 2026).

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 (Zhou et al., 17 Apr 2026). 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 (Zhou et al., 17 Apr 2026). 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 $R_c$0 in the activated subnetwork, while the full model stores $R_c$1 parameters (Zhou et al., 17 Apr 2026). 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:

$R_c$2

This is then embedded in a robust multi-task loss over a batch of $R_c$3 structures:

$R_c$4

$R_c$5

The same section of the paper also restates the MoE aggregation as

$R_c$6

(Zhou et al., 17 Apr 2026).

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 (Zhou et al., 17 Apr 2026). 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 (Zhou et al., 17 Apr 2026).

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 (Zhou et al., 17 Apr 2026). 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 $R_c$7 experts across ranks, prefetching only the top-$R_c$8 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 (Zhou et al., 17 Apr 2026). 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 (Zhou et al., 17 Apr 2026).

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 (Zhou et al., 17 Apr 2026). 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 (Zhou et al., 17 Apr 2026). 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 (Zhou et al., 17 Apr 2026).

The paper states that these system–algorithm co-designs maintain over 90% parallel efficiency at 20K+ ranks while supporting second-order differentiation natively (Zhou et al., 17 Apr 2026). 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 (Zhou et al., 17 Apr 2026).

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 $R_c$9 1,048 PFLOPS (93.8% eff.), sustained 762 PFLOPS (25.8% of 2.952 EF)
LineShine	133 PFLOPS $r_{ij}=\|X_i-X_j\|$0 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 $r_{ij}=\|X_i-X_j\|$1 1.63 M edges/sec with 53.9% efficiency and peak 602 PFLOPS on CNIS, and 0.90 M $r_{ij}=\|X_i-X_j\|$2 1.77 M edges/sec with 50.6% efficiency and peak 672 PFLOPS on LineShine (Zhou et al., 17 Apr 2026).

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 (Zhou et al., 17 Apr 2026). 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 (Zhou et al., 17 Apr 2026).

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 (Zhou et al., 17 Apr 2026). 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 $r_{ij}=\|X_i-X_j\|$3 weak-scaling efficiency (Zhou et al., 17 Apr 2026).

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 (Zhou et al., 17 Apr 2026). 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.