PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2212.10947v3cs.CL

TL;DR

LLMs’ context windows limit their ability to process long text, while specialized approaches often require training dedicated architectures. PCW lets off-the-shelf LLMs process multiple context windows without further training. It outperforms the vanilla single-window approach across a broad set of multi-class classification tasks and shows promising signals for reading multiple retrieved documents.

  • Problem

    LLMs’ context windows limit the number of tokens they can process, constraining long-input tasks and the in-context examples available for tasks with diverse inputs or outputs.

    Source: §1 · p. 1

  • Method

    PCW splits long text into parallel context windows, restricts attention within each window, and reuses positional embeddings so task tokens can attend across windows without further training.

    Source: §1 · p. 2

  • Results

    PCW is more effective than a vanilla single context window across a broad set of multi-class classification tasks, with promising signals for reading multiple retrieved documents.

    Source: §6 · p. 9

  • Takeaways & Limitations

    PCW makes a larger amount of text accessible to off-the-shelf LLMs during decoding, warranting further investigation for applications that require it.

    Source: §1 · p. 3

    PCW does not improve performance on every task; its suitability for long-text generation remains an open question.

    Source: Limitations · p. 10

Abstract

from arXiv · show

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.

1 Introduction

Parallel Context Windows (PCW) addresses context limits that constrain long-input tasks and in-context learning by letting off-the-shelf LLMs process multiple windows without further training. The paper reports gains in multi-class classification and extractive question answering, with additional tests in multi-hop and retrieval-augmented question answering.

1 Introduction

  • LLM context windows are typically limited to several thousand tokens because global attention has computational costs quadratic in context length.This constrains tasks involving long inputs, large sets of retrieved documents, or more in-context examples.Source: §1 · pp. 1–2
  • PCW splits long text into parallel windows, restricts each window’s attention to itself, and reuses positional embeddings so task tokens can attend to all windows without further training.The method makes two post-hoc modifications to a pretrained decoder-based LLM.Source: §1 · p. 2
  • With 3 windows, classification tasks with more than 6 classes gained 6.7 and 7.1 points on LLaMA 32.5B and 65B, and 7.4, 8.2, and 8.7 on Jurassic-1 7B, 17B, and 178B.These are average performance gains over the single-window approach.Source: §1 · p. 2
  • PCW is tested in multi-hop question answering on HotpotQA and retrieval-augmented question answering on Natural Questions, with reported benefits in these settings.The authors report benefits in some multi-hop cases and advantages from reading several retrieved documents in parallel.Source: §1 · p. 3
  • The authors conclude that PCW makes more text accessible to off-the-shelf LLMs during decoding and merits further investigation for applications requiring long text sequences.They also identify long-input tasks and in-context learning as potential applications.Source: §1 · pp. 1, 3

2 Parallel Context Windows

PCW expands the processable context by replicating positional embeddings across windows and restricting attention so each context window reads only itself while task tokens can read across all windows. These modifications allow decoding with more context while keeping computational cost linear in the number of windows.

2 Parallel Context Windows

  • PCW makes two modifications to a pretrained LLM—positional embeddings and the attention mask—so it can process long contexts without further training.The method inserts input into replicas of the original context window and lets end-of-sequence task tokens attend to them.Source: §2 · p. 3

2.1 Positional Embeddings Modification

  • PCW expands processable context from C = N − T tokens to B · C + T by reusing the first C positional embeddings across B windows and retaining the last T embeddings for task tokens.The task tokens then see the window replicas as context in their near past.Source: §2.1 · pp. 3–4

2.2 Attention Mask Modification

  • Context tokens attend autoregressively only within their own window, while task tokens attend to all context windows; this lets decoding use B times more context at linear computational cost in B.Setting B = 1 recovers the vanilla LLM mechanism.Source: §2.2 · p. 4

3 PCW for In-Context Learning

PCW is evaluated for in-context learning across classification and information-extraction tasks, with experiments spanning multiple model families and sizes. It most strongly benefits classification with many output classes and consistently improves the tested information-extraction tasks.

3.1 Experimental Setup

  • The experiments cover 9 models from GPT2, LLaMA, and Jurassic-1, 15 classification datasets, datasets with up to 150 output classes, and four information-extraction datasets.PCW uses B = 3 windows in the main experiments, with B × nmax training examples versus nmax for regular ICL.Source: §3.1 · pp. 4–5

3.2 Classification Tasks Results

  • With more than 6 classes, PCW significantly outperforms ICL in nearly all models and datasets, with average gains of 7.4, 8.2, and 8.7 for J1 and 2.3, 5.3, 6.7, and 7.1 for LLaMA models.The J1 and LLaMA values are ordered by increasing model size.Source: §3.2 · pp. 5–6
  • For J1 models, class count and average PCW improvement have Pearson r = 0.92; for LLaMA models, r = 0.79, with 10–15-point gains in many dozens-of-label datasets.The J1 correlation is between log-number of classes and average improvement; GPT2 models show no significant correlation in the cited discussion.Source: §3.2 · p. 6, Figure 4 · p. 6
  • On matched coarse- and fine-grained datasets, PCW gains 2.6 and 8.1 points on TREC and 2.5 and 9.0 points on NLU, respectively.The paired datasets are used to isolate class-count effects from differences such as domain, style, or genre.Source: §3.2 · p. 7
  • PCW reduces average classification-score standard deviations for J1-Large, J1-Grande, and J1-Jumbo to 3.1, 2.3, and 2.6, versus 3.9, 3.4, and 3.9 with standard ICL.For LLaMA models in increasing size order, the corresponding standard deviations are 4.0, 2.9, 2.9, and 2.2 with PCW, versus 4.6, 3.8, 3.6, and 3.2 with vanilla ICL.Source: §3.2 · p. 7

3.3 Information Extraction Results

  • On information-extraction tasks, including extractive question answering, PCW consistently improves results for J1-Large and J1-Grande, significantly in almost all cases, and also lowers standard deviations.The reported metrics are EM for ATIS and MIT Movies, and F1 for SQuAD and adversarialQA.Source: §3.3 · p. 7, Table 3 · p. 7
  • For multiple-choice QA, J1-Grande significantly improves on OpenBookQA, while other tested cases show no significant improvement or worsening over ICL.The tested tasks are OpenBookQA and StoryCloze.Source: §3.3 · p. 7

4 PCW for Question Answering

This section applies PCW to question-answering settings where answers draw on multiple documents, processing retrieved evidence or multi-hop evidence in parallel windows.

4.1 Retrieval Based Question Answering

  • PCW outperforms the baseline across tested candidate-document counts on J1-Grande NQ, and both methods improve as more documents become available.Figure 5 plots Exact Match against the number of documents in a single window.Source: Figure 5 · p. 8, §4.1 · p. 8
  • For retrieval QA, PCW divides retrieved documents across windows, processing B × N documents—for example, 3 when B=3 and N=1—and evaluates with Exact Match.The single-window baseline uses retrieved Wikipedia evidence and in-context examples; the PCW setup increases the number of retrieved documents per question.Source: §4.1 · pp. 7–8

4.2 Multi-hop Question Answering

  • On HotpotQA, PCW improves comparison-question performance for J1-Large and J1-Grande but degrades bridge-question performance, matching independent versus cross-conditioned evidence processing.Comparison questions allow independent extraction from two texts, whereas bridge questions require interpreting each text in relation to the other.Source: Table 4 · p. 8, §4.2 · p. 8
  • HotpotQA uses 10 evidence documents—2 gold and 8 distractors—with five windows holding two documents each in a zero-shot J1 evaluation.Answering requires reasoning over both gold documents.Source: §4.2 · p. 8
  • Encoding sequential relations between windows remains an open direction for making PCW more suitable for bridge questions.The current approach is less suitable when processing one text depends on the other.Source: §4.2 · pp. 8–9

5 Related Work

The paper relates PCW to methods for improving in-context learning and extending context, emphasizing its compatibility with existing LLMs and its no-training approach.

5.1 In-Context Learning

  • PCW is complementary to calibration, example-ordering, and noisy-channel methods, and supports generative tasks.These methods address few-shot performance issues, while PCW focuses on increasing the number of examples shown to the model.Source: §5.1 · p. 9

5.2 Expanding the Context Window

  • The paper presents PCW as the first approach among those discussed to use existing LLMs for longer inputs without further training.Prior approaches include sparse attention, relative positional factors, and encoder–decoder methods; some require further training or retain quadratic self-attention costs.Source: §5.2 · p. 9
  • Compared with concurrent work by Hao et al. (2022), PCW reports gains without scaling attention weights, including particularly large gains for diverse-output tasks.The paper also reports effectiveness with learned or rotary positional embeddings and gains in both ICL and document retrieval settings.Source: §5.2 · p. 9

6 Conclusion and Future Work

The conclusion reports broad in-context-learning gains and promising retrieval results, then identifies further applications and training as future directions.

6 Conclusion and Future Work

  • Mainstream LLM development continues to advance through scaling, innovation, and data updating while regular context windows remain up to several thousand tokens.Source: §6 · p. 9
  • PCW is more effective than vanilla single-window ICL across a broad set of multi-class classification tasks, with promising signals for reading multiple retrieved documents.The paper frames these results as suggesting potential gains for ICL tasks with diverse input or output spaces.Source: §6 · p. 9
  • Future work includes testing PCW in other long-text settings and assessing whether short additional training could further enhance its demonstrated abilities.The paper describes both as promising directions.Source: §6 · p. 9

Limitations

PCW requires choosing the number of context windows in advance, and its benefits vary by task. For unequal-length windows, the experiments used left indentation and a shared BOS token, with sample allocation and document ordering adapted to the task.

Limitations

  • The number of windows B must be chosen beforehand; results were task dependent, with diminishing returns around B = 5 to 7.Most results used B = 3; the authors recommend tuning B on a development set when available, or conservatively using B = 3 otherwise.Source: Limitations · p. 10
  • PCW shows impressive in-context-learning gains for tasks such as multi-class classification and information extraction, yet it does not improve every task, suggesting some tasks may not suit parallel processing.The authors say PCW is more suitable when input text can be divided into a few independent inputs, while its benefit for long-text generation remains an open question.Source: Limitations · p. 10

Handling context windows of various lengths

  • For unequal-length windows, left indentation was preferred for ICL, padding had little effect, and using one shared BOS token performed better than multiple BOS tokens.The experiments used left indentation throughout.Source: Handling context windows of various… · p. 12
  • ICL samples were allocated equally across windows and balanced by greedily switching long and short samples; the longest percentile of train and test samples was removed.The per-window sample limit used N, the context window size; Tmax, the longest test sample length; and D90, the 90th percentile of train-sample lengths.Source: Handling context windows of various… · p. 12
  • For Section 4.1, documents were ranked within windows so each window ended with its highest-ranked document; Section 4.2 used random window division.Training examples were not parallelized: baseline and PCW used the same randomly selected examples, with new examples drawn for each test sample.Source: Handling context windows of various… · p. 12

A.2 Evaluation Details

Evaluation used constrained greedy decoding for classification and greedy decoding at temperature 0 for information extraction. Restrictive decoding reduced computational costs across most experiments.

A.2 Evaluation Details

  • For classification, constrained greedy decoding allowed only tokens that could form a valid label, reducing computation but potentially changing baseline and PCW performance slightly.The authors expected the effect to be minor because most labels had few tokens and their first token was usually indicative of the label.Source: §A.2 · p. 13
  • Information extraction predictions used greedy decoding at temperature 0, with Exact Match or F1 as the evaluation metric.The decoding setup was similar to Zhao et al. (2021).Source: §A.2 · p. 13
  • Most classification tasks took 1–4 GPU hours across models, while J1-Jumbo experiments took roughly 10–50 GPU hours each.Section 3.3 and Section 4 experiments took up to 20 GPU hours each.Source: §A.2 · p. 13

B.1 Overview

The evaluation covered classification, extraction, multiple-choice, and question-answering datasets, with dataset-specific prompt formats and test-set choices. ICL experiments used input-output examples without instructions.

B.1 Overview

  • The classification evaluation used 15 datasets, including SST-2, RTE, AGNews, SST-5, TREC, DBPedia, BANKING77, and CLINIC150; TREC and NLU used fine- and coarse-grained labels.The passage also lists CR, Subj, CB, YELP, and NLU among the classification datasets.Source: §B.1 · p. 13
  • Six additional datasets, evaluated only with J1 models, covered extraction and multiple-choice: ATIS airlines, MIT Movie Genre, SQuAD, adversarialQA, OpenBookQA, and StoryCloze.The listed per-dataset nmax values range from 8 for SQuAD and adversarialQA to 87 for OpenBookQA.Source: §B.1 · p. 13
  • Section 4 used Natural Questions and HotpotQA; evaluation used standard test sets or validation sets when public test sets were unavailable.Natural Questions used half its 3,610-sample test set, while HotpotQA used its full 7,405-sample validation set; all datasets were in English.Source: §B.1 · p. 13
  • Natural Questions documents were loaded with Pyserini, and its validation set served as the few-shot training set because Pyserini did not support loading the training set directly.The authors considered this reasonable because the validation set was used only for few-shot prompting, not parameter optimization.Source: §B.1 · p. 13

B.2 Preprocessing and Formatting

  • ICL prompts used input-output pairs without instructions; classification formats drew mainly on Lu et al. (2021), while extraction and multiple-choice formats followed Brown et al. (2020).New formats were created for classification datasets with dozens of labels; Section 4 prompts followed Lazaridou et al. (2022).Source: §B.2 · p. 14

C The Effect of the Number of Context Windows on Performance

The optimal number of context windows depends on task and class count. Increasing B continued to help some high-class-count datasets, while performance on SST-2 appeared to worsen.

C The Effect of the Number of Context Windows on Performance

  • The optimal B depended on the number of unique labels: higher-class-count classification tasks improved with larger B, while tasks with few classes tended to favor B = 1 or 2.The main experiments displayed B = 3 results for simplicity.Source: C The Effect of the Number of Context… · p. 14
  • AGNews and DBPedia continued improving as B increased, with convergence around B = 6, while increasing B seemed to harm SST-2.The comparison used a representative three-dataset subset, with B varied from 1 to 8.Source: C The Effect of the Number of Context… · p. 14
  • The authors recommend choosing B on a development set when available, or using a conservative B = 3 when one is unavailable.They considered very large numbers of training samples—and further investigation of larger B—less relevant to most practical ICL cases.Source: C The Effect of the Number of Context… · p. 14

D.1 Replication Experiments with GPT2

The GPT2-Large and GPT2-XL replication experiments mirror the earlier results, and an error analysis suggests PCW can mitigate recency bias in ICL.

D.1 Replication Experiments with GPT2

  • PCW placed a training example with the test label closer to the test example 62% of the time; in those cases it outperformed ICL by 19.4%, versus roughly 10% across the full test set.The analysis suggests PCW provides a solution to the recency bias noted by Zhao et al. (2021).Source: §D.1 · p. 14
  • Qualitative error inspection suggested vanilla ICL fails more often when the test label appears earlier in the prompt.With more context windows, PCW more often presents a training example with the test label near the end of a window.Source: §D.1 · p. 14
  • The GPT2-Large and GPT2-XL replication results mirror the results reported in Table 1.Source: §D.1 · p. 14

D.2 Multiple-Choice QA

The multiple-choice QA experiments evaluate PCW on OpenBookQA and StoryCloze using few-shot prompting and candidate-completion likelihoods. More in-context examples improved J1-Grande on OpenBookQA, but did not significantly affect the other scenarios.

D.2 Multiple-Choice QA

  • The experiments covered OpenBookQA and StoryCloze, evaluating each possible completion by its average per-token language-model likelihood.The setup followed Brown et al. (2020) but did not use their calibration; the J1-Grande OpenBookQA sample sizes were reduced to 15 training sets and 125 test examples.Source: §D.2 · pp. 14–15
  • Increasing in-context training examples under PCW improved J1-Grande performance on OpenBookQA, but did not significantly affect the other scenarios.The authors say further analysis on more datasets is needed to understand potential gains for multiple-choice tasks.Source: §D.2 · p. 15, Table 5 · p. 15
  • The appendix includes results for different B choices for J1-Large, J1-Grande, and J1-Jumbo; the Jumbo experiments tested only B = 2 and B = 3.The best result for each dataset is boldfaced in the tables.Source: Table 7 · p. 16, Table 8 · p. 16, Table 9 · p. 17
  • Figure 7 reports that larger context windows benefit multi-class tasks, while simpler tasks with fewer classes may lose performance.Source: Figure 7 · p. 17

Related papers

Loading 2212.10947v3…