MemLong: Memory-Augmented Retrieval for Long Text Modeling

📄 arXiv: 2408.16967v1 📥 PDF

作者: Weijie Liu, Zecheng Tang, Juntao Li, Kehai Chen, Min Zhang

分类: cs.CL, cs.AI

发布日期: 2024-08-30

🔗 代码/项目: GITHUB

💡 一句话要点

MemLong：提出一种基于记忆增强检索的长文本建模方法，显著扩展LLM的上下文处理能力。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 长文本建模 记忆增强检索 大型语言模型 上下文扩展 外部检索器 可控注意力 低资源需求

📋 核心要点

  1. 现有LLM在处理长文本时面临注意力机制的计算复杂度高和内存消耗大的挑战。
  2. MemLong通过引入外部检索器，结合不可微的检索模块和可训练的解码器，实现记忆增强检索。
  3. 实验表明，MemLong在长文本建模任务上优于现有模型，并显著扩展了上下文处理长度。

📝 摘要（中文）

大型语言模型（LLMs）在各个领域取得了显著的成功。然而，由于注意力机制的二次时间和空间复杂度以及生成过程中键值缓存不断增长的内存消耗，处理长上下文仍然是LLMs面临的重大挑战。本文介绍了MemLong：用于长文本生成的记忆增强检索，该方法旨在通过利用外部检索器进行历史信息检索来增强长上下文语言建模的能力。MemLong结合了一个不可微的“ret-mem”模块和一个部分可训练的仅解码器语言模型，并引入了一种细粒度的、可控的检索注意力机制，该机制利用语义级别的相关块。在多个长上下文语言建模基准上的综合评估表明，MemLong始终优于其他最先进的LLMs。更重要的是，MemLong可以将单个3090 GPU上的上下文长度从4k扩展到80k。

🔬 方法详解

问题定义：现有大型语言模型在处理长文本时，由于自注意力机制的二次方复杂度，计算和内存消耗巨大，限制了其在长文档理解和生成等任务中的应用。此外，生成过程中不断增长的键值缓存也加剧了内存压力。

核心思路：MemLong的核心思路是利用外部检索器来选择性地检索与当前生成内容相关的历史信息，从而避免对整个上下文进行注意力计算。通过这种方式，模型可以专注于更相关的上下文片段，降低计算复杂度，并显著减少内存消耗。这种设计借鉴了人类在处理长篇内容时会回顾和参考相关信息片段的习惯。

技术框架：MemLong的整体框架包含三个主要模块：1) 外部检索器：负责从历史文本中检索相关片段；2) ret-mem模块：一个不可微的模块，用于将检索到的信息整合到模型的输入中；3) 部分可训练的解码器语言模型：利用检索到的信息生成文本。整个流程是，给定输入文本，检索器首先检索相关片段，然后ret-mem模块将这些片段与当前输入组合，最后解码器利用组合后的信息生成输出。

关键创新：MemLong的关键创新在于其细粒度的、可控的检索注意力机制。该机制允许模型在语义层面上选择相关的文本块，而不是简单地使用固定大小的窗口。此外，ret-mem模块的不可微设计允许使用预训练的检索器，而无需端到端训练，从而降低了训练成本。

关键设计：MemLong使用预训练的检索模型（例如，基于sentence embeddings的模型）来检索相关的文本块。检索到的文本块通过ret-mem模块与当前输入连接，然后输入到解码器语言模型中。解码器语言模型采用部分可训练的策略，即只训练部分参数，以进一步降低训练成本。损失函数通常是标准的语言模型损失，例如交叉熵损失。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

MemLong在多个长文本建模基准测试中取得了显著的性能提升，超越了现有的最先进模型。更重要的是，MemLong能够将单个3090 GPU上的上下文长度从4k扩展到80k，这表明其在处理超长文本方面具有巨大的潜力。代码已开源。

🎯 应用场景

MemLong具有广泛的应用前景，包括长文档摘要、长篇小说生成、代码生成、知识库问答等。通过扩展LLM的上下文处理能力，MemLong可以更好地理解和生成复杂的长文本内容，从而提升相关应用的性能和用户体验。此外，MemLong的低资源需求使其更易于部署在资源受限的环境中。

📄 摘要（原文）

Recent advancements in Large Language Models (LLMs) have yielded remarkable success across diverse fields. However, handling long contexts remains a significant challenge for LLMs due to the quadratic time and space complexity of attention mechanisms and the growing memory consumption of the key-value cache during generation. This work introduces MemLong: Memory-Augmented Retrieval for Long Text Generation, a method designed to enhance the capabilities of long-context language modeling by utilizing an external retriever for historical information retrieval. MemLong combines a non-differentiable ``ret-mem'' module with a partially trainable decoder-only language model and introduces a fine-grained, controllable retrieval attention mechanism that leverages semantic-level relevant chunks. Comprehensive evaluations on multiple long-context language modeling benchmarks demonstrate that MemLong consistently outperforms other state-of-the-art LLMs. More importantly, MemLong can extend the context length on a single 3090 GPU from 4k up to 80k. Our code is available at https://github.com/Bui1dMySea/MemLong