Fisher–Rao geodesic router

Integrate the method into the router probability tensor immediately after its softmax and before top-k selection or mixture aggregation. Given router logits \(a\in\mathbb R^m\), compute \(\pi=\operatorname{softmax}(a)\), then \(x=2\sqrt{\pi+\varepsilon}\) with \(\varepsilon=10^{-8}\). For an unconstrained Fisher router, renormalize \(x\) to norm 2 and recover probabilities with \(\pi_i=(x_i/2)^2\). For a constrained experiment, choose an orthonormal matrix \(U\in\mathbb R^{m\times r}\), define \(P_S=UU^{\mathsf T}\), compute \(v=P_Sx\), add a small positive floor if some coordinates are negative, renormalize \(x'=2v/\|v\|_2\), and set \(\pi'_i=(x'_i/2)^2\). A differentiable positivity-preserving alternative is \(x=2\sqrt{\operatorname{softmax}(Uh)}\) for a learned low-dimensional vector \(h\). Keep the standard load-balancing loss unchanged. Compare ordinary softmax routing, Euclidean probability projection, and Fisher projection in a 4- or 8-expert Switch-style MLP on CIFAR-10 or a small language-model dataset. Measure validation loss, expert-load variance, router entropy, gradient spikes, and dead-expert frequency. Success is lower load variance or fewer routing instabilities at equal loss, or better validation loss under large router learning rates. Also test spherical interpolation between learned routing distributions and verify constant Fisher speed numerically.