Fast and Robust Early-Exiting Framework for Autoregressive Language Models with Synchronized Parallel Decoding (2023)

Abstract: To tackle the high inference latency exhibited by autoregressive language models, previous studies have proposed an early-exiting framework that allocates adaptive computation paths for each token based on the complexity of generating the subsequent token. However, we observed several shortcomings, including performance degradation caused by a state copying mechanism or numerous exit paths, and sensitivity to exit confidence thresholds. Consequently, we propose a Fast and Robust Early-Exiting (FREE) framework, which incorporates a shallow-deep module and a synchronized parallel decoding. Our framework enables faster inference by synchronizing the decoding process of the current token with previously stacked early-exited tokens. Furthermore, as parallel decoding allows us to observe predictions from both shallow and deep models, we present a novel adaptive threshold estimator that exploits a Beta mixture model to determine suitable confidence thresholds. We empirically demonstrated the superiority of our proposed framework on extensive generation tasks.

## introduction recent advancements in autoregressive language models @xcite @xcite have revolutionized the quality of language generation in various generative tasks, including question answering @xcite , summarization @xcite , and machine translation @xcite . nevertheless, these large transformer models have shown high inference latency due to the considerable number of layers and the autoregressive decoding step. as the multiple stacks of transformer layers have to be computed sequentially for each individual token, the inference process poses sig- * equal contribution † corresponding authors nificant computational burdens and hinders their real-time adaptability @xcite . in light of the necessity to expedite inference latency, the early-exiting framework @xcite @xcite emerges as a promising approach that dynamically allocates computation paths based on the complexity of generation for each token. as illustrated in figure 1a , tokens that are relatively easy to predict the subsequent token yield consistent predictions with only a few layer computations, while those with higher difficulty require computations across a larger number of layers to generate accurate predictions. in an ideal scenario, the early-exiting method empowers models to achieve notable acceleration in inference without compromising the generation quality when compared to that of a full model. however, our extensive analysis identified four challenges in the early-exiting framework. firstly, despite the potential to exit at earlier layers, key and value states for remaining layers are still required for processing subsequent tokens. while previous works have proposed the state copying mechanism @xcite to efficiently compute these states by reusing hidden states from the early-exited layer, our findings reveal that this method performs poorly with larger models and longer output sequences (see section 4.1). additionally, setting all layers as possible exit positions does not guarantee faster inference due to (1) the defective performance of earlier layers that can generate abnormally long sequence outputs, and (2) the computational overhead from confidence measurement at every layer (see section 4.2 and 4.3). lastly, achieving the desired level of latency and accuracy with early-exiting heavily depends on selecting the appropriate confidence threshold for the target task. this often entails significant efforts and additional computational overhead (see section 4.4). hence, these challenges call for a new approach that consistently (b) fast and robust early-exiting (free) figure 1 : overview of our free framework compared to the conventional early-exiting framework. free exhibits three key differences: (1) free employs a shallow-deep module that utilizes two exit points instead of employing all layers as exit points, (2) free replaces the state copying mechanism (yellow colored) with synchronized parallel decoding (red colored) to prevent performance degradation while accelerating inference speed, and (3) free utilizes an adaptive threshold estimator to determine the appropriate threshold values for each dataset during inference. in this paper, we introduce a fast and robust early-exiting (free) framework that incorporates a shallow-deep module and synchronized parallel decoding. our framework not only offers consistent speedup and performance even for larger models and longer output sequences, but also eliminates the need for the computationally expensive process of finding the appropriate exit threshold. specifically, the shallow-deep module bifurcates the computation paths into a shallow model (with a specified number of early layers) and a deep model (including all layers). our synchronized parallel decoding accumulates consecutive early-exited tokens that only pass through the shallow model until a non-exiting token is encountered. thereby, we synchronize the decoding process of the current non-exiting token with the previously stacked tokens, as shown in the left of figure 1b . this prevents performance degradation by utilizing actual attention computed key and value instead of approximated states through state copying, while it also achieves a more efficient approach compared to decoding each token autoregressively. furthermore, we devise a novel adaptive threshold estimator, as shown in the right of figure 1b , by leveraging the fact that parallel decoding outputs predictions even for early-exited tokens from the deep model. this estimator uses a beta mixture model (bmm) to capture the correlation between confidence scores and prediction alignment of two models, determining the proper confidence threshold for each dataset. in practice, we demonstrate the efficiency of our free framework on extensive generation tasks. ## preliminary a transformer network @xcite is composed of l layers, where each layer consists of two sublayers, a multi-head attention (mha) layer and a feed-forward network (ffn) layer. the computation for hidden states at time step t + 1 via stacked transformer blocks is as follows: after l th layer of the decoder network, the predicted token ŷt+1 , is determined by the probability output from a softmax classifier w l : after the current token is early-exited at the ℓ th layer, we need to calculate the key and value states for all deeper blocks in order to perform the self-attention for the subsequent tokens that pass through deeper blocks. for a more efficient approach of caching key and value states, the earlyexiting frameworks employ the state copying mechanism. it duplicates the hidden states of the earlyexited layer (i.e., h i t+1 = h ℓ t+1 , ∀@xmath0, l]), allowing us to compute the approximate key and value states required for the self-attention of transformer networks. @xcite have verified that state copying from lower layers does not have a detrimental effect on performance in the case of small-sized t5 models @xcite . ## re-evaluating early-exit framework in this section, we present four new findings from our re-evaluation of the early-existing framework. we utilized different model sizes of t5 @xcite on samsum @xcite and cnn/dailymail @xcite , and longt5-base @xcite architectures on multi-news @xcite and big-patent @xcite . ## novel early-exiting framework: free building upon the discoveries in section 4, we introduce a fast and robust early-exiting framework named free, leveraging a shallow-deep module and capitalizing on the structure of parallel decoding. furthermore, we present a confidence estimation algorithm designed to enhance the robustness of early-exiting within the free framework. ## experiments 6.1 experimental setup we conducted experiments on various sequence modeling tasks, including question answering (squad; @xcite ), machine translation (iwslt 2017 en-de; @xcite , and text summarization tasks using samsum, cnn/dailymail, multi-news, and bigpatent datasets. the longt5-base model was used for the multi-news and bigpatent datasets, while the t5-large model was used for the other datasets. all implementations are based on pytorch using huggingface @xcite . further details can be found in appendix b. ## conclusion we proposed free framework to address the challenges of conventional early-exiting frameworks for autoregressive language models. our approach incorporates three key components: (1) shallowdeep module, (2) synchronized parallel decoding, evaluate the quality of summaries written for a news article. rate each summary on four dimensions: {dimension_1}, {dimension_2}, {dimension_3}, and {dimension_4}. you should rate on a scale from 1 (worst) to 5 (best). article: {article} summary: {summary} given a new article, which summary is better? answer "summary 0" or "summary 1". you do not need to explain the reason. article: {article} summary 0: {summary_0} summary 1: {summary_1} figure 8 : the template for pairwise comparison. we measured twice by changing the order of summaries for a fair comparison. and (3) adaptive threshold estimation. through extensive experiments on various generation tasks, we empirically demonstrated the superior performance of free framework, achieving significant acceleration in latency without compromising the quality of the generated output. limitations. our work addressed a fast and robust existing framework that can be efficiently utilized without concerns about performance degradation. however, our approach does have a few limitations which we discuss below: (1) our method requires additional computational resources to finetune the shallow model. however, as we have demonstrated, parameter-efficient fine-tuning methods would be a promising solution to overcome this limitation. (2) while our work demonstrates robustness in the depth of the shallow model, further investigation is required to determine the appropriate depth for various language models. this aspect remains an area for additional research. baseline rouge-l ( ) 40 45 50 normalized latency 0.2 0.4 0.6 0.8 1.0 baseline 34 36 38 40 42 normalized latency 0.2 0.4 0.6 0.8 1.0 wce kd-last kd-unif kd-dyna (a) samsum (b) cnn/dailymail figure 10: the trade-off between performance and normalized latency per sentence. we varied the exit thresholds in the range of {0.0, 0.1, 0.3, 0.5, 0.7, 0.9}. the latency values are normalized by the latency of a baseline, which is a simple fine-tuned full model.