MatRIS-MoE: Invariant Sparse Mixture-of-Experts

• MatRIS-MoE is a unified framework that integrates multi-scale representation learning with sparse expert routing to deliver capacity-efficient and robust neural inference.
• It employs a matryoshka pooling mechanism and elastic expert utilization, enabling dynamic adjustment of token granularity and computational budgets.
• The design incorporates invariance-preserving elements and cross-modal fusion, supporting applications in AVSR, language modeling, and scientific simulation.

MatRIS-MoE (Matryoshka Representation Invariant Sparse Mixture-of-Experts) architectures generalize and unify sparse expert-based model scaling with elastic, multi-scale, and invariant design principles. This class of architectures subsumes advances in language modeling, audio-visual sequence modeling, and large-scale scientific simulation, providing a substrate for capacity-efficient, robust, and dynamically adaptable neural network inference. MatRIS-MoE involves the coordinated application of Mixture-of-Experts (MoE) mechanisms, Matryoshka representation learning (MRL), and invariance-preserving architectural elements, combined with sophisticated parallel and distributed execution frameworks for exascale training. Notable variants are applied to Audio-Visual Speech Recognition (AVSR) in LLMs, transformer-based LLMs with elastic expert utilization, and quantum-accurate interatomic potential prediction (Cappellazzo et al., 5 Oct 2025, Wang et al., 30 Sep 2025, Zhou et al., 17 Apr 2026).

1. Core Architectural Components

MatRIS-MoE models instantiate three principal ingredients:

• Multi-Scale Representation Learning: Tokens from input modalities (e.g., audio, video, molecular graphs) are compressed elastically into variable-granularity sequences using average pooling, stacking, or message-passing protocols. This “matryoshka” pooling allows inference-time adjustment of token granularity without retraining or duplication of weights (Cappellazzo et al., 5 Oct 2025).
• Sparse Mixture-of-Experts Layers: Transformer feed-forward blocks (FFN) or message-/feature-update stages are replaced or augmented with a set of expert subnetworks, with expert selection determined by a learned router, generally activating a top-K routed subset plus a small set of always-on shared experts (Cappellazzo et al., 5 Oct 2025, Zhou et al., 17 Apr 2026). Each expert is a small two-layer bottleneck MLP or similar.
• Invariance and Modality Fusion: Architectural features (distance/angle embeddings, permutation-invariant sum aggregation) ensure strict invariance to symmetry groups as required for the task domain (e.g., 3D translations/rotations, atom permutations) (Zhou et al., 17 Apr 2026). Modalities (e.g., Z^a for audio, Z^v for video) are concatenated in embedding space, with no cross-modality fusion before the LLM or backbone.

2. Routing, Gating, and Expert Sharing

Routing within MatRIS-MoE is handled by lightweight learned routers at each layer or interaction block. For input token hidden state $h_\ell \in \mathbb{R}^{d_h}$:

• Router Logits/Gating: $g = \mathrm{softmax}(W_r h_\ell + b_r)$ produce gating scores for $N_r$ routed experts.
• Top-K Sparse Activation: Only the top-K experts with highest scores per token (or per node/edge in graphs) are activated; all others receive zero weight.
• Shared Experts and Router Consistency: A subset $N_s$ of experts remains always active (shared across all tokens/scales). Critically, the same router parameters and expert pool are used across all matryoshka scales, so learned token-expert assignments are consistently aligned over all compression rates and input granularities (Cappellazzo et al., 5 Oct 2025). This “expert-sharing” is essential for cross-scale generalization, and enables smaller-scale (coarser) sequences to inherit pathways learned at fine resolution.
• Coarse-to-Fine Expert Nesting: In elastic MoE variants, stochastically varying the number of active experts $k_\ell$ across layers and forward passes during training enforces a nested ranking of expert utility—top-1 expert for coarse/vital information, with additional experts providing finer-grained capacity. This ensures that different expert assignments are nested, not disjoint, across inference budgets (Wang et al., 30 Sep 2025).

3. Training Procedures and Objectives

MatRIS-MoE training proceeds by jointly optimizing over all target granularities and expert configurations:

• Multi-Granularity Training: For Matryoshka tokenization, define a set $\{a_1,\dots,a_G\}$ of audio rates and $\{v_1,\dots,v_L\}$ of video rates (or generally modality rates). All $G \times L$ combinations are sampled during training, training the LLM and MoE adapters on all scales in parallel (Cappellazzo et al., 5 Oct 2025).
• Elastic Expert Utilization: At each training step, the number of active experts $k_\ell$ at each layer is randomly sampled from a uniform or capacity-aware range $[k_\text{min},k_\text{max}]$. This enforces router consistency and nested expert utility across all possible inference-time budgets (Wang et al., 30 Sep 2025).
• Loss Function: The canonical objective is a combination of autoregressive cross-entropy (or task-specific loss) averaged over all granularities or pooling rates, plus a load-balancing penalty on the MoE routers to encourage uniform expert utilization. For example,

$g = \mathrm{softmax}(W_r h_\ell + b_r)$0

where $g = \mathrm{softmax}(W_r h_\ell + b_r)$1 is the load-balancing loss as a function of the empirical expert assignment distribution (Cappellazzo et al., 5 Oct 2025, Wang et al., 30 Sep 2025).

4. Invariant and Domain-Specific Backbones

In scientific applications (e.g., interatomic potential prediction), MatRIS-MoE leverages invariant graph-based backbones:

• Input Encoding: Inputs are atom features, pairwise distances $g = \mathrm{softmax}(W_r h_\ell + b_r)$2, and angles $g = \mathrm{softmax}(W_r h_\ell + b_r)$3, encoded via radial and angular basis functions (e.g., Legendre polynomials); no raw coordinates are used (Zhou et al., 17 Apr 2026).
• Interaction Blocks: Each Interaction Block (IB) includes a triangular update for three-body interactions, message-passing (pre-attention) MoE, multi-head self-attention, feature-update (post-attention) MoE, and refinement MLP (Zhou et al., 17 Apr 2026).
• Invariance Properties: Models are constructed to be strictly rotationally, translationally, and permutation invariant via neighborhood sum aggregation, invariant embeddings, and permutation-symmetric readout layers.

5. Scalability, Parallelism, and Elasticity

MatRIS-MoE architectures scale to billions of parameters via hybrid parallelism and just-in-time (JIT) expert dispatch:

• Parallel Execution (Janus FS-3D): A hybrid of data-parallel (DP), graph-parallel (GP), and expert-parallel (EP) sharding is applied. Parameters are distributed across DP and EP ranks; node and edge activations across GP ranks (Zhou et al., 17 Apr 2026). Sparse MoE dispatch is performed using all-to-all communication only over top-K active experts, with FP16 compression for efficiency.
• Second-Order Derivative Support: For scientific simulation, the architecture supports efficient “double backward” differentiation for force- and Hessian-matching (Zhou et al., 17 Apr 2026).
• Empirical Efficiency: The 11.5B-parameter MatRIS-MoE model achieves >90% parallel efficiency on exascale hardware, with training speedups of multiple orders of magnitude (Zhou et al., 17 Apr 2026).

6. Empirical Performance and Trade-Offs

MatRIS-MoE delivers significant performance and efficiency gains across domains:

• Elastic Inference: Single models match the performance of specialist models trained at fixed expert or token budgets across a sweep of active experts (e.g., $g = \mathrm{softmax}(W_r h_\ell + b_r)$4) or compression rates, with sharp reduction in compute at modest loss in fidelity (Wang et al., 30 Sep 2025, Cappellazzo et al., 5 Oct 2025).
• Parameter and Compute Efficiency: In AVSR, 12.7M active parameters at inference suffice for all granularities, with up to 8× reduction in FLOPs as token compression increases, and minimal word error rate degradation (Cappellazzo et al., 5 Oct 2025).
• Robustness: MatRIS-MoE models degrade more gracefully under noise or high compression than LoRA/fixed-scale Matryoshka baselines, and maintain consistent gains in unimodal (ASR/VSR) scenarios (Cappellazzo et al., 5 Oct 2025).
• Domain Extension: In interatomic potential learning, the architecture enables rapid, quantum-accurate simulation across the periodic table, previously infeasible at billion-parameter scale (Zhou et al., 17 Apr 2026).

7. Implementation Considerations

• Expert Budgeting: Modelers must select the number of routed/shared experts (e.g., $g = \mathrm{softmax}(W_r h_\ell + b_r)$5, $g = \mathrm{softmax}(W_r h_\ell + b_r)$6 for speech, $g = \mathrm{softmax}(W_r h_\ell + b_r)$7 for atomic graphs) and top-K for routing as a function of compute.
• Embedding and Aggregation: Projections to common embedding spaces simplify cross-modal fusion, while strict invariance is essential for physical modeling.
• Load Balancing: Router loss and JIT expert planning are required for both efficiency and convergence.
• Elastic Deployment: Practitioners may preselect a menu of $g = \mathrm{softmax}(W_r h_\ell + b_r)$8 or compression parameters to match end-task preferences for accuracy and latency, trading off fine-grained versus coarse expert utilization (Wang et al., 30 Sep 2025).

---

MatRIS-MoE architectures unify multi-scale, capacity-adaptive, and invariance-preserving modeling for both perception (speech, vision, language) and AI-for-science workloads, overcoming prior limitations in fixed-scale, non-elastic MoE, and monolithic backbone architectures (Cappellazzo et al., 5 Oct 2025, Wang et al., 30 Sep 2025, Zhou et al., 17 Apr 2026). They enable robust, efficient, and interpretable deployment of neural models across a wide range of computational and domain-theoretic constraints.