LL o CO : Learning Long Contexts Offline (2024)

Abstract: Processing long contexts remains a challenge for large language models (LLMs) due to the quadratic computational and memory overhead of the self-attention mechanism and the substantial KV cache sizes during generation. We propose LLoCO, a novel approach to address this problem by learning contexts offline through context compression and in-domain parameter-efficient finetuning with LoRA. Our method enables an LLM to create a concise representation of the original context and efficiently retrieve relevant information to answer questions accurately. Our approach extends the effective context window of a 4k token LLaMA2-7B model to handle up to 128k tokens. We evaluate our approach on several long-context question-answering datasets, demonstrating that LLoCO significantly outperforms in-context learning while using 30 \times fewer tokens during inference. LLoCO achieves up to 7.62 \times speed-up during inference and 11.52 \times higher throughput during finetuning, substantially reduces the cost of long document question answering. This makes it a promising solution for efficient long context processing.

## introduction large language models (llms) have demonstrated remarkable performance across a wide range of tasks @xcite . many of these tasks require llms to understand and reason about long contexts. for instance, document question answering is one of the most common applications of llms, where the model is presented with a document as context and asked to respond to related questions or summarize the text. these documents may range from lengthy articles to entire books, potentially surpassing the context window limit of llms. consequently, there is a growing trend in both academia and industry to enhance llms' capability to process longer contexts effectively @xcite @xcite . this need has driven innovation among llm providers like openai and anthropic to develop models that can handle increasingly lengthy texts consisting of several thousands of tokens. despite the impressive progress made by the llm model providers, scaling these models to adeptly manage extended contexts remains a formidable challenge, both technically and financially. due to the self-attention mechanism, transformer-based llms incur a quadratic computational and memory overhead as sequence length increases. many long-context tasks require reusing the same context many times, which incurs extra latency overhead and substantial costs, as most commercial llms operate on a pricing model that is directly tied to the number of tokens processed. for example, a single inference run with a document that has 100k tokens would take 1.5 usd on claude 3 opus foot_0 and 1 usd on gpt-4-turbo foot_1 . orthogonal to the exciting efforts of expanding the context window limit, our study introduces an innovative strategy to tackle the long context challenge. we propose a method where context information is condensed offline through finetuning, enabling the model to provide accurate responses during inference with streamlined contextual representations. to illustrate our idea, consider an analogy: envision an llm as a student preparing for an exam, where we, the researchers, are the examiners providing study materials and questions. traditional in-context learning with full context or retrieval-augmented generation (rag) resembles an open-book exam, where the llm has access to all materials while answering questions. in con-figure 1 : the architecture of regular llm (left) vs @xcite . in regular llms, long contexts are appended directly to the prompt. in contrast, lloco first processes these contexts through a context encoder. the resulting summary token embeddings are then prepended to the llm's prompt, which are significantly shorter. lloco instruction finetunes on embeddings of targeted document groups using a lora module. this aligns the llm's embedding space with the summary embeddings while keeping both the llm and context encoder unchanged. regarding the first question, we find the formulation closely resembles the existing research direction on context compression. prior works @xcite @xcite have proposed compressors that aim to distill the essence of the original texts into compact representations that are aligned with llms. our work primarily focuses on the second question -we've observed that despite the progress in context compression, llms often struggle to accurately read such "cheat sheets" and tend to hallucinate when applying them to answer queries. to address this issue, we employ in-domain parameter-efficient finetuning directly on the compacted context (cheat sheet) without altering its content, which significantly improves the llm's ability to accurately extract and utilize information from these compressed representations. in this way, we can steer the llm to process long context more efficiently and accurately. this approach allows us to extend the effective context window of a 4k llama2-7b model to handle up to 128k tokens. moreover, we achieve state-of-the-art results that match or even surpass the performance of a llama2-7b-32k model with full context on long context benchmarks, while using 30× fewer tokens. this insight has led us to introduce lloco, a pipeline that learns contexts offline through the combination of context compression and parameterefficient finetuning. it consists of three stages: preprocessing, finetuning, and serving. first, we preprocess the documents into "cheat sheets". then, we employ low-rank adaptation (lora) @xcite to perform parameter-efficient finetuning on these "cheat sheets" in groups. for serving system design, we use a standard rag retriever to retrieve the compressed document as well as the most relevant lora module, and apply them to the llm for inference. the contributions of our work could be summarized as follows: long-context llms recently, there have been efforts to increase the llms' context window sizes efficiently with continual pretraining or finetuning. one line of work focuses on scaling the rotary position embeddings (rope) @xcite , achieving longer contexts up to 128k @xcite @xcite . mistral @xcite proposes sliding window attention that only attends to part of tokens from the previous layer, reducing compute and enabling pretraining with long-context to 30k. nevertheless, as the auto-regressive generation in llms is largely memory-bound @xcite , storing the kv cache of longer contexts slows down the inference and requires large gpu vrams. context compression a closely related topic is context compression, which aims to train a general compressor that can compress any input prompts. gist @xcite , autocompressor @xcite , and icae @xcite finetune llms in a "soft prompt-tuning" manner, either applying specific regularization in attention masks or utilizing dedicated "memory tokens" to compress contexts into embeddings with significant shorter lengths. llmlingua family @xcite proposes a questionaware framework to compress prompts in natural languages for black-box llm apis. another line of work employs kv cache compression by eviction @xcite , which only keeps informative keys and values for generation during inference, or quantization @xcite . however, all those approaches aim to compress any inputs, which usually incurs rapid performance drops as the compression ratio exceeds a certain limit (e.g. 2-4×), especially for out-of-distribution texts, hindering their practical applications. in this work, we mainly focus on extreme compression of in-distribution documents of interest, going up to 30× compression rates. retrieval-augmented generation retrieval enhances the performance of llms in knowledgeintensive tasks such as question answering with long documents or in open-domain. rag techniques are effective both in finetuned scenarios and when applied to off-the-shelf llms @xcite @xcite , and have seen many advancements recently, emphasizing improved external knowledge integration, broader query support, and enhanced content clarity @xcite @xcite . despite the improved performance from retrieval, challenges still remain in terms of runtime efficiency @xcite and effective filtering of irrelevant parts @xcite . our proposed method opens up new opportunities for capturing important information while ensuring efficient generation. parameter-efficient finetuning parameterefficient finetuning methods @xcite @xcite freeze the majority of the model weights and only update a small subset to finetune large models efficiently. lora @xcite is one of the most widely adopted techniques, and recent advancements @xcite have focused on enhancing the system performance for deploying lora adaptors. in particular, @xcite has achieved the deployment of thousands of lora instances on a single nvidia a100 gpu. lloco, which utilizes lora for finetuning, could benefit from those improvements to deploy lora adaptors efficiently in long context scenarios. ## experiments in the experiment section, we aim to investigate the following aspects of lloco: (1) its effectiveness in comprehending compressed long contexts, (2) the extent to which summary embeddings can preserve essential information, and (3) the efficiency improvements on the associated system costs. datasets we select four datasets dedicated to question-answering tasks-quality (pang et al., 2021), qasper @xcite , narrativeqa @xcite , and hot-potqa @xcite )-along with one for summarization, qmsum @xcite . all datasets have long documents as contexts. for all the datasets, we use their validation set for evaluation. we follow the official metrics for each dataset. for quality, we report the exact match (em) score. for qmsum, we report the geometric mean of rouge scores. for the remaining datasets (qasper, narrativeqa, and hotpotqa), we report the f1 scores. more details on these datasets can be found in appendix b. ## conclusion we propose lloco, a new paradigm that addresses the long-context challenge by preprocessing the context before inference through parameterefficient finetuning. our approach extends the ef-fective context window of a llama2-7b model with a 4k tokens context size to handle up to 128k tokens. evaluations on extensive long-context benchmarks show that lloco significantly outperforms in-context learning while using 30× fewer tokens during inference, making it a promising solution for efficient long-context processing.