Harder Tasks Need More Experts: Dynamic Routing in MoE Models

Quzhe Huang, Zhenwei An, Nan Zhuang, Mingxu Tao, Chen Zhang +6 more
3/12/2024
cs.LGcs.CL

Official code: ZhenweiAn/Dynamic_MoE (74 stars, last commit 2024-07-30).

Abstract

In this paper, we introduce a novel dynamic expert selection framework for Mixture of Experts (MoE) models, aiming to enhance computational efficiency and model performance by adjusting the number of activated experts based on input difficulty. Unlike traditional MoE approaches that rely on fixed Top-K routing, which activates a predetermined number of experts regardless of the input's complexity, our method dynamically selects experts based on the confidence level in expert selection for each input. This allows for a more efficient utilization of computational resources, activating more experts for complex tasks requiring advanced reasoning and fewer for simpler tasks. Through extensive evaluations, our dynamic routing method demonstrates substantial improvements over conventional Top-2 routing across various benchmarks, achieving an average improvement of 0.7% with less than 90% activated parameters. Further analysis shows our model dispatches more experts to tasks requiring complex reasoning skills, like BBH, confirming its ability to dynamically allocate computational resources in alignment with the input's complexity. Our findings also highlight a variation in the number of experts needed across different layers of the transformer model, offering insights into the potential for designing heterogeneous MoE frameworks. The code and models are available at https://github.com/ZhenweiAn/Dynamic_MoE.

View on arXivView PDF

Code Implementations(1)

Inference Code for Paper "Harder Tasks Need More Experts: Dynamic Routing in MoE Models"

749PythonMar 12, 20242 years ago

Similar Papers

Cite this paper

@article{huang2024harder,
  title  = {Harder Tasks Need More Experts: Dynamic Routing in MoE Models},
  author = {Quzhe Huang and Zhenwei An and Nan Zhuang and Mingxu Tao and Chen Zhang and Yang Jin and Kun Xu and Kun Xu and Liwei Chen and Songfang Huang and Yansong Feng},
  year   = {2024},
  eprint = {2403.07652},
  archivePrefix = {arXiv},
  url    = {http://arxiv.org/abs/2403.07652v1}
}

Add a badge to your README

Shows whether official code is linked for this paper.

Code of Paper badge
[![Code of Paper](https://codeofpaper.com/badge/2403.07652.svg)](https://codeofpaper.com/paper/2403.07652)

Discussion