Paper: Mixture-of-Supernets: Improving Weight-Sharing Supernet Training with Architecture-Routed Mixture-of-Experts (ICLR-2024)

Reviewer comment (Novelty):
**[major]**: As far as I can tell, many of the main contributions are already published in [AutoMoE: Heterogeneous Mixture-of-Experts with Adaptive Computation for Efficient Neural Machine Translation](https://aclanthology.org/2023.findings-acl.580) (Jawahar et al., Findings 2023) for machine translation and the presented work would need to outline how it differs from previous work besides applying it for different tasks.

Author response:
Compared to AutoMoE, our work, Mixture-of-Supernets, have entirely different goals, applications of mixture-of-experts, router specifications, as explained below:
  1. **Goals** Given a search space of dense and mixture-of-expert models, the goal of the AutoMoE framework is to search for high-performing model architectures that satisfy user-defined efficiency constraints. The final architectures can be dense or mixture-of-expert models. On the other hand, given a search space of dense models only, the goal of the Mixture-of-Supernets framework is to search for high-performing dense model architectures that satisfy user-defined efficiency constraints. The final architecture can be a dense model only.
  2. **Applications of mixture-of-experts** The main application of mixture-of-experts idea by the AutoMoE framework is to augment the standard NAS search space of dense models with mixture-of-experts models. To this end, the AutoMoE framework modifies the standard weight sharing supernet to support weight generation for mixture-of-expert models. On the other hand, the Mixture-of-Supernets framework uses the mixture-of-expert design to: (i) increase the capacity of standard weight sharing supernet and (ii) customize weights for each architecture. Post training, the expert weights are collapsed to create a single weight for the dense architecture.
  3. **Router specifications** The router underlying the AutoMoE framework takes token embedding as input, outputs a probability distribution over experts, and passes token embedding to top-k experts. On the other, the router underlying the Mixture-of-Supernets framework takes architecture embedding as input, outputs a probability distribution over experts (layer-wise MoS) / neurons (neuron-wise MoS), uses the probability distribution to combine ALL the expert weights into a single weight, and passes token embedding to the single weight.

Reviewer comment (Experiments):
**[major]**: The machine translation experiments are missing a compute-matched dense and mixture of experts baseline trained from scratch without NAS to make results comparable.

Author response:
**The machine translation experiments are missing a compute-matched dense and mixture of experts baseline trained from scratch without NAS to make results comparable.** For compute-matched mixture of expert baseline, prior work such as AutoMoE has already shown “training from scratch without NAS” baseline (e.g., random search) clearly underperforms NAS baselines such as HAT (e.g., on the WMT’14 En-De task, random search underperforms HAT by 0.9 BLEU points for the same latency constraint of 600ms).

Reviewer comment (Theory):
The decision to use MoE as a method for weight disentanglement is not fully justified. A better explanation and introduction is necessary.

Author response:
Thanks for the suggestion. MoE is an established modeling technique to increase capacity of a model and specialize weights to handle a particular set of inputs well. Our method extends this general notion such that the capacity of supernet can be increased and supernets can also specialize weights to handle a particular set of architectures well. As discussed in Section 3.1, MoE can increase the capacity of supernet and naturally does not alter the subnetwork architecture compared to other ways such as adding hidden layers or hidden neurons, which is desirable for our supernet formulation. We will add this text in the revision.

Reviewer comment (Theory):
In the current MoE formulation the weights for a given architecture are formed of the linear combination of m (2) expert. The combination of experts is handled by the router. Given that all operations appear to be linear, can the router not directly generate the weights?

Author response:
**In the current MoE formulation the weights for a given architecture are formed of the linear combination of m (2) expert. The combination of experts is handled by the router. Given that all operations appear to be linear, can the router not directly generate the weights?** If a router is tasked to generate the weights, the number of learnable parameters increases by a large magnitude, which can significantly increase the optimization difficulty. For instance, if we need to construct a weight matrix of 768 x 768, the architecture embedding size is 10, the router hidden dimension is 64, the router contains two projection matrices (10 x 64) and (64 x (768 x 768)), which amounts to 37M learnable parameters for a single weight matrix only. On the other hand, if the number of experts is 2, the router underlying layer-wise mixture-of-supernets contains two projection matrices (10 x 64) and (64 x 2), which amounts to 768 parameters only.

Reviewer comment (Experiments):
The use of the experts allows the amount of weight sharing between candidate linear operations to be tempered. One architecture may use one set of weights and another can use a completely separate set, or a linear combination of the expert weights. Setting the weights to be completely decoupled (disentangled) would revert the approach to the traditional weight sharing approaches such as in ENAS. Please could the authors explain why a completely disentangled ablation is not performed (m=number of layers). Also can the authors provide an ablation for m=1, i.e. weights fully entangled?

Author response:
**could the authors explain why a completely disentangled ablation is not performed (m=number of layers).** 'm' is the number of experts and is specific to a weight matrix in the original model. For instance, if m=2, our methods represent each weight matrix in the model using two expert weights. We want to understand why m=number of layers is suggested. Can you please provide more context to your question?  
**can the authors provde an ablation for m=1, i.e. weights fully entangled?** When m=1 (dense supernet model), mixture-of-supernets generalize to traditional weight sharing supernet (e.g., HAT in MT), which has sub-optimal performance and high retraining time (e.g., mixture-of-supernets reduce the retraining cost by 60% for WMT'14 En-De task as shown in Table 6).

Reviewer comment (Experiments):
The focus of the paper is on removing the need for from-scratch training of the final architecture. However the purpose of training the supernet is generally to find the optimal architecture, not to find optimal weights. Would stopping the supernet training earlier save more time?

Author response:
**The focus of the paper is on removing the need for from-scratch training of the final architecture. However the purpose of training the supernet is generally to find the optimal architecture, not to find optimal weights. Would stopping the supernet training earlier save more time?** Stopping the supernet training earlier can lead to unreliable performance estimates for search and high retraining time for the final architecture. If we've misunderstood your question, can you please clarify?

Reviewer comment (Writing):
Explanation of neuron-wise MoS needs a better explanation. Elements of its explanation should be in the earlier section on layer-wise MoS, such as the introduction of m value as the number of experts. Making Fig 1 less busy would help with this.

Author response:
**Explanation of neuron-wise MoS needs a better explanation. Elements of its explanation should be in the earlier section on layer-wise MoS, such as the introduction of m value as the number of experts....** In the revised pdf (text in red color), we have added text that defines 'm' in Section 3.2 and connected each subsection in Section 3. Let us know if the new text improves the clarity of the presentation of neuron-wise MoS method.

Reviewer comment (Evaluation):
Can the authors explain why the final architecture performs better for having being trained via their method? Is it because the architecture found is better or because the final model's performance is improved by having a better weight initialisation?

Author response:
**Can the authors explain why the final architecture performs better for having being trained via their method? Is it because the architecture found is better or because the final model's performance is improved by having a better weight initialisation?** Our mixture-of-supernet formulation produces a high-performing supernet thanks to the increase in the supernet’s capacity and reduction in gradient conflict (see A.1.2 for empirical analysis of gradient conflict). Such supernet can enhance the quality and the efficiency of final architecture identified by the search algorithm (e.g., 1.6, 0.2, 10.6, 1.0, 0.2 avg. GLUE points improvements for BERT models of 5M, 10M, 26.8M, 30M, 50M parameters respectively (see Table 3), at least 1 BLEU point improvements (except for WMT’14 En-Fr, 100ms) across different latency constraints and WMT datasets for MT). Also, such a supernet can provide better weight initialization for obtaining the optimal weights for the final architecture and reduce the retraining cost significantly (e.g., the retraining cost of the final model is 0 for BERT evaluation as shown in Table 3 and minimal compared to SOTA NAS framework for MT evaluation as shown in Table 6).

Reviewer comment (Experiments):
The study concentrates on smaller networks, which may not face the challenges of larger models where memory and computational resources are more critical. To fully evaluate the method's applicability, incorporating a broader range of network sizes and model families (e.g., BERT, T5, GPT) would be beneficial for assessing its scalability and generalizability.

Author response:
**Weakness 1** Following the evaluation setup of popular NAS benchmarks for NLP, we explore model sizes (less than 200M parameters) and model families (BERT, MT) due to high computational requirements for supernet training. Our evaluation setups already validate the generalizability of our proposed methods for language model pretraining (BERT) and task-specific training (MT). Extending the evaluation setup to bigger model sizes (e.g., billion parameters) and other model families (T5, GPT) is orthogonal to our main research questions, but will be interesting future directions for our work.

Reviewer comment (Evaluation):
The absence of FLOPs (Floating Point Operations Per Second) as a performance metric is noticeable. Including FLOPs would provide a more complete comparison with other methods. Additionally, the exclusion of the STS-B task, a single regression task in the GLUE benchmark, is not justified; its inclusion could enhance the validation of task generalization.

Author response:
**Weakness 2** Thanks for the suggestions. 
  1. For BERT evaluation, we compare against NAS-BERT and AutoDistil in terms of model size and performance only (see Table 3) due to two reasons: (i) the search algorithm is explicitly optimized to improve the model_size-performance tradeoff only and (ii) NAS-BERT and AutoDistil do not report FLOPs in their paper or provide code to compute the cost. For MT evaluation, we compare against HAT on only latency and performance, as the search algorithm is explicitly optimized to improve the latency-performance tradeoff only. 
  2. In our work, we only include the following GLUE tasks: CoLA, MRPC, SST2, QNLI, QQP, RTE, and MNLI to be consistent with the evaluation setup of NAS-BERT and AutoDistil (as mentioned in Section 4.1). NAS-BERT and AutoDistil do not share their models, for us to obtain their STS-B performance for comparison with our methods.

Reviewer comment (Evaluation):
Although efficiency is purported to be a key advantage of the proposed method, this is not thoroughly discussed in the main body of the paper. The supplementary comparison of memory efficiency between HAT and the proposed method indicates a significant memory efficiency gap. Moreover, the improvements in task performance are relatively marginal and efficiency drop compared to other baselines is significant, suggesting a trade-off that may undermine the method's relative advantage when considering the marginal performance enhancement.

Author response:
**Weakness 3** 
  1. In terms of efficiency, the main advantage of our method is to significantly reduce the resources (e.g., time, steps) for retraining the final architectures for convergence once the search is complete. For BERT evaluation, there is no retraining required for our framework, as shown in Table 2 and 3. For MT evaluation, there is minimal retraining required for our framework, as shown in Table 6. This advantage offsets the increase in the resources (time, memory) for supernet (one-time cost) and search, leading to big savings in the overall time. Due to space constraints, we discuss the breakdown of the overall time savings in the appendix (see Section A.4.10). 
  2. The performance improvements obtained by our methods over baselines in both BERT evaluation and MT evaluation are big (not marginal): e.g., 1.6, 0.2, 10.6, 1.0, 0.2 avg. GLUE points improvements for BERT models of 5M, 10M, 26.8M, 30M, 50M parameters respectively (see Table 3), at least 1 BLEU point improvements (except for WMT’14 En-Fr, 100ms) across different latency constraints and WMT datasets for MT. These strong results are also attested by other reviewers: Reviewer g33E (“strong on language modeling and MT”), Reviewer yeUq (“non-marginal improvement in the final model accuracy”), and Reviewer GCC7 (“BERT and MT models demonstrating significant metric improvements”).

Reviewer comment (Writing):
Paper is hard to follow and took multiple readings to understand the problem setting. Any reader not familiar with NAS area is likely to feel lost. I recommend significantly reworking the problem setup for clarity. This will greatly improve the paper. (Eg. Section 2 is confusing, and extra prose to explain what a supernet is informally and why it is used would be helpful). Connecting each subsection in Section 3 is also needed.

Author response:
* **Weakness 1** In the revised pdf (text in red color), we have added text that introduces supernet informally, usefulness of supernet (see Section 2), and connected each subsection in Section 3. Let us know if the new text improves the clarity.

Reviewer comment (Experiments):
k-shot NAS is extremely similar to the proposed method, and should not be taxing to reimplement. Given that this is an intuitive baseline even if k-shot NAS didn't exist, I believe this, or something similar (iterative vs joint training) is an important baseline. Without this, it is not clear how effective the proposed method is.

Author response:
* **Weakness 2** Neuron-wise MoS has the key difference to k-shot nas, that the smaller subnets are not a subset of larger ones anymore. Also, reimplementing the k-shot NAS methods (originally defined for convolution architectures for image classification tasks) and adapting the training recipe (e.g., identifying supernet settings, search settings) for BERT/MT tasks is non-trivial. k-shot NAS is not an intuitive baseline as the relevant baselines in NAS for NLP uses standard weight-sharing supernet. For baselines, we primarily use the state-of-the-art frameworks in NAS for BERT (e.g., NAS-BERT, AutoDistil) and NAS for machine translation literature (e.g., HAT).

Reviewer comment (Evaluation):
More than training steps, cost (FLOPS and wall clock time) while training both initial model and final model is a more informative metric than what is displayed in Table 3.

Author response:
* **Weakness 3** Thanks for the suggestion. In Table 3, we compare against NAS-BERT and AutoDistil, both of which do not report FLOPs and wall clock time in their paper or provide code to compute the cost. Nevertheless, our work focuses on reducing the retraining cost of the final model. For our methods, the retraining cost of the final model is 0 for BERT evaluation (as shown in Table 3) and minimal compared to SOTA NAS framework for MT evaluation (as shown in Table 6, which includes wall clock time).