Dynamic Mixture of Experts: An Auto-Tuning Approach for Efficient Transformer Models

Thank you for your valuable suggestions. We have addressed the questions you raised below.

More experiments

Thank you for your suggestion, we have added the additional experiments, including

• Finetuning ViT-S on TinyImageNet dataset.
• Train ViT-S from scratch on CIFAR10 dataset. We train the models for 200 epochs due to the time limitation.
• Add an additional baseline [8] on MoE-LLaVA experiment.

Note that our BERT and ViT experiments are already conducted based on [1]. Furthermore, we conduct ViT experiments on TinyImagNet. We investigate various expert initialization strategies to further validate our findings, including:

1. Average: Averaging the parameters of existing experts to initialize the new expert.
2. W-Average: Using weighted averaging of the parameters of existing experts, where the weights correspond to the number of experts to be activated.
3. Most activated: Initializing the new expert using the parameters of the most frequently activated expert.

The results show that:

1. DynMoE converge faster than standard MoE settings (refer to Figure 24 of page 25 for more details);
2. W-Average achieve the best performance in most cases.
3. DynMoE achieves better performance than top-$p$ routing [8] without the requirements of the additional hyper-parameter $p$.

Literature review. To my understanding, the proposed method lies in the field of dynamic neural networks [6]. It is encouraged to discuss more about the relationship to other dynamic paradigms, such as layer skipping [7]. When it allows a token to select 0 expert, the proposed framework would naturally allows layer skipping.

Thank you for the suggestion. We have revised the paper (lines 535-539) to include a discussion about dynamic neural networks [6,7]. It is worth noting that in our current implementation, tokens are not allowed to select 0 experts during inference. However, we believe that this would be a valuable extension for future work.

[8] Harder Tasks Need More Experts: Dynamic Routing in MoE Models. arXiv preprint 2024.

Thank you for your valuable suggestions. We have addressed the questions you raised below.

In the top-any routing, is G a hyperparameter that needs to be predefined? Would this lead to cumbersome tuning?

We are sorry for the misunderstanding. The $G$ are trainable parameters, and are initially set to 0. In Figures 17-22 of the Appendix, we have reported the value of learned threshold $G$. Our findings indicate that there exists a shared expert that more easier to be activated in each MoE layer.

In the top-any routing, the selected k experts considered equally important by default. Would introducing varying importance levels for different experts yield better results? Although the authors mention related content in Remark 3.1, ablation experiments seem necessary.

We would like to clarify the following points:

• In top-k gating, the gating scores are computed using the softmax function, which allows for the comparison of gating scores among different experts. Specifically, a higher score for some experts results in a lower score for others.
• In contrast, in the top-any gating mechanism of DynMoE, the gating scores for different experts are calculated independently using sigmoid function (see figure 2 for a more intuitive understanding). Consequently, these scores are not comparable among experts. As such, we believe that the range of expert scores does not accurately reflect their importance.
• Additionally, we have conducted an ablation study to demonstrate the effectiveness of this design choice. Results show that using different expert scores significantly reduce the model performance.

Why top-any MoE training strategy dones't lead to increased training time?

We apologize for the misunderstanding and would like to clarify that the comparable training FLOPs in Table 4 of our submission is due to DynMoE activating fewer experts compared to MoE-LLaVA. Consequently, DynMoE's inference is more efficient. However, the training of DynMoE is not faster than MoE-LLaVA because of the additional training FLOPs associated with top-any routing and adaptive expert addition/removal.

The author could further elaborate on how the diversity loss in the auxiliary loss prevents activating all experts simultaneously.

Thank you for the suggestion. The diverse loss serves as a weak regularization for sparsity because:

• Our routing method calculates the cosine similarities between tokens and routing parameters, and then uses the similarity score as the gating score. As the router parameters are forced to be diverse, tokens with high similarity to certain experts tend to have low similarity to other experts. This prevents tokens from choosing too many experts, thereby promoting sparsity.
• As shown in the language tasks presented in Table 8 of the attached file, as well as in the vision-language tasks shown in Table 2 of the submission, DynMoE tends to activate only a subset of experts rather than all of them.

How exactly does DYNMOE remove experts? For example, if the interval for adding/removing experts is set to 300 steps and expert x is activated in the first 299 steps but not in the 300th step, will it be removed? Alternatively, if expert x is activated in the first step but not for the remaining 299 steps, will it be removed?

DynMoE removes experts that are never activated during the interval for adding/removing experts. Consequently, in both scenarios mentioned by the reviewer, the experts will not be removed.

For new experts added in DYNMOE, Would averaging the weights of other experts to initialize the new expert’s weights lead to better results? Or would random initialization yield better performance?

We apologize for the missing details in our submission. To clarify:

• The experts in our MoE layer are initialized by copying the weights of pre-trained models.

• In our current implementation, we re-use the weights of the most recently discarded expert when adding a new expert to the MoE layer.

• We conducted additional experiments with various expert initialization strategies when adding new experts, including:

  1. Average: Averaging the parameters of existing experts to initialize the new expert.
  2. W-Average: Using weighted averaging of the parameters of existing experts, where the weights correspond to the number of experts to be activated.
  3. Most activated: Initializing the new expert using the parameters of the most frequently activated expert.

  Results show that (1) DynMoE converges faster than standard MoE settings; (2) W-Average achieves the best performance in most cases (refer to Figure 24 of page 25 for more details).

Unconvincing real latency/throughput evaluation: A major flaw of the work is that authors only briefly mention the actual speedup that the method supposedly provides. The authors write (line 488): "To ensure a fair evaluation, MoE-LLaVA employs the expert dispatching implementation from DynMoE by fixing the top-k values.". This statement is unclear, and for me it suggests the opposite of a fair evaluation, as the alternative implementation may significantly slow down the original MoE-LLaVA. The readers should be aware of any overhead stemming from the introduction of top-any routing. I would at least expect a non-negligible overhead stemming from the non-constant number of token-expert executions.

We apologize for any previous ambiguity in our statement. However, we stand by our claim that we conducted a fair evaluation.

• The forward pass of a MoE layer consists of three key steps:
  1. Forwarding the gating network to generate gating scores.
  2. Dispatching tokens to the appropriate experts based on the gating scores and a top-$k$ selection process.
  3. Passing the tokens through the selected experts and combining their results.

MoE-LLaVA utilizes the Deepspeed-MoE implementation. However, despite its efficiency, the Deepspeed-MoE implementation for step (2) is limited to supporting only top-1 and top-2 gating. Consequently, we developed our own implementation for step (2), which resulted in some efficiency degradation due to additional engineering efforts, i.e., we did not incorporate a CUDA implementation for parallel computation.

To ensure a fair comparison focused on algorithm design, we maintained the same implementation for step (2) but varied the value of $k$ for MoE-LLaVA and DynMoE. For steps (1) and (3), MoE-LLaVA adopted the original implementation provided by Deepspeed-MoE.

In summary, we believe we have presented a fair comparison.

Finally, to provide a complete picture, the authors should also measure and report the throughput of dense models for comparison.

We have provided the throughput of dense models below. Overall, the dense 1.6B model activates fewer parameters and is more efficient than MoE variants, while DynMoE (row 3) is more efficient than the standard MoE (row 2).

No from-scratch training experiments: All experiments are conducted on pre-trained models, which are then converted to sparse MoE models during the fine-tuning. Mixture-of-Experts layer was originally proposed to scale up the number of parameters without slowing down from-scratch training. Although the pre-trained setup is as important, demonstrating that the proposed method is also appropriate for the from-scratch setup would strengthen the paper.

Thank you for your suggestion, we would like to clarify that:

1. MoE is also crucial for the fine-tuning stage in the LLM literature [2, 10]. By using MoE, we can scale up a pretrained LLM to a larger one with minimal additional cost.

2. Additionally, we have conducted an experiment where we trained the model from scratch to further validate our findings. We trained ViT-S on the CIFAR-10 dataset from scratch, and train the models for 200 epochs due to the time limitation. We use various expert initialization strategies while adding the new experts, including:

   1. Average: Averaging the parameters of existing experts to initialize the new expert.
   2. W-Average: Using weighted averaging of the parameters of existing experts, where the weights correspond to the number of experts to be activated.
   3. Most activated: Initializing the new expert using the parameters of the most frequently activated expert.

   The results show that DynMoE (Original), DynMoE (Average), and DynMoE (W-Average) achieve faster convergence rate than the standard MoE (refer to Figure 24 of page 25 for more details). | CIFAR10 (ViT-S, 2 MoE Layer, Acc per 50 Epoch) | 10 | 50 | 100 | 150 | 200 | | --- | --- | --- | --- | --- | --- | | MoE (K = 8, k = 1) | 46.18 | 59.76 | 64.85 | 68.64 | 72.66 | | MoE (K = 8, k = 2) | 47.10 | 60.44 | 66.07 | 70.86 | 73.79 | | MoE (K = 8, k = 4) | 46.71 | 59.40 | 64.73 | 69.07 | 72.77 | | MoE (K = 8, k = 8) | 48.07 | 60.85 | 64.85 | 68.50 | 70.68 | | DynMoE (Original, avg topk=7) | 45.54 | 61.88 | 66.91 | 70.91 | 74.84 | | DynMoE (Average, avg topk=7) | 47.05 | 62.77 | 67.20 | 71.24 | 74.70 | | DynMoE (W-Average, avg topk=6.5) | 48.56 | 60.94 | 65.00 | 69.66 | 74.16 | | DynMoE (Most activated, avg topk=6.5) | 45.79 | 60.84 | 65.83 | 69.48 | 71.71 |

Thank you for your valuable suggestions. We have addressed the questions you raised below.

Limited novelty and lack of proper attribution: While the authors demonstrate some improvement over standard MoE layers, the idea of dynamically selecting the number of executed experts on a per-token basis is not new [1,2,3,4]. Moreover, none of these are mentioned in the related work section. The authors should properly discuss these works and highlight the differences of the proposed method.

We appreciate the reviewer's feedback on the matter. However, we believe that NONE of these studies [1,2,3,4], including several contemporaneous ones, have addressed the issues we intend to tackle in this paper. Consequently, our approach remains valuable to the research community.

• Regarding the use of hyper-parameters for gating, while some studies [1,2,3,4] enable tokens to activate different number of experts, new hyper-parameters are introduced. DynMoE does not require any additional hyper-parameters and truly eliminates the need for the hyper-parameter $k$.
• Some studies [1, 2] rely on dense training during the initial stages, which may not align well with the intention of utilizing MoE architectures in LLMs, namely, scaling up models with fewer computational resources. Consequently, [1] did not conduct experiments on LLMs, and [2] did not provide a comparison of training efficiency. In contrast, DynMoE consistently uses sparse gating for all training and inference stages and offers a comprehensive discussion on efficiency.
• Furthermore, unlike [1,2,3,4], which cannot adjust the number of experts $K$, DynMoE allows for the automatic adjustment of the number of experts during training.
• In addition to its innovative architecture, the results of DynMoE offer valuable insights into the design of MoE architectures. As detailed in Section 4.5, we can observe two key findings: (1) the bottom layers require a greater number of experts; and (2) there are shared experts present within each MoE layer.

In summary, we maintain that DynMoE retains significant value to the community in terms of novelty and contribution.

Lack of comparison to existing MoE methods from the literature: Given the above, does the proposed method provide any advantages over the existing top-any literature? An empirical comparison with these methods would strengthen the contribution. As of right now, the authors compare their method with only a single baseline (one of the top-k MoE variants), which is an extremely low number for an empirical-only work. Would the benefits of the proposed method hold even when stacked against alternative top-k MoE variants? Some of these variants also allow for an effectively different number of experts executed per token in inference [5].

Thank you for your suggestion. We would like to clarify the following points, and have included the discussion about related works in the revised paper (lines 103-107):

1. We compare the DynMoE to the standard top-$k$ MoEs due to the following reasons:
   • The top-$k$ MoE remains a strong baseline and is prevalent in recent MoE-based LLMs [9, 10].
   • DynMoE offers several advantages over existing MoE variants: (1) it does not require hyperparameters during gating, (2) it can automatically adjust the number of experts, (3) it supports sparse training, and (4) it does not rely on specialized training stages.
2. During the rebuttal, we have incorporated additional baselines. For the related studies [1,2,3,4]
   1. [1,2] relies on dense training, thereby difficult to conduct a fair evaluation.
   2. Both [3,4] share the same key idea that changing top-$k$ gating to top-$p$ gating, thereby we use [3] as the baseline algorithm we compared, and using the $p=0.4$ as suggested in their paper.
   3. We can find that DynMoE achieves better performance to [3] (where $p$ is set to 0.4 as in the paper) without the additional parameter $p$.

Is the MoE implementation used in this work dropless [6], i.e. does it drop tokens if the predefined capacity of an expert is exceeded? If it is not dropless, how often are tokens dropped and how does DynMoE compare to standard MoE in terms of load balance?

Thank you for the comments. We would like to provide the following clarification

1. We use a dropless implementation, and the capacity are adaptively calculated. In, detail the capacity is calculated by $c = max_{k}{N_k}$, where $N_k$ is the number of tokens that choose expert $k$.

2. In our submission, we did not enforce an equal load distribution among experts based on the following considerations:

   1. The significance of experts varies, with some being more likely to be activated by all tokens. This issue is recognized in recent MoE studies [9] and is also observed in our research, as discussed in Section 4.5.
   2. Enforcing load balance by including a load-balance loss might not align with our requirements for sparse routing, as activating all experts could be considered a form of load balance.

   We believe that the load balance issue can be addressed through engineering techniques, such as allocating additional computation units to shared experts.

3. To further illustrate our point, we have conducted additional ablation studies on the MoE-LLaVA setting. Specifically, we have: (1) Included a load-balance loss in our model; (2) Introduced an efficiency loss to enforce sparsity, as suggested in [2]. We report the following metrics:

   1. Performance (Table 1): The performance of different settings.
   2. Load Balance (Table 2): The frequency with which each expert is activated, calculated as (expert activation time / total token count).
   3. Efficiency (Table 3): The top-k values per layer.
   4. Efficiency (Table 4): The top-k activation frequency, calculated as (number of tokens that activate k experts / total tokens).

   We can find that

   1. Although it does not explicitly enforce load balancing, the original DynMoE achieves load balancing comparable to that of the standard top-2 MoE (Table 2).
   2. Adding the load balance loss slightly decreases the performance of DynMoE (Table 1) while increasing the number of activated experts (Table 3). However, it improves load balancing (Table 2).
   3. Adding an additional efficiency loss on top of the load balance loss improves performance (Table 1) and helps overcome some extreme cases, such as the reduction of the top-$k$ values in the bottom layer from 2.88 to 2.02 (Table 3), and reduce the number of tokens that activate all 4 experts (Table 4). Moreover, the efficiency loss further enhances load balancing (Table 2).

[9] DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models.

[10] Mistral 7B.

[11] Moe-llava: Mixture of experts for large vision-language models.

[12] Emergent Mixture-of-Experts: Can Dense Pre-trained Transformers Benefit from Emergent Modular Structures?

How are experts initialized at the start given a pre-trained model? Are parameters of experts that are added during training initialized randomly, or are the weights of the most recently discarded expert re-used?

We apologize for the missing details in our submission. To clarify:

• The experts in our MoE layer are initialized by copying the weights of pre-trained models.

• In our current implementation, we re-use the weights of the most recently discarded expert when adding a new expert to the MoE layer.

• We conducted additional experiments with various expert initialization strategies when adding new experts, including:

  1. Average: Averaging the parameters of existing experts to initialize the new expert.
  2. W-Average: Using weighted averaging of the parameters of existing experts, where the weights correspond to the number of experts to be activated.
  3. Most activated: Initializing the new expert using the parameters of the most frequently activated expert.

  Results show that (1) DynMoE converge faster than standard MoE settings; (2) W-Average achieve the best performance in most cases (refer to Figure 24 of page 25 for more details).

Remark 3.1 is not fully convincing to me. Can the authors explain more in detail why different ranges of expert scores are supposed to be undesirable?

We would like to clarify the following points:

• In top-k gating, the gating scores are computed using the softmax function, which allows for the comparison of gating scores among different experts. Specifically, a higher score for some experts results in a lower score for others.
• In contrast, in the top-any gating mechanism of DynMoE, the gating scores for different experts are calculated independently using sigmoid function (see figure 2 for a more intuitive understanding). Consequently, these scores are not comparable among experts. As such, we believe that the range of expert scores does not accurately reflect their importance.
• Additionally, we have conducted an ablation study to demonstrate the effectiveness of this design choice. Results show that using different expert scores significantly reduce the model performance.

Standard deviations not being reported: The authors do not report standard deviations of the results in any of the experiments. This is particularly surprising as the authors write in the appendix that the scores are averaged over three random seeds. Without standard deviations it is hard to estimate the significance of the results, especially since some of the scores in Tables 1, 2, 3, and Figure 4 are really close. The authors also do not present any evidence that training of the proposed DynMoE is stable in any separate experiment.

Thank you for the suggestion!

• We have added the std for language tasks. Results show that DynMoE generally has lower std value, indicating more stable performance.
• Due to the high computation cost, LLMs general do not report this [2, 11, 12].
• Furthermore, we would like to emphasize that the primary objective of this paper is to eliminate the reliance on hyper-parameters, rather than solely improving model performance.

Readability: I would recommend to do a pass through the paper to improve grammar and readability (e.g. "and uses an additional gating network g to predict the scores that the input token embedding assigned to each expert." or "To promising efficiency and avoiding burden communication, we only check if experts required to be added add or removed.").

Thank you for the suggestion. We have revised the first paragraph as follows and removed the additional part from the second sentence.

The traditional top-$k$ gating method takes the token embedding $x$ as input and employs an additional gating network $g$ to predict the gating scores. These gating scores are then used to determine which experts will be activated for the input tokens.

Are the auxiliarry losses somehow weighted? If yes, why is it not reported in the paper?

We are sorry for the missing details. We use the default setting of EMoE and MoE-LLaVA for the weights of auxiliarry losses. For both of these two settings, the weights of auxiliarry losses are set to 1e-2.

In the appendix the authors write: "For each epoch, we begin recording routing at 1/3 of the epoch and complete recording routing and execute the adaptive process at 2/3 of the epoch." Can the authors clarify why is this needed?

Thank you for your detailed comments. Our design choice aims to prevent experts from being frequently removed or added without adequate training and gating information. By using this method, we can ensure:

1. Sufficient gating information is collected during the epoch range from 1/3 to 2/3 to make informed decisions about adding or removing experts.
2. Experts are adequately trained before the next round of expert addition/removal.

We found that this design choice stabilizes the training process.

Currently the method relies on a linear router. However, some works assume a deeper router [7,8]. Can the method be extended to non-linear routers

Thank you for your insightful comments. While using a deeper router is feasible, it may come at the cost of losing some of the benefits associated with DynMoE. Specifically:

• We can still perform top-any gating by passing tokens through the router and then applying the designed threshold sigmoid gating to activate the appropriate number of experts.
• The model removal process remains effective when using deep routers, but the model addition process becomes more challenging. In the current implementation, we can initialize the router parameters of new experts using the embeddings of tokens that do not choose any experts, ensuring that these tokens will activate the new expert. However, this simple initialization method may not be feasible when using a deep router.

Dear Reviewer 4q9o,

Thank you for your time and effort in reviewing our paper. We appreciate your constructive feedback and understand your concerns regarding the potential limitations of DynMoE. We would like to address the following points:

• Potential misunderstanding in comments: We believe there may be some misunderstanding regarding the reviewer's concerns. In both your initial comments and your first response, the focus appears to be on whether we conducted a fair comparison (to MoE-LLaVA). Specifically, your feedback includes:

  • "Unconvincing real latency/throughput evaluation"
  • "I am still leaning towards rejection of this work. I am not convinced by the authors' response that the throughput evaluation is fair."

  In response to these comments, we made sure to address the fairness of our evaluation in our previous reply, and we believe that we adequately clarified this point. This is why we referenced the ‘Old Request’ and ‘New Request’ in our earlier response, as we felt this issue had already been resolved.

• Overhead of the MoE Algorithm: We agree with your observation that the newly designed MoE algorithm may introduce additional overhead. As discussed in our previous response, the main source of overhead comes from the routing process, which is likely the reason why DynMoE (with forced top-1) is clearly slower than MoE (also with forced top-1). We have also provided more details on this aspect below

  Time (per sample and MoE layer)	Routing (ms)	Gating (ms)	Expert Passing (ms)
  MoE-LLaVA (enforced top-1)	0.04	0.60	1.08
  DynMoE (enforced top-1)	0.45	0.75	1.10

• Future Work on Efficient MoE Implementations: We acknowledge that further optimizations are needed to reduce the overhead of our approach. In the revised version of the paper, we will

  • provide more detailed experimental settings for each table;
  • include the discussion of efficiency as outlined in our rebuttal;
  • clearly discuss the limitations regarding if the dynamic mechanism can be implemented efficiently on GPUs without introducing any overhead.

While we recognize the potential limitations, we believe our key findings—particularly regarding the auto-tuning approach—are significant and offer valuable new insights for the field.

Once again, we appreciate your detailed review and constructive feedback. Thank you.

Unclear on what mechanisms are at play to optimise the gating mechanisms. In particular it appears to this reviewer that the only gradient to the gating and gating thresholds is via the (1/k) factor of eq (6).. If so one can tend to understand that no per-expert specific gradient is observed.. By that I mean, for a token, all of the params used to compute its gating values will observe similar gradient directions regardless of which expert they refer to.

We would like to clarify the following points:

1. The gradients of the gating network are dependent on the expert outputs $E_{e}(x)$, which vary among experts. Consequently, the gradients of the gating parameters corresponding to different experts will also differ.
2. In Figures 17-22 of the Appendix, we have reported the similarities of the gating parameters and the learned threshold $G$. Our findings indicate that the learned parameters and $G$ corresponding to different experts are significantly different, even though $G$ is initially assigned the same value.

[1] DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models.

[2] Learn To be Efficient: Build Structured Sparsity in Large Language Models.

Thank you for your valuable suggestions! We have addressed the concerns you raised below.

In my understanding, early design of MoEs [Shazeer et al., 2017] had many engineering challenges in consideration such as efficient use of compute units and network, which led to many of the fixed top-k routing decisions (e.g. why should the compute unit that processes expert A process only 10 tokens.. if it has to wait for another compute unit that has to process 100 tokens?) which is also similar to the concerns of equal utilisation of experts with the purpose to efficiently utilise hardware. With that in mind I find it hard to understand from the text in this paper: (a) why those concerns don’t apply here, (b) how has hardware changed overtime such they are no longer important, (c) if the approach here would not scale to utilise hardware as efficiently as a fixed top-k routing as one scales the method.

Thank you for your valuable feedback! We would like to provide the following explanation regarding your concerns:

• In our submission, we did not enforce an equal load distribution among experts based on the following considerations:

  1. The significance of experts varies, with some being more likely to be activated by all tokens. This issue is recognized in recent Mixture of Experts (MoE) studies [1] and is also observed in our research, as discussed in Section 4.5.
  2. Enforcing load balance by including a load-balance loss might not align with our requirements for sparse routing, as activating all experts could be considered a form of load balance.

  We believe that this issue can be addressed through engineering techniques, such as allocating additional computation units to shared experts.

• To further illustrate our point, we have conducted additional ablation studies on the MoE-LLaVA setting. Specifically, we have: (1) Included a load-balance loss in our model; (2) Introduced an efficiency loss to enforce sparsity, as suggested in [2]. We report the following metrics:

  1. Performance (Table 1): The performance of different settings.
  2. Load Balance (Table 2): The frequency with which each expert is activated, calculated as (expert activation time / total token count).
  3. Efficiency (Table 3): The top-k values per layer.
  4. Efficiency (Table 4): The top-k activation frequency, calculated as (number of tokens that activate k experts / total tokens).

  We can find that

  1. Although it does not explicitly enforce load balancing, the original DynMoE achieves load balancing comparable to that of the standard top-2 MoE (Table 2).
  2. Adding the load balance loss slightly decreases the performance of DynMoE (Table 1) while increasing the number of activated experts (Table 3). However, it improves load balancing (Table 2).
  3. Adding an additional efficiency loss on top of the load balance loss improves performance (Table 1) and helps overcome some extreme cases, such as the reduction of the top-$k$ values in the bottom layer from 2.88 to 2.02 (Table 3), and reduce the number of tokens that activate all 4 experts (Table 4). Moreover, the efficiency loss further enhances load balancing (Table 2).