Expansion-balanced MoE routing
Implementation & benchmark of arXiv:2607.26049 — Recent progress in graph theory using expansion
Source paper: Recent progress in graph theory using expansion arXiv:2607.26049 ⓘ · analyzed Aug 31, 2026
AI-generated research hypothesis, automatically tested. Not peer-reviewed.
Idea description
Treat token-to-expert assignments as a bipartite routing graph and penalize patterns that leave token groups sharing too few experts. Enforcing sublinear expansion makes routing diversity explicit: even a moderately sized group of tokens must reach many distinct experts, reducing expert collapse and improving load balance without requiring every token to use every expert.
Formulas
Mathematical statement
Apply Definition 3.1 to a bipartite routing graph with token vertices on the left and expert vertices on the right. For a token subset U with k <= |U| <= n/2, define N_G(U) as the experts selected by at least one token in U; require |N_G(U)| >= [epsilon/log^2(3|U|/k)]|U|. Here n is the number of tokens in the routing batch, epsilon is the desired expansion coefficient, k is the minimum group size, and each token has at most r expert assignments. The condition prevents a group of many tokens from collapsing onto one or two experts. A differentiable proxy uses sampled subsets U and a soft count of distinct experts, c_j(U)=1-prod_{i in U}(1-p_{ij}), where p_{ij} is the router probability for token i and expert j. The resulting soft neighborhood size is C(U)=sum_j c_j(U), which approximates the hard cardinality |N_G(U)| and permits a hinge penalty when it falls below the paper's lower bound.
Implementation notes
Integrate the method into the router logits before top-r expert selection. Let p_ij=softmax_j(z_ij/t) for token i and expert j, with E experts and routing batch size n. At each training step, sample subsets S containing random groups of sizes k, 2k, 4k, ..., min(n/2, 32k), together with groups formed by tokens currently assigned to the same expert. For each U, compute c[j]=1-prod(i in U)(1-p[i,j]), set C=sum(c), and add lambda_expand * relu(eps*len(U)/log(3*len(U)/k)^2 - C)^2 to the router loss. Then perform standard top-r routing, capacity enforcement, and expert computation. The exact paper-derived quantity is the target lower bound; the product-based distinct count is a differentiable surrogate because hard neighborhoods are nondifferentiable. Log the hard diagnostic |N_G(U)| by counting experts selected by at least one token in U, and report the fraction of sampled subsets violating the bound. Warm epsilon from zero during the first 10% of training and combine this term with the usual per-expert load-balancing loss. Test an 8-expert Switch-Transformer-style model on C4 or a synthetic clustered-token task against standard top-1 and top-2 routing at identical FLOPs and capacity factors. Success means lower maximum expert load, fewer dropped tokens, reduced router collapse, and equal or better validation loss at the same active-expert count.
Verification
Stage 1 · Toy mechanism gate: Passed ✓
Stage 2 · Mechanism transferred to benchmark: Confirmed ✓
Stage 2 · Practical benchmark result: No measurable effect
Methodology: 8 paired seeds, permutation test (p<0.05) against a learning-rate-tuned baseline. How verification works
Stage 1 — Mechanism check agent confidence 8/10
Built a differentiable expansion penalty, hard neighborhood diagnostics, and a matched 8-expert toy router benchmark in experiment.py. Uniform soft-count and finite-group predictions matched closed forms to below 5e-7 error, lambda scaling was exact, and the hinge activated at the predicted epsilon threshold. The regularizer reduced load standard deviation from 16.0 to 13.08, dropped-token fraction from 0.50 to 0.391, and subset violations from 1.0 to 0.292, while cross-entropy increased from 0.0013 to 0.182; this is a real load-balance signal but not an accuracy win.
- Agent confidence
- 8/10
- Baseline
- CE 0.00128; max hard load 32; min hard load 0; load std 16.0; dropped fraction 0.500; violation fraction 1.000; mean hard neighborhood 3.896; soft expansion penalty 4.335
- Idea
- CE 0.182; max hard load 32; min hard load 0; load std 13.077; dropped fraction 0.391; violation fraction 0.292; mean hard neighborhood 6.000; soft expansion penalty 0.714; formula errors <=4.7e-7; lambda scaling exact; predicted epsilon threshold 2.107 and observed hinge transition there
Limitations: Only a synthetic 128-token, 8-expert top-1 router was tested with one random seed and a fixed subset schedule. No top-2 routing, capacity-factor sweep, FLOP-matched language-model training, C4 validation, multiple-seed statistics, or comparison against stronger modern load-balancing objectives was performed.
How to run: python3 experiment.py
Stage 2 — Benchmark vs. tuned baseline (arXiv:2607.26049)
Ahead of the baseline (−0.0%) but not statistically significant (p=0.0863, wins 6 / 8 paired seeds).
- Benchmark
- Sequence modeling
- Model
- transformer_tiny
- Paired seeds
- 8
- Baseline mean
- 0.9597
- Idea mean
- 0.9594
- Effect (Δ)
- -0.0002 (−0.0%; negative = idea better)
- Wins
- 6 / 8 paired seeds
- p-value
- 0.0863 (permutation test, 20 000 shuffles)
- Smallest detectable effect
- ±0.0%
- Mechanism
- Confirmed ✓
- Practical effect
- No effect
- Baseline tuning
- swept over 4 configs
The initial larger configuration exceeded the 40-minute budget, so the final full protocol used a compact one-layer, 16-wide transformer, 200 training samples, 100 test samples, and 3 epochs. Dense soft MoE routing was used rather than hard top-1/top-2 sparse dispatch with capacity enforcement; no C4-scale language-model experiment was performed.
How to run: python3 bench_expansion.py
Verdict computed by deterministic test code from paired-seed statistics — not by the language model.
Stage 2 — Benchmark latest Failed ✗
Agent confidence: 9/10
Expansion-balanced routing produced a small directional improvement but not a statistically significant task-metric win. The benchmark mechanism signature was confirmed directionally, with slightly larger token-group expert neighborhoods and lower routing load dispersion, but bench.make_report() returned no measurable effect because the paired permutation p-value was 0.0863.
- Baseline
- Full 8-seed sequence MSE mean 0.9596702233, std 0.0851384368; best baseline config lr=0.003, temperature=0.7, epochs=3.
- Idea
- Full 8-seed sequence MSE mean 0.9594315290, std 0.0850028169; paired delta -0.0002386943, 6/8 wins, permutation p=0.0863.
Benchmark result
Ahead of the baseline (−0.0%) but not statistically significant (p=0.0863, wins 6 / 8 paired seeds).
- Benchmark
- Sequence modeling
- Model
- transformer_tiny
- Paired seeds
- 8
- Baseline mean
- 0.9597
- Idea mean
- 0.9594
- Effect (Δ)
- -0.0002 (−0.0%; negative = idea better)
- Wins
- 6 / 8 paired seeds
- p-value
- 0.0863 (permutation test, 20 000 shuffles)
- Smallest detectable effect
- ±0.0%
- Mechanism
- Confirmed ✓
- Practical effect
- No effect
- Baseline tuning
- swept over 4 configs
The initial larger configuration exceeded the 40-minute budget, so the final full protocol used a compact one-layer, 16-wide transformer, 200 training samples, 100 test samples, and 3 epochs. Dense soft MoE routing was used rather than hard top-1/top-2 sparse dispatch with capacity enforcement; no C4-scale language-model experiment was performed.
How to run: python3 bench_expansion.py
Artifacts
- bench_expansion.py 4.8 KB View Raw
- bench_report.json 3.7 KB View Raw
- experiment.py 5.3 KB View Raw
- report.md 1.6 KB View
- report_bench_2026-08-31T180832.md 5.5 KB View
- results.json 3.7 KB View Raw