Pushing Mixture of Experts to the Limit: Extremely Parameter Efficient MoE for Instruction Tuning

We appreciate R Bdev for their valuable feedback and their observations, noting that the “proposed method is well described and the paper is easy to follow” and that our proposed method “can achieve better quality.” We acknowledge the concern raised for the preliminary analysis of the efficiency metric being reported for our proposed method, and we used the rebuttal period to run additional, more comprehensive benchmarking

We have now curated a comprehensive table detailing the efficiency metrics, which include parameter number, training and inference time, memory consumption, and overall performance for various methods applied to the T5-3B. By using these metrics, we aim to provide a more nuanced understanding of the performance-to-cost ratio, a crucial aspect of our study's focus on efficiency. This addition not only addresses the specific concerns raised but also enhances the overall comprehensiveness of our paper. We are confident that this addition substantially enhances the quality of our work. Furthermore, we will apply the recommended revisions addressing semantic and spelling errors in the final version of the manuscript. We respectfully request that the reviewer consider revising their evaluation in light of these improvements.

Training time and inference cost: We evaluated $(IA)^3$, MoV, and LoRA for training and inference times. $(IA)^3$ trains the fastest but with the worst-performing method. MoV (10 and 30 experts) offers a balance of significantly higher training speed relative to full fine-tuning and higher performance, outperforming LoRA. In inference, $(IA)^3$ and LoRA add no extra time when their weights are integrated into the base model, while MoV increases it by only 1.1x to 1.15x depending on the number of experts.

The number of fine-tuned parameters and memory consumption: In the table, we provide the number of fine-tuned parameters (# Params) for each method. In terms of updated parameters vs. performance trade-off, our MoV method outperforms others while balancing parameter count and performance effectively. Specifically, MoV with 30 experts surpasses LoRA rank 16 in performance while updating a smaller number of parameters.

We also include the maximum memory usage during training for each method. Note that this memory usage is influenced by the total parameters, fine-tuned number of parameters, optimizer, batch size, and hardware. Notably, our MoV method at the 3B scale offers optimal performance without significant efficiency costs during training and inference. It also scales effectively to 11B, maintaining competitive fine-tuning performance, as shown in Figure 1 and Table 4. We plan to include this table and a detailed cost-benefit analysis in the final paper in response to reviewer feedback.

Do you have any insights why more number of experts doesn't always give better results?

We hypothesize that as the number of experts increases, each additional expert contributes less to overall performance improvement because the most significant gains are often achieved early on when key expertise areas are covered. Additional experts might overlap in their specialization, leading to diminishing returns; this notion is reflected in the routing probabilities of experts in Figure 5. The same empirical observations were illustrated in the following [1, 2].

[1] Fedus, William, et al. “Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity.” arXiv preprint arXiv:2101.03961 (2022).

[2] Shen, Sheng, et al. "Mixture-of-experts meets instruction tuning: A winning combination for large language models." arXiv preprint arXiv:2305.14705 (2023).

We would like to thank R hesi for their detailed review. We are particularly flattered that they mentioned we “utilize extensive accuracy vs. compute figures to demonstrate the superior efficiency.” Also, we appreciate them noting the “scalability of the proposed methods, which is important for large models.” We address their remaining concerns below:

In Sec. 3, the authors mention to apply load balance loss specifically for discrete merging. Is load balancing applied also for soft merging? And does soft merging rely on load balancing loss to prevent expert collapse? If not, how to specilize the abilities of each expert to obtain the results in Fig. 5 with simple token-level routing without any explicit control?

Thanks for pointing out an unclear description. Our soft merging does not rely on any external load-balancing loss. The main reason that simple soft merging achieves expert specialization is the distinct task and dataset composition in the training, which is characterized by instruction tuning. Since we start with a pre-trained model, token representations are already diversified from each other based on different tasks and datasets, enabling routing to be specialized.

In our experiments, we found that load balancing loss decreases performance in both soft and discrete routing cases. We attribute this to two potential factors: (1) highly different data sizes for each dataset in our instruction tuning mixture, and (2) the misalignment between pretraining loss and fine-tuning loss when we add an auxiliary load balancing loss. This is similar to the experimental result shown in [1; Section 4.1, FLAN-ST_base].

If indeed we can get specialized experts, does the similar distribution shown in Fig. 5 suggest that the evaluated tasks still have connection with the several training tasks, otherwise it is hard to understand why a linear router can generate such a low entropy distribution on the evaluated tasks with a zero-shot setting?

Thanks for giving us the opportunity to highlight this important point. Indeed, there are cross-task relationships between multiple training tasks and held-out evaluation tasks. For example, [2] has shown that question-answering tasks such as “cosmos_qa” and “quail” have strong cross-task transfer to NLI tasks (e.g., super glue cb, super glue rte).

To validate this, during the rebuttal period, we fine-tuned task-specific $(IA)^3$ vectors for each task category in training and evaluated them for held-out tasks. As shown in the table below, different training tasks perform best on different held-out tasks. For example, the multiple-choice_qa task is most beneficial for super_glue_cb and super_glue_copa. Conversely, paraphrase and structure_to_text super_glue_rte tasks work best on super_glue_wsc.

As validated by our experimental results, MoE-style architecture better exploits such latent relationships with expert specialization compared to other parameter-efficient fine-tuning methods $(IA)^3$, and LoRA, resulting in higher performance with similar efficiency.

Thanks for the detailed response. However, there are several questions raised in my initial review that the authors do not response directly.

1. About the usage of loading balancing loss.

The authors specifically claim that they do not adopt loading balancing loss, suggesting that the learnt experts can not be sparsely activated during training, otherwise according to [1], even initialization with a pre-trained model would still meet the risk of the expert degradation problem.

2. About the performance ceiling for fully fine-tuning.

If this so called performance ceiling does exist and 1) it is close with the performance of full fine-tuning, how to explain why in Table 1 MoLORA-15 (rank 4) exceeds the T0-3B baseline? 2) If this ceiling is way higher than the performance of full fine-tuning (and also MoLORA), then the authors' response do not make sense any more.

3. About the novelty of the architecture.

Instead of a novel architecture, the proposed method is just an another implementation of the tranditional MoE framework. The tranditional MoE is defined in the Equation (MoE) on top of Page 4 in the main paper, and there are lots of different implementation for the $E_i(\cdot)$ function, including 1) MLPs, 2) IAs and 3) LoRAs.

Therefore, I do agree with Reviewer EZ7y that MoE is utilized in this paper only because the authors insist on using it without a solid research motivation, and as shown in the experiment tables in the authors' reply to Reviewer EZ7y, the less than 2.0% improvement is also far away from "Pushing Mixture of Experts to the Limit" in the paper title. Thus, my score stands.

[1] Wu, Lemeng, et al. "Residual mixture of experts." arXiv preprint arXiv:2204.09636 (2022).

We thank R EZ7y for their very positive review of our work, noting that our paper is “well-written, presenting a clear and sound approach” and emphasizing strongly the importance of the research problem. We greatly appreciate their claim that our proposed work is a “significant and valuable contribution to the field.” We conducted an additional experiment during the rebuttal period to evaluate our method on additional summarization tasks.

Regarding the use of the T5 model family in our experiments and extending to decoder-only architectures:

We thank R EZ7y for suggesting an additional experimental axis to further evaluate our proposed method. We note that our choice to focus on the T5 model family was partly dictated by how extensive and expensive our experiments were. We trained 100+ models, the majority of them on 3B and 11B scales, including preliminary experiments, final model runs, ablations, and analysis. T5 is also heavily used in prior work [1, 2], which allows for a clean comparison.

Our primary focus in the given work is to show the effectiveness of our MoE-style parameter-efficient architecture and propose an efficient alternative to the full fine-tuning baseline for instruction tuning. However, we agree that an important direction of future work is extending these results to other architectures.

Regarding our evaluation setup and evaluating our proposed methods on more complex auto-regressive generation tasks:

We clarify that we chose these tasks to fairly compare with T0 by mirroring its exact fine-tuning and evaluation setup, thereby showcasing our method's relative efficacy against a key benchmark in instruction-based tuning. These tasks are the comprehensive set of prompt instructions from the (P3), so we also avoid any bias introduced by selecting a subset that favors our method. We follow the same procedure as [4] where each task is converted into the format provided by the templates [1]

However, to validate our methods’ efficacy in an evaluation setup that includes auto-regressive generation, we conduct an additional evaluation during the rebuttal period on a summarization dataset, namely SAMSum [3]. Below, we compare our MoV and MoLoRA with PEFT baselines ($(IA)^3$, LoRA) and full fine-tuning using the T5-3B base model. Note that this is not a held-out task, given that its training split is included in the P3 dataset. However, we believe it may be a valuable data point for the generalization of our methods in-distribution. We calculate the rouge scores for each model.

The below results confirm that MoV and MoLoRA achieve very competitive results with full fine-tuning in a more complex generative task, similar to the trend we showed in the paper’s main evaluation setting.

Regarding further exploration of routing input, particularly with the use of representation derived from the last encoder layer of T5:

In this work, we limit our exploration in terms of routers’ input by only comparing pre-computed sentence embeddings and token representation. We are heartened that R EZ7y acknowledges we have already conducted extensive experiments (doubling the number of variants) to compare token routing to sentence routing. Our results show token routing achieves better performance since routers can exploit similarities between tokens for expert decisions, offering a better generalization. This is at least a preliminary signal that a higher degree of inductive bias for task datasets (the representation derived from the last encoder layer would take this even further) is not necessarily beneficial, as one can acquire a diverse set of task knowledge directly from the hidden representations of tokens. We leave further exploration to future work.

[1] Sanh, Victor, et al. “Multitask Prompted Training Enables Zero-Shot Task Generalization” arXiv preprint arXiv:2110.08207 (2022).

[2] Liu, Haokun, et al. "Few-shot parameter-efficient fine-tuning is better and cheaper than in-context learning." Advances in Neural Information Processing Systems 35 (2022): 1950-1965.

[3] Gliwa, Bogdan, et al. "SAMSum corpus: A human-annotated dialogue dataset for abstractive summarization." arXiv preprint arXiv:1911.12237 (2019).

[4] Colin, Raffel, et al. “Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer.” arXiv preprint arXiv:1910.10683 (2020).

We would like to thank R s33v for their positive feedback and for highlighting the “in-depth ablation study” that shows the “capabilities and scalabilities'' of our proposed method of combining PEFT with MoEs. We also express gratitude towards them for emphasizing positively that our approach, “leverage the advantages and avoid the disadvantages of the MoE structure when combined with PEFT methods.” We address some of the individual concerns below:

MoE takes advantage of the inherent heterogeneity of the data, allowing different parameters to handle different distributions in the data. However, fine-tuning usually focuses on one or a few downstream tasks. In this case, the motivation for using ten or even dozens of experts for learning requires further justification.

R s33v is indeed correct that if there were only one or a few downstream tasks, we wouldn’t use as many experts. However, we are interested in the multi-task fine-tuning setting that has led to recent breakthroughs in NLP [1], where the fine-tuning dataset is deliberately structured to represent a diverse and large set of tasks, with a high degree of heterogeneity in the data. Specifically in our experimental setup, the P3 dataset we use for fine-tuning contains 62 datasets, each containing data points with 5 to 20 unique prompt templates, where there are 8 distinct tasks such as paraphrasing, QA closed book, QA extractive, QA multiple-choice, sentiment, summarization, topic classification, and word-sense disambiguation. Hence, multiple experts are more appropriate here to gain from the benefits of modular specialization given the diversity in the number of datasets, prompts, and tasks.

The article mentioned that the larger the batch size, the easier it is for MoE to collapse to an expert. However, the difference between fine-tuning and pretrain is that the model has entered a relatively stable and better-performing landscape before training begins. Why under such conditions, the more common gradient direction brought by a larger batch size will still cause a certain expert to dominate the gradient descent direction of the parameter space?

We clarify to R s33v that our setting differs from “traditional” finetuning where all weights have already been calibrated during pretraining. Even though the model is indeed pre-trained, both the router layers and PEFT experts are randomly initialized before finetuning which increases instability during finetuning relative to updating all pre-trained weights.

Our finding that larger batch sizes lead to the learning collapse in MoE is consistent with other findings that MoE models benefit from smaller batch sizes due to the instability of training and fine-tuning MoE-style architectures using large batch sizes [2, 3]. We highlight this finding to contribute to the understanding of limitations for MoE approaches.

[1] Sanh, Victor, et al. “Multitask Prompted Training Enables Zero-Shot Task Generalization” arXiv preprint arXiv:2110.08207 (2022). [2] Zoph, Barret, et al. “ST-MoE: Designing Stable and Transferable Sparse Expert Models.” arXiv preprint arXiv:2202.08906 (2022). [3] Shen, Sheng, et al. "Mixture-of-experts meets instruction tuning: A winning combination for large language models." arXiv preprint arXiv:2305.14705 (2023).