Smoothing the Top-k Exposure Boundary for Sparse Mixture-of-Experts

1 Introduction

• Static top-k turns continuous routing scores into a hard boundary: experts just inside the cutoff receive full task-loss feedback, while adjacent experts outside receive none.The paper describes this boundary as brittle because small score fluctuations can separate highly competitive experts.Source: §1 · pp. 1–2
• Dynamic-routing alternatives can remove the top-k threshold, but may change deployment budgets, alter model structure, or impair task performance at extreme efficiency.The paper notes these changes can make such approaches difficult to apply on model-serving infrastructure.Source: §1 · p. 2
• Elastic Expert Routing samples the active expert budget from a localized discrete distribution centered at k, smoothing near-boundary selection over training iterations.Its symmetric sampling neighborhood matches the expected compute of static top-k training, while inference retains static top-k with no structural modification or latency overhead.Source: §1 · p. 2
• +0.84 and +2.02 points on OLMoE-1B-7B and Qwen3-30B-A3B, respectively, and +1.6 points on average across downstream pretraining tasks.The reported fine-tuning gains are over static top-k baselines; pretraining is evaluated under compute-matched and token-matched conditions.Source: §1 · p. 2
• Routing diagnostics associate the gains with middle-layer routing changes and improved expert prioritization.The introduction reports that elastic training refines prioritization rankings and changes routing distributions in foundational middle layers.Source: §1 · p. 2

2 Preliminary and Related Work

• Sparse MoE scales model capacity while keeping per-token computation sparse, but fixed k creates a sharp exposure boundary between selected experts and nearby unselected experts.The paper reframes k as both a compute constraint and a choice that shapes training exposure.Source: §2 · p. 2
• MoE training uses auxiliary objectives, capacity constraints, and routing regularization, while other approaches change token–expert assignment or use alternative selection mechanisms.The paper contrasts these approaches with changing the training distribution over the active expert budget while preserving the routing architecture and assignment mechanism.Source: §2 · pp. 2–3
• Dynamic-routing methods vary token-level expert budgets through thresholds, proposer modules, or architectural designs, whereas this work varies the budget only during training and keeps static top-k inference.The paper characterizes existing dynamic approaches as requiring structural modifications or complex training techniques that can hinder deployment and stable training.Source: §2 · p. 3

3.1 Budget-Neighborhood Training

• Sampling budget schedules from a distribution centered at k0 changes training-time configurations while preserving the expected active-expert count.The fixed-budget objective is recovered when the distribution is a point mass at k0.Source: §3.1 · p. 3

3.2 Instantiating the Budget Neighborhood

• An a–b schedule samples budgets from {a, ..., b} for each token and layer, then dispatches each token to its top-κ experts.The experiments use σ = 1.5 unless otherwise stated; load-balancing statistics use realized assignments.Source: §3.2 · p. 3

3.3 Rank-Neighborhood Smoothing

• Elastic routing replaces the static top-k step mask with a budget-distribution survival function, smoothing exposure locally across router ranks.Ranks at or below the minimum budget are always exposed, ranks within the neighborhood have fractional exposure, and ranks above its maximum are never exposed. Centering the distribution at k0 preserves expected active-expert count; wider neighborhoods shift more exposure to lower-ranked candidates and shrink the always-on core.Source: §3.3 · pp. 3–4

4.1 Setup

• The setup compares static top-8 with token-level elastic schedules across two SFT settings and one pretraining setting, matching expected expert count and training FLOPs and using static top-8 inference.SFT reports Avg(9); pretraining reports an average across downstream tasks.Source: §4.1 · p. 4

4.2 Supervised Fine-Tuning Results

• 6–10 raises OLMoE Avg(9) from 30.57 to 31.41 and Qwen3-30B-A3B Avg(9) from 56.42 to 58.44.An elastic schedule also scores best on all nine OLMoE tasks and eight of nine Qwen3 tasks; 6–10 gives the strongest aggregate result for both backbones.Source: §4.2 · p. 4, Table 1 · p. 5

4.3 From-Scratch Pretraining Validation

• Under matched compute, 6–10 raises the from-scratch pretraining average from 41.3 to 42.9; 7–9 reaches 42.8 and 5–11 reaches 42.4.The 6–10 schedule also reduces LAMBADA perplexity from 54.2 to 42.7.Source: Table 2 · p. 5, §4.3 · p. 5

4.4 Depth-Localized Routing Changes

• L3–L7 had the largest non-boundary router divergence; restricting 6–10 there slightly exceeded global 6–10, while L8–L12 gained less.L0 and L15 also had elevated divergence but were excluded from the diagnostic block because of input-adaptation and final-readout effects; L3–L7 was exploratory, not established as optimal.Source: Figure 2 · p. 6, §4.4 · p. 6, Table 3 · p. 7

4.5 Expert-Selection Diagnostics and Router Intervention

• 6–10 changed the top-1 expert for 13.62% of tokens versus 7.08% of top-8 sets, or about 0.57 experts out of eight.Router swaps did not recover full 6–10 performance, suggesting gains involve router–expert co-adaptation rather than the router alone.Source: Figure 3 · p. 6, §4.5 · p. 7, Table 4 · p. 7

4.6 Performance Under Different Top-k Values

• At inference top-k values across the scanned range, 6–10 remains above static top-8 and is strongest at the default k = 8 point.The 7–9 and 5–11 schedules improve some parts of the curve but are less consistently strong.Source: §4.6 · p. 7, Figure 4 · p. 8

4.7 Ablations and Budget Controls

• With 6–10 support fixed, σ = 1.5 outperforms σ = 1.0 and σ = 3.0; static top-9, top-10, and top-p also trail 6–10 despite using more expert computation.The shape results favor neither a distribution concentrated near one cutoff nor one flattened across ranks 6–10; the budget controls suggest gains are not explained solely by activating more experts.Source: Table 5 · p. 8, §4.7 · p. 8, Table 6 · p. 8

5 Conclusion

• Elastic Expert Routing smooths the selection boundary by sampling a discrete expert budget during training while preserving static top-k inference.The conclusion describes it as a drop-in training strategy for fixed-budget sparse MoE architectures.Source: §5 · p. 8
• The method consistently improves MoE performance in supervised fine-tuning and from-scratch pretraining without increasing inference cost.Source: §5 · p. 8
• Elastic training primarily changes routing policy in foundational middle layers and refines expert rankings near the deployment cutoff.Source: §5 · p. 8

Limitations

• From-scratch pretraining uses a 1.427B-parameter MoE with 0.365B active parameters, smaller than the fine-tuning backbones.Source: Limitations · p. 9
• Fine-tuning uses OLMoE-1B-7B with 7B total parameters and 1B active per token, and Qwen3-30B-A3B with 30.5B total and 3.3B activated.Source: Limitations · p. 9
• The authors frame pretraining findings as controlled validation at moderate scale; validation under large-scale from-scratch pretraining remains future work.Source: Limitations · p. 9

A Experimental Details

• Table A.1 is labeled as defining shared evaluation and routing conventions.Source: Table A.1 · p. 11
• Within each setting, routing schedules share the model, tokenizer, data, optimizer, evaluation harness, and decoding protocol.Source: A Experimental Details · p. 11
• The training-time routing schedule is the only experimental component that changes within a setting.Source: A Experimental Details · p. 11

A.1 Shared Protocol

• Avg(9) is computed over the nine reported benchmark tasks listed in the section.Source: §A.1 · p. 11
• Table A.1 defines the common evaluation surface for the supervised fine-tuning experiments.Source: §A.1 · p. 11
• The protocol’s aggregate measure is labeled Avg(9).Source: §A.1 · p. 11

• Measured training costs are reported for each routing condition in Table A.2.The measurements complement expected expert-side FLOP matching used in the main experiments.Source: §A.2 · p. 11
• Elastic schedules are centered at eight experts.Source: §A.2 · p. 11
• Measured wall-clock behavior also depends on dispatcher, communication, padding, and memory-allocation details.Source: §A.2 · p. 11

• Each elastic schedule is compared only with the static top-8 baseline within the same backbone and training recipe.The two SFT settings use different infrastructure because they start from different sparse MoE backbones.Source: §A.3 · p. 11
• OLMoE fine-tuning uses Dolci-Instruct-SFT for four epochs, with AdamW-family optimization, bf16 precision, batch size 256, sequence length 4096, and 16 GPUs.Dynamic budget sampling uses σ = 1.5.Source: §A.3 · pp. 11–12
• Qwen3 fine-tuning uses Infinity-Instruct 7M Core for three epochs with ms-swift and Megatron-LM, and both SFT settings use ListenEval over the shared Avg(9) suite.Qwen3 uses batch size 16, sequence length 8192, and 16 GPUs; dynamic budget sampling uses σ = 1.5.Source: §A.3 · p. 12

A.4 From-Scratch Pretraining

• The pretraining setup uses C ≈ 6.0×10^19, with learning rate 1.09516×10^-3 and global batch size 256.The compute budget is defined as C = M · D, where M is non-embedding forward FLOPs per token and D is training tokens; the selected batch corresponds to 1.05M tokens per update.Source: §A.4 · p. 12
• Pretraining compares static top-8 routing with elastic schedules 7–9, 6–10, and 5–11, using σ = 1.5 and static top-8 evaluation.Evaluation uses the frozen validation split and validation suite in Table 2.Source: §A.4 · p. 12

A.5 Full Task-Level Ablation Results

• Tables A.4 and A.5 provide full task-level breakdowns for the distribution-shape and alternative-budget-control ablations, respectively.The breakdowns use the evaluation suite and averaging protocol of their corresponding main-text tables.Source: §A.5 · p. 12, Table A.4 · p. 13, Table A.5 · p. 13