△ Mechanism confirmed, baseline not beaten 2026

Shell-Wise Balanced MoE Routing

Implementation & benchmark of arXiv:2607.23641 — Optimal Dynamic Cooling of Multiple Qubits

Usefulness7/10

Difficulty5/10

Novelty5/10

Architecture Scalability Stability Accuracy Moe routing Transformer Training Regularization Combinatorics Probability Linear algebra

Source paper: Optimal Dynamic Cooling of Multiple Qubits arXiv:2607.23641 ⓘ · analyzed Aug 30, 2026

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

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.

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}$$

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.

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

Mechanism confirmed, baseline not beaten

Stage 1 · Toy mechanism gate: Passed ✓

Stage 2 · Mechanism transferred to benchmark: Confirmed ✓

Stage 2 · Practical benchmark result: Inconclusive

Methodology: 8 paired seeds, permutation test (p<0.05) against a learning-rate-tuned baseline. How verification works

Stage 1 — Mechanism check agent confidence 9/10

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.

Agent confidence: 9/10

Baseline: Aggregate expert loads [1627,1320,182,211,199,186,199,172], load CV=1.0948; Eq.66-grid maximum marginal spread=17.53-20.35 in divisible-shell sweeps.
Idea: Aggregate loads [512,512,512,512,512,512,512,512], load CV=0.0000; per-shell count spread <=1 in all nondivisible sweeps; exact Eq.66-grid spread=0.0 for divisible shells; reassigned fraction=0.4712; router-score loss/token=0.5094.

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.

How to run: python3 shell_moe_experiment.py

Stage 2 — Benchmark vs. tuned baseline (arXiv:2607.23641) Sep 2, 2026

Ahead of the baseline (−4.2%) but not statistically significant (p=0.3804, wins 5 / 8 paired seeds).

Benchmark: Correlated_token_moe_regression
Model: custom_moe
Paired seeds: 8
Baseline mean: 0.0566
Idea mean: 0.0542
Effect (Δ): -0.0024 (−4.2%; negative = idea better)
Wins: 5 / 8 paired seeds
p-value: 0.3804 (permutation test, 20 000 shuffles)
Smallest detectable effect: ±10.1%
Mechanism: Confirmed ✓
Practical effect: Inconclusive
Baseline tuning: swept over 9 configs

Limitations:

The accepted custom track is small and synthetic; no built-in language-model track exists for MoE routing. Top-2 routing, capacity overflow, wall-clock dispatch overhead, and larger-scale perplexity were not tested. The benchmark uses a greedy quota assignment and 12 training epochs.

How to run: python3 stage2_shell_moe.py

Verdict computed by deterministic test code from paired-seed statistics — not by the language model.

Stage 2 — Benchmark Aug 31, 2026 latest Failed ✗

Agent confidence: 10/10

The registered custom track correlated_token_moe_regression was benchmarked end-to-end with eight paired seeds. Shell-wise balancing exactly achieved zero measured load CV and zero per-shell spread, but the best test MSE improvement from 0.0566415 to 0.0542420 was not significant (paired permutation p=0.38035), so the idea did not produce the required significant win.

Baseline: Registered-track baseline, best lr=0.006 and auxiliary loss=0.01; full 8-seed mean test MSE 0.056641495786607265; mean load CV 0.11686608614250746; mean shell spread 258.375.
Idea: Registered-track shell-balanced MoE, best lr=0.006 and K=8; full 8-seed mean test MSE 0.05424201488494873; mean load CV 0.0; mean shell spread 0.0; mean reassigned fraction 0.503828125.
Paired delta: -0.002399480901658535, 5/8 idea wins, permutation p=0.38035

Benchmark result