Paper: Mixture-of-Experts Meets Instruction Tuning: A Winning Combination for Large Language Models (ICLR-2024)

Reviewer comment (Evaluation):
There is little insight into what is causing the failure of direct finetuning of MoE models on downstream task data. It could be that MoE models have higher capacity to overfit, however it is unclear if instruction tuning is preventing this or if there are other factors at hand. More conceptual insight would be nice, however this is not viewed as a major weakness.

Author response:
Thanks for pointing out, we included a discussion regarding the motivation quantitative and qualitative in general response 2. 

Also, as highlighted in [1], MoE models often suffer from greater generalization challenges than their Dense counterparts, despite potentially having equivalent or lower pretraining losses. This is primarily because (1) their significantly larger total parameter count tends to predispose them to overfitting, and (2) the complexity introduced by different routing strategies, auxiliary losses, expert dropout rates, and sensitivities to batch size and precision, which result in a higher number of hyperparameters needing fine-tuning. Given this context, direct task-specific fine-tuning might render MoE models more prone to suboptimal performance compared to dense models. This paper aims to demonstrate that instruction tuning can act as an intermediary phase for pre-trained MoE models. It helps in stabilizing hyperparameters and bridging the gap between pre-training and task-specific fine-tuning, thereby enhancing their generalizability.

[1] Artetxe, Mikel, et al. "Efficient Large Scale Language Modeling with Mixtures of Experts." Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing. 2022.

Reviewer comment (Experiments):
*Table 1:* Why not show Switch and GS performance at the 32G FLOP scale?

Author response:
This is a very insightful ablation to have! However, training a model with 32G FLOPs could span several months, and considering our current computational resource limitations, we have focused on scaling the most effective ST-MoE-32B models. We acknowledge the importance of this ablation and plan to address it in future work.

Reviewer comment (Presentation):
*Figure 3:* Why not label the orange and blue curves by model size, at least in the caption?

Author response:
Thanks a lot for the suggestions, we have included the model sizes in the captions in the updated version.

Reviewer comment (Evaluation):
*Figure 6:* How is expert utilization measured?

Author response:
Thank you for emphasizing this important aspect. We have expanded on this point in the revised version of our paper. We calculate the utilization of each expert by dividing the number of tokens dispatched to it by the total capacity available to that expert. This calculation is averaged across all experts, providing a measure of the actual usage of each expert in the model.

Reviewer comment (Writing):
*Figure 2 caption:* “Average zero performance” should perhaps be corrected to “Average zero shot performance?”.

Author response:
We corrected all the typos and added more clarifications of the capations in the updated version.

Reviewer comment (Presentation):
Add a period after the paragraph title “Routing Strategy” for formatting consistency.

Author response:
We corrected all the typos and added more clarifications of the capations in the updated version.

Reviewer comment (Writing):
*Appendix p. 15:* “We leave the study of scaling decoder-only FLAN-MOE as future works.” should be corrected to “We leave the study of scaling decoder-only FLAN-MOE to future work.”

Author response:
We corrected all the typos and added more clarifications of the capations in the updated version.

Reviewer comment (Experiments):
While several MoE models have been tested, the conclusion about the necessity of the instruction-tuning stage is not convincingly demonstrated. For instance, the addition of this instruction-tuning stage can introduce additional training and tuning costs, e.g., in comparison to using just task-specific fine-tuning. Is it possible that the performance improvement can also come from this extra training cost?

Author response:
Thanks for pointing this out. We added more clarification regarding the instruction tuning stage in the updated version. 

Specifically, the model of Dense / MoE will not be changed in the instruction tuning stage, and we conduct instruct-tuning for 100k steps following [6]. The instruction-tuned models introduce 10% of the pre-training cost, or 10x of the single-task finetuning cost regardless of Dense or MoE models. We also include a detailed discussion regarding the training / inference cost of MoE models versus Dense models in general response 1.

In addition, we used instruction-tuned dense models as baseline approaches. We showed that instruction-tuning benefits MoE models more than dense models. Thus, the gain is not only from the added instruction-tuning cost.

Reviewer comment (Writing):
Some statements made by the paper are rather confusing. For example, the paper states, “However, we show that conventional, task-specific finetuning MoE models lead to suboptimal performance, often even worse than finetuning dense models with the same computational cost. One of the possible reasons is the discrepancy between general pretraining and task-specific finetuning.” However, regardless of whether the model architecture is dense or sparse, isn't there always a discrepancy between pre-training and task-specific fine-tuning?

Author response:
Thanks for raising this important clarification problem. We also updated the clarification in the introduction. Yes, the discrepancy between pre-training and task-specific fine-tuning always exists for both dense and MoE models. 

However, as pointed out in [1, 5, 2], MoE models may encounter even more generalization challenges even if they have the same or lower pertaining loss as Dense models given (1) the large total number of parameters makes it inclined to overfit; (2) different routing strategies and auxiliary losses, expert dropout rates, sensitivity to batch size and precision introduce more hyperparameters for MoE models to tune. 

Considering this, direct task-specific finetuning may lead MoE models more easily to suboptimal performance compared to dense counterparts. In this paper, we are trying to show that instruction tuning could serve as a transition phase for pre-trained MoE models to fix the hyperparameters and mitigate the additional discrepancy between pre-training and task-specific fine-tuning, and therefore improve its generalizability.

Reviewer comment (Writing):
When the paper says the FLAN-MoE “does not come from increased computation resources or memory requirements,” what does it mean? Does it refer to computation and memory requirements during training/inference compared to compute-equivalent dense/MoE models?

Author response:
Thank you for pointing out this essential aspect. For a detailed exploration of training and inference expenses, along with the FLOPs comparisons for MoE and Dense models, kindly refer to the thorough explanation provided in our general response 1.

Reviewer comment (Writing):
What is the difference between FLAN-MoE and MoE in Section 4.1?

Author response:
The architecture of Flan-MoE is the same as MoE in Section 4.1, while the Flan-MoE refers to the instruction-tuned MoE models and MoEs refer to the pre-trained MoEs.

Reviewer comment (Experiments):
The paper does not seem to describe the model architecture of FLAN-MoE adequately. The only one mentioned is ST-MoE-32B. It would be interesting to see how different pre-trained MoE models would affect the conclusion.

Author response:
Thanks for pointing this out. We added more details in the updated paper. 

Specifically, all the pre-trained MoE models follow previous studies where Switch Transformer (Switch) has only one activated expert while GShard (GS), Expert-Choice (EC) [4] and ST-MoE use two experts per token. GS, EC, and Switch have the same pre-training objectives and model training architecture data as equivalent T5 models, while they differ in routing strategy and auxiliary losses. 

ST-MoE uses the same data as PaLM models and the pre-training objectives and model architecture are as T5 models. From Table 1, we can see that the advanced routing strategy benefits as EC/ST-MoE outperforms GS, and using more experts also benefits as GS outperforms Switch. 

Moreover, we include a preliminary study regarding the decode-only GLAM model [3] in Appendix A.3, it has the equivalent architecture as GPT3 except for the MoE layers which follow the design of GS at xl scale in Table 8. It can be seen that decoder-only MoE model benefits more from instruction tuning, which shows the potential of Flan-MoE at more generalized architecture and objective setting. We leave the study of scaling decoder-only Flan-MoE to future works.

Reviewer comment (Novelty):
Lack of clear motivation. The motivation behind the combination of MoE with instruction tuning requires further discussion. While it is acknowledged that instruction tuning and MoE models can outperform dense models or fine-tuning MoE, it will be beneficial to provide some insights into why these approaches were chosen.

Author response:
Thanks for the key clarification question regarding the motivation. We have updated our clarification in the introduction in the revised version: ‘’As shown in the ST-MoE study [Zoph et al], MoE models are more prone to overfitting on downstream tasks. Instruction-tuning adds multi-task losses which can be considered as regularization terms.’’ 

We also provide a more detailed analysis in general response 2 for your reference.

Reviewer comment (Evaluation):
I understand the performance of instruction tuning on MoE models, but can you please provide any analysis or insight about the reasons behind such good performance? Does it help improve the routing strategy, expert specialization, or something else?

Author response:
Thanks for the key clarification question regarding the motivation. We have updated our clarification in the introduction in the revised version: ‘’As shown in the ST-MoE study [Zoph et al], MoE models are more prone to overfitting on downstream tasks. Instruction-tuning adds multi-task losses which can be considered as regularization terms.’’ 

We also provide a more detailed analysis in general response 2 for your reference.

Reviewer comment (Evaluation):
The author claims that these advancements are attained without necessitating an increase in computational resources or even reducing the resource requirements. I am confused about how it can reduce resource requirements in the training and inference time. Can you discuss more about the details of the process?

Author response:
Thank you for bringing this crucial point to our attention. We also present an in-depth discussion regarding the training and inference costs, as well as the FLOPs for MoE and Dense models in general response 1 for your reference.