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

Abstract: Weight-sharing supernets are crucial for performance estimation in cutting-edge neural architecture search (NAS) frameworks. Despite their ability to generate diverse subnetworks without retraining, the quality of these subnetworks is not guaranteed due to weight sharing. In NLP tasks like machine translation and pre-trained language modeling, there is a significant performance gap between supernet and training from scratch for the same model architecture, necessitating retraining post optimal architecture identification. This study introduces a solution called mixtureof-supernets, a generalized supernet formulation leveraging mixture-of-experts (MoE) to enhance supernet model expressiveness with minimal training overhead. Unlike conventional supernets, this method employs an architecturebased routing mechanism, enabling indirect sharing of model weights among subnetworks. This customization of weights for specific architectures, learned through gradient descent, minimizes retraining time, significantly enhancing training efficiency in NLP. The proposed method attains state-of-the-art (SoTA) performance in NAS for fast machine translation models, exhibiting a superior latency-BLEU tradeoff compared to HAT, the SoTA NAS framework for machine translation. Furthermore, it excels in NAS for building memoryefficient task-agnostic BERT models, surpassing NAS-BERT and AutoDistil across various model sizes. The code can be found at: https://github.com/UBC-NLP/MoS.

text: Findings of the Association for Computational Linguistics ACL 2024, pages 10424–10443 August 11-16, 2024 ©2024 Association for Computational Linguistics
This study introduces a solution called mixtureof-supernets, a generalized supernet formulation leveraging mixture-of-experts (MoE) to enhance supernet model expressiveness with minimal training overhead. Unlike conventional supernets, this method employs an architecturebased routing mechanism, enabling indirect sharing of model weights among subnetworks. This customization of weights for specific architectures, learned through gradient descent, minimizes retraining time, significantly enhancing training efficiency in NLP. The proposed method attains state-of-the-art (SoTA) performance in NAS for fast machine translation models, exhibiting a superior latency-BLEU tradeoff compared to HAT, the SoTA NAS framework for machine translation. Furthermore, it excels in NAS for building memoryefficient task-agnostic BERT models, surpassing NAS-BERT and AutoDistil across various model sizes. The code can be found at: https://github.com/UBC-NLP/MoS.

heading: 1 Introduction
text: Neural architecture search (NAS) automates the design of high-quality architectures for natural language processing (NLP) tasks while meeting specified efficiency constraints (Wang et al., 2020a; Xu
*Some of the work was completed while Ganesh was interning at Meta.
et al., 2021, 2022a). NAS is commonly treated as a black-box optimization (Zoph et al., 2018; Pham et al., 2018), but obtaining the best accuracy requires repetitive training and evaluation, which is impractical for large datasets. To address this, weight sharing is applied via a supernet, where subnetworks represent different model architectures (Pham et al., 2018).
Recent studies demonstrate successful direct use of subnetworks for image classification with performance comparable to training from scratch (Cai et al., 2020; Yu et al., 2020). However, applying this supernet approach to NLP tasks is more challenging, revealing a significant performance gap when using subnetworks directly. This aligns with recent NAS works in NLP (Wang et al., 2020a; Xu et al., 2021), which address the gap by retraining or finetuning the identified architecture candidates. This situation introduces uncertainties about the optimality of selected architectures and requires repeated training for obtaining final accuracy on the Pareto front, i.e., the best models for different efficiency (e.g., model size or inference latency) budgets. This work aims to enhance the weight-sharing mechanism among subnetworks to minimize the observed performance gap in NLP tasks.
The weight-sharing supernet is trained by iteratively sampling architectures from the search space and training their specific weights from the supernet. Standard weight-sharing (Yu et al., 2020; Cai et al., 2020) involves directly extracting the first few output neurons to create a smaller subnetwork (see Figure 1 (a)), posing two challenges due to limited model capacity. First, the supernet imposes strict weight sharing among architectures, causing co-adaptation (Bender et al., 2018; Zhao et al., 2021) and gradient conflicts (Gong et al., 2021). For example, in standard weight-sharing, if a 5M-parameters model is a subnetwork of a 90Mparameters model, 5M weights are directly shared. However, the optimal shared weights for the 5M
10424
model may not be optimal for the 90M model, leading to significant gradient conflicts during optimization (Gong et al., 2021). Second, the supernet’s overall capacity is constrained by the parameters of a single deep neural network (DNN), i.e., the largest subnetwork in the search space. However, when dealing with a potentially vast number of subnetworks (e.g., billions), relying on a single set of weights to parameterize all of them could be insufficient (Zhao et al., 2021).
To address these challenges, we propose a Mixture-of-Supernets (MoS) framework. MoS enables architecture-specific weight extraction, allowing smaller architectures to avoid sharing some output neurons with larger ones. Additionally, it allocates large capacity without being constrained by the number of parameters in a single DNN. MoS includes two variants: layer-wise MoS, where architecture-specific weight matrices are constructed based on weighted combinations of expert weight matrices at the level of sets of neurons, and neuron-wise MoS, which operates at the level of individual neurons in each expert weight matrix. Our proposed NAS method proves effective in constructing efficient task-agnostic BERT models (Devlin et al., 2019) and machine translation
(MT) models. For efficient BERT, our best supernet outperforms SuperShaper (Ganesan et al., 2021) by 0.85 GLUE points, surpasses NAS-BERT (Xu et al., 2021) and AutoDistil (Xu et al., 2022a) in various model sizes (≤ 50M parameters). Compared to HAT (Wang et al., 2020a), our top supernet reduces the supernet vs. standalone model gap by 26.5%, provides a superior pareto front for latencyBLEU tradeoff (100 to 200 ms), and decreases the steps needed to close the gap by 39.8%. A summary in the Table 1 illustrates the time savings and BLEU improvements of MoS supernets for the WMT’14 En-De task.
We summarize our key contributions:
1. We propose a formulation that generalizes weight-sharing methods, encompassing direct weight sharing (e.g., once-for-all network (Cai et al., 2020), BigNAS (Yu et al., 2020)) and flexible weight sharing (e.g., fewshot NAS (Zhao et al., 2021)). This formulation enhances the expressive power of the supernet.
2. We apply the MoE concept to enhance model capability. The model’s weights are dynamically generated based on the activated subnetwork architecture. Post-training, the MoE
can be converted into equivalent static models as our supernets solely depend on the fixed subnetwork architecture after training. (3)
3. Our experiments show that our supernets achieve SoTA NAS results in building efficient task-agnostic BERT and MT models.

heading: 2 Supernet - Fundamentals
text: A supernet, utilizing weight sharing, parameterizes weights for millions of architectures, offering rapid performance predictions and significantly reducing NAS search costs. The training objective can be formalized as follows. LetXtr denote the training data distribution. Let x, y denote the training sample and label respectively, i.e., x, y ∼ Xtr. Let arand denote an architecture uniformly sampled from the search spaceA. Let fa denote the subnetwork with architecture a, and f be parameterized by the supernet model weights W . Then, the training objective of the supernet can be given by,
min W
Ex,y∼XtrEarand∼A[L(farand(x;W ), y)]. (1)
The mentioned formulation is termed single path one-shot (SPOS) optimization (Guo et al., 2020) for supernet training. Another popular technique is sandwich training (Yu et al., 2020), where the largest (abig), smallest (asmall), and uniformly sampled architectures (arand) from the search space are jointly optimized.

heading: 3 Mixture-of-Supernets
text: Existing supernets typically have limited model capacity to extract architecture-specific weights. For simplicity, assume the model function fa(x;W ) is a fully connected layer (output o = Wx, omitting bias term for brevity), where x ∈ nin × 1, W ∈ nout × nin, and o ∈ nout × 1. nin and nout correspond to the number of input and output features respectively. Then, the weights (Wa ∈ nouta × nin) specific to architecture a with nouta output features are typically extracted by taking the first nouta rows
1 (as shown in Figure 1 (a)) from the supernet weight W . Assume one samples two architectures (a and b) from the search space with the number of output features nouta and noutb respectively. Then, the weights corresponding to the architecture with the smallest number of output features will be a subset of those of the
1We assume the number of input features remains constant. If it changes, only the initial columns of Wa are extracted.
other architecture, sharing the first |nouta − noutb | output features exactly. This weight extraction technique enforces strict weight sharing between architectures, irrespective of their global architecture information (e.g., different features in other layers). For example, architectures a and b may have vastly different capacities (e.g., 5M vs 90M parameters). The smaller architecture (e.g., 5M ) must share all weights with the larger one (e.g., 90M ), and the supernet (modeled by fa(x;W )) cannot allocate weights specific to the smaller architecture. Another issue with fa(x;W ) is that the supernet’s overall capacity is constrained by the parameters in the largest subnetwork (W ) in the search space. Yet, these supernet weights W must parameterize numerous diverse subnetworks. This represents a fundamental limitation of the standard weight-sharing mechanism. Section 3.1 proposes a reformulation to overcome this limitation, implemented through two methods (Layer-wise MoS, Section 3.2, Neuron-wise MoS, Section 3.3), suitable for integration into Transformers (see Section 3.4).

heading: 3.1 Generalized Model Function
text: We can reformulate the function fa(x;W ) to a generalized form g(x, a;E), which takes 2 inputs: the input data x, and the activated architecture a. E includes the learnable parameters of g. Then, the training objective of the proposed supernet becomes,
min E Ex,y∼XtrEarand∼A[L(g(x, arand;E), y)]. (2)
For the standard weight sharing mechanism mentioned above, E = W and function g just uses a to perform the “trimming” operation on the weight matrix W , and forwards the subnetwork. To further minimize objective equation 2, enhancing the capacity of the model function g is a potential approach. However, conventional methods like adding hidden layers or neurons are impractical here since the final subnetwork architecture of mapping x to fa(x;W ) cannot be altered. This work introduces the concept of Mixture-of-Experts (MoE) (Fedus et al., 2022) to enhance the capacity of g. Specifically, we dynamically generate weights Wa for a specific architecture a by routing to certain weight matrices from a set of expert weights. We term this architecture-routed MoE-based supernet as Mixture-of-Supernets (MoS) and design two routing mechanisms for function g(x, a;E). Due
to lack of space, the detailed algorithm for supernet training and search is shown in A.2.

heading: 3.2 Layer-wise MoS
text: Assume there are m (number of experts) unique weight matrices ({Ei ∈ Rnoutbig×ninbig }mi=1, or expert weights), which are learnable parameters. For simplicity, we only use a single linear layer as the example. For an architecture a with nouta output features, we propose the layer-wise MoS that computes the weights specific to the architecture a (i.e. Wa ∈ Rnouta×nin) by performing a weighted combination of expert weights, Wa = ∑ i α i aE i a. Here, Eia ∈ Rnouta×nin corresponds to the standard top rows extraction from the ith expert weights. The alignment vector (αa ∈ [0, 1]m, ∑ i α i a = 1) captures the alignment scores of the architecture a with respect to each expert (weights matrix). We encode the architecture a as a numeric vector Enc(a) ∈ Rnenc×1 (e.g., a list of the number of output features for different layers), and apply a learnable router r(·) (an MLP with softmax) to produce such scores, i.e. αa = r(Enc(a)). Thus, the generalized model function for the linear layer (as shown in Figure 1 (b)) can be defined as (omitting bias for brevity):
g(x, a;E) = Wax = ∑
i
r(Enc(a))iEiax. (3)
The router r(·) governs the degree of weight sharing between two architectures through modulation of alignment scores (αa). For instance, if m = 2 and a is a subnetwork of architecture b, the supernet can allocate weights specific to the smaller architecture a by setting αa = (1, 0) and αb = (0, 1). In this scenario, g(x, a;E) exclusively utilizes weights from E1, and g(x, b;E) uses weights from E2, enabling updates to E1 and E2 towards the loss from architectures a and b without conflicts. It’s worth noting that few-shot NAS (Zhao et al., 2021) is a special case of our framework when the router r is rule-based. Moreover, g(·) functions as an MoE, enhancing expressive power and reducing the objective equation 2. Once supernet training is done, for a given architecture a, the score αa = r(Enc(a)) can be generated offline. Expert weights collapse, reducing the number of parameters for architecture a to nouta × nina . Layer-wise MoS results in a lower degree of weight sharing between differently sized architectures, as evidenced by a higher Jensen-Shannon distance between their
alignment probability vectors compared to similarly sized architectures. Refer to A.1.1 for more details.

heading: 3.3 Neuron-wise MoS
text: Layer-wise MoS employs a standard MoE setup, where each expert is a linear layer/module. The router determines the combination of experts to use for forwarding the input x based on a. In this setup, the degree of freedom for weight generation is m, and the parameter count grows by m× |W |, with |W | being the parameters in the standard supernet. Therefore, a sufficiently large m is needed for flexibility in subnetwork weight generation, but it also introduces too many parameters into the supernet, making layer-wise MoS challenging to train. To address this, we opt for a smaller granularity of weights to represent each expert, using neurons in DNN as experts. In terms of the weight matrix, neuron-wise MoS represents an individual expert with one row, whereas layer-wise MoS uses an entire weight matrix. For neuron-wise MoS, the router output βa = r(·) ∈ [0, 1]noutbig×m for each layer, and the sum of each row in βa is 1. Similar to layer-wise MoS, we use an MLP to produce the noutbig ×m matrix and apply softmax on each row. We formulate the function g(x, a;E) for neuronwise MoS as
Wa = ∑
i
diag(βia)E i a, (4)
where diag(β) constructs a noutbig × noutbig diagonal matrix by putting β on the diagonal, and βia is the i-th column of βa. Ei is still an noutbig × nin matrix as in layer-wise MoS. Compared to layerwise MoS, neuron-wise MoS offers increased flexibility (m×nouta instead of only m) to manage the degree of weight sharing between different architectures, with the number of parameters remaining proportional to m. Neuron-wise MoS provides finer control over weight sharing between subnetworks. Gradient conflict, computed using cosine similarity between the supernet and smallest subnet gradients following NASVIT (Gong et al., 2021), is lowest for neuron-wise MoS compared to layerwise MoS and HAT, as shown by the highest cosine similarity (see A.1.2).
3.4 Adding g(x, a;E) to Transformer MoS is adaptable to a single linear layer, multiple linear layers, and other parameterized layers (e.g., layer-norm). Given that the linear layer dominates
the number of parameters, we adopt the approach used in most MoE work (Fedus et al., 2022). We apply MoS to the standard weight-sharing-based Transformer (fa(x;W )) by replacing the two linear layers in every feed-forward network block with g(x, a;E).

heading: 4 Experiments - Efficient BERT

heading: 4.1 Experiment Setup
text: We explore the application of our proposed supernet in constructing efficient task-agnostic BERT (Devlin et al., 2019) models, focusing on the BERT pretraining task. This involves pretraining a language model from scratch to learn task-agnostic text representations using a masked language modeling objective. The pretrained BERT model is then finetuned on various downstream NLP tasks. Emphasis is on building highly accurate yet small BERT models (e.g., 5M − 50M parameters). Both BERT supernet and standalone models are pretrained from scratch on Wikipedia and Books Corpus (Zhu et al., 2015). Performance evaluation is conducted by finetuning on seven tasks from the GLUE benchmark (Wang et al., 2018), chosen by AutoDistil (Xu et al., 2022a). The architecture encoding, data preprocessing, pretraining settings, and finetuning settings are detailed in A.4.1. Baseline models include standalone and standard supernet models proposed in SuperShaper (Ganesan et al., 2021). Our proposed models are layer-wise and neuron-wise MoS. All supernets undergo sandwich training 2. Parameters m and router’s hidden dimension are set to 2 and 128, respectively, for MoS supernets.

heading: 4.2 Supernet vs. standalone gap
text: For investigating the supernet vs. standalone gap, the search space is derived from SuperShaper (Ganesan et al., 2021), encompassing BERT architectures differing only in hidden size at each layer (120, 240, 360, 480, 540, 600, 768) with fixed numbers of layers (12) and attention heads (12). This search space includes around 14 billion architectures. We examine the supernet vs. standalone model gap for the top model architecture from the pareto front of Supernet (Sandwich) (Ganesan et al., 2021). Table 2 illustrates the GLUE benchmark performance of standalone
2SuperShaper notes that SPOS performs poorly compared to sandwich training; hence, we do not study SPOS for building BERT models. The learning curve is shown in A.4.2.
training for the architecture (1x pretraining budget, equivalent to 2048 batch size * 125,000 steps) as well as architecture-specific weights from different supernets (0 additional pretraining steps; i.e., only supernet pretraining). MoS (layer-wise or neuron-wise) bridges the gap between task-specific supernet and standalone performance for 6 out of 7 tasks, including MNLI, a widely used task for evaluating pretrained language models (Liu et al., 2019; Xu et al., 2022b). The average GLUE gap between the standalone model and standard supernet is 0.13 points. Remarkably, with customization and expressivity properties, layer-wise and neuron-wise MoS significantly improve standalone training by 0.75 and 0.85 average GLUE points, respectively. 3 Table 2 demonstrates that MoS imposes a computational overhead of under 22% for BERT, resulting in a minimum of 0.8 average GLUE improvement compared to the standard supernet. This overhead may not be significant, as it represents a one-time investment that eliminates the need for additional training after the search process.

heading: 4.3 Comparison with SoTA NAS
text: The SoTA NAS frameworks for constructing a taskagnostic BERT model are NAS-BERT (Xu et al., 2021) and AutoDistil (Xu et al., 2022a).4 The NASBERT pipeline comprises: (1) supernet training (with a Transformer stack containing multi-head attention, feed-forward network [FFN], and convolutional layers at arbitrary positions), (2) search based on the distillation (task-agnostic) loss, and (3) pretraining the best architecture from scratch (1x pretraining budget, equivalent to 2048 batch size * 125,000 steps). The third step needs to be repeated for every constraint change and hardware change, incurring substantial costs. The AutoDistil pipeline involves: (1) constructing K search spaces and training supernets for each search space independently, (2a) agnostic-search mode: searching based on the self-attention distillation (taskagnostic) loss, (2b) proxy-search mode: searching based on the MNLI validation score, and (3) extracting architecture-specific weights from the supernet without additional training. The first step can be costly as pretraining K supernets requires K times training compute and memory, compared
3Consistency of these results across different seeds is discussed in A.4.5.
4AutoDistil (proxy) outperforms SoTA distillation approaches such as TinyBERT (Jiao et al., 2020) and MINILM (Wang et al., 2020b) by 0.7 average GLUE points. Hence, we do not compare against these works.
to training a single supernet. The proxy-search mode may favor AutoDistil unfairly, as it finetunes all architectures in its search space on MNLI and utilizes the MNLI score for ranking. For a fair comparison with SoTA, MNLI task is excluded from evaluation. 5
Our proposed NAS pipeline addresses challenges in NAS-BERT and AutoDistil. In comparison to SoTA NAS, our search space includes BERT architectures with uniform Transformer layers: hidden size (120 to 768 in increments of 12), attention heads (6, 12), intermediate FFN hidden dimension ratio (2, 2.5, 3, 3.5, 4). This search space comprises 550 architectures, similar to AutoDistil. The supernet is based on neuron-wise MoS, and the search uses perplexity (task-agnostic) to rank architectures. Unlike NAS-BERT, our final architecture weights are directly extracted from the supernet without additional pretraining. Unlike
5Refer to A.4.4 for a comparison of neuron-wise MoS against baselines that don’t directly tune on the MNLI task. Neuron-wise MoS consistently outperforms baselines in terms of both average GLUE and MNLI task performance.
AutoDistil, our pipeline pretrains only one supernet, significantly reducing training compute and memory. We use only task-agnostic metrics for search, similar to AutoDistil’s agnostic setting. Table 3 compares neuron-wise MoS supernet with NAS-BERT and AutoDistil for various model sizes. NAS-BERT and AutoDistil performances are obtained from respective papers. On average GLUE, our pipeline improves over NAS-BERT for 5M , 10M , and 30M model sizes, with no additional training (100% additional training compute savings, equivalent to 2048 batch size * 125,000 steps). On average GLUE, our pipeline: (i) surpasses AutoDistil-proxy for 6.88M and 50M model sizes with 1.88M and 0.1M fewer parameters respectively, and (ii) outperforms both AutoDistil-proxy and AutoDistil-agnostic for 26M model size. Besides achieving SoTA results, our method significantly reduces the heavy workload of training multiple models in subnetwork retraining (NAS-BERT) or supernet training (AutoDistil).

heading: 5 Experiments - Efficient MT

heading: 5.1 Experiment setup
text: We discuss the application of proposed supernets for building efficient MT models following the setup by Hardware-aware Transformers (HAT (Wang et al., 2020a)), the SoTA NAS framework for MT models with good latency-BLEU tradeoffs. Focusing on WMT’14 En-De, WMT’14 En-Fr, and WMT’19 En-De benchmarks, we maintain consistent architecture encoding and training settings for supernet and standalone models (details in A.5.2). Baseline supernets include HAT and Supernet (Sandwich). Proposed supernets are Layer-wise MoS and Neuron-wise MoS, both using sandwich training, with m and router’s hidden dimension set to 2 and 128, respectively. Refer to A.5.8 for the rationale behind choosing ‘m’.

heading: 5.2 Supernet vs. standalone gap
text: In HAT’s search space of 6M encoder-decoder architectures, featuring flexible parameters like embedding size (512 or 640), decoder layers (1 to 6), self/cross attention heads (4 or 8), and number of top encoder layers for decoder attention (1 to 3), good supernets should exhibit minimal mean absolute error (MAE) and high rank correlation between supernet and standalone performance for a given architecture. Table 4 presents MAE and Kendall rank correlation for 15 random architectures, showcasing that sandwich training yields better MAE and rank quality compared to HAT. While our proposed supernets achieve comparable rank quality for WMT’14 En-Fr and WMT’19 En-De, and slightly underperform for WMT’14 En-De, they exhibit minimal MAE across all tasks. Particularly, neuron-wise MoS achieves substantial MAE improvements, suggesting lower additional training steps needed to make MAE negligible (as detailed in Section 5.4). Supernet and standalone performance plots reveal neuron-wise MoS excelling for almost all top-performing architectures (see A.5.3). The training overhead for MoS is generally negligible, e.g., for WMT’14 En-De, supernet training takes 248 hours, with neuron-wise MoS and layer-wise MoS requiring 14 and 18 additional hours, respectively (less than 8% overhead, see Section A.5.10 for details).

heading: 5.3 Comparison with the SoTA NAS
text: The pareto front from the supernet is obtained using an evolutionary search algorithm that leverages
the supernet for quickly identifying top-performing candidate architectures and a latency estimator for promptly discarding candidates with latencies surpassing a threshold. Settings for the evolutionary search algorithm and latency estimator are detailed in A.5.4. Three latency thresholds are explored: 100 ms, 150 ms, and 200 ms. Table 5 illustrates the latency vs. supernet performance tradeoff for models in the pareto front from different supernets. Compared to HAT, the proposed supernets consistently achieve significantly higher BLEU for each latency threshold across all datasets, emphasizing the importance of architecture specialization and expressiveness of the supernet. See A.5.6 for the consistency of these trends across different seeds.

heading: 5.4 Additional training to close the gap
text: The proposed supernets significantly minimize the gap between the supernet and standalone MAE (as discussed in Section 5.2), yet the gap remains non-negligible. Closing the gap for an architecture involves extracting architecture-specific weights from the supernet and conducting additional training until the standalone performance is reached (achieving a gap of 0). An effective supernet should demand a minimal number of additional steps and time for the extracted architectures to close the gap. In the context of additional training, we evaluate the test BLEU for each architecture after every 10K steps, stopping when the test BLEU matches or exceeds the test BLEU of the standalone model. Table 6 presents the average number of additional training steps required for all models on the pareto front from each supernet to close the gap. Compared to HAT, layer-wise MoS achieves an impressive reduction of 9% to 51% in training steps, while neuron-wise MoS delivers the most substantial reduction of 21% to 60%. For the WMT’14 EnFr task, both MoS supernets require at least 2.7% more time than HAT to achieve SoTA BLEU across different constraints. These results underscore the importance of architecture specialization and supernet expressivity in significantly improving the training efficiency of subnets extracted from the supernet.

heading: 5.5 Comparison to AutoMoE
text: Although AutoMoE (Jawahar et al., 2023) and MoS pursue distinct objectives (as discussed in Appendix A.3), we proceed to compare the supernet BLEU scores of HAT, AutoMoE, and MoS under a latency constraint of 200 ms on the NVIDIA
V100 GPU across the three WMT benchmarks. Table 7 shows that MoS consistently outperforms AutoMoE and HAT across all datasets. Interestingly, AutoMoE falls behind HAT, suggesting a potential discrepancy between the performance of AutoMoE’s supernet and standalone models.

heading: 6 Related Work
text: In this section, we briefly review existing research on NAS in NLP. Evolved Transformer (ET) (So et al., 2019) is an initial work that explores NAS for efficient MT models. It employs evolutionary search and dynamically allocating training resources for promising candidates., HAT (Wang et al., 2020a) introduces a weight-sharing supernet as a performance estimator, amortizing the training cost for candidate MT evaluations in evolutionary search. NAS-BERT (Xu et al., 2021) partitions the BERT-Base model into blocks and trains a weight-sharing supernet to distill each block. NASBERT uses progressive shrinking during supernet training to prune less promising candidates, identifying top architectures for each efficiency constraint quickly. However, NAS-BERT requires pretraining the top architecture from scratch for every constraint change, incurring high costs. SuperShaper (Ganesan et al., 2021) pretrains a weightsharing supernet for BERT using a masked language modeling objective with sandwich training. AutoDistil (Xu et al., 2022a) employs few-shot NAS (Zhao et al., 2021): constructing K search spaces of non-overlapping BERT architectures and
training a weight-sharing BERT supernet for each search space. The search is based on self-attention distillation loss with BERT-Base (task-agnostic search) and MNLI score (proxy search). AutoMoE (Jawahar et al., 2023) augments the search space of HAT with mixture-of-expert models to design efficient translation models. Refer to A.3 for the main differences between our framework and the AutoMoE framework.
In the computer vision community, K-shot NAS (Su et al., 2021) generates weights for each subnet as a convex combination of different supernet weights in a dictionary using a simplex code. While K-shot NAS shares similarities with layerwise MoS, there are key distinctions. K-shot NAS trains the architecture code generator and supernet iteratively due to training difficulty, whereas layerwise MoS trains all its components jointly. K-shot NAS has been applied specifically to convolutional architectures for image classification tasks. However, it introduces too many parameters with an increase in the number of supernets (K), a concern alleviated by neuron-wise MoS due to its granular weight specialization. In our work, we focus on NLP tasks and relevant baselines, noting that supernets in NLP tend to lag significantly behind standalone models in terms of performance. Additionally, the authors of K-shot NAS have not released the code to reproduce their results, preventing a direct evaluation against their method.

heading: 7 Conclusion
text: We introduced Mixture-of-Supernets, a formulation aimed at enhancing the expressive power of supernets. By adopting the idea of MoE, we demonstrated the ability to generate flexible weights for subnetworks. Through extensive evaluations for constructing efficient BERT and MT models, our supernets showcased the capacity to: (i) minimize retraining time, thereby significantly improving NAS efficiency, and (ii) produce high-quality architectures that meet user-defined constraints.

heading: 8 Limitations
text: The limitations of this work are as follows:
1. Applying Mixture-of-Supernet (MoS) to popular benchmarks in NLP, focusing on efficient machine translation and BERT, offers valuable insights. A potential impactful future direction could involve extending the application of MoS to build efficient autoregressive decoder-only language models, such as GPT4 (OpenAI, 2023).
2. Introducing MoE architecture potentially need more training budget. In our work, we do not use large number of training iteration for fair comparison and fixing the number of expert weights (m) to 2 works well. We will investigate the full potential of the proposed supernets by combining larger training budget (e.g., ≥ 200K steps) and larger number of expert weights (e.g., ≥ 16 expert weights) in the future work.
3. Due to the high computational requirements for pretraining BERT, we only investigate the gap between the supernet and standalone models for the top model from the pareto front of the Supernet (Sandwich) (see Table 2). It would be interesting to explore this gap for a larger number of architectures from the search space, as shown in Table 4 for MT tasks.

heading: Acknowledgments
text: MAM acknowledges support from Canada Research Chairs (CRC), the Natural Sciences and Engineering Research Council of Canada (NSERC; RGPIN-2018-04267), Canadian Foundation for Innovation (CFI; 37771), and Digital Research Alliance of Canada.6 Lakshmanan’s research was supported in part by a grant from NSERC (Canada). We used ChatGPT for rephrasing and grammar checking of the paper.

heading: A Appendix

heading: A.1 Weight Sharing and Gradient Conflict Analysis
text: A.1.1 Jensen-Shannon distance of alignment vector as a weight sharing measure
We use the Jensen-Shannon distance of alignment vector generated by Layer-wise MoS for two architectures as a proxy to quantify the degree of weight sharing. Ideally, the lower the Jensen-Shannon distance, the higher the degree of weight sharing and vice-versa. We experiment with two architectures of 23M parameters (Smallest A and Smallest B) and two architectures of 118M parameters (Largest A and Largest B). From Table 8, it is clear that
Layer-wise MoS induces low degree of weight sharing between differently sized architectures shown by higher Jensen-Shannon distance between their alignment vectors. On the other hand, there is a high degree of weight sharing between similarly sized architectures where Jensen-Shannon distance is significantly low.
A.1.2 Cosine similarity between the supernet gradient and the smallest subnet gradient as a gradient conflict measure.
We compute gradient conflict using cosine similarity between the supernet gradient and the smallest subnet gradient, following NASVIT work (Gong et al., 2021). In Table 9, we show that Neuron-wise MoS enjoys lowest gradient conflict compared to Layer-wise MoS and HAT, shown by highest cosine similarity.
A.2 Detailed algorithm for Supernet training and Search
A.2.1 Supernet training algorithm The detailed algorithm for supernet training is shown in Algorithm 1.

heading: A.2.2 Search algorithm
text: The detailed algorithm for search is shown in Algorithm 2.

heading: A.3 Comparison to the AutoMoE work
text: Goals: Given a search space of dense and mixtureof-expert models, the goal of the AutoMoE framework (Jawahar et al., 2023) is to search for highperforming model architectures that satisfy userdefined 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. In addition, the MoS framework minimizes the retraining compute required for the searched architecture to approach the standalone performance. The MoS framework designs the supernet with flexible weight sharing and high capacity. On the other hand, the supernet underlying the AutoMoE framework suffers from strict weight sharing and limited capacity. Applications of mixture-of-experts: The main application of mixture-of-experts idea by the AutoMoE framework is to augment the standard NAS
Algorithm 1 Training algorithm for Mixture-of-Supernets used in MT. Input: Training data: Xtr, Search space: A, No. of training steps: num-train-steps, Type of MoS: mos-type Output: Training Supernet Weights: E
1: E← Random weights from Normal Distribution. 2: for iter ← 1 to num-train-steps do 3: // sample data 4: x, y ∼ Xtr 5: // perform sandwich sampling 6: for a in [arand ∼ A, abig, asmall] do 7: Enc(a) // create the architecture encoding 8: // generate architecture-specific weights 9: if mos-type == Layer wise MoS then
10: Wa = ∑ i r(Enc(a)) iEia 11: else if mos-type == Neuron wise MoS then 12: Wa = ∑ i diag(β i a)E i a 13: // compute task-specific loss 14: loss← L(Wax, y) 15: loss.backward() // compute gradients 16: Update E using accumulated gradients // learning rule 17: return E
Algorithm 2 Evolutionary search algorithm for Neural architecture search used in MT. Input: supernet, latency-predictor, num-iterations, num-population, num-parents, num-mutations, num-crossover, mutate-prob, latency-constraint Output: best-architecture
1: // create initial population 2: popu← num-population random samples from the search space 3: for iter ← 1 to num-iterations do 4: // generate parents by picking top candidates 5: cur-parents← top ‘num-parents’ architectures from popu by MoS validation loss 6: // generate candidates via mutation 7: cur-mutate-popu = {} 8: for mi← 1 to num-mutations do 9: cur-mutate-gene ← mutate a random example from popu with mutation probability
mutate-prob 10: if cur-mutate-gene satisfies latency-constraint via latency-predictor then 11: cur-mutate-popu = cur-mutate-popu ∪ cur-mutate-gene 12: // generate candidates via cross-over 13: cur-crossover-popu = {} 14: for ci← 1 to num-crossover do 15: cur-crossover-gene← crossover two random examples from popu 16: if cur-crossover-gene satisfies latency-constraint via latency-predictor then 17: cur-crossover-popu = cur-crossover-popu ∪ cur-crossover-gene 18: // update population 19: popu = cur-parents ∪ cur-mutate-popu ∪ cur-crossover-popu 20: return top architecture from popu by MoS’s validation loss
search space of dense models with mixture-ofexperts models. To this end, the AutoMoE framework modifies the standard weight sharing supernet to support weight generation for mixture-ofexpert models. On the other hand, the Mixture-ofSupernets 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. 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 hand, the router underlying the Mixtureof-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 (all experts).

heading: A.4 Additional Experiments - Efficient BERT
text: A.4.1 BERT pretraining / finetuning settings
Pretraining data: The pretraining data consists of text from Wikipedia and Books Corpus (Zhu et al., 2015). We use the data preprocessing scripts provided by Izsak et al. to construct the tokenized text. Supernet and standalone pretraining settings: The pretraining settings for supernet and standalone models are taken from SuperShaper (Ganesan et al., 2021): batch size of 2048, maximum sequence length of 128, training steps of 125K, learning rate of 5e−4, weight decay of 0.01, and warmup steps of 10K (0 for standalone). For experiments with the search space from SuperShaper (Ganesan et al., 2021) (Section 4.2), the architecture encoding a is a list of hidden size at each layer of the architecture (12 elements since the supernet is a 12 layer model). For experiments with the search space on par with AutoDistil (Xu et al., 2022a) (Section 4.3), the architecture encoding a is a list of four elastic hyperparameters of the homogeneous BERT architecture: number of layers, hidden size of all layers, feedforward network (FFN) expansion ratio of all layers and number of attention heads of all layers (see Table 10 for sample homogeneous BERT architectures).
Finetuning settings: We evaluate the performance of the BERT model by finetuning on each of the seven tasks (chosen by AutoDistil (Xu et al., 2022a)) in the GLUE benchmark (Wang et al., 2018). The evaluation metric is the average accuracy (Matthews’s correlation coefficient for CoLA only) on all the tasks (GLUE average). The finetuning settings are taken from the BERT paper (Devlin et al., 2019): learning rate from {5e−5, 3e−5, 2e−5}, batch size from {16, 32}, and epochs from {2, 3, 4}.
A.4.2 Learning curve for BERT supernet variants
Figure 2 shows the training steps versus validation MLM loss (learning curve) for the standalone BERT model and different supernet based BERT variants. The standalone model and the supernet are compared for the biggest architecture (big) and the smallest architecture (small) from the search space of SuperShaper (Ganesan et al., 2021). For the biggest architecture, the standalone model performs the best. For the smallest architecture, the standalone model is outperformed by all the supernet variants. In both cases, the proposed supernets (especially neuron-wise MoS) perform much better than the standard supernet.

heading: A.4.3 Architecture comparison of Neuron-wise MoS vs. AutoDistil
text: Table 10 shows the comparison of the BERT architecture designed by our proposed neuron-wise MoS with AutoDistil.
A.4.4 Fair comparison of Neuron-wise MoS w.r.t SoTA with MNLI
We compare neuron-wise MoS with NAS-BERT and AutoDistil (agnostic) for different model sizes (≤ 50M parameters) based on GLUE validation performance. In Table 11, we include results on MNLI task. For fair comparison, we drop AutoDistil (proxy), which directly uses MNLI task for architecture selection. Neuron-wise MoS improves over the baselines in all model sizes, in terms of average GLUE. For MNLI task, neuron-wise MoS improves over the baselines in most model sizes.
A.4.5 BERT results with different random seeds
Table 12 displays BERT results on CoLA and RTE with various random seeds. Layer-wise MoS consistently enhances performance over baselines in RTE and diminishes performance compared to
baselines in CoLA for both seeds. The BERT architecture (67M parameters) corresponds to the top model from the pareto front of Supernet (Sandwich) in SuperShaper’s search space (consistent with Table 2).

heading: A.5 Additional Experiments - Efficient Machine Translation
text: A.5.1 Machine translation benchmark data Table 13 shows the statistics of three machine translation datasets: WMT’14 En-De, WMT’14 En-Fr, and WMT’19 En-De.
A.5.2 Training settings and metrics The training settings for both supernet and standalone models are the same: 40K training steps, Adam optimizer, a cosine learning rate scheduler, and a warmup of learning rate from 10−7 to 10−3 with cosine annealing. The best checkpoint is selected based on the validation loss, while the performance of the MT model is evaluated based on BLEU. The beam size is four with length penalty of 0.6. The architecture encoding a is a list of following 10 values:
1. Encoder embedding dimension corresponds to embedding dimension of the encoder.
2. Encoder #layers corresponds to number of encoder layers.
3. Average encoder FFN. intermediate dimension corresponds to average of FFN intermediate dimension across encoder layers.
4. Average encoder self attention heads corresponds to average of number of self attention heads across encoder layers.
5. Decoder embedding dimension corresponds to embedding dimension of the decoder.
6. Decoder #Layers corresponds to number of decoder layers.
7. Average Decoder FFN. Intermediate Dimension corresponds to average of FFN intermediate dimension across decoder layers.
8. Average decoder self attention heads corresponds to average of number of self attention heads across decoder layers.
9. Average decoder cross attention heads corresponds to average of number of cross attention heads across decoder layers.
10. Average arbitrary encoder decoder attention corresponds to average number of encoder layers attended by cross-attention heads in each decoder layer (-1 means only attend to the last layer, 1 means attend to the last two layers, 2 means attend to the last three layers).
A.5.3 Supernet vs. Standalone performance plot
Figure 3 displays the supernet vs. the standalone performance for 15 randomly sampled architectures on all the three tasks. Neuron-wise MoS excel for almost all the top performing architectures (≥ 26.5 and ≥ 42.5 standalone BLEU for WMT’14 En-De and WMT’19 En-De respectively), which indicates that the models especially in the pareto front can benefit immensely from neuron level specialization.

heading: A.5.4 HAT Settings
text: Evolutionary search: The settings for the evolutionary search algorithm include: 30 iterations, population size of 125, parents population of 25, crossover population of 50, and mutation population of 50 with 0.3 mutation probability. Latency estimator: The latency estimator is developed in two stages. First, the latency dataset is constructed by measuring the latency of 2000 randomly sampled architectures directly on the userdefined hardware (NVIDIA V100 GPU). Latency is the time taken to translate a source sentence to a target sentence (source and target sentence lengths of 30 tokens each). For each architecture, 300 latency measurements are taken, outliers (top 10% and bottom 10%) are removed, and the rest (80%) is averaged. Second, the latency estimator is a 3 layer multi-layer neural network based regressor, which is trained using encoding and latency of the architecture as features and labels respectively.
A.5.5 Additional training steps to close the gap vs. performance
Figure 4, Figure 5, and Figure 6 show the additional training steps vs. BLEU for different latency constraints on the WMT’14 En-De task, WMT’14 En-Fr and WMT’19 En-De tasks respectively.

heading: A.5.6 Evolutionary Search - Stability
text: We study the initialization effects on the stability of the pareto front outputted by the evolutionary search for different supernets. Table 14 displays sampled (direct) BLEU and latency of the models in the pareto front for different seeds on the
WMT’14 En-Fr task. The differences in the latency and BLEU across seeds are mostly marginal. This result highlights that the pareto front outputted by the evolutionary search is largely stable for all the supernet variants.
A.5.7 Impact of different router function Table 15 displays the impact of varying the number of hidden layers in the router function for neuronwise MoS on the WMT’14 En-De task. Two hidden layers provide the right amount of router capacity, while adding more hidden layers results in steady performance drop.
A.5.8 Impact of increasing the number of expert weights ‘m’
Table 16 displays the impact of increasing the number of expert weights ‘m’ for the WMT’14 En-Fr task, where the architecture for all the supernets is the top architecture from the pareto front of HAT for the latency constraint of 200 ms. Under the standard training budget (40K steps for MT), the performance of layer-wise MoS does not seem to improve by increasing ‘m’ from 2 to 4. Increasing ‘m’ introduces too many parameters, which might necessitate a significant increase in the training budget (e.g., 2 times more training steps than the standard training budget). For fair comparison with existing literature, we use the standard training budget for all the experiments. We will investigate the full potential of the proposed supernets by combining larger training budget (e.g., ≥ 200K steps) and larger number of expert weights (e.g., ≥ 16 expert weights) in future work.

heading: A.5.9 SacreBLEU vs. BLEU
text: We use the standard BLEU (Papineni et al., 2002) to quantify the performance of supernet following HAT for a fair comparison. In Table 17, we also experiment with SacreBLEU (Post, 2018), where the similar trend of MoS yielding better performance for a given latency constraint holds true.
A.5.10 Breakdown of the overall time savings Table 18 shows the breakdown of the overall time savings of MoS supernets versus HAT for comput-
ing pareto front for the WMT’14 En-De task. The latency constraints include 100 ms, 150 ms, 200 ms. MoS have an overall GPU hours savings of at least 20% w.r.t. HAT, thanks to significant savings in additional training time (45%-51%).

heading: A.5.11 Codebase
text: We share the codebase at: https://github.com/ UBC-NLP/MoS, which can be used to reproduce all the results in this paper. For both BERT and machine translation evaluation benchmarks, we add a
README file that contains the following instructions: (i) environment setup (e.g., software dependencies), (ii) data download, (iii) supernet training, (iv) search, and (v) subnet retraining.

heading: Future Work
text: , ≥ 16 expert weights) in future work