Parallel Context Windows for Large Language Models (2023)

Abstract: When applied to processing long text, Large Language Models (LLMs) are limited by their context window. Existing efforts to address this limitation involve training specialized architectures, and cannot be easily applied to off- the-shelf LLMs. We present Parallel Context Windows (PCW), a method that alleviates the context window restriction for any off-the-shelf LLM without further training. The key to the approach is to carve a long context into chunks (“windows”), restrict the attention mechanism to apply only within each window, and re-use the positional embeddings across the windows. Our main results test the PCW approach on in-context learning with models that range in size between 750 million and 178 billion parameters, and show substantial improvements for tasks with diverse input and output spaces. We show additional benefits in other settings where long context windows may be beneficial: multi-hop questions and retrieval-augmented question answering with multiple retrieved documents. Our results highlight Parallel Context Windows as a promising method for applying off-the-shelf LLMs in a range of settings that require long text sequences. We make our code publicly available at https://github.com/ai21labs/parallel-context-windows .

## introduction a key parameter of a large language model (llm) is its context window, the number of text tokens it can process in a forward pass. current llm architectures limit the context window size-typically up to several thousand tokens-because the global nature of the attention mechanism imposes computational costs quadratic in context length. this presents an obstacle to use cases where the llm needs to process a lot of text, e.g., tackling tasks that require long inputs @xcite , considering large sets of retrieved documents for open-book question answering @xcite , figure 1 : in-context learning (icl) accuracy against n-shot training examples for the banking77 intent classification dataset @xcite using the model jurassic-1-grande (17b). the blue line shows the improvement in performance as the context window is filled with examples; the orange line shows how our parallel context windows method, which adds up to four times more training examples, provides a significant boost in performance. the error bars represent the standard deviation across multiple runs, as explained in section 3.1. or performing in-context learning @xcite when the desired input-output relationship cannot be adequately characterized within the context window. previous work has addressed such obstacles by training dedicated architectures, e.g., training sparse attention mechanisms for long inputs @xcite and fusion-in-decoder readers for retrieved documents @xcite . however, these architectures are often tailored to specific use cases, and they are often constrained in terms of their size as a tradeoff, in order to facilitate long text consumption. it remains an open problem to find an effective way to allow off-the-shelf llms to process text longer than its original context window, without dedicated training. in this paper, we introduce parallel context wintest example tokens attend to the tokens of all context windows. we conducted an in-depth investigation of the extent to which parallel context windows can improve llms' ability to perform in-context learning @xcite : when a pretrained llm is given an input sequence of concatenated "training" input-output pairs representing a task, followed by 1 we will use llm to refer to decoder-only language models. a single "test" input, it is able to supply the corresponding test output with high accuracy. crucially, in the setting of in-context learning, the context window limitation inherently caps the number of training examples that can be inserted before the test example. this significantly limits the applicability of in-context learning for tasks with long or highly diverse inputs or outputs. we focus on these types of tasks, showing that parallel context windows significantly aid incontext learning of two task families that tend to suffer from low in-context learning performance: classification tasks that have many classes and extractive question answering tasks. we experiment with jurassic-1 models @xcite having between 7b and 178b parameters and gpt2 models @xcite having between 750m and 1.5b parameters. notably, using 3 parallel context windows leads to average performance gains of 6.7, 7.3, and 7.9 points in the in-context learning scores of classification tasks with over 5 classes for jurassic-1 models of sizes 7b, 17b, and 178b, respectively (see example in figure 1 ). our results show that parallel context windows broadens the scope of tasks that can be learned via the ## parallel context windows this section provides the details of our parallel context windows method. the high-level idea of pcw is to insert a long input sequence into multiple replicas of the llm's original context window, and to allow for a small amount of tokens at the end of the sequence to attend to all of the context windows simultaneously. we design pcw so that the modifications made to the off-the-shelf llm are minimal, such that processing long contexts remains effective even without further training of the llm. a side advantage is that the llm modifications required for pcw are quite simple to implement. specifically, pcw applies two modifications to two mechanisms in common autoregressive llms: the positional embeddings (section 2.1) and the attention mask (section 2.2). figure 3 illustrates both changes. ## pcw for question answering in this section, we explore potential usages of pcw in other settings than in-context learning. specifically, we examined two question-answering settings where pcw is expected to help aggregate information from multiple texts. firstly, we consider the case of question answering based on retrieved documents. secondly, we experiment with multi-hop reasoning, where the model is required to utilize more than one text while answering a question. importantly, while in section 3 the parallel context windows were used for processing training examples for icl, in this section the windows are used for parallel processing of documents related to the test example. ## conclusion and future work in recent years, a multitude of successful approaches have been proposed for allowing transformer-based language models to leverage large amounts of text during inference, leading to a variety of dedicated architectures. in parallel, however, the mainstream llm production line of new models with "regular"-up to several thousand tokens-context window sizes enjoys faster progress in the form of scaling, innovation, and data updating. this paper introduced parallel context windows (pcw): a simple approach for allowing any offthe-shelf llm to broaden the scope of text it can access during inference. we showed the effectiveness of pcw in the framework of in-context learning, where access to a context that is larger by a factor of b implies learning from b times more training examples. our results show that pcw is more effective than the vanilla single context window approach for in-context learning over a broad set of multi-class classification tasks, suggesting that pcw could improve in-context learning in tasks with diverse input or output spaces. we also showed promising signals for applying pcw for multiple retrieved document reading. two key directions of future work strike us as particularly promising. first, by demonstrating that an off-the-shelf llm can attend to substantially larger quantities of text via pcw, our results motivate further investigation of the pcw method in other settings in which it would be desirable to apply mainstream llms over long text sequences. second, though our results suggest that pcw is effective without further training, we believe that further (short) training of an llm with parallel context windows could further enhance the abilities demonstrated in this work.