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 mixture-of-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 architecture-based 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 memory-efficient 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.

## introduction neural architecture search (nas) automates the design of high-quality architectures for natural language processing (nlp) tasks while meeting specified efficiency constraints @xcite @xcite . nas is commonly treated as a black-box optimization @xcite , 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 @xcite . recent studies demonstrate successful direct use of subnetworks for image classification with performance comparable to training from scratch @xcite . 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 @xcite , 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 @xcite 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 @xcite and gradient conflicts @xcite . 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 model may not be optimal for the 90m model, leading to significant gradient conflicts during optimization @xcite . 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 @xcite . 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 @xcite and machine translation (mt) models. for efficient bert, our best supernet outperforms supershaper @xcite by 0.85 glue points, surpasses nas-bert @xcite and autodistil @xcite in various model sizes (≤ 50m parameters). compared to hat @xcite , our top supernet reduces the supernet vs. standalone model gap by 26.5%, provides a superior pareto front for latency-bleu tradeoff (100 to 200 ms), and decreases the steps needed to close the gap by 39.8%. a summary in the we summarize our key contributions: the mentioned formulation is termed single path one-shot (spos) optimization @xcite for supernet training. another popular technique is sandwich training @xcite , where the largest (a big ), smallest (a small ), and uniformly sampled architectures (a rand ) from the search space are jointly optimized. existing supernets typically have limited model capacity to extract architecture-specific weights. for simplicity, assume the model function f a (x; w ) is a fully connected layer (output @xmath0, omitting bias term for brevity), where @xmath1, @xmath2, and @xmath3. n in and n out correspond to the number of input and output features respectively. then, the weights (w @xmath4. assume one samples two architectures (a and b) from the search space with the number of output features n outa and n out b respectively. then, the weights corresponding to the architecture with the smallest number of output features will be a subset of those of the other architecture, sharing the first |n outan out b | 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 f a (x; w )) cannot allocate weights specific to the smaller architecture. another issue with f a (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). ## related work in this section, we briefly review existing research on nas in nlp. evolved transformer (et) @xcite ) is an initial work that explores nas for efficient mt models. it employs evolutionary search and dynamically allocating training resources for promising candidates., hat @xcite introduces a weight-sharing supernet as a performance estimator, amortizing the training cost for candidate mt evaluations in evolutionary search. nas-bert @xcite partitions the bert-base model into blocks and trains a weight-sharing supernet to distill each block. nas-bert 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. super-shaper @xcite pretrains a weightsharing supernet for bert using a masked language modeling objective with sandwich training. autodistil @xcite employs few-shot nas @xcite : 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). auto-moe (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 @xcite 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. dataset additional training steps (↓) additional training time (nvidia v100 hours) (↓) supernet wmt'14 en-de wmt'14 en-fr wmt'19 en-de wmt'14 en-de wmt'14 en-fr wmt'19 en-de hat 33k 33k 26k 63.9 60.1 52.3 laye. mos 16k (51.5%) 30k (9%) 20k (23%) 35.5 (44.4%) 66.5 (-10.6%) 45.2 (13.5%) neur. mos 13k (60%) 26k (21%) 16k (38.4%) 31.0 (51.4%) 61.7 (-2.7%) 39.5 (24.5%) ## conclusion 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.