Parallel Context Windows for Large Language Models

Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics Volume 1: Long Papers, pages 6383–6402 July 9-14, 2023 ©2023 Association for Computational Linguistics Parallel Context Windows for Large Language Models Nir Ratner Yoav Levine Yonatan Belinkov Ori Ram Inbal Magar Omri Abend Ehud Karpas Amnon Shashua Kevin Leyton-Brown Yoav Shoham AI21 Labs nirr@ai21.com 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 ar- chitectures, 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 pa- rameters, 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 doc- uments. 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. 1 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. Cur- rent LLM architectures limit the context window size—typically up to several thousand tokens— be- cause 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., tack- ling tasks that require long inputs (Tay et al., 2020; Shaham et al., 2022), considering large sets of re- trieved documents for open-book question answer- ing (Karpukhin et al., 2020; Levine et al., 2022a,b), 10 0 10 1 10 2 n-shot Training Examples 10 20 30 40 50 60 70 BANKING77 Accuracy(%) Single Context Window ICL (Regular) Parallel Context Windows ICL (Ours) Figure 1: In-context learning (ICL) accuracy against n-shot training examples for the BANKING77 intent classification dataset (Casanueva et al., 2020) 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 signifi- cant boost in performance. The error bars represent the standard deviation across multiple runs, as explained in Section 3.1. or performing in-context learning (Brown et al., 2020) when the desired input-output relationship cannot be adequately characterized within the con- text window. Previous work has addressed such obstacles by training dedicated architectures, e.g., training sparse attention mechanisms for long inputs …

Content truncated to 500 words