Three Phases of Expert Routing

Code and data for the paper:

Three Phases of Expert Routing: How Load Balance Evolves During Mixture-of-Experts Training Charafeddine Mouzouni. April 2026.

Summary

We model MoE token routing as a congestion game and track the effective congestion parameter across training. The trajectory reveals three phases: surge (the router learns to balance), stabilization (experts specialize under steady balance), and relaxation (the router trades balance for quality). This non-monotone trajectory is invisible to analysis of converged models.

Repository structure

moe/                  Core MFG solver library
  params.py           MoE-MFG parameter definitions
  solver.py           Policy iteration solver
  theory.py           Analytical formulas (anti-concentration, gamma_c)
  costs.py            Cost structure definitions

experiments/          Experiment scripts (each is self-contained)
  exp8_*              Effective congestion decomposition
  exp9_*              Per-token analysis, continuation spread, multi-type MFG
  exp10_*             Static validation (119 texts, three-way split)
  exp10b_*            Clustering vs game structure ablation
  exp11b_*            Training dynamics with bootstrap CIs (14 checkpoints)
  exp11c_*            Bootstrap CIs at step 35K (surge peak)
  exp12_*             Quality estimator robustness
  aws_*               Scripts run on GPU instances (OLMoE dense, OpenMoE, robustness)

results/              Cached experiment outputs (JSON)
  exp11b_dynamics_scaled.json    OLMoE 14-checkpoint trajectory with bootstrap CIs
  exp11c_bootstrap_35k.json      Step 35K bootstrap CIs
  olmoe_dense_results.json       6 dense surge checkpoints + 7 annealing
  openmoe_dynamics_v2.json       OpenMoE-8B 6-checkpoint dynamics
  exp10_definitive.json          Static analysis (Tables 2-4 in paper)
  exp10b_cluster_baselines.json  Clustering ablation
  exp12_quality_robustness.json  Quality estimator robustness (Figure 2)
  exp8_effective_congestion.json Decomposition (Table 5)
  exp9_per_token_analysis.json   Continuation spread + multi-type
  exp4_scaled_validation.json    Intermediate data for exp8/exp9
  exp3b_multi_model.json         Cross-architecture scope (Table 6)
  synthetic_recovery.json        Identification validation
  k_type_ablation.json           K=2,4,8 cluster count robustness

figures/              Figure generation scripts
paper/                LaTeX source and compiled figures

Models used

Model	HuggingFace	M	K	Checkpoints
OLMoE-1B-7B	allenai/OLMoE-1B-7B-0924	64	8	20 (14 coarse + 6 dense)
OpenMoE-8B	OrionZheng/openmoe-8b	32	2	6

Reproducing results

Requirements

pip install torch transformers scipy numpy scikit-learn sentencepiece protobuf matplotlib

Running experiments

Each experiment script is self-contained. From the repository root:

# Static analysis (Table 2, 3, 4 in paper) — ~30 min on CPU
python -m experiments.exp10_definitive

# Training dynamics with bootstrap CIs (Table 1) — ~3 hours on CPU
python -m experiments.exp11b_dynamics_scaled

# Quality estimator robustness (Figure 2) — ~25 min on GPU
python experiments/aws_quality_robustness.py

# Clustering ablation — ~45 min on CPU
python -m experiments.exp10b_cluster_baselines

GPU scripts (aws_*.py) are standalone and do not import from moe/. They can be run directly on any machine with PyTorch and Transformers installed.

Generating figures

python figures/fig_dynamics.py    # Figure 1: three-phase trajectory
python figures/fig_robustness.py  # Figure 2: quality estimator robustness

Compiling the paper

cd paper && pdflatex main.tex && pdflatex main.tex

Citation

@article{mouzouni2026threephases,
  title={Three Phases of Expert Routing: How Load Balance Evolves During Mixture-of-Experts Training},
  author={Mouzouni, Charafeddine},
  journal={arXiv preprint},
  year={2026}
}

License

MIT License. See LICENSE.