Investigating Acceleration of LL a MA Inference by Enabling Intermediate Layer Decoding via Instruction Tuning with ‘ LITE ’ (2024)

Abstract: Large Language Models (LLMs) have achieved remarkable performance across a wide variety of tasks; however, their large size makes their inference slow and computationally expensive. Focusing on this problem, we study instruction tuning LLMs with additional explicit Losses from the Intermediate layers (LITE) and show that it enables these layers to acquire ‘good’ generation ability without affecting the generation ability of the final layer. We then perform ‘dynamic confidence-based early exiting’ at token level from the intermediate layers which improves the computational efficiency of text generation without sacrificing the quality of the generation. We conduct comprehensive experiments by instruction tuning LLaMA-2 models on the Alpaca dataset and evaluate on four different instruction test sets. We show that dynamic early exiting achieves consistent and considerable inference cost improvements (37.86% for 7B and 46.35% for 13B model) while maintaining the generation quality. We further conduct a thorough analysis of the results and dissect the efficiency improvements which reveals several important findings.

## introduction recently developed llms @xcite @xcite @xcite have revolutionized the field of natural language processing and achieved remarkable performance across a wide variety of tasks. 'instruction tuning' further teaches these models to follow the user's instruction provided in natural language @xcite @xcite @xcite . despite all the notable abilities of these models, their large size (number of parameters) makes their inference slow and computationally expensive which poses a practical challenge limiting their widespread adoption in resource constrained applications. focusing on the above problem, in this work, we investigate instruction tuning llms in a way that enables intermediate layer decoding for efficiently generating text without compromising the quality of the generation. we first show that in standard instruction tuning, only the final layer of the model acquires the ability to generate 'quality' text while the representations of the intermediate layers (when passed through the language modeling head) fail to do so. this restricts decoding from these intermediate layers without degrading the generation quality. addressing this point, we instruction tune llms with additional explicit losses from the intermediate layers (lite) and show that it enables these layers to acquire 'good' generation ability. importantly, we show that these layers acquire this ability without affecting the generation ability of the final layer; however, as expected, their generation ability still remains slightly inferior to the generation ability of the final layer. thus, decoding the complete response from intermediate layers improves the efficiency of inference but still results in degradation in the quality of the response. addressing the above limitation, we show that (a) lite greatly aligns the intermediate layers' token prediction with that of the final layer and (b) the intermediate layers' token prediction probabilities provide a strong signal of this alignment. building on these findings, we perform 'dynamic confidencebased early exiting' at token level from the intermediate layers which improves the efficiency of inference while maintaining the generation quality. we conduct comprehensive experiments by instruction tuning llama-2 models @xcite on the widely used alpaca dataset @xcite and holistically evaluate on four different human-instruction test sets including vicuna, wizardlm, koala, and self-instruct. figure 1 compares the quality of responses (evaluated using the claude model as detailed in section 5) and the inference cost (measured in flops) of the (i) standard generation method from the final layer with (ii) the dynamic early exiting method. it shows that dynamic early exiting achieves consistent and considerable inference cost improvements (37.86% for 7b and 46.35% for 13b model on average) while maintaining the generation quality. we further perform a thorough analysis of the results over several important aspects, such as, comparing the semantic similarity between the responses generated from the final layer and the early exiting method, and dissecting the efficiency improvements by comparing the number of tokens generated in the outputs. we also discuss the potential of intermediate layer decoding in 'speculative sampling' and 'hallucination detection'. in summary, our work contributes to improving the efficiency of llm inference while maintaining the generation quality, a crucial step en route to enabling their widespread adoption. ## related work improving the inference efficiency of llms is an important research direction and is receiving considerable attention from the nlp community. in this section, we review some of the existing methods and differentiate our work from them. reducing model size: since model size plays a crucial role in increasing the inference cost and latency, techniques like quantization @xcite @xcite , knowledge distillation @xcite @xcite , model compression and network pruning @xcite have been shown to be effective in improving the inference efficiency. furthermore, during sampling, a cache of the keys and values can be maintained for every attention layer which reduces the computations at inference time (kv caching). however, it increases the gpu vram memory requirement of inference. another technique speculative sampling @xcite first generates a draft of k tokens from a smaller auto-regressive model and then scores the draft using the target model. this results in generation of more than one token (on average) from the target model in a single pass. early exiting and cascading based inference techniques have been shown to be effective for classification tasks with bert-style models, such as deebert @xcite ) that speeds up bert inference by inserting extra classification layers between each encoder layer, power-bert @xcite ) that focuses on progressive wordvector elimination (based on significance computed using self-attention) along the encoder pipeline, dynabert @xcite ) that adjusts the size of the model by selecting adaptive width and depth, and cascading @xcite @xcite @xcite in which sequential inference is done through models of bigger and bigger size with conditional exiting to output predictions efficiently. our work is also related to confident adaptive language modeling (calm) @xcite and depth-adaptive transformers @xcite in which early exiting is performed by learning additional classifiers attached to the decoder layers. @xcite proposed to short-cut away transformer inference in between certain layers by learning linear transformations across layers in the network, i.e., casting internal representations. o' @xcite ; @xcite explore leveraging the intermediate layers for contrastive decoding to improve reasoning. we discuss other related and concurrent work in appendix h. our work differs in the following aspects: (1) firstly, most of the existing work in early exiting focuses on improving the efficiency of encoderonly models (like bert) or encoder-decoder models (like t5); our work focuses on the current stateof-the-art decoder-only llms (llama-2). furthermore, we focus on the instruction tuning setting with text generation, unlike prior work that focused on solving simpler tasks like glue classification or qa. (2) early exiting methods typically require training additional classifiers for the intermediate layers, however, in this work, we use the same shared language modeling head at all the layers; thus, we do not introduce new model weights. (3) for leveraging the intermediate layers for decoding, we enable them to acquire generation ability by instruction tuning with lite, unlike other methods that use a pre finetuned model in which these layers have poor generation ability as we show in section 6.1. (4) existing methods typically require complex architectural modifications, pruning, saliency quantification, or training new parameters. in contrast, our method (both for tuning and inference) is simple and easy to implement and yet achieves considerable benefits. (5) existing methods typically require training a separate model for each computation budget; however, in our method, the same model can be adapted to meet all the computation constraints (by varying the exiting confidence thresholds). (6) the computational efficiency often comes with a compromise in performance. however, our method maintains the generation quality while providing efficiency benefits. ## instruction tuning with lite instruction tuning (it): one of the major reasons that necessitate instruction tuning of llms is the mismatch between their pre-training objective and the users' objective, i.e., llms are typically trained on minimizing the word prediction error while users want the model to follow their instructions. to this end, an instruction tuning dataset is collected and a pre-trained model is fine-tuned in a supervised manner @xcite @xcite . loss calculation during instruction tuning of a typical decoder-only llm (llama in this case) is shown in figure 2 (left). the model consists of a stack of decoder layers followed by a language modeling head which outputs the probability distribution over the vocabulary tokens as its prediction. during the supervised fine-tuning, the loss over the output tokens is backpropagated from the final layer of the model: it with lite: we show that in standard instruction tuning, only the final layer of the model acquires the ability to generate 'quality' text while the representations of the intermediate layers (when passed through the language modeling head) fail to do so (section 6.1). in other words, it does not explicitly teach the intermediate layers of the tuned llm to generate tokens. this restricts decoding from these intermediate layers without degrading the generation quality. we note that during tuning, the same language modeling head (that is used with the final layer) can also be used with the intermediate layers to obtain the losses of those layers. thus, this does not impact the number of parameters of the model. to this end, we calculate a weighted aggregation of the losses from the intermediate layers (including the final) to calculate the overall loss value: during training, we use the representations of the intermediate layers and calculate the loss from these layers at the end. we note that this is a general formulation as it captures a variety of scenarios including the standard fine-tuning in which the loss is calculated only from the last layer (i.e., w 1:n -1 = 0 and w @xmath0. furthermore, this formulation also allows aggregating losses from only the selected intermediate layers instead of all the layers by accordingly defining the lm head pathways and the w i values. in section 6.2, we will show that this formulation while enabling the intermediate layers with 'good' generation ability does not adversely affect the final layer's generation ability. furthermore, as expected, the quality of generation typically improves with the layer number as the later layers have more capacity to learn. ## making inference efficient in this section, we first detail auto-regressive inference and then describe early exiting techniques. auto-regressive inference: it refers to the process of generating a sequence of tokens where each token is generated based on the preceding tokens in the sequence. for generating a token, the model takes the input (including the previously generated tokens) and runs a forward pass in which the input is fed to the model and passed sequentially along its layers until the probabilities for the next token are predicted (called as logits). chaining model forward passes with next token selection iteratively leads to the generation of text. in greedy decoding, the token with the highest probability is selected as the next word prediction at each timestep. ## experimental setup instruction tuning: we instruction tune the llama-2 models @xcite (7b and 13b ) with the widely used alpaca dataset @xcite . alpaca consists of 52k instructionfollowing demonstrations generated using the selfinstruct @xcite technique. in it with lite for 7b model (32 total layers), we aggregate losses from the following selected intermediate lay-ers: @xcite @xcite @xcite along with the final layer and use equal weights in loss calculation. similarly, for the 13b model (40 total layers), we use @xcite @xcite @xcite @xcite layers. we perform full parameter fine-tuning on 4 a100 gpus. we skip selecting the initial layers because they have a limited capacity to learn and thus can not give good token predictions. furthermore, we select layers at an interval of 4 so that at inference time, the model can do enough reasoning/interactions between two consecutive checkpoints. otherwise, checking at every layer can result in computational overhead. we train this model for 5 epochs so that it achieves training loss comparable to standard tuning. we present all the results corresponding to this tuning configuration in the main paper and present the study corresponding to weighted lite in the appendix c. ## results and analysis in this section, we first demonstrate the inability of the intermediate layers of the model tuned with standard it to generate 'quality' text (6.1). then, we show the impact of it with lite: it does not adversely affect the generation quality of the final layer (6.2), it aligns the intermediate layers' token predictions with the final layer (6.3), and the corresponding prediction confidence values provide a strong signal of the alignment (6.4). these findings motivate dynamic confidence-based early exiting. finally, we show the effectiveness of the method in improving the efficiency of inference while maintaining the generation quality (6.5). to avoid repetition, we present results for the 7b model in the main paper and for the 13b model in b.8. ## conclusion and discussion in this work, we proposed instruction tuning with additional explicit losses from the intermediate layers and showed that it enables these layers to acquire 'good' generation ability without affecting the final layer's generation ability. we performed 'dynamic confidence-based early exiting' at token level from the intermediate layers and showed that 8 12 16 20 24 28 32 exit layer 0 5 10 15 20 25 30 percentage figure 8: percentage of token outputs from different exiting layers in the proposed method. looking forward, our work additionally opens up several other avenues for new research, such as speculative sampling from the intermediate layers to improve the inference efficiency and checking information consistency from the output of intermediate layers to detect hallucinations (discussed in appendix i). furthermore, this approach is complementary to some existing efficiency methods described in section 2 and h, i.e., they can be used in conjunction to achieve even more efficiency gains.