[Paper Review] In Search of Needles in a 11M Haystack: Recurrent Memory Finds What LLMs Miss

Yuri Kuratov, Aydar Bulatov+4 more|arXiv (Cornell University)|Feb 16, 2024

Genetics, Bioinformatics, and Biomedical Research7 citations

TL;DR

The paper introduces the BABILong benchmark for ultra-long documents and shows that a recurrent memory transformer (RMT/RMT-R) can process up to 11 million tokens, outperforming GPT-4 and RAG on long-context tasks.

ABSTRACT

This paper addresses the challenge of processing long documents using generative transformer models. To evaluate different approaches, we introduce BABILong, a new benchmark designed to assess model capabilities in extracting and processing distributed facts within extensive texts. Our evaluation, which includes benchmarks for GPT-4 and RAG, reveals that common methods are effective only for sequences up to $10^4$ elements. In contrast, fine-tuning GPT-2 with recurrent memory augmentations enables it to handle tasks involving up to $11 imes 10^6$ elements. This achievement marks a substantial leap, as it is by far the longest input processed by any neural network model to date, demonstrating a significant improvement in the processing capabilities for long sequences.

Motivation & Objective

• Motivate the need to evaluate NLP models on extremely long contexts beyond current capabilities.
• Propose a scalable benchmark (BABILong) for long-context QA with distributed facts.
• Evaluate GPT-4, RAG, and recurrent memory augmented transformers on ultra-long inputs.
• Demonstrate that recurrence with self-retrieval can extend context processing well beyond existing models.

Proposed method

• Introduce BABILong, an ultra-long context benchmark by hiding task sentences among background text from PG19/Wiki.
• Use GPT-4-Turbo (128k window) and Mistral (32k window) to assess baseline LLM performance on growing context sizes.
• Fine-tune GPT-3.5 on 16k-token tasks and evaluate its performance with context expansion.
• Implement Recurrent Memory Transformer (RMT) and RMT with self-retrieval (RMT-R) to process long sequences with linear scaling.
• Augment RMT/RMT-R with retrieval of past memory states to emulate attention-like access over past segments.
• Experiment with segmenting inputs into 512-token chunks and maintaining memory tokens across segments; analyze attention/memory usage.
• Compare retrieval-augmented generation (RAG) with GPT-4 alongside memory-augmented approaches, using FAISS/LangChain for embeddings (text-embedding-ada-002).

Experimental results

Research questions

• RQ1 Can current LLMs solve needle-in-a-haystack tasks as context length grows to 10^4–10^5 tokens and beyond?
• RQ2 Do retrieval-augmented and memory-augmented architectures sustain performance on ultra-long contexts where standard attention becomes inefficient?
• RQ3 How do RMT and RMT-R scale with increasing sequence length, and how do they compare to GPT-4 and RAG on long-context QA?
• RQ4 What is the impact of memory-based retrieval on multi-hop reasoning over very long documents?
• RQ5 Is there a regime where recurrent memory approaches outperform large LLMs on extremely long sequences (millions of tokens)?

Key findings

• GPT-4 and Mistral performance degrades as context length increases, even with large windows.
• RMT and RMT-R substantially outperform GPT-4 on long sequences and can handle up to 11 million tokens in experiments.
• RMT-R (with memory state retrieval) yields further improvements due to retrieved past states, maintaining performance across very long inputs.
• Retrieval augmentation (RAG) shows limited or task-dependent benefits for the BABILong tasks and can be less effective than memory-based methods for certain long-context settings.
• The models trained with memory mechanisms demonstrate persistent performance and generalize to sequence lengths well beyond their training horizon (up to 128k tokens and beyond, with RMT-R keeping pace up to 10M+ tokens).