[论文解读] 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 Research被引用 7

一句话总结

本文介绍了 BABILong 基准，用于超长文档，并展示一种循环记忆 transformer（RMT/RMT-R）能够处理高达 11 million tokens，在长上下文任务中优于 GPT-4 和 RAG。

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.

研究动机与目标

• 促使评估 NLP 模型在超出当前能力范围的极长上下文上的必要性。
• 提出一个可扩展的基准（BABILong）用于具有分布事实的长上下文问答。
• 评估 GPT-4、RAG 与具备循环记忆的变换器在超长输入上的表现。
• 证明带自检索的循环性可以将上下文处理扩展到远超现有模型的水平。

提出的方法

• 引入 BABILong，这是一个极长上下文基准，通过将任务句隐藏在来自 PG19/Wiki 的背景文本中。
• 使用 GPT-4-Turbo（128k 窗口）和 Mistral（32k 窗口）来评估基线 LLM 在不断增长的上下文规模上的表现。
• 在 16k-token 任务上微调 GPT-3.5，并评估其在扩展上下文后的表现。
• 实现 Recurrent Memory Transformer (RMT) 与带自检索的 RMT-R，以线性扩展处理长序列。
• 为 RMT/RMT-R 增强过去记忆状态的检索，以模拟对过去片段的类注意力访问。
• 尝试将输入分割为 512-token 的块，并在各块之间维持记忆标记；分析注意力/记忆的使用。
• 将检索增强生成（RAG）与 GPT-4 一起，与内存增强方法比较，使用 FAISS/LangChain 进行嵌入（text-embedding-ada-002）。

实验结果

研究问题

• RQ1 在上下文长度增长到 10^4–10^5 tokens 及更高时，当前的 LLM 能否解决 needle-in-a-haystack 任务？
• RQ2 在标准注意力变得低效的超长上下文中，检索增强和记忆增强的架构是否能维持性能？
• RQ3 RMT 与 RMT-R 如何随序列长度增加而扩展，与 GPT-4 及 RAG 在长上下文问答中的表现相比如何？
• RQ4 基于记忆的检索对极长文档的多跳推理有何影响？
• RQ5 是否存在一个阶段，循环记忆方法在极长序列（数百万个 token）上超越大型 LLM？

主要发现

• GPT-4 与 Mistral 的性能随着上下文长度增加而下降，即使窗宽很大。
• RMT 和 RMT-R 在长序列上显著优于 GPT-4，并且在实验中能够处理高达 11 million tokens。
• RMT-R（带记忆状态检索）由于检索到的过去状态而带来进一步提升，在极长输入中保持了性能。
• 检索增强（RAG）在 BABILong 任务中显示出有限或任务依赖的收益，在某些长上下文设置下可能不如基于记忆的方法有效。
• 具有记忆机制的模型表现出持久的性能并能泛化到远超其训练时长的序列长度（高达 128k tokens 及以上，且 RMT-R 与之同步，达到 10M+ tokens）。