Control-Cost Laguerre MoE Routing

(1) Integration point: replace the top-1 router in a sparse MoE layer. Let the preceding transformer block produce token vectors $h_b\in\mathbb R^d$. Add a projection $x_b=P h_b$ into a lower-dimensional routing space, maintain one prototype $y_i\in\mathbb R^d$ per expert, and compute the control cost to every expert. Use $A,B,\tau$ as fixed small matrices or learn them from a stable parameterization; precompute $F=e^{A\tau}$ and $G=W_\tau^{-1}$, so the cost is a batched quadratic form. (2) Pseudocode: initialize $\psi_i=0$ and capacities $\nu_i=1/r$; for each batch compute $q_b=F x_b$, $d_{bi}=0.5(y_i-q_b)^\top G(y_i-q_b)-\psi_i$; assign $a_b=\arg\min_i d_{bi}$; dispatch tokens by $a_b$; estimate $\hat m_i=B^{-1}\sum_b1[a_b=i]$; update $\psi_i\leftarrow\psi_i+\eta_\psi(\nu_i-\hat m_i)$ and optionally center $\psi\leftarrow\psi-\mathrm{mean}(\psi)$. Train the experts and projection with the ordinary task loss. Use a straight-through assignment or stop-gradient through routing. (3) Computed from the mathematics: $F$ and $G$ come from matrix exponential and controllability-Gramian calculations; cells and assignments use the exact Control Laguerre rule; only minibatch masses and learned prototypes are estimated or optimized empirically. If $W_\tau$ is ill-conditioned, use $(W_\tau+\epsilon I)^{-1}$. (4) First experiment: compare this router against Switch top-1 routing with its standard load-balancing loss on a 4- or 8-expert small Transformer trained on CIFAR-100 text labels, WikiText-2, or a synthetic mixture with clusters connected by anisotropic dynamics. Measure validation loss or perplexity, coefficient of variation of expert loads, dropped-token fraction, router overhead, and training stability across seeds. Success means lower load variance and fewer dropped tokens at equal capacity, with equal or better validation loss and no extra all-to-all communication.