General Purpose Text Embeddings from Pre-trained Language Models for Scalable Inference (2020)

Abstract: The state of the art on many NLP tasks is currently achieved by large pre-trained language models, which require a considerable amount of computation. We aim to reduce the inference cost in a setting where many different predictions are made on a single piece of text. In that case, computational cost during inference can be amortized over the different predictions (tasks) using a shared text encoder. We compare approaches for training such an encoder and show that encoders pre-trained over multiple tasks generalize well to unseen tasks. We also compare ways of extracting fixed- and limited-size representations from this encoder, including pooling features extracted from multiple layers or positions. Our best approach compares favorably to knowledge distillation, achieving higher accuracy and lower computational cost once the system is handling around 7 tasks. Further, we show that through binary quantization, we can reduce the size of the extracted representations by a factor of 16 to store them for later use. The resulting method offers a compelling solution for using large-scale pre-trained models at a fraction of the computational cost when multiple tasks are performed on the same text.

## introduction large pre-trained language models achieve stateof-the-art performance on many natural language processing (nlp) tasks @xcite @xcite . however, inference for these models requires significant computational resources, which limits their practical use. recent trends show that scaling models up @xcite @xcite in terms of computation still improves end task performance, raising questions this computational burden is exacerbated by the need to fine-tune end-to-end a separate model for each task. since each model has a new set of parameters, none of the computation can be shared by models for different tasks during inference. this is particularly inefficient in real-world settings that require multiple predictions about each input. for example, given a news article, we may want to predict its topic @xcite , sentiment @xcite @xcite , overall text quality @xcite , whether it is humorous @xcite or offensive @xcite and so on. knowledge distillation (kd) is a way of reducing the computation required by large pre-trained lms @xcite . however, there is a sizeable gap in accuracy between the best models using knowledge distillation and the full fine-tuned models. another way of speeding up computation is through system optimizations such as quantization and operator fusion @xcite . these techniques can reduce the amount of computation significantly, but may not be sufficient by themselves and can be combined with the methods we discuss. in this paper we look at new ways to make inference computationally efficient focusing on the case where different models (models for different tasks) are run over the same piece of text. we propose new methods to run multiple task-specific models in a way that amortizes the computation over the different tasks. the central idea is to compute the activations for the full model once and use smaller task-specific models on top of it. we explore three possible ways for sharing computation. the first solution is inspired by work on general purpose text encoders @xcite ; hill task i single-task encoder feature extractor ! ! task 1 multi-task encoder feature extractor ! ! task k-1 ... multi-task encoder task k feature extractor ( or ) classification head i ! classif. head 1 ! classif. head k-1 ! ... classification head k ! ❄ ❄ ! (a) single-task finetuning. task i single-task encoder feature extractor ! ! task 1 multi-task encoder feature extractor ! ! task k-1 ... multi-task encoder task k feature extractor ( or ) classification head i ! classif. head 1 ! classif. head k-1 ! ... classification head k ! ❄ ❄ ! (b) multi-task pre-training. task i single-task encoder feature extractor ! ! task 1 multi-task encoder feature extractor ! ! task k-1 ... multi-task encoder task k feature extractor ( or ) classification head i ! classif. head 1 ! classif. head k-1 ! ... classification head k ! ❄ ❄ ! (c) leave-one-task-out finetuning. et al., 2016; conneau et al., 2017; subramanian et al., 2018), which produce fixed-size representations (i.e. the second solution is a multi-task system @xcite @xcite , where a single model is jointly trained to handle many tasks (see figure 1b ). if most layers are shared, the overall inference cost can be nearly k times less than for k separate single-task models, while providing competitive task accuracy @xcite @xcite . however, multi-task systems work best when the set of tasks is known in advance. adding new tasks requires retraining the multi-task model and re-incurring training costs, thus limiting the utility of this approach in real-world systems where new classification tasks may be introduced periodically. we propose a third solution: a multi-task encoder that is shared across tasks and produces limited-size representations that grow with the length of the input, similar to contextualized word representations @xcite . we evaluate our representations on 14 text classification tasks using a leave-one-task-out evaluation protocol (see figure 1c ), where a multi-task encoder model is trained on k -1 tasks, frozen and used as a static feature extractor for an unseen k th task. 1 we find an important ingredient to performing well on an unseen (k th ) task is to extract features from multiple layers and positions of the encoder. ultimately, our general purpose encoders offer a better tradeoff between task accuracy and inference cost than either fixed-size representations or distilled models, while requiring minimal additional computation to handle new tasks. we also consider the case in which not all of the predictions can be done at the same time and intermediate representations have to saved. in that context, we study the relationship between representation size and end-task performance. we find that features extracted by our encoders are amenable to heavy quantization enabling a 16x reduction in the size of the extracted features with negligible impact on unseen task performance. multi-task learning (mtl) has a rich history in machine learning @xcite and nlp @xcite . multi-task models can potentially leverage similarities across tasks to achieve higher endtask accuracy than single-task models @xcite @xcite . compared to single-task models, a multi-task model can also be more efficient during inference by sharing computation across tasks. most work in multi-task learning assumes that the set of end-tasks is fixed and known in advance and training is performed for all tasks together. this setup can present challenges in the real world where tasks may require different retraining schedules and new tasks may be frequently added or removed. general purpose text encoders are usually pre-trained with a mix of supervised and selfsupervised training objectives and produce fixedsize representations @xcite @xcite . unlike multi-task learning, general purpose text encoders are typically evaluated on unseen tasks, which is more representative of real-world settings in which new tasks may be added periodically. unfortunately, these approaches often underperform single-task baselines @xcite @xcite . another line of work has explored adapting pretrained models by adding additional task-specific capacity at each layer @xcite , however these methods do not improve inference efficiency since there is no task-independent computation that can be shared across tasks. knowledge distillation @xcite is a technique where a more efficient student model is trained to mimic the behaviour of a larger or ensembled teacher model. a knowledge distilled version of bert @xcite has been proposed to reduce the computation required by large pre-trained language models. distilroberta reaches 95% of roberta-base's performance on glue while being twice faster. quantization and other compression techniques have been explored for word embeddings @xcite and sentence embeddings @xcite . recent work has also explored quantization for contextualized word representations, generally showing that quantization-aware training is necessary to achieve reasonable end task performance (zafrir @xcite . quantization is complementary to the approaches we consider and is explored more in section 5. ## experimental setup our goal is to develop text encoders that produce representations which achieve high accuracy for multiple task with little task-specific processing. we first introduce our tasks, encoder models and finetuning framework. ## feature extraction and pooling a common way to extract features from bertlike models is to take the representation in the last transformer layer corresponding to a special cls token prepended to the input sequence @xcite . recent work has also explored extracting features from every position and layer, then linearly combining the layers with task-specific weights @xcite . we propose a more general framework for extracting features, shown in figure 2 . we extract features from several layers of the encoder and k, v task-specific query ! classification head ! multi-head attention ! q ... ❄ ❄ ❄ figure 3: our proposed multi-head attention pooler. the extracted features are frozen ( ) and used as both the keys (k) and values (v). each task has its own query (q), multi-head attention module and classification head, all of which are fine-tuned ( ). layer-wise pooling approaches: position-wise pooling approaches: ## storage considerations and quantization in a real-world settings it may be necessary to store extracted features for later use, such as when new tasks are introduced that require "backfilling" classifications for older content @xcite . storage costs quickly become impractical when pooling over multiple hidden layers and positions, with some approaches (section 4) requiring features from every layer and position in the encoder. for roberta large , with 24 layers and 1024 dimension representations, a 50 token input would thus emit 50 * 24 * 1024 half-precision floating point numbers and require 2.3mb of storage! we consider quantization methods, described below, for reducing the storage of extracted features. with quantization, we replace floating point numbers with alternative representation formats that have reduced bit width. we will show in section 6 that extracted features are surprisingly robust: they show little degredation in end-task accuracy even with binary quantization. recent work has made similar observations in the context of 8-bit integer quantization for bert model weights and activations @xcite . we explore both 8-bit (uint8) and 1-bit (boolean) quantization of extracted features (see ## results and discussion ## conclusion we study how to improve the efficiency of largescale pre-trained models so that can be used in practical settings. we show that when several tasks are performed on a single piece of text, the computation can be effectively amortized reducing the amount of computation per task. compared to distillation, the shared computation method achieves higher accuracy and reduces computational cost after 7 tasks need to be performed on the same piece of text. we show that the shared features can be quantized with very little loss in accuracy, which means that the intermediate computation can be stored for later use. in total, the techniques that we present provide a compelling solution for running large-scale pre-trained models in applications where multiple predictions are made on the same piece of text.