Summary (Overview)

• This paper provides a rigorous mathematical analysis of softmax routing dynamics in mixture-of-experts (MoE) models, focusing on three key questions: which simplex face is selected during routing concentration, why concentration persists with trainable experts, and what determines recovery time for suppressed experts.
• The authors establish a stationary-loss barrier theorem showing that fixed-expert routing converges to the optimal projection of the target onto the expert convex hull, with the selected face determined by loss-level thresholds.
• For the common-target square-loss model with trainable experts, they prove that slow expert adaptation tracks fixed-expert dynamics on finite horizons, while every fixed positive adaptation rate yields an interior routing limit—reconciling transient concentration with long-time interior convergence.
• The main recovery result provides an exact normal-velocity decomposition separating complementary-expert covariance from relative gradient advantage, yielding recovery times of $\Theta(\log(1/\varepsilon))$ for $\alpha=0$ and $\Theta(\varepsilon^{-\alpha})$ for $\alpha>0$.
• Paired Transformer-MoE experiments confirm that removing a class-conditional router-gradient contrast delays confirmed recovery and reduces accumulated log-odds on fixed tokens under both standard and capacity-constrained dispatch.

Introduction and Theoretical Foundation

Sparse mixture-of-experts (MoE) models increase capacity by activating only a subset of expert subnetworks per input, enabling conditional computation at scale [12, 13, 23, 36]. During training, self-reinforcing routing preferences can concentrate assignments on a few experts—a form of routing collapse. The paper studies routing concentration as small pre-dispatch probability mass outside a proper coordinate face of the simplex $\Delta^{n-1}$.

The theoretical foundation rests on the geometry of softmax gradient flow. For a routing vector $p \in \Delta^{n-1}$, the softmax Jacobian is:

$$J(p) = \mathrm{diag}(p) - pp^\top$$

The Euclidean logit-gradient flow and its probability-space form are:

$$\dot{q} = -J(p)g, \quad \dot{p} = -J(p)^2 g$$

where $g$ is the probability gradient of the half-loss. At $p_j = 0$, the $j$-th row of $J(p)$ vanishes, creating boundary degeneracy—normal probability motion vanishes as $p_j \downarrow 0$ under bounded gradient signals. This degeneracy is a property of the vector field, while concentration describes the routing state.

The paper addresses three questions:

1. Which simplex face is selected as routing concentrates?
2. Why can concentration persist when experts are trainable?
3. What determines the time required for a suppressed expert to recover?

Methodology

Problem Setup

For $n$ experts with outputs $E_i \in \mathbb{R}^d$ stacked as rows of $E \in \mathbb{R}^{n \times d}$, and common target $y \in \mathbb{R}^d$, the routed output is $E^\top p$ with square loss:

$$\mathcal{L}(p, E) = \|E^\top p - y\|^2$$

The half-loss gradients are:

$$g(p, E, y) := \nabla_p \mathcal{H} = ER, \quad s(p, E, y) := \nabla_q \mathcal{H} = J(p)g, \quad \nabla_{E_i}\mathcal{H} = p_i R$$

where $R := E^\top p - y$. The coupled dynamics with expert-to-router learning-rate ratio $\eta \geq 0$ are:

$$\begin{cases} \dot{p} = -M(p)E(E^\top p - y), \\ \dot{E}_i = -\eta p_i (E^\top p - y), \quad i = 1, \ldots, n. \end{cases} \tag{2.2}$$

Key Theoretical Tools

1. Stationary-loss barrier: For each nonempty support set $S$, define the convex hull $\mathcal{C}_S$, its projection $z_S = \mathrm{Proj}_{\mathcal{C}_S} y$, and loss level $\ell_S = \|z_S - y\|^2$. The barrier separates optimal from suboptimal stationary supports.

2. Normal-velocity decomposition (Proposition 5.1): For suppressed probability $x = p_j$, with conditional probability vector $w_k = p_k/(1-x)$ on the complementary face:

$$\dot{x} = x(1-x)^2 \Gamma_j$$

where $\Gamma_j = \Lambda_j + x(1 + \|w\|_2^2)\Delta_j$, with $\Lambda_j$ the weighted covariance $\mathrm{Cov}_{k\sim w}(w_k, g_k)$ and $\Delta_j$ the relative gradient advantage.

3. Discrete log-odds identity (Proposition 6.2): For arbitrary logit increments $u^m = q^{m+1} - q^m$:

$$\mathrm{logit}(p_j^N) = \mathrm{logit}(p_j^0) + \sum_{m=0}^{N-1}(u_j^m - c_j^m)$$

where $c_j^m = \log\left(\sum_{k \neq j} w_k^m e^{u_k^m}\right)$.

Empirical Validation / Results

Face Selection (Theorem 3.1)

For fixed experts, if a trajectory enters the strict sublevel below the lowest suboptimal stationary loss $\ell_{\mathrm{comp}}$, then:

$$\mathrm{dist}(p(t), \mathcal{P}^*) \to 0, \quad E^\top p(t) \to z^*$$

where $z^* = \mathrm{Proj}_{\mathcal{C}(E)} y$ and $\mathcal{P}^*$ is the set of routing vectors representing $z^*$. The outside-face mass satisfies $D_{F^*}(p(t)) \to 0$.

Table 1: Fixed-expert face selection (measured at $T = 2000$)

Expert Adaptation (Theorems 4.1–4.2)

The finite-horizon slow-expert approximation gives:

$$\sup_{0 \leq t \leq T}\left(\|p_\eta(t) - p_0(t)\| + \|E_\eta(t) - E^0\|_F\right) \leq C_T \eta \tag{4.1}$$

For every fixed $\eta > 0$, the joint flow converges to an interior routing vector with explicit tails:

$$\|R(t)\|^2 \leq \|R(0)\|^2 e^{-2\eta t/n}, \quad \|q(t) - q_\infty\| \leq C_0 e^{-\eta t/n} \tag{4.2, 4.3}$$

with probability floor:

$$p_i(t) \geq \frac{p_i(0)}{p_i(0) + e^{2C_0}(1 - p_i(0))} > 0 \tag{4.4}$$

Recovery-Time Scaling (Theorem 5.2)

Under positive two-sided power bounds $c x(t)^\alpha \leq \Gamma_j(t) \leq c'x(t)^\alpha$ along the recovery path:

$$T_\theta = \Theta(\log(1/\varepsilon)) \text{ when } \alpha = 0, \quad T_\theta = \Theta(\varepsilon^{-\alpha}) \text{ when } \alpha > 0$$

For two experts with advantage bounded above and away from zero:

Transformer-MoE Interventions

Table 2: Effect of router-gradient projection on restricted mean time to confirmed recovery (through 1,024 steps)

Table 3: Log-odds response of initially subthreshold tokens at step 1,024

Theoretical and Practical Implications

Theoretical Contributions

1. Face selection mechanism: The stationary-loss barrier provides a complete characterization of which simplex face attracts routing dynamics, connecting convex projection theory with gradient-flow dissipation.

2. Two-timescale resolution: The paper resolves an apparent paradox—slow expert adaptation preserves finite-horizon concentration toward faces, while every fixed positive rate ensures interior long-time limits. This reconciles transient collapse with asymptotic interior behavior.

3. Recovery-time classification: The normal-velocity decomposition $\dot{x} = x(1-x)^2\Gamma_j$ provides a unified framework for understanding how boundary degeneracy ($x$ factor), optimization geometry ($(1-x)^2$ factor), and gradient structure ($\Gamma_j$) jointly determine recovery times.

4. Discrete-time exactness: The log-odds identity (6.2) holds for arbitrary optimizer increments, making it applicable to real neural network training with AdamW, clipping, and other discrete update rules.

Practical Implications

• Load-balancing losses encourage balanced aggregate expert use, but individual inputs can concentrate differently while maintaining aggregate balance—the paper's analysis applies at the individual-input level.
• Router-gradient interventions can systematically delay or accelerate recovery; removing class-conditional contrasts measurably slows recovery in real Transformer-MoE training.
• Recovery time scales inform training schedules: if a suppressed expert requires $\Theta(\varepsilon^{-\alpha})$ steps to recover, early intervention (before suppression becomes severe) is critical.

Conclusion

The paper provides a comprehensive mathematical framework for understanding routing collapse and recovery in MoE models. Key takeaways:

1. Fixed experts: A stationary-loss barrier identifies the selected simplex face, with geometric bounds controlling outside-face mass and entry times.

2. Trainable experts: Slow adaptation tracks fixed-expert dynamics on finite horizons, but every fixed positive rate yields interior convergence—concentration is transient for the common-target square-loss model.

3. Recovery dynamics: The normal-velocity decomposition $\dot{x} = x(1-x)^2\Gamma_j$ separates covariance from gradient advantage, with the order of vanishing of $\Gamma_j$ determining logarithmic vs. power-law recovery times.

4. Neural validation: Removing class-conditional router-gradient contrasts delays confirmed recovery and reduces accumulated log-odds on fixed tokens, confirming the theoretical predictions in real Transformer-MoE training.

Future directions include estimating recovery coefficients in naturally evolving systems, testing predictions on new trajectories, and exploring independent evaluation inputs with norm-matched gradient controls to further assess how neural responses transfer across inputs and depend on gradient direction.