△ Mechanism confirmed, baseline not beaten 2026

Expansion-balanced MoE routing

Implementation & benchmark of arXiv:2607.26049 — Recent progress in graph theory using expansion

Usefulness7/10
Difficulty4/10
Novelty6/10

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

$$|N_G(U)|\geq\frac{\varepsilon}{\log^{2}(3|U|/k)}\,|U|,\qquad k\leq|U|\leq n/2.$$
$$c_j(U)=1-\prod_{i\in U}(1-p_{ij}),\qquad C(U)=\sum_{j=1}^{E}c_j(U).$$
$$\mathcal{L}_{\mathrm{expand}}=\frac{1}{|\mathcal{S}|}\sum_{U\in\mathcal{S}}\left[\max\left(0,\frac{\varepsilon|U|}{\log^{2}(3|U|/k)}-C(U)\right)\right]^2.$$

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

Mechanism confirmed, baseline not beaten

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
Limitations:

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
Limitations:

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

Implementation overview ⬇ Download all as ZIP 6 files · code, reports and structured results