Venue: NAACL 2024

# Efficiently Distilling LLMs for Edge Applications

Achintya Kundu, Fabian Lim, Aaron Chew, Laura Wynter, Penny Chong, Rhui Dih Lee

IBM Research, Singapore

###### Abstract

Supernet training of LLMs is of great interest in industrial applications as it confers the ability to produce a palette of smaller models at constant cost, regardless of the number of models (of different size / latency) produced. We propose a new method called Multistage Low-rank Fine-tuning of Super-transformers (MLFS) for parameter-efficient supernet training. We show that it is possible to obtain high-quality encoder models that are suitable for commercial edge applications, and that while decoder-only models are resistant to a comparable degree of compression, decoders can be effectively sliced for a significant reduction in training time.

## 1 Introduction

Given their sizes up to billions of parameters, [23, 4], it is challenging for enterprises to fine-tune Large Language Models (LLMs), and furthermore they are not suitable for deployment on edge devices with limited memory and computational power. We wish to enable LLMs on edge environments for enterprise use cases. This requires the following two capabilities. (1) Accommodating a variety of edge device hardware: A single fine-tuned model is not optimal across the spectrum of devices. For industrial applications, a palette of fine-tuned LLMs is required for different hardware. (2) Dynamically changing resource levels: At run-time, the available resources on edge devices evolve over time, and appropriate model should be dynamically selected based on the available resources of each device.

A considerable amount of research has focused on compressing LLMs [56, 78, 90, 103]. Methods that train a single small model guided by a large teacher model such as DistilBERT [103] and BERT-PKD [104], either achieve limited compression or do not scale to a large number of deployment devices. Supernet training methods [23, 105, 106, 107, 108, 109, 110] were introduced to address these limitations: multiple smaller subnets within the supernet are trained simultaneously with weight-sharing. This one-time training approach produces a palette of smaller models, helping mitigate the computational cost of fine-tuning a model for each deployment scenario. However, the full-parameter supernet training approach is impractical when fine-tuning of an LLM is required for multiple deployment scenarios, limiting its utility for enterprises.

Parameter-efficient fine-tuning (PEFT) methods such as Low-Rank Adaptation (LoRA) reduces the number of trainable parameters by allowing only rank-decomposition matrices to be trained while freezing the pre-trained weights of the model. PEFT methods, however, are not applicable to supernet training due to the implications on the weight-shared sub-networks. Our work bridges this gap to enable efficient fine-tuning of LLMs for edge devices. Our contributions are:

1. We propose a parameter-efficient, distillation-based approach for supernet training of LLMs.
2. We devise a gradient scaling scheme to improve convergence speed of any form of supernet training.
3. We demonstrate significant compression of encoder models for edge. We highlight the limits of comparable compression for decoder models, while demonstrating a huge reduction in the steps needed for convergence.

## 2 Related Work

Classical compression methods have been used for LLMs including pruning [103, 104], low rank approximation [105],

et al., 2019; Lan et al., 2019), and quantization Shen et al. (2020); Zafrir et al. (2019); Bhandare et al. (2019). Knowledge distillation (KD) is adopted in BERT-PKD Sun et al. (2019), tiny-BERT Jiao et al. (2020), and distilBERT Sanh et al. (2019) and Gu et al. (2023) in MiniLLM to distill knowledge from the layers of a large transformer model to a smaller one. See also the survey Zhu et al. (2023). All these existing methods produce a single compressed model, unsuitable for edge scenarios with multiple deployment devices having varying computational capability.

Neural architecture search (NAS) based on reinforcement learning Zoph and Le (2016) and evolutionary algorithms Real et al. (2019); Zhu et al. (2019) trains every possible architecture and is very slow. Weight-sharing NAS was thus developed: in Guo et al. (2020); Cai et al. (2018), the building blocks in the same layer are isolated as all architectures are single paths. Weight-sharing NAS does not scale well to large architecture search spaces, hence, weight-entangled NAS, where subnets with common parts share weights, was introduced.

For resource-constrained edge deployment, supernet training Cai et al. (2019); Kundu et al. (2023); Chen et al. (2021); Xu et al. (2021); Gao et al. (2022); Dong et al. (2022) was developed as a mode of jointly training multiple sub-networks (subnets) with entangled weights: one trains the supernet only once for all deployment scenarios. Cai et al. (2019) introduced an elastic convolutional neural network with "progressive shrinkage", where larger subnets are trained first. Recent works have improved sampling strategies, e.g. the sandwich rule with in-place distillation Yu et al. (2020), attentive sampling Wang et al. (2021), stochastic nature gradient Zhang et al. (2021), or post-training sampling Lou et al. (2021). Our work is related to supernet training for transformer models Hou et al. (2020); Zhang et al. (2021); Wang et al. (2022, 2020); Chen et al. (2021b). This gradient scaling technique can be used with any of the above supernet methods.

Parameter-efficient fine-tuning (PEFT) has been of great benefit in fine tuning LLMs. BitFit Ben Zaken et al. (2022) updates the bias terms in pre-trained models while freezing the remaining parameters. LoRA Hu et al. (2022) decomposes attention weight gradients into low-rank matrices to reduce the number of trainable parameters. AdaLoRA Zhang et al. (2023) and QLoRA Dettmers et al. (2023) further improve LoRA Hu et al. (2022). Note that PEFT allows fine-tuning a base model on a single GPU but does not produce smaller models. None of the PEFT methods can be used for weight-sharing supernet training.

## 3 Solution Design

For use in enterprise settings, the solution must allow fine-tuning of models on a small GPU footprint. In addition, inference cost in terms of storage must be minimised. We therefore design a solution which does not store the full size model checkpoint for every downstream task but only the frozen weights of the pre-trained base model and the low rank matrices. For inference in commercial edge use cases, we wish to enable storing the desired models locally for a wide variety of edge device resource requirements. We thus develop an approach where storage is minimised, storing only one base model and as many low rank adapter matrices as there are target model size variations, where low-rank adapters are very small. If the model is stored locally on an edge device, our proposed slicing operation takes place where the supernet fine-tuning is performed and the desired model is downloaded for inference. The slicing operation takes place for each model size-task combination and each resulting subset can be cached for inference.

## 4 Problem Formulation

First, we provide notation. Given a transformer model with architectural configuration \(\Phi\) and weights \({\rm W}\), we denote its forward-pass mapping by \({\rm f_{\Phi}}(\cdot;{\rm W}):\mathcal{X}\rightarrow\mathcal{Y}\). We consider the output space \(\mathcal{Y}\) to be the set of all non-negative vectors in \(\mathbb{R}^{\nu}\) with elements summing to \(1\), where \(\nu\) denotes the number of classes / vocabulary size). With slight abuse of notation, we write the forward-pass mapping of an input \(x\in\mathcal{X}\) through a transformer model \(\Phi\) as \(\hat{y},z,{\bf h}~{}=~{}{\rm f_{\Phi}}(x;{\rm W}),\) where \(\hat{y}\in\mathcal{Y}\) denotes the predicted probability distribution over the (class labels) vocabulary, \(z\) denotes the vector of logits, and \({\bf h}\) represents a tuple of features such as hidden state vectors and attention values from different transformer layers. Note that \(\hat{y}=\boldsymbol{\sigma}(z)\), where \(\boldsymbol{\sigma}\) is the standard soft-max function that maps a vector of logits into a probability vector. Given a training data set \(\mathcal{D}_{train}\subset\mathcal{X}\times\mathcal{Y}\), model weights \({\rm W}\) are learnt by minimizing training loss:

\[\underset{{\rm W}}{\operatorname*{argmin}}\,\Big{[}\,\mathcal{L}_{\Phi}({\rm W })\,:=\,\mathbb{E}\big{[}\ell\,{\rm f_{\Phi}}(x;{\rm W})\,,y\,]\ \big{]}\,\Big{]}, \tag{1}\]

where \(\mathbb{E}\) denotes expectation over training example \((x,y)\) drawn uniformly at random from \(\mathcal{D}_{train}\) and 

\(\ell\) denotes a loss function. Most commonly, \(\ell\) is chosen to be a task specific loss function, \(\ell_{\text{task}}\), such as cross-entropy (i.e., \(\mathtt{CE}[\cdot,\cdot]\)) for classification or causal language modeling loss for generative models.

Next, we introduce the super-transformer and related terminologies. We define three types of networks - _Teacher network_, _Super-transformer_ (supernet) and _Sub-transformer_ (subnet). The teacher is a fixed network with the same configuration as the pre-trained transformer. A super-transformer is a dynamic model whose architectural dimensions (embedding dimension, number of heads, number of layers, etc.) are configurable at run time. The _maxnet_ (resp. _minnet_) is the largest (resp. smallest) network in the super-transformer's architecture space. Weight entanglement (weight-sharing) allows super-transformer weights to be used across sub-transformers, which are subsets of the super-transformer. Pre-trained transformer weights initialise the super-transformer.

The dynamic nature of a super-transformer is explicitly specified via a set \(\mathcal{A}\), called _configuration space_, consisting of architectural configurations of all sub-transformer models under consideration. The definition of a super-transformer also includes how the configuration \(\Phi\in\mathcal{A}\) is to be mapped to a unique transformer model \(\mathrm{f}_{\Phi}\). A weight-sharing super-transformer uses a set of shared weights \(\mathrm{W}_{\mathsf{Sup}}\) to define all sub-transformer models' weights. This is done through a weight projection operator \(\mathbf{\Pi}\) that slices (selects an appropriate subset of) the super-transformer's weights \(\mathrm{W}_{\mathsf{Supp}}\) into weights of a sub-transformer model:

\[\mathrm{W}_{\Phi}:=\mathbf{\Pi}_{\Phi}(\mathrm{W}_{\mathsf{Sup}})\,,\;\forall \Phi\in\mathcal{A}. \tag{2}\]

The aim of a weight-sharing super-transformer is to simultaneously train all the transformer models \(\left\{\mathrm{f}_{\Phi}(\cdot;\mathbf{\Pi}_{\Phi}(\mathrm{W})):\mathcal{X} \rightarrow\mathcal{Y}\,|\,\Phi\in\mathcal{A}\right\}\) through the shared weights \(\mathrm{W}_{\mathsf{Sup}}\). A typical training objective for super-transformers is the training loss averaged over all model configurations in \(\mathcal{A}\):

\[\underset{\mathrm{W}_{\mathsf{Sup}}}{\operatorname*{argmin}}\Big{[}\mathcal{L }_{\mathsf{Sup}}(\mathrm{W}_{\mathsf{Sup}}):=\mathbb{E}\big{[}\mathcal{L}_{ \Phi}\!\big{(}\mathbf{\Pi}_{\Phi}(\mathrm{W}_{\mathsf{Sup}})\big{)}\big{)} \Big{]}, \tag{3}\]

where \(\mathbb{E}\) denotes expectation over model configuration \(\Phi\) drawn uniformly at random from \(\mathcal{A}\) and \(\mathcal{L}_{\Phi}\), as defined in (1), is averaged training loss for configuration \(\Phi\). Super-transformer weights, \(\mathrm{W}_{\mathsf{Sup}}\), are learnt with stochastic gradient (denoted \(\hat{\nabla}\)) of the super-transformer's loss \(\mathcal{L}_{\mathsf{Sup}}\) estimated as

\[\hat{\nabla}_{\mathrm{W}}\mathcal{L}_{\mathsf{Sup}}(\mathrm{W}_{ \mathsf{Sup}})\!=\!\frac{1}{K}\sum_{j=1}^{K}\hat{\nabla}_{\mathrm{W}}\mathcal{ L}_{\Phi_{j}}\big{(}\mathbf{\Pi}_{\Phi_{j}}(\mathrm{W}_{\mathsf{Sup}})\big{)}, \tag{4}\] \[\hat{\nabla}_{\mathrm{W}}\mathcal{L}_{\Phi}(\mathrm{W}_{\Phi})\!= \!\frac{1}{|\mathcal{B}|}\!\sum_{i\in\mathcal{B}}\!\nabla_{\mathrm{W}}\ell \big{[}\mathrm{f}_{\Phi}\big{(}x^{i};\mathrm{W}_{\Phi}\big{)}\,,y^{i}\big{]}, \tag{5}\]

where \(\{\Phi_{1},\cdots,\Phi_{K}\}\) are \(K\) sub-transformer configurations sampled from \(\mathcal{A}\) to approximate the expectation in (3) and \(\mathcal{B}\) is a mini-batch of training examples sampled from \(\mathcal{D}_{train}\) to approximate the expectation in (1). Fine-tuning LLM super-transformers is computationally challenging in enterprise use cases as it involves computing gradients of sub-transformers' loss functions with respect to a huge number of parameters.

## 5 MLfs

We therefore developed Multistage Low-rank Fine-tuning of Super-transformers (MLFS). Given a teacher model with configuration \(\Phi_{\mathsf{Tch}}\) and pretrained weights \(\mathrm{W}_{\mathsf{Tch}}^{\mathrm{pretrain}}\), we assume that its weights (denoted \(\mathrm{W}_{\mathsf{Tch}}\)) can be fine-tuned on the given task by learning low-rank matrices \(A_{0},B_{0}\) on top of pre-trained weights \(\mathrm{W}_{\mathsf{Tch}}^{\mathrm{pretrain}}\):

\[\mathrm{W}_{\mathsf{Tch}}:=\mathrm{W}_{\mathsf{Tch}}^{\mathrm{pretrain}}+A_{0} \ast B_{0}, \tag{7}\]

where \(A_{0},B_{0}\) are of (low) rank \(r\). Note that pre-trained weights \(\mathrm{W}_{\mathsf{Tch}}^{\mathrm{pretrain}}\) remain unchanged during super-transformer fine-tuning. The low-rank matrices, \(A_{0}\) and \(B_{0}\), are learnt by minimizing the cross-entropy loss of the teacher model \(\mathrm{f}_{\Phi_{\mathsf{Tch}}}(\cdot;\mathrm{W}_{\mathsf{Tch}}):\mathcal{X} \rightarrow\mathcal{Y}\) over the training data set \(\mathcal{D}_{train}\). Specifically, we perform \(E_{0}\) epochs of fine-tuning on the teacher to learn \(A_{0},B_{0}\). This is stage-\(0\) of the multistage fine-tuning algorithm. We denote the teacher weights obtained at the end of stage-\(0\) by \(\mathrm{W}_{\mathsf{Tch}}\). We now define a super-transformer with maxnet configuration the same as the teacher's. Thus the super-transformer's weights \(\mathrm{W}_{\mathsf{Sup}}\) are of the same size as the teacher weights \(\mathrm{W}_{\mathsf{Tch}}\)). To fine-tune the super-transformer weights \(\mathrm{W}_{\mathsf{Sup}}\), in each of the subsequent stages, we freeze \(\mathrm{W}_{\mathsf{Tch}}\) and propose learning two stage-specific low-rank matrices \(A_{s},B_{s}\), of the same rank, \(r\), as \(A_{0},B_{0}\), that are shared across all sub-transformer models in that stage. To be precise, we impose the following structure on the weights of the sub-transformers at stage-\(s\):

\[\begin{array}{l}\mathrm{W}_{\mathsf{Sup}}:=\mathrm{W}_{\mathsf{ Tch}}+\sum_{s=1}^{2}A_{s}\ast B_{s},\\ \mathrm{W}_{\Phi}=\mathbf{\Pi}_{\Phi}(\mathrm{W}_{\mathsf{Sup}}),\;\forall\Phi \in\mathcal{A}.\end{array}\] (8) 



**Input:** Transformer model (teacher) with configuration \(\Phi_{\tt Tch}\) & off-the-shelf pre-trained weights \({\rm W^{pretrain}_{\sf Tch}}\), model configuration space \(\mathcal{A}\) consisting of smaller (than \(\Phi_{\tt Tch}\)) transformer architectures of interest, \(\mathcal{D}_{train}\): fine-tuning data set for the target task, \(r\): rank of the low-rank matrices and distillation factor \(\alpha\in[0,1]\).

**Loss functions:** Target task loss \(\ell_{\tt task}\), knowledge distillation loss \(\ell_{\sf KD}\), feature distillation loss \(\ell_{\sf FD}\).

**Multistage Training:**

```
1:for stage \(s=0,1,2\)do
2:Initialize the low-rank matrices \(\{A_{s},B_{s}\}\) to be learned at stage \(s\) .
3:foriteration\(=1,...\)do
4: Get a mini-batch \(\mathcal{B}\) of training examples from data set \(\mathcal{D}_{train}\): \(\{(x^{i},y^{i})\in\mathcal{D}_{train}\,|\,i\in\mathcal{B}\}\).
5: Load the super-transformer model with weights \(~{}{\rm W_{Sup}}\leftarrow{\rm W^{pretrain}_{\sf Tch}}+\sum_{l=0}^{s}A_{l}*B_{ l}\).
6:\(\mathcal{A}_{s}:=\{\Phi_{1},\Phi_{2},\cdots\}\leftarrow\texttt{sample\_ sub-transformers}(\mathcal{A},\text{stage}=s)\). [\(\Phi_{1}\) is the maxnet].
7:for each \(\Phi_{j}\in\mathcal{A}_{s}\)do
8: Load the sub-transformer model \(\Phi_{j}\) with weights \({\rm W_{\Phi_{j}}}:=\boldsymbol{\Pi}_{\Phi_{j}}({\rm W_{Sup}})\).
9:\(n_{j}:=\#\) of fine-tuning weights in model configuration \(\Phi_{j}\) .
10: Compute forward-pass on the sub-transformer \(\Phi_{j}\): \(\hat{y}_{j}^{i},z_{j}^{i},\mathbf{h}_{\Phi_{j}}\gets\mathrm{f}_{\Phi_{j}} (x^{i};{\rm W_{\Phi_{j}}}),\,\forall i\in\mathcal{B}\).
11: For the case of maxnet (\(\Phi_{1}\)) set the distillation factor \(\alpha\) to \(0\).
12: Find the loss: \(\texttt{loss}_{j}^{i}\leftarrow(1-\alpha)\,\ell_{\tt task}[\hat{y}_{j}^{i}, y^{i}]+\alpha\left(\ell_{\sf RD}[z_{j}^{i},z_{1}^{i}]\,+\,\ell_{\sf FD}[\mathbf{ h}_{\Phi_{j}}^{i},\mathbf{h}_{\Phi_{1}}^{i}]\right),\,\forall i\in\mathcal{ B}\).
13: Compute gradients \(({\nabla_{\!\!A_{s}}}\texttt{loss}_{j}^{i},{\nabla_{\!\!B_{s}}}\texttt{loss}_ {j}^{i})\) using backward-pass on sub-transformer \(\Phi_{j}\).
14:endfor
15: Update \(A_{s},B_{s}\) using the gradients \((\widehat{{\nabla}_{\!\!A_{s}}}\mathcal{L_{\mathsf{Sup}}},\widehat{{\nabla}_ {\!\!B_{s}}}\mathcal{L_{\mathsf{Sup}}})\) of the super-transformer's loss: \[\widehat{{\nabla}_{\!\!\mathrm{W}}}\mathcal{L_{\mathsf{Sup}}}=\frac{1}{| \mathcal{A}_{s}|}{\sum_{\Phi_{j}\in A_{s}}}\!\!\bigg{(}\frac{n_{1}}{n_{j}} \bigg{)}^{\gamma}\,\hat{{\nabla}_{\!\!\mathrm{W}}}\mathcal{L_{\Phi_{j}}},~{}~ {}\hat{{\nabla}_{\!\!\mathrm{W}}}\mathcal{L_{\Phi_{j}}}=\frac{1}{|\mathcal{B} |}{\sum_{i\in\mathcal{B}}{\nabla_{\!\!\mathrm{W}}}\texttt{loss}_{ j}^{i}},~{}~{}\forall\,{\rm W}\in\{A_{s},B_{s}\}.\] (6)
16:endfor
17:endfor
18:Output:\(\{A_{s},B_{s}\}_{s=0}^{2}\) and fine-tuned super-transformer weights: \({\rm W_{Sup}}={\rm W^{pretrain}_{\sf Tch}}+\sum_{s=0}^{2}A_{s}*B_{s}\).
```

**Algorithm 1** Multistage Low-rank Fine-tuning of Super-transformers (MLFS)

Stage-\(s\) of the fine-tuning process involves learning only the low-rank matrices, \(A_{s}\), \(B_{s}\), by minimizing the super-transform loss as in (3). In stage-\(1\), we sample sub-transformer models by sampling different widths from the super-transformer keeping the depth (number of layers) same as the maxnet. In stage-\(2\), we sample sub-transformer models by sampling different widths as well as depths. We always sample the maxnet model from the super-transformer as the \(1^{st}\) sub-transformer model, \(\Phi_{1}\), at every iteration. We call this **M**ultistage **L**ow-rank **F**ine-tuning of **S**uper-transformers (MLFS) and present it in Algorithm 1.

**Proposition 1**: _Let the individually fine-tuned weights of a subnet, \(\Phi\), be expressed as \({\rm W_{\Phi}}=\boldsymbol{\Pi}_{\Phi}({\rm W^{pretrain}_{\sf Tch}})+\Delta{ \rm W_{\Phi}}\). Then, MLFS has the following structure on \(\Delta{\rm W_{\Phi}}\):_

\[\Delta{\rm W_{\Phi}}=\boldsymbol{\Pi}_{\Phi}\left(\sum_{s=0}^{2}A_{s}*B_{s} \right),\,\forall\Phi\in\mathcal{A}, \tag{9}\]

_where \(\{A_{s},B_{s}\}_{s=0,1,2}\) are low-rank matrices shared across all sub-transformers \(\Phi\in\mathcal{A}\)._

To illustrate the computational savings, recall \({\rm W^{pretrain}_{\sf Tch}}\in\mathbb{R}^{d\times d}\), where \(d\) is typically of the order \(10^{4}-10^{6}\). For rank \(r\) (typically \(<10\)) for the low-rank matrices: \(A_{s}\in\mathbb{R}^{d\times r}\), \(B_{s}\in\mathbb{R}^{r\times d},~{}s=0,1,2\), where \(r\ll d\). Then, the number of parameters to be learnt in the MLFS approach is \(6rd\). In contrast, full fine-tuning requires updating \(d^{2}\) parameters at every iteration.

Gradient ScalingFor faster convergence of the smaller sub-transformers within a super- transformer, we propose a novel weighted-combination of the gradients of the sampled sub-transformers.

**Proposition 2**: _Let \(1^{st}\) sampled sub-transformer, \(\Phi_{1}\), be the maxnet be in every iteration. Then the scaled gradient of the super-transformer training

loss, \({\cal L}_{\mathsf{sup}}\), in Algorithm 1 is given by_

\[\sum_{j=1}^{K}(n_{1}/n_{j})^{\gamma}\,\nabla_{\mathrm{W}}{\cal L}_{\Phi_{j}}, \tag{10}\]

_where \(\nabla_{\mathrm{W}}\) denotes gradient w.r.t. only those weights that are being fine-tuned (in this case only the LoRA matrices), \(n_{j}\) denotes the actual number of trainable weights in model configuration \(\Phi_{j}\) and \(\gamma\geq 1\) is a hyper-parameter._

**Proof:** Each sub-transformer gradient in (10), \(\mathsf{grad}^{j}\), is scaled by \((n_{1}/n_{j})\), which is obtained from the relative weighting of the loss functions. Let \({\cal L}_{j}(\mathrm{W})\) denote the \(j\)-th sub-transformer's loss. Using first-order Taylor expansion, we get:

\[{\cal L}_{\Phi_{j}}(\mathrm{W+\boldsymbol{\delta}})\approx{\cal L}_{\Phi_{j}} (\mathrm{W})+\langle\nabla_{\mathrm{W}}{\cal L}_{\Phi_{j}}(\mathrm{W}), \boldsymbol{\delta}\rangle,\]

where \(\langle\cdot,\cdot\rangle\) denotes inner (dot) product operation. Therefore, the steepest possible decrease in the loss function \({\cal L}_{\Phi_{j}}\) can be approximated as:

\[\Delta{\cal L}_{\Phi_{j}}\approx\|\nabla_{\mathrm{W}}{\cal L}_{\Phi_{j}}( \mathrm{W})\|_{1}\,|\boldsymbol{\delta}|_{max}\approx O(n_{j})|\boldsymbol{ \delta}|_{max},\]

where we approximate the \(\|\cdot\|_{1}\) norm using the zero-th norm, i.e., number of non-zero elements and \(n_{j}\) stands for the actual number of trainable parameters in sub-transformer configuration \(\Phi_{j}\). Since the decrease in the loss of a sub-transformer model \(\Phi_{j}\) is approximately proportional to the number of trainable model parameters \((n_{j})\), we scale the losses using \((n_{1}/n_{j})^{\gamma},\gamma\geq 1\) so that training losses of smaller sub-transformer models converge at a rate similar to that of larger sub-transformer configurations. Recall that \(n_{1}\) is the maximum number of trainable parameters as \(1^{st}\) sampled sub-transformer \(\Phi_{1}\) is always the maxnet. \(\square\)

Distillation Loss for Super-transformers:Knowledge distillation is straightforward in a fixed-network fine-tuning setting. However, it is less so when fine-tuning a supernet, and in particular, fine-tuning a supernet using the proposed multistage LoRA based approach. Specifically, the subnets receive two types of knowledge distillation (KD) from the teacher: (a) the usual KD loss that utilizes the output logits of the teacher and (b) distillation of features from transformer layers (Jiao et al., 2020) of the teacher.

To define the distillation based losses precisely, let the forward-pass mapping of an input training sample \(x^{i}\) through sub-transformer \(\Phi_{j}\) be \(\hat{y}^{i}_{j},z^{i}_{j},\mathbf{h}^{i}_{\Phi_{j}}\leftarrow\mathrm{f}_{\Phi _{j}}(x^{i};\mathrm{W}_{\Phi_{j}})\), where \(\mathbf{h}^{i}_{j}:=(\mathbf{h}^{i,1}_{j},\,\ldots,\,\mathbf{h}^{i,l}_{j},\,\ldots)\) with \(\mathbf{h}^{i,l}_{j}\) denoting the feature vector from \(l\)-th layer of sub-transformer \(\Phi_{j}\). In super- transformers, the model (maxnet) having the largest configuration, \(\Phi_{1}\), acts as the teacher and knowledge distillation loss for all other sub-transformers w.r.t the teacher is defined as

\[\ell_{\mathcal{RD}}[z^{i}_{\Phi_{j}},z^{i}_{\Phi_{1}}]=\text{KL}\lfloor\, \boldsymbol{\sigma}(z^{i}_{\Phi_{j}}/t),~{}\boldsymbol{\sigma}(z^{i}_{\Phi_{1} }/t)\,],\,\forall j>1,\]

where \(\text{KL}\lfloor\cdot,\cdot\rfloor\) denotes the standard KL divergence between two probability vectors, and \(t\geq 1\) is a hyper-parameter called the temperature. Let \(d_{j}\) denote the embedding dimension (hidden size) in sub-transformer \(\Phi_{j}\). We compute feature based distillation loss by projecting features \(\mathbf{h}^{i,l}_{\Phi_{j}}\in\mathbb{R}^{d_{j}}\) to a low-dimensional space \(\mathbb{R}^{d_{\rm low}}\):

\[\ell_{\mathcal{FD}}[\mathbf{h}^{i}_{\Phi_{j}},\mathbf{h}^{i}_{\Phi_{1}}]=\sum_ {l}\beta^{l}_{j}\,\|\mathbf{U}^{l}_{j}\mathbf{h}^{i,l}_{\Phi_{j}}-\mathbf{U}^{ l}_{1}\mathbf{h}^{i,g_{j}(l)}_{\Phi_{1}}\|^{2}_{2},\]

where \(g_{j}\) maps each layer index of the sub-transformer configuration \(\Phi_{j}\) to that of the super-transformer (/ maxnet \(\Phi_{1}\)). In this paper, we propose to share the maxnet's feature projection matrices \(\{\mathbf{U}^{l}_{1}\in\mathbb{R}^{d_{\rm low}\times d_{1}}\}\) across all sub-transformer models. We do so by slicing the matrices \(\{\mathbf{U}^{l}_{1}\}\):

\[\mathbf{U}^{l}_{j}\;:=\;[\,\mathbf{U}^{g_{j}(l)}_{1}]_{\Phi_{j}}\in\mathbb{R}^ {d_{\rm low}\times d_{j}}, \tag{11}\]

where the operation \([\,\,]_{\Phi_{j}}\) selects appropriate subset of columns depending on the configuration \(\Phi_{j}\). To reduce the number of user-chosen hyper-parameters, we propose the following hyperparameter sharing: \(\beta^{l}_{j}:=\beta_{g_{j}(l)},\,\forall j,\,l=1,2,\ldots\). Thus, apart from setting fewer hyper-parameters, one needs to learn only maxnet's feature projection matrices \(\{\mathbf{U}^{l}_{1}\,:\,l=1,2,\ldots\}\), making feature distillation in a super-transformer setting computationally efficient. Additionally, we save computation through use of features only from a fixed subset of maxnet layers for distillation across all sub-transformers: i.e., we use the following subset of maxnet layers: \(\{\,g_{\min}(l)\,:\,l=1,\dots,\,L_{\min}\}\), where \(L_{\min}\) denotes the number of transformer layers in the smallest sub-transformer \(\Phi_{\min}\) and \(g_{\min}\) maps layer indices of \(\Phi_{\min}\) to that of maxnet \(\Phi_{1}\).

## 6 Results on Encoder and Decoder LLMs

We report performance on encoder tasks using GLUE (Wang et al., 2018) with BERT\({}_{base}\) as the teacher model \(\Phi_{\texttt{Tch}}\). For decoder LLMs, we use Santacodet (Allal et al., 2023) and Codellama7B (Roziere et al., 2023) on a python coding task using _bigcode/the-stack_ data (Koectkov et al., 2022).

 

report performance of the sub-transformer models at the end of stage \(s=2\). On GLUE, we use the train set for fine-tuning and the dev set for accuracy evaluation. For santacoder, we evaluate performance using HumanEval Chen et al. (2021) and report pass@1 scores. All experiments were conducted using PyTorch on a single Nvidia A100 (40GB) GPU. Additional details on the experiment settings are provided in the Appendix.

### Performance of Encoder Models

We compare performance of encoder models obtained with the MLFS approach against a static, fixed model (BERT base) from Zhang et al. (2021); Hou et al. (2020), two popular distilled variants of the fixed model: TinyBERT Jiao et al. (2020) and DistilBERT Sanh et al. (2019), and models trained using existing super-transformer methods DynaBERT Hou et al. (2020). Figure 1 shows the performance of the palette of models, from a 45M param. minnet to full-size 110M maxnet. Encoder models produced by MLFS are at par or better than much costlier methods. Results of PDERT, BERT-PKD are from Zhang et al. (2021), static BERT from Zhang et al. (2021) for all except MRPC for which we use Hou et al. (2020). Note that TinyBERT performs data augmentation leading to higher accuracy but much longer computation time. We do not perform data augmentation for fairness of the comparison to the other methods. The main observation is that MLFS provides accurate, smaller encoder models at 1/4 the size of the teacher and 1/3 its runtime latency on a single GPU.

Ablation Study on Gradient ScalingIn super-net training, the weights of maxnet and subnets are shared and trained simultaneously. The maxnet tends to converge and overfit earlier than smaller subnets. The different convergence rates renders selecting a single supernet checkpoint for all networks difficult. Gradient scaling solves this by speeding up convergence of the smaller subnets to match that of the larger subnets or the maxnet. Fig. 2 shows that gradient scaling improves minnet convergence, indicated by lower minnet loss.

Ablation Study on Rank in MLFSFinally, in Fig. 3, we examines the impact of rank \(r\) of the

Figure 1: Performance of task-specific BERT models produced by MLFS vs. other methods on 3 GLUE data sets.

Figure 3: Ablation study on MLFS rank of \(A,B\). Maxnet (top: blue), minnet (bottom: green), and average of two medium-sized subnets (middle: orange). Rank \(r=8\) is optimal for small and medium subnets.

Figure 2: Ablation study on gradient scaling: MLFS minnet convergence is improved using gradient scaling.

 

matrices \(A,B\) on performance. Note that the actual number of parameters fine-tuned vary as we vary the rank \(r\). The aim is to provide good results for the smaller networks. Here, rank \(r=8\) works well across the GLUE data sets. Therefore, we use rank \(r=8\) for \(A,B\) for all other MLFS experiments. From the scale of the y-axis in 3, observe that MLFS is not overly sensitive to the chosen rank.

### Performance of Decoder Models

Turning now to decoder models, we consider two code-pre-trained LLMs, Santacoder Allal et al. (2023) and Codellama7B Roziere et al. (2023). We evaluate a custom 0.7B parameter Santacoder model obtained from the 1.1B teacher. Due to an inability to fine-tune on the full 24M coding examples, we use up to 1.2M. Fig. 4 shows that MLFS pass@1 improves rapidly as number of tokens increases from a low 10k to 400k to 1.2M examples, only 5% of the 24M examples. Table 1 shows analogous results with 3 small MLFS models. The improvement in pass@1 indicates that the smaller models retain the ability to learn from the larger teacher. Again, from Table 2, we see that smaller models produced by MLFS from _CodeLlama-7B-Python_ retain their ability to learn and improve quickly as the number of examples increases. Note that the full data set includes 24M examples; MLFS achieves nearly 75% of the performance of fullsize CodeLlama after less than 2% of the examples.

Contrary to encoder models, the compression levels that retain sufficient performance of the teacher with decoders is less. While MLFS retains accuracy performance of encoder models at 1/4 the size of the teacher, the decoder models are reduced to at most 2/3 the teacher's size.

MLFS slicing of the teacher model can, however, benefit decoder models by reducing substantially the training/fine-tuning time needed compared to a randomly-initialised model, as shown in Fig. 5 on Santacodor sliced from 1.1B to 0.7B. In other words, when a smaller model is required for edge inference, one can train it from a random initialisation, or slice from a teacher as does MLFS, and train starting from the sliced weights. The latter significantly reduces training time as seen in the validation loss curves. See Samragh et al. (2023) for a similar observation.

## 7 Perspectives

Enterprise users require an efficient way to fine-tune LLMs for inference on edge devices of many sizes. We developed MLFS for such edge deployment scenarios. We demonstrate its benefits on encoder LLMs. We show the limitation of compressing decoder LLMs to a comparable degree; however, MLFS offers significant gains for smaller decoder training/fine-tuning by slicing from a larger pre-trained teacher.

Figure 4: Performance of MLFS on a custom Santacodor 0.7B model using 10K/400K/1.2M training examples.

Figure 5: Convergence comparison of validation loss while fine-tuning a custom model from random vs using MLFS. MLFS achieves low validation loss much faster.

 



## References

* L. Ben Allal, R. Li, D. Koccetkov, C. Mou, C. Akiki, C. Munoz Ferrandis, N. Muennighoff, M. Mishra, A. Gu, M. Dey, et al. (2023)Santacoder: don't reach for the stars!. arXiv preprint arXiv:2301.03988. Cited by: SS4.
* E. B. Zaken, Y. Goldberg, and S. Ravfogel (2022)BitFit: simple parameter-efficient fine-tuning for transformer-based masked language-models. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers), Dublin, Ireland, pp. 1-9. External Links: Link Cited by: SS5.
* A. Brandare, V. Sripathi, D. Karkada, V. Menon, S. Choi, K. Datta, and V. Saletore (2019)Efficient 8-bit quantization of transformer neural machine language translation model. arXiv preprint arXiv:1906.00532. Cited by: SS1.
* T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, et al. (2020)Language models are few-shot learners. Advances in neural information processing systems33, pp. 1877-1901. Cited by: SS1.
* H. Cai, C. Gan, T. Wang, Zhekai Zhang, and S. Han (2019)Once-for-all: train one network and specialize it for efficient deployment. arXiv preprint arXiv:1908.09791. Cited by: SS1.
* H. Cai, L. Zhu, and S. Han (2018)Proxylessnas: direct neural architecture search on target task and hardware. arXiv preprint arXiv:1812.00332. Cited by: SS1.
* M. Chen, J. Tworek, H. Jun, Q. Yuan, H. Pande de Oliveira Pinto, J. Kaplan, H. Edwards, Y. Burda, N. Joseph, G. Brockman, A. Ray, R. Puri, G. Krueger, M. Petrov, H. Khlaaf, G. Sastry, P. Mishkin, B. Chan, S. Gray, N. Ryder, M. Pavlov, A. Power, L. Kaiser, M. Bavarian, C. Winter, P. Tillet, F. Such, D. Cummings, M. Plappert, F. Chantzis, E. Barnes, A. Herbert-Voss, W. Hebigen Guss, A. Nichol, A. Paino, N. Tezak, J. Tang, I. Babuschkin, S. Balaji, S. Jain, W. Saunders, C. Hesse, A. N. Carr, J. Leike, J. Achimam, V. Misra, E. Morikawa, A. Radford, M. Knight, M. Brundage, M. Murati, K. Mayer, P. Welinder, B. McGrew, D. Amodei, S. McCandlish, I. Sutskever, and Wojciech Zaremba (2021)Evaluating large language models trained on code. arXiv preprint arXiv:2107.03374. Cited by: SS1.
* M. Chen, H. Peng, J. Fu, and H. Ling (2021)Autoformer: searching transformers for visual recognition. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 12270-12280. Cited by: SS1.
* T. Dettmers, A. Pagnoni, A. Holtzman, and L. Zettlemoyer (2023)Qlora: efficient finentuning of quantized lllms. arXiv preprint arXiv:2305.14314. Cited by: SS1.
* P. Dong, X. Niu, L. Li, L. Xie, W. Zou, T. Ye, Z. Wei, and H. Pan (2022)Prior-guided one-shot neural architecture search. arXiv preprint arXiv:2206.13329. Cited by: SS1.
* J. Gao, H. Xu, H. Shi, X. Ren, L. Philip, X. Liang, X. Jiang, and Z. Li (2022)AutoBERT-zero: evolving bert backbone from scratch. In Proceedings of the AAAI Conference on Artificial Intelligence, Vol. 36-10, pp. 10663-10671. Cited by: SS1.
* Y. Gu, L. Dong, F. Wei, and M. Huang (2023)Knowledge distillation of large language models. arXiv preprint arXiv:2306.08543. Cited by: SS1.
* Z. Guo, X. Zhang, H. Mu, W. Heng, Z. Liu, Y. Wei, and J. Sun (2020)Single path one-shot neural architecture search with uniform sampling. In Computer Vision-ECCV 2020: 16th European Conference, Glasgow, UK, August 23-28, 2020, Proceedings, Part XVI 16, pp. 544-560. Cited by: SS1.
* L. Hou, Z. Huang, L. Shang, X. Jiang, X. Chen, and Q. Liu (2020)Dynabert: dynamic bert with adaptive width and depth. Advances in Neural Information Processing Systems33, pp. 9782-9793. Cited by: SS1.
* C. Hsieh, C. Li, C. Yeh, H. Nakhost, Y. Fujii, A. Ratner, R. Krishna, C. Lee, and T. Pfister (2023)Distilling step-by-step! outperforming larger language models with less training data and smaller model sizes. arXiv preprint arXiv:2305.02301. Cited by: SS1.
* E. J. Hu, Y. Shen, P. Wallis, Z. Allen-Zhu, Y. Li, S. Wang, L. Wang, and W. Chen (2022)LoRA: low-rank adaptation of large language models. In International Conference on Learning Representations, Cited by: SS1.
* G. Jawahar, H. Yang, Y. Xiong, Z. Liu, D. Wang, F. Sun, M. Li, A. Pappu, B. Oguz, M. Abdulhg-Maged, L. V. S. Lakshmanan, R. Krishnamoorhi, and V. Chandra (2023)Mixture-of-supernets: improving weight-sharing supernet training with architecture-routed mixture-of-experts. arXiv preprint arXiv:2306.04845. Cited by: SS1.
* Xiaoqi Jiao, Y. Yin, L. Shang, X. Jiang, X. Chen, L. Li, F. Wang, and Q. Liu (2020)TinyBERT: distilling BERT for natural language understanding. In Findings of the Association for Computational Linguistics: EMNLP 2020, Online, pp. 4163-4174. Cited by: SS1.
* D. Kocetkov, R. Li, L. Ben Allal, J. Li, C. Mou, C. Munoz Ferrandis, Y. Jernite, M. Mitchell, S. Hughes, T. Wolf, D. Bahdanau, L. von Werra, and

Harm de Vries. 2022. The stack: 3 tb of permissively licensed source code. _arXiv preprint arXiv:2211.15533_.
* Kundu et al. (2023) Achintya Kundu, Laura Wynter, Rhui Dih Lee, and Luis Angel D. Bathen. 2023. Transfer-once-for-all: AI model optimization for edge. In _IEEE International Conference on Edge Computing and Communications, EDGE 2023, Chicago, IL, USA, July 2-8, 2023_, pages 26-35. IEEE.
* Lan et al. (2019) Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Soricut. 2019. Albert: A lite bert for self-supervised learning of language representations. _arXiv preprint arXiv:1909.11942_.
* Lou et al. (2021) Wei Lou, Lei Xun, Amin Sabet, Jia Bi, Jonathon Hare, and Geoff V Merrett. 2021. Dynamic-ofa: Runtime dnn architecture switching for performance scaling on heterogeneous embedded platforms. In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_, pages 3110-3118.
* Ma et al. (2019) Xindian Ma, Peng Zhang, Shuai Zhang, Nan Duan, Yuexian Hou, Ming Zhou, and Dawei Song. 2019. A tensorized transformer for language modeling. _Advances in neural information processing systems_, 32.
* McCarley et al. (2019) JS McCarley, Rishav Chakravarti, and Avirup Sil. 2019. Structured pruning of a bert-based question answering model. _arXiv preprint arXiv:1910.06360_.
* Mukherjee and Awadallah (2020) Subhabrata Mukherjee and Ahmed Awadallah. 2020. Xtremedistil: Multi-stage distillation for massive multilingual models. _arXiv preprint arXiv:2004.05686_.
* Mukherjee et al. (2021) Subhabrata Mukherjee, Ahmed Hassan Awadallah, and Jianfeng Gao. 2021. Xtremeditsitransformers: Task transfer for task-agnostic distillation. _arXiv preprint arXiv:2106.04563_.
* Raffel et al. (2020) Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. 2020. Exploring the limits of transfer learning with a unified text-to-text transformer. _The Journal of Machine Learning Research_, 21(1):5485-5551.
* Real et al. (2019) Esteban Real, Alok Aggarwal, Yanping Huang, and Quoc V Le. 2019. Regularized evolution for image classifier architecture search. In _Proceedings of the aaai conference on artificial intelligence_, volume 33-01, pages 4780-4789.
* Roziere et al. (2023) Baptiste Roziere, Jonas Gehring, Fabian Gloeckle, Sten Sootla, Itai Get, Xiaoqing Ellen Tan, Yossi Adi, Jingyu Liu, Tal Remez, Jeremy Rapin, Artym Kozhevnikov, Ivan Evtimov, Joanna Bitton, Manish Bhatt, Cristian Canton Ferrer, Aaron Grattafiori, Wenhan Xiong, Alexandre Defossez, Jade Copet, Faisal Azhar, Hugo Touvron, Louis Martin, Nicolas Usunier, Thomas Scialom, and Gabriel Synnaeve. 2023. Code llama: Open foundation models for code.
* Samragh et al. (2023) Mohammad Samragh, Mehrdal Farajtabar, Sachin Mehta, Raviteja Vemulapalli, Fatash Faghri, Devang Naik, Oncel Tuzel, and Mohammad Rastegari. 2023. Weight subcloning: direct initialization of transformers using larger pretrained ones.
* Sanh et al. (2019) Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. 2019. Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter. _arXiv preprint arXiv:1910.01108_.
* Shen et al. (2020) Sheng Shen, Zhen Dong, Jiayu Ye, Linjian Ma, Zhewei Yao, Amir Gholami, Michael W Mahoney, and Kurt Keutzer. 2020. Q-bert: Hessian based ultra low precision quantization of bert. In _Proceedings of the AAAI Conference on Artificial Intelligence_, volume 34-05, pages 8815-8821.
* Sun et al. (2019) Siqi Sun, Yu Cheng, Zhe Gan, and Jingjing Liu. 2019. Patient knowledge distillation for bert model compression. _arXiv preprint arXiv:1908.09355_.
* Voita et al. (2019) Elena Voita, David Talbot, Fedor Moiseev, Rico Senr nich, and Ivan Titov. 2019. Analyzing multihead self-attention: Specialized heads do the heavy lifting, the rest can be pruned. _arXiv preprint arXiv:1905.09418_.
* Wang et al. (2018) Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman. 2018. GLUE: A multi-task benchmark and analysis platform for natural language understanding. In _Proceedings of the 2018 EMNLP Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP_, pages 353-355, Brussels, Belgium. Association for Computational Linguistics.
* Wang et al. (2021) Dilim Wang, Meng Li, Chengyue Gong, and Vikas Chandra. 2021. Attentiveness: Improving neural architecture search via attentive sampling. In _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_, pages 6418-6427.
* Wang et al. (2020) Hanrui Wang, Zhaghao Wu, Zhijian Liu, Han Cai, Ligeng Zhu, Chuang Gan, and Song Han. 2020. Hat: Hardware-aware transformers for efficient natural language processing. _arXiv preprint arXiv:2005.14187_.
* Wang et al. (2022) Rui Wang, Qibing Bai, Junyi Ao, Long Zhou, Zhixiang Xiong, Zhihua Wei, Yu Zhang, Tom Ko, and Haizhou Li. 2022. Lighthubert: Lightweight and configurable speech representation learning with once-for-all hidden-unit bert. _arXiv preprint arXiv:2203.15610_.
* Xu et al. (2021) Jin Xu, Xu Tan, Renqian Luo, Kaitao Song, Jian Li, Tao Qin, and Tie-Yan Liu. 2021. Nas-bert: task-agnostic and adaptive-size bert compression with neural architecture search. In _Proceedings of the 27th ACM SIGKDD Conference on Knowledge Discovery & Data Mining_, pages 1933-1943.
* Yu et al. (2020) Jiahui Yu, Pengchong Jin, Hanxiao Liu, Gabriel Bender, Pieter-Jan Kindermans, Mingxing Tan, Thomas Huang, Xiaodan Song, Ruoming Pang, and Quoc Le. 2020. Bignas: Scaling up neural architecture search 

with big single-stage models. In _Computer Vision-ECCV 2020: 16th European Conference, Glasgow, UK, August 23-28, 2020, Proceedings, Part VII 16_, pages 702-717. Springer.
* Zafrir et al. (2019) Ofir Zafrir, Guy Boudoukh, Peter Izsak, and Moshe Wasserblat. 2019. Q8bert: Quantized 8bit bert. In _2019 Fifth Workshop on Energy Efficient Machine Learning and Cognitive Computing-NeurIPS Edition (EMC2-NIPS)_, pages 36-39. IEEE.
* Zhang et al. (2023) Qingru Zhang, Minshuo Chen, Alexander Bukharin, Pengcheng He, Yu Cheng, Weizhu Chen, and Tuo Zhao. 2023. Adaptive budget allocation for parameter-efficient fine-tuning. _arXiv preprint arXiv:2303.10512_.
* Zhang et al. (2021) Shaokun Zhang, Xiawu Zheng, Chenyi Yang, Yuchao Li, Yan Wang, Fei Chao, Mengdi Wang, Shen Li, Jun Yang, and Rongrong Ji. 2021. You only compress once: Towards effective and elastic bert compression via exploit-explore stochastic nature gradient. _arXiv preprint arXiv:2106.02435_.
* Zhu et al. (2019) Hui Zhu, Zhulin An, Chuanguang Yang, Kaiqiang Xu, Erhu Zhao, and Yongjun Xu. 2019. Eena: efficient evolution of neural architecture. In _Proceedings of the IEEE/CVF International Conference on Computer Vision Workshops_, pages 0-0.
* Zhu et al. (2023) Xunyu Zhu, Jian Li, Yong Liu, Can Ma, and Weiping Wang. 2023. A survey on model compression for large language models. _arXiv preprint arXiv:2308.07633_.
* Zoph and Le (2016) Barret Zoph and Quoc V Le. 2016. Neural architecture search with reinforcement learning. _arXiv preprint arXiv:1611.01578_.

## Appendix A Details of Experimental Set-up

Following Hu et al. (2022), we use the fine-tuned MNLI checkpoint to initialize the model weights for experiments on small data sets such as RTE and MRPC. In MLFS, the Low rank matrices are added on the QKV vectors and the intermediate size of feed-forward network (FFN) layers. We set \(\beta_{l}=0.1~{}\forall l\) in feature distillation loss and choose distillation factor \(\alpha=0.9\). For training, we use a maximum sequence length of 128; effective batch size of 128 for QQP, MNLI, QNLI, and 64 for the other data sets. Training is done for a maximum of 8 epochs for all GLUE data sets except SST-2 for which we allocate maximum 3 epochs. We set an initial learning rate of \(5e^{-4}\) for QNLI & MNLI, and \(1e^{-3}\) for other GLUE data sets. We use rank \(r=8\) for the low rank matrices \(A,B\) unless mentioned otherwise. We choose gradient scaling hyper-parameter \(\gamma=1\) for SST-2 and \(\gamma=2\) for all other data sets.

## Appendix B Additional Experimental Results

First, we present additional results on distilling Santacoder-1.1B model. In Fig. 6, we compare HumanEval performance of a 0.7B Santacoder model fine-tuned through full fine-tuning (FT) from random initialisation vs. full-rank (non-LoRA) MLFS with (\(\alpha=0.9\)) and without (\(\alpha=0\)) distillation. The improvement in the evaluation numbers is remarkable even after fine-tuning on up to only 5% of the examples. In Fig. 7, we also show better convergence of validation loss on the Santacoder 0.7B for MLFS with distillation loss (\(\alpha>0\)). This demonstrates the benefit of MLFS distillation as compared to full MLFS fine tuning of the sliced model.

Finally, in Fig. 8, we show performance of a spectrum of models distilled from BERT\({}_{base}\) using MLFS on 3 more GLUE data sets: QNLI, QQP, and MNLI.

Figure 6: Superior performance of supernet training compared to other full fine-tuning based approaches on three data sets with 10K/400K/1.2M examples.

Figure 7: Convergence comparison of validation loss while fine-tuning a custom model using MLFS with/without distillation.

 



## 6 Conclusion

Figure 8: Performance of task-specific BERT models produced by MLFS vs. other methods on 3 GLUE data sets.