Stochastic Online Linear Programming for Resource-Aware Expert Routing

Mixture-of-experts (MoE) models decouple parameter count from per-token compute, but require a router that matches tokens to experts while balancing the load across experts. We formulate MoE routing as a stochastic online linear program (SOLP), which allows a flexible formulation of matching-quality versus expert-balance tradeoffs through linear constraints. Specifically, while our formulation easily admits the traditional cardinality constraints on experts-per-token or tokens-per-expert, it can also capture substantially more flexible aggregate constraints on expert utilization across sequences, layers or even a whole training batch. Furthermore, we provide an elegant dual update algorithm, which finds a near-optimal solution as we establish in our main theoretical guarantee. Empirically, we evaluate various instantiations of our general SOLP framework across a seven-scale ladder built on the Gemma architecture, ranging from $183$M to $6.59$B parameters ($6$ to $30$ layers), and show that the flexibility afforded by our formulation results in better validation perplexity across this ladder, when compared with standard Top-$k$ routing approaches. We further show additional benefits of our dual update algorithm, where the dual parameters can be cheaply retrained for a new set of constraints at inference-time, yielding a powerful mechanism for inference time compute optimization, without any significant quality losses.