On the Representation Collapse of Sparse Mixture of Experts

By Zewen Chi, Li Dong, Shaohan Huang, Damai Dai, Shuming Ma, Barun Patra, Saksham Singhal, Payal Bajaj, Xia Song, Furu Wei. 2022

199 citations recorded by the source corpus.

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.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Information Systems, Model Architectures, Efficiency & Compression

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

1. Furu Wei
2. Li Dong
3. Shaohan Huang
4. Shuming Ma
5. Wenhui Wang
6. Jianfeng Gao
7. Saksham Singhal
8. Zewen Chi

In Computer Vision and Pattern Recognition

1. Furu Wei
2. Luc Van Gool
3. Li Dong
4. Dacheng Tao
5. Wenhui Wang
6. Ming-Hsuan Yang
7. Jianfeng Gao
8. Alan L. Yuille

In Information Systems

1. Furu Wei
2. Li Dong
3. Shuming Ma
4. Shaohan Huang
5. Vishrav Chaudhary
6. Lei Cui
7. Qiang Liu
8. Fan Yang

In Model Architectures

1. Furu Wei
2. Wenhui Wang
3. Li Dong
4. Shuming Ma
5. Shaohan Huang
6. Ze Liu
7. Han Hu
8. Vishrav Chaudhary

In Efficiency & Compression

1. Vishrav Chaudhary
2. Li Dong
3. Han Hu
4. Shuming Ma
5. Shaohan Huang
6. Qiang Liu
7. Fan Yang
8. Furu Wei

Citations in the corpus

• Language Is Not All You Need: Aligning Perception with Language Models (2023, methodology): “In the future, we would like to scale up K OSMOS -1 in terms of model size [13, 14, 60], and integrate the speech [12] capability into K OSMOS -1.”
• Language Models are General-Purpose Interfaces (2022, methodology): “After we obtain the input representations, we feed them to the causal language model, which has shown promising performance on commonsense reasoning (Chowdhery et al., 2022) and planning (Huang et al., 2022).”
• Tutel: Adaptive Mixture-of-Experts at Scale (2022, background): “because: 1) its normalization effect on input may lead to more stable routing when the amplitude or dimension of the input feature is scaled; 2) There is a concurrent work showing that the cosine router is more accurate in cross-lingual language tasks [6].”
• TorchScale: Transformers at Scale (2022, methodology): “T ORCH S CALE implement X-MoE [Chi et al., 2022], a variant of sparse MoE model.”
• A Survey on Mixture of Experts in Large Language Models (2024, unknown): “This challenge has led to the development of parameter-efficient fine-tuning (PEFT) methods, which update only a small set of parameters, achieving remarkable performances in many NLP tasks with less computational cost [150], [151].”