Grassmannian Mixture-of-Experts: Concentration-Controlled Routing on Subspace Manifolds

TL;DR

This paper introduces Grassmannian MoE, a novel routing framework on subspace manifolds that offers a continuous, interpretable control over sparsity and routing entropy, improving load balancing and interpretability in mixture-of-experts models.

Contribution

We propose a new Grassmannian MoE framework using concentration parameters of Matrix Bingham distributions for principled, smooth sparsity control and uncertainty-aware expert routing.

Findings

Achieves 0% routing collapse across all seeds.

Improves load balance by 15-30%.

Provides interpretable, heterogeneous expert specialization.

Abstract

Mixture-of-Experts models rely on learned routers to assign tokens to experts, yet standard softmax gating provides no principled mechanism to control the tradeoff between sparsity and utilization. We propose Grassmannian MoE (GrMoE), a routing framework that operates on the Grassmannian manifold of subspaces, where gating weights arise from the concentration parameters of Matrix Bingham distributions. This construction yields a single, interpretable knob -- the concentration matrix $\Lambda$ -- that continuously controls routing entropy, replacing discrete top-$k$ selection with a smooth, geometrically principled sparsity mechanism. We further develop an amortized variational inference procedure for posterior routing distributions, enabling uncertainty-aware expert assignment that naturally resists expert collapse. We formally prove tight bounds relating the Bingham concentration spectrum to routing entropy, expected top-$k$ mass, and an exponential bound on expert collapse, establishing the first formal theory of concentration-controlled sparsity. On synthetic routing tasks, a 350M-parameter MoE language model with 8 experts, a 1.3B-parameter model with 16 experts, and a 2.7B-parameter model with 32 experts, GrMoE achieves 0\% routing collapse across all seeds, comparable or better perplexity with 15--30\% improved load balance, and a smooth monotonic relationship between concentration and effective sparsity that enables post-hoc sparsity tuning without retraining. Token-level analysis reveals that experts learn heterogeneous concentration values that correlate with linguistic specialization, providing interpretable routing behavior.