# Expansion-balanced MoE routing

- ID: 1430
- Canonical URL: https://synthcore.org/idea/1430/expansion-balanced-moe-routing
- API JSON: https://synthcore.org/api/idea/1430.json
- API Markdown: https://synthcore.org/api/idea/1430.md
- Verification status: mech_ok_no_baseline
- Source: [arXiv:2607.26049](https://arxiv.org/abs/2607.26049)
- Category: regularization
- Solves: scalability, stability, accuracy
- ML areas: moe-routing, transformer, regularization, training
- Math tags: graph-theory, combinatorics, optimization
- Ratings: usefulness 7/10; difficulty 4/10; novelty 6/10

## 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.

## 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.

## Key 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.$$

## 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

- Status: mech_ok_no_baseline
- Mechanism evidence: yes
- Mechanism confirmed: yes
- Practical verdict: no_effect
- Verdict: 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.

### Mechanism check

- Verdict: 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.
- Confidence: 8/10
- 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.

### Practical benchmark

- Paired seeds: 8
- Baseline mean: 0
- Idea mean: 0
- p-value: 0.0863
- Paired wins: 6/8
- Benchmark verdict: no measurable effect

## Artifacts

- [bench_expansion.py](https://synthcore.org/code/559/bench_expansion.py)
- [bench_report.json](https://synthcore.org/code/559/bench_report.json)
- [experiment.py](https://synthcore.org/code/559/experiment.py)
- [report.md](https://synthcore.org/code/559/report.md)
- [report_bench_2026-08-31T180832.md](https://synthcore.org/code/559/report_bench_2026-08-31T180832.md)
- [results.json](https://synthcore.org/code/559/results.json)
- [Download all files as ZIP](https://synthcore.org/download/559)

## Disclaimer

AI-generated research hypothesis, automatically tested. Not peer-reviewed.