Stuck in Collapse: Understanding Routing Collapse in Sparse Mixture-of-Experts
Abstract
Mixture-of-Experts is a key architecture for scaling model capacity with limited computation. However, routing collapse, where token assignments concentrate on a small subset of experts, remains a long-standing failure mode that degrades both model performance and computational efficiency. To understand how routing collapse emerges and persists, we conduct a controlled pretraining study across two model scales and eight model variants. We find that routing collapse is not a gradual accumulation of imbalance, but a failure to correct routing perturbations, with correction dynamics that differ across routing operators. Tracing the underlying optimization mechanism, we observe that these perturbations coincide with rapid early reorganization of hidden representations, which can act as a common input-side source. Under sparse routing, selection-dependent asymmetry in router gradients can fail to correct the resulting bias. For unselected experts, these gradients either vanish entirely or lose expert-specific information and share a common scalar component. To validate this mechanism, we introduce a targeted router gradient intervention that restores expert-specific information for unselected experts while preserving sparse forward computation. It substantially alleviates routing imbalance without an auxiliary loss, providing evidence that router gradient dynamics for unselected experts play a key role in sustaining routing collapse.
est. 32% chance this paper gets accepted at ICLR 2027.
What do you think this paper will get?
All positions stay anonymous.