地址

论文地址：https://arxiv.org/pdf/2408.16967

时间

2024 年8月30

作者

苏州大学、哈尔滨工业大学

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 this https URL

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

我们的代码可在该链接获取。

Large Language Models (LLMs) have achieved re-markable success in various fields. However, due to the quadratic time and space complexity of vanilla attention mechanisms (Vaswani et al., 2017), it is challenging to extend the context length consider-ably, which poses significant limitations for applica-tions involving long-sequence tasks, such as long-document summarization (Koh et al., 2022) and multiple rounds of dialogue (Wang et al., 2024a). As a result, LLMs are often expected to maintain a long working capability (a.k.a. long context LLMs) to effectively handle these demanding scenarios.

To tackle the computational bottleneck, numer-ous efforts have been made. The first line of work focuses on reducing the computation of vanilla at-tention mechanisms (Vaswani et al., 2017) by em-ploying sparse attention operations (Beltagy et al., 2020; Wang et al., 2020; Kitaev et al., 2020; Xiao et al., 2023a; Chen et al., 2023b; Lu et al., 2024). Although these types of works can reduce com-putational complexity to approximately O(n), it often comes with trade-offs in model capability. Therefore, Some works shift their focus to mem-ory selection (Dai et al., 2019; Bertsch et al., 2024; Yu et al., 2023). These approaches, as token-level memory selection, can result in the truncation of se-mantic information. Another recent line of work is Retrieval-Augment Language Modeling (Wu et al., 2022; Wang et al., 2024b; Rubin and Berant, 2023). These works usually introduce a retrieval mecha-nism to enhance the model’s ability to handle long texts. However, these methods have several draw-backs. Firstly, the information stored in memory may experience distribution shifts due to changes in model parameters during training. Secondly, these methods often require retraining, which is im-practical in the era of large models. Finally, these models are often prone to processing long text in-puts at the expense of the original capabilities of the pre-trained model. To address the limitations of previous research, we posed the following ques-tion: Can we utilize the explicit retrieval capa-bilities of a retriever to approximate the implicit retrieval processes within the model?

大型语言模型（LLMs）在各种领域取得了显著的成功。然而，由于基础注意力机制（Vaswani等人, 2017）的时间和空间复杂度为二次方，很难大幅延长上下文长度，这给涉及长序列任务的应用带来了重大限制，如长文档摘要（Koh等人, 2022）和多轮对话（Wang等人, 2024a）。因此，人们期望LLMs能够保持较长的工作能力（即所谓的长上下文LLMs），以有效应对这些需求场景。

为了克服计算瓶颈，已经做出了许多努力。第一类工作集中在通过使用稀疏注意力操作（Beltagy等人, 2020; Wang等人, 2020; Kitaev等人, 2020; Xiao等人, 2023a; Chen等人, 2023b; Lu等人, 2024）减少基础注意力机制（Vaswani等人, 2017）的计算量。尽管这类工作可以将计算复杂度降低至大约O(n)，但它通常伴随着模型能力的折衷。因此，有些工作转向了记忆选择（Dai等人, 2019; Bertsch等人, 2024; Yu等人, 2023）。作为词元级记忆选择的方法可能会导致语义信息被截断。另一条近期的研究方向是检索增强的语言建模（Wu等人, 2022; Wang等人, 2024b; Rubin和Berant, 2023）。这些工作通常会引入一种检索机制来增强模型处理长文本的能力。但是，这些方法有几个缺点。首先，在训练过程中模型参数的变化可能导致存储在记忆中的信息经历分布偏移。其次，这些方法往往需要重新训练，在大模型时代这是不切实际的。最后，这些模型常常以牺牲预训练模型的原始能力为代价来处理长文本输入。为了解决先前研究的局限性，我们提出了以下问题：我们能否利用检索器的显式检索能力来近似模型内部的隐式检索过程？

In this work, we propose MemLong, an efficient and lightweight method to extending the context window of LLMs. The key idea is to store past contexts and knowledge in a non-trainable mem-ory bank and further leverages these stored em-beddings to retrieve chunk-level key-value (K-V) pairs for input into the model.. MemLong is ap-plicable to any decoder-only pretrained language models by incorporating (1) an additional ret-mem component for memory and retrieval, and (2) a retrieval causal attention module for integrating local and memory information. The memory and retrieval process of MemLong is illustrated in Fig-ure 1(b). During generation,one text that exceeds the model’s maximum processing length is stored as context information in a Memory Bank. Sub-sequently, given a recently generated text chunk in a long document, we use the retriever to explic-itly retrieve past information, obtaining additional context information through index alignment.

在这项工作中，我们提出了一种高效且轻量级的方法MemLong，用以扩展LLMs的上下文窗口。其核心思想是在非训练的记忆库中存储过去的上下文和知识，并进一步利用这些存储的嵌入来检索块级键值（K-V）对以供模型输入。通过添加（1）额外的ret-mem组件用于记忆和检索，以及（2）一个检索因果注意模块来整合局部和记忆信息，MemLong适用于任何仅解码器的预训练语言模型。图1(b)展示了MemLong的记忆和检索过程。在生成过程中，超过模型最大处理长度的文本被作为上下文信息存储在记忆库中。随后，对于长文档中最近生成的文本段落，我们使用检索器明确地检索过去的信息，通过索引对齐获得额外的上下文信息。

MemLong offers several benefits: (1) Distribu-tional Consistency: Unlike previous models that experienced a distribution shift when information was stored in memory, MemLong ensures the dis-tribution of cached information remains consistent.(2) Training Efficient: We freeze the lower layers of the model and only need to finetune the upper layers which greatly reduced computational cost. In our experiments, finetuning a 3B parameter ver-sion of MemLong on 0.5B tokens requires only eight 3090 GPUs for eight hours. (3) Extensive Context Window: Since only a single layer’s K-V pairs need to be memorized, MemLong is capable of extending the context window up to 80k tokens easily on a single 3090 GPU.

Extensive experiments have demonstrated that MemLong exhibits superior performance in several aspects when compared with other leading LLMs. MemLong outperforms OpenLLaMA (Touvron et al., 2023) and other retrieval-based models on several long-context language modeling datasets. In retrieval-augmented in-context learning tasks, MemLong achieves an improvement of up to 10.2 percentage points over OpenLLaMA.

MemLong提供了几个优点：(1) 分布一致性：与以前在信息存储于记忆时经历分布偏移的模型不同，MemLong确保了缓存信息的分布保持一致。(2) 训练效率：我们冻结了模型的下层并只需要微调上层，这大大降低了计算成本。在实验中，只需八个3090 GPU八小时即可完成对具有30亿参数版本的MemLong进行0.5亿个标记的微调。(3) 广泛的上下文窗口：因为只需要记住单一层的K-V对，MemLong能够在单个3090 GPU上轻松将上下文窗口扩展至80k个标记。

广泛的实验表明，与其他领先的LLMs相比，MemLong在多个方面表现出优越的性能。在几个长上下文语言建模数据集上，MemLong的表现超过了OpenLLaMA（Touvron等人, 2023）和其他基于检索的模型。在检索增强的上下文学习任务中，MemLong相对于OpenLLaMA实现了高达10.2个百分点的改进。

We introduce MemLong, an innovative approach that significantly enhances the capability of lan-guage models to process long texts by leveraging an external retriever. MemLong utilizes a profi-cient retriever to swiftly and accurately access text relevant to the distant context with minimal mem-ory overhead. MemLong successfully expands the model’s context window from 2k to 80k tokens. We demonstrate that MemLong exhibits consider-able competitive advantages in long-distance text modeling and comprehension tasks. MemLong can achieve up to a 10.4 percentage point improvement in performance compared to the full-context model.

我们介绍了一种创新的方法MemLong，它通过利用外部检索器显著增强了语言模型处理长文本的能力。MemLong使用高效的检索器快速准确地访问与远距离上下文相关的文本，同时最小化内存开销。MemLong成功地将模型的上下文窗口从2k扩展到了80k个标记。我们证明了MemLong在长距离文本建模和理解任务中展现出相当大的竞争优势。与全上下文模型相比，MemLong的性能最多可以提高10.4个百分点。