PaperlayerView exact revision on arXiv

Source-linked AI summary

BABILong: Testing the Limits of LLMs with Long Context Reasoning-in-a-Haystack

Yuri Kuratov, Aydar Bulatov, Petr Anokhin, Ivan Rodkin, Dmitry Sorokin, Artyom Sorokin, Mikhail Burtsev

arXiv:2406.10149v2cs.CLcs.AI

TL;DR

Existing benchmarks do not adequately evaluate models across the very long contexts they increasingly support. BABILong addresses this gap with scalable reasoning tasks whose facts are distributed through natural documents. Popular LLMs use only 10-20% of context, while fine-tuned recurrent memory transformers process up to 50 million tokens.

  • Problem

    Existing long-context benchmarks scale only to 40,000 tokens while models support hundreds of thousands or millions, limiting evaluation at relevant lengths.

  • Method

    BABILong embeds facts from 20 reasoning tasks in natural PG19 documents and provides scalable evaluation splits reaching 10 million tokens.

  • Results

    Popular LLMs use only 10-20% of context and degrade as reasoning complexity increases, whereas fine-tuned recurrent memory transformers process up to 50 million tokens.

  • Takeaways & Limitations

    BABILong provides an extensible evaluation framework for testing long-context reasoning as document length and fact distribution increase.

  • Takeaways & Limitations

    The dataset reuses limited bAbI vocabularies, producing low variety in names and objects and making tasks easier for fine-tuned models.

Abstract

from arXiv · show

In recent years, the input context sizes of large language models (LLMs) have increased dramatically. However, existing evaluation methods have not kept pace, failing to comprehensively assess the efficiency of models in handling long contexts. To bridge this gap, we introduce the BABILong benchmark, designed to test language models' ability to reason across facts distributed in extremely long documents. BABILong includes a diverse set of 20 reasoning tasks, including fact chaining, simple induction, deduction, counting, and handling lists/sets. These tasks are challenging on their own, and even more demanding when the required facts are scattered across long natural text. Our evaluations show that popular LLMs effectively utilize only 10-20\% of the context and their performance declines sharply with increased reasoning complexity. Among alternatives to in-context reasoning, Retrieval-Augmented Generation methods achieve a modest 60\% accuracy on single-fact question answering, independent of context length. Among context extension methods, the highest performance is demonstrated by recurrent memory transformers after fine-tuning, enabling the processing of lengths up to 50 million tokens. The BABILong benchmark is extendable to any length to support the evaluation of new upcoming models with increased capabilities, and we provide splits up to 10 million token lengths.

1 Introduction

BABILong addresses the mismatch between rapidly expanding model context windows and benchmarks that inadequately test long-context reasoning. It evaluates reasoning over distributed facts and reports sharp limitations in current models alongside stronger results from recurrent memory approaches.

  • Existing long-context benchmarks scale only to 40,000 tokens, despite models supporting hundreds of thousands or millions of tokens.
  • BABILong hides facts from 20 reasoning tasks inside natural PG19 background texts to test reasoning across extremely long documents.Tasks include fact chaining, induction, deduction, counting, and list/set handling.
  • BABILong supports predefined splits up to 10 million tokens and evaluates samples reaching 50 million tokens.Its construction is intended to support evaluation at almost arbitrary lengths.
  • Over 30 recent long-input language models are evaluated across different sizes, architectures, and context-extension methods.
  • Popular LLMs effectively use only 10-20% of context, with performance degrading sharply as reasoning complexity increases.Retrieval-augmented generation performs poorly, while task-specific fine-tuning helps.
  • Recurrent memory transformers answer single-fact questions in texts up to 50 million tokens when retrieval-augmented generation fails.The figure describes this as a record sequence length processed by a single model.

2 The BABILong Benchmark for Long Context Processing

BABILong extends basic reasoning tasks by embedding their facts among long, irrelevant natural-text backgrounds. This construction tests information selection, distractor resistance, and reasoning as context length grows.

  • Relevant task sentences are hidden among irrelevant background sentences until the augmented sample reaches a target length.Background sentences are added in their natural order, allowing lengths beyond the original task.
  • BABILong extends 20 bAbI reasoning tasks by mixing their facts, distractors, and questions with natural background text.The original tasks cover skills including spatial and temporal reasoning, deduction, and coreference resolution.
  • Models must identify supporting facts among distractors and perform reasoning such as counting, which is especially difficult for language models.
  • Table 2 reports average QA1-QA5 accuracy and groups models by their claimed context length.It compares standard models, retrieval-augmented methods, and fine-tuned small models including ARMT, RMT, and Mamba.
  • Generated benchmarks such as bAbI and BABILong are described as immune to contamination from training-set data leakage.

3 Benchmarking Results

BABILong evaluations show that current models use only a fraction of available context, with performance worsening as context length and reasoning complexity increase. Retrieval is weak on multi-fact tasks, whereas task-specific fine-tuning, especially recurrent memory models, supports much longer inputs.

  • Effective Context Size: Only 23 of 34 tested LLMs reached at least 85% accuracy on any QA1–QA3 task without background distractors.The benchmark defines satisfactory performance as accuracy above 85%.
  • Effective Context Size: For QA1, most models used only up to 4K tokens effectively, while the strongest models reached 16K–64K depending on the model.The measured full-context utilization range on QA1 was 5% to 50%.
  • Reasoning Complexity: When answers required two or three supporting facts, models generally failed to achieve satisfactory accuracy, especially when facts were embedded in background text.Only GPT-4 and Gemini Pro 1.5 solved QA2 without background text, while the best QA3 scores remained below 80%.
  • Retrieval-Augmented Generation: RAG using sentence retrieval outperformed 512-token chunk retrieval on QA1, but accuracy decreased after 16K tokens and smaller chunks risk information loss.The advantage was task-specific and may not transfer directly to real-world applications.
  • Retrieval-Augmented Generation: RAG performed weakly on BABILong because multi-hop tasks require retrieving multiple facts, while temporal dependencies and similarity limitations can cause relevant facts to be missed.In QA2 and QA3, retrieval accuracy fell below random guessing.
  • Fine-Tuning: Fine-tuned recurrent memory models processed substantially longer sequences than RAG, with RMT reaching 11.1M tokens and ARMT reaching 50M tokens.Mamba performed well but became extremely slow beyond 128K tokens; recurrent models remained faster at very long lengths.
  • Comparison with Other Benchmarks: BABILong differentiates model behavior from lengths as short as 2K tokens, whereas RULER requires at least 128K tokens to diverge significantly from short-context MMLU behavior.BABILong’s correlation with MMLU decreases as length increases, unlike RULER’s nearly constant correlation.

4 Related Work on Long Context Benchmarks and Datasets

Existing long-context benchmarks cover varied tasks and length ranges but often remain limited in scale, specificity, or suitability for modern LLMs. BABILong instead emphasizes natural-language reasoning over many distributed facts in very large corpora.

  • Earlier Benchmarks: Long Range Arena uses 1K–16K-token tasks, but many are specialized or weakly related to NLP and require task-specific fine-tuning.Its tasks include ListOps, byte-level classification, and byte-level retrieval.
  • Existing Long-Context Datasets: LongBench combines six real and synthetic task types, but its average sample lengths are 6K tokens in English and 13K in Chinese, with a 40K maximum.Its tasks include summarization, multidocument QA, and code completion.
  • Existing Long-Context Datasets: Other datasets evaluate QA, classification, summarization, retrieval, or related tasks over sources including wikis, papers, novels, and scripts.Examples include InfinityBench, Loogle, Bamboo, LVEval, NovelQA, Marathon, XL2-Bench, and DocFinQA.
  • Synthetic Long-Context Tests: Needle-in-a-haystack benchmarks enable very long synthetic evaluations, but magic-number retrieval tasks are overly simplistic and often yield perfect performance.Related variants include passkey, key-value retrieval, and counting-star insertion tasks.
  • Length and Language Coverage: Several benchmarks support predefined or generated length bins up to 256K tokens, but most are English-only or cover Chinese only in some datasets.Examples include LongBench, Ada-LEval, LVEval, Bamboo, S3Eval, RULER, and CountingStars.
  • BABILong: BABILong provides 20 tasks for natural-language reasoning over multiple facts distributed through very large textual corpora, with more deceptive information mixing than existing approaches.Its capabilities include multi-hop tracing, aggregation, deduction, induction, temporal, positional, size, and path reasoning.

Conclusions

BABILong is presented as a diverse, scalable benchmark for evaluating long-context reasoning, with adaptable lengths and distributed facts. Experiments reveal limited effective context use in many models, while recurrent transformers process up to 50 million tokens.

  • BABILong provides 20 reasoning tasks, adaptable document lengths and fact placement, predefined bins from 0k to 10M tokens, and potential leak resistance.
  • Popular open-source LLMs, GPT-4, Gemini 1.5 Pro, and RAG solutions show limitations in effective long-context utilization, relying heavily on the first 5-25% of input.
  • Among evaluated models, Mamba and recurrent transformers achieve the strongest results, while Mamba is difficult to infer beyond 128K tokens.
  • RMT and ARMT enable processing lengths up to 50 million tokens, whereas fine-tuned GPT-3.5-Turbo and Mistral-7B remain limited to 16K and 32K tokens.

Limitations

The benchmark’s scope and results depend on background-text choices, retrieval settings, dataset vocabulary, and recurrent-model trade-offs. Several limitations are identified as addressable through improved retrieval, richer vocabularies, or expanded task sources.

  • Only PG19 and Wiki were tested as background sources, and similar facts can interfere with one another, increasing benchmark difficulty.
  • The GPT-4 and Llama-3 RAG experiments did not optimize retrieval, and selected prompts may have been suboptimal.
  • The current dataset reuses bAbI fact-generation parameters, producing limited name and object variety that can make tasks easier for fine-tuned models.
  • Other fact and question sources, domain-matched documents, or document-grounded question generation could extend BABILong toward more complex real-world scenarios.
  • Recurrent approaches trade reduced parallelizability for constant memory requirements, but their finite storage capacity remains a limitation.

A Code and Data Availability

The project releases code, evaluation data, model predictions, and licensing information to support reproducible benchmarking. It also describes related long-context approaches and a recurrent transformer extended with trainable retrieval.

  • Availability: Code for generating data and evaluating models is available in the BABILong GitHub repository.
  • Availability: Pre-generated evaluation sets provide 100 samples per length and task up to 10 million tokens, plus 1000 samples per length and task up to 128K tokens.
  • Availability: The released code is licensed under Apache 2.0 and uses PG19 data under Apache 2.0 and bAbI data under BSD licensing.
  • Availability: The release includes data generation, metrics, the evaluation pipeline, and model predictions for reproducing and verifying reported results.
  • Long-context approaches: The work augments Recurrent Memory Transformer with retrieval of its own past memory tokens, described as the first combination of a recurrent transformer with trainable retrieval.

C Details on RMT, ARMT, and Mamba fine-tuning and evaluation on BABILong

RMT, ARMT, and Mamba are fine-tuned with progressively longer segment curricula and evaluated on BABILong at extended lengths. RMT results vary across training seeds, and evaluation time grows linearly with context length.

  • Fine-tuning setup: RMT and ARMT use a GPT-2 137M backbone with 512-token segments, augmented by 16 and 10 memory tokens respectively.
  • Fine-tuning setup: RMT and ARMT curricula progressively increase sequence lengths across randomly selected segment counts, with ARMT using 2-3-5-8-16-32 segments.
  • Evaluation: Figure 5 reports average performance and standard deviation across three training runs, showing variation between training seeds on five BABILong tasks.
  • Evaluation: Evaluation uses the full test set through 1M tokens and averaged results over 100 samples at 10M tokens.
  • Evaluation: Evaluation time grows linearly with context length when processing BABILong samples with recurrent models.

D Detailed LLMs evaluation on BABILong QA1-5 tasks

Across the first five BABILong tasks, model performance varies sharply with task complexity and context length. Larger models lead on short contexts, but recurrent and state-space approaches perform strongly on longer sequences.

  • Task complexity: QA1 and QA5 are easiest, with most models exceeding 70% accuracy on the 0k split.QA4 is substantially harder, while QA2 and QA3 are harder still.
  • Model scale: Parameter count improves accuracy on the shortest 0k split, where several larger models consistently outperform smaller models.The named models include GPT-4, Phi-3-medium, Qwen, Jamba, Command-R, Yi-34B, and Mixtral 8x22B.
  • Context scaling: As context length increases, some of the largest models lose their advantage over smaller models.
  • Retrieval augmentation: Retrieval-augmented Llama-3 works across context lengths up to 10M tokens but fails dramatically on QA2 and QA3.Retrieval can match or surpass weaker competitors on longer QA4 and QA5 contexts, but sentence ordering and semantic mismatch hinder reasoning on QA2 and QA3.
  • Context-extension methods: Fine-tuned RMT, ARMT, and Mamba outperform LLMs across many sequence lengths despite having up to 100 times fewer parameters.Mamba is strongest on medium-length sequences, while RMT and ARMT handle much larger sequences up to 10M tokens.

E Gemini Evaluation

Gemini 1.5 Pro 002 was evaluated on BABILong QA1 through API requests, but built-in safety filtering affected the observed evaluation results. Refusals increased with context size, reaching up to 14% of requests.

  • Evaluation conditions: Gemini 1.5 Pro 002 refused to respond to up to 14% of QA1 requests as context size increased.The evaluation used API requests and excluded refused responses from the reported results in Figure 1b.
  • Evaluation conditions: The reported results omitted requests rejected by built-in content safety filtering, even with BLOCK_NONE enabled.Full results including rejected requests were presented separately.
  • Evaluation conditions: The evaluation used 1000 samples for lengths up to 32K tokens.

F BABILong Dataset Statistics

BABILong scales bAbI reasoning tasks by distributing task facts through long natural documents with distractor text. It supports diverse task complexities and sequence lengths up to 10M tokens, while tokenizer-specific lengths can vary.

  • Task statistics: BABILong contains 20 diverse tasks spanning distractor-based needle finding, counting, logical reasoning, and spatial reasoning.QA1, QA5, and QA10 are generally easier, whereas QA7, QA15, and QA19 are among the most challenging.
  • Task statistics: The number of facts required for reasoning significantly affects task difficulty.
  • Sequence-length statistics: The benchmark scales tasks by adding distractor sentences, with prepared splits from 0k through 10M tokens.Listed splits include 0k, 4k, 8k, 16k, 32k, 64k, 128k, 512k, 1M, and 10M tokens.
  • Task statistics: Figure 7 compares model accuracy across bAbI tasks, with each dot averaging 1000 samples and black stars marking median accuracy.
  • Tokenization: Sample token counts may differ across models because BABILong lengths are measured with the GPT-2 tokenizer.The authors describe the comparison as fair despite trade-offs involving sequence length and embedding-layer size.

H Recurrent Memory Transformer Analysis

RMT memory states remain stable without relevant facts and change when new facts appear, at both 16k and 128k contexts. Its attention patterns indicate writing facts to memory and later reading them to answer questions.

  • Memory-state dynamics: At both 16k and 128k contexts, RMT memory states change visibly when a fact enters the input but remain similar to their initial states otherwise.This pattern indicates learned detection and storage of relevant facts in memory.
  • Memory attention: RMT uses memory attention first to write a fact into memory and later to read it while answering the question.

I LLMs fine-tuning results

Fine-tuning substantially improves long-context performance for GPT-3.5 and Mistral-7B on QA1, while transfer to other tasks depends on the model and training setup. GPT-4-Turbo also shows location-dependent performance, with supporting facts hardest to identify in the middle of the context.

  • Fine-tuning effects: 90%+ accuracy on QA1 was achieved by fine-tuned GPT-3.5 and Mistral-7B across context lengths.Both models significantly improved their scores after fine-tuning.
  • Fine-tuning effects: Fine-tuned GPT-3.5 improved on QA2-QA5 after training on QA1.
  • Fine-tuning effects: Full fine-tuning of Mistral-7B on QA1 degraded its scores on QA2-QA5.The evaluation used no distractor text for these comparisons.
  • Fact location: GPT-4-Turbo performed differently depending on where supporting facts appeared, with the middle of the context being most difficult.The middle corresponds to depth = 50 in Figure 10.
Loading 2406.10149v2…