Embedding Recycling for Language Models (2023)

Abstract: Real-world applications of neural language models often involve running many different models over the same corpus. The high computational cost of these runs has led to interest in techniques that can reuse the contextualized embeddings produced in previous runs to speed training and inference of future ones. We refer to this approach as embedding recycling (ER). While multiple ER techniques have been proposed, their practical effectiveness is still unknown because existing evaluations consider very few models and do not adequately account for overhead costs. We perform an extensive evaluation of ER across eight different models (17 to 900 million parameters) and fourteen tasks in English. We show how a simple ER technique that caches activations from an intermediate layer of a pretrained model, and learns task-specific adapters on the later layers, is broadly effective. For the best-performing baseline in our experiments (DeBERTa-v2 XL), adding a precomputed cache results in a 90% speedup during training and 87-91% speedup for inference, with negligible impact on accuracy. Our analysis reveals important areas of future work.

## introduction large pretrained language models form the foundation of modern nlp, and continue to push the state-of-the-art on a wide range of natural language processing tasks @xcite @xcite . larger models tend to offer superior accuracies @xcite , but also entail higher computational costs. the steep computational cost associated with large neural language models slows down experimentation, increases financial barriers to the technology, and contributes to global climate change @xcite . our work studies how to reduce computational cost for workloads in which many distinct models are run over the same text. for example, a scholarly search tool that helps users find and understand relevant literature may run separate models for entity recognition, topic classification, relation extraction, summarization, question answering, and so on over a large corpus of papers. new and improved models for the tasks are developed frequently, necessitating additional runs. the need for repeated model runs has also been noted for other applications in previous work, including news applications @xcite and virtual assistants @xcite . further, repeated runs also occur very frequently during model development, when exploring model variants or executing multiple training epochs. recent work has introduced ways to reduce computational cost in such settings by re-using model activations from one task to speed up other ones @xcite . a pretrained language model's internal activations form a contextualized embedding, which reflects syntactic and semantic knowledge about the input text @xcite @xcite which can be useful across a variety of downstream tasks. we define embedding recycling (er) as the technique of caching certain activations from a previous model run, and re-using them to improve the efficiency of future training and inference. recycling imposes a small computation time cost the first time a model processes a text, in order to compute and populate the cache. thereafter, all subsequent runs on the text can use the precomputed cache, improving efficiency. while previous work has shown the promise of er approaches, the existing evaluations are limited. for example, @xcite and @xcite each evaluate er for only one or two base models. likewise, for er techniques that cache activations on persistent storage, the storage and time cost of the cache itself has yet to be quan-tified. in this paper, we present a more comprehensive evaluation of er with several models and tasks, along with a thorough efficiency analysis. we study a simple layer-recycling er method that caches the activations from an intermediate layer of a pretrained model, and uses those cached activations as the starting point when the same input sequence is seen again during fine-tuning or inference. we show that even this simple method yields substantial improvements to throughput at small or no cost to accuracy on average. our results provide the strongest evidence to date that er can be a practically important technique for reducing costs for nlp systems, but as we discuss in section 6, they also suggest important challenges that must be addressed in future work. our contributions are summarized below: ## related work the embedding recycling technique we investigate is based on findings from prior work suggesting that not all layers of a pretrained transformer are equally important for end-task finetuning. shallower layers tend to converge earlier in training than deeper layers @xcite , and weights of later layers change more than earlier ones @xcite , suggesting that earlier layers tend to extract universal features whereas later layers focus on task-specific modeling. @xcite find that 90% of fully fine-tuned performance can be reached when finetuning only the final quarter of a transformer's layers and leaving the rest frozen. several proposed methods vary the number of frozen layers over the course of training, approaching or exceeding the performance of fully finetuned models while substantially speeding up the training process @xcite @xcite . similar to our approach, some dynamic freezing methods also employed caching mechanisms @xcite , but the dynamic number of frozen layers means the cache applies only at training time and only for a single task. in contrast, we cache embeddings from the pretrained model, which can then be reused across multiple downstream tasks and applied at inference time as well. other recent studies have sought to improve model inference speed by skipping computations in later layers. @xcite found that in some cases up to half of the layers can be removed from the model with only a 1-3% drop in task performance. early exit strategies have also been proposed, which allow the model to dynamically decide when to skip later layers @xcite . skipbert @xcite combined early exiting with an approach in which cached n-gram embeddings approximate the intermediate activations of new inputs. @xcite explored prompt-tuning as a parameterefficient approach for adapting frozen language models without adjusting model weights, conditioning language models with soft prompts to perform downstream tasks. precomputing text representations to speed up future processing on the same data is commonly done when creating fixed-size document-level embeddings for use on document-level tasks @xcite ; in contrast, we study contextualized token-level embeddings that can be used for tasks such as named entity recognition (ner) and question answering. readonce transformers @xcite do consider multitask variable-length document representations, but do so in a setting where a cached document representation is paired with a query text (such as a question or prompt); the approach is pretrained with qa data and evaluated on qa and summarization, rather than tasks such as text classification or ner where the entire input can be cached. @xcite propose an approach similar to ours that caches general-purpose token-level model representations, trained in a multi-task setting; however, that approach only applies a small mlp to the stored representations and reports a meaningful drop in accuracy (greater than 2% on average) compared to fully fine-tuned models. we find that reusing the later layer parameters of a pretrained transformer in addition to the cached activations often enables us to essentially match fully fine-tuned model accuracy while reducing computational cost. @xcite combine layer freezing and knowledge distillation to create a multi-task model. they do not consider caching activations on persistent storage as we do, but instead re-use activations across tasks at inference time via a branching multitask model. they use a two stage process where 12-n layers are fine-tuned for each individual task keeping n frozen layers. this is followed by distillation of the n layers for further computational gains. we take advantage of the parameter efficient adapter modules @xcite , and replace this process with a single step of fine-tuning a frozen base model that has adapters attached only to the deeper layers. our work also has connections to work on memory-and retrieval-augmented language modeling. prior work on using memory (e.g., @xcite ; @xcite ; rae and razavi (2020); @xcite ) generally focuses on modeling long-range context and caching representations of older history in a sequence, while work on retrieval (e.g., @xcite al. (2020); karpukhin et al. (2020) ) focuses on fetching text from a knowledge base or corpus to serve as additional context. in both cases, the aim is to use representations of additional text (from earlier in a document or from a knowledge base) to improve modeling of new inputs. in contrast, our work focuses on caching the representations of an entire sequence to speed up computation for new tasks. ## methods in the transformer architecture @xcite , an input sequence x of length s and dimension d is transformed with a function f : r s×@xmath0, ..., f (n ) as follows: our approach is to cache the output representations h @xmath1. we refer to this process of caching and reusing the output representations of a layer as layer recycling. this enables us to reduce the size of the transformer model from n layers to nk layers, reducing the computational cost during fine-tuning and inference. note that the key requirement of layer recycling is that we first need to process the entire data with the transformer model and cache the representations, so that we could later reuse these representations many times during fine-tuning and inference on new tasks. we experiment with two types of layer recycling approaches as explained next. we start with a pretrained transformer f (e.g., bert) consisting of f ( foot_0 ) , ..., f (k) , ..., f (n ) layers. during the first epoch of fine-tuning for any given task, we run the transformer over a corpus c and cache the output representations of layer k for each instance c in c, i.e., h k @xmath2. however, for every subsequent epoch of fine-tuning using the same transformer model, we only run and fine-tune the latter n -k layers f (k+1) , ..., f (n ) . we can either train all of the weights in the layers (which we refer to as reduced models), or only train adapter modules added on the layers (discussed below). in either case, for the instance c in the dataset c we simply retrieve and use the previously cached representation h k @xmath3. this avoids the extra computation through layers f (1) , ..., f (k) but adds a small cost for retrieving the representation from storage (see subsection 5.4 for efficiency analysis). ## experimental setup we now present our experiments evaluating whether recycled embeddings can be paired with reduced large language models to maintain accuracy while improving training and inference speed. we explore the effectiveness of embedding recycling across a variety of different tasks, datasets, and transformer models. ## discussion and future work our experiments raise several questions and suggest multiple avenues for future work, including: ## conclusion we have presented embedding recycling, a general technique for reusing previous activations of neural language models to improve the efficiency of future training and inference. we show how a simple technique of caching a layer of activations in a pretrained model is effective. we validate our approach in experiments across fourteen tasks and eight model architectures. we find that recycling typically has small or no impacts to accuracy on average, but does yield substantial throughput increases demonstrated through a careful efficiency analysis. we also discuss several open challenges for future work.