Long-Context Language Modeling with Parallel Context Encoding (2024)

Abstract: Extending large language models (LLMs) to process longer inputs is crucial for a wide range of applications. However, the substantial computational cost of transformers and limited generalization of positional encoding restrict the size of their context window. We introduce Context Expansion with Parallel Encoding (CEPE), a framework that can be applied to any existing decoder-only LLMs to extend their context window. CEPE employs a small encoder to process long inputs chunk by chunk, enabling the frozen decoder to utilize additional contexts via cross-attention. CEPE is efficient, generalizable, and versatile: trained with 8K-token documents, it extends the context window of LLAMA-2 to 128K tokens, offering 10x the throughput with only 1/6 of the memory. CEPE yields strong performance on language modeling and in-context learning. CEPE also excels in retrieval-augmented applications, while existing long-context models degenerate with retrieved contexts. We further introduce a CEPE variant that can extend the context window of instruction-tuned models using only unlabeled data, and showcase its effectiveness on LLAMA-2-CHAT, leading to a strong instruction-following model that can leverage very long contexts on downstream tasks.

## introduction enabling long and extensible context is crucial for large language models (llms) to effectively perform complex tasks, such as summarizing a book or answering questions with hundreds of retrieved web pages. however, several challenges limit the ability of llms to leverage long context: (1) llms and popular positional encodings @xcite do not generalize to sequence lengths longer than the lengths seen during training @xcite , streamingllm @xcite , and replug @xcite . cepe trained on 8k tokens can generalize to 128k tokens with minimal computational and memory costs. @xcite , even after additional fine-tuning @xcite @xcite . cepe offers several benefits: (1) length generalization: cepe is not limited by positional encoding constraints as the long context is encoded in chunks, each with its own positional encoding; we apply cepe to llama-2 @xcite and train it on a filtered version of red-pajama (together, 2023b) for 20b tokens-only 1% of the pre-training budget of llama-2. we first show that cepe-llama-2, trained with 8k input length, continues to improve perplexity on longer input up to 128k tokens. then, we apply cepe to a retrieval-augmented setting, as our larger context window allows incorporating more retrieved documents. compared to existing methods, cepe achieves better performance on both retrieval-augmented language modeling and open-domain question answering. additionally, we demonstrate that cepe can effectively leverage more demonstrations for in-context learning @xcite . all the above is achieved with a much lower memory and computational cost than most previous solutions. finally, we propose cepe-distilled (ceped), which extends the context window of instructiontuned models, using only unlabeled data. ceped distills the behavior of the original instructiontuned model to the new architecture through an auxiliary kl divergence loss, which eliminates the need to curate expensive long-context instructionfollowing data @xcite . we apply ceped to llama-2-chat and show that while preserving their instruction understanding ability, ceped-llama-2-chat can incorporate more context and improve performance on long-text understanding tasks @xcite . to conclude, cepe is a lightweight framework that can extend context windows of any base or instruction-tuned lms. we hope cepe can empower future llm research with cheap and effective long-context abilities. ## long-context language modeling we start by evaluating cepe on long-context language modeling benchmarks to assess basic lm abilities, and compare with existing approaches in terms of perplexity, memory, and throughput. datasets. we evaluate on arxiv and books from our redpajama test split, as well as three longcontext datasets: pg19 @xcite , proof-pile @xcite , and codeparrot @xcite . we filter all documents to have at least 32, 768 tokens, and sample 5, 000 sequences for each dataset. we calculate the perplexity on the last 256 tokens of each sequence. following @xcite , for the experiments with 128k tokens, we filter documents to have at least 131, 072 tokens and evaluate on only 10 sequences. 4 models. we focus on 7b parameters model. our baseline includes llama-2 and its long-sequence fine-tuned versions: llama-2-32k (together, 2023a), yarn-64k, and yarn-128k @xcite . llama-2-32k was trained on rp with upsampled arxiv and books domains, and a mixture of other data , and yarn was trained on pg19. while yarn-64k and yarn-128k were fine-tuned for only 1.7b and 3.4b tokens, respectively, fine-tuning all parameters requires much more memory relative to cepe. we also evaluate on training-free long-context methods: streamingllm @xcite and replug @xcite with llama-2-7b. although replug was originally evaluated in a retrieval-augmented setting, we found that it also works well in long-context modeling, when view-4 this is due to computational cost and scarcity of long documents. all 10 sequences are from different documents. for cepe, we put 2k tokens in the decoder when @xmath0, 096, and 4k tokens in the decoder in other settings. additional tokens are split into chunks of 256 tokens and fed into the encoder. results. we show the results in we also outperform replug across all domains while achieving much higher throughput-replug encodes additional context in chunks but each chunk requires a forward pass of the main input, incurring slow speed, due to which we omit replug at longer length. streamingllm maintains a low memory usage and a reasonable throughput, but the perplexity does not always decrease as the sequence length increases, likely due to the model only using a limited number of cached key-value pairs. compared to streamingllm, cepe can leverage all input tokens and achieve better perplexity with better throughput. ## retrieval-augmented applications retrieval-augmented settings naturally benefit from long-context lms, as models can leverage the ad-ditional context to include more retrieval results. thus, we test if cepe trained on long-context data can transfer to retrieval-augmented settings. ## in-context learning in-context learning (icl; brown et al., 2020) is one of the most important emerging qualities of llms. in this experiment, we examine whether cepe can effectively utilize demonstrations from the encoder context to improve performance. specifically, we use a range of classification tasks that contains a 5 snapshot from 2018-12-20, and each passage is 100 words @xcite . models. our baseline is llama-2 with 2 demonstrations. for cepe, we add additional demonstrations in the encoder. we also compare with an "oracle", where the llama-2 decoder takes 40 demonstrations. note that the oracle is significantly more expensive. more details are in §c.2. results. the results are shown in ## instruction-tuned models for long text understanding in this section, we show that applying ceped on llama-2-chat produces instruction-following models that can leverage long inputs. datasets. zeroscrolls @xcite is a collection of zero-shot long-context understanding tasks that require instruction-following abilities. specifically, we test on narrativeqa, qasper, quality, govreport, summscreenfd, and qm-sum, which all have large validation sets made available. we follow the formats and instructions of @xcite for each dataset, except the long text is placed before the instructions. ## ablation studies we conduct comprehensive ablations to show the effectiveness of our training data mixture, pretraining and fine-tuning the encoder, and the warmup training stage. we also ablate to verify the effectiveness of the kl divergence loss in ceped. lastly, we evaluate on needle in the haystack @xcite , and analyze the results in §d. ## conclusion we propose cepe to extend the context window of existing language models. the key idea is to leverage a small encoder and cross-attention modules to process long inputs and achieve low memory and computational complexity. compared to existing methods, cepe extrapolates to input lengths well beyond the training length, while remaining efficient and effective. consequently, cepe augments pre-trained models to be performant on both long-context and retrieval-augmented applications. we also show that ceped can be applied to instruction-tuned models with additional contexts using an auxiliary loss with only unlabeled data. we believe that there is still room for improvement through better data to train flexible and robust models. we hope our work can be a useful and accessible tool for the community to study long-context models in diverse applications.