Three Phases of Expert Routing: How Load Balance Evolves During Mixture-of-Experts Training

1. Problem Statement

Mixture-of-Experts (MoE) architectures scale model capacity by routing tokens to a subset of specialized expert networks. A critical engineering challenge is load balancing: without intervention, tokens tend to concentrate on a few high-quality experts, leaving others idle (expert collapse).

• Current State: Standard remedies use an auxiliary balance loss (e.g., Switch Transformers) with a tunable coefficient $\alpha$ to penalize load concentration.
• The Gap: While these methods work in practice, there is no theoretical understanding of how the tradeoff between load balance and expert quality evolves during the training process. Existing analyses are typically static (post-hoc) and fail to capture the dynamic trajectory of the routing mechanism.

2. Methodology

The authors model MoE token routing as a congestion game within the framework of Mean-Field Games (MFG).

• Game Formulation:
  • Players: Tokens.
  • Resources: Expert networks.
  • Cost Function: A token incurs a cost based on expert quality ($q_i$) and congestion ($\gamma \mu_i$), where $\mu_i$ is the load on expert $i$ and $\gamma$ is the congestion coefficient.
  • Equilibrium: The routing distribution is modeled as the unique equilibrium of a potential game with entropy regularization (softmax).
• Key Parameter: The Effective Congestion Coefficient ($\gamma_{eff}$). This single parameter quantifies the strength of the balance-quality tradeoff at any given training checkpoint. It captures both the explicit auxiliary loss and the implicit balance learned through gradient descent.
• Identification: $\gamma_{eff}$ is estimated by fitting the theoretical MFG equilibrium to observed load distributions from training checkpoints. The authors minimize the $L_1$ distance between the observed load and the theoretical best-response map.
• Models Analyzed:
  • OLMoE-1B-7B: 20 checkpoints (dense sampling in the early "surge" region).
  • OpenMoE-8B: 6 checkpoints (different architecture: 32 experts, 4 MoE layers).
• Extensions:
  • Multi-type MFG: Clusters tokens into $K$ types to model heterogeneity, improving load prediction.
  • Decomposition: Separates $\gamma_{eff}$ into explicit ($\gamma_{explicit} = \alpha M$) and implicit ($\gamma_{implicit}$) components.

3. Key Contributions

1. The Three-Phase Trajectory: The discovery of a non-monotone, inverted-U trajectory for $\gamma_{eff}$ during pretraining, invisible to static analysis of converged models.
2. Theoretical Honesty: Proving that for well-balanced models, the single-type MFG equilibrium is mathematically equivalent to temperature-scaled softmax. The game theory lens does not improve static prediction accuracy but provides a structural explanation for why softmax works and what the "temperature" represents (the congestion coefficient).
3. Effective Congestion Decomposition: Demonstrating that the optimizer internalizes balance far beyond the explicit loss signal. At convergence, $\gamma_{eff}$ is $\approx 13\times$ larger than $\gamma_{explicit}$.
4. Scope Diagnostics: Developing metrics (e.g., $K/M$ ratio, continuation spread $\epsilon_l$) to determine where the MFG model applies and where it breaks down (e.g., late layers with high token differentiation).

4. Key Results

A. The Three-Phase Training Dynamics

Tracking $\gamma_{eff}$ across 20 checkpoints of OLMoE-1B-7B reveals three distinct phases:

1. Phase 1: Surge (Steps 5K–50K):
   • $\gamma_{eff}$ rises sharply from $\approx 14$ to a peak of 36–39 (at steps 30K–40K).
   • Interpretation: The router aggressively learns to balance the load. Routing entropy increases, and the system prioritizes utilization over raw expert quality.
2. Phase 2: Stabilization (Steps 100K–400K):
   • $\gamma_{eff}$ plateaus at $\approx 24–28$.
   • Interpretation: The router has found its operating point. Experts begin to specialize (quality spread $B_0$ drops), but the balance constraint remains steady.
3. Phase 3: Relaxation (Steps 400K–1.2M):
   • $\gamma_{eff}$ declines significantly from $\approx 27$ to 8.5.
   • Interpretation: Once experts have differentiated sufficiently, the router "relaxes" the balance constraint to prioritize quality. The model trades uniformity for selectivity.

• Replication: This inverted-U pattern was replicated on OpenMoE-8B, confirming it is not an artifact of a specific architecture.
• Fine-Tuning: The three-phase pattern is specific to pretraining. During fine-tuning (annealing), $\gamma_{eff}$ remains stable, indicating the routing equilibrium has already settled.

B. Static vs. Dynamic Performance

• Static Equivalence: On the converged model, the single-type MFG and temperature-scaled softmax are statistically indistinguishable (Held-out $L_1$: 0.199 vs. 0.200). The game theory model adds no predictive power for a single snapshot.
• Dynamic Value: The value lies entirely in the dynamics. The non-monotone trajectory reveals that early training prioritizes balance, while late training prioritizes quality.

C. Implicit Balance Dominance

• At convergence, the explicit auxiliary loss contributes $\gamma_{explicit} = 0.64$.
• The fitted effective congestion is $\gamma_{eff} = 8.5$.
• Result: The optimizer internalizes 13$\times$ more effective congestion than the explicit loss provides. The auxiliary loss acts as a "seed," but gradient dynamics grow the balance internally.

D. Multi-Type Extension

• By clustering tokens into types (e.g., $K=4$), the Multi-type MFG improves load prediction accuracy by 30% on average across all 16 layers compared to the single-type model.
• This confirms that token heterogeneity is significant, though the coupling effect is less pronounced in well-balanced models.

5. Significance and Implications

• New Diagnostic Tool: $\gamma_{eff}$ serves as a real-time monitor for MoE training. A premature drop in $\gamma_{eff}$ or a failure to reach the "surge" peak could signal expert collapse or training instability.
• Reframing the Auxiliary Loss: The paper challenges the view that the auxiliary loss is the primary driver of balance. Instead, it suggests the loss is a catalyst for a much larger internal balancing mechanism driven by the optimizer.
• Optimization Strategy: The findings suggest a potential for adaptive balance schedules. Practitioners could theoretically reduce the explicit balance coefficient ($\alpha$) during the "Relaxation" phase (Phase 3) to allow the router to focus more on quality without destabilizing the system.
• Theoretical Insight: It provides a rigorous game-theoretic justification for the ubiquity of softmax in MoE routers, linking the "temperature" parameter directly to the congestion coefficient of a potential game.

Conclusion

The paper establishes that MoE training is not a static optimization of a fixed objective but a dynamic evolution through three distinct regimes: Surge, Stabilization, and Relaxation. By modeling routing as a congestion game, the authors reveal that the "temperature" of the router is a dynamic parameter reflecting the tension between load balancing and expert specialization, a phenomenon invisible to standard post-training analysis.