SortedNet: A Scalable and Generalized Framework for Training Modular
  Deep Neural Networks
Deep neural networks (DNNs) must cater to a variety of users with different
performance needs and budgets, leading to the costly practice of training,
storing, and maintaining numerous user/task-specific models. There are
solutions in the literature to deal with single dynamic or many-in-one models
instead of many individual networks; however, they suffer from significant
drops in performance, lack of generalization across different model
architectures or different dimensions (e.g. depth, width, attention blocks),
heavy model search requirements during training, and training a limited number
of sub-models. To address these limitations, we propose SortedNet, a
generalized and scalable training solution to harness the inherent modularity
of DNNs. Thanks to a generalized nested architecture (which we refer as
\textit{sorted} architecture in this paper) with shared parameters and its
novel update scheme combining random sub-model sampling and a new gradient
accumulation mechanism, SortedNet enables the training of sub-models
simultaneously along with the training of the main model (without any
significant extra training or inference overhead), simplifies dynamic model
selection, customizes deployment during inference, and reduces the model
storage requirement significantly. The versatility and scalability of SortedNet
are validated through various architectures and tasks, including LLaMA, BERT,
RoBERTa (NLP tasks), ResNet and MobileNet (image classification) demonstrating
its superiority over existing dynamic training methods. For example, we
introduce a novel adaptive self-speculative approach based on sorted-training
to accelerate large language models decoding. Moreover, SortedNet is able to
train 160 sub-models at once, achieving at least 96\% of the original model's
performance.
Mojtaba Valipour, Mehdi Rezagholizadeh, Hossein Rajabzadeh, Parsa Kavehzadeh, Marzieh Tahaei, Boxing Chen, Ali Ghodsi
# SortedNet, a Place for Every Network and Every Network in its Place:

###### Abstract

As the size of deep learning models continues to grow, finding optimal models under memory and computation constraints becomes increasingly more important. Although usually the architecture and constituent building blocks of neural networks allow them to be used in a modular way (i.e. using the sub-networks of a given network after training), their training process is not aware of this modularity. Consequently, conventional neural network training lacks the flexibility to adapt the computational load of the model during inference. This paper proposes, SortedNet, a generalized and scalable solution to harness the inherent modularity of deep neural networks across various dimensions (e.g. width, depth, blocks) for efficient dynamic inference. Our training considers a nested architecture for the sub-models with shared parameters and train them together with the main model in a sorted and probabilistic manner. This sorted training of sub-networks enables us to scale the number of sub-networks to hundreds using a single round of training. We utilize a novel updating scheme during training that combines random sampling of sub-networks with gradient accumulation to improve training efficiency. Furthermore, the sorted nature of our training leads to a search-free sub-network selection at inference time; and the nested architecture of the resulting sub-networks leads to minimal storage requirement and efficient switching between sub-networks at inference. Our general dynamic training approach is demonstrated across various architectures and tasks, including large language models and pre-trained vision models. Experimental results show the efficacy of the proposed approach in achieving efficient sub-networks while outperforming state-of-the-art dynamic training approaches. Our findings demonstrate the feasibility of training up to 160 different sub-models simultaneously, showcasing the extensive scalability of our proposed method while maintaining 96% of the model performance.

1University of Waterloo

2Huawei Noah's Arc Lab

{mojtaba.valipour, hossein.rajabzadeh, ali.ghodsi}@uwaterloo.ca,

{mehdi.rezagholizadeh, marzieh.tahaei, boxing.chen}@huawei.com

## 1 Introduction

_"For every minute spent organizing, an hour is earned." - Benjamin Franklin._

There has been a remarkable growth in the size of deep neural networks. Nevertheless, the computation/memory resources allocated to a model at inference depends on the specific hardware availability and the accuracy/time requirements of applications, whether deployed in the cloud or on edge devices.

In particular, the computational burden from concurrent processes and battery limitations can significantly impact the resources allocated to a neural network. Moreover, in the era of gigantic pre-trained models the computational demand can vary from task to task. Therefore there is a growing demand for models that can adapt themselves to such dynamic conditions. Unfortunately, conventional neural network training, with its fixed architecture, falls short in adaptive adjusting of the computational load at inference time.

On the other hand, deep neural networks demonstrate modular architectures along various dimensions, like layers and blocks across depth, and neurons and channels and even attention heads along width. This inherent modularity enables the extraction of sub-networks with the same shape as the original model. However, the current training methods fail to effectively leverage this modularity, resulting in limited practical advantages of sub-networks in the final product. Consequently, the performance of these sub-networks falls short compared to the main model, making their deployment during inference impractical.

Hence, the challenge lies in harnessing the full potential of modularity in deep neural networks, allowing for the efficient utilization of sub-networks to enhance performance and enable practical deployment in real-world scenarios.

Recent works have proposed a variety of approaches for training dynamic models. These approaches While effective often use a sophisticated training process combined knowledge distillation Hou et al. (2020), architecture modification Nunez et al. (2023), and redundant subnetwork optimization Fan et al. (2019). Although not explicitly mentioned, an important ingredient shared by all these methods is the attempt to implicitly sort sub-networks along a specific dimension with respect to computation/accuracy. Limiting dynamicity to one or two dimension while leaving other dimensions intact can lead to suboptimal sub-networks. Inspired by these works (Valipour et al., 2023), Rippel et al. (2014)), in this paper we explore how sorting generalized to all dimensions can provide many in one efficient dynamic models. Our solution take advantage from an intrinsic nested sub-networks which are sorted monotonically from bottom to top (i.e. we have only a single instance from each sub-network). This sorted configuration with shared parameters enforces a regular order and consistency in the knowledge learned by sub-networks. Sorting these sub-networks based
on their computation/accuracy characteristics presents the most optimal solution. By organizing the model in this manner, extracting the desired sub-network becomes a search-free process. The use of a predefined sorting order ensures that each targeted sub-network possesses a unique computation overhead, effectively removing optimization of redundant sub-networks from training. In the resulting nested architecture with shared parameters in a _sorted_ manner, each smaller (shallower/narrower) model is a sub-networks of a larger (deeper/wider) model. This will lead to models with sorted accuracy, latency and importance.

To achieve this sorted architecture, during training, we propose a novel updating scheme that combines random sampling of sub-networks with gradient accumulation in order to further reduce the cost of training. With one single training our proposed method can yield multiple models with different capacities.

Our general dynamic training approach is applicable to any architecture without necessitating any modifications to the original model. The proposed nested architecture offers several benefits, including minimal storage requirements and efficient switching between various computation budgets during inference.

Through comprehensive empirical studies across different architectures, tasks and dynamicity along various dimensions ranging from width and depth to attention head and embedding layer we show the superiority and generalizability of our proposed methods over state of the art dynamic training methods.

To summarize, the main contributions of this paper are:

* Introduction of nested sorted network training, which leverages efficient sub-network training through a combination of gradient accumulation and sub-network sampling.
* Generalization of dynamicity to multiple dimensions through stochastic sampling of sub-networks during training.
* Outperforming state-of-the-art methods in dynamic training on CIFAR10 [11]. Furthermore, scaling the number of sub-networks to 160 and showcasing the efficacy of our single round training method.
* Demonstrating the effectiveness of the proposed method on Large pre-trained language models by dynamic training of the BERT model.

Figure 1: SortedNet: The overall diagram of our proposed method. During training, in each iteration, we sample from a pre-defined random distribution which will help us to optimize the sub-models as well as the original model.

Figure 2: Comparing SortedNet and Once For All: on a hypothetical 5-layer network, we show how the sub-network selection strategy of SortedNet differs from the Once-for-All [10] approach.
## 2 Related Work

In this section, we briefly review the most relevant existing works to our SortedNet idea. A summary of these solutions and how they are different from each other can be found in Table 1.

**Slimmable Networks [23]**Slimmable networks is an idea to train a single neural network in a way that it can be deployed with adjustable width at the inference time. This solution was proposed particularly for CNN architectures and thus, careful consideration of the batch normalization module for various width sizes is necessary. In this regard, they use switchable batch normalization which leads to additional trainable parameters. In contrast to slimmable networks, our SortedNet are architecture agnostic and work in both depth and width dimensions.

**Early Exit [20]**Early exit refers to a technique which adds a classifier to intermediate layers of an already trained neural network. While the parameters of the main model are frozen, the parameters of the classifiers are updated in a separate fine-tuning process. In this approach, each of the classifiers and their subsequent network can be treated as an independent sub-model. While this solution is relatively straightforward, the performance of the sub-models lags significantly behind that of the main model.

**Dyna-BERT [1]**Dyna-BERT is a dynamic compression method for pre-trained BERT models, allowing flexible adjustments of the size of the model along depth and width at the inference time. While the introduced objective in the DynaBERT paper has some overlap with ours, there are several key differences: first, in DynaBERT, Only a few subsets of the model are functional (while our SortedNet do not have such an assumption); second, DynaBERT needs an already trained teacher and uses knowledge distillation but our technique does not need KD; third, DynaBERT needs search during both training and inference while our solution is _search-free_; Lastly, DynaBERT's applicability is architecture-dependent, whereas our approach is not.

**Layer-drop [20]**Layer-drop is a structured dropout solution at the training time which allows layer pruning at the inference time. Similar to DynaBERT, this solution is applied to pre-trained language models; however, in contrast to DynaBERT, Layer-drop only targets the depth of neural networks and not their width. In Layer-drop, there is no fixed training pattern and any layer can be dropped with a certain probability, which is referred to as drop rate. At the inference time, the number of active layers can be adjusted by the drop-rates that are seen during the training time of that network (i.e. to achieve the best performance on any other drop-rate value, the network needs to be re-trained.). Layer-drop works only in depth while our solution works for both depth and width. Moreover, Layer-Drop requires specific search patterns for dropping layers at the inference time and training time, whereas our solution is search free.

**Once-for-All [1]**Once-for-all(OFA) targets efficient inference across different devices by first training an OFA network which supports many sub-networks with varying latency/accuracy characteristics ; it then searches among the feasible sub-networks according to the accuracy and latency requirements of their target device. OFA has a progressive training nature i.e. it goes from the largest model to the smaller sub-networks. OFA is different from our solution from the following aspects: first, it needs teacher and knowledge distillation; second, OFA requires a separate Neural Architecture Search (NAS) at the inference time; third, OFA is not architecture agnostic (their solution is for CNN-based neural networks while our SortedNet works for both CNNs and Transformers). Moreover, OFA is different from our solution in terms of the sub-network selection strategy. While our SortedNet selects sub-networks in a sorted manner, OFA does not have any particular assumption for sorting sub-networks (see Fig. 2 for more details).

**Learning Compressible Subspace [20]**Learning Compressible Subspace (LCS) is an adaptive compression technique based on training compressible subspace of neural networks (using a linear convex combination of two sets of weights for the network). While LCS does not require any re-training at the inference time, this solution has some other limitations including: first, it needs double memory at the training time; second, the choices of initial weights and the compression function are unclear and arbitrary (left as a hyper-parameter); third, it is only tried on CNNs; forth,

\begin{table}
\begin{tabular}{|l|c c c c c c c|} \hline \multirow{2}{*}{**Method**} & \multicolumn{2}{c|}{**Sub-**} & \multicolumn{2}{c|}{**\# of**} & \multicolumn{2}{c|}{**No**} & \multirow{2}{*}{**Target Dim.**} & \multirow{2}{*}{**Architecture**} \\  & **Networks: Config.(\#)** & & & & & **Free** & **Params** & **Re-training** \\ \hline \hline Early Exit [20] & Sorted (Few) & Low & ✓ & ✓ & \(\left|\theta\right|\) & ✗ & Depth & Transformer \\ \hline Layer Drop [20] & Random & \multirow{2}{*}{Low} & \multirow{2}{*}{\(\mathcal{K}\)} & \multirow{2}{*}{\(\mathcal{K}\)} & \multirow{2}{*}{\(\left|\theta\right|\)} & ✗ & Depth & Transformer \\ Joulin [2] & & & & & & & & \\ \hline DynaBERT [1] & Sorted \& & & & & & & \\ Random (Few) & High & ✗ & ✗ & \(2\left|\theta\right|\) & ✗ & Depth \& Width & Transformer \\ \hline Once for All [1] & Nested & High & ✗ & ✗ & \(\left|\theta\right|\,\mathrm{or}\,2\left|\theta\right|\) & ✗ & General & CNN \\ \hline LCS [20] & Arbitrary & \multirow{2}{*}{High} & ✓ & ✓ & \(\left|\theta\right|\,\mathrm{or}\,2\left|\theta\right|\) & ✓ & General & CNN \\  & (Many) & & & & & & & \\ \hline Slimmable [23] & Sorted (Few) & Moderate & ✓ & ✓ & \(\left|\theta\right|\) & ✓ & Width & CNN \\ \hline
**SortedNet (Ours)** & Sorted (Many) & High & ✓ & ✓ & \(\left|\theta\right|\) & ✓ & General & 
\begin{tabular}{c} CNN \& Transformer \\ \end{tabular} \\ \hline \end{tabular}
\end{table}
Table 1: Comparison of different existing related work and distinguishing our solution
similar to Layer-drop intermediate sub-networks are trained randomly which will make the performance of the target model sub-optimal.

## 3 Methodology

### A Generalized and Scalable View

In the related work section, we have discussed several approaches concerning the training of many-in-one networks. These approaches differ in terms of their target architecture, training loss, number of training parameters, configuration of the sub-networks (random or sorted), the number of trained sub-models, and reliance on search or re-training before deployment. Our SortedNet method can be viewed as a simple general and scalable version of these existing solutions. These generalization and scalability have been mostly resulted from the sorted configuration of sub-networks with their shared parameters and our stochastic training. To the best of our knowledge this is the first work which has scaled training of sorted nested sub-networks to various dimensions.

### SortedNet: Towards a Generalized Solution for Training Many-in-One Networks

While many of deep neural network architectures are modular in design (using similar layers such as in Transformers [20] or blocks such as in MobileNet [2]) but this modularity is not preserved during training. In this subsection, we introduce our SortedNet solution which aims at training a generalized and scalable many-in-one networks. In order to train many-in-one networks, we need to specify a few design choices: first, how to form the sub-networks and their configurations; second, what are the target architectures; and third, how to train the sub-networks along with the main model.

Designing the Sub-networksSortedNet impose an inductive bias on the training based on the assumption that the parameters of sub-networks across each dimension have a concentric (or onion shape) architecture with shared parameters in a _sorted_ manner. This sorted configuration with shared parameters enforces a regular order and consistency in the knowledge learned by sub-networks (see Fig. 1).

Let's consider a many-in-one neural network \(f(x;\theta(n))\) with the parameters \(\theta(n)\) and the input \(x\) which is comprised of \(n\) sub-networks \(f(x;\theta(i))|_{i=0}^{n-1}\), where \(\theta(i)\) represents the weights of the \(i^{\text{th}}\) sub-model. We define a universal set which contains all unique sub-models: \(\Theta=\{\theta(0),\theta(1),...,\theta(n)\}\).

How to build the sub-models?Suppose that we would like to target \(D=\{Dim_{1},Dim_{2},...,Dim_{K}\}\) many-in-one dimensions in our model. Then, let's start with \(\Theta=\varnothing\) and build the sub-models iteratively. In this regard, at each iteration \(t\) during training, we have a sampling and truncation procedures along any of the targeted dimensions:

\[\theta_{t}^{*}=\cap_{j=1}^{|D|}\theta_{Dim_{j}\downarrow b_{j}^{t }}(n) \tag{1}\] \[\text{where }b_{j}^{t}\sim P_{B_{j}}\] \[\text{IF }\theta_{t}^{*}\notin\Theta:\Theta\leftarrow\Theta\cup \{\theta_{t}^{*}\}\]

where \(Dim_{j}\downarrow b_{j}^{t}\) indicates that we have truncated \(\theta(n)\) along the \(Dim_{j}\) dimension up to the index \(b_{j}^{t}\) at the iteration \(t\). \(b_{j}^{t}\) is sampled from a distribution \(P_{B_{j}}\) with the support set of \(B_{j}=\{1,2,...,d_{j}\}\) to form the \(i^{\text{th}}\) sub-model. \(d_{j}\) refers to the maximum index of the \(j^{\text{th}}\) dimension. This iterative process will be done during training and the set of \(n\) unique sub-models \(\Theta\) will be built.

To illustrate the process better, let's see a simple case such as BERT\({}_{base}\) where we want to make a many-in-one network across the width and depth dimensions, \(D=\{\text{Depth, Width}\}\). In this case, we have 12 layers and the hidden dimension size of 768. Suppose that \(Depth\) corresponds to \(j=1\) and \(Width\) corresponds to \(j=2\) in Eq. 1. For simplicity, let's use a discrete uniform distribution for sampling indices across these two dimensions. To create the first sub-network (\(i=1\)), we need to sample \(b_{1}^{1}\) uniformly from the set of natural numbers in the range of 1 to 12: \(B_{1}=\{1,2,...,12\}\); and we need to sample \(b_{2}^{1}\) from the range of 1 to 768: \(B_{2}=\{1,2,3,...,768\}\). Bear in mind that we can even choose a subset of \(B_{1}\) and \(B_{2}\) as the support set for sampling probability distribution. After these two samplings, we will have two truncated sets of parameters: \(\theta_{Depth\downarrow b_{1}^{1}}\) and \(\theta_{Width\downarrow b_{2}^{1}}\). The intersection of these two truncated parameters will give us the first sub-network: \(\theta_{1}=\theta_{Depth\downarrow b_{1}^{1}}\cap\theta_{Width\downarrow b_{2 }^{1}}\).

Training ParadigmRegular training of neural networks concerns improving the performance of the whole model and usually this training is not aware of the performance of the sub-networks. In fact, in this scenario, if we extract and deploy the sub-models of the trained large model on a target task, we would experience a significant drop in the performance of these sub-networks compared with the main model. However in SortedNet, we propose a training method that allows for training sub-networks together with the main model in a stochastic way. The SortedNet paradigm leads to the following benefits:

* Search-free sub-model extraction: after training, by importance sorting of sub-models the best sub-model for a given budget can be selected without the need for search.
* Anytime: Each smaller sub-model is a subset of a larger one which makes switching between different sub-models efficient. This leads to an important feature of our SortedNet which is so-called _anytime_ that is a network which can produce its output at any stage of its computation.
* Memory efficient: we train a many-in-one network where sub-models are all part of a single checkpoint, which minimizes storage requirement.

For efficiency purposes, in our training, the last layer, e.g. the classification layer, is shared between all sub-models; alternatively we can add a separate classification layer to each sub-model. For simplicity and efficiency, we chose to use a shared classification layer.

### SortedNet Algorithm

In this subsection, we describe our proposed training algorithm. For training a SortedNet with \(n\) sub-models, at each iteration during training, a random index needs to be sampled
from a pre-defined distribution: \(b^{i}_{j}\sim P_{B_{j}}\). After finding the target sub-model \(\theta^{*}_{t}\) at each iteration, we can use one of the following objectives to update the parameters of the selected sub-model:

* (Stochastic Loss) Only train the selected sub-model \(f(x,\theta^{*}_{t})\) : \(\min\limits_{\theta^{*}_{t}}\mathcal{L}\triangleq\mathrm{L}(y,f(x,\theta^{*}_{ t}))\) where \(L\) is the loss function for training the model on a given task (e.g. \(L\) can be a regular cross entropy loss) and \(y\) refers to the ground-truth labels.
* (Stochastic Summation) Train the sub-model \(f(x,\theta^{*}_{t})\) and all its targeted sub-models along each dimension. Let's assume that \(\Theta^{\perp}(\theta^{*}_{t})\) is the universal set for all targeted sub-networks of \(\theta^{*}_{t}\). Then the loss function can be defined as: \(\min\limits_{\Theta^{\perp}(\theta^{*}_{t})}\mathcal{L}\triangleq\sum_{\theta \in\Theta^{\perp}(\theta^{*}_{t})}\mathrm{L}(y,f(x,\theta))\)

This way, one sub-model or a subset of sub-models are updated in each iteration. Alternatively, one can choose to train all the sub-models at each iteration which is costly in the large scale.

## 4 Experiments

In this section, we discuss a set of experiments that we conducted to show the effectiveness and importance of sorting information and fixing a nested property. The details of the hyperparameters for each experiment can be found in Appendix A.3.
### Is SortedNet scalable?

To show that our proposed method is scalable, we designed an experiment that try to train 160 different models across multiple dimensions (width and depth) all at once. As baseline, we trained the largest network (a MobileNetV2), and reported the best performance of the model. Because the performance of the model was poor for all the other sub-models (less than 12%), we trained the classifier layer for 5 more epochs before evaluating each sub-model for the baseline and reported the best performance. As the results suggests in figure 3, our method was able to capture the maximum performance for many of these sub-models in a zero-shot manner. In each cell, we reported the performance of the sub-network on top and the recovery percentage of the model with respect to the largest model (in this example, 95.45). Despite sharing the weights across all models, sharing the classifier and zero-shot evaluation, the proposed method preserved up to 96% of the performance of the largest model which is highly encouraging. Further training of the classifier for our proposed method will lead to even better performance as shown in appendix A.4 (between \(\sim\)2 to 15% improvement for different sub-models). In addition, we also tried to sort the depth and width using proposed method individually which has been reported in the figure 3 as D. Only, and W. Only, respectively. Across width, SortedNet successfully preserved up to 99% of the largest network's performance.

### Can we find the best sub-models using SortedNet?

As shown in figure 4, based on the performance of the models in the previous experiment which has been shown in figure 3, we selected a subset of best-performing networks (width \(>\) 60% and depth \(>\) 13 blocks), and retrained the network from scratch using SortedNet to show the success rate of our proposed method. As shown, SortedNet successfully preserved up to 99% of the performance of the ordinary training of the largest network.

We can also make this selection process fully automated by sorting the performance of all sub-models after evaluation and filtering out a subset of best-performing models that perform better than a desired threshold. As can be seen in figure 6, there is a set of sub-networks which perform better than 80%. To better understand the pattern, we annotated some of the points using "\({}^{D}_{W}\)" as template which shows for each model the corresponding width and depth.

### Can we generalize SortedNet?

In another experiment, as shown in table 2, we wanted to show our stochastic approach can outperform the state-of-the-art methods such as LCS (shown as \(LCS_{p}\) in table) (Nunez et al., 2023), Slimmable Neural Network (NS) (Yu et al.

\begin{table}
\begin{tabular}{l c c|c c c|c c} \hline Network & Width & FLOPs & NS-IN & LCS-p-IN & SortedNet-IN & NS-BN & LCS-p-BN (aka US) & SortedNet-BN \\ \hline  & 100\% & 301M & 88.67 & 87.61 & **89.14** & 79.84 & 65.87 & **85.24** \\  & 75\% & 209M & 87.86 & 85.73 & **88.46** & 78.59 & **85.67** & 85.29 \\ \multirow{-1}{*}{cpreresnet20 (He et al., 2015) (CIFAR10)} & 50\% & 97M & 84.46 & 81.54 & **85.51** & 69.44 & 65.58 & **70.98** \\  & 25\% & 59M & 75.42 & **76.17** & 75.10 & 10.96 & **15.78** & 12.59 \\ \hline avg. & - & - & 84.10 & 82.76 & **84.55** & 59.70 & 58.22 & **63.52** \\ \hline \end{tabular}
\end{table}
Table 2: Comparing the performance of state-of- the-art methods with Sorted-Net over CIFAR10 in terms of test accuracies.

Figure 5: Comparing the training loss trajectory of SortedNet on CIFAR10 for different gradient accumulation values with LCS_p. Each subfigure demonstrates the results in different widths. The rightmost subfigure reports the average across the widths. The underlying network (cPreResNet20) and hyperparameters are fixed.

Figure 6: Finding best sub-models automatically using a desired threshold bar to eliminate the worst performing models.
2018), and Universally Slimmable Networks (US) (Yu and Huang 2019). To make the comparisons fair, we equalized the number of gradient updates for all models. We also tried to remove the impact of architecture design such as the choice of the normalization layers. Therefore, we tried to compare methods by different layer normalization techniques such as BatchNorm (Ioffe and Szegedy 2015) and InstanceNorm (Ulyanov, Vedaldi, and Lempitsky 2016). In addition, we made sure that the orthogonal methods such as Knowledge Distillation will not impact the results as these methods can be applied and improve the results independent of the method. As shown in the table, SortedNet demonstrates a superior average performance compared to other methods, indicating its generalization across various settings such as different norms.

### Can we extend SortedNet to language models?

In this experiment, the goal is to apply SortedNet for a pre-trained transformer model and evaluate the performance on the GLUE benchmark (Wang et al. 2018). As the baseline, we chose RoBERTa (Liu et al. 2019) to demonstrate the flexibility of our algorithm. In table 3, we sorted all the layers of RoBERTa-base. As the results demonstrate, our proposed method in average perform better than the baseline by a significant margin (\(\sim\) 23%). However, the largest model has a small drop in performance (less than 2%). It is interesting that the transformer architecture can preserve the performance of sub-models up to some extent without additional training. However, our algorithm improve the performance of these sub-models between 10 to 40% approximately.

## 5 Can we extend SortedNet to complex dimensions?

In this section, we are interested to investigate whether SortedNet is applicable to more complex dimensions other than width and depth. For example, can we utilize the SortedNet for sorting the Attention Heads (Vaswani et al. 2017)? To achieve this, we conducted an experiment over BERT-large (Devlin et al. 2019) which we tried to sort the information across multiple dimensions at once including, number of layers, hidden dimension, and number of attention heads. In other words, we tried to sort information over Bert-large and Bert-base as Bert-base can be a subset of the Bert-large and therefore respect the nested property. As shown in table 4, in addition to the reported performance of Bert-base and Bert-large according to the original paper (Devlin et al. 2019), we reported the performance of these models in the paper experimental setting. The performance of randomly initialized Bert-base has been reported as well. We also extracted a Bert-base from a Bert-large model, and we reported the performance of such model in the same table. Additionally, we highlighted the number of training updates with respect to each objective function in front of each model. For example, in the last row (Sorted \(BERT_{LARGE}\)), we approximately trained our Sorted model half of the times (\(\sim 3Epochs\)) over the objective function of Bert-base (\(\mathcal{L}_{B}\)) and the other half of the times over the objective function of Bert-large (\(\mathcal{L}_{L}\)) in an iterative random manner as introduced in the section 3. The learned Bert-base performance with these methods is still around 10% behind a pre-trained base but we argue that this is the value of pre-training. To investigate the impact, one should apply the SortedNet during pre-training which we will leave for future research. However, the performance of the learned Bert-large is on-par with an individual Bert-large which suggests sharing the weights does not necessarily have a negative impact over learning. It seems, however, the secret sauce to achieve a similar performance is that we should keep the number of updates for each objective the same as the individual training of Bert-large and Bert-base.

### Ablation Study

Convergence (Training Time) AnalysisBeing sorted and randomly selecting one sub-network at the time from a pre-defined set of the sub-networks empowers SortedNet with a higher convergence rate and a faster training time. Figure 5 empirically certifies this claim and compares the training convergence of SortedNet against LCP_p, which, to the best of our knowledge, LCP_p stands as the most recent state-of-the-art method. As LCS_p uses summation loss over four

\begin{table}
\begin{tabular}{l c c c c c c c c|c} \hline \hline  & Acc. & Acc. & F1 & Mathews Corr. & Acc. & Acc. & Acc. & Pearson & & \\
**Model** & **MNLI** & **SST-2** & **MRPC** & **CoLA** & **QNLI** & **QQP** & **RTE** & **STS-B** & **AVG** & **AVG w/o ours** \\ \hline Sorted-RoBERTa (11) & 60.07 & 70.76 & 81.22 & 0.00 & 59.64 & 77.80 & 47.65 & 9.36 & **50.81** & 40.33 \\ Sorted-RoBERTa (21) & 71.98 & 80.28 & 81.22 & 0.00 & 81.49 & 87.09 & 47.29 & 70.37 & **64.97** & 40.86 \\ Sorted-RoBERTa (41) & 76.74 & 80.50 & 81.22 & 0.00 & 85.21 & 88.82 & 46.93 & 75.07 & **66.81** & 41.06 \\ Sorted-RoBERTa (41) & 79.13 & 84.75 & 81.22 & 44.51 & 86.60 & 90.11 & 49.10 & 84.94 & **75.04** & 42.95 \\ Sorted-RoBERTa (5L) & 81.14 & 89.91 & 81.22 & 48.41 & 87.88 & 90.86 & 55.96 & 88.22 & **77.95** & 43.80 \\ Sorted-RoBERTa (6L) & 82.21 & 92.09 & 86.67 & 53.41 & 88.83 & 91.12 & 67.87 & 89.09 & **81.41** & 46.13 \\ Sorted-RoBERTa (7L) & 82.99 & 92.78 & 89.13 & 56.42 & 89.29 & **91.29** & 73.29 & 89.58 & **83.10** & 44.80 \\ Sorted-RoBERTa (8L) & 83.33 & 93.23 & 89.78 & 57.22 & 89.40 & 91.29 & 75.09 & 89.67 & **83.63** & 55.17 \\ Sorted-RoBERTa (9L) & 83.39 & 92.66 & 89.66 & 58.69 & 89.40 & 91.25 & **77.26** & 89.72 & **84.00** & 61.36 \\ Sorted-RoBERTa (10L) & **87.42** & 93.12 & **91.64** & **61.21** & **91.87** & 91.19 & 74.01 & **89.74** & **85.02** & 54.30 \\ Sorted-RoBERTa (11L) & 87.34 & **93.35** & 91.45 & 60.72 & 91.74 & 91.17 & 74.01 & 89.72 & **84.94** & 77.48 \\ Sorted-RoBERTa (12L) & 83.35 & 92.89 & 90.81 & 59.20 & 89.44 & 91.28 & 76.53 & 89.77 & 84.16 & **86.13** \\ \hline avg. & 79.26 & 87.93 & 86.09 & 41.25 & 85.50 & 89.45 & 64.26 & 79.61 & **76.67** & 52.86 \\ \hline \hline \end{tabular}
\end{table}
Table 3: A comparison of the performance of different sub-models with and without the SortedNet. The model’s performance will improve if we have more budgets and calculate the representation of deeper layers.
sub-networks in every training steps and to have a fair comparison, we therefore report the performance of SortedNet in different values of gradient accumulation (\(g_{acc}\)), where \(g_{acc}=4\) provides a fair comparison with LCS_p. As shown in the figure, SortedNet with \(g_{acc}=4\) converges either faster or competitive across different sub-networks. Moreover, SortedNet does not require any for-loop in its implementation; thus tailoring parallel computation and resulting in faster running time. We empirically investigate this feature and found that in the same settings, SortedNet runs at least one third faster than LCS_p (details in Appendix A.2).

The impact of gradient accumulationThe goal of this experiment is to examine the impact of gradient accumulation (\(g_{acc}\)) on the performance of SortedNet within an equal number of parameters updates. Table 5 presents the results obtained in terms of accuracies for 4 different gradient accumulation values. To ensure an equal number of updates, the maximum number of epochs is adjusted for each scenario, e.g. \(g_{acc}=k\) receives \(k\) times more epochs than \(g_{acc}=1\). As the results explains, increasing gradient accumulation values results in a higher performance for SortedNet. This observation can be attributed to the increase in training stochasticity when gradient accumulation is raised. Consequently, each sub-network in SortedNet contributes more equally to the updating of weight parameters, leading to a faster convergence rate. More details are provided in Appendix A.1.

In addition, we highlighted the details of each experiment hyperparameters in appendix A.3 and further analysis has been provided in appendix A.5 to better understand the behavior of sortedNet methodology.

## Conclusion

In summary, this paper proposes a new approach for training dynamic neural networks that leverages the modularity of deep neural networks to efficiently switch between sub-networks during inference. Our method sorts sub-networks based on their computation/accuracy and trains them using an efficient updating scheme that randomly samples sub-networks while accumulating gradients. The stochastic nature of our proposed method is helping our algorithm to generalize better and avoid greedy choices to robustly optimize many networks at once. We demonstrate through extensive experiments that our method outperforms previous dynamic training methods and yields more accurate sub-networks across various architectures and tasks. The sorted architecture of the dynamic model proposed in this work aligns with sample efficient inference by allowing easier samples to exit the inference process at intermediate layers. Exploring this direction could be interesting for future.

## Limitations

It is good to note that our proposed method might be sensitive to the randomness as the chosen trajectory at the moment is random uniform. Further research is necessary to investigate the effect of choosing more optimal strategies for choosing the next model at each iteration.

\begin{table}
\begin{tabular}{l c c c c c c c c c c c} \hline \hline \multirow{2}{*}{Model} & \multirow{2}{*}{Flops} & \multirow{2}{*}{Weights} & \multirow{2}{*}{MNLI} & \multirow{2}{*}{QQP} & \multirow{2}{*}{QNLI} & \multirow{2}{*}{SST-2} & \multirow{2}{*}{CoLA} & \multirow{2}{*}{STS-B} & \multirow{2}{*}{MRPC} & \multirow{2}{*}{RTE} & \multirow{2}{*}{avg.} \\ \cline{5-10} \cline{6-11