Budgeted Worst-Case MoE Routing

Integrate this at the MoE router's dispatch stage, before top-1 or top-2 expert assignment. Partition tokens into demand groups \(k\) using a cheap categorical feature such as source batch, language, sequence-position bucket, or a learned token cluster. Estimate nominal group counts \(\bar d_k\) in each minibatch and set \(\delta_k=\rho\bar d_k\), with \(\rho\) initially 0.1--0.3. For each training step, construct one or several budgeted perturbations \(d_k=\bar d_k+\delta_kz_k\). A simple projected-ascent procedure is: initialize \(z=0\); repeatedly compute router logits and soft assignments \(x_{ke}=\operatorname{softmax}_e(g_{ke}/T)\), compute loads \(\ell_e=\sum_kd_kx_{ke}\), evaluate task loss plus \(\lambda\max_e\ell_e/C_e\), ascend in \(z\), clip every coordinate to \([-1,1]\), and project onto the \(\ell_1\) ball of radius \(\Gamma\). Use 1--3 ascent steps so the overhead is measurable. In a stronger version, replace raw softmax routing by 3--10 Sinkhorn or projected-gradient steps minimizing the regularized system-optimal objective in the fourth formula; the resulting \(x^*(d)\) is the inner response. Backpropagate through these unrolled steps and update router and expert parameters with the robust loss, optionally mixing it with nominal loss using coefficient \(\alpha\). The paper supplies the uncertainty geometry and worst-case bilevel structure; the engineer estimates \(\bar d_k\), \(\delta_k\), and expert capacities empirically. First test on a small 8-expert Switch-Transformer-style model on C4 or synthetic mixtures with deliberately shifted token-group frequencies. Compare ordinary auxiliary load-balancing loss, independent coordinate perturbations, and budgeted robust routing at equal FLOPs. Success means lower maximum expert utilization and fewer overflow or dropped tokens under held-out correlated demand shifts, with no more than 1--2% nominal perplexity degradation and ideally improved loss under distribution shift.