# Shell-Wise Balanced MoE Routing

- ID: 1327
- Canonical URL: https://synthcore.org/idea/1327/shell-wise-balanced-moe-routing
- API JSON: https://synthcore.org/api/idea/1327.json
- API Markdown: https://synthcore.org/api/idea/1327.md
- Verification status: mech_ok_no_baseline
- Source: [arXiv:2607.23641](https://arxiv.org/abs/2607.23641)
- Category: architecture
- Solves: scalability, stability, accuracy
- ML areas: moe-routing, transformer, training, regularization
- Math tags: combinatorics, probability, linear-algebra
- Ratings: usefulness 7/10; difficulty 5/10; novelty 5/10

## Idea description

Replace a soft auxiliary load-balancing loss with a deterministic router correction that balances expert assignments separately within confidence or score shells. The correction preserves the number of tokens assigned to each shell while making every expert receive equal or nearly equal mass inside each shell, preventing high-confidence tokens from monopolizing a subset of experts.

## Mathematical statement

The paper defines shell weights q_k(P_1)=P_1^k(1-P_1)^{N-k} and writes the marginal of target position alpha as p'_{1,alpha}(P_1)=sum_{k=0}^N (sum_x x_alpha n_{x,k})q_k(P_1), where x is a binary configuration, n_{x,k} counts configurations of type x in shell k, and P_1 is an excitation probability. Because the functions q_k are linearly independent, equality of all marginals for every P_1 is equivalent to the finite constraints sum_x x_alpha n_{x,k}=sum_x x_gamma n_{x,k} for every pair alpha,gamma and every shell k (Eq. 66). In the ML adaptation, a shell k is a bin of router confidence or token difficulty, x_e indicates assignment to expert e, and n_{x,k} counts assignment masks in that shell. The transferable property is exact per-shell balance, stronger than aggregate batch balancing and independent of a tuned loss coefficient.

## Key formulas

- $$q_k(P_1)=P_1^k(1-P_1)^{N-k}$$
- $$p^{\prime}_{1,\alpha}(P_1)=\sum_x x_\alpha\operatorname{Tr}A_x=\sum_{k=0}^{N}\left(\sum_x x_\alpha n_{x,k}\right)q_k(P_1)$$
- $$\sum_x x_\alpha n_{x,k}=\sum_x x_\gamma n_{x,k}\qquad\text{for all }\alpha,\gamma,k$$
- $$\sum_{e=1}^{E}x_e=r,\qquad \sum_{x:\,x_e=1}n_{x,k}\approx \frac{r}{E}\sum_x n_{x,k}$$

## Implementation notes

Integrate this after the MoE router produces logits and before token dispatch. Use top-r routing with E experts. Compute a scalar shell statistic for each token, such as top-1 confidence, selected-probability sum, or router entropy, and quantize it into K fixed bins; bin k is the mathematical shell. Within each bin, retain the preferred expert mask whenever possible, but repair assignments so every expert receives either floor(rB_k/E) or ceil(rB_k/E) selections, where B_k is the number of tokens in shell k. A practical procedure is: compute logits z_i, top-r masks x_i, and shell indices k_i; for each shell initialize expert deficits d_e=rB_k/E; sort tokens by the score margin between their preferred and alternative experts; process low-margin tokens first; choose the feasible mask minimizing -sum_e x_{i,e}z_{i,e}+lambda sum_e x_{i,e}d_e; update d_e. The paper's Eq. (66) is enforced approximately or exactly inside each shell. Shell boundaries can be batch quantiles, and reassignment cost is measured empirically. Test an 8-expert top-1 and top-2 language model on a small corpus against standard auxiliary load balancing at equal FLOPs. Measure load coefficient of variation, capacity overflows, router-loss sensitivity, training throughput, and perplexity. Success is fewer overflows and lower load variance without perplexity degradation or significant routing overhead.

## Verification

- Status: mech_ok_no_baseline
- Mechanism evidence: yes
- Mechanism confirmed: yes
- Practical verdict: inconclusive
- Verdict: Built a deterministic top-1 shell-wise balanced router with quantile confidence shells and exact per-shell expert quotas. Across E={3,5,8} and K={2,4,7}, the observed maximum per-shell count spread was 1, matching the floor/ceil prediction; for divisible shell sizes, the Eq. 66 polynomial expert-marginal spread was exactly 0 over 99 excitation probabilities, while baseline spread was 17.5–20.3. In the synthetic imbalanced-routing comparison, CV fell from 1.095 to 0.000 and every shell reached zero CV, but 47.1% of tokens were reassigned with average router-score loss 0.509 per token; this demonstrates the mechanism, not an end-to-end perplexity win.

### Mechanism check

- Verdict: Built a deterministic top-1 shell-wise balanced router with quantile confidence shells and exact per-shell expert quotas. Across E={3,5,8} and K={2,4,7}, the observed maximum per-shell count spread was 1, matching the floor/ceil prediction; for divisible shell sizes, the Eq. 66 polynomial expert-marginal spread was exactly 0 over 99 excitation probabilities, while baseline spread was 17.5–20.3. In the synthetic imbalanced-routing comparison, CV fell from 1.095 to 0.000 and every shell reached zero CV, but 47.1% of tokens were reassigned with average router-score loss 0.509 per token; this demonstrates the mechanism, not an end-to-end perplexity win.
- Confidence: 9/10
- Limitations: Only a synthetic top-1 router was tested; top-2 routing, capacity overflow simulation, throughput/dispatch overhead, auxiliary-loss sensitivity, language-model training, perplexity, and GPU execution were not tested. The greedy assignment implementation does not yet implement the proposed top-r deficit-aware optimizer.

### Practical benchmark

- Paired seeds: 8
- Baseline mean: 0
- Idea mean: 0
- p-value: 0.38035
- Paired wins: 5/8
- Benchmark verdict: no significant win

## Artifacts

- [bench_report.json](https://synthcore.org/code/522/bench_report.json)
- [report.md](https://synthcore.org/code/522/report.md)
- [report_bench_2026-08-31T165337.md](https://synthcore.org/code/522/report_bench_2026-08-31T165337.md)
- [results.json](https://synthcore.org/code/522/results.json)
- [shell_moe_bench_track.py](https://synthcore.org/code/522/shell_moe_bench_track.py)
- [shell_moe_experiment.py](https://synthcore.org/code/522/shell_moe_experiment.py)
- [stage2_shell_moe.py](https://synthcore.org/code/522/stage2_shell_moe.py)
- [Download all files as ZIP](https://synthcore.org/download/522)

## Disclaimer

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