Elastic & Matryoshka MoE Training

• The paper introduces elastic MoE training by sampling varied K values per layer, ensuring a nested expert hierarchy that maintains performance without retraining.
• Both frameworks use coarse-to-fine expert ranking and scale-wise token compression combined with load-balancing penalties to optimize resource-aware inference.
• Empirical results show improved accuracy and reduced word error rates under low-resource conditions, demonstrating practical benefits for large-scale models.

Elastic and Matryoshka Mixture-of-Experts (MoE) training encompasses a class of techniques designed to endow MoE-based Transformer models with the ability to dynamically adjust expert utilization at inference time. By embedding a coarse-to-fine structure either in token granularity or expert selection, these frameworks achieve robust, resource-aware performance across a range of model capacities, with direct applications to LLMs and audio-visual speech recognition (AVSR). Two principal frameworks exemplify this approach: Matryoshka MoE (M-MoE) for general LLMs (Wang et al., 30 Sep 2025) and MoME (Mixture of Matryoshka Experts) for AVSR (Cappellazzo et al., 5 Oct 2025).

1. Model Foundations and Matryoshka Principle

Standard Mixture-of-Experts models embed modular subnetworks (“experts”) within each Transformer layer, selected per token by a routing network. Conventionally, fixed top-K routing constrains the model to a constant number of experts per token, which severely limits elasticity: varying K at inference leads to unstable or poor performance (Wang et al., 30 Sep 2025).

The Matryoshka principle—adapted from Matryoshka Representation Learning (MRL)—imposes a nested structure by systematically varying granularity during training. In the expert dimension (M-MoE), this means exposing the router to a range of K values, compelling it to learn a globally coherent ranking of experts so that coarse-to-fine capacity can be engaged as needed. In the sequence granularity dimension (MoME), elastic token compression across S scales enables the model to learn to decode from multiple information densities, supported by shared or sparsely-routed experts (Cappellazzo et al., 5 Oct 2025).

2. Mathematical Formalism and Routing Mechanisms

Mathematically, the core operation in both frameworks is a sparsely-gated sum over expert outputs. For expert selection, given an input $h \in \mathbb{R}^d$ and N total experts $\{E_i\}$, the output is:

$y = \sum_{i=1}^{N} w_i(h; K) E_i(h)$

where $w_i(h; K)$ is nonzero for only the Top-K experts, determined by softmaxed router logits $s = h W_g$. For M-MoE, $K$ is sampled per layer (or batch) during training:

$K_\ell \sim \mathcal{U}[K_{min}, K_{max}]$

All loss functions are averaged across chosen K and include a load-balancing penalty to encourage even expert utilization:

$$\mathcal{L}_{M-MoE} = \mathbb{E}_{k_{dyn}} \mathbb{E}_{(x,y^*)}\left[\ell(y(x; k_{dyn}), y^*)\right] + \lambda \mathcal{L}_{balance}$$

For MoME, the model operates across S sequence granularities. For each scale $s$, the MoE output involves both shared and routed (Top-K) experts:

$$MoE^{(s)}(h) = \sum_{i=1}^{N_s} E_i(h) + \sum_{i=N_s+1}^{N_s+N_r} g_i^{(s)}(h) E_i(h)$$

$\{E_i\}$0

Here, $\{E_i\}$1 and $\{E_i\}$2 are shared router weights, invariant across scales (Cappellazzo et al., 5 Oct 2025).

3. Training Procedures and Regularization Strategies

Training in both frameworks cycles through multiple inference scenarios within each batch. In M-MoE, each Transformer layer samples its own K per minibatch; in MoME, each batch is processed at S different token compression rates.

Loss functions comprise a standard cross-entropy term (averaged over scenarios) and crucial load-balancing regularizers [Shazeer et al. 2017]. For MoME, the regularizer for routed experts at each scale is:

$\{E_i\}$3

where $\{E_i\}$4 is the fraction of tokens and $\{E_i\}$5 the mean routing probability for expert $\{E_i\}$6. Only MoE-specific parameters (experts, routers, and input projectors) are updated, with the base LLM typically frozen during fine-tuning (Cappellazzo et al., 5 Oct 2025). Layer-wise or scale-wise cycling ensures all levels of system elasticity are properly calibrated.

A high-level training step for MoME is:

• For each scale $\{E_i\}$7 in $\{E_i\}$8:
  • Compress input tokens to scale $\{E_i\}$9: $y = \sum_{i=1}^{N} w_i(h; K) E_i(h)$0
  • Project and concatenate text prefix: $y = \sum_{i=1}^{N} w_i(h; K) E_i(h)$1
  • Forward pass through LLM + MoME
  • Accumulate cross-entropy loss and $y = \sum_{i=1}^{N} w_i(h; K) E_i(h)$2
• Average over S scales, update only MoME parameters

For M-MoE, pseudocode involves sampling $y = \sum_{i=1}^{N} w_i(h; K) E_i(h)$3 for each layer per step, then applying Top-$y = \sum_{i=1}^{N} w_i(h; K) E_i(h)$4 routing, loss computation, and parameter updates as usual (Wang et al., 30 Sep 2025).

4. Expert Ranking, Generalization, and Elasticity

The critical property enforced by Matryoshka training is a consistent, hierarchical ranking over experts (or granularity levels). The router must ensure that:

• The highest-ranked expert(s) carry the most general (“coarse”) information
• Additional experts successively refine and specialize the output

Empirical verification uses the “Focused Spearman Correlation” between expert ranks at different K: M-MoE achieves $y = \sum_{i=1}^{N} w_i(h; K) E_i(h)$5 between, for example, $y = \sum_{i=1}^{N} w_i(h; K) E_i(h)$6 and $y = \sum_{i=1}^{N} w_i(h; K) E_i(h)$7, indicating a robust nested ordering. In contrast, standard fixed-K MoE models exhibit $y = \sum_{i=1}^{N} w_i(h; K) E_i(h)$8 outside their trained K (Wang et al., 30 Sep 2025). MoME extends this to elastic sequence granularities: sharing experts and a router across scales ensures that representations learned for fine-grained (low-compression) inputs benefit compressed, coarse-grained regimes, improving generalization and robustness under token scarcity (Cappellazzo et al., 5 Oct 2025).

5. Practical Implications and Performance

Both frameworks deliver strong empirical performance with high elasticity:

Inference K	Top-K Specialists	M-MoE-Layer
1	52.01	51.69
2	52.16	52.71
4	53.43	53.77
6	54.32	53.56

(Table: Average accuracy over seven LLM benchmarks (Wang et al., 30 Sep 2025).)

A single M-MoE model matches or exceeds fixed-K “specialists” at each value of K, without retraining. Flexibility is further enhanced by layer-wise expert budgets: assigning higher K to early layers yields the greatest performance return under constrained compute budgets.

For MoME, on LRS2/LRS3, models with only 12–14M activated parameters outperform fixed-rate baselines requiring >27M, and exhibit significantly less word error rate (WER) degradation under severe noise (e.g., 32.6% at SNR = –5 dB, vs. 41.8% for the fixed baseline) (Cappellazzo et al., 5 Oct 2025).

6. Implementation Guidelines and Limitations

Robust elastic MoE training requires careful attention to regularization, sampling, and memory optimization:

• Layer-wise K sampling surpasses batch- or microbatch-level schemes for elasticity (Wang et al., 30 Sep 2025).
• Use activation budgets to cap per-token expert count and redistribute unused capacity—saving ~10% memory.
• Both uniform and capacity-aware K sampling regimes are viable; mild capacity bias (e.g., $y = \sum_{i=1}^{N} w_i(h; K) E_i(h)$9) can refine high-K performance.
• Always incorporate per-layer load-balancing terms to avert expert collapse.
• Health of the nested ranking can be monitored via Focused Spearman and MODS metrics.

Limitations: The studied range of K is modest (typically up to K=6); scaling to higher orders or more experts per layer may require adaptive or capacity-aware schedules. Integrating token-level dynamic budgets or top-p routing is an open area for further research (Wang et al., 30 Sep 2025). For MoME, fixed encoder/LLM backbones limit global adaptation; end-to-end finetuning may enhance performance where resources permit.

7. Broader Context and Future Directions

Elastic and Matryoshka MoE frameworks represent a significant advance in adaptive, resource-aware deep learning. By embedding expert and scale nesting, such models support:

• Dynamic inference adaptation to varying compute/resource constraints without retraining
• Fine-grained tradeoffs between inference speed and accuracy
• Improved robustness to information loss (e.g., aggressive token compression, noise)

A plausible implication is that these techniques could generalize across other domains where multi-scale or conditional computation is advantageous, including vision, multimodal fusion, and structured prediction. Future work is anticipated to address larger K regimes, more expressive routing strategies, and tight integration with adaptive computation primitives (Cappellazzo et al., 5 Oct 2025, Wang et al., 30 Sep 2025).