Characterizing Prompt Compression Methods for Long Context Inference

Token pruning methods perform compression by discarding irrelevant tokens. Selective-Context (Li et al., 2023) uses a small language model to judge self-information of tokens. Then, tokens with low self-information are pruned from the original prompt. LLMLingua (Jiang et al., 2023b) is a similar method to Selective-Context but uses perplexity to determine the importance of tokens. LLMLingua first performs coarse-grained pruning by removing entire in-context examples and then performs fine-grained token pruning on the prompt. LongLLMLingua (Jiang et al., 2023c) is a modification of LLMLingua designed for long context prompt compression. Unlike LLMLingua, LongLLMLingua considers the perplexity of the question when conditioned on supporting documents to determine which documents are most relevant. After performing coarse-grained compression by removing irrelevant documents, fine-grained token pruning is performed by considering the perplexity of tokens before and after being conditioned on the question. The drop in perplexity after conditioning on the question is used to judge the relevance of a token. Tokens with low relevance are pruned. PCRL (Jung & Kim, 2023) uses reinforcement learning to train a policy network to remove tokens in the original context. Specifically, the state seen by the policy is the original context and the action taken by the policy is a binary string denoting whether each token in the original context is kept or removed. The ROUGE (Lin, 2004) between the output from the original context and the output from the compressed context is considered as the reward to maximize. The policy network is a frozen pre-trained small language model (e.g. GPT-2) a MLP head for binary classification. There has also been extensive research on token pruning methods for white-box Transformer models (Goyal et al., 2020; Kim & Cho, 2020; Kim et al., 2022; Wang et al., 2021). Such methods utilize the Transformer model’s attention map at each layer in order to determine which tokens are least attended to by other tokens. These tokens are pruned before the sequence proceeds to the next layer in the Transformer. For the purposes of black-box prompt compression, a smaller white-box model may be used for token pruning, with the unpruned tokens from the white-box model being sent to the black-box LLM.

Abstractive compression techniques rely on summarization techniques to reduce the length of the original context. RECOMP’s (Xu et al., 2023) abstractive compressor is a fine-tuned T5-Large (775M) model (Raffel et al., 2020) that summarizes the initial context into a more compact form. By prompting the summarizer with the question at inference time, they generate query-aware summaries. In the fine-tuning training data, they drive the summarization model to produce an empty string if a summarized context leads to performance degradation on the downstream task. To omit the fine-tuning process in RECOMP, it is also possible to use a larger LLM that can perform summarization. Prompt-SAW (Ali et al., 2024) uses a 7B Vicuna model (Chiang et al., 2023) to create a knowledge graph with the key entities and their relationships. Then, each entity-relation pair is encoded with an embedding model and similarity search is performed with the question embedding to determine the most relevant information to keep. PRCA (Yang et al., 2023a) uses a small language model, such as T5-Large, to generate a smaller context from the original context. In order to train the small language model, PRCA uses a two-stage training process. In the first stage, supervised training is performed so that the small language model can learn to produce summaries well. Then, in the second stage, proximal policy optimization is applied to train the small language to produce distilled contexts that perform well when given to the downstream LLM. Similarly to PCRL, the ROUGE score between the output from the policy’s compressed context and the output from the original context is used to form the reward for training.

Extractive compression selects relevant documents, sentences, or phrases from the original context. RECOMP also has an extractive compression method that is used to extract the most relevant sentences given the initial context and question. RECOMP trains an encoder model so that useful sentences have higher inner product with the question in the embedding space. In their evaluation, the encoder is fine-tuned from a contriever (110M) checkpoint (Izacard et al., 2021). Document rerankers perform a similar function to RECOMP’s extractive compressor. Reranker models take a question and document and output a relevance score for the document to the query. Rerankers are typically applied in RAG pipelines after an initial retrieval step to further refine the document set. Prior work (Nogueira & Cho, 2019) fine-tunes a BERT model (Devlin et al., 2019) for passage rereranking. There is also a line of work (Pradeep et al., 2023a, b) that fine-tunes 7B language models to perform zero-shot listwise reranking. An illustration of extractive compression and its comparison to abstrative compression and token pruning can be found in Figure 2.

Prompt compression methods may further be classified as query-aware or query-agnostic. Query-aware compression methods compress contexts differently depending on the question or task. On the other hand, query-agnostic compression methods do not rely on the question or task and thus compression may be performed offline only once. Since such methods do not have access to the downstream task, they operate by exploiting redundancy in natural language. LLMLingua-2 (Pan et al., 2024) performs query-agnostic compression by training a classifier model to identify and remove redundant tokens. Prompt-SAW also has a query-agnostic variant in which similar information elements in the constructed knowledge graph are de-duplicated. An illustration of query-aware and query-agnostic prompt compression is shown in  Figure 3. Furthermore, Table 1 gives a categorization of existing methods.

Extractive compression methods are represented via the reranker (blue). The reranker model has very strong performance across all models and datasets. There are many example data points where compression is performed and accuracy significantly increases. For example, on 2WikiMultihopQA with GPT-3.5-Turbo, the reranker is able to compress $7.75\times$ while increasing accuracy by 7.89 points. Similarly, on MuSiQue with Mixtral 8x7B, the reranker is able to compress $4.14\times$ while increasing accuracy by 7.16 points. On HotpotQA with Mixtral 8x7B, it is able to compress $3.55\times$ while increasing accuracy by 4.54 points. A significant advantage of extractive compression is that grammatical constructs are preserved, as pruning occurs at a coarse granularity. Retrieval based methods are a widely used extractive compression methodology in which relevant chunks are retrieved via similarity search on embeddings. However, as shown in Section 4.2.1 we see significant improvements in extractive compression when using a reranker model over standard retrieval. This is because reranker models use language models that take in both the query and context to assign relevance. In contrast, retrieval methods perform light-weight similarity search over embeddings. Therefore, the precise method used for extractive compression plays a significant role.

Abstractive compression methods (pink) often exhibits inferior performance compared to extractive compression. The primary challenge with abstractive compression arises from the use of smaller, potentially weaker models, which may omit crucial information or introduce hallucinations. This is particularly problematic in summarization tasks where the large model has to generate a summary based only on the weaker model’s summaries, which can potentially discard information that the large model would have preferred to keep. Concretely, on summarization datasets with GPT-3.5-Turbo and Mixtral 8x7B, query-agnostic abstractive compression lags behind extractive compression by 3-5 points. Additionally, on MultifieldQA, query-agnostic abstractive compression is typically 10-15 accuracy points below the reranker at the same compression ratio. Therefore online query-aware abstractive compression, as shown in Section 4.2.3, or fine-tuned summarizers may perform better than prompting out-of-the-box LMs for offline summarization.

There are three token pruning methods: LongLLMLingua (orange), reranker + LongLLMLingua (purple), reranker + token pruning (green). We observe that LongLLMLingua and reranker + LongLLMLingua typically exhibit the worst behavior across datasets. In Section B.1, we perform a sweep over LongLLMLingua hyper-parameters but do not see any significant improvement. Reranker + token pruning generally trails slightly behind the plain reranker method. We hypothesize that the lackluster performance of token pruning is due to the disruption of grammar and sentence comprehension caused by unstructured pruning. However, we notice that reranker + token pruning outperforms the reranker model for GovReport and QMSum on Mixtral 8x7B at higher compression rates. Similarly, on GPT-3.5-Turbo, reranker + token pruning is competitive with the plain reranker on summarization datasets at high compression rates. Nonetheless, the performance of the reranker + token pruning method trails the reranker on question-answering tasks. In general, token pruning methods appear better suited for aggregation-style tasks that require pieces of knowledge from all segments of the initial context. Furthermore, rather than using out-of-the-box language models, practitioners may see better results by training language models specifically for token pruning (Jung & Kim, 2023; Pan et al., 2024).

As discussed earlier in Section 4.1.1, instead of using a reranker for chunk-level compression, it is also possible to prune irrelevant chunks by using similarity search between the question and chunk embeddings. We conduct the study using OpenAI’s text-embedding-3-small as the embedding model. As shown in Figure 5, the reranker outperforms the retriever model. However, the retriever model has the advantage of requiring less resources at inference time, since document embeddings are computed offline. In many settings, reranking is applied after an initial retrieval step to reduce the number of documents that need to be reranked.

For the token pruning methods in Section 4.1, the reranker selects $25\%$ more chunks than originally and then applied a token pruning rate of $20\%$ to achieve each compression ratio. Here, we perform a study where the reranker selects $2\times$ more chunks and an aggressive token pruning rate of $50\%$ is applied. As shown in  Figure 6, such aggressive token pruning leads to accuracy degradation. After observing the pruned context, we hypothesize that this is because aggressive token pruning leads to unstructured text that does not respect grammatical constructs, making it difficult for the downstream model to correctly reason over it.

The abstractive compression method presented in  Section 4.1.2 performs query-agnostic abstractive compression. This is largely beneficial for applications that need low-latency responses, as summaries are precomputed offline. However, it is also possible to perform query-aware abstractive compression, in which summaries are generated by conditioning on the question. Specifically, we use the reranker model to first select relevant chunks and then use a small language model to summarize the concatenation of selected chunks. We show the results in Table 2 with 16 selected chunks and include more results with 8 and 32 chunks in  Section B.4. We experiment with both Mistral 7B and Llama 8B (AI@Meta, 2024) as summarizers. After observing difficulties in prompting such models to produce summaries of specific lengths, we used prompting methods similar to RECOMP (Xu et al., 2023), which allows the Mistral model to freely choose the summarization length. In general, our experience with abstractive compression indicates that strong prompt engineering is necessary to achieve desired performance. The summarization prompts are shown in Section B.5. As shown in Table 2, query-aware abstractive compression demonstrates stronger performance than query-agnostic abstractive compression. For example, on NarrativeQA, MultiFieldQA, and HotpotQA, query-aware compression performs 3-6 points better than query-agnostic. These trends persist across both Mistral 7B and Llama 8B. It is possible that more detailed summarization prompt engineering can further improve performance. Therefore, query-aware abstractive compression may be a promising technique applications willing to handle the overhead of performing on-the-fly summarization.

To determine the impact of chunk size, we run a set of experiments after changing chunk size from 128 to 512 tokens. The results are shown in Figure 7 and Figure B.7. We notice that large chunk sizes do not perform well at large compression ratios when compared to smaller chunk sizes. We hypothesize that this is because there are very few chunks being provided to the model when the chunk size is large. As a result, the model does not have the ability to see text from varying regions of the initial context. In contrast, using smaller chunk sizes allows more chunks to be used, alleviating this issue. At smaller compression ratios, the chosen chunk size has lesser impact. Ultimately, chunk size should be carefully determined after examining an application’s data source as well as the desired compression ratio. Additionally, as we demonstrate in Section 4.3, certain applications may require application-specific chunking techniques.