Paper: Soft Merging of Experts with Adaptive Routing (ICLR-2024)

Reviewer comment (Experiments):
The main weakness is the small scale of the data and models studied in the paper. I believe the challenge of reducing computational cost with mixture-of-expert models is more relevant to larger models. The authors however only presented results on ResNet and T5-base (with only 200M parameters). Experiment results with larger models are appreciated.

Author response:
While we agree that past MoE work has primarily focused on reducing computational costs of large models, our work instead focuses on the benefits of modularity in MoEs with the performance advantages they offer. This includes exploiting similar experts for examples from similar tasks to foster positive transfer and reduce negative interference among tasks. We specifically chose multitask learning as our benchmark because the MoE framework is intuitively appropriate for this setting. It's crucial to mention that each experiment, such as those with T5-GLUE, required approximately 100 hours on a single A6000 GPU, and the ResNet-DomainNet experiments took about 11 hours. This is because these benchmarks contain a set of datasets (eight in GLUE and six in DomainNet), making them quite extensive. Given these considerations, we opted for medium-sized models for our experiments. While we recognize the value of testing larger models, the current scope and our computational resources guided our choice of model scale.

Reviewer comment (Experiments):
If experiments with larger models are not feasible, I hope authors can discuss potential limitations of the study under those larger-scale scenarios. Do you expect the findings in Figure 2 change in larger scale setups?

Author response:
Our method is designed to be agnostic to the backbone architecture or model size. We incorporate MoE blocks into the existing architecture and focus on learning the router and experts within these blocks. Consequently, we anticipate that our method should effectively scale to larger-sized models. As for scaling the number of experts, as detailed in Section 4.1 of our paper, we have found that SMEAR does scale well. In the ResNet-DomainNet setting, it successfully narrows the gap to Ensemble routing. In contrast, in the T5-GLUE setting, the results show no significant difference, which might be attributed to performance saturation. If the reviewer has any specific large-scale experiment in mind, we would be willing to run them, resources permitting.

Reviewer comment (Evaluation):
The author's hypothesis about the shortcomings of learned routings in existing works that "stem from the gradient estimation techniques used to train modular models that use non-differentiable discrete routing decisions" is not supported with evidence other than the performance of final models in Figure 2. If you visualize the learned routing of these baselines (as in Figure 3), will you notice flaws in these learned routing? Do they yield degenerated routing, or routing that is uniform across all the experts?

Author response:
Past work on gradient estimators such as  [https://arxiv.org/abs/1711.00123], [https://openreview.net/forum?id=r1lgTGL5DE, https://arxiv.org/abs/2109.11817] supports our hypothesis regarding the challenges associated with learning using gradient estimator techniques. In our paper, we adopt an empirical approach to compare learned routing methods, and therefore use final model performance as the metric for evaluating routing efficacy. Additionally, we have provided a qualitative analysis of learned routing methods in Appendix G. Here's a summary of their characteristics: In ResNet-DomainNet, Top-k demonstrates uniform routing in the initial layers but chooses a single expert in the last layer.  REINFORCE, ST-Gumbel, and Dselect-k tend to mostly exhibit degenerate single-expert routing policies. Interestingly, all these gradient estimators learn to assign a distinct expert for the Quickdraw dataset. However, this degree of specialization is insufficient for achieving superior performance scores. In T5-GLUE, these estimators display degenerate routing in some layers and a tendency to share a few experts (approximately 3 out of 8) in other layers across tasks. Methods like Latent Skills, Ensemble, and SMEAR, which utilize most experts in the MoE layer, generally outperform these estimators. It is important to note that Latent Skills requires task metadata, and Ensemble routing incurs significant additional computational costs. SMEAR not only performs better than Latent Skills but also approaches the effectiveness of Ensemble, all without needing task metadata and while maintaining a computational cost nearly equivalent to that of discrete routing methods. We incorporated this summary into the main text.

Reviewer comment (Evaluation):
How does the routing learned by reported baselines look like?

Author response:
Past work on gradient estimators such as  [https://arxiv.org/abs/1711.00123], [https://openreview.net/forum?id=r1lgTGL5DE, https://arxiv.org/abs/2109.11817] supports our hypothesis regarding the challenges associated with learning using gradient estimator techniques. In our paper, we adopt an empirical approach to compare learned routing methods, and therefore use final model performance as the metric for evaluating routing efficacy. Additionally, we have provided a qualitative analysis of learned routing methods in Appendix G. Here's a summary of their characteristics: In ResNet-DomainNet, Top-k demonstrates uniform routing in the initial layers but chooses a single expert in the last layer.  REINFORCE, ST-Gumbel, and Dselect-k tend to mostly exhibit degenerate single-expert routing policies. Interestingly, all these gradient estimators learn to assign a distinct expert for the Quickdraw dataset. However, this degree of specialization is insufficient for achieving superior performance scores. In T5-GLUE, these estimators display degenerate routing in some layers and a tendency to share a few experts (approximately 3 out of 8) in other layers across tasks. Methods like Latent Skills, Ensemble, and SMEAR, which utilize most experts in the MoE layer, generally outperform these estimators. It is important to note that Latent Skills requires task metadata, and Ensemble routing incurs significant additional computational costs. SMEAR not only performs better than Latent Skills but also approaches the effectiveness of Ensemble, all without needing task metadata and while maintaining a computational cost nearly equivalent to that of discrete routing methods. We incorporated this summary into the main text.

Reviewer comment (Presentation):
The appendix and full text does not have to be in separate pdfs.

Author response:
Thank you for your suggestion. We have now combined them into a single PDF for improved readability and accessibility.

Reviewer comment (Reproducibility):
The implementation details of the proposed method are not clear enough. For example, what routing function $R(\cdot)$ do you use? And what is the specific form of the objective function in the training process?

Author response:
For the routing function R(.), we follow a standard approach of using a single linear layer which is a weight matrix of shape d x N, where d is the model dimension and N is the number of experts in the MoE layer. It takes in as input the average hidden states of tokens at the MoE layer and generates a probability distribution over N experts by applying softmax to the outputs. We included these details in our updated draft for greater clarity.  

Regarding training objectives, there are no additional loss functions added for the router. We use the language modeling loss for T5-GLUE and cross-entropy classification loss for ResNet-DomainNet and train expert and router parameters end-to-end with standard backpropagation. The benefit of our method being fully differentiable with these loss objectives is that it allows the router to be trained directly without any need for approximate gradient estimators or custom losses.

Reviewer comment (Novelty):
This paper lacks sufficient novelty. Can you explain what are the differences and advantages of the proposed method, compared to the $\pi$-Tuning method proposed in [1] (Section 2.2 and Section 3.5)?

Author response:
We want to emphasize that our method is simple, eliminating the need for gradient estimation and it outperforms previous routing methods. Moreover, our approach involves per-example adaptive routing without requiring any task meta-data, which is often difficult to obtain in real-world scenarios. In contrast,  $\pi$-Tuning employs a set of existing task specific experts, retrieving the top k experts for a downstream task and learns to interpolate among these experts for the downstream task. In our case, firstly, we do not assume access to a pool of task specific experts and secondly, we do not have access to task metadata. While  $\pi$-Tuning enables transfer learning to a new downstream task by learning to interpolate, our focus is on developing a routing algorithm among the combined dataset of tasks without distinct task boundaries and achieve positive transfer among existing tasks and reduce negative transfer which can be seen by the improved performance of our method. We added a more explicit discussion of $\pi$-Tuning in our updated draft.

Reviewer comment (Reproducibility):
What routing function $R(\cdot)$ do you use?

Author response:
For the routing function R(.), we follow a standard approach of using a single linear layer which is a weight matrix of shape d x N, where d is the model dimension and N is the number of experts in the MoE layer. It takes in as input the average hidden states of tokens at the MoE layer and generates a probability distribution over N experts by applying softmax to the outputs. We included these details in our updated draft for greater clarity.

Reviewer comment (Reproducibility):
What is the specific form of the objective function in the training process?

Author response:
Regarding training objectives, there are no additional loss functions added for the router. We use the language modeling loss for T5-GLUE and cross-entropy classification loss for ResNet-DomainNet and train expert and router parameters end-to-end with standard backpropagation. The benefit of our method being fully differentiable with these loss objectives is that it allows the router to be trained directly without any need for approximate gradient estimators or custom losses.

Reviewer comment (Theory):
Many MoE methods often see instabilities in training - mode collapse i.e. an expert being chosen much more, non-optimal solutions etc. The presented method is a weighted average of parameters. How does such a method regularise against particularly known failure modes of MoEs which often require explicit regularisation such as through load, importance, entropy, or Mutual Information?

Author response:
> The presented method is a weighted average of parameters. How does such a method regularise against particularly known failure modes of MoEs which often require explicit regularisation such as through load, importance, entropy, or Mutual Information? How would you go about encouraging exploration across experts in SMEAR to help determine which expert weighting is most beneficial per token/image etc.? In a discrete scenario, one might look at dropout and/or jitter noise on incoming representations?

Our method optimizes the routing policy using the loss function's gradient, focusing on minimizing loss. All past load balancing techniques are equally applicable to SMEAR, but they are inappropriate in the experimental context we consider because dataset sizes in the multitask mixture vary by multiple orders of magnitude and the load should therefore not necessarily be balanced. As concrete evidence, note that tag routing is an extremely unbalanced routing strategy but nevertheless has higher performance than most gradient estimators we considered. So, to promote exploration, we implemented expert dropout, randomly removing some experts in the MoE block and averaging the weights of the remaining ones. This strategy proved beneficial in specific contexts, notably in the T5-GLUE setting, as evidenced in our ablation study (Table 1). In contrast, we found that jitter noise did not offer additional benefits. In the case of other routing methods, we use expert dropout regularizer if it was found to improve performance (Table 1).

Reviewer comment (Experiments):
How would you go about encouraging exploration across experts in SMEAR to help determine which expert weighting is most beneficial per token/image etc.? In a discrete scenario, one might look at dropout and/or jitter noise on incoming representations? Did the authors try something similar to investigate performance?

Author response:
> The presented method is a weighted average of parameters. How does such a method regularise against particularly known failure modes of MoEs which often require explicit regularisation such as through load, importance, entropy, or Mutual Information? How would you go about encouraging exploration across experts in SMEAR to help determine which expert weighting is most beneficial per token/image etc.? In a discrete scenario, one might look at dropout and/or jitter noise on incoming representations?

Our method optimizes the routing policy using the loss function's gradient, focusing on minimizing loss. All past load balancing techniques are equally applicable to SMEAR, but they are inappropriate in the experimental context we consider because dataset sizes in the multitask mixture vary by multiple orders of magnitude and the load should therefore not necessarily be balanced. As concrete evidence, note that tag routing is an extremely unbalanced routing strategy but nevertheless has higher performance than most gradient estimators we considered. So, to promote exploration, we implemented expert dropout, randomly removing some experts in the MoE block and averaging the weights of the remaining ones. This strategy proved beneficial in specific contexts, notably in the T5-GLUE setting, as evidenced in our ablation study (Table 1). In contrast, we found that jitter noise did not offer additional benefits. In the case of other routing methods, we use expert dropout regularizer if it was found to improve performance (Table 1).

Reviewer comment (Experiments):
The computational cost of the averaging works _"as long as there is a modest number of expert"_. In papers such as Switch Transformer, they use 128 experts compared to the 16 used in Scaling experiments. Would this method work in this context?

Author response:
>The computational cost of the averaging works "as long as there is a modest number of expert". In papers such as Switch Transformer, they use 128 experts compared to the 16 used in Scaling experiments. Would this method work in this context?

Our reference to a “modest number of experts” is based on the computational ratio for discrete routing versus SMEAR, which is calculated as (1 + N x 2 / L x 4). Here, N represents the number of experts and L the number of tokens per example. As long as the second term of this ratio remains less than 1, the computational requirements are close. In our experiments with T5-GLUE, where L is 128 and N is approximately 8, the runtime difference was minimal. Moreover, tensor averaging typically benefits from GPU optimization. Therefore, we believe scaling up to N ~ 100 is feasible if the sequence length is around 1000, a scenario similar to what is observed in Switch models. This approach should maintain a computational cost comparable to discrete routing. We added this example to the paper for greater clarity.

Reviewer comment (Theory):
Can you foresee any pitfalls from where weight averaging might not be beneficial? What are the main limitations of the method?

Author response:
>Can you foresee any pitfalls from where weight averaging might not be beneficial? What are the main limitations of the method?

While we recognize that SMEAR shows improved performance compared to other routing methods, a performance gap w.r.t. Ensemble routing still exists in the ResNet-DomainNet setting. Additionally, our approach necessitates that the experts be homogeneous, since the method of weight averaging works only in this case. Exploring better merging techniques to resolve this gap and allow for heterogeneous expert architectures would be exciting directions for future research.

Reviewer comment (Evaluation):
How much specialisation do you learn as one gets deeper into the network? Can you control this and would it get better results if you learned? There is a lot of learned sparsity in Figure 3a but is this necessarily good? It could be a limitation that not enough exploration occurs through the proposed learning scheme.

Author response:
>How much specialisation do you learn as get one gets deeper into the network? Can you control this and would it get better results if you learned? There is a lot of learned sparsity in Figure3a but is this necessarily good? It could be a limitation that not enough exploration occurs through the proposed learning scheme.

While we did not compute a quantitative measure of specialization as one progresses through the network's layers, our qualitative observations offer some insights. In the ResNet-DomainNet setting, we observed no significant difference in specialization across layers. Similarly, in the T5-GLUE setting, we did not find any obvious patterns in terms of layer-specific specialization.
Regarding sparsity, a possible advantage is that it allows experts to specialize to certain tasks, but it is not obvious a priori that sparsity is valuable. Our method does not explicitly constrain sparsity and therefore can trade-off sparsity as appropriate. In the T5-GLUE experiments, we observed considerable sparsity, while in ResNet-DomainNet, despite the presence of specialization, sparsity was less pronounced. We believe exploring how to effectively regulate this specialization could be a promising direction for future research, potentially leading to improved routing algorithms.

Reviewer comment (Presentation):
Formatting – not a question, but a comment on legibility. Section 4.1 needs to be broken up into paragraphs.

Author response:
> Formatting - not a question. but a comment on legibility. Section 4.1. needs to be broken up into paragraphs.

Thanks for the comment. We have updated the draft accordingly.