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

We can reformulate the function $f_{a}(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,

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 $f_{a}(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 $W_{a}$ 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.

Assume there are $m$ (number of experts) unique weight matrices ($\{{E^{i}\in\mathcal{R}^{n_{{out}_{big}}\times n_{{in}_{big}}}}\}^{m}_{i=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 $n_{out_{a}}$ output features, we propose the layer-wise MoS that computes the weights specific to the architecture $a$ (i.e. $W_{a}\in\mathcal{R}^{n_{out_{a}}\times n_{in}}$) by performing a weighted combination of expert weights, $W_{a}=\sum_{i}\alpha^{i}_{a}E^{i}_{a}$. Here, $E^{i}_{a}\in\mathcal{R}^{n_{out_{a}}\times n_{in}}$ corresponds to the standard top rows extraction from the $i^{th}$ expert weights. The alignment vector ($\alpha_{a}\in[0,1]^{m},\sum_{i}\alpha^{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 $\text{Enc}(a)\in\mathcal{R}^{n_{enc}\times 1}$ (e.g., a list of the number of output features for different layers), and apply a learnable router $r(\cdot)$ (an MLP with softmax) to produce such scores, i.e. $\alpha_{a}=r(\text{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):

The router $r(\cdot)$ governs the degree of weight sharing between two architectures through modulation of alignment scores ($\alpha_{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 $\alpha_{a}=(1,0)$ and $\alpha_{b}=(0,1)$. In this scenario, $g(x,a;E)$ exclusively utilizes weights from $E^{1}$, and $g(x,b;E)$ uses weights from $E^{2}$, enabling updates to $E^{1}$ and $E^{2}$ 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(\cdot)$ 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 $\alpha_{a}=r(\text{Enc}(a))$ can be generated offline. Expert weights collapse, reducing the number of parameters for architecture $a$ to $n_{out_{a}}\times n_{in_{a}}$. 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.

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}}\times|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 $\beta_{a}=r(\cdot)\in[0,1]^{n_{out_{big}}\times m}$ for each layer, and the sum of each row in $\beta_{a}$ is 1. Similar to layer-wise MoS, we use an MLP to produce the $n_{out_{big}}\times m$ matrix and apply softmax on each row. We formulate the function $g(x,a;E)$ for neuron-wise MoS as

where $\text{diag}(\beta)$ constructs a $n_{out_{big}}\times n_{out_{big}}$ diagonal matrix by putting $\beta$ on the diagonal, and $\beta_{a}^{i}$ is the $i$-th column of $\beta_{a}$. $E^{i}$ is still an $n_{out_{big}}\times n_{in}$ matrix as in layer-wise MoS. Compared to layer-wise MoS, neuron-wise MoS offers increased flexibility ($m\times{n_{out_{a}}}$ 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 layer-wise MoS and HAT, as shown by the highest cosine similarity (see A.1.2).

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 ($f_{a}(x;W)$) by replacing the two linear layers in every feed-forward network block with $g(x,a;E)$.

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 ²²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.. Parameters $m$ and router’s hidden dimension are set to $2$ and $128$, respectively, for MoS supernets.

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 training for the architecture ($1$x 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. ³³3Consistency of these results across different seeds is discussed in A.4.5. 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.

The SoTA NAS frameworks for constructing a task-agnostic BERT model are NAS-BERT (Xu et al., 2021) and AutoDistil (Xu et al., 2022a).⁴⁴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. The NAS-BERT 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 ($1$x 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 (task-agnostic) 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 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. ⁵⁵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.

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 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).

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$’.

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).

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.

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 En-Fr 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.

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.