PaperlayerView exact revision on arXiv

Source-linked AI summary

When Load-Balancing Goes Too Far: Expert Pruning in Over-Dispersed Mixture-of-Experts Models

Berkcan Kapusuzoglu, Connor Pryor, Sangwoo Cho, Supriyo Chakraborty, Shi-Xiong Zhang, Sambit Sahu, Milind Naphade

arXiv:2609.04453v1cs.LGcs.AIcs.CL

TL;DR

Under over-dispersed routing, router-based importance signals and perplexity can fail to reflect downstream task performance, while pruning exposes trade-offs across domains. MESA adjusts expert scores to protect the currently worst-affected domain; at 25% pruning, it achieves the smallest worst-case degradation and outperforms REAP on 7 of 11 benchmarks.

  • Problem

    Under over-dispersed routing, near-uniform token distribution makes router-based importance signals unreliable, and perplexity may not predict downstream task accuracy.

    Source: §1 · p. 2

  • Method

    MESA iteratively boosts importance scores for experts serving whichever domain is currently worst-affected, minimizing worst-case domain degradation.

    Source: §1 · p. 2, §3.2 · p. 4

  • Results

    At 25% pruning, MESA achieves the smallest worst-case domain degradation and outperforms REAP on 7 of 11 benchmarks, with 5 statistically significant gains.

    Source: §1 · p. 2, §6 · p. 12

  • Takeaways & Limitations

    Over-dispersed routing is a distinct pruning regime in which standard importance assumptions fail, while under standard routing methods degrade proportionally without domain imbalance.

    Source: §6 · p. 12

    MESA’s algorithm is a greedy heuristic without formal convergence guarantees.

    Source: §3.2 · p. 4

Abstract

from arXiv · show

Expert pruning reduces the memory and serving cost of Mixture-of-Experts (MoE) models by removing low-importance experts identified by the router, assuming router probabilities provide a reliable importance signal. We observe that this assumption breaks down under over-dispersed routing, a regime associated with aggressive load-balancing during training, in which tokens are distributed nearly uniformly across experts and importance signals collapse. In this regime, perplexity does not predict downstream task accuracy: on gpt-oss-20B, the lowest-perplexity pruning configuration yields the worst mathematical reasoning, while the highest-perplexity configuration preserves it. This does not occur under standard routing (e.g., Mixtral-8x7B-Instruct), where perplexity and accuracy degrade together. Pruning under over-dispersed routing also exposes a capability trade-off in which no single scoring metric dominates: activation-aware scoring preserves mathematical reasoning but severely degrades knowledge-intensive science (an 18-point gap on GPQA), whereas frequency-based scoring exhibits the reverse. We propose Minimax Expert Score Allocation (MESA), a domain-aware method that iteratively boosts importance scores for experts serving whichever domain is currently worst-affected, minimizing worst-case domain degradation rather than average accuracy. At 25% expert pruning MESA achieves the smallest worst-case degradation across domains, outperforming activation-aware baselines on 7 of 11 benchmarks at a correspondingly reduced memory footprint, and it generalizes to gpt-oss-120B, Gemma-4-26B-A4B, and OLMoE-1B-7B. Our results indicate that over-dispersed routing is a qualitatively distinct pruning regime in which standard assumptions fail, and that recognizing it is a prerequisite for principled expert pruning of load-balanced MoE models.

1 Introduction

Expert pruning can reduce MoE memory and serving costs, but over-dispersed routing makes router-based importance signals unreliable and exposes conflicts between perplexity and task accuracy. MESA addresses domain-level degradation by boosting scores for experts serving the worst-affected domain.

1 Introduction

  • MoE expert pruning removes entire expert modules to reduce the substantial memory and serving costs created by models’ large total parameter counts.MoE models activate only a subset of parameters per token.Source: §1 · p. 1
  • Aggressive load-balancing can distribute tokens nearly uniformly across experts, causing global importance scores to average away subtle domain-specific signals and disproportionately degrade some domains.The cited regime is associated with unusually high auxiliary load-balancing loss coefficients, such as λaux = 0.9.Source: §1 · p. 2
  • Across five pruning configurations, the lowest-PPL random-pruning result (PPL=40.95) has the worst task accuracy, while highest-PPL REAP (PPL=145.04) retains strong performance.Spearman ρ is 0.70 on GSM8K and −0.10 on GPQA.Source: §1 · p. 2
  • Over-dispersed routing is relevant to deployed models because production systems use high auxiliary loss coefficients for hardware load balance and gpt-oss-20b ships with λaux = 0.9.The authors argue pruning methods should account for routing distribution rather than assume standard expert specialization.Source: §1 · p. 2
  • Under over-dispersed routing, activation-aware REAP preserves AIME 2025 at 73% but scores 47.72% on GPQA versus 65.48%, while frequency-based methods show the reverse pattern.No single pruning metric dominates across tasks.Source: §1 · p. 2
  • At 25% pruning, MESA achieves the smallest worst-case domain degradation and outperforms REAP on 7 of 11 benchmarks, with 5 statistically significant gains.Its objective prioritizes the worst-affected domain rather than average-case accuracy.Source: §1 · p. 2

2 Related Work

Prior work studies expert scoring, pruning allocation, and domain-fair compression, but the cited approaches leave a gap in domain-aware importance scoring for MoE experts.

Expert importance scoring

  • REAP [15] combines routing probability with activation ℓ2 norm; DiEP [7] finds router probability mass competitive, while SHAPE [9] reports a gpt-oss-20b result using a REAP baseline that reduces to router probability alone.The authors compute actual activation norms in their REAP implementation.Source: §2 · p. 2

Expert pruning and allocation

  • EvoESAP [14] uses genetic search for per-layer pruning counts, while uniform allocation and global ranking serve as common baselines.Prior work also reports a multiple-choice versus generative-task trade-off: merging favors multiple-choice tasks, while pruning better preserves generative capabilities.Source: §2 · p. 3

Domain fairness in pruning

  • Minimax fairness has been applied to group-level training objectives [20] and dense-model data selection, but not to MoE expert importance scoring.MESA applies the principle to expert scores without retraining.Source: §2 · p. 3

Over-dispersed routing and domain-fair compression

  • Prior work studies load-balancing effects on training stability [6] and SimBal specialization recovery [17], but leaves post-training pruning under near-uniform routing unexplored.DRPruning upweights underrepresented data domains during dense-model fine-tuning; MESA instead adjusts expert scores on frozen models without fine-tuning.Source: §2 · p. 3

3 Method

The method characterizes over-dispersed routing and introduces MESA, which uses per-domain router mass to protect the currently worst-affected domain during pruning.

3.1 Over-Dispersed Routing and Its Consequences

  • In gpt-oss-20b, λaux = 0.9 pushes routing probabilities toward uniformity, with π_i ≈ 1/E for all experts.The auxiliary loss uses the token fraction routed to each expert and its mean routing probability.Source: §3.1 · p. 3
  • Near-uniform top-k routing can let pruning decrease perplexity while degrading downstream task performance, and global scores can erase subtle domain-specific expert affinities.For example, the passage contrasts 3.8% of Math tokens routed to Expert A with 2.4% of Code tokens, against 3.1% global traffic.Source: §3.1 · p. 3
  • gpt-oss-20b has mean per-token entropy of 4.54 bits, or 90.8% of maximum, and stays above 82% across layers; Mixtral averages 74.3% and falls to 64–68% in later layers.gpt-oss uses λaux = 0.9 and top-4 routing; Mixtral uses λaux ≈ 0.001 and top-2 routing.Source: §3.1 · p. 4, Figure 1 · p. 4
  • High per-token entropy leaves no distinctly unimportant expert for an individual token, making global scores noisy and domain-conditional measurement essential.Pruning forces the router to redistribute mass among remaining experts, de-dispersing routing.Source: §3.1 · p. 4
  • The authors use mean entropy above 0.85 · log2 E as a soft diagnostic, not a sharp threshold, for regimes where global importance signals degrade.The cited comparison is 74.3% for Mixtral and 90.8% for gpt-oss, with an 11% margin between the two test architectures.Source: §3.1 · p. 4

3.2 MESA: Minimax Expert Score Allocation

  • MESA seeks a pruning set that minimizes the worst-case fractional routing loss across domains by iteratively adjusting importance scores for the currently worst-affected domain.The pipeline profiles domain-specific routing mass, runs the adjustment, then globally ranks experts for removal.Source: §3.2 · pp. 4–5
  • MESA is a greedy heuristic for the minimax objective and has no formal convergence guarantees.It boosts scores for experts serving the domain with the highest fractional routing loss and records the best solution found so far.Source: §3.2 · pp. 4–5

4 Experiments

Across pruning experiments, MESA most consistently limits worst-domain accuracy loss, while over-dispersed routing makes perplexity an unreliable guide. Cross-architecture results show that MESA's benefits depend on routing and evaluation conditions, and expert removal reduces memory rather than per-request latency.

4.1 Setup

  • The experiments prune 25% of experts and evaluate 11 benchmarks, using 1,000 calibration examples per domain across six domains.Some benchmarks, including MMLU-Pro and WinoGrande, lack a dedicated matching calibration domain.Source: §4.1 · p. 5

4.2 PPL–Accuracy Disconnect

  • Random pruning's PPL of 40.95 accompanies near-worst task accuracy, while MESA improves GSM8K above the unpruned baseline despite moderate PPL.On gpt-oss-20b, MESA scores 92.87% on GSM8K versus 88.86% unpruned; random pruning scores 78.54%.Source: Table 1 · p. 6, §4.2 · p. 6

4.3 Main Results

  • At 25% pruning, MESA achieves the smallest worst-case degradation and outperforms REAP on 7 of 11 benchmarks, including a 17.8-point GPQA advantage.The trade-off remains: REAP's largest advantage is 10.0 points on AIME25, and the GPQA gap is statistically significant (z = 3.57, p < 0.001 on 198 questions).Source: §4.3 · pp. 6–7, Table 2 · p. 7

4.4 Deployment Footprint

  • Removing 25% of experts reduces resident memory by approximately the same fraction, but leaves per-request decode latency essentially unaffected.Each token still routes through k expert FFNs; reducing per-token compute would require changing k or expert internals.Source: §4.4 · p. 7

4.5 Cross-Architecture Validation

  • On standard-routing Mixtral, REAP retains 98% of baseline average accuracy versus MESA's 63%, while pruning increases perplexity as expected.MESA also has the largest PPL degradation on Mixtral, unlike its performance under over-dispersed routing.Source: Table 5 · p. 8, §4.5 · p. 8

4.6 Generality to a Third Over-Dispersed Family: OLMoE-1B-7B

  • On over-dispersed OLMoE, eval-matched prefill calibration makes MESA the leading pruned method on six commonsense tasks, retaining 92.8% of unpruned accuracy.Decode-time calibration made domain routing similar and MESA's pruned sets nearly identical to RM-UNIFORM's; prefill calibration restored per-domain structure.Source: Table 6 · p. 9, §4.6 · p. 9

4.7 Generality to a Third Architecture: Gemma-4-26B-A4B-it

  • On Gemma-4-26B-A4B-it at r = 0.25 (128 routed experts, top-8; 128→96 experts per layer) and temperature 0.6, MESA matches or exceeds REAP and RM-UNIFORM on all 11 benchmarks, leading GPQA-Diamond by 14.3pp over REAP.Gemma is not over-dispersed (native per-token routing entropy ≈0.69); under deterministic temperature-0, 3-seed decoding, MESA scored 65.66 ± 1.82 versus REAP’s 55.56 ± 1.95 on GPQA-Diamond, with non-overlapping confidence intervals.Source: Table 8 · p. 10, §4.7 · p. 10

5 Analysis

The analysis finds opposing method rankings across reasoning and knowledge tasks, while MESA offers a more balanced profile and reaches stable prune maps with substantially less calibration compute than REAP.

5.1 Capability Trade-off

  • Method rankings invert between reasoning/generative and knowledge/pattern-matching task clusters (permutation test p = 0.0004), indicating a structured capability trade-off.The accuracy-matrix PCA assigns GSM8K and AIME’25 to PC1, and MMLU and ARC to PC2.Source: §5.1 · p. 11
  • MESA limits knowledge degradation to −1.68% and reasoning degradation to −4.67%, while REAP preserves reasoning (−1.45%) but degrades knowledge more (−3.48%).RM-Uniform and Random have reasoning degradation below −17%; MESA’s GSM8K accuracy improves by 4.01%.Source: §5.1 · p. 11
  • MESA’s allocation step uses a precomputed router-mass matrix rather than additional model forward passes, making it structurally cheaper than activation-based scoring such as REAP.The passage makes no wall-clock timing claim.Source: Table 10 · p. 12

5.2 Calibration Efficiency

  • MESA reaches 97.9% prune-map agreement with the reference at 1,000 samples and 8,192 tokens, using approximately 8× less calibration compute than REAP.At 20 samples per domain, agreement is 96.4%; sequence length drives stability more than sample count.Source: §5.2 · p. 12

6 Conclusion

The paper characterizes over-dispersed routing as a distinct pruning regime with unreliable perplexity and task-dependent scoring trade-offs. It presents MESA’s worst-domain objective and identifies limits in evaluation scope and pruning severity.

6 Conclusion

  • At r = 0.25, MESA minimizes worst-domain degradation and outperforms REAP on 7 of 11 benchmarks, but does not retain that advantage at r = 0.50.Under standard Mixtral routing, methods degrade proportionally and domain imbalance does not emerge.Source: §6 · p. 12
  • The evaluation covers two architectures, uses a single main-results pruning ratio of r = 0.25, and leaves more aggressive ratios and broader generalization for further study.Algorithm 1 is a greedy heuristic without formal convergence guarantees, though empirical ablations show stable convergence within 15 iterations.Source: §6 · p. 12

A PPL Confound: Top-k-Then-Softmax Mechanism

  • With gpt-oss topk-then-softmax routing, pruning reduces a near-random top-4 selection pool, improving per-token predictions while potentially sacrificing task-specific capabilities not measured by perplexity.At λaux = 0.9, all 32 expert logits are nearly equal; after pruning, the pool shrinks from 32 to 24 experts.Source: A PPL Confound: Top-k-Then-Softmax… · p. 14

B Permutation Test Details

The permutation test assesses whether opposing task-group rankings could arise by chance. Its reported results support the identified reasoning-versus-knowledge split among pruned methods.

Procedure

  • The procedure partitions 12 benchmark tasks into equal-sized groups, computes each method’s mean accuracy per group, then compares method rankings across groups.The Spearman rank correlation is averaged over all 924 partitions to form the null distribution.Source: B Permutation Test Details · pp. 14–15

Results

  • Only 0.04% of the 924 task groupings produce correlations as negative as the observed reasoning-versus-knowledge clusters (p = 0.0004).The test compares GSM8K and AIME25 with MMLU, ARC-Challenge, and ARC-Easy among pruned methods.Source: B Permutation Test Details · p. 15
  • Including the unpruned baseline raises the average cross-group correlation, but the ranking inversion among pruned methods remains statistically significant.The baseline dominates both cluster means and compresses the anti-correlation.Source: B Permutation Test Details · p. 15

C Extended Capability Frontier Analysis

Task-loading patterns reveal a sharp split between knowledge/pattern-matching and reasoning/generative benchmarks, with pruning methods trading off across those clusters. MESA balances the trade-off, while the following ablation section provides details on its calibration and hyperparameter results.

C Extended Capability Frontier Analysis

  • Task loadings separate into knowledge/pattern-matching tasks and reasoning/generative tasks, with the latter cluster including GSM8K, AIME25, GPQA, IFEval, WinoGrande, and OpenBookQA.The clusters emerge from per-method accuracy patterns rather than being imposed in advance.Source: C Extended Capability Frontier Analysis · p. 15
  • Pruned-method rankings reverse across clusters: REAP leads reasoning/generative tasks, while EvoESAP leads knowledge/pattern-matching tasks.The cluster-mean correlation is strongly negative, though the anti-correlation weakens when the unpruned baseline is included.Source: C Extended Capability Frontier Analysis · p. 15
  • On GPQA-Diamond, MESA scores 65.48% (+1.0% over baseline), while REAP scores 47.72% (−16.8%).GPQA-Diamond combines graduate-level factual knowledge with multi-step reasoning and has the second-largest accuracy range across pruning methods, at 21%.Source: C Extended Capability Frontier Analysis · p. 15
  • With router mass nearly uniform, REAP’s activation-magnitude scoring favors computation-heavy experts, while router-mass scoring favors high-volume “highway” experts.The paper attributes this divergence to λaux = 0.9 decorrelating activation norms from routing frequencies; MESA instead optimizes worst-case routing perturbation across domains.Source: C Extended Capability Frontier Analysis · p. 15

D MESA Hyperparameter Ablations

D.1 Calibration Data Requirements

MESA’s calibration selections depend more on sequence length than sample count, and calibration budgets can shift the capability trade-off. Its iterations and boost rate are robust over tested settings, while near-uniform routing and substantial cross-domain expert overlap make domain-fair pruning challenging.

D.1 Calibration Data Requirements

  • At r=0.25, calibration requires sequence lengths of at least 1024 tokens for stable routing statistics, while 20 samples suffice when sequence length is adequate.Prune-map agreement measures 192 binary expert decisions; calibration budgets can explore different regions of the capability Pareto frontier.Source: §D.1 · p. 15, Table 11 · p. 16
  • At 1000 samples, increasing sequence length from 4096 to 8192 raises prune-map agreement by 15.0%, exceeding the 5.8% gain from 20 to 200 samples at length 2048.Source: §D.1 · p. 16
  • Higher calibration shifts MESA toward domain balance, while lower calibration preserves broader knowledge; the passage recommends choosing the budget by downstream priority.The reported high-calibration configuration gains 2% on PIQA and loses 4% on MMLU-Pro relative to the lower-calibration setting.Source: §D.1 · p. 16

D.2 Iterations and Boost Rate

  • At r=0.125 and r=0.25, MESA converges within 6–12 steps; at r=0.375, it takes 20 steps, and higher pruning ratios produce higher loss plateaus.Linear, cosine, and constant decay schedules reach identical final losses within numerical precision, converging in 8–17 steps; convergence is effectively reached by T=15.Source: §D.2 · p. 16, Table 12 · p. 17
  • At r=0.25, boost rates η0 in [0.05, 0.60] reach the same worst-domain loss of 0.0491, with at least 95.9% Jaccard agreement against the default.Oscillation increases with η0, but the best-snapshot mechanism preserves the final solution; the recommended default is η0 = 0.15.Source: §D.2 · p. 16

E.1 Per-Domain Routing Mass

  • Figure 5’s near-uniform importance grid compresses domain-specific routing variation into a 0.01 dynamic range, which global scores average away.MESA’s per-domain measurement preserves these weak signals.Source: §E.1 · p. 17

E.2 Cross-Domain Expert Overlap

  • Top-8 expert sets for pairs of domains overlap by about 45%, making separate per-domain top-expert selections nearly identical.MESA targets marginal experts in the 0.41–0.47 overlap zone, where domains partially diverge.Source: §E.2 · p. 17

E.3 Per-Task Degradation Patterns

At 25% pruning, per-task accuracy changes show a structured, task-specific degradation pattern rather than a uniform benchmark effect. The analysis also varies routing entropy to examine how global importance-signal quality changes between routing regimes.

E.3 Per-Task Degradation Patterns

  • At r = 0.25, per-task accuracy changes from baseline across methods show that capability degradation is task-specific rather than uniform across benchmarks.Source: §E.3 · p. 17
  • MESA converges within 6–20 iterations for r ≤0.375; higher pruning ratios yield higher loss plateaus despite similar convergence dynamics.The figure caption suggests the higher-ratio difficulty lies in the objective landscape rather than the optimizer.Source: Figure 4 · p. 18

F Regime Transition: Importance Signal Degradation

G Qualitative Analysis: Knowledge Degradation Under Pruning

A qualitative electrostatics example illustrates a specific knowledge failure in a REAP-pruned model despite correct problem setup. A separate Gemma-4 replication tests whether a worst-domain protection result depends on decoding temperature.

G Qualitative Analysis: Knowledge Degradation Under Pruning

  • In the electrostatics example, the baseline and REAP-pruned model set up the method of images correctly, but REAP misses the factor of 1/2 for induced-charge energy and selects the wrong formula.Source: G Qualitative Analysis: Knowledge… · p. 18

H Gemma-4 Temperature-0 Replication

  • The Gemma-4-26B-A4B-it flagship result uses temperature 0.6, matching the paper’s other results and preserving chain-of-thought.Source: H Gemma-4 Temperature-0 Replication · p. 18
  • A temperature-0 replication independently tests the same r=0.25 head-to-head using 3-seed mean ± std on the available task subset.The authors disclose both temperatures rather than reporting only the more favorable one.Source: H Gemma-4 Temperature-0 Replication · p. 18

I gpt-oss-120B at r = 0.5: Scale and Aggressive-Ratio Generality

At 50% pruning, MESA retains worst-domain protection on gpt-oss-120B but does not win every benchmark. The results frame this as an aggressive-pruning stress test, not evidence of average-case superiority.

I gpt-oss-120B at r = 0.5: Scale and Aggressive-Ratio Generality

  • Across gpt-oss-20B routing, no expert dominates any domain, yet weak domain-specific mass differences remain; domains share 0.41–0.47 of their top-eight experts per layer.The weak residual routing structure is visible alongside near-uniform routing.Source: Figure 5 · p. 19, Figure 6 · p. 20
  • At r = 0.5, the authors claim only worst-case protection, not average-case superiority, and report the comparison as a stress test.This is substantially heavier pruning and larger scale than the main setting.Source: I gpt-oss-120B at r = 0.5: Scale and… · pp. 18–19
  • Gemma-4-26B-A4B-it has a separate temperature-0 replication over three seeds on five tasks; additional benchmarks await temperature-0 replication.This replication is reported alongside, not instead of, the temperature-0.6 result.Source: Table 13 · p. 22
  • A domain-aware REAP control scores GPQA-Diamond 58.4 versus MESA’s 65.66 at temperature 0 and 67.85 at temperature 0.6, attributing the gain to the minimax objective.The comparison tests domain conditioning separately from MESA’s worst-case objective.Source: I gpt-oss-120B at r = 0.5: Scale and… · p. 22
  • At r = 0.5 on gpt-oss-120B, MESA leads REAP on GPQA-Diamond (58.46 vs. 51.85) and five other benchmarks, while REAP wins five.MESA’s worst-domain protection holds under this heavier pruning ratio, but it does not dominate every benchmark.Source: I gpt-oss-120B at r = 0.5: Scale and… · pp. 18–19, Table 14 · p. 22
  • RM-UNIFORM scores 0.00 on every task at r = 0.5 on gpt-oss-120B, a genuine method failure rather than missing measurements.The table reports one measurement per cell.Source: I gpt-oss-120B at r = 0.5: Scale and… · p. 22, Table 14 · p. 22
  • The gpt-oss-120B r = 0.5 comparison is single-seed; multi-seed confirmation remains future work.Table 14 reports a single measurement per cell.Source: I gpt-oss-120B at r = 0.5: Scale and… · p. 22, Table 14 · p. 22

Related papers

Loading 2609.04453v1…