Length-Adaptive Transformer: Train Once with Length Drop, Use Anytime with Search (2021)

Abstract: Despite transformers’ impressive accuracy, their computational cost is often prohibitive to use with limited computational resources. Most previous approaches to improve inference efficiency require a separate model for each possible computational budget. In this paper, we extend PoWER-BERT (Goyal et al., 2020) and propose Length-Adaptive Transformer that can be used for various inference scenarios after one-shot training. We train a transformer with LengthDrop, a structural variant of dropout, which stochastically determines a sequence length at each layer. We then conduct a multi-objective evolutionary search to find a length configuration that maximizes the accuracy and minimizes the efficiency metric under any given computational budget. Additionally, we significantly extend the applicability of PoWER-BERT beyond sequence-level classification into token-level classification with Drop-and-Restore process that drops word-vectors temporarily in intermediate layers and restores at the last layer if necessary. We empirically verify the utility of the proposed approach by demonstrating the superior accuracy-efficiency trade-off under various setups, including span-based question answering and text classification. Code is available at https://github.com/clovaai/lengthadaptive-transformer .

## introduction pre-trained language models @xcite @xcite @xcite have achieved notable improvements in various natural language processing (nlp) tasks. most of them rely on transformers @xcite , and the number of model parameters ranges from hundreds of millions to billions @xcite @xcite . despite this high accuracy, excessive computational overhead during inference, both in terms of time and memory, has hindered its use in real applications. this level of excessive computation has further raised the concern over energy consumption as well @xcite @xcite . recent studies have attempted to address these concerns regarding large-scale transformers' computational and energy efficiency (see §7 for a more extensive discussion.) among these, we focus on power-bert @xcite which progressively reduces sequence length by eliminating word-vectors based on the attention values as passing layers. power-bert establishes the superiority of accuracy-time trade-off over earlier approaches @xcite @xcite . however, it requires us to train a separate model for each efficiency constraint. in this paper, we thus develop a framework based on power-bert such that we can train a single model that can be adapted in the inference time to meet any given efficiency target. in order to train a transformer to cope with a diverse set of computational budgets in the inference time, we propose to train once while reducing the sequence length with a random proportion at each layer. we refer to this procedure as lengthdrop, which was motivated by the nested dropout @xcite . we can extract sub-models of shared weights with any length configuration without requiring extra post-processing nor additional fine-tuning. it is not trivial to find an optimal length configuration given the inference-time computational budget, although it is extremely important in order to deploy these large-scale transformers in practice. once a transformer is trained with the proposed lengthdrop, we search for the length configuration that maximizes the accuracy given a computational budget. because this search is combinatorial and has multiple objectives (accuracy and efficiency), in this work, we propose to use an evolutionary search algorithm, which further allows us to obtain a full pareto frontier of accuracy-efficiency trade-off of each model. power-bert, which forms the foundation of the proposed two-stage procedure, is only applicable to sequence-level classification, because it eliminates some of the word vectors at each layer by design. in other words, it cannot be used for token-level tasks such as span-based question answering @xcite because these tasks require hidden representations of the entire input sequence at the final layer. we thus propose to extend power-bert with a novel drop-and-restore process ( §3.3), which eliminates this inherent limitation. word vectors are dropped and set aside, rather than eliminated, in intermediate layers to maintain the saving of computational cost, as was with the original power-bert. these set-aside vectors are then restored at the final hidden layer and provided as an input to a subsequent task-specific layer, unlike the original power-bert. the main contributions of this work are twofold. first, we introduce lengthdrop, a structured variant of dropout for training a single length-adaptive transformer model that allows us to automatically derive multiple sub-models with different length configurations in the inference time using evolutionary search, without requiring any re-training. second, we design drop-and-restore process that makes power-bert applicable beyond classification, which enables power-bert to be applicable to a wider range of nlp tasks such as span-based question answering. we empirically verify length-adaptive transformer works quite well using the variants of bert on a diverse set of nlp tasks, including squad 1.1 @xcite and two sequence-level classification tasks in glue benchmark @xcite . our experiments reveal that the proposed approach grants us fine-grained control of computational efficiency and a superior accuracy-efficiency trade-off in the inference time compared to existing approaches. ## background in this section, we review some of the building blocks of our main approach. in particular, we review transformers, which are a standard backbone used in natural language processing these days, and power-bert, which was recently proposed as an effective way to train a large-scale, but highly effi-cient transformer for sequence-level classification. ## length-adaptive transformer in this section, we explain our proposed framework which results in a transformer that reduces the length of a sequence at each layer with an arbitrary rate. we call such a resulting transformer a length-adaptive transformer. we train length-adaptive transformer with lengthdrop which randomly samples the number of hidden vectors to be dropped at each layer with the goal of making the final model robust to such drop in the inference time. once the model is trained, we search for the optimal trade-off between accuracy and efficiency using multi-objective evolutionary search, which allows us to use the model for any given computational budget without fine-tuning nor retraining. at the end of this section, we describe drop-and-restore process as a way to greatly increase the applicability of power-bert which forms a building block of the proposed framework. in short, we train a length-adaptive transformer once with lengthdrop and drop-and-restore, and use it with an automatically determined length configuration for inference with any target computational budget, on both sequencelevel and token-level tasks. ## experiment setup datasets we test the proposed approach on both sequence-level and token-level tasks, the latter of which could not have been done with the original power-bert unless for the proposed drop-and-restore. we use mnli-m and sst-2 from glue benchmark @xcite , as was done to test power-bert earlier, for sequence-level classification. we choose them because consistent accuracy scores from standard training on them due to their sufficiently large training set imply that they are reliable to verify our approach. we use squad 1.1 @xcite for token-level classification. evaluation metrics we use the number of floating operations (flops) as a main metric to measure the inference efficiency given any length configuration, as it is agnostic to the choice of the underlying hardware, unlike other alternatives such as hardware-aware latency @xcite or energy consumption @xcite . we later demonstrate that flops and wall-clock time on gpu and cpu correlate well with the proposed approach, which is not necessarily the case for other approaches, such as unstructured weight pruning @xcite . pre-trained transformers since bert was introduced by @xcite , it has become a standard practice to start from a pre-trained (masked) language model and fine-tune it for each downstream task. we follow the same strategy in this paper and test two pre-trained transformerbased language models; bert base @xcite and distilbert @xcite , which allows us to demonstrate that the usefulness and applicability of our approach are not tied to any specific architectural choice, such as the number of layers and the maximum input sequence length. although we focus on bert-based masked language models here, the proposed approach is readily applicable to any transformer-based models. learning we train a length-adaptive transformer with lengthdrop probability and layer-drop probability both set to 0.2. we use n @xmath0. we start fine-tuning the pre-trained transformer without drop-and-restore first, just as goyal et al. (2020) did with power-bert. we then continue fine-tuning it for another five epochs with dropand-restore. this is unlike the recommended three epochs by @xcite , as learning progresses slower due to a higher level of stochasticity introduced by lengthdrop and layerdrop. we use the batch size of 32, the learning rate of 5e -5 for squad 1.1 and 2e -5 for mnli-m and sst, and the maximum sequence length of 384 for squad 1.1 and 128 for mnli-m and sst. ## results and analysis efficiency-accuracy trade-off we use squad 1.1 to examine the effect of the proposed approach on the efficiency-accuracy trade-off. when the underlying classifier was not trained with length-drop, as proposed in this paper, the accuracy drops even more dramatically as more word vectors are dropped at each layer. the difference between standard transformer and length-adaptive transformer is stark in figure 2 . this verifies the importance of training a transformer in a way that makes it malleable for inference-time re-configuration. when the model was trained with the proposed lengthdrop, we notice the efficacy of the proposed approach of using evolutionary search to find the optimal trade-off between inference efficiency and accuracy. the trade-off curve from the proposed search strategy has a larger area-under-curve (auc) than when constant-rate length reduction was used to meet a target computational budget. it demonstrates the importance of using both lengthdrop and evolutionary search. we make a minor observation that the proposed approach ends up with a significantly higher accuracy than distillbert when enough computational budget is allowed for inference (log @xmath1. this makes our approach desirable in a wide array of scenarios, as it does not require any additional pre-training stage, as does distilbert. with a severe constraint on the computational budget, the proposed approach could be used on distilbert to significantly improve the efficiency without compromising the accuracy. ## comparison with other works our framework allows a novel method for anytime prediction with adaptive sequence length given any transformers. thus, our goal is not state-of-the-art classification accuracy, although our experimental results ( §5) demonstrate that our method still attains a good accuracy level. we emphasize that other adaptive computation works ( §7) are orthogonal with ours, meaning that various adaptive dimensions (sequence length, depth, attention head, hidden dimension, etc.) can be jointly used. in other words, even if other adaptive methods show better curves than ours, our method and theirs can boost each other when combined. we provide some comparison results with power-bert (not anytime prediction method) and dynabert @xcite (concurrent adaptive computation method) as follows. comparison with power-bert according to @xcite , power-bert achieves 2.6x speedup for mnli-m and 2.4x speedup for sst-2 by losing 1% of their accuracy. length-adaptive transformer obtains a 2.9x speedup in terms of flops without losing accuracy on mnlim and sst-2. considering figure 3 , our speedup in execution time would be close to 2.9x in the same setting of power-bert where the time measurement is done with a batch size of 128 on gpu. it indicates that our model offers a better trade-off than power-bert, even with a single model. comparison with dynabert according to @xcite , dyanbert obtains a gain of +1.0, +0.1, +0.4 for the best accuracy in squad 1.1, mnli-m, and sst-2, respectively, while length-adaptive transformer achieves a gain of +1.1, +0.6, +0.3. these results imply that length-adaptive transformer can give a comparable (or better) performance with dynabert. ## related work the main purpose of the proposed algorithm is to improve the inference efficiency of a large-scale transformer. this goal has been pursued from various directions, and here we provide a brief overview of these earlier and some concurrent attempts in the context of the proposed approach. weight pruning weight pruning @xcite focuses on reducing the number of parameters that directly reflects the memory footprint of a model and indirectly correlates with inference speed. however, their actual speed-up in runtime is usually not significant, especially while executing a model with parallel computation using gpu devices @xcite . adaptive architecture there are three major axes along which computation can be reduced in a neural network; (1) input size/length, (2) network depth, and (3) network width. the proposed approach, based on power-bert, adaptively reduces the input length as the input sequence is pro-cessed by the transformer layers. in our knowledge, @xcite is the first work in this direction for transformers. funnel-transformer @xcite and multi-scale transformer language models @xcite also successfully reduce sequence length in the middle and rescale to full length for the final computation. however, their inference complexity is fixed differently with power-bert because they are not designed to control efficiency. more recently, tr-bert @xcite introduces a policy network trained via reinforcement learning to decide which vectors to skip. layerdrop @xcite drops random layers during the training to be robust to pruning inspired by @xcite . word-level adaptive depth in @xcite and @xcite might seemingly resemble with length reduction, but word vectors that reached the maximal layer are used for self-attention computation without updating themselves. escaping a network early @xcite based on the confidence of the prediction @xcite @xcite @xcite also offers a control over accuracyefficiency trade-off by changing a threshold, but it is difficult to tune a threshold for a desired computational budget because of the example-wise adaptive computation. slimmable neural networks @xcite reduce the hidden dimension for the any-time prediction. dynabert @xcite ) can run at adaptive width (the number of attention heads and intermediate hidden dimension) and depth. hardware-aware transformers @xcite construct a design space with arbitrary encoder-decoder attention and heterogeneous layers in terms of different numbers of layers, attention heads, hidden dimension, and embedding dimension. spatten @xcite performs cascade token and head pruning for an efficient algorithm-architecture co-design. structured dropout a major innovation we introduce over the existing power-bert is the use of stochastic, structured regularization to make a transformer robust to the choice of length configuration in the inference time. @xcite proposes a nested dropout to learn ordered representations. similar to lengthdrop, it samples an index from a prior distribution and drops all units with a larger index than the sampled one. search there have been a series of attempts at finding the optimal network configuration by solving a combinatorial optimization problem. in computer vision, once-for-all @xcite use an evolutionary search @xcite to find a better configuration in dimensions of depth, width, kernel size, and resolution given computational budget. similarly but differently, our evolutionary search is mutli-objective to find length configurations on the pareto accuracy-efficiency frontier to cope with any possible computational budgets. moreover, we only change the sequence length of hidden vectors instead of architectural model size like dimensions. sequence length shortformer @xcite initially trained on shorter subsequences and then moved to longer ones achieves improved perplexity than a standard transformer with normal training while reducing overall training time. novel architectures with the efficient attention mechanism @xcite @xcite @xcite are suggested to reduce the transformer's quadratic computational complexity in the input sequence length. @xcite and @xcite provide a survey of these efficient transformers and their benchmark comparison, respectively. ## conclusion and future work in this work, we propose a new framework for training a transformer once and using it for efficient inference under any computational budget. with the help of training with lengthdrop and dropand-restore process followed by the evolutionary search, our proposed length-adaptive transformer allows any given transformer models to be used with any inference-time computational budget for both sequence-level and token-level classification tasks. our experiments, on squad 1.1, mnlim and sst-2, have revealed that the proposed algorithmic framework significantly pushes a better pareto frontier on the trade-off between inference efficiency and accuracy. furthermore, we have observed that the proposed length-adaptive transformer could achieve up to 3x speed-up over the standard transformer without sacrificing accuracy, both in terms of flops and wallclock time. although our approach finds an optimal length configuration of a trained classifier per computational budget, it leaves a open question whether the proposed approach could be further extended to support per-instance length configuration by for instance training a small, auxiliary neural network for each computational budget. yet another aspect we have not investigated in this paper is the applicability of the proposed approach to sequence generation, such as machine translation. we leave both of these research directions for the future. our approach is effective, as we have shown in this paper, and also quite simple to implement on top of existing language models. we release our implementation at https://github.com/clovaai/lengthadaptive-transformer , which is based on hugging-face's transformers library @xcite , and plan to adapt it for a broader set of transformerbased models and downstream tasks, including other modalities @xcite @xcite .