SortedNet: A Scalable and Generalized Framework for Training Modular Deep Neural Networks

TL;DR

SortedNet introduces a generalized, architecture-agnostic framework for training many sub-models in parallel by leveraging a sorted, shared-parameter architecture and a gradient-accumulation updating scheme. It enables seamless dynamic inference with a single checkpoint, supports both depth and width (and beyond) across CNNs and Transformers, and scales to hundreds of sub-models while retaining high performance (e.g., up to $96\%$ of the largest model with 160 sub-models). The authors provide convergence and performance-bounded theory and demonstrate practical benefits, including automatic sub-model selection, reduced storage, and significant inference speedups for large language models via speculative decoding. The results span vision and NLP tasks (CIFAR-10, BERT/RoBERTa, LLaMA) and extend to decoder-based LLMs, highlighting SortedNet’s potential to streamline deployment under varying budgets and hardware constraints. Overall, the work offers a new, scalable paradigm for training many-in-one networks that preserves performance while enabling flexible, budget-aware inference in real-world systems.

Abstract

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.

Figures (7)

• Figure 1: (a) The overall diagram of our SortedNet training approach. First, we need to define the pool of sub-models of interest including the main model as well. During training, at each iteration, we sample from the pool of sub-models (given a pre-defined random distribution) to be trained for the target loss function (for one step). (b) The generalizability of Sorted configuration across more complex dimensions, supporting blocks and functions iin addition to width and depth. (c) Illustrating the difference between the nested and sorted sub-models. In nested architectures, smaller sub-models are encapsulated by larger sub-models, which is not necessarily the case for what we refer to as sorted models. Moreover, sorted models are tied to the origin (i.e. starting index) of each dimension which might not be the case in nested models.
• Figure 2: (a) CIFAR10 classification accuracy (and recovery percentage) for Sorted-Net (160 Models) and the baseline. In each cell, we reported the performance of the sub-model (top) and the relative performance of the model (in percentage) with respect to the baseline largest model performance (bottom). W. Only: Sorting only the widths, D. Only: Sorting only the depth. More black the better. (b) CIFAR10 classification performance for the best-performing subset of sub-models trained by SortedNet from scratch. More black the better. (c) Finding best sub-models automatically using a desired threshold bar to eliminate the worst performing models.
• Figure 3: (left) Confidence-based Early-Exit, exiting from intermediate sub-models whenever the confidence passes the determined threshold (right) Sorted Self-Speculative Decoding, verifying the adaptively exited draft tokens from intermediate sub-models by the full model.
• Figure 4: Comparing SortedNet and Once For All: on a hypothetical 5-layer network, we show how the sub-model selection strategy of SortedNet differs from the Once-for-All cai_once-for-all_2020 approach.
• 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.