STAR: Rethinking MoE Routing as Structure-Aware Subspace Learning

Mixture-of-Experts (MoE) scales model capacity efficiently by selectively routing inputs to a specialized subset of experts. However, input-expert specialization, the core motivation of MoE, critically depends on whether the router is actually aware of input structure. In practice, MoE routing is typically implemented as a shallow linear projection with limited awareness of input representation, which often leads to unstable routing. We propose STAR, a Structure Aware Routing that rethinks MoE routing as a subspace learning problem by augmenting standard learnable routing with an evolving principal subspace that tracks dominant input structure via Generalized Hebbian Algorithm (GHA). By aligning routing decisions directly with input structure, STAR enables stable expert specialization. We evaluate STAR on controlled synthetic setup and large-scale language and vision tasks, where it consistently improves routing quality and downstream performance over strong MoE baselines. Moreover, optional test-time subspace updates further enhance routing robustness and generalization under input distribution shifts. Code is available at \url{https://github.com/psmiz/STAR}.

Modern AI models are growing so large that running them is becoming prohibitively expensive. Mixture-of-Experts (MoE) is a popular workaround: the model is split into many smaller experts, and a routing component decides which experts handle each input, so only a small fraction of the model is used at a time. The catch is that the router is usually a simple linear component that does not really see the structure of its inputs, so experts often specialize poorly and some get overused while others sit idle. We propose STAR, a router that explicitly tracks the dominant patterns in the data it sees. As training proceeds, STAR continuously estimates the main directions along which inputs vary, and uses these to make more informed routing decisions. The router can even keep adapting to new data after training, without any extra labels or fine-tuning. Across language and vision benchmarks, STAR produces more reliable expert specialization, better accuracy, and stronger robustness when test data differs from training. More broadly, we show that grounding routing in the actual structure of the data, rather than relying on a simple learned shortcut, is a promising direction as AI models continue to scale.