PaperlayerView exact revision on arXiv

Source-linked AI summary

Retrieval meets Long Context Large Language Models

Peng Xu, Wei Ping, Xianchao Wu, Lawrence McAfee, Chen Zhu, Zihan Liu, Sandeep Subramanian, Evelina Bakhturina, Mohammad Shoeybi, Bryan Catanzaro

arXiv:2310.03025v2cs.CLcs.AIcs.IRcs.LG

TL;DR

Practitioners still lack clear evidence on whether retrieval augmentation or longer context performs better, and whether the approaches work well together. The paper compares them across two LLMs and nine long-context tasks, finding that retrieval-augmented Llama2-70B-32k achieves a higher average score than GPT-3.5-turbo-16k and its non-retrieval baseline.

  • Problem

    It remains unclear whether extending context windows or retrieval augmentation better serves downstream tasks, and whether combining them can improve accuracy.

    Source: §1 · pp. 1–2

  • Method

    The authors systematically study retrieval augmentation and context-window extension using two pretrained LLMs across nine downstream long-context tasks.

    Source: §1 · p. 2

  • Results

    On average, retrieval-augmented Llama2-70B-32k scores 43.6, above GPT-3.5-turbo-16k (42.8) and its non-retrieval baseline (40.9).

    Source: §1 · p. 2

  • Takeaways & Limitations

    The study suggests retrieval-augmented 4K models may offer practitioners comparable performance to 16K models with greater inference efficiency.

    Source: §5 · p. 9

    The study focuses on decoder-only LLMs with 43B and 70B parameters, selected for their zero-shot ability to incorporate context after instruction tuning.

    Source: §2.4 · p. 3

Abstract

from arXiv · show

Extending the context window of large language models (LLMs) is getting popular recently, while the solution of augmenting LLMs with retrieval has existed for years. The natural questions are: i) Retrieval-augmentation versus long context window, which one is better for downstream tasks? ii) Can both methods be combined to get the best of both worlds? In this work, we answer these questions by studying both solutions using two state-of-the-art pretrained LLMs, i.e., a proprietary 43B GPT and Llama2-70B. Perhaps surprisingly, we find that LLM with 4K context window using simple retrieval-augmentation at generation can achieve comparable performance to finetuned LLM with 16K context window via positional interpolation on long context tasks, while taking much less computation. More importantly, we demonstrate that retrieval can significantly improve the performance of LLMs regardless of their extended context window sizes. Our best model, retrieval-augmented Llama2-70B with 32K context window, outperforms GPT-3.5-turbo-16k and Davinci003 in terms of average score on nine long context tasks including question answering, query-based summarization, and in-context few-shot learning tasks. It also outperforms its non-retrieval Llama2-70B-32k baseline by a margin, while being much faster at generation. Our study provides general insights on the choice of retrieval-augmentation versus long context extension of LLM for practitioners.

1 INTRODUCTION

The introduction frames retrieval augmentation and context-window extension as approaches to compare and combine, and reports that retrieval improves long-context performance, with the strongest result from retrieval-augmented Llama2-70B-32K.

Introduction

  • The work asks whether context-window extension outperforms retrieval for downstream tasks with informative queries and whether combining the methods can further improve accuracy.Retrieval uses a standalone retriever to select relevant context for the LLM.Source: §1 · pp. 1–2
  • The study evaluates a proprietary 43B pretrained GPT and Llama2-70B on nine long-context tasks spanning single- and multi-document QA, query-based summarization, and in-context few-shot learning.Source: §1 · p. 2
  • Using much less computation, retrieval-augmented 4K GPT-43B scores 29.32 versus 29.45 for 16K, while Llama2-70B scores 36.02 versus 36.78.The introduction reports these results as comparable performance between the retrieval-augmented 4K models and 16K long-context models.Source: §1 · p. 2
  • 43.6 average score for retrieval-augmented Llama2-70B-32K exceeds GPT-3.5-turbo-16k’s 42.8, Davinci-003, and its non-retrieval baseline’s 40.9.The model is also 4× faster on NarrativeQA.Source: §1 · p. 2

2 RELATED WORK

The related work surveys long-context LLMs, efficient attention, and retrieval-augmented language models, including prior studies of how retrieval interacts with long context.

2.1 PARALLEL WORK

  • Bai et al. (2023) found retrieval helpful only for Llama2-7B-chat-4k, whereas our Llama2-70B-32k performs as well as Davinci003 and GPT-3.5-turbo-16k and improves with retrieval.The authors hypothesize that black-box APIs make controlled experiments challenging and that smaller white-box models have limited zero-shot ability to incorporate retrieved context.Source: §2.1 · p. 2

2.2 LONG CONTEXT LARGE LANGUAGE MODELS

  • Long-context pretraining faces quadratic attention costs and predominantly short documents, motivating continued-training and positional-extension methods.This work applies positional interpolation to extend 4K context windows of proprietary 43B and Llama2-70B models to 16K and 32K; prior black-box evaluations also report a “lost in the middle” phenomenon.Source: §2.2 · pp. 2–3

2.3 EFFICIENT ATTENTION METHODS

  • Approximate attention methods trade inductive biases for reduced quadratic-complexity bottlenecks, while FlashAttention speeds up exact attention through GPU-memory-aware computation.Prior approaches include sparse, recurrence-based, low-rank, memory-based, and similarity- or clustering-based methods; their inductive biases may reduce general LLM training quality.Source: §2.3 · p. 3

2.4 RETRIEVAL-AUGMENTED LANGUAGE MODELS

  • Retrieval augmentation has been used to improve language-model perplexity, factual accuracy, downstream accuracy, and in-context learning through inference, fine-tuning, pretraining, or end-to-end integration.Most prior studies focus on models around 10 billion parameters; this work studies decoder-only 43B and 70B LLMs because their scale supports strong zero-shot context incorporation after instruction tuning.Source: §2.4 · p. 3

3 EXPERIMENTAL SETUP

The experimental setup compares retrieval with extended-context self-attention for zero-shot long-context question answering and summarization. It specifies the models, datasets and metrics, context-extension procedure, retrievers, and instruction tuning used in the evaluations.

3.1 LARGE LANGUAGE MODELS

  • The study experimented with two pretrained GPT models: proprietary Nemo GPT-43B, trained with a 4,096-token sequence length and RoPE embeddings, and publicly available Llama2-70B, with a 4,096-token context window and RoPE embeddings.Source: §3.1 · p. 4

3.2 DATASETS AND METRICS

  • Four Scroll validation datasets span meeting summarization (QMSum; Zhong et al., 2021), paper QA (Qasper; Dasigi et al., 2021), noisy long-text QA (NarrativeQA; Koˇciský et al., 2018), and multiple-choice QA (QuALITY; Pang et al., 2022).The study frames the comparison as zero-shot integration of long-context information through retrieval or self-attention, using models larger than 40B after instruction tuning.Source: §3.2 · p. 4, §3.1 · p. 4
  • Three LongBench datasets add multi-hop QA (HotpotQA; Yang et al., 2018; MuSiQue; Trivedi et al., 2022) and cross-field long-context QA (MultiFieldQA-en; Bai et al., 2023).MuSiQue is designed to reduce reasoning shortcuts and train-test leakage, while MultiFieldQA-en distributes evidence across diverse documents.Source: §3.2 · pp. 4–5
  • Evaluation documents range from 4.9k tokens for Qasper to 84k for NarrativeQA; non-retrieval baselines truncate inputs to fit their context windows.The study reports ROUGE-1/2/L geometric mean for QMSum, exact match for QuALITY, and F1 for the other five datasets.Source: §3.2 · p. 5

3.3 CONTEXT WINDOW EXTENSION

  • Position interpolation extends GPT-43B from 4K to 16K and Llama2-70B to 16K and 32K, with finetuning on the Pile to adapt position embeddings.Finetuning follows Chen et al. (2023), using batch size 128 and learning rate 5e-6.Source: §3.3 · p. 5

3.4 RETRIEVAL

  • Retrieval experiments use Dragon (Lin et al., 2023), Contriever (Izacard et al., 2021), and OpenAI text-embedding-ada-002 to rank document chunks for generation.Documents are split into 300-word chunks; the top N are ranked by embedding dot product and concatenated in relevance order. Top-5 chunks generally fit 4K, while top-10 and top-20 fit 16K.Source: §3.4 · pp. 5–6

3.5 INSTRUCTION TUNING

  • Instruction tuning uses 102K samples from six datasets and a proprietary conversational corpus, then trains on answer tokens for 1,000 steps.The prompt places context before the dialogue; tuning uses batch size 128 and learning rate 5e-6, and subsequent results use instruction-tuned chat models.Source: §3.5 · p. 6

4 RESULTS

Retrieval improves long-context LLM results, and retrieval-augmented Llama2-70B-32k outperforms GPT-3.5-turbo-16k and Davinci-003 on the reported average comparisons. The results also show that performance depends on context length, model size, retriever, and the number of retrieved chunks.

4.1 MAIN RESULTS

  • Retrieval raises Llama2-70B-4K from 31.61 to 35.73 and GPT-43B-4K from 26.44 to 29.32.For HotpotQA, extending 4K to 16K raises scores from 34.64 to 43.97 for Llama2-70B and from 28.91 to 37.48 for GPT-43B. The seven tasks' minimum average sequence length exceeds 4096, causing truncation; HotpotQA benefits from longer contexts because its multi-hop questions require intermediate hops.Source: §4.1 · p. 6
  • With the same top-five retrieved chunks, 16K and 32K retrieval models outperform 4K; Llama2-70B exhibits lost-in-the-middle, and 32K better incorporates retrieved context.The authors confirm the lost-in-the-middle pattern across Llama2-70B context lengths and report higher accuracy for 32K when incorporating the top-five retrieved context.Source: Figure 1 · p. 7, §4.1 · p. 7

4.2 COMPARING TO OPENAI MODELS

  • Llama2-70B-32k-ret outperforms GPT-3.5-turbo-16k on Avg-7 and Davinci-003 on Avg-4*.Avg-7 averages seven datasets; Avg-4* averages four ZeroSCROLLS datasets.Source: Table 3 · p. 7, §4.2 · p. 8

4.1 MAIN RESULTS

  • Retrieval improves Llama2-70B across context lengths, with its 32K retrieval model scoring 39.60 versus 37.36 without retrieval; gains are greater than for GPT-43B.The authors attribute the difference from smaller-model results to stronger zero-shot ability to incorporate retrieved evidence: retrieval did not help Llama2-7B-32k or ChatGLM2-6B-32k, while it helped Llama2-7B-chat-4k, whose context length remained a bottleneck.Source: §4.1 · pp. 7–8

4.3 ABLATION ON DIFFERENT RETRIEVERS

  • Dragon, Contriever, and OpenAI embeddings all show retrieval boosting Llama2-70B at both 4K and 32K.The reported consistency holds across the tested retrievers and both context lengths.Source: Table 4 · p. 8, §4.3 · p. 9

4.4 INCREASING THE NUMBER OF RETRIEVED CHUNKS

  • For each tested sequence length, top 5 or top 10 chunks yield the best average results; adding up to 20 is unhelpful and can hurt performance.The authors suggest lost-in-the-middle effects or distraction from irrelevant information as possible explanations.Source: Table 5 · p. 8, §4.4 · p. 9

4.2 COMPARING TO OPENAI MODELS

  • On MSQ, HQA, and MFQA, GPT-3.5-turbo-4k retrieval raises the average from 37.08 to 41.15; GPT-3.5-turbo-16k retrieval and non-retrieval averages are 43.27 and 43.60, below Llama2-70B-32k-ret's 44.51.GPT-3.5-turbo-16k is a black-box API, so its implementation, model size, and preprocessing are unknown.Source: §4.2 · p. 8

4.5 RETRIEVAL FOR FEW-SHOT TASKS

  • On two LongBench few-shot tasks, Llama2-70B-32k-ret outperforms its non-retrieval baseline and GPT-3.5-turbo-16k by a large margin.The method uses each dataset question to retrieve relevant QA pairs from the provided few-shot examples.Source: Table 6 · p. 9, §4.5 · p. 9

5 CONCLUSION

The study finds that retrieval improves instruction-tuned LLMs with both short and extended context windows, and that combining retrieval with context extension yields a model that outperforms GPT-3.5-turbo-16k and Davinci003 on average across downstream tasks with informative queries.

  • Llama2-70B-32k-ret outperforms GPT-3.5-turbo-16k and Davinci003 in average score across downstream tasks with informative queries.Source: §5 · p. 9
  • Retrieval largely boosts performance for both 4K short-context and 16K/32K long-context LLMs.Source: §5 · p. 9
  • 4K context LLMs with simple retrieval augmentation perform comparably to 16K context LLMs while being more efficient at inference.Source: §5 · p. 9

6 FUTURE DIRECTIONS

The paper identifies three future directions: developing advanced methods for pretrained LLMs, extending context windows beyond 32K, and mitigating the lost-in-the-middle phenomenon.

6 FUTURE DIRECTIONS

  • Developing memory or hierarchical-attention methods for pretrained LLMs such as Llama2-70B remains non-trivial.Source: §6 · p. 9
  • Extending context windows to 64K or longer for 70B-parameter models is an interesting direction, though longer-sequence pretraining requires much more computation.Source: §6 · p. 9
  • Mitigating the open lost-in-the-middle problem is another direction; continued pretraining with UL2 loss (Tay et al., 2022b) is a potential solution.Source: §6 · p. 9

A APPENDIX

In the Table 7 example, LLaMA2-70b-32k-ret is the only model that derives the correct answer, 2006. Figure 2 shows token-length distributions for full documents and selected chunks across seven datasets.

A APPENDIX

  • 2006 is answered correctly only by LLaMA2-70b-32k-ret; all other models fail.The table caption attributes the result to retrieval combined with long context.Source: Table 7 · p. 15
  • Figure 2 presents token-length distributions for full documents and the top-5, 10, and 20 chunks across seven datasets.Source: Figure 2 · p. 16
Loading 2310.03025v2…