DEED : Dynamic Early Exit on Decoder for Accelerating Encoder-Decoder Transformer Models (2024)

Abstract: Encoder-decoder transformer models have achieved great success on various vision-language (VL) and language tasks, but they suffer from high inference latency. Typically, the decoder takes up most of the latency because of the auto-regressive decoding. To accelerate the inference, we propose an approach of performing Dynamic Early Exit on Decoder (DEED). We build a multi-exit encoder-decoder transformer model which is trained with deep supervision so that each of its decoder layers is capable of generating plausible predictions. In addition, we leverage simple yet practical techniques, including shared generation head and adaptation modules, to keep accuracy when exiting at shallow decoder layers. Based on the multi-exit model, we perform step-level dynamic early exit during inference, where the model may decide to use fewer decoder layers based on its confidence of the current layer at each individual decoding step. Considering different number of decoder layers may be used at different decoding steps, we compute deeper-layer decoder features of previous decoding steps just-in-time, which ensures the features from different decoding steps are semantically aligned. We evaluate our approach with three state-of-the-art encoder-decoder transformer models on various VL and language tasks. We show our approach can reduce overall inference latency by 20%-74% with comparable or even higher accuracy compared to baselines.

## introduction transformer models with auto-regressive decoders have shown great success on vision-language (vl) @xcite @xcite @xcite and language tasks @xcite @xcite . among many successful models tackling these * work conducted during an internship at amazon. in addition, even using one decoder layer can achieve decent accuracy (77.5% vs. 81.5%), implying that most examples do not need all decoder layers during inference. nevertheless, encoder-decoder models rely on the auto-regressive decoding to bring its ability into full play at inference. with auto-regressive decoding, each output token is generated conditioned on previous tokens. therefore, it has to generate tokens one after another, and repeat the feed-forward in each layer as many times. this mechanism leads to high inference latency in the decoder, and makes the decoder take up most of the total inference latency, as shown in figure 1 . interestingly, even using only one decoder layer, an encoder-decoder model can still get decent prediction accuracy (see figure 1 ), which means samples got correct by the one decoder layer do not need the excessive computations in the deeper decoder layers. inspired by these facts, we propose an approach to dynamically allocate adequate amount of computation at a particular decoding step in order to speed-up inference without sacrificing accuracy. specifically, we build dynamic early exit on decoder (deed), a multi-exit model with an early exit strategy to let the model decide whether or not to exit at a specific decoder layer at each decoding step dynamically. following existing work @xcite @xcite @xcite @xcite , we employ confidence-based dynamic early exit where the decoder may decide to exit when it is confident about its prediction. unlike encoder acceleration, the dynamic early exit for auto-regressive decoder is more challenging. the challenge is two-fold: to address the second challenge, we propose a novel algorithm that dynamically computes required deeper-layer features for previous decoding steps just-in-time. this algorithm effectively resolves the semantic misalignment among different layers at different generation steps. in contrast, the existing work, dat @xcite and calm @xcite , which uses the shallow-layer features as the deeper-layer features directly for later decoding steps, failed to mitigate such semantic misalignment and thus substantially undermine the generative ability of the model. our contributions are summarized as follows: ## related work encoder-decoder models for vision-language and language tasks transformer models, when first proposed, consist of encoder and decoder and are mainly for language tasks @xcite . following researches improve encoderdecoder transformer models by introducing better unsupervised pre-training strategies and small model architecture changes @xcite . recently, encoder-decoder transformer models have pushed the edge for vision-language (vl) tasks @xcite @xcite @xcite @xcite because of strong representation ability of encoder and generative ability of decoder. for example, flamingo @xcite uses a vision encoder to encode input images and a text decoder to generate text predictions for various vl tasks. latr @xcite utilizes the sequence generation ability in decoder and layout in multi-modality learning and achieves state-ofthe-art accuracy on text-based vqa tasks. ofa @xcite proposes a unified sequence-tosequence learning framework to incorporate various vl tasks into the encoder-decoder scheme. our work focuses on accelerating the decoder inference for this type of encoder-decoder transformer models on vl and language tasks. dynamic early exit using dynamic early exit (dee) is a popular strategy to reduce the inference latency of transformer models @xcite @xcite @xcite @xcite @xcite @xcite . for example, deebert @xcite and romebert @xcite applies dee to bert (kenton and toutanova, 2019) based on classification confidence scores from different encoder layers. berxit @xcite learns a policy for dynamic early exit. tokee @xcite introduces a token-level early exit approach for sequence labelling. however, these encoder-focused approaches cannot be applied to transformer decoders directly, due to the challenges imposed by the auto-regressive mechanism in decoder models. dat @xcite is one approach tackling decoder early exit. it introduces a haltand-copy approach, which halts the computation at a layer if the prediction is confident, and copies the feature from shallow decoder layers to deeper layers in later decoding steps when needed. calm @xcite follows the same haltand-copy approach for decoder early exit. however, this approach suffers strong semantic misalignment because the semantic information from different decoder layers are not compatible. thus the later decoding step at deeper layers cannot obtain meaningful features from previous steps, leading to significant accuracy drops. in contrast, our approach dynamically computes the deeper-layer features from earlier steps just-in-time to resolve the semantic misalignment and to achieve high accuracy. ## approach we propose deed, a dynamic early exit on decoder approach to accelerate encoder-decoder transformer models for vl and language tasks. specifically, we leverage confidence-based step-level dynamic early exit to decide which decoder layer to exit based on how confident we are at each decoding step. at training, we train our multi-exit model with deep supervision @xcite , where the output features of each decoder layer are input to a shared generation head and supervised using the ground truth. at inference, we apply dynamic early exit on the auto-regressive decoder. at each decoding step, the model decides how many decoder layers to use based on its confidence about the output token -hence different number of layers may be used at different decoding steps. in the follow sections, we first introduce the auto-regressive decoding process and the challenge of semantic misalignment in dynamic early exit on decoder in section 3.1. then we describe our multiexit model architecture and our training strategy in section 3.2. finally we show how we resolve the semantic misalignment problem with just-in-time computation of decoder features in section 3.3. ## conclusions we propose deed, a multi-exit model with step-level dynamic early exit on decoder for encoder-decoder transformer model acceleration. deed leverages confidence-based step-level dynamic early exit to reduce the computation at each decoding step. to improve the accuracy when exiting at shallow layers, we build a multi-exit model leveraging multiple techniques including deep supervision, shared generation head, adaptation modules, and emphasizing the learning of the final decoder layer. we apply our approach to three state-of-the-art encoder-decoder transformer models. results on various vision-language and language tasks show that our approach significantly reduces the inference latency with comparable or even higher accuracy compared to baselines. in the future, we will explore deed for decoder-only models. this is the appendix for the main deed paper. here we discuss the architecture and results of latr++, and the results of our reproduced ofa vs. the original ofa results.