Superlinear Reinforced MoE Routing

Integrate this into the token-to-expert assignment operation of a top-1 or top-2 MoE layer, operating independently for each microbatch \(b\). First compute the usual router affinity logits \(s_{b,u,i}\) for token \(u\) and expert \(i\), and perform a cheap initial assignment using softmax or uniform expert probabilities. For rounds \(r=1,\ldots,R\), maintain expert loads \(L_{b,i}^{(r-1)}\), form the probability in the fourth formula with a small \(\varepsilon=10^{-3}\), sample or take the argmax assignment for every token, and recompute the loads. Stop early if \(\max_i L_{b,i}^{(r)}\ge \rho n_b\), where \(\rho\) can be 0.9 during routing or 1.0 for a hard consensus test. Use \(\alpha\in[1.2,2.0]\); the paper's two-leader phase rule gives an initial choice of \(\alpha\) for a desired number of rounds, for example \(\alpha\in(2^{1/3},2^{1/2})\) for a three-round transition. During training, use a straight-through estimator for sampled assignments, or replace indicators by router probabilities when computing gradients. The mathematical component is the normalized power reinforcement and its phase-regime prediction; estimate the actual concentration rate, overflow probability, and number of active experts empirically because learned logits violate the paper's exchangeability assumptions. The first experiment should use a 4-expert Switch-Transformer-style layer on CIFAR-100 or a small language-model corpus, comparing static top-1 routing, expert-choice routing, and this method at equal FLOPs. Measure loss versus optimizer steps, tokens dropped by capacity limits, active experts per batch, all-to-all bytes, and load variance. Success means fewer active experts or communication groups at matched validation loss, with no increase in dropped-token rate and faster concentration across routing rounds.