Enhancing and Accelerating Large Language Models via Instruction-Aware Contextual Compression

Retrieval has been integrated into language models for years to enhance various aspects such as perplexity [4, 32], factual accuracy [24], downstream task accuracy [8, 12, 13, 18], and in-context learning capability [11]. Combined with a standalone retriever [16, 34, 20], retrieval-augmented LLM is a well-established approach for addressing question answering with long documents in an open-domain context. In previous studies, language models have been augmented with retrieval during inference [17, 41], fine-tuning [13, 18, 8], and pretraining [4, 13, 32]. There are also some methods that aim to integrate LLM and retriever into a single model, creating an end-to-end solution [15, 28].

Many approaches have sought to improve the handling of longer contexts in Large Language Models (LLMs) through modifications to their underlying architectures. Notably, the Longformer [3] employs a linear attention mechanism that scales with sequence length, allowing it to accommodate longer contexts effectively. CoLT5 [1] introduces conditional computation techniques that enable the model to focus more on crucial tokens in both feedforward and attention layers. However, it’s worth noting that many existing works have not yet adopted such architectural modifications, mainly due to the high computational cost associated with training LLMs. Another category of approaches addresses the context length limitation by employing context chunking strategies. The Parallel Context Windows [26] proposes a parallel context window method, which calculates attention within individual chunks and incorporates positional embeddings between these chunks.

Prompt engineering is a relatively emerging discipline focused on crafting and refining prompts to harness the power of language models (LMs) for diverse applications and research endeavors. Proficiency in prompt engineering aids in gaining a deeper insight into the capabilities and constraints of large language models (LLMs). Researchers employ prompt engineering to enhance the performance of LLMs across an array of common and intricate tasks, including question answering and arithmetic reasoning. Developers leverage prompt engineering to devise resilient and efficient prompting strategies that interact seamlessly with LLMs and other associated tools. Prompt engineering encompasses two key directions: text-to-text and text-to-image [33, 25] interactions. This area of research has witnessed significant manual efforts, exemplified by A Prompt Pattern Catalog [37], where a comprehensive collection of handcrafted prompt techniques is meticulously documented. On the other hand, there are automated approaches to prompt generation, such as the work on ”Automatic Prompt Engineer” by Zhou and ”LM-BFF” by Gao [42, 7].

Context compression can be considered a form of prompt engineering, although their emphases are slightly different. Several techniques aim to compress prompts effectively while maintaining context relevance. The Selective Context [19] approach leverages concepts from information theory, specifically self-information, to compress the context. Another approach, Learning to Compress Prompts with Gist Tokens [23], trains Gist models to compress prompt words into ”gist” tokens before inputting them into the LLMs. Additionally, LeanContext [2] extracts a dynamic number $k$ of key sentences from prompts and uses a reinforcement learning mechanism to determine the optimal value of $k$ for compression.

We introduce Instruction-Aware Contextual Compressor as a trainable module to implement the Instruction-Aware Contextual Compression method. Instruction-Aware Contextual Compressor adopts an encoder-decoder architecture, consisting of both an encoder and a decoder, as illustrated in Figure 3.

The document encoder is a standard multi-layer transformer encoder and utilized to extract features from the input documents.

The decoder is a standard multi-layer transformer decoder, which equipped with two distinct functionalities, which can be toggled by modifying the masking:

1. 1.

   The ranking decoder performs re-ranking based on the output features obtained from the document encoder. In this mode, interaction occurs between instruction features and document features within cross-attention layers. The decoder employs bidirectional self-attention layers without any masking.

2. 2.

   The generation decoder, on the other hand, replaces the bidirectional self-attention layers in the decoder with causal self-attention layers. It uses a [BOS] token to denote the start of a sequence and an end-of-sequence token to signify its conclusion.

The model consists of 8 encoder and 8 decoder layers, which affect its depth and ability to capture complex patterns. The primary model dimension is 512 and it uses a feed-forward dimension of 1024 for its inner layers. The model has 6 attention heads, allowing it to focus on different aspects of the input data. The model has a total of 0.18 billion trainable parameters, which is significantly smaller than the current mainstream large language models like Llama  [30, 31], which have 7 billion, 13 billion, and 70 billion parameters, respectively. We initialize Instruction-Aware Contextual Compressor with the pre-trained weights of umT5[6], which has been pre-trained on a multilingual corpus, enabling it to handle multilingual tasks effectively.

The model’s architecture, in conjunction with its training objectives, empowers it to capture the intricate interplay between instructions and documents, facilitating the extraction of the most pertinent information from the documents.

We jointly optimize two objectives during training, with one ranking-based objective and one generation based objective. Each instruction-document pair only requires one forward pass through the document encoder, and two forward passes through the decoder, where different functionalities are activated to compute the two losses as delineated below

Ranking Loss activates the ranking decoder. It aims to learn instruction-document representation that captures the fine-grained alignment between instruction and document. In ranking task, the model uses a ranking head (a linear layer) to predict a instruction-document matching score given their instruction feature and document feature. All positive samples are placed in the first position, and an additional 19 hard negative samples are retrieved by the retriever. Then, the model is trained as a 20-class classification task. This approach effectively boosts the scores of positive samples while suppressing the scores of negative samples. The specific formula for calculating the loss is as follows:

where $s_{i}$ represents the ranking score of $i$-th sample. $t_{i}$ is the $i$-th target. $c$ represents the number of class.

The Language Modeling Loss activates the generation decoder, with the goal of generating useful response based on the provided context and instruction. It is optimized using a cross-entropy loss that trains the model to maximize the likelihood of the text in an autoregressive manner. We incorporate a label smoothing factor of 0.1 when calculating the loss. In comparison to the ranking loss, the Language Modeling loss equips the model with the ability to generalize for following instructions. This empowers the model to gain a deeper understanding of the potential correct context location and to effectively model fine-grained correlations. the language modeling loss can be written as:

where $p(y_{t}|y_{<t})$ denotes the output probability for the correct token $y_{t}$ given the previous context $y_{<t}$. $N$ denotes the sequence length.

Instruction-Aware contextual compression by ranking is a fairly straightforward process. First, an appropriate text-splitting strategy is employed, which can involve splitting based on specific character or by length, effectively converting the document into multiple chunks. Next, the model’s ranking capability is applied to score chunks with instruction, followed by sorting them. Chunks are then retained based on a specified percentage. Throughout the compression process, care is taken to maintain the order of the chunks.

Compressing context by generation leverages the ability of Grad-CAM [27] to capture fine-grained relevance. The calculation process of Grad-CAM can be summarized as follows: Firstly, perform a forward pass to obtain the final classification probabilities. Calculate the gradients of the token with respect to the target class. Average the gradients for each token and extract an attention map within a specific cross-attention layer. Grad-CAM scores can be obtained by multiplying the attention map and the gradient weight vector. These Grad-CAM scores can be considered as the contribution of each token to the classification result.

There’s no need to pre-split the context; instead, the entire context is input into the model. Following that, k-step responses are generated based on the context and instruction. In the Instruction-Aware Contextual Compressor, attention maps and gradients are recorded for specific cross-attention layers, which are used to compute token-level Grad-CAM scores. These token-level Grad-CAM scores are then averaged to obtain chunk/sentence-level Grad-CAM scores. The chunks are subsequently sorted based on their Grad-CAM scores, and a specified percentage of the highest-scoring chunks is retained. Similarly, throughout the compression process, care is taken to maintain the order of the chunks.

Effectively ensemble these two methods of Instruction-Aware contextual compression can yield better compression results. The magnitude difference between the ranking score and Grad-CAM score makes it challenging to determine a suitable weighting parameter for fusion. Therefore, in the end, we opted for a non-parametric approach. Specifically, we individually rank the two types of information and then use the average of the two rankings to compress the context.

In this study, we utilized two types of datasets. The first are ranking datasets, designed to empower the model with robust re-ranking capabilities. The second are generation datasets, intended to equip the model with generative abilities. The ranking datasets comprise 15 million samples, while the generation datasets consist of 1.63 million samples.

During our experimentation, we conducted tests using Instruction-Aware Contextual Compressor on ChatGPT, which is based on the GPT-3.5-turbo-0613 architecture. ChatGPT represents an Instruct-tuned language model that has undergone further enhancement through Reinforcement Learning from Human Feedback (RLHF) and boasts an impressive 175 billion parameters. The foundational language model of ChatGPT appears to be code-davinci-0022, and previously, davinci, as outlined in [5]. Our objective was to compare the performance of ChatGPT with and without the application of Instruction-Aware Contextual Compressor to gain insights into its impact on the model’s efficiency and accuracy. The settings of ChatGPT are all set to their default values, except for the $top\_p$ parameter, which is set to 0.1. This adjustment is made to reduce the impact of randomness on the evaluation results.

We conduct a comparative analysis to assess the effectiveness of Instruction-Aware Contextual Compressor and analyze the associated trade-offs.

RAG Baseline: The table 1 displays several RAG baselines, including scenarios where RAG is not used, direct feeding of the correct context to the large language model, the scenario where only a retriever is used in pipeline, and the scenario where a retriever is used in combination with Instruction-Aware Contextual Compressor for re-ranking.

Compression Baseline: Our evaluation involves a comparison between Instruction-Aware Contextual Compressor and Selective Context [19], which employs a basic approach to filter out an equivalent amount of data based on self-information. Selective Context utilizes the GPT-2 model with 124 million parameters. Some readers may doubt whether the GPT-2 124M model is too small to be considered a sufficiently robust baseline for comparison. To this end, we used Baichuan-7B with 7 billion parameters to run the Selective Context, which has parameters 56 times larger than GPT-2. The results at a retention ratio of 0.5 are presented in the table below:

We found that Baichuan-7B did not perform better than GPT-2. This also indicates that the selective context method is not scalable. This therefore indicates that the baseline we used is sufficiently robust.

Retention Ratios: In our experiments, we explore various content retention ratios: 0.2, 0.35, 0.5, 0.65, and 0.8. These ratios determine the proportion of content to be retained. This exploration allows us to examine the trade-off between efficiency and performance as the amount of retained information varies.

Setting for Inference Measure: To measure the inference acceleration and memory savings brought about by context compression, we conducted practical measurements using an NVIDIA 4090 GPU with 24GB of VRAM. The LLM model used was Baichuan-7B [40], and the data format employed was bfloat16.

We initially compare the performance of Instruction-Aware Contextual Compressor with varying context retention ratios to the reranked original context, which utilizes the original context after reranking but no compression at all. All results are shown in table 3, and the ”diff” column represents the difference in performance compared to uncompressed text.

As shown in the table 3, at retention rate of 0.8, the performance loss is minimal, with Rouge-1 showing only a marginal decrease in the range of 0.003 to 0.008. This demonstrates a high level of consistency between answers provided in compressed contexts and those in original contexts. Surprisingly, the Rouge-2 and Rouge-L score with the generation method is even higher than the original text, which was unexpected. This indicates that our method successfully filtered unrelated or even noisy content, improving the LLM’s performance.

As the retention ratio decreases, the effectiveness of all methods declines, which is expected since there is less valuable information provided to the LLM. Overall, the generation method outperforms the ranking method, with a slower rate of performance decline from 0.8 to 0.35 compared to the ranking method. However, it was unexpected that at a retention ratio of 0.2, there was a sudden significant drop in performance, indicating a rapid loss of effectiveness for the generation method at low retention ratios.

In traditional machine learning, ensemble learning is widely regarded as a robust and effective method for improving model performance. Therefore, we propose using the ”average rank” to combine the generation method and the ranking method. This involves taking the average of the ranks assigned by the generation method and the ranking method, resulting in a new ranking score for a given text. Overall, the ”average rank” method outperforms both the generation and ranking methods. It shows significant improvement from 0.2 to 0.65 retention ratios, with fewer losses compared to the original context. At a retention ratio of 0.8, while it may not outperform the generation method, it still surpasses the ranking method.

In this section, we compare our method to Selective Context baseline, and the results are presented in Figure 4. Selective Context is a context compression method based on text self-information and represents the state-of-the-art as of the time of writing this paper. Comparing our method to Selective Context, which serves as a baseline, can effectively demonstrate the validity of our approach. As shown in Figure 4, our proposed method, Instruction-Aware Contextual Compressor, is even more effective compared to Selective Context. Both methods, whether purely ranking-based or generation-based, outperform Selective Context, and the lead becomes more significant as the retention ratio decreases. This indicates that our proposed method excels in selecting more informative content even when only a limited amount of information can be retained.

For context compression using generative information, intuitively, if the number of generation steps is too few, it might not have generated a complete response. Consequently, the effectiveness at this stage may be suboptimal. As the number of generation steps increases, the effectiveness of compression is expected to improve. To explore this, we conducted experiments with a fixed retention ratio of 0.5, testing a series of generation step values ranging from 4 to 64. The results, as shown in Figure 5, indeed demonstrate that for generation-based context compression, the performance gradually improves with an increase in the number of steps. However, after reaching 32 steps, it reaches a plateau, indicating a diminishing marginal return with further increases in the number of generation steps.

We also measured the impact of context compression on the Large Language Model (LLM). As demonstrated in the table 4, when the retention ratio is set to 0.5, the inference speed per token increases by a factor of 2.2, while memory usage decreases by 5.05%.

The ability to filter out irrelevant documents through re-ranking is also a crucial capability of Instruction-Aware Contextual Compression. Following a hierarchical design approach, a large number of initially retrieved documents are first re-ordered using Instruction-Aware Contextual Compressor to select the top-k documents, which are then further compressed. Therefore, we present the retrieval performance of Instruction-Aware Contextual Compressor on extensive datasets, measured by Recall@1, 5, and 10, as shown in 5.

Guangdong Laboratory of Artificial Intelligence and Digital Economy(SZ), Yutang, Shenzhen, 518000, Guangdong, China.

Haowen Hou, Fei Ma, Binwen Bai, Xinxin Zhu & Fei Yu

All authors have contributed to the ideas and design of this research. Haowen Hou was in charge of developing and training the model, and also wrote this manuscript. Fei Ma carried out the experiments and the analysis of the outcomes. Binwen Bai and Xinxin Zhu managed the data collection, arrangement, cleansing, and its ultimate transformation into the training format. Fei Yu offered feedback on the preliminary drafts of the manuscript. All authors have reviewed and approved the final manuscript.

Correspondence to Haowen Hou.

The authors confirm that there are no financial conflicts or personal connections that might be perceived as affecting the findings presented in this study.

There are no concerns regarding the ethical use of the data or the informed consent process. All necessary ethical guidelines and protocols have been followed, and informed consent was obtained from all participants involved in the study.