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

Summary (Overview)

• Core Problem: Traditional sparse Mixture-of-Experts (MoE) models use a static top-k expert selection that creates a rigid step-function boundary, arbitrarily separating highly competitive experts into "full-supervision" and "zero-feedback" zones based on minor score fluctuations.

• Proposed Solution: The authors introduce Elastic Expert Routing, which stochastically samples the active expert budget from a localized discrete distribution centered at k during training, softening the sharp threshold into a gradual selection probability distribution while preserving the expected computational cost and inference-time budget.

• Key Results: Elastic routing improves downstream macro-averages by +0.84 points on OLMoE-1B-7B and +2.02 points on Qwen3-30B-A3B during supervised fine-tuning. In from-scratch pretraining, it reduces validation LM loss by up to 0.04 and outperforms static top-k by 1.6 points on average across downstream tasks.

• Robustness: The method outperforms alternative MoE training approaches (Noisy top-k, expert FFN dropout, and DefaultMoE) under the same static top-8 inference rule.

• Mechanistic Insights: Diagnostics reveal that elastic routing primarily refines expert prioritization rankings (changing top-1 expert for ~13.6% of tokens while only altering ~7% of the top-8 set composition) and produces the most significant routing changes in middle layers (L3–L7).

---

Introduction and Theoretical Foundation

Background and Motivation

Sparse Mixture-of-Experts (MoE) models scale model capacity efficiently by activating only a restricted subset of experts per token. In the common top-k architecture, a routing network ranks experts for each token and dispatches to the first k candidates. The choice of k serves as a static computational hyperparameter, providing intuitive control over per-token computational requirements.

However, this fixed threshold introduces a critical problem: it converts a continuous routing distribution into a rigid step function at the selection threshold. This creates a brittle boundary where:

• Experts ranked just inside the cutoff receive full gradient supervision
• Adjacent competitors ranked just outside receive zero task-loss feedback
• Minor score fluctuations can arbitrarily separate highly competitive experts

Theoretical Foundation

The paper formalizes this problem through rank-exposure accounting. For a token representation $h$, let $\pi_j(h)$ denote the expert at rank $j$. The top-r prefix mask is defined as:

$$m_j^{(r)} = \mathbf{1}[j \leq r]$$

For a scalar budget distribution $q$, the task-loss exposure probability of rank $j$ is:

$$a_q(j) = \mathbb{E}_{K \sim q}\left[m_j^{(K)}\right] = \Pr(K \geq j) = \sum_{r=j}^{K_{\max}} q(r)$$

Static top-k training corresponds to the point-mass case $a_{\delta_{k_0}}(j) = \mathbf{1}[j \leq k_0]$, whereas elastic routing replaces this step mask with the survival function of $q$.

The budget distribution is the discrete derivative of the exposure curve:

$$a_q(j) - a_q(j+1) = q(j), \quad a_q(M+1) = 0$$

This yields expected exposure conservation:

$$\sum_{j=1}^{M} a_q(j) = \sum_{j=1}^{M} \Pr(K \geq j) = \mathbb{E}[K]$$

When $q$ is exactly centered at the deployment budget, $\mathbb{E}[K] = k_0$, elastic routing preserves the expected active-expert count while changing which boundary ranks enter the task-loss path.

---

Methodology

Budget-Neighborhood Training

The paper formulates elastic routing as a modification to the training objective. Let $\mathcal{D}$ denote the training distribution and $g(\hat{y}, y)$ the task loss. For an input $x$, let $f_\theta(x; \kappa)$ denote the MoE forward pass under a budget schedule $\kappa$, where $\kappa_{i,\ell}$ denotes the number of experts used by token $i$ in MoE layer $\ell$.

Fixed-budget baseline (static top-k):

$$\mathcal{L}_{\text{fixed}}(\theta) = \mathbb{E}_{(x,y) \sim \mathcal{D}}\left[g(f_\theta(x; \mathbf{k}_0), y)\right], \quad \mathbf{k}_0 \equiv (k_0, \ldots, k_0)$$

Elastic routing objective:

$$\mathcal{L}_{\text{elastic}}(\theta) = \mathbb{E}_{(x,y) \sim \mathcal{D}} \mathbb{E}_{\boldsymbol{\kappa} \sim Q}\left[g(f_\theta(x; \boldsymbol{\kappa}), y)\right]$$

The fixed-budget objective is recovered when $Q$ degenerates to a point mass at $\mathbf{k}_0$.

Instantiating the Budget Neighborhood

The budget distribution uses a truncated discrete Gaussian. An $a$–$b$ schedule restricts the sampled budget to $\mathcal{K}_{a:b} = \{a, \ldots, b\}$ and assigns probability:

$$q_{a:b}(r) = \frac{1}{Z_{a:b}} \exp\left(-\frac{(r - k_0)^2}{2\sigma^2}\right), \quad r \in \mathcal{K}_{a:b}$$ $$Z_{a:b} = \sum_{s=a}^{b} \exp\left(-\frac{(s - k_0)^2}{2\sigma^2}\right)$$

Here $r$ denotes a candidate active-expert budget, $k_0$ is the center, and $\sigma$ controls the neighborhood width. Unless otherwise stated, all experiments use $\sigma = 1.5$. For each token $i$ and MoE layer $\ell$, the method samples $\kappa_{i,\ell} \sim q_{a:b}$ and dispatches the token to the top-$\kappa_{i,\ell}$ experts under the current router ranking.

Key Properties

1. Expected computational cost matching: Since sampling neighborhoods remain symmetric around $k_0$, the expected active-expert count equals $k_0$, matching deterministic training costs.

2. Inference compatibility: The model preserves the standard static top-k routing rule at inference time with zero structural modification or latency overhead.

3. Local smoothing: The smoothing effect is local in router rank space, not a differentiable relaxation of top-K selection. If $\text{supp}(q) = [K_{\min}, K_{\max}]$, then ranks $j \leq K_{\min}$ are always exposed, ranks $K_{\min} < j \leq K_{\max}$ have fractional exposure, and ranks $j > K_{\max}$ are never exposed.

---

Empirical Validation / Results

Experimental Setup

• Supervised Fine-Tuning (SFT): OLMoE-1B-7B-0924 on Dolci-Instruct-SFT and Qwen3-30B-A3B-Base on Infinity-Instruct 7M Core
• Pretraining: 1.42B-A365M MoE models with 128 experts per layer, trained on FineWeb-Edu for 26.25B tokens with a 6e19 FLOPs budget
• Elastic schedules compared: 7–9, 6–10, and 5–11 (centered at top-8, $\sigma = 1.5$)
• Evaluation: All main evaluations use static top-8 inference

Supervised Fine-Tuning Results

Key findings:

• All three elastic schedules outperform the static top-8 baseline on both backbones
• The 6–10 schedule gives the best Avg(9) for both backbones
• An elastic schedule obtains the best score on all nine OLMoE tasks and on eight of nine Qwen3 tasks

From-Scratch Pretraining Results

Comparison with Other MoE Training Methods

Expert-Selection Diagnostics

The paper defines two disagreement metrics. For a given layer $\ell$ and token $t$, let $e_{\ell,t}^{\text{elastic}}$ and $e_{\ell,t}^{\text{static}}$ denote the top-1 expert chosen by each configuration, and $S_{\ell,t}^{\text{elastic}}$ and $S_{\ell,t}^{\text{static}}$ their active top-8 expert sets:

$$D_1 = \frac{1}{L \cdot T} \sum_{\ell=1}^{L} \sum_{t=1}^{T} \mathbf{1}\left[e_{\ell,t}^{\text{elastic}} \neq e_{\ell,t}^{\text{static}}\right]$$ $$D_8 = \frac{1}{L \cdot T} \sum_{\ell=1}^{L} \sum_{t=1}^{T} \left(1 - \frac{\left|S_{\ell,t}^{\text{elastic}} \cap S_{\ell,t}^{\text{static}}\right|}{8}\right)$$

Results show that 6–10 changes the top-1 routed expert for 13.62% of tokens on average, while only 7.08% of the selected top-8 set differs (about 0.57 experts out of 8), indicating that elastic routing primarily refines expert rankings rather than altering the global composition of the active expert pool.

Router-Swap Intervention

Optimal performance gains require mutual alignment through joint router–expert co-adaptation.

Ablations and Budget Controls

Distribution shape ablation (fixed 6–10 support):

Budget controls (increasing compute without smoothing):

Static top-9, static top-10, and top-p all remain below 6–10 despite using more expert computation, confirming that gains come from smoothing the boundary rather than simply activating more experts.

---

Theoretical and Practical Implications

Theoretical Implications

1. Reformulation of static top-k training: The paper formalizes static top-k MoE training as a hard rank-exposure boundary problem, providing a principled framework (rank-exposure accounting) for understanding and addressing boundary rigidity.

2. Expected exposure conservation: The mathematical framework shows that elastic routing preserves expected active-expert count ($\mathbb{E}[K] = k_0$) while changing which boundary ranks enter the task-loss path. The exposure above the cutoff is balanced by exposure from ranks at or below it:

$$\sum_{j > k_0} a_q(j) = \sum_{j \leq k_0} \left(1 - a_q(j)\right) = \frac{1}{2}\mathbb{E}|K - k_0|$$

3. Differentiation from existing methods: Unlike dynamic routing approaches that alter deployment budgets or model structures, elastic routing preserves the standard fixed-budget inference constraints while softening the sharp threshold during training only.

Practical Implications

1. Zero-cost deployment: The method requires no structural modification or latency overhead at inference time, making it directly applicable to existing serving infrastructures.

2. Consistent gains across settings: The improvement holds across supervised fine-tuning (both small and large backbones) and from-scratch pretraining, suggesting broad applicability.

3. Middle-layer routing refinement: The depth-localized diagnostic (L3–L7 achieving Avg(9) 31.50 vs. 31.41 for all-layer elastic) suggests that practitioners could potentially target elasticity to specific layers for computational efficiency.

4. Robustness across inference budgets: Elastic routing improves performance not just at k=8 but throughout a range of inference top-k values, indicating more robust learned expert cooperation patterns.

---

Conclusion

Summary of Contributions

1. Problem formulation: The paper formulates static top-k MoE training as a hard rank-exposure boundary and introduces Elastic Expert Routing to soften it into a gradual selection probability distribution without modifying inference mechanics.

2. Consistent performance gains: Elastic routing demonstrates consistent improvements during supervised fine-tuning on OLMoE-1B-7B and Qwen3-30B-A3B backbones (+0.84 and +2.02 points respectively), as well as during from-scratch pretraining under fair settings (up to 0.04 reduction in validation LM loss, +1.6 points average downstream).

3. Mechanistic understanding: Routing diagnostics and router-swap interventions show that gains are associated with middle-layer routing changes and improved expert prioritization, not simply increased compute or altered expert pool composition.

Key Takeaways

• The moderate 6–10 interval gives the strongest aggregate result across both SFT and pretraining settings, suggesting that useful smoothing should expose experts near the cutoff threshold without allocating excessive routing probability to low-ranked candidates.
• The default $\sigma = 1.5$ outperforms both more concentrated ($\sigma = 1.0$) and flatter ($\sigma = 3.0$) distributions, indicating the smoothing distribution should be neither too concentrated nor too flat.
• Elastic routing primarily refines expert rankings (13.62% top-1 disagreement) rather than altering the global composition of the active expert pool (7.08% set disagreement).

Future Directions

The main limitation is the scale of from-scratch pretraining. While the pretraining results (1.42B parameters) provide a controlled validation of the proposed training objective at moderate scale, validating the approach under large-scale from-scratch pretraining remains an important direction for future work.