Elamur Research⌕
☰
home / papers / 2024 / Not All Experts are Equal: Efficient Expert Pruning and Skipping for Mixture-of-Experts Large Language Models

Not All Experts are Equal: Efficient Expert Pruning and Skipping for Mixture-of-Experts Large Language Models

Paper · 2024 · Xudong Lu, Qi Liu, Yuhui Xu et al.

Summary

Mixture-of-Experts LLMs such as Mixtral 8x7B activate only a few experts per token, yet every expert must stay in memory. The experts make up roughly 96% of the weights, so deployment is expensive. Unstructured or N:M weight pruning can shrink models, but realising its gains requires special hardware. The authors instead sparsify at the level of whole experts, with two post-training, plug-and-play techniques:

  1. Expert pruning. For each MoE layer, they enumerate every subset of experts that could be kept. Using a small calibration set, they pick the subset whose outputs best reconstruct the original layer outputs. This works in a general (task-agnostic) setting and, for the first time, a task-specific one.
  2. Dynamic skipping. At inference, a token's secondary expert is dropped when its routing weight is small relative to the top expert's.

On Mixtral 8x7B and Mixtral 8x7B Instruct, removing two experts per layer lets the model fit on one 80G GPU instead of two, with moderate accuracy loss. Combining pruning with skipping improves the speed–accuracy trade-off further. This makes large MoE models cheaper to run on ordinary hardware.

Key points

  • Loading Mixtral 8x7B in bf16 needs at least two A100-80G GPUs, and the experts account for about 45B of its 47B parameters.
  • Keeping 6 of 8 experts per layer halves the GPU count and gives a ~1.2× inference speedup. The accuracy drop is ~2.9 points for the task-agnostic model and ~6.2 points for the task-specific one, which falls to 1.6 after task-specific fine-tuning.
  • Evaluation uses calibration data from C4 and zero-shot accuracy on 8 tasks from EleutherAI LM Harness.
  • Pruning plus dynamic skipping keeps nearly 90% of Mixtral 8x7B Instruct's performance with half the parameters and a 1.33× token-generation speedup. Pruning to 6 experts plus skipping matches the speed of pruning to 4 experts alone, with clearly higher accuracy.
  • With 6 experts kept, the subsets chosen from C4 and from MATH calibration data coincide in only four layers (2, 4, 16, 31). This shows that which experts matter depends strongly on the domain.

Interesting ideas

  • I1. Brute-force search when the space is small. When a layer has only a handful of components (8 experts), exhaustively testing all subsets against a reconstruction loss is cheap and avoids heuristic importance scores. The same idea could help prune attention heads or adapters in other small modular setups. The paper itself notes that it stops scaling beyond a few dozen experts.
  • I2. Structured removal of whole units. Dropping entire modules, rather than scattered individual weights, gives memory and speed savings on standard GPUs without custom kernels. This is a useful principle when compressing models for cheap deployment.
  • I3. Calibration data as a domain selector. Because the retained experts differ sharply between C4 and MATH, a small in-domain calibration set could yield a specialised, smaller model. For example, Ukrainian-text calibration might keep the experts most useful for Ukrainian NLP.
  • I4. Ratio-based skipping of a weak second choice. Skipping the second expert whenever its routing weight is small relative to the top expert's is an almost free runtime trick. Similar thresholds could serve any top-k routing or ensemble.
  • I5. Greedy layer-wise search over-fitting at high sparsity. In the ablation, the layer-by-layer progressive search did slightly better at mild pruning but worse at 50% pruning, which the authors attribute to over-fitting the small calibration set. This is a reminder that more adaptive search is not automatically better when calibration data is limited.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Institutions

🏢SalesforceSan Francisco · United States37 works · 15,019 citations🎓Shanghai Jiao Tong UniversityShanghai · China63 works · 15,047 citations🔬Shanghai Artificial Intelligence LaboratoryShanghai · China44 works · 10,162 citations🔬Centre for Perceptual and Interactive IntelligenceHong Kong · China5 works · 1,085 citations

📚 Shelves 5The field of science this work belongs to (OpenAlex fields).Computer Science · 9,086 worksArtificial Intelligence · 8,889 worksArts and Humanities · 9,296 worksLanguage and Linguistics · 9,206 worksComputer Vision and Pattern Recognition · 1,282 works

🧭 Topics 2Research subjects and languages the work is about.Language models & pretraining · 1,343 worksEfficient models & compression · 77 works

Cite

Cite as Xudong Lu, Qi Liu, Yuhui Xu, Aojun Zhou, Siyuan Huang, Bo Zhang, Junchi Yan and Hongsheng Li. 2024. Not All Experts are Equal: Efficient Expert Pruning and Skipping for Mixture-of-Experts Large Language Models. Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers).

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.18653/v1/2024.acl-long.334

@inproceedings{lu2024not,
  title = {{Not All Experts are Equal: Efficient Expert Pruning and Skipping for Mixture-of-Experts Large Language Models}},
  author = {Xudong Lu, Qi Liu, Yuhui Xu, Aojun Zhou, Siyuan Huang, Bo Zhang, Junchi Yan and Hongsheng Li},
  booktitle = {Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)},
  year = {2024},
  doi = {10.18653/v1/2024.acl-long.334},
  url = {https://doi.org/10.18653/v1/2024.acl-long.334}
}

References

36 references, 12 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Keivan Alizadeh, Iman Mirzadeh, Dmitry Belenko, Karen Khatamifard, Minsik Cho, Carlo C Del Mundo, Mohammad Rastegari, and Mehrdad Farajtabar. 2023. Llm in a flash: Efficient large language model inference with limited memory.
  2. Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. 2020. Language models are few-shot learners. Advances in neural information processing systems, 33:1877–1901.
  3. Tianyu Chen, Shaohan Huang, Yuan Xie, Binxing Jiao, Daxin Jiang, Haoyi Zhou, Jianxin Li, and Furu Wei. 2022. Task-specific expert pruning for sparse mixture-of-experts. arXiv preprint arXiv:2206.00277. — [arXiv:2206.00277](https://arxiv.org/abs/2206.00277) find
  4. Zewen Chi, Li Dong, Shaohan Huang, Damai Dai, Shuming Ma, Barun Patra, Saksham Singhal, Payal Bajaj, Xia Song, Xian-Ling Mao, Heyan Huang, and Furu Wei. 2022. On the representation collapse of sparse mixture of experts. In Advances in Neural Information Processing Systems. find
  5. Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et al. 2021. Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14168.
  6. Casas, Emma Bou Hanna, Florian Bressand, et al. 2024. Mixtral of experts. arXiv preprint arXiv:2401.04088. find
  7. Shuangrui Ding, Peisen Zhao, Xiaopeng Zhang, Rui Qian, Hongkai Xiong, and Qi Tian. 2023. Prune spatio-temporal tokens by semantic-aware temporal accumulation. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 16945–16956. find
  8. Sehoon Kim, Sheng Shen, David Thorsley, Amir Gholami, Woosuk Kwon, Joseph Hassoun, and Kurt Keutzer. 2022. Learned token pruning for transformers. In Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pages 784–794. find
  9. William Fedus, Barret Zoph, and Noam Shazeer. 2022. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. The Journal of Machine Learning Research, 23(1):5232– 5270. find
  10. Young Jin Kim, Ammar Ahmad Awan, Alexandre Muzio, Andres Felipe Cruz Salinas, Liyang Lu, Amr Hendy, Samyam Rajbhandari, Yuxiong He, and Hany Hassan Awadalla. 2021. Scalable and efficient moe training for multitask multilingual models. arXiv preprint arXiv:2109.10465. — [arXiv:2109.10465](https://arxiv.org/abs/2109.10465) find
  11. Elias Frantar and Dan Alistarh. 2023. Sparsegpt: Massive language models can be accurately pruned in one-shot. In International Conference on Machine Learning, pages 10323–10337. PMLR. find
  12. Elias Frantar, Saleh Ashkboos, Torsten Hoefler, and Dan Alistarh. 2022. Gptq: Accurate post-training quantization for generative pre-trained transformers. arXiv preprint arXiv:2210.17323. find
  13. Trevor Gale, Deepak Narayanan, Cliff Young, and Matei Zaharia. 2023. Megablocks: Efficient sparse training with mixture-of-experts. Proceedings of Machine Learning and Systems, 5. find
  14. Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac’h, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Ani… find
  15. Yihui He, Xiangyu Zhang, and Jian Sun. 2017. Channel pruning for accelerating very deep neural networks. In Proceedings of the IEEE international conference on computer vision, pages 1389–1397. find
  16. Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. 2021. Measuring mathematical problem solving with the math dataset. NeurIPS.
  17. Itay Hubara, Brian Chmiel, Moshe Island, Ron Banner, Joseph Naor, and Daniel Soudry. 2021. Accelerated sparse neural training: A provable and efficient method to find n: m transposable masks. Advances in neural information processing systems, 34:21099– 21111. find
  18. Robert A Jacobs, Michael I Jordan, Steven J Nowlan, and Geoffrey E Hinton. 1991. Adaptive mixtures of local experts. Neural computation, 3(1):79–87. find
  19. Albert Q Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Yeskendir Koishekenov, Alexandre Berard, and Vassilina Nikoulina. 2022. Memory-efficient nllb-200: Language-specific expert pruning of a massively multilingual mac… find
  20. Woosuk Kwon, Sehoon Kim, Michael W Mahoney, Joseph Hassoun, Kurt Keutzer, and Amir Gholami. 2022. A fast post-training pruning framework for transformers. Advances in Neural Information Processing Systems, 35:24101–24116. find
  21. Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. 2020. Gshard: Scaling giant models with conditional computation and automatic sharding. arXiv preprint arXiv:2006.16668. find
  22. Ji Lin, Jiaming Tang, Haotian Tang, Shang Yang, Xingyu Dang, and Song Han. 2023. Awq: Activation-aware weight quantization for llm compression and acceleration. arXiv preprint arXiv:2306.00978. find
  23. Asit Mishra, Jorge Albericio Latorre, Jeff Pool, Darko Stosic, Dusan Stosic, Ganesh Venkatesh, Chong Yu, and Paulius Micikevicius. 2021. Accelerating sparse deep neural networks. arXiv preprint arXiv:2104.08378. find
  24. OpenAI. 2023. Gpt-4 technical report. abs/2303.08774.
  25. ArXiv, Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. 2019. Exploring the limits of transfer learning with a unified text-to-text transformer. arXiv e-prints.
  26. Omar Sanseviero, Lewis Tunstall, Philipp Schmid, Sourab Mangrulkar, Younes Belkada, and Pedro Cuenca. 2023. Mixture of experts explained. find
  27. Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. 2017. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. arXiv preprint arXiv:1701.06538. find
  28. Mingjie Sun, Zhuang Liu, Anna Bair, and J Zico Kolter. 2023. A simple and effective pruning approach for large language models. arXiv preprint arXiv:2306.11695. find
  29. Gemini Team, Rohan Anil, Sebastian Borgeaud, Yonghui Wu, Jean-Baptiste Alayrac, Jiahui Yu, Radu Soricut, Johan Schalkwyk, Andrew M Dai, Anja Hauth, et al. 2023. Gemini: a family of highly capable multimodal models. arXiv preprint arXiv:2312.11805.
  30. Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. 2023a. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971.
  31. Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. 2023b. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288. find
  32. Ke Wang, Houxing Ren, Aojun Zhou, Zimu Lu, Sichun Luo, Weikang Shi, Renrui Zhang, Linqi Song, Mingjie Zhan, and Hongsheng Li. 2024. Mathcoder: Seamless code integration in LLMs for enhanced mathematical reasoning. In The Twelfth International Conference on Learning Representations.
  33. Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pier-ric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, et al. 2020. Transformers: State-of-the-art natural language processing. In Proceedings of the 2020 conference on empirical methods in natural language p…
  34. Longhui Yu, Weisen Jiang, Han Shi, Jincheng Yu, Zhengying Liu, Yu Zhang, James T Kwok, Zhenguo Li, Adrian Weller, and Weiyang Liu. 2023. Metamath: Bootstrap your own mathematical questions for large language models. arXiv preprint arXiv:2309.12284.
  35. Aojun Zhou, Yukun Ma, Junnan Zhu, Jianbo Liu, Zhijie Zhang, Kun Yuan, Wenxiu Sun, and Hongsheng Li. 2021. Learning n: m fine-grained structured sparse neural networks from scratch. arXiv preprint arXiv:2102.04010. find
  36. Aojun Zhou, Ke Wang, Zimu Lu, Weikang Shi, Sichun Luo, Zipeng Qin, Shaoqing Lu, Anya Jia, Linqi Song, Mingjie Zhan, and Hongsheng Li. 2024. Solving challenging math word problems using GPT-4 code interpreter with code-based self-verification. In The Twelfth International Conference on Learning Repre…

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/lu-2024-not-all-experts-are-equal/