Sorted LL a MA : Unlocking the Potential of Intermediate Layers of Large Language Models for Dynamic Inference (2024)

Abstract: Large language models (LLMs) have revolutionized natural language processing (NLP) by excelling at understanding and generating human-like text. However, their widespread deployment can be prohibitively expensive. SortedNet is a recent training technique for enabling dynamic inference by leveraging the modularity in networks and sorting sub-models based on computation/accuracy in a nested manner. We extend SortedNet to generative NLP tasks, making large language models dynamic without any Pre-Training and by only replacing Standard Fine-Tuning (SFT) with Sorted Fine-Tuning (SoFT). Our approach boosts model efficiency, eliminating the need for multiple models for various scenarios during inference. We show that this approach can unlock the potential of intermediate layers of transformers in generating the target output. Our sub-models remain integral components of the original model, minimizing storage requirements and transition costs between different computational/latency budgets. The efficacy of our proposed method was demonstrated by applying it to tune LLaMA 2 13B on the Stanford Alpaca dataset for instruction following and TriviaQA for closed-book question answering. Our results show the superior performance of sub-models in comparison to Standard Fine-Tuning and SFT+ICT (Early-Exit), all achieved with very efficient tuning and without additional memory usage during inference.

## introduction large language models are revolutionizing the way we interact with information in today's world @xcite @xcite . new models are continually emerging, demonstrating their capabilities in understanding and, more importantly, in generating human-like text. notably, models such as chatgpt, llama 2 70b @xcite , and falcon 180b @xcite have had a profound impact on the applicability of large language models (llms). however, deploying these expansive language models can become prohibitively expensive. what distinguishes this new era of chatgpt-like models is their ability to perform an extraordinarily wide array of tasks in natural language processing (nlp), reasoning, and more, all through behavior cloning @xcite . in fact, a single model can leverage the strong contextual learning ability offered by standard fine-tuning to address numerous tasks, spanning from language comprehension to complex reasoning. while this unified usage simplifies the deployment of these models as general assistants, it remains highly inefficient. enabling dynamic inference, where the computational resources allocated to a given query vary at inference time, can significantly enhance the practicality of employing such models in real-time scenarios. this enables the use of smaller models when the budget is limited or latency is critical. it is important to note that dynamic inference strategies for large models with a substantial number of parameters should not require loading different models during inference. previous research has explored methods for training dynamic models capable of adapting to evolving resource constraints @xcite @xcite . however, existing approaches often rely on complex training procedures or necessitate modifications to the original model architecture. sortednet @xcite introduces a novel approach to training deep neural networks that leverages the inherent modularity of these networks to construct sub-models with varying computational loads. this method sorts sub-models hierarchically based on their computation/accuracy characteristics, facilitating efficient deployment during inference. furthermore, it employs an efficient updating scheme combining random sub-model sampling with gradient accumu-lation to minimize the training cost. consequently, with a single round of training, numerous models can be obtained within a single model. while the sortednet approach has primarily been applied to vision and language understanding tasks, given the significant impact of generative language models in today's ai landscape, the efficacy of this method for generative tasks in nlp is of considerable interest. in fact, being able to make a large language model dynamic without the need for pre-training and only at the cost of a round of standard fine-tuning can open doors to efficient inference of these models without incurring additional expenses associated with common model compression methods like knowledge distillation and pruning, among others. moreover, since all the resultant models are components of the original model, the storage requirements and the cost associated with transitioning between different computation demands become minimal. otherwise, managing multiple models for various scenarios during inference becomes impractical. in this study, we challenge the conventional approach of relying solely on the last layer's contextual embeddings and use sorted fine-tuning (soft) in place of standard fine-tuning to enhance the performance of these models across multiple layers. by doing so, we aim to provide new insights into the efficiency and effectiveness of middle layers in producing high-quality results for specific downstream tasks. our proposed approach can potentially optimize these sub-models in addition to the main model, ultimately enhancing their overall performance. in this paper, we seek to answer the following questions through systematic evaluation: to answer these questions, we employ llama 2 13b and perform both standard fine-tuning (sft) and sorted fine-tuning (soft) on the stanford alpaca @xcite and triviaqa @xcite datasets. for sorted fine-tuning, we target 8 sub-models and share the llm head among them to ensure cost parity. we utilize the pandalm benchmark @xcite to assess the perfor-mance of the sub-models on alpaca dataset. our findings demonstrate the superior performance of soft in comparison to sft and even to memorydemanding methods like early exit @xcite . the contributions of this paper can be summarized as follows: ## related work this section briefly introduces the most relevant papers to our work. many-in-one models deep neural networks (dnns) are often overparameterized, motivating researchers to explore ways to use the parameters of the models more efficiently. more number of parameters lead to higher costs of deployment for neural networks. moreover, in practice, these overparametrized dnns are expected to accommodate customers with varying requirements and computational resources. to address these diverse demands, one can think of training models of different sizes, which can be prohibitively costly (in terms of training and memory), or another alternative is to train many-in-one networks @xcite . manyin-one solutions aim to train a network along with some of its sub-networks simultaneously for specific tasks. for example, we can consider the early-exit method @xcite , wherein a prediction head is fine-tuned on top of specific intermediate layers within a network. another approach is layer drop @xcite , which trains a network in any depth by randomly dropping the layers during training. while both early-exit and layer drop are simple solutions, they are not state-of-the-art in terms of performance. in early-exit, we only train the output prediction layer on top of each intermediate layer, and this layer might not have enough capacity to retain a good performance. layer drop, conversely, suffers from the abundant number of possible sub-models in training, which makes the training process exhaustive and sub-optimal. furthermore, this approach requires tuning the extent of dropping layers during training. this additional hyper-parameter, layer drop rate during training determines the best size and setting of the model at the inference time. deviating from the training drop rate at the inference time can result in a significant drop in performance. @xcite in once for all (ofa) proposed an alternative solution to neural architecture search (nas). ofa requires training the model and all possible sub-models in an arbitrary progressive way followed by a separate search phase. dyna-bert @xcite is another work that targets training dynamic pre-trained manyin-one bert models in two stages: first, distilling from the main network to the width adaptive networks and then distilling from the width adaptive networks to depth adaptive networks. both width adaptive and depth adaptive networks have a limited pre-defined set of width and depth for the sub-models. while both ofa and dynabert have shown successful results, their solutions are hardly applicable to multi-billion-parameter llms because of their complicated multi-stage training process and their search and knowledge distillation requirements. sortednet @xcite is a recent method that forms and trains sub-models of a network in a sorted manner while not requiring any search during training or inference. sortednet has shown superior performance compared to other previously mentioned methods in terms of simplicity, performance, scalability, and generalization. considering these benefits, we target deploying the sortednet training algorithm for developing manyin-one llms. ## methodology this paper focuses on making generative llms many-in-one by unlocking the potential of intermediate layers through the sortednet approach @xcite . let's consider a language model f (x; θ) with the parameters θ and the input x. the following is the sorted training procedure: forming sub-networks first, we need to form the sub-networks of the llm. for the sake of simplicity and without loss of generality, we focus on the depth-wise sub-networks. supposed that the sub-network f n (x; θ n ) refers to the first n layers of f (x; θ). in this paper, the language model is considered to be llama2 13b. since llama2 comprises 40 layers, we define the sub-networks as @xmath0, 16, 20, 24, 28, 32, 36, 40}. calculating the output of sub-networks each sub-model's output will be predicted using the shared output prediction head from the last layer (original network). remember that in the llama query: summarize a meeting from the given list of bullet points. be sure to convert shorthand into a first-hand account. rose: analyze data and presents findings john: propose new idea jane: appointed to head project tom: need more time to fix software bug ## experiments this section delves into the experiments' specifics and the analysis provided to understand better the effect of sorted fine-tuning over the performance of a large language model like llama2 @xcite . before diving into results, we are going to define certain notations that we used for different setups in our experiments: we first train the model with sft paradigm and then further fine-tune the language model head exclusively for each sub-model while keeping their weights frozen. the sft+ict is also known as early-exit @xcite in the literature. ## conclusion this work presents sorted llama, a many-in-one language model for dynamic inference obtained using sorted fine-tuning (soft) instead of standard fine-tuning. sorted llama unlocks the potential capability of intermediate layers, offering dynamic adaptation without pre-training or additional expenses related to model compression. it presents a promising avenue for optimizing generative language models in nlp. our approach makes the deployment of these models far more efficient. as all sub-models remain integral components of the original model, the burden of storage requirements and transition costs between different computational demands is minimized, making the management of multiple models during inference a practical reality. our systematic evaluation of instruction following and questions answering benchmarks challenged conventional wisdom by empowering middle layers to produce high-quality results. this, in turn, enables dynamic inference of llms with a highly efficient tuning method (soft), ultimately optimizing the usage of llms. our encouraging results show the promising capability of sortednet @xcite to train multiple language models with different sizes at once without incurring expensive costs.