论文题目:InfLLM: Unveiling the Intrinsic Capacity of LLMs for Understanding Extremely Long Sequences with Training-Free Memory
论文链接:https://arxiv.org/pdf/2402.04617.pdf
InfLLM旨在揭示大型语言模型(LLMs)在处理极长序列时的内在能力,而无需额外训练。现有的LLMs在处理长序列时,由于预训练时序列长度受限,无法泛化到更长的序列,这导致了领域外和干扰问题。为了解决这些问题,现有的方法采用了滑动注意力窗口和丢弃远距离标记来处理极长序列,但这些方法无法深入理解序列中的长距离依赖关系。
InfLLM通过 将远距离上下文存储到额外的内存单元 中,并采 用 高效的机制来查找与当前处理的标记相关的单元进行注意力计算 ,从而允许LLMs在保持捕捉长距离依赖关系的同时高效处理长序列。 InfLLM不需要任何训练 ,就能使在几千个标记序列上预训练的LLMs在性能上超过那些持续在长序列上训练的LLMs。即使在序列长度扩展到 1024K 时,InfLLM仍然能够有效地捕捉长距离依赖关系。
InfLLM的示意图
InfLLM几个 关键组件 :
- 滑动窗口注意力机制:为了处理长序列,InfLLM采用了滑动窗口注意力机制。这意味着在每个计算步骤中,模型只关注输入序列的一部分(称为局部窗口),而不是整个序列。这样可以避免超出模型处理能力的长序列问题。
- 额外的上下文内存单元:为了解决滑动窗口机制无法捕捉长距离依赖的问题,InfLLM构建了一个额外的上下文内存模块。这个内存模块存储了被滑动窗口机制丢弃的远距离上下文信息。
- 块级上下文内存单元:为了提高内存查找的效率,InfLLM将过去的键值向量(key-value vectors)分成多个块(block-level memory units),每个块包含一定数量的标记。这样可以减少内存单元的数量,同时保持性能。
- 代表性标记(Representative Tokens):在每个内存单元中,InfLLM选择几个代表性标记(representative tokens),这些标记在局部窗口内具有最高的注意力分数,代表了整个单元的语义。这有助于在计算相关性分数时避免不重要标记的干扰。
- 内存查找和注意力计算:在每个计算步骤中,InfLLM将当前的局部窗口、从上下文内存中选择的相关内存单元以及当前的标记隐藏向量(current tokens hidden vectors)合并,形成当前的键值缓存(current key-value cache)。然后,使用这个缓存来计算注意力输出。
- 缓存管理:为了高效地处理极长序列,InfLLM采用了一种缓存管理策略。它将大部分内存单元存储在CPU内存中,只将当前步骤所需的代表性标记和内存单元保留在GPU内存中。这样可以在有限的GPU内存下有效地编码极长序列。
基于Mistral-7B-inst-v0.2与Vicuna-7B-v1.5在两个广泛使用的长文档基准测试(∞-Bench和LongBench)上进行了评估,证明了InfLLM的有效性。
能够处理极长流式输入的模型用♣表示。持续训练的模型用∗表示。中文本长度的95%分位数是214K和31K。InfLLM的上下文窗口大小表示为“局部窗口大小 + 选择的内存大小”
此外,论文还探讨了InfLLM与检索增强生成(RAG)的联系: RAG更适合于知识驱动的任务,以获取额外的外部知识。值得注意的是,我们的方法可以直接与RAG结合,以推进知识驱动的任务。RAG用于从互联网检索外部知识,检索到的多个文档被拼接成长输入,供InfLLM生成高质量的答案, 并指出InfLLM在无需额外训练和更广泛的适用性方面具有优势。
最后,论文提出了未来的研究方向,包括如何有效地训练上下文内存模块以及如何结合关键值缓存压缩方法来降低内存管理的计算和内存成本。