Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Harder Tasks Need More Experts: Dynamic Routing in MoE Models

Quzhe Huang, Zhenwei An, Nan Zhuang, Mingxu Tao, Chen Zhang, Yang Jin, Kun Xu, Kun Xu, Liwei Chen, Songfang Huang, Yansong Feng

TL;DR

This work targets the inefficiency of fixed Top-K routing in Mixture of Experts (MoE) models by introducing a dynamic routing mechanism that activates experts based on input difficulty, controlled by a threshold p. It couples a confidence-based routing strategy with two auxiliary losses (load balance and dynamic entropy) to encourage minimal yet sufficient expert activation, and demonstrates that dynamic routing yields gains over Top-2 routing while using fewer parameters. Empirically, the approach achieves average improvements on downstream tasks, with notable gains on BBH, and reveals meaningful layer-wise and task-specific patterns in expert allocation. The method also reports efficiency benefits in both training and inference and provides insights for designing heterogeneous MoE architectures. Code and models are open-sourced to foster further research in dynamic MoE routing.

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.

Harder Tasks Need More Experts: Dynamic Routing in MoE Models

TL;DR

This work targets the inefficiency of fixed Top-K routing in Mixture of Experts (MoE) models by introducing a dynamic routing mechanism that activates experts based on input difficulty, controlled by a threshold p. It couples a confidence-based routing strategy with two auxiliary losses (load balance and dynamic entropy) to encourage minimal yet sufficient expert activation, and demonstrates that dynamic routing yields gains over Top-2 routing while using fewer parameters. Empirically, the approach achieves average improvements on downstream tasks, with notable gains on BBH, and reveals meaningful layer-wise and task-specific patterns in expert allocation. The method also reports efficiency benefits in both training and inference and provides insights for designing heterogeneous MoE architectures. Code and models are open-sourced to foster further research in dynamic MoE routing.

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.
Paper Structure (31 sections, 11 equations, 4 figures, 4 tables)

This paper contains 31 sections, 11 equations, 4 figures, 4 tables.

Table of Contents

  1. Introduction
  2. Method
  3. Top-K Routing MoE
  4. Dynamic Routing MoE
  5. Loss
  6. Dynamic Loss
  7. Load Balance Loss
  8. Final Loss
  9. Experiments
  10. Settings
  11. Training data
  12. Model Settings
  13. Evaluation
  14. Experiment Models
  15. Dense
  16. ...and 16 more sections

Figures (4)

  • Figure 1: Comparison between Top-K routing mechanism and Top-P routing mechanism. (a) Each token selects fixed K=2 experts with Top-K routing probabilities. (b) In Top-P routing mechanism, each token selects experts with higher routing probabilities until the cumulative probability exceeds threshold.
  • Figure 2: Average scores of MoE-Dynamic with different threshold $p$ on downstream tasks
  • Figure 3: Average activated experts number across training procedure.
  • Figure 4: Activated experts in different layers