Source-linked AI summary
Distribution-Consistent Inference for Dynamic Sparse Mixture-of-Experts
Dohyeon Kim, Bedionita Soro, Sung Ju Hwang
TL;DR
Fixed expert budgets can waste computation, while reducing activation may shift SMoE representations away from the training-time distribution. The paper proposes LDA, a training-free layer-wise correction, and reports improved efficiency–performance trade-offs across routing strategies, while noting that capacity loss and manually selected thresholds remain limitations.
Problem
Fixed top-k routing gives every token the same expert budget, while reduced routing can deviate from the training-time representation distribution.
Method
LDA uses layer-wise calibration statistics to align reduced-routing representations with the default top-k0 distribution without retraining or architectural changes.
Results
LDA consistently improves the performance–efficiency trade-off across diverse tasks, models, and top-k, top-p, and pruning strategies.
Takeaways & Limitations
Representation-level distribution shift is a correctable component of degradation under reduced expert activation, complementing routing-based efficiency methods.
Takeaways & Limitations
LDA cannot recover expert capacity lost through fewer active experts, and its dynamic top-p experiments require manually selected thresholds per task.
Abstract
from arXiv · showhide
Mixture-of-Experts (MoE) architectures have emerged as a powerful paradigm for scaling model capacity while preserving efficient inference in large foundation models. However, most MoE models use a fixed top-$k$ expert selection policy, assigning the same expert budget to every token even when fewer experts may be sufficient. Inference-time dynamic top-$k$ routing can reduce computation without retraining, but existing methods often overlook the distributional shift caused by deviating from the training-time routing configuration. We show that reducing the number of activated experts consistently increases the RMS scale and variance of SMoE outputs, inducing a representation mismatch that contributes to downstream performance degradation in addition to the loss of expert capacity. To address this correctable component, we propose Layer-wise Distribution Alignment (LDA), a lightweight inference-time correction that uses layer-wise calibration statistics to align reduced-routing representations with the default configuration. Across multiple SMoE LLMs, benchmarks, and routing strategies, LDA recovers much of the performance lost induced by the distributional shift under reduced routing while preserving sparse-inference efficiency with negligible overhead.
1 Introduction
SMoE models commonly use fixed top-k routing, although dynamic inference-time routing can reduce unnecessary computation. This paper identifies representation-level distributional shift under reduced routing and proposes LDA to correct it without retraining.
- Fixed top-k routing assigns every token the same expert budget, even when fewer experts may suffice.
- Inference-time dynamic routing reduces activated experts without retraining, but existing methods overlook shifts from the training-time routing distribution.
- Reducing top-k consistently increases SMoE output scale and variance, creating a representation mismatch beyond diminished expert capacity.
- Layer-wise Distribution Alignment uses calibration statistics to align reduced-routing representations with the default routing distribution.
- LDA consistently improves the performance–efficiency trade-off across tasks, models, and top-k, top-p, and pruning strategies.
2 Related Works
Prior SMoE work mainly develops routing strategies for reducing expert computation. This paper instead studies the output-distribution mismatch caused by reduced activation and introduces a complementary training-free correction.
- SMoE selectively activates a small expert subset to expand total parameters while controlling per-token computation.
- Existing methods improve efficiency through confidence-based allocation, null experts, adaptive top-k selection, or training-free expert pruning.
- Prior routing strategies leave underexplored how reducing activated experts changes SMoE output distributions.
- The proposed correction addresses representation mismatch and complements existing training-free dynamic routing strategies.
3 Observation and Analysis
Reducing top-k routing changes SMoE output distributions: variance and RMS scale increase, residual-stream balance shifts, and routing trajectories diverge from the default configuration. The analysis attributes these effects to score re-normalization and heterogeneous per-dimension distribution changes.
- 3.1 Preliminaries: Under top-k routing, selected experts are re-normalized before their weighted outputs are summed into the SMoE layer output.The selected set is determined by the highest routing scores, and the resulting layer output is denoted y.
- 3.2 Observation: Reduced top-k Amplifies Representation Variance and Scale: Reducing top-k monotonically increases SMoE-output variance and RMS magnitude, changing the scale of the resulting representation.The output distribution becomes more dispersed and larger relative to the residual stream as fewer experts are activated.
- 3.2 Observation: Reduced top-k Amplifies Representation Variance and Scale: Score re-normalization explains RMS amplification because removing lower-ranked experts increases the sum of squared normalized routing scores.Under the stated routing-score assumptions, this produces a larger expected squared norm for the SMoE output.
- 3.3 Structural Analysis: In residual blocks, an enlarged SMoE output increases RMSNorm’s denominator and can relatively suppress useful residual-stream components.The combined representation z = r + y becomes more strongly influenced by the weighted expert output before normalization.
- 3.3 Structural Analysis: Variance amplification is heterogeneous across dimensions, and Table 1 identifies variance as the moment with the most consistent increase as k decreases.The per-dimension distributions therefore cannot be fully characterized by a single global scaling factor.
- 3.3 Structural Analysis: Reduced top-k lowers routing match rate across layers, causing the subsequent routing trajectory to diverge from default top-k0 routing.The divergence is especially pronounced at smaller k.
4 Layer-wise Distribution Alignment (LDA)
LDA aligns reduced-routing SMoE representations with default-routing layer-wise statistics using a lightweight, per-dimension affine correction. It targets heterogeneous distribution shifts while requiring calibration but no additional training or parameter updates.
- Alignment design: LDA uses default top-k0 layer-wise representation statistics as references for aligning reduced top-k outputs.Calibration estimates per-dimension statistics under default and reduced routing, with preceding layers propagated under default routing.
- Alignment design: The correction operates in the original representation space because default and reduced-routing outputs share that space.It avoids an additional projection or nonlinear transformation.
- Alignment design: Per-dimension correction addresses heterogeneous variance shifts that global scaling cannot capture.The method aligns per-dimension means and standard deviations rather than full covariance or higher-order statistics.
- Alignment design: LDA applies an element-wise affine transformation that restores first- and diagonal second-order statistics toward the default configuration.The operation requires negligible inference overhead and no additional training or parameter updates.
- Compatibility: LDA is compatible with existing training-free dynamic routing because it operates independently of how reduced expert counts are selected.The method is defined layer-wise and can be combined with alternative routing policies.
5 Experiments
Experiments across multiple SMoE models, routing strategies, benchmarks, and serving settings show that LDA mitigates reduced-routing degradation and improves the performance–efficiency trade-off with negligible overhead.
- Evaluation scope: Across fixed top-k, dynamic top-p, and PESF routing, performance degradation becomes more pronounced as the average number of activated experts decreases.The evaluation covers multiple models and 13 downstream benchmarks.
- Routing strategies: LDA substantially improves reduced top-k performance, particularly when k is 4 or 2, while retaining the same expert budget.The results associate remaining degradation with representation variance and scale mismatch in addition to reduced capacity.
- Routing strategies: With an appropriate top-p threshold, LDA can match or exceed default top-k0 performance while using fewer experts on average.LDA also improves the corresponding PESF dynamic-pruning baseline.
- Representation effects: LDA aligns reduced top-k output distributions, suppresses increased output scale relative to the residual stream, and raises routing match rates across layers.Figure 5 uses 2,048 held-out calibration tokens and samples 100,000 values for panel (a).
- Cross-model results: The same representation-level tendency appears across the other evaluated SMoE architectures.The paper reports consistent improvements under reduced expert activation budgets beyond the primary model.
- Efficiency trade-off: Reducing activated experts lowers FLOPs and raises throughput, but vanilla routing incurs substantial performance loss, especially at smaller k.LDA preserves much of the task performance at comparable inference cost, with additional cost O(LD).
6 Conclusion
The paper identifies scale and distribution mismatch as a correctable component of reduced-activation degradation and proposes LDA as a lightweight, training-free correction. Across models, tasks, and routing strategies, LDA consistently mitigates degradation and improves the performance–efficiency trade-off.
- Conclusion: Reducing top-k amplifies SMoE output scale and variance, disturbing the balance between weighted expert outputs and the residual stream.The paper distinguishes this routing-induced mismatch from the loss of expert capacity.
- Conclusion: LDA aligns reduced top-k representations with default layer-wise statistics using an element-wise affine transformation with additional cost O(LD).It is training-free and compatible with training-free dynamic routing strategies.
- Conclusion: Experiments across multiple SMoE LLMs and downstream tasks show consistent mitigation of reduced-activation degradation and improved performance–efficiency trade-offs.The conclusion covers multiple routing strategies.
Limitations
LDA corrects routing-induced representation mismatch but cannot recover lost expert capacity, uses limited per-dimension statistics, and requires calibration. Dynamic top-p experiments additionally rely on manually selected task-specific thresholds.
- Method limitations: LDA cannot recover expert capacity lost by activating fewer experts, so extreme reduction may leave a capacity gap and only partially recover default performance.This limitation becomes relevant when expert activation is reduced substantially.
- Method limitations: Per-dimension mean and standard deviation alignment ignores covariance structure, asymmetry, and higher-order statistics.More expressive alignment may help in certain layers or tasks.
- Method limitations: LDA requires a calibration stage to estimate layer-wise reference and target statistics.The authors report robustness to calibration-set choice and size, but the method is not calibration-free.
- Routing limitations: Dynamic top-p experiments use manually selected thresholds that must be chosen per task.Automatically determining activated expert counts with training-free routing remains future work.
A.2 Empirical Support for Mild Assumptions
The analysis supports the assumptions underlying the RMS amplification mechanism: expert-output norms remain comparable, while normalized routing scores become more concentrated and cross terms remain subordinate as k decreases.
- Empirical support for assumptions: Activated expert-output norms remain within a comparable range across top-k settings, supporting limited variation in expert-output scale.This directly supports the assumption that activated expert outputs have comparable scales.
- Empirical support for assumptions: The sum of squared normalized routing scores increases substantially as k decreases, indicating greater concentration of normalized routing scores.
- Empirical support for assumptions: The cross term remains substantially smaller than the first diagonal term across evaluated layers and top-k settings.Its relative contribution further decreases as k decreases, supporting dominance of the first diagonal term in the approximation.
- Experimental setup: The experiments compare fixed top-k, dynamic top-p, and PESF routing, evaluating reduced activation with and without LDA.Top-p thresholds preserve the inference-efficiency setting by capping activated experts at the default top-k0 configuration.
- Experimental setup: LDA is evaluated across recent SMoE LLMs with diverse architectures, including Qwen3-30B-A3B, kanana-2-30b-a3b-instruct, and GLM-4.7-Flash.The models differ in expert configurations and gating functions, while all use routing-score re-normalization.
B.3 Benchmarks
The evaluation spans general knowledge, mathematical reasoning, code generation, and instruction-following benchmarks, using deterministic option-based and generation-based protocols across routing strategies and models.
- Benchmark coverage: The benchmark suite covers general knowledge and commonsense reasoning, mathematical reasoning, code generation, and instruction-following.Tasks include MMLU, GSM8K, MATH, MBPP, HumanEval, and IFEval, among others.
- Evaluation protocols: Option-based tasks use deterministic candidate selection from final-logit probabilities, while generation-based tasks use greedy decoding and task-specific answer extraction.
- Evaluation protocols: Top-p routing thresholds are swept over p ∈{0.1, 0.2, . . . , 0.9}, with p* selected as the best-performing threshold for each task.
- Statistical tests: McNemar’s test is used for option-based tasks, whereas paired t-tests compare generation-based performance across five random seeds.
- Cross-model evaluation: Across additional architectures, reduced expert activation degrades vanilla routing performance while LDA mitigates the degradation under different routing strategies.The additional models differ in expert configurations, shared experts, and gating functions.
D.1 kanana-2-30b-a3b-instruct
On kanana-2-30b-a3b-instruct, reducing expert activation degrades vanilla performance, while LDA consistently improves reduced-routing results across fixed, dynamic, and pruning-based strategies.
- Fixed top-k routing: Under fixed top-k routing, performance decreases as activated experts are reduced, with larger degradation at smaller k values.
- Fixed top-k routing: LDA consistently improves performance over the corresponding reduced top-k baseline across all task categories.
- Dynamic top-p routing: Under top-p routing, lowering p reduces average activated experts and degrades vanilla performance, whereas LDA improves performance at the same threshold.
- Dynamic top-p routing: At the best-performing top-p threshold, LDA achieves higher performance than the default top-k0 baseline across diverse tasks.
- Dynamic expert pruning: LDA also consistently improves performance when combined with PESF dynamic expert pruning.This extends the observed mitigation beyond fixed reduced top-k and dynamic top-p routing.
- Overall pattern: Across different architectures and routing strategies, LDA improves reduced-routing performance, supporting representation-scale and distribution consistency under reduced activation.
E Comparison with Simple Scaling Baselines
Comparisons with global scaling baselines indicate that reduced-routing mismatch is not only a scalar scale shift: LDA’s per-dimension alignment generally performs better, although its calibration depends on calibration statistics.
- Scaling baselines: Expert-Count Scaling applies a multiplier based on the ratio of default to reduced expert counts, testing whether expert-count correction alone can recover performance.
- Scaling baselines: RMS Scaling uses layer- and top-k-specific average RMS statistics to correct output scale, but remains a single-scalar correction.
- Baseline results: Expert-Count Scaling substantially degrades performance relative to vanilla reduced top-k, indicating that routing effects are not simply proportional to expert count.
- Baseline results: RMS Scaling substantially mitigates reduced-routing degradation, especially in low-k settings, supporting the role of SMoE output scale.
- LDA comparison: LDA generally outperforms simple scaling baselines by aligning per-dimension means and standard deviations with the default top-k0 configuration.This suggests reduced activation produces dimension-wise distributional changes beyond a global scale shift.
- Calibration dependence: LDA’s effectiveness may depend on calibration-dataset choice and sample count because its inference-time correction uses calibration statistics.The paper evaluates robustness across calibration domains and sample sizes.
F.1 Effect of Calibration Dataset Choice
LDA performance remains stable across calibration datasets and sample sizes, indicating that calibration need not be task-specific or large. The method estimates layer-wise routing statistics and aligns reduced-top-k outputs toward the default distribution.
- Calibration dataset choice: Calibration with C4 performs comparably to task-specific or domain-specific datasets across MMLU, GSM8K, and MBPP.No calibration dataset consistently outperforms the others.
- Cross-architecture evidence: Across kanana-2-30b-a3b-instruct and GLM-4.7-Flash, LDA restores output variance and average scale ratio toward the default configuration while improving routing match rate.These architectures show the same overall tendency as Qwen3-30B-A3B.
- Calibration dataset choice: Using same-domain calibration data does not necessarily improve downstream performance.This supports calibration as distribution estimation rather than task-specific adaptation.
- Calibration sample size: Small calibration sets provide competitive results, while increasing calibration tokens yields no consistent performance improvement.Layer-wise means and standard deviations can be estimated sufficiently well from a small calibration set.
- LDA procedure: LDA estimates layer-wise statistics for each top-k setting and applies them during inference when fewer experts than k0 are activated.The procedure uses a calibration dataset to estimate statistics and applies per-layer alignment under reduced routing.
- Distribution alignment: Reduced top-k changes SMoE output distributions across layers, while LDA aligns them toward the default top-k0 configuration.The before-and-after distributions are shown for Qwen3-30B-A3B using held-out calibration tokens.