Multi-Head Mixture-of-Experts

Authors: xun  wu
shaohan  huang
wenhui  wang
shuming  ma
li  dong
furu  wei
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-03 07:34:23

Abstract: Sparse Mixtures of Experts (SMoE) scales model capacity without significant increases in computational costs. However, it exhibits the low expert activation issue, i.e., only a small subset of experts are activated for optimization, leading to suboptimal performance and limiting its effectiveness in learning a larger number of experts in complex tasks. In this paper, we propose Multi-Head Mixture-of-Experts (MH-MoE). MH-MoE split each input token into multiple sub-tokens, then these sub-tokens are assigned to and processed by a diverse set of experts in parallel, and seamlessly reintegrated into the original token form. The above operations enables MH-MoE to significantly enhance expert activation while collectively attend to information from various representation spaces within different experts to deepen context understanding. Besides, it's worth noting that our MH-MoE is straightforward to implement and decouples from other SMoE frameworks, making it easy to integrate with these frameworks for enhanced performance. Extensive experimental results across different parameter scales (300M to 7B) and three pre-training tasks—English-focused language modeling, multi-lingual language modeling and masked multi-modality modeling—along with multiple downstream validation tasks, demonstrate the effectiveness of MH-MoE.

Main paper: 
 introduction: 
 Large capacity models, such as Large Language Models (LLMs) [39, 28, 6, 25] and Large Multimodal Models (LMMs) [37, 27] , have demonstrated their efficacy across various domains and tasks. To further enhance performance, a reliable approach involves scaling up these models by augmenting the parameter count [13] . But for most of these densely-activated large-capacity models (referred to as Dense models), which utilize all their parameters to process all inputs, the extremely large size of these models significantly reduces inference speed, further limiting their practicality.A promising alternative, facilitating model scalability while mitigating the burdensome computational costs, resides in Sparse Mixtures of Experts (SMoE) [31, 12, 5, 7] . In contrast to Dense model, SMoE contains parallel feed-forward neural networks (referred to as experts) within each building block, and strategically activates distinct experts for specific input tokens via a router, thereby yielding noteworthy efficiency enhancements. For instance, GShard [21] scales a Dense model from 2B to 600B parameters with lower training costs than a 100B Dense model. And recently, Mixtral 8×7B [16] , a SMoE model containing 8 experts is shown to outperform or matches LLaMA-2 70B [34] and GPT-3.5.Despite its success, SMoE exhibits the low experts activation issue, which means that only a small subset of experts are activated during optimization and inference, e.g., 8.33% activation ratio 1 shown in Figure 1 (a), while the majority of them are not used at all (see the dark area). As a result, SMoE fails to utilize the full expressive power of these experts, especially when the number of experts is large, which significantly limits effectiveness and scalability of SMoE. Our aim in this paper is to achieve denser expert activation (i.e., better utilization of "dead" experts) without increase in computational cost. To achieve this, we propose Multi-Head Mixture-of-Experts (MH-MoE). The workflow of MH-MoE is illustrated in Figure 2 . Inspired by the multi-head mechanism utilized in Multi-Head Self-Attention (MHSA) block, MH-MoE splits each input token into multiple sub-tokens and distribute them to different experts. After expert processing, subtokens are seamlessly reintegrated into the original token form, thereby achieving denser expert activation, e.g., 90.71% activation in Figure 1 (a), while also circumventing any additional computational burden in subsequent non-parallel layers, e.g., MHSA block. Specifically, as shown in Figure 2 , when provided with a single input token, MH-MoE activates four experts by splitting it into four sub-tokens, whereas SMoE only activates a single expert.Furthermore, we observe an interesting phenomenon: for tokens with richer semantic information, the sub-tokens split from these tokens are more likely to be allocated to distinct experts. For example, refer to the bright area in Figure 1 (b), where sub-tokens split from these patches are allocated to a greater number of different experts, facilitating the capture of semantically-rich information (e.g., the eagle in the figure). We therefore speculate that the allocation of sub-tokens to distinct experts enables MH-MoE to simultaneously focus on information from various representation spaces within different experts, ensuring a more granular understanding for subtle differences in both vision and language patterns, finally achieving better finer-grained understanding ability. See in Figure 2 , sub-tokens assigned to Experts 3 and 2 capture a detailed understanding of each character's actions within an image patch, while those assigned to Experts 1 and 4 explicitly model the semantics of the false cognate 'camera'.MH-MoE maintains following strengths: (1) Higher experts activation & better scalability. MH-MoE can alleviate lower expert activation problem and significantly enhance the usage of larger experts by enabling optimization of almost all of experts, allowing for more efficient scaling of model capacity.(2) Finer-grained understanding ability. By adaptively assigning sub-tokens to different experts based on the semantic richness of the token, MH-MoE enabling to jointly attend to information from different representation spaces at different experts, and finally achieving better finer-grained understanding ability. (3) Seamless integration. The implementation of MH-MoE is remarkably straightforward and decoupled from other SMoE optimization methods (e.g., GShard [21] ), making it very easy to integrate them together to achieve better performance.We evaluate the proposed MH-MoE on three model pre-training and fine-tuning setting: Englishfocused language modeling, multi-lingual language modeling and masked multi-modality modeling, across different parameter scales (300M to 7B). Extensive experimental among these three tasks demonstrate the effectiveness of MH-MoE.English： during the wildlife safari, I used my high-resolution camera to capture stunning photographs of lions, elephants, and giraffes in their natural habitat, creating unforgettable memories of the adventure.  
background: 
 Sparse Mixture-of-Experts (SMoE) [31, 12, 5, 7] enhances model capacity while maintaining a constant computational demand, thus achieving better performance than densely-activated models on various tasks [22, 19, 39, 28] .Different from densely-activated models, each MoE layer consists of N independent Feed-Forward Networks (FFN) {f FFN i } N i=0 as the experts, along with a gating function g (•) to model a probability distribution indicating the weights over these experts' outputs. For the hidden representation h ∈ R d of each input token, the gating value of routing h to expert f FFN i is denoted as:EQUATIONwhere e i denotes the trainable embedding of the i-th expert and As described above, the most commonly used routing mechanism involves selecting the top-k experts from N experts, where k ≪ N [32] , e.g., k = 2 and N = 2048 in GShard [21] . Such a routing mechanism allows the combination of data parallelism and expert parallelism. Some works [38, 21] suggest that larger values of k often contribute to better model performance. However, with the increase in the value of k, training models with conventional top-k routing implementation becomes much less efficient [21] . In this paper, we introduce MH-MoE, a simple but efficient manner to make denser expert activation without an increase in computational complexity. 
multi-head mixture-of-experts: 
 Concretely, we denote a sequence of inputs tokens by X ∈ R l×d , where l is the number of tokens and d represents the length of token dimension. In MH-MoE, each parallel layer contains a set of N experts, each presented as {f FFN i : Rd h → R d h } N i=0, where h denotes the number of heads (i.e., the number of sub-tokens a single token is split into), which is decoupled from the head in the multi-head self-attention layer. For clarity, we describe the operation of a single MH-MoE layer here only.The full architecture of MH-MoE can be seen in Figure 3 . First, X is projected by a multi-head layer with parameter matricesW head ∈ R d×d , X = X • W ⊤ head (2)where X ∈ R l×d . After that, every token in X is split into h sub-tokens along the token dimensions, and these sub-tokens are arranged in parallel according to the original token sequence, forming a new feature space Ẍ ∈ R (l×h)× d h as2 :EQUATIONwhere function 𭟋 s denotes the token splitting operation: R l×d → R (l×h)× d h , and each sub-token is presented asx i j ∈ R d h, meaning it is the the j th sub-token split from the i th token. Then all these sub-tokens are fed into the gating function g(•). The gating value of routing a certain sub-token x i j into the p th expert is computed asEQUATIONwhere e p ∈ R d h is the learnable embedding of the p th expert. In this paper, we mainly focus on top-k routing, i.e., only the experts with the largest top-k routing score is activated. Φ = Top k g f FFN denote the set of activated experts and |Φ| = k. Then x i j is processed by these activated experts as following,EQUATIONAfter that, all obtained o i j are rearranged in the original order of sub-tokens and concatenated together as 2 :EQUATIONwhere O ∈ R (l×h)× d h . After that, O is transformed back the into original token form by a token merging operation 𭟋 m : R (l×h)× d h → R l×d and then projected by a merge layer with parameter matrices W merge ∈ R d×d to effective integration of multiple features o i j capturing detailed information from different expert representation spaces:EQUATIONThen we get the final output X of the single MH-MoE layer.By implementing the aforementioned operations, we effectively increase the average volume of data routed to a specific expert by a factor of h (as demonstrated in Eq. 3), thus achieving denser expert activation. Besides, the shapes of the input and output in the MH-MoE layer remain unchanged, thus no additional computational cost is introduced in the subsequent block. Specifically, we introduce a hyperparameter β to scale the inner dimensions of each expert, aiming to balance the parameters introduced by the multi-head layer and merge layer, aligning the model's parameters and computational complexity with the original SMoE. Furthermore, the allocation of sub-tokens to distinct experts within MH-MoE enables us to collectively capture semantic information from diverse feature spaces across these experts, thereby enhancing the model's ability to achieve a finer-grained understanding.As the Pytorch-like style pseudocode of MH-MoE shown in Appendix D, MH-MoE is characterized by its overall simplicity of implementation and decoupled from other SMoE optimization strategies [21, 5] , making it easy to integrate with other optimized SMoE frameworks to enhance performance.Load balancing loss. To mitigate the expert load imbalance issue, given the sub-token set Ẍ (depicted in Eq. 3) and the frequency t p of how many sub-tokens are routed to the p th expert, we follow existing works [21, 13] to compute the load balancing loss L balance via:EQUATIONwhere N denotes the number of experts, | Ẍ| is the number of sub-tokens contained in Ẍ. g f FFN p denotes the gating value of routing a certain sub-token x i j into the p th expert (see in Eq. 4). Task specific loss. The term L task is dependent on the particular task that MH-MoE is designed to learn. For instance, during pre-training in the English-focused Language Modeling task, we utilize the language modeling loss [29] , whereas the model predicts the next word in a sequence. So, the overall training objective of MH-MoE is to minimize:EQUATIONwhere α is a coefficient for load balancing. 
experimental setup: 
 Compared Baselines. We include two baseline models for comparison: (1) Dense, which represents a Transformer decoder without the incorporation of sparsely-activated parallel modules (i.e., SMoE layer). (2) X-MoE, which is our implementation based on the SMoE proposed by [5] . We build MH-MoE upon X-MoE and use identical settings. Note that the all models are pre-trained using the same training data and loss (Eq. 9) as MH-MoE, and we ensure that the parameter count of MH-MoE remains consistent with or lower than that of X-MoE, ensuring a fair and equitable comparison. A detailed comparison about parameter and computational complexity can be found in Table 11 .Pre-training Data. We detail the pre-training data of MH-MoE in three areas: (1) English-focused experiments use the RedPajama dataset [8] , which is an open-source pre-training dataset. (2) multilingual tasks follow XLM [20] and use the multilingual Wikipedia as training data.(3) multimodal tasks use a masked multi-modality modeling task with a large dataset of images, documents, and image-text pairs. Further details are available in the Appendix A.Model Architecture and Hyperparameters. For all experiments, we use the X-MoE [5] as our backbone architecture to build our MH-MoE, which has shown better performance than prior SMoE models such as Switch Transformers [13] on cross-lingual understanding benchmarks. For Englishfocused Language Modeling and Multi-lingual Language Modeling, we construct Dense, X-MoE and MH-MoE using the Transformer [36] decoder (L = 12, H = 768, A = 12) with the GPT-43 vocabulary Table 2 : Accuracy / accuracy-normalization scores for language understanding tasks using the LLM Evaluation Harness [14] . N denotes the number of experts. We examined the validation perplexity curves for all pretrained models and pre-training tasks under two expert settings (8 experts and 32 experts). The perplexity trends are depicted in Figure 4 , with the final perplexity values listed in Table 1 . We can observe that as training progresses: 1) the perplexity of our MH-MoE remained lower in comparison to the compared baselines, indicating more effective learning; 2) MH-MoE achieved the lowest perplexity across three distinct experimental setups; 3) an increase in the number of experts led to a corresponding decrease in the perplexity of MH-MoE, suggesting that the model benefits from enhanced representation learning capabilities as more experts are incorporated. These results collectively demonstrate the superiority of MH-MoE in terms of learning efficiency and language representation across multiple pre-training paradigms.English-focused Language Modeling. We evaluate our models on a total of 9 different zeroshot benchmarks to assess their abilities across various natural language tasks like common sense reasoning, general language understanding and knowledge understanding using the LLM Evaluation Harness [14] . As shown in Table 2 , comparing X-MoE with the Dense model, X-MoE show notable improvement, indicating that SMoE models (e.g., X-MoE) benefit from the large model capacity.Overall, for all benchmarks, our MH-MoE obtains the best performance, achieving an average performance gain of 1.1 for 8 experts setting and 1.5 for 32 experts setting compared to X-MoE, demonstrating the effectiveness of our proposed MH-MoE on modeling English-focused language.Multi-lingual Language Modeling. We evaluate our multi-lingual language models on the crosslingual natural language inference (XNLI) corpus [10] , which is the extension of the multi-genre NLI Masked Multi-modal Modeling. We evaluate on the widely used vision-language understanding and generation benchmarks, including visual question answering [15] , visual reasoning [33] and image captioning [23] . We report vqa-score on VQAv2, accuracy for NLVR2. For COCO image captioning, we report BLEU@4 (B@4), METEOR (M), CIDEr (C), and SPICE (S). Table 4 presents the evaluation results. For VQA task, MH-MoE outperforms both Dense and X-MoE by a large margin, e.g., 4.24 and 1.69 points gain on test-dev split, respectively. For visual reasoning task, MH-MoE beats all these two baselines on both dev (1.5 points gain than X-MoE) and test-P split (1.7 points gain than X-MoE). For image captioning task, MH-MoE surpasses X-MoE by 4.2%, 10.2%, 9.4% in terms of B@4, M and S, respectively. Above results show that X-MoE exhibits enhanced visual information comprehension, which also validates the effectiveness of our proposed MH-MoE in capturing diverse semantic information within text-image pair data.This section presents experimental analysis to demonstrate the functionality of MH-MoE. In all comparative experiments, we ensure parameter equality across models by adjusting the internal dimensions of the experts.Number of Heads h. We conduct experiments by adjusting the number of heads (h = 2, 4, 6, 8, and 12) in MH-MoE. As shown in Figure 5 , we find that across all settings of h, our model consistently outperforms the X-MoE. Besides, as the value of h increases, we observe an initial improvement followed by a decline in our model's performance. This leads us to hypothesize that as h initially increases, the enhancement in performance benefits from MH-MoE by activating a greater number of experts, thereby enhancing the model's effectiveness and capturing a wider range of fine-grained token information. However, as h becomes too large, the excessive subdivision of tokens may impair their original semantic content, leading to decreased model performance.The key motivation of MH-MoE is to split each token into several sub-tokens and then merge the sub-tokens after processing by experts. We use two MLP layers for the splitting and merging processes. We conduct a detailed analysis to determine whether the performance improvement is due to the Token-Split-Merging (TSM) operation or the additional MLP layers. In contrast, a significant enhancement in model performance is only achieved when both MLP and TSM are incorporated simultaneously. We hypothesize that introduction of TSM, without the integration of MLP, activates more experts, but the segmentation and merging of the model appears overly straightforward and abrupt in its execution. This limitation hinders the model's ability to meaningfully segment tokens into sub-tokens and effectively merge the diverse information gathered from different expert spaces.We also conducted additional ablation studies to examine the impact of varying the number of MLP layers, different splitting-merging methods and different balance loss. The details of these experiments are provided in Appendix E. 
analysis: 
 Experts Activation. We visualize the activation of each expert varies across parallel expert layers for X-MoE and MH-MoE at Figure 6 . It can be observed that: 1) X-MoE demonstrate a more skewed distribution, wherein a significant portion of experts remain inactivated all the time. 2) Our MH-MoE achieves a denser expert activation compared to X-MoE, effectively mitigating the issue of low expert utilization. 3) As the number of heads h increases, the expert activation frequency in MH-MoE also rises.Scalability. We explore the scalability for both X-MoE and MH-MoE by scaling up the number of experts from 8 to 256 (about 7B parameters). For upstream performance, as shown in Figure 7 (a), with the increase of experts, our MH-MoE could bring more gains. It is because MH-MoE could mitigate the low expert activation problem effectively. With this ability, the superiority of the large-scale SMoE model will be better exerted, thereby achieving the improvement of the upper bound Brighter regions indicate sub-tokens split from this patch are distributed to a greater number of diverse experts. of SMoE with more experts. Detailed validation perplexity curves for these scaling up experiments can be found in Figure 10 at Appendix F. For downstream performance shown in Figure 7 (b), for X-MoE, expert number = 64 is the upper bound, meaning that continuing to increase the number of experts will not bring performance gain. Our MH-MoE not only has a performance advantage over the X-MoE with the same number of experts, but also improves the upper bound from 64 to 256, which demonstrates the effectiveness of the scalability of our MH-MoE on downstream tasks.Experts Assign within Token. We delve into a more granular analysis to validate how the multihead mechanism aids MH-MoE in capturing diverse and intricate semantic information that is often challenging to comprehend, e.g., polysemous and false cognates words (denoted as PF tokens) in languages, and semantically-rich areas in images.For languages data, we compute and compare the divergence levels (i.e., the number of different experts these sub-tokens are routed to) of sub-tokens split from PF tokens and Non-PF tokens. The results, presented in Figure 7 (c), clearly demonstrate that the distribution of divergence for PF tokens is significantly skewed towards the right when compared to that of Non-PF tokens. This indicates that during the MH-MoE's inference process, PF tokens route their sub-tokens to a greater number of different experts, thereby capturing diverse semantic information in contrast to Non-PF tokens for a better polysemous and false cognates word modeling. Note that we utilized the GPT-4 API [25] to extract polysemous words and false cognates from the XNLI [10] corpus, and the corresponding prompt can be found in Table 12 .For image data, we analyzed how the divergence levels of different patches evolve during the training process, as illustrated in Figure 7 . Notably, as training progresses, divergence levels gradually increase in high-frequency texture regions (or regions with rich semantics), while they decrease in low-frequency texture regions. This suggests that MH-MoE tends to route tokens from complex texture areas to a wider variety of experts, thereby enhancing the finer-grained understanding of semantics in those regions. For more visualization examples, please refer to Figure 11 in Appendix G.Experiments on Pure Vision Tasks. We conduct small-scale experiments on pure vision tasks to further validate the effectiveness of MH-MoE. We utilize ViT-B(ase) [11] as the Dense model and AdaMV-MoE [4] as the SMoE baseline. AdaMV-MoE is a multi-task vision SMoE model demonstrating a superior performance across various vision tasks. We build our MH-MoE upon AdaMV-MoE. The comparison results are shown in Table 6 , which demonstrates that MH-MoE exhibits corresponding effectiveness and versatility in pure vision tasks.Model Collapses. Concerns may arise that MH-MoE could revert to the original SMoE approach, routing all sub-tokens from the same token to a single expert. In MH-MoE's training framework, the load balancing loss L balance (Eq 8) serves as a weak constraint, treating all sub-tokens as independent entities and uniformly distributing them among experts, promoting a balanced allocation. Our observations indicate that most sub-tokens from the same token are distributed among 3-5 different experts (see Figures 7 and Figures 7 (c) ). 
conclusion: 
 In this paper, we study how we can to achieve a denser experts activation without introducing additional cost, while improving the fine-grained understanding ability. With the proposed Multi-Head Mixture-of-Experts, we can easily implement the aforementioned functionality. Furthermore, the simplicity of MH-MoE allows it to integrate with other SMoE frameworks to enhance performance easily. Extensive empirical results across three tasks demonstrate the effectiveness of MH-MoE. 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA] Justification: Our model serves as the foundational architecture for large models rather than their socio-economic applications. As such, it operates on a different level, and we do not provide an analysis of societal impact.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.• Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [No] Justification: We will conduct this before our model is publicly released. Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. 12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected? Answer: [Yes] Justification: Please refer to Section 4. Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.• For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: Table 8 presents of pre-training data in Masked multi-modal modeling task. For multi-modal data, there are about 15M images and 21M image-text pairs collected from five public datasets: Conceptual 12M (CC12M) [3] , Conceptual Captions (CC3M) [30] , SBU Captions (SBU) [26] , COCO [23] and Visual Genome (VG) [17] . For monomodal data, we use 14M images from ImageNet-21K and 160GB text corpora [1] from English Wikipedia, BookCorpus [40] , OpenWebText4 , CC-News [24] , and Stories [35] . Table 9 presents the model hyperparameters of X-MoE and MH-MoE for both English-focused language modeling and Multi-lingual language modeling tasks. The gating temperature τ 0 is initialized as 0.3 and 0.07 for the softmax gating and sigmoid gating, respectively. We use the same vocabulary as XLM-R [9] with 250K subwords tokenized by SentencePiece [18] .Table 9 : Model hyperparameters of Dense, X-MoE and MH-MoE. The SMoE frequency refers to how many experts each token will be assigned to, i.e., the value of k in the Top-expert selection. Table 10 presents the hyperparameters for pre-training across three tasks: Language modeling tasks (English-focused language modeling and Multi-lingual language modeling tasks) and Masked multi-modal modeling task.Table 10 : Pre-training hyperparameters for Language modeling tasks (English-focused language modeling and Multi-lingual language modeling tasks) and Masked multi-modal modeling task tasks. We also provide the PyTorch-style code in Algorithm 1 to explain our MH-MoE, which including two main aspects: 1) Stage 1. The creation and initialization of multi-head layer and merge layer.2) Stage 2. The segmentation of tokens, followed by processing through an expert network, and ultimately merging. Your role is to identify polysemous and false cognates in different languages from the given textual input (### Input). Note that "Polysemous" refers to a word having two or more completely different meanings (for example, "grouse" has meanings related to complaining and also refers to a type of bird), while "false cognates in different languages" refers to words in different languages that have similar forms but carry different meanings (for example, in English and Italian, "camera" looks similar but represents different semantic concepts). AlgorithmNumber of MLP layers. We explore the impact of varying the number of layers (n = 0, 1, 2, 3) in MLP on MH-MoE performance. For configurations exceeding a single layer, ReLU activation functions were incorporated between MLP layers to ensure the non-linearity of transformations. The parameter quantities of the models being compared are equal. Upon analyzing the results in Table 13 , we observe that increasing the number of MLP layers beyond one had a negligible impact on the model's performance. This indicates that a single-layer MLP is sufficient for accomplishing token segmentation and fusion.What Makes a Good Splitting-Merging Method? We designed ablation experiments to investigate whether the tokenization method affects the MH-MoE's performance. We replace the multi-head layer and merge layer (both are FC layers) with Conv1d layers (denoted as MH-MoE ♯ ).We conducted experiments in two settings (English-focused language modeling and Multi-lingual language modeling tasks) and explored different levels of randomness by using different kernel sizes and strides for the Conv1d layers. We ensure that the parameters are kept as consistent as possible by adjusting the number of Conv1d layers and the dimensions of experts. Through experimental results presented at Figure 8 , we observe that replacing splitting method with Conv1d (denoted as MH-MoE ♯ ) resulted in significantly inferior performance compared to the original MH-MoE. This underscores the importance of tokenization methods, indicating that tokenization methods such as Conv1d, which disrupt the original input sequence features, are not suitable for this context.We conducted experiments with three groups to assess the impact of different constraints on sub-token allocation: (1) MH-MoE ¶ : MH-MoE trained without the L balance , (2) MH-MoE † : MH-MoE trained with the L balance , and (3) MH-MoE ‡ : MH-MoE trained with both L balance and diversity loss L div . L div is designed to enhance sub-token allocation diversity and defined as: given a set of sub-tokens Υ split from the same token, with P i representing the gating probability distribution of the i th sub-token, The experimental results (perplexity) are summarized in the Figure 9 . Key findings include: (1) Without any constraints (MH-MoE ¶ ), the model's performance degrades, likely due to severe imbalance, causing some experts to receive insufficient data, resulting in suboptimal performance.L div = - |Υ| i=1 |Υ| j=i+1 KL(P i ||P j ), where(2) The addition of L div did not significantly improve model performance compared to using only Lbalance. We hypothesize that L balance alone provides sufficient diversity in sub-token allocation.We provide the training perplexity curve for model training in the experimental setting of increasing the number of experts (from 8 to 256) in Figure 10 .We provide more visualization of variation in assign diversity for sub-tokens split from different patches in vision data at Figure 11 . Justification:Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).Experts activation ratio shown in Figre 1 (a) is 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.Sub-tokens within regions marked with the same color are split from the same token.https://github.com/openai/tiktokenhttp://skylion007.github.io/OpenWebTextCorpus 


Decision: This paper presents a mixture-of-experts approach that routes different sub-tokens to different experts. Reviewers found the idea to be simple and effective, the results to be strong compared to existing baselines, the paper clear, and the experiments comprehensive. Concerns regarding missing experiments, unsupported claims and clarity issues have largely been addressed. Some concerns about XXL-scale experiments seem out of scope in my view, as there is plenty of value in models much smaller than 10T parameters. Other requests for additional experiments and baselines seem not essential either.

Reviews: review: paper_summary: paper_summary: The paper proposed a new architectural extension to large pre-trained models (e.g. LLMs, or Multi-modal models). Specifically, the focus is on Mixture of Experts (MoE) models and the paper proposes a new method by introducing subtokens and routing each subtoken to experts.

The resulting architectural change performs better on down-stream tasks and demonstrates a higher use of a variety of experts.

The author's also conduct a variety of ablation studies to test various aspects and changes of their proposed architectural change.
strength_weakness: strength_weakness: Weaknesses:
1. No actual comparison in terms of "cost" is given except for parameter count and complexity calculations in the appendix. It would be interesting to see actual evaluation of cost (e.g. computational cost of training & computational cost of inference vs. MoE alone).
2. Evaluation and comparison is not done on publicly available and popular (& stat-of-the-art) models of relevant sizes (e.g. Mistral, Aya, Gemma ...). Therefore, it is not clear if this method actually performs better or just on the models and training runs presented in this paper.
3. Clarity
a. some parts are unclear, e.g. equations 2,3,6,7 speak about the new additions of the method, however, it is not clear at all that there is MLP layer present (only later in the paper it is mentioned), however, it is "still" not clear whether and how it applied during equations 2,3,6,7.
b. some mistakes in the paper e.g. Paper checklist 7, the answer is N/A. However, N/A means no experiments were run, the correct answer should be No.)
4. The core claim that the problem with MoE is: "low experts activation issue" has a few problems:
a. It is never clearly explained why this is an issue (and no evidence is given that this is a real issue).
b. On one hand the paper presents that the method increases activations, however, the paper also shows that increasing the number of heads (and therefore activations) seems to also decrease performance after a specific number of heads [Figure 5, Section 4.4]. (Great for reporting this, however, this also challenges the introduction / proposition of the method). I.e. it seems that "low experts activation" is not always a problem, or only a problem to a certain degree. => Therefore, it would be interesting to revise this hypothesis, and also to provide a counter hypothesis (with evidence).
5. Reporting on sharding of the models on the DGX-2 system is not quite described. (Which is quite essential to run it sensibly.)
questions: questions: Questions:
1. How does your model compare against some popular & state-of-the-art models (e.g. MoE LLMs like Mistral) in terms of performance?
2. Referring to "Weakness 4.b", what is your explanation of why your method actually works better? 
3. What is the actual cost of running standard MoE vs. yours (in terms of practical training or inference time) - given the same computation budget.
limitations: limitations: Limitations:
1. Limitations are not really considered in the paper.
2. It would be really interesting to understand what does NOT work in the method in practice?
a.  What did not work during developing of the method?
b. When is the method worse than previous other methods?
c. Computational limitations?
d. etc. etc.
paper_summary: paper_summary: This paper introduces Multi-Head Mixture-of-Experts (MH-MoE), a method for training MoE models with enhanced expert activation. In particular, MH-MoE splits each input token into multiple sub-tokens, which are processed by a set of experts in parallel, and then merges them back into their original token form. The results show improved performance over dense and X-MoE models in parameter-matching settings.
strength_weakness: strength_weakness: 1. Experiments are performed on a limited number of models and baselines, making the results less convincing.

- The experiments are limited by a single model architecture. The authors apply their approach by modifying the previously released X-MoE model.

- The authors compare their performance with the mentioned X-MoE model, whereas there are a variety of MoE approaches that have been proposed and released recently.

- The ablation studies are limited. The performance of SMoE heavily depends on the choice of hyper-parameters, one of the most important being the number of experts to be activated (referred to as top-k). How does the activation map change with an increase in k? Would MH-MoE still perform better, or is there a saturation point? In other words, does MH-MoE seem superior to X-MoE only because a small k=2 was chosen?

2. There is a lack of discussion on limitations. The authors claim that their approach does not have any limitations. I would urge the authors to reflect on the guidelines and reconsider their answer. For example, the weaknesses mentioned above (but not limited to) should be considered.
questions: questions: Please see the reasons to reject.
limitations: limitations: No. authors did not disclose any limitations of their method. See "Weaknesses" section for more details.
paper_summary: paper_summary: This paper proposes Multi-Head Mixture-of-Experts (MH-MoE), a simple yet effective routing strategy that splits each input token into multiple sub-tokens for expert routing. This operation significantly enhances the ratio of activated experts for each token, enabling more fine-grained assigning of tokens. Through extensive experimental results across different parameter scales (300M to 7B) and three 15 pre-training tasks, the paper validates the effectiveness of the proposed MH-MoE strategy.
strength_weakness: strength_weakness: 1. My major concern is about the experimental settings. Implementing MH-MoE based on the X-MoE instead of the vanilla SMoE is strange. The improvements shown by the experiments may not be generalizable to other MoE structures, which makes the proposed method less convincing.
2. Lack of experiments on scaling computations (activated experts). Though the paper conducts scaling experiments on scaling total parameters (number of experts), the computation remains static. The information conveyed by this kind of scaling is limited as the computation is also important in affecting the performance of MoE. I’m not expecting additional experiments as the time is limited during rebuttal. However, a computational scaling experiment would make the method more convincing.
questions: questions: Can you provide the balance loss curves for your method and the baselines? I’m curious about the effect of your method on load balancing as it introduces more sub-tokens for routing.
limitations: limitations: None
paper_summary: paper_summary: The work presents an advanced MoE model architecture to further improve the expert utilization. Authors conduct comprehensive experiments on three pre-training tasks, English-focused language modeling, multi-lingual language modeling and masked multi-modality modeling with models at different scales (300M to 7B).
strength_weakness: strength_weakness: I only have one but important concern: the end-to-end training and inference throughput. I understand that the theoretical computation and communication cost is not high, but since we are conducting a more complicated/fine-grained routing decision, the routing process may take longer in practice. Considering that MoE models are sometimes trained and deployed at very large scale, the intra-node all2all is sometimes required. Will the proposed approach slow down this?

I can increase my score to positive if this concern can be solved.
questions: questions: See weakness.
limitations: limitations: NA

Comments: title: Response to Reviewer Y9UT
Kindly Reminder by Submission1607 Authors
Thank you for your response
Thanks
Response to Reviewer Rrmq
comment: Dear  Reviewer Y9UT,

Thank you for your response and reviewing our work.

---

>Q1: Baselining against SOTA models (SMEAR [1], BTX [2], DeepSeekMoE [3], DEMix [4], etc).

We understand that comparing our work against SOTA MoE models to better understand where SOTA currently stands could indeed enrich our paper. However, the experiments conducted on MH-MoE across models ranging from 300M to 7B parameters were aimed at demonstrating that our proposed MH-MoE can improve the performance of MoE models at various parameter scales, especially when there are many experts, as shown in Figure 7. **This does not imply that our experimental results can be directly compared with the current SOTA models of similar scale (e.g., 7B)**, as these SOTA MoE models differ significantly from our model in terms of training data, training steps, and **especially the scale of activated experts** (e.g., some SOTA models activate 2B parameters, whereas our MH-MoE activates only a few hundred million). Therefore, introducing these results for comparison may not provide meaningful insights.

Furthermore, we want to reiterate that our experiments were conducted under two different MoE frameworks, **with a total of four pretraining tasks and 29 downstream tasks** to validate the effectiveness of MH-MoE. This experimental setup is more extensive than some previous works [1-3], and we believe that our results sufficiently demonstrate the effectiveness of MH-MoE.

**Reference**

[1] Chi, Zewen, et al. "On the representation collapse of sparse mixture of experts." NeurIPS 2022.

[2] Chen, Tianlong, et al. "Adamv-moe: Adaptive multi-task vision mixture-of-experts." ICCV 2023.

[3] Soft Merging of Experts with Adaptive Routing

---

>Q2: Compare with architectures that aimed to solve the same issue.

Thank you for your suggestion. As far as we know, we are the first to identify the issue of low expert activation ratio in MoE models, so we are currently unable to introduce comparisons with similar models in the paper. However, we look forward to more research focusing on this issue in the future, as it is a critical challenge, especially when aiming to increase the number of experts in MoE models significantly. Low expert activation could become a bottleneck limiting the further improvement of MoE models, and we would be very interested in comparing our work with future studies addressing this issue.

---

>Q3: Initial results shared are promising, and I encourage authors to continue investigations.

Thank you for your positive feedback. We will continue our experiments and update the latest results in the upcoming version of the paper.

---
**If you have any further questions or concerns, please feel free to contact us at any time. We are always available and look forward to further discussions with you. :)**

Best regards,

All Authors
Thank the authors for providing more results, I will keep my relatively positive rating.
Dear Reviewer 84mb,

Thank you again for reviewing our work. As the discussion period ends shortly, we wanted to check if you have any further questions or found our responses helpful? We are more than willing to extend our conversation and eagerly anticipate any further discussions that may arise. Please let us know and thanks for your time :).

Best regards,

All Authors
-> A1:
Some recent SOTA MoE models include SMEAR [1], BTX [2], DeepSeekMoE [3], DEMix [4], etc
There were two points raise in the question: 1) building on top of other models; 2) baselining against SOTA models. 
I understand that it is impossible to compare performance with all the variety of MoE models that where developed over the last two years (since 2022 for XMoE and 2023 for AdaMV), but it is always beneficial to select those most recent models of similar architectures that 1) were evaluated on the same datasets used in the paper to understand where SOTA is; 2) architectures that aimed to solve same issue (activation sparsity).
Comparison against two selected baselines creates false premise that activation problem has not been challenged, and doesn't give enough information about where the proposed method lies against SOTA results (even though it is not necessary that the proposed method must improve them).

1. Soft Merging of Experts with Adaptive Routing
2. Branch-Train-MiX: Mixing Expert LLMs into a Mixture-of-Experts LLM
3. Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models
4. DEMix Layers: Disentangling Domains for Modular Language Modeling

-> A2
I understand that in the limited time of rebuttal, it is impossible to perform convincing ablation studies to answer the raised questions. Initial results shared are promising, and I encourage authors to continue investigations.
Thank you very much for your positive feedback and your willingness to raise the score and confidence. We will ensure that the revised version of our paper incorporates all the suggested improvements to enhance its clarity and contribution :).
Thank you for your feedback! Achieving a model with 1T or even more parameters is indeed a significant milestone.

There are a few points we would like to clarify:

1. Increasing the number of experts does indeed increase the cost of communication. However, in our work, the MH-MoE has the same number of experts as the baseline SMoE. Therefore, when scaling to models with a large number of experts, MH-MoE faces the same communication challenges as SMoE.

2. The slower speed of static routing in MH-MoE is due to the need to create matrices of size $N(S^2)$ to store the routing status, where $S$ is the number of tokens needing routing. (Refer to [this code](https://github.com/facebookresearch/fairseq/blob/moe/fairseq/modules/moe/top2gate.py#L193): `combine1_sec` is the matrix, $s$ is the number of tokens, $e$ is the number of experts, and \(c\) is the capacity where $c = \frac{2s}{e}$, so the matrix size is $2e^2$. For MH-MoE, the number of tokens requiring routing is proportional to the number of heads, which affects the speed. In the implementation of MegaBlock, scatter and gather operations are optimized through kernel fusion (Refer to [this code](https://github.com/databricks/megablocks/blob/main/megablocks/layers/moe.py#L185)), eliminating the need to create this matrix and reducing the required memory to $N(S \times D)$, where $S$ is the number of tokens needing routing and $D$ is the dimension of the hidden state. Although the number of tokens increases to $S \times H$, where $H$ is the number of head, the token dimension decreases to $\frac{D}{H}$, so the overall memory requirement does not increase significantly. Regarding the amount of data transferred in communication, MH-MoE and SMoE are the same, both requiring $N(S \times D)$.
Dear Authors thank you for taking the time and effort to do those additional experiments.

Things are much clearer now and the score is now 7 (from 6) with confidence 4 (from 3).

We would highly encourage to add all these changes into the paper (incl. appendix) and to make certain things clearer. As such this paper becomes a good contribution.
Thank you for the explanation. I will maintain my score.
Megablock can alleviate the problem as the expert is small. Imagine that you are training a model with 1000 experts and 10T parameters, or even slightly smaller (maybe 1T). You finally need to do expert parallel on the slow network (inter-node communication). So the problem is still there.
I personally believe you should always think about a very large and sparse model when you want to push the boundary of MoE research further.
Thanks for your detailed feedback!

----

>**Q1**: Would dynamic routing improve the Tok/s for the baseline?

**A1**: In Table A, we compare the speed of the baseline and MH-MoE using static routing and dynamic routing (marked with #). We observed the following:

1. Dynamic routing also increases the tokens per second (Tok/s) for the baseline.
2. The improvement dynamic routing brings to MH-MoE is significantly greater (6.16 toks/s) than the improvement it brings to the baseline (1.77 toks/s).
3. Although dynamic routing increases the Tok/s for the baseline, the improved baseline speed is not much faster than MH-MoE#.

The above results indicate that dynamic routing is an effective way to address the throughput limitations of MH-MoE.

***Table A. End-to-End Throughput on eight GeForce RTX 3090 GPUs, each with 24 GB of GPU memory.*** 

| Models         | End-to-End Throughput (toks/s) |
| -------------- | ------------------------------ |
| SMoE Baseline  | 32.21                          |
| SMoE Baseline# | 33.98                          |
| MH-MoE         | 25.77                          |
| MH-MoE#        | 31.93                          |

---

>**Q2**: Thank you for clarifying why low activation of experts might be an issue. This indeed makes it clearer - perhaps the paper could reflect this answer and make this point a bit clearer. 

**A2**: Certainly, we will clarify this point in the latest version of our paper. Thank you very much for your suggestion.



---

> **Q3**: It would be interesting to add some simple baseline experiment that would somehow "force" the model to activate more experts and therefore increase performance - to really demonstrate that this is the matter. 

**A3**: We designed a small experiment using an MoE model based on the SMoE baseline with a 3-layer transformer, trained for 20k steps. The experiment includes two groups:

1. An MoE with 128 experts and top-1 activation (denoted as SMoE-128-top1).
2. An MoE with 256 experts and top-2 activation (denoted as SMoE-256-top2).

The corresponding perplexity results are shown in Table B. We observed that SMoE-256-top2, which activates more experts (two experts per token), performs better. This demonstrates that activating more experts can provide greater benefits to the model. We will include this small experiment and its explanation in the revised version of the paper to clearly demonstrate that activating more experts makes a significant difference.

***Table B.*** 

|               | PPL$\downarrow$ |
| ------------- | --------------- |
| SMoE-128-top1 | 16.31           |
| SMoE-256-top2 | **15.33**       |



---

> **Q4**: Alternatively (or rather additionally) it would be interesting to add additional hypothesis. (E.g. perhaps learning dynamics become very unstable, etc. - as why would SMoE perform worse with more experts, if it could just have a few "dead" experts and then it would be equivalent to an SMoE with less experts)

**A4**: Thank you for your interesting suggestion. One hypothesis we have is that when there are more experts, insufficient training of the experts may lead to underfitting, resulting in reduced model performance. We will continue to explore possible underlying reasons and include these hypotheses and explanations in the revised version of the paper to enrich the discussion.

---

> **Q5**: Similarly, with regards to number of heads and why performance starts dropping. This is an excellent observation, some validation of this hypothesis might be interesting.

**A5**: This is indeed a very interesting phenomenon. We also noted that previous studies have reported similar observations in multi-head self-attention, where increasing the number of heads does not necessarily lead to better performance. We are conducting further experiments to explore the essence of this phenomenon and the potential insights it might offer.

We can offer a hypothesis here: when the number of heads increases, the dimension of each head decreases, which may harm the representation power of each head. This can cause each head to suffer from a **low-rank bottleneck**, leading to reduced model performance. The impact of this low-rank bottleneck has been studied in the context of multi-head self-attention [1], and we are conducting similar research on MH-MoE in the future to further explore the reasons behind this phenomenon.

**Reference**

[1] Bhojanapalli, Srinadh, et al. "Low-rank bottleneck in multi-head attention models." PMLR, 2020.

---

>**Q6**: Are there any references for your particular implementation?

**A6**: We referenced the following libraries to implement our model:

[1] https://github.com/databricks/megablocks

[2] https://github.com/facebookresearch/fairseq

[3] https://github.com/microsoft/torchscale

---
Your suggestions have provided us with a lot of new inspiration and insights. If you have any further questions, please feel free to let us know. We would be happy to discuss them further with you! :)
Dear Author(s),

Thank you for the detailed rebuttal and response. 

> A1: 

Comment: Thank you for providing such clear and detailed response to this question.

Questions: 
a) Would dynamic routing improve the Tok/s for the baseline?

> A2: 

Comment: thank you for such a clear answer. This is very clear now.

> A3:

Comment: Thank you for clarifying this point. Yes, please add a revised version to make this clearer.

> A4:

Comment: 
a) Thank you for clarifying why low activation of experts might be an issue. This indeed makes it clearer - perhaps the paper could reflect this answer and make this point a bit clearer.
b) Interesting observation. Thanks for commenting this as well.

Questions:
a). While low-activation of experts seems quite plausible to be the cause of low performance for SMoE, it is not entirely proven. It would be interesting to add some simple baseline experiment that would somehow "force" the model to activate more experts and therefore increase performance - to really demonstrate that this is the matter. Alternatively (or rather additionally) it would be interesting to add additional hypothesis. (E.g. perhaps learning dynamics become very unstable, etc. - as why would SMoE perform worse with more experts, if it could just have a few "dead" experts and then it would be equivalent to an SMoE with less experts).
b). Similarly, with regards to number of heads and why performance starts dropping. This is an excellent observation, some validation of this hypothesis might be interesting.

> A5:

Comment: Thank you for providing initial information. Indeed this would be interesting to see in practice. Are there any references for your particular implementation?

Thank you again.