Adapting Language Models to Compress Contexts (2023)

Abstract: Transformer-based language models (LMs) are powerful and widely-applicable tools, but their usefulness is constrained by a finite context window and the expensive computational cost of processing long text documents. We propose to adapt pre-trained LMs into AutoCompressors. These language models are capable of compressing long contexts into summary vectors, which are then accessible to the model as soft prompts. Summary vectors are trained with an unsupervised objective, whereby long documents are processed in segments, and summary vectors from all previous segments are used in language modeling. We fine-tune OPT and Llama-2 models on sequences of up to 30,720 tokens and show that AutoCompressors can utilize long contexts to improve perplexity. We evaluate AutoCompressors on in-context learning by compressing task demonstrations and find that summary vectors are good substitutes for plain-text demonstrations, increasing accuracy while reducing inference costs. Finally, we explore the benefits of pre-computing summary vectors for large corpora by applying summary vectors to retrieval-augmented language modeling and a passage re-ranking task. Overall, AutoCompressors emerge as a simple and inexpensive solution to extend the context window of LMs while speeding up inference over long contexts.

## introduction transformer-based @xcite language models (lms) have recently seen a sharp rise in popularity and are now receiving millions of queries, processing billions of tokens, and generating text for a wide variety of applications @xcite @xcite figure 1 : autocompressors process long documents by recursively generating summary vectors which are passed as soft prompts to all subsequent segments. 2022). with this rise in popularity comes the challenge for researchers to make lms more efficient, to speed up inference and to deploy lms at scale, while increasing their versatility, thus allowing users to process more data in new ways. with these goals in mind, we propose to teach pre-trained lms the ability to compress text into summary vectors. summary vectors are short soft prompts @xcite , one or two orders of magnitude shorter than the pre-compressed plain text, that are obtained from the output states of a language model. summary vectors serve two general purposes: they can help extend the language model's context window to very long documents with minimal computational overhead, and they help speed up inference on text for which summary vectors have been pre-computed and cached. our models, which we call autocompressors, are trained with a simple unsupervised learning objective that encourages the model to store essential information in the summary vectors. summary vectors are produced segment by segment from long documents and are used to improve language modeling in future segments (figure 1 ). our work builds on the recently proposed rmt architecture @xcite with a crucial difference: we introduce summary accumulation, in which summary vectors from all segments are concatenated to produce the summary of the entire document. we also train autocompressors with randomly segmented inputs so they can better compress contexts of variable lengths in downstream tasks. we show that these innovations improve long-range information retention and enable new ways of reasoning over multiple passages. autocompressors can be initialized with pretrained lms to produce powerful and versatile models. we fine-tune autocompressors from opt-2.7b @xcite and llama-2-7b @xcite models on sequences from 6,144 up to 30,720 tokens with a single nvidia a100 gpu of 80gb memory. we show that summary vectors are effective for improving perplexity over long documents and that these compression capabilities are robust to domain generalization. our analysis suggests that autocompressors are able to reason over summary vectors, making them useful for a diverse set of downstream applications. we apply autocompressors to in-context learning (icl) by compressing up to 90 in-context demonstrations. we consider 11 classification tasks, including 7 superglue tasks @xcite , and we find that summary vectors outperform few-shot icl with a comparable number of in-context tokens on 8 out of 11 tasks. finally, we explore two applications where autocompressors can reduce inference costs by pre-computing summary vectors for large corpora. first, we adopt a setting for retrieval-augmented language modeling @xcite . we find that for equal sequence lengths, using summary vectors achieves 1.5× the perplexity gains compared to plain-text passages, and outperforms retrievalaugmented methods for similar computational budgets. secondly, we consider a zero-shot passage re-ranking task @xcite . we establish that re-ranking passages based on their summary vectors achieves the best trade-off between re-ranking performance and inference throughput. in summary, our main contributions are the following: (1) we introduce a method for extending lms to long context windows under small-scale computational requirements by learning to generate summary vectors. we propose summary accumulation and training with randomized segmenting as key features of autocompressors. (2) we show that summary vectors encode useful information for downstream tasks and can be used to reduce the inference cost of in-context learning. (3) we demonstrate the benefits of pre-computing summary vectors for large corpora and using auto-compressors in conjunction with retrievers. most of these architectures typically require expensive training from scratch, or will deviate substantially from a pre-trained initialization. 2 moreover, many language models lack the inductive bias to extrapolate to longer sequences @xcite . while autocompressors could in principle be trained from scratch, we show that they offer a straightforward solution for extending the context window of pre-trained models to longer sequences. ## method we describe how we adapt a pre-trained language model to compress text into summary vectors. an overview of our architecture is shown in figure 1 . ## language modeling evaluation in this section, we train autocompressors and evaluate their long-range language modeling capabilities by sampling long sequences which we split into segments of 2,048 tokens. we fix the final segment and compress the previous n segments. we track the perplexity of the final segment when conditioning on the summary vectors for each n. we conduct our main experiments and ablations with opt models @xcite of 1.3b or 2.7b parameters, fine-tuned on 2b tokens from the pile @xcite . in section 4.1, we evaluate an autocompressor on sequences of 8,000 tokens and compare to an equivalent rmt model and an extended full attention baseline. in section 4.2, we fine-tune an autocompressor on sequences of 30,000 tokens to demonstrate the feasibility on very long sequences. finally, in section 4.3, we scale up autocompressors by fine-tuning a llama-2-7b model on 15b tokens from redpajama (togeth-erai, 2023). full model hyperparameters and data information can be found in appendix a. ## compressing demonstrations for in-context learning in this section, we study the usefulness of summary vectors for performing downstream tasks. we show that in-context demonstrations can reliably be compressed down into summary vectors to improve performance while also increasing efficiency on a diverse set of nlp benchmarks. evaluation we evaluate the in-context learning abilities of the autocompressor based on llama-2-7b from section 4.3 on eleven classification and multiple-choice question-answering datasets. for each dataset, we evaluate the effect of compressing 1, 2 or 3 segments of demonstrations into 50, 100 or 150 summary vectors. for each segment, we include as many demonstrations as possible until we reach 750 tokens. for sst-2, this corresponds to 30 demonstrations per segment on average. we compare this compression approach with the results obtained by prompting the model using 150 and 750 tokens' worth of plain-text demonstrations. we use contextual calibration @xcite ) and class-balanced sampling when these techniques improve performance on a validation set. for each dataset, we report the mean accuracy and standard deviation over 7 random seeds. the detailed settings for each dataset can be found in ## compressing retrieval corpora for efficient inference we study the usefulness of pre-computing summary vectors for large collections of documents. these can be stored and later retrieved for efficient inference. since inference is typically more expensive than storage, this approach has the potential to achieve good practical trade-offs. ## conclusion we have introduced a training strategy for adapting pre-trained lms into autocompressors, which recursively compress contexts into summary vectors. our experiments indicate that summary vectors retain important contextual information, that they can encode in-context demonstrations, and that they can be used in retrieval settings. summary vectors can also be pre-computed, cached and re-used. this offers practical efficiency gains by reducing the size of the attention window. significant future work remains in scaling autocompressors to bigger models and improving the quality of summary vectors to further close the gap with full attention over long-range contexts. limitations 1. we only apply autocompressors to opt models of up to 2.7b parameters and a llama model of 7b parameters. future work needs to establish how autocompressors perform for even larger models. as the summary vector dimension grows, there is promise for retaining more information per vector. 2. our results suggest that summary vectors ignore some useful information that is accessible via full attention. additionally, models do not always benefit from increasing the number of summary vectors. we suspect that the training signal for learning summary vectors efficiently might be limited by pre-trained models being very good at making predictions from the plaintext tokens in the current segment. future work is needed to improve this optimization. 3. summary accumulation still leads to quadratic complexity with increasing number of segments, albeit at a much lower rate than full attention. future work may explore ways to combine many summary vectors more efficiently.