Mean-Payoff Tropical Router

Integrate the controller after the MoE router computes token-to-expert assignments and before capacity truncation. Use one scalar bias \(b_e\) per expert and maintain a small set of constraint rows \(s\): global load balance, local capacity, and minimum router-margin constraints. At each training step: (1) collect batch assignment counts and average selected-versus-second-best margins; compute \(\ell_e\) and \(m_e\); (2) construct offsets \(w_{s,e}=\alpha(c_e-\ell_e)+\beta m_e\), where \(c_e\) is the target log-load and \(\alpha,\beta\) are fixed gains; (3) for every row, select \(a(s)=\arg\max_{e\in L_s}(b_e+w_{s,e})\); (4) build the directed graph from each constraint state \(p(s)\) to its selected expert or state, find strongly connected components, and estimate the highest-mean cycle using Karp's algorithm or exhaustive cycle enumeration when the graph is tiny; (5) if a cycle has positive mean \(W(C)\), subtract \(\eta W(C)\) from the source biases and repeat for two to five policy-improvement passes; (6) add the final \(b_e\) to router logits before top-1 or top-2 selection. The tropical part is the argmax policy and cycle correction; batch statistics and gains \(\alpha,\beta,\eta\) are estimated or tuned empirically. Start with a 12-layer, 8-expert Switch-Transformer-sized model on C4 or synthetic token classification, comparing the standard router auxiliary loss, Sinkhorn balancing, and this controller at equal model FLOPs. Measure expert-load coefficient of variation, dropped-token rate, router entropy, step-to-step assignment oscillation, and validation perplexity. Success means lower drop rate and oscillation at equal perplexity, or equal load balance with a smaller auxiliary-loss coefficient and no additional forward-pass cost.