ERMoE: Eigen-Reparameterized Mixture-of-Experts for Stable Routing and Interpretable Specialization

摘要

Mixture-of-Experts (MoE) models expand capacity via sparse expert activation, but routing logits can misalign with expert structure (unstable routing, underutilization) and load imbalance can create stragglers. Auxiliary load-balancing losses reduce disparity but often weaken specialization and downstream accuracy. We propose ERMoE, a sparse MoE transformer that reparameterizes each expert in a learned orthonormal eigenbasis and routes with an Eigenbasis Score (cosine similarity between token features and an expert basis) instead of learned gating logits. By tying assignments to each expert's representation space, ERMoE stabilizes utilization, improves interpretability, and removes explicit balancing losses and their gradient interference. ERMoE reaches state-of-the-art results on ImageNet and image-text retrieval (COCO, Flickr30K) with flatter expert loads. A 3D MRI variant (ERMoE-ba) improves brain age prediction by over 7% and yields anatomically interpretable expert specializations.