Extending Context Window of Large Language Models via Semantic Compression

Introduction

The recent successful release of large language models (LLMs) such as ChatGPT (Radford et al., 2019) and LLaMA (Touvron et al., 2023) has sparked significant research efforts from both industry and academia. These LLMs have demonstrated the ability to engage in fluent and coherent conversations with human users, and have shown exceptional performance across various tasks, including document summarization, question-answering, dialogue bots, and code generation copilots.

One critical issue faced by state-of-the-art (SoTA) LLMs is the restriction on the length of text that can be inputted into the model at once. When the input context exceeds the limit of the context window, the performance of these models rapidly declines. This limitation poses a challenge when it comes to handling long texts such as scientific papers, novels, and legal contracts with current LLMs. As a result, there has been a growing interest in finding ways to extend the input length without significantly compromising the model’s performance.

The limitation on the context window primarily stems from the quadratic computation of the self-attention mechanism in the transformer. Handling lengthy texts significantly increases the computational costs in terms of memory and time. Typically, models are trained on short contexts, and the maximum sequence length (i.e., the context window) is determined. If the models are compelled to generate contexts that exceed the context window, they tend to compromise the quality of the output due to the lack of position encoding information during the training process. Furthermore, generating long sequences imposes substantial memory requirements on the computational device. This accumulation of memory requirements and the lack of effective position encoding can result in length generalization failure (Anil et al., 2022), where the models struggle to generate meaningful and coherent text beyond a certain context window size.

Some approaches have been developed to address the aforementioned challenges. One approach is to devise architectures with nearly linear complexity, which enables efficient scaling to handle very long sequences. However, training a large model from scratch incurs substantial cost. Another strategy involves employing interpolation and fine-tuning techniques to adapt the position encoding to unseen sequence lengths. While this method has the potential to compromise the overall performance of LLMs, it still demands significant time and GPU resources for fine-tuning and inference on long sequences. Therefore, it is more efficient and resource-friendly to design methods that do not necessitate altering the parameters of the pre-trained model.

While most previous algorithms relied on modifying the pre-trained model, we instead exploit the statistical properties of input natural language. One empirical phenomenon, known as Zipf’s law (Zipf, 2016), observes that a small set of the most frequent word tokens in a large corpus of natural language account for almost all occurrences. This pattern arises from the tendency of language users to minimize effort in their daily conversations. Consequently, by utilizing an expanded vocabulary, sentences can be significantly shortened while preserving the same semantic meaning. Moreover, it is common for language users to include redundant words during communication (Strunk Jr, 2007). These language habits are prevalent among users, and we propose to include a semantic compression module to mitigate the redundancy associated with these habits.

Our proposed semantic compression method, reminiscent of lossy source coding in information theory, extends the context window by equivalently shortening the long text while preserving the semantic meaning. This procedure is conducted before inputting the tokens into the pre-trained LLMs. As illustrated in Fig. 1, the input undergoes compression before being transmitted to the LLM for various potential tasks. The semantic compression method can be customized and optimized for downstream tasks, taking into consideration practical constraints such as time and memory resources. The implementation of the semantic compression module is straightforward and can easily be incorporated into other interpolation-based context window extension methods and black box APIs. It demonstrates enhanced performance compared to SoTA interpolation-based methods on a range of tasks, including single-document question answering, multi-document question answering, summarization, few-shot learning, and information retrieval, using real-world datasets while incurring no extra parameter updates or memory consumption. Empirically, the proposed method is computational efficient and achieves 6-8 times context window extension.

We introduce a context window extension framework for LLMs that utilizes semantic compression. This framework serves as a plug-and-play tool to mitigate redundancy in input texts by efficiently performing topic modeling.

We construct a graph representation of the input to identify distinct sections of the text that pertain to different topics. The result is the segmentation of long texts into separate chunks, each focusing on a specific topic. We then conquer each chunk independently, resulting in a concise version of the original texts. This compression technique helps to condense the information while preserving the key ideas and context.

We demonstrate the applicability of our proposed semantic compression method through extensive experiments. The results highlight the advantages of our method in several key applications, including single-document question answering, multi-document question answering, summarization, few-shot learning, and information retrieval.

Related work

With the advancement of SoTA LLMs, significant progress has been made in extending the context window lengths.

The mainstream line of research aims to adapt existing language models trained on short texts to accommodate longer ones during inference (Anil et al., 2022). The key idea is to modify the positional embedding, which has only been trained on short texts. Several studies are based on the Rotary Position Embeddings (RoPE) of LLaMA and methods of adjusting it to the longer sequences. Chen et al. (2023a) develops the Position Interpolation (PI) method to linearly scale the input positional indices. Peng et al. (2023) presents YaRN, an efficient extrapolate mechanism inspired by the neural tangent kernel, to extend the context window to $64$k and $128$k.

2 Efficient Attention Operations

Due to the self-attention mechanism, the inference cost of LLMs grows quadratically with the sequence length. Many methods have been proposed to decrease the complexity. Dai et al. (2019) present Transformer-XL which utilize segment-level recurrence agency and a novel positional encoding scheme. Beltagy et al. (2020) introduce Longformer with a sparse attention mechanism that scales linearly with sequence length. Bo (2021) provides a faster transformer, RWKV, which combines the strength of RNN and has linear complexity during inference. Dao et al. (2022) propose FlashAttention, a chunking strategy for the input, and utilize recomputation to avoid the quadratic complexity of attention computation. While these methods have the potential to handle longer input sequences (Ding et al., 2023), training new models can be costly. Moreover, these methods are not effective when dealing with out-of-distribution content lengths.

The introduction of new positional embeddings requires fine-tuning on long sequences to adapt to the increased length, which can be computationally expensive. To address this, LongLoRA is introduced by Chen et al. (2023b), offering an efficient fine-tuning method with limited computational costs. More details on several other chunking strategies are provided in the survey by Huang et al. (2023).

3 Prompting

There are ongoing efforts to extend the context window through smart prompting designs. Wingate et al. (2022) utilize soft prompts to encode more information using fewer tokens. Chevalier et al. (2023) present AutoCompressor, which utilizes soft prompts to compress the input sequence and then extends the original length of the base model. Both Zhou et al. (2023) and Wang et al. (2023) recurrently apply LLMs to summarize the input texts to maintain long short-term memory for specific purposes such as story writing and dialogue generation, respectively.

Methodology

We propose our semantic compression method for extending the context window. The core idea is to compress the input into shorter texts without losing the key information and important details. This enables us to effectively include more content within the fixed input length constraint of the LLM. Fig. 2 provides an overview of our method, which leverages pre-trained summarization models commonly used in Natural Language Processing (NLP).

Existing summarization methods also have limitations regarding the length of the input. Here, we propose a divide-and-conquer based approach that takes into account the structure of the text. By identifying the topic structure of lengthy texts and dividing them into blocks that exhibit a certain level of mutual independence, the content within each block can be compressed efficiently due to their statistical correlation. Each block is then processed in parallel using pre-trained models, and the results are combined to create a condensed textual input that can be processed by the LLM. This approach aims to provide a more efficient and effective way of summarizing long texts by leveraging both the structure and content of the original text.

Real-world textual content, such as speech and book, frequently displays hierarchical structures, wherein each section is structured around a particular topic, and different sections differ in topic in a sequential manner. This hierarchical structure, based on topics, bears resemblance to cliques in graphs. To identify this structure within long texts, we utilize weighted graphs to represent them and employ clustering methods to detect cliques in these graphs. The cliques can then be utilized to represent the topic-based content of the text, allowing us to obtain chunks based on the semantic relevance of the topics.

We begin by sequentially constructing sentence-level blocks within given lengths and representing them as nodes in our graph. In this step, we parse the text into different sentences or sub-sentences based on punctuation marks. Next, we sequentially fill the sentence-level blocks until they exceed the desired length before proceeding to the next blocks. Once we have obtained the sentence-level blocks, we connect the graph representation of long text $\mathcal{G}$ based on a pre-trained sentence embedding model (e.g., MiniLM (Wang et al., 2020)), where the weight $\mathcal{G}[i][j]$ represents the semantic similarity between the $i$-th and $j$-th sentence-level blocks. Typically, this similarity is computed using cosine similarity, which measures the cosine of the angle between two embeddings. If the similarity between two blocks is higher, it indicates that they are closer in topics.

2 Topic-Based Chunking

We then apply clustering algorithms on the graph to identify the underlying topic structure. Within each cluster, we group the sentence-level blocks sequentially to obtain the topic-based chunks, which can then be handled simultaneously by the pre-trained model chosen according to the downstream task. The number of clusters can be adjusted to regulate the length of the text following semantic compression. If these semantic chunks still surpass the predetermined length, the identical procedure is repeated to acquire sub-level topic structures.

The obtained topic structures are tree-like, which can be flattened in accordance with the order of the original content. As per the model, each chunk is semantically centered around a specific topic, and these topics are mutually exclusive. Consequently, these chunks can be compressed independently by utilizing a pre-trained summarization model. Choosing from different pre-trained summarization models allows a trade-off between efficiency and effectiveness. Consequently, we can opt to selectively substitute the original chunks with the output of these pre-trained models to ensure the preservation of the underlying topic structure. The semantic compressed text can be forwarded to the LLM directly or in combination with other extension schemes to further enhance the overall outcome.

Experiments

We demonstrate that the proposed method of semantic compression can effectively extend the context window by up to 7-8 times without modifying the parameters of the pre-trained models. Furthermore, the semantic compression module can be seamlessly integrated with existing methods, allowing for further extension of the context window. This versatility enables our approach to be adapted and combined with other techniques, enhancing the overall performance and flexibility. To evaluate the performance of our method, we conduct experiments on several language tasks that require understanding of long contexts. These tasks include passkey retrieval, single-document question answering, multi-document question answering, summarization, and few-shot learning. In each task, the model is provided with a sequence of context $C$ (typically lengthy texts) and a sequence of text $Q$ (e.g., a prompt), and it is expected to generate the output answer $A$. Additionally, we also investigate the perplexity metric (Peng et al., 2023), which measures the model’s ability to predict the text and serves as an indicator of the fluency of the generated output. This analysis allows us to assess not only the effectiveness but also the quality of the generated output.

We begin by evaluating the proposed semantic compression method on various standard benchmark tasks, utilizing the pre-trained 7B LLaMA model (Touvron et al., 2023). The original context window size of this model is $4096$. The tasks and datasets employed in our evaluation are sourced from the SCROLLS benchmark (Shaham et al., 2022) and LongBench (Bai et al., 2023). These datasets provide comprehensive and diverse contexts for our analysis.

Retrieval has been an important application of LLMs. We evaluate the proposed method using a synthetic task for passkey retrieval introduced by Mohtashami & Jaggi (2023), where prompts are synthesized to conceal a generated passkey within a randomly chosen section of a long document. The passkey retrieval task assesses the model’s capacity to extract important information from any position within lengthy contexts. An illustration of the task is shown in Fig. 3. The synthetic long text incorporates the passkey digits, and the task for the LLM is to retrieve these digits from the input text. Further specifics can be found in Appendix A.

General NLP Tasks

LongBench (Bai et al., 2023) is a multi-task benchmark designed for long text scenarios, consisting of six distinct tasks. In this study, we focus on the three English tasks from the set of four natural language tasks, namely single-document question answering, multi-document question answering, summarization, and few-shot learning. Each of the selected datasets contains 200 instances. Further information can be found in Appendix A.

Fluency

We evaluate the fluency of our semantic compression method using the perplexity score, which is defined as the exponential of the average negative log-likelihood of the probabilistic model $P$ on the distribution $D,$ i.e.,

A smaller perplexity score indicates more fluent sequences that are consistent with the model.

2 Baselines

We choose SoTA solutions from each mainstream approach as our baselines.

To accommodate long context within a fixed-size context window, chunking is a straightforward yet efficient approach. In NLP related applications, large pieces of text are usually broken down into smaller segments for targeted applications. When the input length exceeds the context window, the fixed-size chunking method (Bai et al., 2023) truncates the input sequence from the middle. This is because the most significant information typically resides at the beginning and end of the sequence.

Interpolation-based method

YaRN (Peng et al., 2023) is a computationally efficient method for interpolating position encoding, which dynamically adjusts the Relative Positional Encoding (RoPE) over dimensions and scales the attention. YaRN offers multiple length-extended models for different versions of Llama2, with the models being trained on a total of 64 GPUs from 8 $\times$ A100 machines. In order to ensure a fair comparison, we choose the model based on Llama2 7B, adjusted from 4k to 64k, as our baseline.

Fine-tuning approach

LongLoRA (Chen et al., 2023b) is an efficient approach for fine-tuning that combines LoRA and shifts sparse attention to reduce computational costs. LongLoRA applies this technique to Llama2 models of different sizes, ranging from Llama2 7B, Llama2 13B, to Llama2 70B, with token lengths extended from 4k to 32k on a single 8 $\times\text{A100}$ device. In order to ensure a fair and unbiased comparison, we choose the Llama2 7B model with context extension achieved through improved LoRA fine-tuning as our baseline.

Results

We report the main results along with a comprehensive analysis.

We utilize the Llama2 model as our baseline to evaluate the fluency of generated texts by calculating the perplexity (PPL) score. Samples from the GovReport dataset are selected at varying lengths, and the reference texts are compared to the generated texts during the computation. In cases where the length of the input text exceeds the context window of Llama2, our semantic compression module shortens the input, thereby allowing the model to continue generating new content fluently. The resulting scores are depicted in Fig. 5. The plots indicate that the perplexity of Llama2 initially decreases, but once it surpasses the window length, it rapidly increases. However, when our semantic compression method is employed, the PPL remains consistently low. This suggests that our approach successfully extends the context window up to three times without compromising the generation quality of the language model.

Passkey Retrieval

We present the results of the passkey retrieval task in Fig. 5. When employing Llama2 for passkey retrieval, we observe a rapid drop in accuracy to zero once the input length surpasses the window size of $4096$. However, by utilizing our method, the retrieval accuracy of the Llama2 model remains above 90% even for inputs with lengths of up to 30,000. This indicates that the semantic compression method extends the context window size of the language model by approximately 7-8 times. Furthermore, we combine our method with the SoTA interpolation-based method, YaRN, to further expand the context window size to up to 60,000, while consistently maintaining an accuracy above 90%.

General NLP Tasks

We present our results on various general NLP tasks in Table 1, including single-document question answering, multi-document question answering, summarization, and few-shot learning. When the token length is less than 4k, there is no need to compress the context, and our method performs at the same level as the original Llama2 model. However, both the interpolation-based method YaRN and the fine-tuning approach LongLora negatively impact the performance of the Llama2 model across almost all tasks. In the 4k-8k range, our method outperforms others in 8 out of 11 tasks. It is worth noting that our model performs slightly worse in the few-shot learning task. This can be attributed to the fact that few-shot learning necessitates more detailed information, whereas our compression scheme maintains information within a fixed window. Moving on to the 8k-16k range, our method achieves the best results in 9 out of 12 tasks, exhibiting similar performance to the 4k-8k range. In the 16k-32k range, our method outperforms others in 6 out of 11 tasks. In the 32k+ range, other methods fail due to out-of-memory issues, while our method still maintains 70% of the performance achieved in the 4k range.

Conclusion

In this work, we propose a novel approach to addressing the limitation of input length in large language models using semantic compression. By leveraging the statistical properties of natural language and exploiting redundancy in communication, we are able to significantly shorten texts while preserving their semantic meaning. This allows for a 6-8 time extension of the context window without the need for modifying the parameters of the pre-trained model or incurring additional computational costs. Furthermore, the implementation of our semantic compression module is straightforward and can be easily integrated into other interpolation-based methods and black box APIs. This provides flexibility and adaptability to different downstream tasks, considering practical constraints such as time and memory resources. We believe our work can lead to simpler context window extension method to be used in practice, thereby reducing the cost of large language models.

References

Appendix A Datasets

NarrativeQA (Kočiskỳ et al., 2018) is a standard question-answering dataset that includes books from Project Gutenberg3 and movie screenplays from a list of websites. Question-answer pairs were provided by annotators, so that each of the 1,567 books and scripts has about 30 questions and answers, and two reference answers are given for each question.

Qasper (Dasigi et al., 2021) is a question-answering dataset of NLP publications containing abstractive, extractive, and yes/no questions.

MultiFieldQA-en (Bai et al., 2023) is a dataset created from multiple sources including legal documents, government reports, encyclopedias, and academic publications. Doctoral students were requested to annotate each article’s queries and responses.

Multi-Doc QA

HotpotQA (Yang et al., 2018) includes many 2-hop questions written by native speakers based on two related paragraphs.

2WikiMultihopQA (Ho et al., 2020) involves up-to 5-hop questions systematacially constructed by manual templates. Answering these questions requires reasoning paths and can not be solved by local content.

MuSiQue(Trivedi et al., 2022) consists of up to 4-hop questions and removes shortcuts and naturalness questions. Each question contains 2-4 supplement paragraphs which present the reasoning path and related paragraphs.

Summarization

GovReport (Huang et al., 2021) collects detailed reports containing human-written summaries from the U.S. Government Accountability Office and Congressional Research Service. These reports span a wide variety of national policy issues.

QMSum (Zhong et al., 2021) contains annotated meeting-summary pairs across many domains including including product, academic, and committee meetings.

MultiNews(Fabbri et al., 2019) is a multi-document summarization dataset. (Bai et al., 2023) cluster 2-10 news articles discussing the same event or topic, each paired with a human-written summary and form a new long text summarization task.

Few-Shot Learning

To construct few-shot learning with long text, (Bai et al., 2023) select a range of training examples in the following datasets to concatenate the context in LongBench.

TREC (Li & Roth, 2002) is a classification dataset with fine-grained class label.

TriviaQA (Zhong et al., 2021) is a classification dataset and involves messenger-like conversations with human-written summaries.

SAMSum (Fabbri et al., 2019) reading comprehension dataset and consists of question-answer pairs annotated with evidence passages.

Passkey

The randomly generated prompts of the passkey retrieval task is in the format of Fig. 6.

Appendix B Implementation Details

In this section, we provide details of our algorithm implementation. Our algorithm utilizes several mature open-source models. For graph representation, we make use of the sentence similarity models all-MiniLM-L6-v2 provided by the Sentence Transformer platform, which can be found at the following link: https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2. For semantic compression, we employ the pre-trained model distilbart-cnn-12-6Available at: https://huggingface.co/sshleifer/distilbart-cnn-12-6. In most of our experiments, we utilize Llama2-7B-chat-4k as the base large language model (Touvron et al., 2023). The experiments were conducted on a single A40 GPU with 48GB memory.

Appendix C Complexity

Given a context with length $L$, the origin complexity is $\mathcal{O}(L^{2})$. Considering the length limitations of the compression module, we assume it has a minimum input length $\gamma_{1}$ and a maximum input length $\gamma_{2}$. We denote the compression ratio as $\alpha$. Our method utilizes a divide-and-conquer strategy, dividing the long text into chunks where the total length is represented as $L=l_{1}+\cdots+l_{k}$, and each chunk’s length, $l_{i}$, satisfies the condition $\gamma_{1}\leq\l_{i}\leq\gamma_{2}$. By $k\gamma_{1}\leq L$, we can bound the complexity of the compression module

The complexity of inferring the compressed context is

Thus the main complexity of our algorithms can be bounded by $\frac{\gamma_{2}^{2}}{\gamma_{1}}L+\alpha^{2}L^{2}$.

The result suggests that our algorithm can reduce the computational complexity by a factor of the square of the compression ratio during the inference stage. The compression module exhibits linear growth and can be processed in parallel.