Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

DA-MoE: Towards Dynamic Expert Allocation for Mixture-of-Experts Models

Maryam Akhavan Aghdam, Hongpeng Jin, Yanzhao Wu

TL;DR

This study proposes a novel dynamic router mechanism that Dynamically Allocates a variable number of experts for Mixture-of-Experts (DA-MoE) models based on an effective token importance measure and demonstrates that the DA-MoE approach consistently outperforms the state-of-the-art Transformer based MoE model on the popular GLUE benchmark.

Abstract

Transformer-based Mixture-of-Experts (MoE) models have been driving several recent technological advancements in Natural Language Processing (NLP). These MoE models adopt a router mechanism to determine which experts to activate for routing input tokens. However, existing router mechanisms allocate a fixed number of experts to each token, which neglects the varying importance of different input tokens. In this study, we propose a novel dynamic router mechanism that Dynamically Allocates a variable number of experts for Mixture-of-Experts (DA-MoE) models based on an effective token importance measure. First, we show that the Transformer attention mechanism provides a natural and effective way of calculating token importance. Second, we propose a dynamic router mechanism that effectively decides the optimal number of experts (K) and allocates the top-K experts for each input token. Third, comprehensive experiments on several benchmark datasets demonstrate that our DA-MoE approach consistently outperforms the state-of-the-art Transformer based MoE model on the popular GLUE benchmark.

DA-MoE: Towards Dynamic Expert Allocation for Mixture-of-Experts Models

TL;DR

This study proposes a novel dynamic router mechanism that Dynamically Allocates a variable number of experts for Mixture-of-Experts (DA-MoE) models based on an effective token importance measure and demonstrates that the DA-MoE approach consistently outperforms the state-of-the-art Transformer based MoE model on the popular GLUE benchmark.

Abstract

Transformer-based Mixture-of-Experts (MoE) models have been driving several recent technological advancements in Natural Language Processing (NLP). These MoE models adopt a router mechanism to determine which experts to activate for routing input tokens. However, existing router mechanisms allocate a fixed number of experts to each token, which neglects the varying importance of different input tokens. In this study, we propose a novel dynamic router mechanism that Dynamically Allocates a variable number of experts for Mixture-of-Experts (DA-MoE) models based on an effective token importance measure. First, we show that the Transformer attention mechanism provides a natural and effective way of calculating token importance. Second, we propose a dynamic router mechanism that effectively decides the optimal number of experts (K) and allocates the top-K experts for each input token. Third, comprehensive experiments on several benchmark datasets demonstrate that our DA-MoE approach consistently outperforms the state-of-the-art Transformer based MoE model on the popular GLUE benchmark.
Paper Structure (14 sections, 10 equations, 7 figures, 2 tables, 1 algorithm)

This paper contains 14 sections, 10 equations, 7 figures, 2 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Problem Statement
  4. DA-MoE Overview
  5. DA-MoE Model Architecture
  6. Dynamic Expert Allocation Algorithm
  7. Experimental Analysis
  8. Experimental Setup
  9. Pre-training Settings
  10. Fine-tuning Settings
  11. Pre-training Evaluation
  12. Fine-tuning Evaluation
  13. Token Importance Analysis
  14. Conclusion

Figures (7)

  • Figure 1: An example input sentence ("The movie was incredibly inspiring.") for performing sentiment analysis
  • Figure 2: Illustration of a DA-MoE encoder block with a dynamic routing mechanism. DA-MoE introduces a dynamic routing mechanism, allowing each token to be assigned to the top-K experts based on token importance. For example, token one is assigned to four experts, while token two is assigned to three experts.
  • Figure 3: Training log perplexity scaling comparison between DA-MoE and ST (Switch Transformer) base models.
  • Figure 4: Attention weights for 12 heads in the last layer for sentiment analysis task for the sentence ("The movie was incredibly inspiring.")
  • Figure 5: Attention weights for 12 heads in the last layer for paraphrasing task for the sentence ("He said he would come. He mentioned he was coming.")
  • ...and 2 more figures