Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Multi-Head Mixture-of-Experts

Xun Wu, Shaohan Huang, Wenhui Wang, Furu Wei

TL;DR

Sparse Mixture-of-Experts (SMoE) scales model capacity but suffers from low expert activation and coarse token semantics. The authors propose Multi-Head Mixture-of-Experts (MH-MoE), which splits each token into h sub-tokens and routes them to multiple experts via a multi-head layer and a Token-Splitting-Merging (TSM) process, enabling denser activation and finer-grained representation without increasing computation. The approach is straightforward to implement and decoupled from other SMoE optimizations, and it yields consistent improvements in perplexity and downstream tasks across English-focused language modeling, multilingual language modeling, and masked multimodal modeling, with extensive ablations and analyses revealing denser activation and polysemous-token-specific routing. These results highlight MH-MoE as a practical path to scaling SMoE while preserving efficiency and enabling richer, cross-space representations.

Abstract

Sparse Mixtures of Experts (SMoE) scales model capacity without significant increases in training and inference costs, but exhibits the following two issues: (1) Low expert activation, where only a small subset of experts are activated for optimization. (2) Lacking fine-grained analytical capabilities for multiple semantic concepts within individual tokens. We propose Multi-Head Mixture-of-Experts (MH-MoE), which employs a multi-head mechanism to split each token into multiple sub-tokens. These sub-tokens are then assigned to and processed by a diverse set of experts in parallel, and seamlessly reintegrated into the original token form. The multi-head mechanism enables the model to collectively attend to information from various representation spaces within different experts, while significantly enhances expert activation, thus deepens context understanding and alleviate overfitting. Moreover, our MH-MoE is straightforward to implement and decouples from other SMoE optimization methods, making it easy to integrate with other SMoE models for enhanced performance. Extensive experimental results across three tasks: English-focused language modeling, Multi-lingual language modeling and Masked multi-modality modeling tasks, demonstrate the effectiveness of MH-MoE.

Multi-Head Mixture-of-Experts

TL;DR

Sparse Mixture-of-Experts (SMoE) scales model capacity but suffers from low expert activation and coarse token semantics. The authors propose Multi-Head Mixture-of-Experts (MH-MoE), which splits each token into h sub-tokens and routes them to multiple experts via a multi-head layer and a Token-Splitting-Merging (TSM) process, enabling denser activation and finer-grained representation without increasing computation. The approach is straightforward to implement and decoupled from other SMoE optimizations, and it yields consistent improvements in perplexity and downstream tasks across English-focused language modeling, multilingual language modeling, and masked multimodal modeling, with extensive ablations and analyses revealing denser activation and polysemous-token-specific routing. These results highlight MH-MoE as a practical path to scaling SMoE while preserving efficiency and enabling richer, cross-space representations.

Abstract

Sparse Mixtures of Experts (SMoE) scales model capacity without significant increases in training and inference costs, but exhibits the following two issues: (1) Low expert activation, where only a small subset of experts are activated for optimization. (2) Lacking fine-grained analytical capabilities for multiple semantic concepts within individual tokens. We propose Multi-Head Mixture-of-Experts (MH-MoE), which employs a multi-head mechanism to split each token into multiple sub-tokens. These sub-tokens are then assigned to and processed by a diverse set of experts in parallel, and seamlessly reintegrated into the original token form. The multi-head mechanism enables the model to collectively attend to information from various representation spaces within different experts, while significantly enhances expert activation, thus deepens context understanding and alleviate overfitting. Moreover, our MH-MoE is straightforward to implement and decouples from other SMoE optimization methods, making it easy to integrate with other SMoE models for enhanced performance. Extensive experimental results across three tasks: English-focused language modeling, Multi-lingual language modeling and Masked multi-modality modeling tasks, demonstrate the effectiveness of MH-MoE.
Paper Structure (25 sections, 16 equations, 11 figures, 11 tables, 1 algorithm)

This paper contains 25 sections, 16 equations, 11 figures, 11 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Background
  3. Method
  4. Multi-Head Mixture-of-Experts
  5. Training Objectives
  6. Experiments
  7. Experimental Setup
  8. Perplexity Evaluation
  9. Downstream Evaluation
  10. Ablation Studies
  11. Analysis
  12. Experts Activation Analysis
  13. Fine-grained understanding Analysis
  14. Complexity & Parameter Analysis.
  15. Conclusion
  16. ...and 10 more sections

Figures (11)

  • Figure 1: (a) Expert activation distribution on XNLI xnli, encompassing $6$ parallel expert layers with $32$ experts per layer. SMoE has many "dead" experts (dark) which are not activated, while MH-MoE leading to significantly increased usage of these experts. Experts activation ratio is determined by calculating the ratio of each expert's selection frequency in each MoE layer to the total number of tokens, where those exceeding a threshold ($<$1) are considered activated. (b) MH-MoE showcases finer-grained understanding by distributing sub-tokens split from semantically-rich patches to more distinct experts to capture semantic information. Brighter regions indicate that sub-tokens from this patch are distributed to a greater number of diverse experts, while darker regions indicate that sub-tokens are assigned to more of the same experts.
  • Figure 2: Workflow for MH-MoE on both vision and language data. For vision data, different heads routed to different experts try to capture different aspects of details within patches and relations between patches. For language data, different heads attend to capture the varying contexts of false cognates across different languages (e.g., Italian and English) or polysemous words within a single language.
  • Figure 3: Illustration of a typical SMoE layer and the proposed MH-MoE layer. (a) An SMoE layer consists of a router and expert networks, where the experts are sparsely activated according to dot-product token-expert routing scores. (b) MH-MoE introduces additional two MLP layers, namely the multi-head layer and merge layer, and a Token-Splitting-Merging (TSM, Eq. \ref{['EQ. sub-tokens']} and Eq. \ref{['EQ: merge tokens']}) operation incorporated between these two MLPs.
  • Figure 4: Perplexity on validation dataset during the training phase reported for Dense, X-MoE and MH-MoE across three pre-training tasks. (a) English-focused language modeling. (b) Multi-lingual language modeling. (c) Masked multi-modal modeling
  • Figure 5: Distribution of expert activation in X-MoE and MH-MoE on both Harnesseval-harness and XNLIxnli corpus, encompassing $6$ SMoE layers with $32$ experts per layer. The top of the heatmap is the first SMoE layer while the bottom is the last. Experts activation ratio is determined by calculating the ratio of each expert's selection frequency in each MoE layer to the total number of tokens.
  • ...and 6 more figures