MAESTRO: Markov-Chain-Based Expert Pruning for MoE Models
Imagine managing a company with 128 departments, each fully staffed—but each project actually only uses 4–5 of them. This is the deployment dilemma of Mixture-of-Experts (MoE) models: GPT-OSS-20B has 32 experts per layer, Qwen3-30B has 128, yet each token only activates 4–8 of them. Computationally efficient, memory-wise extravagant: all experts must stay resident in VRAM whether used or not. A 30B-parameter MoE model may only run 3B of computation per inference, but you still need 30B of memory for it.
The obvious fix is pruning unimportant experts—but how do you define "unimportant"?
The Blind Spot of Existing Methods: Local Heuristics
Existing MoE pruning methods, largely borrowed from dense Transformer pruning, evaluate expert importance with local heuristics: selection frequency, parameter norms, token distributions. These approaches treat each expert independently, ignoring the cross-layer dependencies in MoE routing.
Example: if expert A in layer 5 is activated, expert B in layer 6 may be activated 80% of the time. A and B form a functional pair—removing either breaks this processing path. But looking at per-layer statistics alone, both may seem unimportant due to low individual frequency.
The Core Innovation: Modeling Routing as a Markov Chain
MAESTRO (Markov-chain Approximated Expert Sparsification via Transition-based Routing) from IIT Delhi and NVIDIA models the token's expert-activation trajectory across layers as an ergodic Markov chain:
1. State space: all L×E expert slots (L layers × E experts) are states of the chain; a token traveling from layer 1 to L moves between states. 2. Transition matrices: run autoregressive inference (not teacher-forcing) on a small calibration set, counting "expert i → expert j" co-occurrences between adjacent layers to build each layer's transition matrix C^(l). 3. Ergodicity guarantee: connect the last layer back to the first (irreducibility) and add small self-loops (aperiodicity) so a unique stationary distribution exists. 4. Stationary distribution: solve for π, where π[i] represents the long-run probability of expert i being routed. 5. Pruning: experts with the smallest stationary probability contribute least and are removed first.
The key insight: the stationary distribution is derived from global trajectories and inherently encodes cross-layer dependencies. Bridge experts with low individual activation frequencies but critical connectivity receive higher stationary probabilities and are preserved.
Engineering Details
- Autoregressive calibration vs teacher-forcing: MAESTRO collects transition statistics via autoregressive rollout, since teacher-forcing input distributions differ from real inference. Removing autoregression drops average performance by 1.15% and increases cross-task variance by 10.53%.
- Uniform pruning constraint: pruning the same fraction per layer rather than a global top-k avoids over-pruned bottleneck layers. Flexibility is traded for stability.
- Recovery Fine-Tuning (RFT): lightweight post-pruning fine-tuning to repair the output distribution. Without RFT, average performance drops 20.59% and cross-task variance explodes by 1289%.
- 50% compression: up to 10.61% higher average performance retention than the strongest baseline, with significantly lower cross-task variance.
- 25% compression: on GPT-OSS-20B, MAESTRO beats the uncompressed model on 4 tasks—Wikitext perplexity 2.35 vs 2.52, Lambada perplexity 1.38 vs 1.42, PIQA 81.18 vs 80.69, CommonsenseQA 80.43 vs 79.12. Pruning improved things—a regularization effect.
- 50% compression: Lambada perplexity 1.46, only 0.04 above the original—no other method comes close.
- vs HC-SMoE: the clustering-based method collapses at 50% compression, with average retention of 74.58% and a standard deviation of 16.60%.
Experimental Results
Tested on GPT-OSS-20B (32 experts/layer) and Qwen3-30B (128 experts/layer) across 17 tasks in 5 domains (including safety, bias, and ethics evaluation):
Why Lower Variance?
Local heuristic pruning may accidentally keep task-relevant experts for some tasks while cutting them for others, causing volatile results. MAESTRO's global view preserves experts serving as bridges in the overall routing structure—these contribute to many tasks, so performance is more stable across them.
Analysis
The elegance of MAESTRO is converting an engineering problem (which experts to cut) into a probability modeling problem (the Markov chain stationary distribution), giving "expert importance" a precise mathematical definition: not "how often selected" but "share of the long-run routing distribution."
The bigger takeaway: MoE experts are not independent parts but nodes in a routing network. A node's value lies in its role within network connectivity—analogous to betweenness centrality in social network analysis. Some nodes have little traffic, but removing them breaks the network.
RFT's necessity is notable: even the smartest pruning strategy perturbs the output distribution. MAESTRO without RFT still beats REAP without RFT by ~1%, showing higher pruning quality, but RFT's gains are of a different magnitude—compression and recovery are orthogonal problems requiring simultaneous attention.
Finally, the phenomenon of beating the original model at 25% compression deserves study. If pruning acts as regularization, the original model may be over-parameterized—some experts memorizing training noise rather than learning useful patterns. This suggests a "sweet spot" for expert count in MoE architectures, beyond which more experts may actually hurt.
---
Paper: https://arxiv.org/abs/2607.08601 HTML full text: https://arxiv.org/html/2607.08601