Multi-Head Mixture-of-Experts

Thank you for your feedback :)

Q1: 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.

A1: In fact, if traditional static routing [1] (i.e., using zero-padding or dropping tokens to force a uniform distribution of tokens among experts) were applied in MH-MoE, the situation would be different. Since MH-MoE uses sub-tokens for more complex and fine-grained routing decisions, the matrix dimensions required to record routing information expand accordingly during both training and inference. This results in more time-consuming matrix operations and higher all-to-all communication costs within nodes, thereby affecting the end-to-end throughput of deploying MH-MoE. We present a comparison of throughput between XMoE and MH-MoE using static routing (denoted as MH-MoE*) in Table A. Our results show that MH-MoE with static routing is slower than the baseline.

To address this issue, we have implemented dynamic routing [2] (refer to the implementation: databricks/megablocks (github.com)) in our MH-MoE model. Dynamic routing alleviates the constraints imposed by static routing on both software and hardware, enabling a more efficient and unrestricted routing strategy. Table A shows a comparison of the throughput taken by MH-MoE with dynamic routing (denoted as MH-MoE) and the baseline. We observe that using the more efficient dynamic routing approach results in end-to-end throughput for MH-MoE that is nearly identical to the baseline.

This is the approach we have taken to address the issue of high end-to-end training and inference throughput in MH-MoE. If you have any further questions, we welcome continued discussion during the discussion phase. :)

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

Reference

[1] Lepikhin, Dmitry, et al. "Gshard: Scaling giant models with conditional computation and automatic sharding." arXiv 2020.

[2] Gale, Trevor, et al. "Megablocks: Efficient sparse training with mixture-of-experts." PMLR 2023.

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: 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), 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)$.

Thank you for your detailed feedback. We address your feedback point by point below.

Q1: 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.

A1: In fact, we conducted experiments on two MoE structures:

1. For English-focused language modeling, multilingual language modeling, and masked multimodality modeling tasks, we implemented our MH-MoE on XMoE [1] and conducted extensive experiments across these three tasks (ranging from 300M to 7B parameters, with 3 upstream tasks and 26 downstream tasks) to validate the effectiveness of MH-MoE.
2. For pure vision tasks, as detailed in Section 5 of the main text, we implemented MH-MoE on AdaMV-MoE [2]. We validated that MH-MoE can further enhance AdaMV-MoE's performance in vision tasks such as classification, object detection, and instance segmentation.

We chose XMoE and AdaMV-MoE because they have been proven superior to vanilla SMoE in extensive experiments [1, 2]. We believe that the extensive experiments on these multiple upstream & downstream tasks with the two MoE structures demonstrate the effectiveness of MH-MoE. If you have any MoE structures of interest, please let us know during the discussion stage, and we would be happy to conduct further experiments :)

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.

Q2: 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.

A2: This is a meaningful suggestion. To explore the issue you mentioned, we conducted additional experiments with top-k=3 and top-k=4 to investigate the effects of increasing activated experts on MH-MoE.

We performed a language model experiment based on the RedPajama dataset, utilizing a 6-layer transformer architecture with each layer having a hidden dimension of 768. For the MoE settings, the model selects the top-k=3 or top-k=4 experts for each input. We employ a Multi-Head Mixture of Experts (MH-MoE) configuration with 2 heads. The learning rate is set to 6e-4, token number of each batch is 0.5M. Due to time and resource constraints, we present results for the model at 25k steps.

The experimental results are shown in the Table A below. We find that, regardless of whether k=2 （see in the maintext）, k=3, or k=4, MH-MoE consistently outperforms the baseline MoE model.

Table A.

Q3: 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.

A3: In the PDF file located in the Author Rebuttal at the top of the page, we provide a comparison of the changes in balance loss during training for XMoE Baseline and MH-MoE under two training settings, 32 experts and 8 experts, in the English-focused language modeling experiment. If you have any questions, please feel free to reach out to us during the discussion phase :).

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

Thank you for your feedback. We address your feedback point by point below.

Q1: 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.

A1: In fact, we conducted experiments on two MoE structures:

1. For English-focused language modeling, multilingual language modeling, and masked multimodality modeling tasks, we implemented our MH-MoE on XMoE [1] and conducted extensive experiments across these three tasks (ranging from 300M to 7B parameters, with 3 upstream tasks and 26 downstream tasks) to validate the effectiveness of MH-MoE.
2. For pure vision tasks, as detailed in Section 5 of the main text, we implemented MH-MoE on AdaMV-MoE [2]. We validated that MH-MoE can further enhance AdaMV-MoE's performance in vision tasks such as classification, object detection, and instance segmentation.

We chose XMoE and AdaMV-MoE because they have been proven superior to vanilla SMoE in extensive experiments [1, 2]. We believe that the extensive experiments on these multiple upstream & downstream tasks with the two MoE structures demonstrate the effectiveness of MH-MoE. If you have any MoE structures of interest, please let us know during the discussion stage, and we would be happy to conduct further experiments :)

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.

Q2: 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?

A2: This is a meaningful suggestion. To explore the issue you mentioned, we conducted additional experiments with k=3 and k=4 to investigate the effects of increasing k on MH-MoE.

We performed a language model experiment based on the RedPajama dataset, utilizing a 6-layer transformer architecture with each layer having a hidden dimension of 768. For the MoE settings, the model selects the k=3 or k=4 experts for each input. We employ a Multi-Head Mixture of Experts (MH-MoE) configuration with 2 heads. The learning rate is set to 6e-4, token number of each batch is 0.5M. Due to time and resource constraints, we present results for the model at 25k steps.

The experimental results are shown in the Table A below. We find that, regardless of whether k=2 （see in the maintext）, k=3, or k=4, MH-MoE consistently outperforms the baseline MoE model. Additionally, there is no indication of a saturation point in MH-MoE’s performance as k increases.

The activation map's behavior with an increase in k shows effects similar to those observed in Figure 6 of the main text. Specifically, as k increases, the activation map becomes brighter, indicating that experts are activated more frequently. This is because an increase in k means that each expert's probability of being assigned tokens also increases with each token assignment. Consequently, within a fixed number of training steps, the number of times experts are activated naturally increases.

Table A.

Q3: 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.

A3: Thank you for your suggestion. We will include an explicit limitation section in the latest version, summarized as follows:

An interesting observation is that the performance of our MH-MoE does not consistently improve with an increasing number of heads. We hypothesize that this may be due to each sub-token's dimension becoming too small when the number of heads is too large, leading to information loss within the tokens. This, in turn, makes meaningful and effective routing more challenging, thereby affecting the overall performance of the model. This is an intriguing area for further exploration, and we plan to investigate it in future research.

We are grateful to the reviewer for the extensive review. Due to space constraints, we address your questions in the Rebuttal below as well as at the top of this page in the Author Rebuttal:

Q1: What is the actual practical training or inference time of running MoE vs. MH-MoE.

A1: For training cost evaluation: In fact, if traditional static routing [1] (i.e., using zero-padding or dropping tokens to force a uniform distribution of tokens among experts) were applied in MH-MoE, the situation would be different. Since MH-MoE uses sub-tokens for more complex and fine-grained routing decisions, the matrix dimensions required to record routing information expand accordingly during training. This results in more time-consuming matrix operations and higher all-to-all communication costs within nodes, thereby affecting the end-to-end throughput of training MH-MoE. We present a comparison of throughput between XMoE and MH-MoE using static routing (denoted as MH-MoE*) in Table A. Our results show that MH-MoE with static routing is slower than the baseline.

To address this issue, we have implemented dynamic routing [2] (refer to the implementation: databricks/megablocks (github.com)) in our MH-MoE model. Dynamic routing alleviates the constraints imposed by static routing on both software and hardware, enabling a more efficient and unrestricted routing strategy. Table A shows a comparison of the training throughput taken by MH-MoE with dynamic routing (denoted as MH-MoE) and the baseline. We observe that using the more efficient dynamic routing approach results in end-to-end training throughput for MH-MoE that is nearly identical to the baseline.

Table A. End-to-End Throughput on eight GeForce RTX 3090 GPUs.

For inference cost evaluation, we implemented our MHMoE model on the Mistral 7*8B architecture and used vLLM (https://github.com/vllm-project/vllm) to test the inference speed. vLLM is a widely used LLM inference framework. We tested the inference speed on four A6000 GPUs. All GPUs processed the same prompts and used the same sampling parameters (temperature=0.8, top_p=0.95), with the tensor parallel size set to 4. To ensure the fairness of the inference speed test, given that the MHMoE model outputs garbled text, we set a consistent output length. The results are shown in Table B. We can observe that the inference speed of the MHMoE model is slightly slower than that of the SMoE model. Additionally, as the number of heads in MHMoE increases, the inference speed decreases slightly. However, the difference is not significant, and the MHMoE model still achieves a high inference speed.

Table B.

Reference

[1] Lepikhin, Dmitry, et al. "Gshard: Scaling giant models with conditional computation and automatic sharding." arXiv 2020.

[2] Gale, Trevor, et al. "Megablocks: Efficient sparse training with mixture-of-experts." PMLR 2023.

Q2: Evaluation and comparison is not done on publicly available models of relevant sizes (e.g. Mistral).

A2: The training data and steps for the models you mentioned are either lengthy or not well-defined, making them difficult to compare directly. To ensure fairness, we have used our own implementation of the two MoE methods for comparison:

1. For English-focused language modeling, multilingual language modeling, and masked multimodality modeling tasks, we implemented our MH-MoE on XMoE [1] and conducted extensive experiments across these three tasks (ranging from 300M to 7B parameters, with 3 upstream tasks and 26 downstream tasks) to validate the effectiveness of MH-MoE.
2. For pure vision tasks, as detailed in Section 5 of the main text, we implemented MH-MoE on AdaMV-MoE [2]. We validated that MH-MoE can further enhance AdaMV-MoE's performance in vision tasks such as classification, object detection, and instance segmentation.

We chose XMoE and AdaMV-MoE because they have been proven superior to vanilla SMoE in extensive experiments. We believe that the extensive experiments on these multiple upstream & downstream tasks with the two MoE structures demonstrate the effectiveness of MH-MoE.

Q3: Clarity a. some parts are unclear.

A3: Thank you for pointing this out. We apologize for the unclear or incorrect descriptions in a and b. In fact, in all experimental setups (except for the ablation experiments on MLP layers in Table 5 and the number of MLP layers in Table 13), the MH-MoE model uses both the multi-head layer and the merge layer. These two MLP layers are represented by the $W_\{head\}$ and $W_{merge}$ matrices in Equations 2 and 7, respectively. The multi-head layer and the merge layer perform matrix transformations on the last dimension (hidden state dimension) of the input tokens and the merged output tokens for feature modeling. We will revise these sections in the updated version of the paper to ensure a clearer and more accurate presentation.

Q4: 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. b. 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?

A4: Pleasee see in the Author Rebuttal section.

Q5: Reporting on sharding of the models.

A5: Pleasee see in the Author Rebuttal section