On the Representation Collapse of Sparse Mixture of Experts
This addresses a fundamental limitation in scaling mixture-of-experts models for practitioners, though it appears incremental as it builds on existing routing mechanisms.
The paper tackles the representation collapse problem in sparse mixture of experts models, where token clustering around expert centroids occurs during routing mechanism learning. The proposed method estimates routing scores on a low-dimensional hypersphere, achieving consistent gains across seven multilingual benchmarks and alleviating representation collapse.
Sparse mixture of experts provides larger model capacity while requiring a constant computational overhead. It employs the routing mechanism to distribute input tokens to the best-matched experts according to their hidden representations. However, learning such a routing mechanism encourages token clustering around expert centroids, implying a trend toward representation collapse. In this work, we propose to estimate the routing scores between tokens and experts on a low-dimensional hypersphere. We conduct extensive experiments on cross-lingual language model pre-training and fine-tuning on downstream tasks. Experimental results across seven multilingual benchmarks show that our method achieves consistent gains. We also present a comprehensive analysis on the representation and routing behaviors of our models. Our method alleviates the representation collapse issue and achieves more consistent routing than the baseline mixture-of-experts methods.