Towards Understanding the Mixture-of-Experts Layer in Deep Learning
Zixiang Chen, Yihe Deng, Yue Wu, Quanquan Gu, Yuanzhi Li
Abstract
The Mixture-of-Experts (MoE) layer, a sparsely-activated model controlled by a router, has achieved great success in deep learning. However, the understanding of such architecture remains elusive. In this paper, we formally study how the MoE layer improves the performance of neural network learning and why the mixture model will not collapse into a single model. Our empirical results suggest that the cluster structure of the underlying problem and the non-linearity of the expert are pivotal to the success of MoE. This motivates us to consider a challenging classification problem with intrinsic cluster structures. Theoretically, we proved that this problem is hard to solve by a single expert such as a two-layer convolutional neural network (CNN). Yet with the MoE layer with each expert being a two-layer CNN, the problem can be solved successfully. In particular, our theory shows that the router can learn the cluster-center features, which helps divide the input complex problem into simpler classification sub-problems that individual experts can conquer. To our knowledge, this is the first theoretical result toward formally understanding the mechanism of the MoE layer for deep learning.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Cited by top-tier papers47
- Pushing Mixture of Experts to the Limit: Extremely Parameter Efficient MoE for Instruction TuningTed Zadouri, Ahmet Üstün, Arash Ahmadian, Beyza Ermis et al.ICLR 2024 · 169 citations
- Localizing Task Information for Improved Model Merging and CompressionKe Wang, Nikolaos Dimitriadis, Guillermo Ortiz-Jiménez, François Fleuret et al.ICML 2024 · 107 citations
- MVMoE: Multi-Task Vehicle Routing Solver with Mixture-of-ExpertsJianan Zhou, Zhiguang Cao, Yaoxin Wu, Wen Song et al.ICML 2024 · 74 citations
- Robust Learning with Progressive Data Expansion Against Spurious CorrelationYihe Deng, Yu Yang, Baharan Mirzasoleiman, Quanquan GuNeurIPS 2023 · 53 citations
- Morphological Prototyping for Unsupervised Slide Representation Learning in Computational PathologyAndrew H. Song, Richard J. Chen, Tong Ding, Drew F. K. Williamson et al.CVPR 2024 · 51 citations
Builds on8
- BASE Layers: Simplifying Training of Large, Sparse ModelsMike Lewis, Shruti Bhosale, Tim Dettmers, Naman Goyal et al.ICML 2021 · 382 citations
- Hash Layers For Large Sparse ModelsStephen Roller, Sainbayar Sukhbaatar, Arthur Szlam, Jason WestonNeurIPS 2021 · 316 citations
- Toward Understanding the Feature Learning Process of Self-supervised Contrastive LearningZixin Wen, Yuanzhi LiICML 2021 · 162 citations
- Towards Understanding Ensemble, Knowledge Distillation and Self-Distillation in Deep LearningZeyuan Allen-Zhu, Yuanzhi LiICLR 2023 · 151 citations
- Beyond Linearization: On Quadratic and Higher-Order Approximation of Wide Neural NetworksYu Bai, Jason D. LeeICLR 2020 · 128 citations
Related papers
- Patch-level Routing in Mixture-of-Experts is Provably Sample-efficient for Convolutional Neural NetworksMohammed Nowaz Rabbani Chowdhury, Shuai Zhang, Meng Wang, Sijia Liu et al.ICML 2023 · 45 citations
- Mixture of Experts Provably Detect and Learn the Latent Cluster Structure in Gradient-Based LearningRyotaro Kawata, Kohsei Matsutani, Yuri Kinoshita, Naoki Nishikawa et al.ICML 2025
- Tight Clusters Make Specialized ExpertsStefan K. Nielsen, Rachel S. Y. Teo, Laziz U. Abdullaev, Tan Minh NguyenICLR 2025
- Theory on Mixture-of-Experts in Continual LearningHongbo Li, Sen Lin, Lingjie Duan, Yingbin Liang et al.ICLR 2025
- MoEC: Mixture of Expert ClustersYuan Xie, Shaohan Huang, Tianyu Chen, Furu WeiAAAI 2023 · 27 citations
