InfLLM: Training-Free Long-Context Extrapolation for LLMs with an Efficient Context Memory
TL;DR 精炼摘要
本文介绍了InfLLM,一种无需训练即可有效处理长序列的记忆机制。通过将远程上下文存储于额外记忆单元,InfLLM允许大型语言模型(LLMs)在有限上下文窗口内高效捕捉长距离依赖。研究发现,InfLLM的表现可与在长序列上进行持续训练的基线相媲美,其优势在于无需昂贵的微调过程。
摘要
Large language models (LLMs) have emerged as a cornerstone in real-world applications with lengthy streaming inputs (e.g., LLM-driven agents). However, existing LLMs, pre-trained on sequences with a restricted maximum length, cannot process longer sequences due to the out-of-domain and distraction issues. Common solutions often involve continual pre-training on longer sequences, which will introduce expensive computational overhead and uncontrollable change in model capabilities. In this paper, we unveil the intrinsic capacity of LLMs for understanding extremely long sequences without any fine-tuning. To this end, we introduce a training-free memory-based method, InfLLM. Specifically, InfLLM stores distant contexts into additional memory units and employs an efficient mechanism to lookup token-relevant units for attention computation. Thereby, InfLLM allows LLMs to efficiently process long sequences with a limited context window and well capture long-distance dependencies. Without any training, InfLLM enables LLMs that are pre-trained on sequences consisting of a few thousand tokens to achieve comparable performance with competitive baselines that continually train these LLMs on long sequences. Even when the sequence length is scaled to K, InfLLM still effectively captures long-distance dependencies. Our code can be found in \url{https://github.com/thunlp/InfLLM}.
思维导图
论文精读
中文精读
1. 论文基本信息
1.1. 标题
InfLLM: Training-Free Long-Context Extrapolation for LLMs with an Efficient Context Memory
1.2. 作者
Chaojun Xiao*, Pengle Zhang*, Xu Han†, Guangxuan Xiao, Yankai Lin, Zhengyan Zhang, Zhiyuan Liu†, Maosong Sun (*共同第一作者,†通讯作者)
1.3. 发表期刊/会议
arXiv。arXiv 是一个预印本(preprint)服务器,论文在此发布通常在正式同行评审和发表之前。它在学术界,尤其是在人工智能和计算机科学领域,是一个非常重要的信息交流平台,允许研究人员快速分享最新成果。
1.4. 发表年份
2024
1.5. 摘要
大型语言模型(LLMs)在处理长篇流式输入(例如,LLM 驱动的智能体(agents))的实际应用中扮演着基石作用。然而,由于现有 LLMs 是在受限的最大长度序列上预训练的,它们在处理更长序列时会面临域外(out-of-domain)和注意力分散(distraction)问题。常见的解决方案通常涉及在更长序列上进行持续预训练(continual pre-training),但这会带来昂贵的计算开销和模型能力不可控的变化。
本文揭示了 LLMs 在无需任何微调(fine-tuning)的情况下理解极长序列的内在能力。为此,我们引入了一种无需训练(training-free)的基于记忆(memory-based)的方法,名为 InfLLM。具体来说,InfLLM 将远距离上下文存储到额外的记忆单元(memory units)中,并采用一种高效的机制来查找与当前词元(token)相关的单元进行注意力计算。通过这种方式,InfLLM 允许 LLMs 在有限的上下文窗口(context window)内高效处理长序列,并有效捕获长距离依赖(long-distance dependencies)。
在无需任何训练的情况下,InfLLM 使那些在包含数千个词元序列上预训练的 LLMs 达到了与那些在长序列上持续训练的竞争基线(competitive baselines)相当的性能。即使序列长度扩展到 1,024K,InfLLM 仍能有效地捕获长距离依赖。我们的代码可以在 https://github.com/thunlp/InfLLM 找到。
1.6. 原文链接
原文链接: https://arxiv.org/abs/2402.04617
PDF 链接: https://arxiv.org/pdf/2402.04617v2.pdf
发布状态:预印本(preprint)
2. 整体概括
2.1. 研究背景与动机
当前的大型语言模型(LLMs)在多种任务中取得了显著成就,并被广泛应用于各种需要处理长篇流式输入的场景,例如 LLM 驱动的智能体(agents)。然而,这些 LLMs 通常在相对较短的序列上进行预训练(例如,几千个词元),这导致它们在处理远超训练长度的序列时面临严峻挑战。
主要的挑战包括:
-
域外问题 (Out-of-domain Issue):
LLMs未在长序列上进行训练,导致其在遇到超出其训练范围的上下文长度时,性能会急剧下降。这通常与位置编码(positional encoding)机制的限制有关。 -
注意力分散问题 (Distraction Issue): 长序列中往往包含大量噪声或不相关的上下文信息,这些信息会分散
LLM的注意力,使其难以聚焦于关键信息,从而影响其理解和推理能力。为了解决这些问题,现有研究提出了多种方案:
-
持续预训练 (Continual Pre-training): 在更长的序列上对模型进行额外的训练。动机:通过暴露模型于更长的序列来直接提升其处理能力。存在问题:计算成本极其高昂,需要大规模高质量的长序列数据集,且可能导致模型在短上下文上的性能下降(即“灾难性遗忘”或
catastrophic forgetting)。 -
修改位置编码 (Modifying Positional Encoding): 调整或扩展现有的位置编码机制(如
Rotary Position Embedding, RoPE),以支持更长的序列。动机:解决out-of-domain问题。存在问题:通常无法解决distraction问题,并且仍然可能需要一定的微调或导致性能折衷。 -
滑动窗口注意力 (Sliding Window Attention): 仅关注当前词元周围的局部上下文。动机:降低计算复杂度,处理流式输入。存在问题:直接丢弃了所有远距离的上下文信息,无法捕获长距离依赖。
鉴于上述挑战和现有解决方案的局限性,本文的动机在于寻找一种无需额外训练、计算高效且能有效处理长距离依赖的解决方案,从而使在短序列上训练的
LLMs能够直接应用于极长序列。
2.2. 核心贡献/主要发现
本文提出了 InfLLM 方法,其核心贡献和主要发现如下:
- 揭示
LLMs的内在长序列理解能力 (Unveiling Intrinsic Capacity): 论文的一个关键发现是,LLMs本身具备在无需微调的情况下理解极长序列的内在能力。InfLLM通过提供相关的长距离上下文,有效地利用了这种能力。 - 提出训练无关的基于记忆的方法
InfLLM(Training-Free Memory-Based Method):InfLLM是一种无需任何额外训练即可扩展LLMs上下文窗口的方法。这显著降低了计算开销,避免了持续训练可能引入的性能下降和模型能力变化。 - 结合滑动窗口注意力与高效上下文记忆 (Efficient Context Memory with Sliding Window Attention):
InfLLM在sliding window attention的基础上,额外构建了一个高效的context memory模块。这使得模型既能处理流式长输入(通过滑动窗口),又能通过记忆模块捕获被滑动窗口忽略的远距离依赖。 - 创新性块级记忆单元机制 (Novel Block-Level Memory Units): 为了解决大规模噪声上下文和内存查找效率问题,
InfLLM将历史key-value向量组织成块(blocks),每个块作为一个记忆单元。并通过选择块内最具代表性的词元(representative tokens)作为单元表示,来高效计算相关性并进行查找。- 有效查找 (Effective Lookup): 块级语义一致性提高了信息检索的准确性。
- 高效查找 (Efficient Lookup): 避免了逐词元计算,减少了计算开销,并通过连续内存访问提高了效率。
- 高效缓存管理与卸载机制 (Efficient Cache Management and Offloading):
InfLLM采用将不常用记忆单元卸载到CPU memory、只在GPU memory中保留常用单元的策略,并结合LRU(最近最少使用)缓存策略,大大降低了GPU memory占用,实现了在有限资源下处理极长序列的能力。 - 在
∞-Bench和LongBench上的卓越性能 (Superior Performance on Benchmarks):- 在无需任何训练的情况下,
InfLLM使Mistral-7B和Llama-3-8B等基础模型在处理几千个词元的训练长度时,达到了与在长序列上持续训练的竞争基线相当甚至更优的性能。 InfLLM成功将Llama-3的上下文长度从8K扩展到128K,甚至在1,024K的超长序列上也能有效捕获长距离依赖,并在Retrieve.Passkey任务上保持100%的准确率。
- 在无需任何训练的情况下,
- 效率优势 (Efficiency Advantage): 相较于持续训练模型,
InfLLM在时间消耗上降低了34%,且GPU memory使用量仅为34%。
3. 预备知识与相关工作
3.1. 基础概念
-
大型语言模型 (Large Language Models, LLMs): 这是一类基于神经网络的模型,通常采用
Transformer架构,通过在海量文本数据上进行预训练来学习语言的统计规律和知识。它们能够理解、生成人类语言,并执行如问答、摘要、翻译等多种自然语言处理任务。LLMs的核心在于其能够捕捉长距离依赖关系,并根据上下文生成连贯且有意义的文本。 -
上下文窗口 (Context Window): 指
LLMs在生成下一个词元(token)时可以“看到”或考虑的输入序列的最大长度。在训练阶段,模型通常被限制在一个固定的上下文窗口内进行学习。如果输入序列长度超出此窗口,超出部分的信息就无法被模型直接利用。这导致了LLMs在处理超长文本时性能下降的“域外问题(out-of-domain)”。 -
词元 (Token): 在自然语言处理中,词元是文本的最小有意义单元。它可以是一个单词、一个子词(
subword)、一个字符,甚至是标点符号。LLMs的输入和输出都是词元序列。 -
位置编码 (Positional Encoding):
Transformer架构本身是并行处理的,不包含序列的顺序信息。为了让模型能够理解词元在序列中的相对或绝对位置,需要引入位置编码。常见的位置编码包括绝对位置编码(如正弦波位置编码)和相对位置编码(如Rotary Position Embedding, RoPE)。RoPE是目前许多流行LLMs(如Llama、Mistral)使用的位置编码,它通过旋转矩阵的方式将位置信息融入到query和key向量中,使其能够更好地处理相对位置关系。 -
自注意力机制 (Self-Attention Mechanism):
Transformer架构的核心组件,允许模型在处理序列中的每个词元时,都能关注到序列中的其他所有词元,并根据它们的重要性分配不同的权重。其计算涉及三个核心向量:- 查询向量 (Query, Q): 代表当前词元在寻找相关信息时发出的“查询”。
- 键向量 (Key, K): 代表序列中所有词元的“索引”或“特征”,用于响应查询。
- 值向量 (Value, V): 包含序列中所有词元的实际信息,根据注意力权重进行加权求和后输出。
基本的
Self-Attention计算公式如下: 其中,、、 分别是query、key、value矩阵, 是key向量的维度,用于缩放点积以防止梯度过小。softmax函数将注意力分数归一化为概率分布。
-
Key-Value Cache (KV Cache): 在
LLMs的推理(inference)阶段,尤其是在自回归(autoregressive)生成时,为了避免重复计算已经处理过的词元对应的Key和Value向量,这些向量会被缓存起来。当生成下一个词元时,只需要计算新词元的Key和Value,然后与KV Cache中已有的向量拼接起来进行注意力计算,这大大提高了生成效率。
3.2. 前人工作
论文将现有工作分为三大类:上下文长度外推(Context Length Extrapolation)、高效上下文计算(Efficient Context Computation)和基于记忆的模型(Memory-based Models)。
3.2.1. 上下文长度外推 (Context Length Extrapolation)
这类方法旨在使在短序列上训练的 LLMs 能够处理更长的序列,而无需额外的预训练。
- 早期方法: 设计新的相对位置编码机制(如
Press et al., 2022; Sun et al., 2023),在预训练阶段就考虑长序列。 - 基于
RoPE的方法: 针对广泛使用的RoPE(Su et al., 2021),通过降尺度(downscaling)或重用原始位置索引(Chen et al., 2023b; Peng et al., 2023; Chen et al., 2023a; Jin et al., 2024; An et al., 2024)来实现长度外推。这些方法可以缓解由未见长度引起的“域外问题”,但通常无法解决由噪声上下文引起的“注意力分散问题”。 - 滑动窗口注意力 (Sliding Window Attention): 例如
LM-Infinite(Han et al., 2023)和Streaming-LLM(Xiao et al., 2023)。这些方法通过仅关注局部上下文来流式处理极长序列,直接丢弃所有远距离上下文。它们在处理效率上有优势,但缺点是无法捕获长距离依赖。
3.2.2. 高效上下文计算 (Efficient Context Computation)
这类方法主要关注提高注意力层的计算效率,通常需要修改模型架构并重新训练模型。
- 稀疏注意力 (Sparse Attention): 通过限制每个词元只关注序列中的一小部分词元来降低计算复杂度(
Zaheer et al., 2020; Beltagy et al., 2020; Child et al., 2019; Ainslie et al., 2020; Zhao et al., 2019)。 - 核函数近似 (Kernel Functions Approximation): 使用核函数近似注意力计算,将二次复杂度降低到线性复杂度(
Kitaev et al., 2020; Wang et al., 2020; Katharopoulos et al., 2020)。 - 状态空间模型 (State-Space Models, SSMs): 用具有线性复杂度的状态空间模型替代注意力层(
Gu et al., 2022; Gu & Dao, 2023),如Mamba模型。 - KV 向量淘汰 (KV Eviction): 在推理过程中淘汰无用的
key-value向量以减少计算量(Zhang et al., 2023b; Li et al., 2024; Ge et al., 2023)。这些方法可以提高效率,但由于未见位置编码导致的“域外问题”,通常不能在不额外训练的情况下外推上下文窗口。
3.2.3. 基于记忆的模型 (Memory-based Models)
这类方法通过引入外部记忆模块来增强模型的知识存储和长距离信息处理能力。
- 早期记忆网络 (Memory Networks): 如
Graves et al., 2014; Weston et al., 2015; Sukhbaatar et al., 2015; Miller et al., 2016,通过可寻址的记忆模块为模型提供额外知识。 - 循环
Transformer记忆层 (Recurrent Transformer Memory Layers): 如Dai et al., 2019; Rae et al., 2020; Khandelwal et al., 2020; Wu et al., 2022; Bertsch et al., 2023; Munkhdalai et al., 2024。这些工作将序列分成段,单独编码每段,并使用记忆来存储前一段的上下文信息,以递归地处理长序列。 - 与
InfLLM的区别: 尽管概念相似,但上述记忆模型通常需要修改模型架构并进行额外的训练。相比之下,InfLLM专注于探索LLMs的固有特性,提出了一个无需训练的记忆模块,用于长文本理解。
3.3. 技术演进
LLMs 的长上下文处理能力经历了从固定短上下文到尝试扩展、再到更精细化管理的演进:
- 初始阶段:
Transformer模型由于其二次计算复杂度,通常在几千个词元(例如512到4096)的固定上下文窗口内进行训练。超过这个长度,计算和内存成本都难以承受。 - 早期扩展尝试: 研究者尝试通过修改位置编码机制(如
RoPE的缩放和插值方法)来“欺骗”模型,使其能够泛化到更长的位置。同时,稀疏注意力等高效注意力机制开始被提出,以降低计算复杂度,允许在预训练阶段处理更长的序列。 - 流式处理与记忆引入: 随着
LLMs应用场景对实时、流式输入的需求增加,sliding window attention等机制被提出,以固定有限的窗口处理无限长的流。然而,这种方式丢失了长距离信息。为了弥补这一缺陷,将外部记忆引入Transformer成为一个自然的方向。早期的记忆增强Transformer模型通常需要专门设计和训练。 InfLLM的贡献:InfLLM代表了这一演进中的一个重要步骤,它结合了sliding window的效率和外部记忆捕获长距离依赖的能力,但最关键的是,它以无需训练的方式实现了这一点,从而降低了应用门槛,并揭示了现有LLMs潜在的长上下文处理能力。
3.4. 差异化分析
InfLLM 与相关工作的主要区别和创新点在于:
- 与持续预训练模型 (
Llama-1M等) 相比:InfLLM无需任何额外的训练,避免了昂贵的计算成本、数据需求和对模型原有能力的潜在损害。它通过利用LLMs的内在能力实现长上下文处理,而持续预训练模型则通过修改模型本身并进行再训练来适应长上下文。 - 与位置编码外推方法 (
NTK,SelfExtend) 相比:InfLLM不仅解决了out-of-domain问题,更通过context memory有效解决了长序列中的distraction问题,选择性地引入相关上下文,避免噪声。位置编码外推方法虽然能扩展长度,但无法有效过滤噪声。 - 与滑动窗口模型 (
LM-Infinite,StreamingLLM) 相比:InfLLM在使用sliding window attention的同时,通过context memory弥补了滑动窗口丢失远距离关键信息的缺点,使其能够捕获长距离依赖。滑动窗口模型则直接丢弃了远距离上下文。 - 与高效上下文计算方法 (稀疏注意力、
SSMs等) 相比: 这些方法通常需要修改模型架构并重新训练,而InfLLM是一个训练无关的即插即用模块,可以直接应用于现有预训练LLMs。 - 与传统基于记忆的模型 (
Transformer-XL等) 相比: 传统记忆模型通常需要修改架构并进行专门训练。InfLLM的创新之处在于其“无需训练”的记忆模块设计,直接利用现有LLM的注意力能力来构建和查找记忆,使其更具通用性和易用性。 - 与检索增强生成 (
RAG) 相比:InfLLM是一个内部增强机制,直接利用LLM自身的注意力能力来检索和整合上下文,无需额外的检索器训练和微调LLM来适应检索结果的输入格式。RAG则依赖外部检索模块,需要额外的训练数据和模型,且其性能受检索器影响,可能存在out-of-distribution问题。
4. 方法论
4.1. 方法原理
InfLLM 的核心思想是,在处理极长序列时,LLMs 仅需要关注当前词元的局部上下文,并从远距离的、已被处理过的历史信息中,高效地检索出与当前任务高度相关的上下文信息。这样既能利用 sliding window attention 机制处理流式输入并控制计算量,又能通过检索到的相关记忆单元捕获长距离依赖,从而克服 out-of-domain 和 distraction issues。
具体而言,它构建了一个无需训练的上下文记忆(context memory)模块。当模型处理长序列时,一部分历史 key-value 向量会被“驱逐”并存储到这个记忆中。在每个计算步骤,InfLLM 会智能地从这个记忆中查找并加载与当前词元最相关的记忆单元,然后将这些记忆单元与当前局部上下文拼接起来,形成扩展的上下文窗口,供 LLM 进行注意力计算。
4.2. 核心方法详解
InfLLM 的整体框架和具体组件设计如下:
4.2.1. 整体框架 (Overall Framework)
为了处理远超 LLM 训练长度的输入序列 ,InfLLM 采取了分块(chunk-by-chunk)编码和逐词元(token-by-token)生成的方式。
-
序列分段: 对于每个计算步骤,输入由过去
key-value向量 和当前词元 组成。在编码阶段, 等于块大小(chunk size);在解码阶段, 等于 1。 -
过去上下文的划分: 过去
key-value向量 根据其与当前词元的距离,被划分为三个部分:- 初始词元 (Initial Tokens), : 最早的 个词元,通常包含系统提示(
system prompts)和任务描述,这些信息通常非常重要,因此被保留。 - 驱逐词元 (Evicted Tokens), : 介于初始词元和局部词元之间的词元。这些词元由于距离当前词元较远,通常会被滑动窗口注意力机制忽略。
InfLLM将这部分词元存储到context memory中。 - 局部词元 (Local Tokens), : 距离当前词元最近的 个词元,它们构成了
sliding window attention的核心部分。
- 初始词元 (Initial Tokens), : 最早的 个词元,通常包含系统提示(
-
构建当前 Key-Value 缓存: 在每个计算步骤中,
InfLLM会动态地构建一个用于注意力计算的key-value缓存 。这个缓存由三部分拼接而成:初始词元。- 从
context memory中通过查找操作 获取的相关记忆单元。 局部词元。 因此,当前key-value缓存表示为: 其中, 表示从evicted tokens构成的context memory中,根据当前词元 查找并返回相关的记忆单元。
-
注意力计算: 构建好
key-value缓存 后,LLM会使用它与当前词元的query向量 进行注意力计算,生成输出 :-
, , 是
LLM注意力层中的参数矩阵,用于将输入特征映射为query、key和value向量。 -
是当前需要编码的词元序列。
-
和 是通过 和 矩阵对当前词元 线性变换得到的
key和value向量。 -
和 分别表示由
InfLLM动态构建的key-value缓存 中的key和value向量部分。 -
Concat函数用于将key向量 ( 和 ) 拼接起来,以及将value向量 ( 和 ) 拼接起来,形成完整的key和value序列,供Attn函数使用。 -
Attn函数代表标准的自注意力计算,它会计算 与拼接后的key序列的点积,并通过softmax得到注意力权重,然后与拼接后的value序列进行加权求和,得到最终的注意力输出 。特别说明: 如果查找操作 总是返回空集,那么
InfLLM将退化为LM-Infinite和Streaming-LLM这类直接丢弃远距离上下文的模型。这强调了context memory在捕获长距离依赖中的关键作用。
-
4.2.2. 上下文记忆 (Context Memory)
LLMs 的注意力分数矩阵通常是稀疏的,这意味着只需要保留一小部分 key-value 向量就可以生成相同的输出。受此启发,InfLLM 设计了一个 context memory 来高效地从大量 evicted tokens 中查找相关上下文,并忽略不相关的部分以节省计算成本。
传统的记忆构建方法可能面临内存单元过多、计算开销大、非连续内存访问等问题。InfLLM 采用块级记忆单元(Block-Level Memory Units)来解决这些问题。
4.2.2.1. 块级记忆单元 (Block-Level Memory Units)
InfLLM 将 evicted tokens 划分为多个记忆单元,每个单元包含 个连续词元。为了高效地进行相关性查找,每个块需要一个能够代表其语义内容的单元表示。直接使用整个块的 key-value 向量作为表示会导致计算量过大。因此,InfLLM 采取了选择块内“代表性词元”的方法。
-
代表性词元选择:
InfLLM不引入额外的训练参数,而是利用LLM内部已有的注意力机制来识别块内最重要的词元。对于第 个词元,其代表性分数(representative score) 定义为:- 是第 个词元的
query向量。 - 是第 个词元的
key向量。 - 是局部窗口的大小。
这个公式直观地表示了第 个词元在其对应局部窗口内对其他词元的“影响程度”或“重要性”。分数越高,表示该词元在该局部窗口内越显著。通过计算每个词元的 ,并选择分数最高的 个词元作为该块的代表性词元。
对于一个记忆块 ,其包含 个
key-value对 。通过上述方法,从 中选出 个代表性词元,形成该块的代表性key-value对集合R(\mathbf{B}) = \{ (\mathbf{k}_{b_j}^B, \mathbf{v}_{b_j}^B) \}_{j=1}^{r_k}。这些代表性词元将用于后续的相关性计算。
- 是第 个词元的
-
记忆查找与相关性计算: 在记忆查找阶段,
InfLLM需要从所有存储的记忆单元中选择 个与当前词元最相关的单元加载。为了衡量记忆块 与当前词元 之间的相关性,InfLLM计算它们之间的相似度分数(relevance score):- 是当前需要编码的词元序列(长度为 )。
- 是一个记忆块。
- 是当前词元序列中第 个词元的
query向量。 - 是记忆块 中第 个代表性词元的
key向量。 这个公式通过计算当前词元的query向量与记忆块中所有代表性词元的key向量的点积之和,来衡量它们的语义相关性。选择相关性分数最高的 个记忆单元加载到当前key-value cache中。
4.2.2.2. 位置编码 (Positional Encoding)
现有的 LLM 训练通常使用有限数量的位置编码,导致直接应用于长序列时出现 out-of-domain 问题。此外,InfLLM 的当前 key-value cache 是由不连续的文本块(初始词元、相关记忆单元、局部词元)组成的,直接分配连续的位置编码会造成信息不匹配。
受先前工作启发,InfLLM 采取了一种简化的位置编码策略:对于所有超出局部窗口大小的词元,它们被赋予相同的(固定的)位置编码。
- 具体来说,
context memory单元中的词元与当前词元之间的距离被统一设置为局部窗口大小。 - 解释: 这种策略旨在避免模型因看到训练时未见过的超大位置索引而崩溃。通过将所有远距离的记忆单元都视为与当前词元距离为 ,模型可以在其熟悉的相对位置范围内进行计算。尽管这种处理方式似乎忽略了远距离记忆单元之间的相对位置信息,但论文在附录中通过实验(
Retrieve.Passkey任务)表明,模型仍能通过decoder-only模型的单向特性和key-value隐藏状态的顺序嵌入,间接捕获上下文的相对位置信息。
4.2.2.3. 缓存管理 (Cache Management)
为了在保持高效性的同时处理极长序列流,InfLLM 需要保留所有记忆单元并在每个计算步骤中进行查找。然而,将所有记忆单元都存储在 GPU memory 中是不可行的。考虑到大多数记忆单元的使用频率较低,InfLLM 采用了卸载机制 (Offloading Mechanism):
- 存储位置: 大部分记忆单元存储在
CPU memory中。 - GPU 缓存: 只有在当前步骤中需要使用的以及频繁使用的记忆单元才保留在
GPU memory中。 - LRU 策略:
GPU memory中的缓存空间是固定的,通过最近最少使用(Least Recently Used, LRU)策略进行管理。- 加载: 在每个计算步骤,首先计算所有记忆单元的相关性分数。对于被选中的需要使用的记忆单元,首先检查它是否已在
GPU cache中。如果不在,则从CPU memory传输到GPU memory。 - 卸载: 在注意力计算后,
GPU cache中的记忆单元会根据其频率分数 进行排序。分数最低的单元会被卸载回CPU memory。频率分数 的更新公式如下:- :记忆块 的频率分数。
- :衰减系数(
decay coefficient),用于反映之前查找的影响。 - :当前计算步骤中涉及的词元数量。
- :记忆单元的块大小。
- :在执行注意力计算后,
query向量 与key向量 之间的注意力分数(值在 0 到 1 之间)。 通过这种方式,InfLLM能够高效地利用有限的GPU memory,同时确保常用和相关记忆单元的快速访问。实验结果显示,这种卸载机制的miss rate较低,因此不会引入显著的时间开销。
- 加载: 在每个计算步骤,首先计算所有记忆单元的相关性分数。对于被选中的需要使用的记忆单元,首先检查它是否已在
- 进一步优化: 对于极长的序列,记忆单元的代表性词元也可以卸载到
CPU memory中,构建一个高效的k-nearest-neighbor索引,从而进一步降低计算复杂度。
5. 实验设置
5.1. 数据集
为了评估 InfLLM 的有效性,论文采用了两个广泛使用的长文本基准测试:
-
∞-Bench(InfiniteBench) (Zhang et al., 2023a):- 特点: 这是一个专门为评估
LLMs长上下文理解能力设计的基准,包含多样化的任务。 - 任务类型: 包括问答(
question answering)、摘要(summarization)、上下文检索(context retrieval)和数学计算(mathematic computing)。 - 语言: 主要使用英文数据集,因为基础模型(
Mistral,Llama-3)也主要在英文语料库上预训练。 - 长度特点:
∞-Bench中的平均序列长度高达145.1K词元,95%的序列长度甚至达到214K词元,远超大多数现有LLMs的最大上下文长度,因此对模型提出了严峻挑战。 - 具体任务示例 (在实验结果部分有体现):
Retrieve.PassKey (R.PK):在冗长的文本中查找一个特定的5位数字序列。Retrieve.Number (R.Num):检索文本中的一个特定数字。Retrieve.KV (R.KV):检索一个key-value对。Choice:选择题任务。QA:问答。Sum:摘要。Math.Find (Math.F):数学查找任务。
- 特点: 这是一个专门为评估
-
LongBench(Bai et al., 2023):- 特点: 这是一个双语、多任务的长期上下文理解基准测试。
- 长度特点:
95%的序列长度为31K词元。尽管比∞-Bench短,但对于大多数LLMs来说仍然是长序列。 - 具体任务示例 (在附录 Table 5 中有体现):
- 问答:
NQA,Qasper,MFQA,HQA,2WikiMQA,Musique。 - 摘要:
GovReport,QMSum,MultiNews。 - 代码补全:
RepoBench-P。 - 其他:
TREC,TQA,SAMSum,PsgCount,PsgRetrieval,LCC。
- 问答:
5.2. 评估指标
论文中使用的评估指标并未在正文中详细定义,但从任务类型和结果表格中的列名可以推断。根据 ∞-Bench 和 LongBench 的官方说明以及常见的 LLM 评估实践,以下是对可能使用的评估指标的概念定义、数学公式和符号解释的补充:
5.2.1. 准确率 (Accuracy)
- 概念定义: 准确率是最常见的分类和检索任务评估指标,它衡量模型正确预测的样本数量占总样本数量的比例。对于
Retrieve.PassKey、Retrieve.Number、Retrieve.KV、Choice、Math.Find等具有明确正确答案的离散型任务,准确率是直接衡量模型性能的指标。 - 数学公式:
- 符号解释:
Number of Correct Predictions:模型做出正确预测的样本数量。Total Number of Predictions:总共进行预测的样本数量。
5.2.2. Rouge Score (Recall-Oriented Understudy for Gisting Evaluation)
- 概念定义:
Rouge是一组用于评估自动文本摘要和机器翻译质量的指标。它通过比较模型生成的文本与参考(人工)摘要之间的重叠(overlap),主要关注召回率(recall),即有多少参考摘要中的信息被模型捕捉到。常见的Rouge变体包括ROUGE-N(基于 N-gram 重叠)、ROUGE-L(基于最长公共子序列LCS)和ROUGE-S(基于跳跃二元组skip-bigram)。论文中的Sum(摘要)任务通常使用ROUGE-L或ROUGE-1/2的F1 score。 - 数学公式 (以
ROUGE-N Recall为例):- 其中,
Count_match(n-gram)是指模型生成摘要和参考摘要中共同出现的n-gram的数量。 Count(n-gram)是指参考摘要中n-gram的数量。ROUGE-N Precision和ROUGE-N F1-score也可以类似计算。
- 其中,
- 符号解释:
n-gram:由文本中 个连续词元组成的序列。Reference:人工编写的参考摘要。Count_match(n-gram):模型生成文本与参考文本中匹配的n-gram计数。Count(n-gram):参考文本中n-gram的总计数。
5.2.3. F1 Score / Exact Match (EM)
- 概念定义: 对于问答(
QA)任务,通常使用Exact Match (EM)和F1 Score。Exact Match (EM):如果模型生成的答案与参考答案完全一致,则EM为 1,否则为 0。它是一个非常严格的指标。F1 Score:结合了精确率(Precision)和召回率(Recall)的调和平均值。它对答案的重叠程度进行评分,即使答案不完全匹配,也能给予部分分数。
- 数学公式 (F1 Score): 其中,
- 符号解释:
-
Precision:模型答案中正确词元的比例。 -
Recall:参考答案中被模型捕获的正确词元的比例。 -
Number of Correct Tokens in Prediction:模型预测答案中与参考答案重叠的词元数量。 -
Total Number of Tokens in Prediction:模型预测答案的总词元数量。 -
Total Number of Tokens in Reference:参考答案的总词元数量。由于论文中表格直接给出了
R.PK、R.Num、R.KV、Choice、QA、Sum、Math.F的数值,且大部分任务(除Sum外)是离散型判断或抽取任务,因此可以合理推断,这些任务的评估指标很可能就是准确率 (Accuracy),而Sum任务则会是ROUGE分数,QA可能是F1或EM。在论文中,这些数值通常被统一为百分比或分数,以方便比较。
-
5.3. 对比基线
论文将 InfLLM 与多种具有代表性的长上下文处理方法进行了比较:
-
原始模型 (Original Models):
Mistral-7B-inst-v0.2(最大上下文32K词元)Llama-3-8B-Instruct(最大上下文8K词元)- 这些模型代表了在不进行任何长上下文扩展时
LLM的基准性能,它们在遇到超长序列时通常表现不佳。
-
位置降尺度和重用方法 (Position Downscaling and Reusing):
NTK-aware scaled RoPE (NTK)(LocalLLaMA, 2023):通过非线性插值方法调整RoPE的旋转基,以扩展上下文窗口。SelfExtend(Jin et al., 2024):通过重用相邻词元的位置ID,使扩展的相对位置保持在训练上下文窗口的范围内。- 这些方法旨在解决
LLM在未见长度上的“域外问题”。
-
滑动窗口方法 (Sliding Window):
LM-Infinite (Infinite)(Han et al., 2023):采用滑动窗口机制直接丢弃远距离上下文,流式读取超长序列。StreamingLLM (Stream)(Xiao et al., 2023):同样采用滑动窗口机制,只关注局部上下文。- 这类方法效率高,但无法捕获长距离依赖。
-
Key-Value 淘汰方法 (Key-Value Eviction):
H2O(Zhang et al., 2023b):一种广泛使用的key-value淘汰方法,旨在通过丢弃无用的key-value向量来减少计算复杂度。- 论文指出,这类方法主要用于提高计算效率,但由于未见位置编码问题,通常无法在不额外训练的情况下外推上下文窗口。
-
持续训练模型 (Models with Continual Training):
Llama-3-8B-Instruct-Gradient-1048k (Llama-1M):一个在长文本数据和聊天数据集上进行了额外微调的模型,其上下文窗口扩展到1048K。- 这个基线用于与
InfLLM进行性能和效率的对比,以证明InfLLM无需训练即可达到竞争性能的优势。
-
检索增强生成 (Retrieval-Augmented Generation, RAG):
RAG-E5(使用E5-mistral-7B-instruct作为检索模型):RAG模型通过从外部数据库检索相关信息来增强LLM的生成能力。- 这个基线用于对比
InfLLM作为内部上下文增强机制的泛化能力和训练成本优势。
5.4. 实现细节
- 基础模型:
Mistral-7B-Instruct-v0.2(Jiang et al., 2023),其最大训练长度为32K词元。Llama-3-8B-Instruct(Meta, 2024),其最大训练长度为8K词元。
InfLLM参数设置:- 编码块大小(
encoding chunk size):512词元。 - 过去
key-value向量的记忆单元大小(memory unit size):128词元。 - 代表性词元数量(
number of representative tokens):4。 - 局部窗口大小(
local window size):4K词元(对于Mistral-based和Llama-3-based InfLLM均是)。 - 加载的相关记忆单元数量:
Mistral-based InfLLM加载96个,Llama-3-based InfLLM加载32个。 - 初始词元数量(
initial tokens):128词元(LM-Infinite,StreamingLLM,InfLLM均采用此设置,用于覆盖系统提示和任务描述)。
- 编码块大小(
- 计算加速: 采用
FlashAttention(Dao, 2023) 加速所有基线模型的实验。 - 硬件:
NVIDIA A100或A800GPU。 - 精度: 所有实验均采用半浮点精度(
half-float precision)。 - 附录中的额外细节:
context memory在所有LLM层构建。GPU cache大小设置为32(是每个步骤加载单元数量的两倍)。- 频率分数衰减系数(
frequency score decay coefficient):0.1。 - 在扩展到
1,024K上下文的实验中,编码块大小设置为2048,代表性词元数量设置为1,以加速实验。
6. 实验结果与分析
6.1. 核心结果分析
6.1.1. 在 ∞-Bench 上的表现
以下是原文 Table 1 的结果,展示了 InfLLM 和基线模型在 ∞-Bench 上的表现。
以下是原文 Table 1 的结果:
| Window | Streaming | R.PK | R.Num | R.KV | Choice | QA | Sum | Math.F | Avg. | |
| Mistral-based Models (7B) | ||||||||||
| Mistral | 32K | 28.8 | 28.8 | 14.8 | 44.5 | 12.9 | 25.9 | 20.6 | 25.2 | |
| NTK | 128K | × | 100.0 | 86.8 | 19.2 | 40.2 | 16.9 | 20.3 | 26.9 | 44.3 |
| SelfExtend | 128K | 100.0 | 100.0 | 15.6 | 42.8 | 17.3 | 18.8 | 19.1 | 44.8 | |
| Infinite | 32K | ✗ | 28.8 | 28.8 | 0.4 | 42.8 | 11.4 | 22.5 | 16.3 | 21.6 |
| Streaming | 32K | ✓ | 28.8 | 28.5 | 0.2 | 42.4 | 11.5 | 22.1 | 16.9 | 21.5 |
| H20 | 32K | ✓ | 8.6 | 4.8 | 2.6 | 48.0 | 15.6 | 24.4 | 26.9 | 18.7 |
| InfLLM | 16K | ✓ | 100.0 | 96.1 | 96.8 | 43.7 | 15.7 | 25.8 | 25.7 | 57.7 |
| Llama-3-based Models (8B) | ||||||||||
| Llama-3 | 8K | 8.5 | 7.8 | 6.2 | 44.1 | 15.5 | 24.7 | 21.7 | 18.4 | |
| NTK | 128K | ×X | 0.0 | 0.0 | 0.0 | 0.0 | 0.4 | 6.4 | 2.6 | 1.3 |
| SelfExtend | 128K | × | 100.0 | 100.0 | 0.2 | 19.7 | 8.6 | 14.7 | 22.6 | 38.0 |
| Infinite | 8K | ✓ | 6.8 | 7.6 | 0.2 | 41.5 | 14.6 | 20.8 | 20.6 | 16.0 |
| Streaming | 8K | ✓ | 8.5 | 8.3 | 0.4 | 40.6 | 14.3 | 20.4 | 21.4 | 16.3 |
| H20 | 8K | ✓ | 2.5 | 2.4 | 0.0 | 0.0 | 0.7 | 2.8 | 6.0 | 2.1 |
| InfLLM | 8K | ✓ | 100.0 | 99.0 | 5.0 | 43.7 | 19.5 | 24.3 | 23.7 | 45.0 |
分析:
InfLLM的显著优势: 无论是基于Mistral还是Llama-3,InfLLM在总体平均性能 (Avg.) 上都远超所有其他基线。尤其是在Mistral上达到了57.7,Llama-3上达到了45.0。这表明InfLLM能够有效弥补原始LLM在处理长上下文时的不足。- 滑动窗口机制的局限性:
Infinite和Streaming模型(如Mistral-based的Infinite21.6,Streaming21.5)虽然能处理长序列,但由于直接丢弃了远距离上下文,其性能甚至低于原始Mistral模型(25.2)。这凸显了InfLLM中context memory在捕获长距离依赖方面的关键作用。 - 位置编码外推方法的不足:
NTK和SelfExtend方法虽然能够将上下文窗口扩展到128K,但在某些任务上表现不佳,例如Mistral-based的SelfExtend在Sum任务上表现为18.8,低于原始Mistral的25.9。这可能是因为它们主要解决了out-of-domain问题,但未能有效解决噪声上下文引起的distraction问题。Llama-3-based NTK甚至在多个任务上得到0.0的极低分数,表明其稳定性可能不足。 KV淘汰方法的局限性:H2O的表现普遍非常差(Mistral-based为18.7,Llama-3-based为2.1),这验证了论文的观点,即KV淘汰方法虽然旨在提高效率,但不能在不额外训练的情况下有效外推上下文窗口。InfLLM在特定任务上的表现:InfLLM在Retrieve.PassKey、Retrieve.Number、Retrieve.KV等检索任务上表现尤为突出,例如Mistral-based InfLLM在R.PK和R.KV上分别达到100.0和96.8。这直接证明了其context memory能够准确地从海量信息中检索出关键信息。
6.1.2. 与持续训练模型的比较
以下是原文 Table 2 的结果,比较了 InfLLM 与经过持续训练的 Llama-3-8B-Instruct-Gradient-1048k (Llama-1M) 的性能和效率。
以下是原文 Table 2 的结果:
| Train-Free | R.PK | R.Num | R.KV | Choice | QA | Sum | Math.F | VRAM | Time | |
| Llama-1M | ✗ | 100.0 | 99.8 | 23.2 | 51.5 | 13.6 | 18.5 | 18.3 | 76.6G | 40.4s |
| InfLLM | ✓ | 100.0 | 99.0 | 5.0 | 43.7 | 19.5 | 24.3 | 23.7 | 26.3G | 26.7s |
| Llama-1M+InfLLM | ✗ | 100.0 | 100.0 | 55.8 | 39.3 | 20.3 | 17.1 | 31.4 | 26.3G | 26.7s |
分析:
- 性能对比:
InfLLM(训练无关) 在R.PK和R.Num上与Llama-1M(持续训练) 表现相当,甚至在QA、Sum、Math.F上超越了Llama-1M。这有力地支持了论文的核心论点:LLMs即使未经额外训练,也具有理解极长序列的内在能力,InfLLM只是有效地激活了这种能力。Llama-1M需要512个GPU进行持续训练,这对于大多数研究者来说是无法承受的,而InfLLM无需训练,具有极高的实用性。
- 效率优势:
InfLLM的VRAM使用量仅为Llama-1M的34%(26.3Gvs76.6G)。InfLLM的时间消耗降低了34%(26.7svs40.4s)。- 论文还提到,在
256K词元长度时,full-attention基线模型会因为内存不足而失败,而InfLLM可以在单个GPU上高效处理高达1024K词元的序列。这证明了InfLLM在资源受限环境下的强大实用性。
InfLLM与持续训练模型的结合:Llama-1M+InfLLM的结果表明,InfLLM也可以与持续训练的模型结合,并且在某些任务(如QA,Math.F)上带来进一步的性能提升,同时保持了InfLLM的效率优势。这表明InfLLM还可以作为一种提高推理速度的通用方法。
6.1.3. 与检索增强生成 (RAG) 的比较
以下是原文 Table 3 的结果,比较了 InfLLM 和 RAG 在三个上下文检索任务上的表现。
以下是原文 Table 3 的结果:
| Task | | R.PK | R.Num | R.KV |
| RAG-E5 | 89.2 | 65.4 | 13.2 |
| InfLLM | 100.0 | 96.1 | 96.8 |
分析:
InfLLM性能优势:InfLLM在所有三个检索任务上都显著优于RAG-E5模型。尤其在R.PK和R.KV任务上,InfLLM接近或达到满分(100.0和96.8),而RAG-E5分别只有89.2和13.2。InfLLM的泛化能力: 结果表明,即使没有额外数据或训练,InfLLM也表现出卓越的泛化能力。RAG的局限性:RAG需要额外的检索数据来训练检索模型,并需要微调LLM以适应检索增强的输入,而InfLLM是完全训练无关的。RAG模型的性能受其检索组件的限制,并且可能面临out-of-distribution问题,导致在训练分布之外的任务上表现不佳。InfLLM则无特定任务要求,具有更广泛的适用性。
6.1.4. 扩展到 1,024K 上下文
下图是原文 Figure 3 的结果,展示了 InfLLM 在不同长度序列上的表现。

分析:
InfLLM的超长上下文处理能力: 在Retrieve.PassKey任务中,即使序列长度扩展到1,024K词元,InfLLM仍能保持100%的准确率。这充分证明了InfLLM在极端长序列中有效捕获长距离依赖和识别关键信息的能力。- 滑动窗口方法的局限性: 相比之下,
LM-Infinite的性能随着序列长度的增加而迅速下降。这证实了仅依赖局部窗口而丢弃远距离上下文的方法在长距离依赖任务中的根本性缺陷。 - 对噪声的鲁棒性: 这个实验也间接表明
InfLLM能够有效过滤长序列中的噪声,精准定位所需信息。
6.2. 消融实验/参数分析
6.2.1. 记忆设置的影响 (The Impact of Memory Settings)
下图是原文 Figure 2 的结果,展示了 InfLLM 在不同代表性词元数量、选择单元数量和记忆单元大小下的性能。

分析:
- 不同代表性词元数量 (Figure 2a):
- 随着代表性词元数量的增加(从
1到4),模型性能呈现上升趋势,表明更多的代表性词元能更好地捕捉记忆单元的语义内容。 - 当数量达到
8时,性能略有下降。这可能是由于引入了语义不相关的词元作为单元表示,反而增加了噪声。这暗示了存在一个最优的代表性词元数量,并且未来的工作可以探索更有效和强大的单元表示方法。
- 随着代表性词元数量的增加(从
- 不同选择单元数量 (Figure 2b):
- 随着选择的记忆单元数量增加(从
1到32),模型性能显著提升。这表明更多的记忆单元能提供更全面的相关上下文,提高了信息召回率。 - 然而,更多的单元也意味着更高的内存调度时间和注意力计算时间。因此,提升查找准确性以减少所需单元数量,是提高效率的关键。
- 随着选择的记忆单元数量增加(从
- 不同记忆单元大小 (Figure 2c):
- 记忆单元大小对性能的影响因任务而异。例如,对于
Retrieve.KV任务,64或128可能是较好的选择;对于Math.Find任务,32可能是最优的。 - 这说明记忆单元的语义连贯性很重要。过大的单元可能导致查找不精确,过小的单元则增加查找开销。
- 目前使用启发式规则分割上下文可能导致次优性能,未来研究方向是探索动态分割上下文的方法。
- 记忆单元大小对性能的影响因任务而异。例如,对于
6.2.2. 消融研究 (Ablation Study)
以下是原文 Table 4 的结果,展示了关于动态记忆查找和单元表示的消融研究。
以下是原文 Table 4 的结果:
| Task | R.KV | Math.F | QA |
| InfLLM | 96.8 | 25.7 | 15.7 |
| Decoding-Only | 85.2 | 26.3 | 12.0 |
| w/o Lookup | 0.4 | 16.3 | 11.4 |
| Mean Repr | 84.6 | 25.1 | 14.9 |
分析:
- 上下文记忆查找 (Context Memory Lookup):
w/o Lookup(无任何记忆查找):性能急剧下降,尤其在R.KV任务上仅为0.4。这验证了context memory对于长文本理解的不可或缺性,以及在编码和解码阶段都集成远距离上下文的重要性。模型需要整合长距离上下文来生成连贯的记忆,并从中收集有用信息以生成正确答案。Decoding-Only(仅在解码阶段查找):性能相较于完整InfLLM有所下降(R.KV从96.8降至85.2,QA从15.7降至12.0)。这表明远距离上下文信息对于输入编码和输出生成两个阶段都至关重要,模型需要全面的长距离上下文才能进行有效的理解和推理。
- 单元表示 (Unit Representation):
Mean Repr(使用平均key向量作为单元表示):虽然性能略低于InfLLM的代表性词元选择方法(R.KV从96.8降至84.6),但仍能达到竞争性水平。这说明LLM原始的注意力向量本身就具有有效进行相关性计算的能力。这也暗示了未来可以进一步探索更高效的单元表示方法来提升模型性能。
6.3. 在 LongBench 上的表现 (附录 C.2)
以下是原文 Table 5 的结果,展示了 InfLLM 和基线模型在 LongBench 上的表现。
以下是原文 Table 5 的结果:
| Window | NQA | Qasper | MFQA | HQA | 2WikiMQA | Musique | |||
| Mistral-based Models (7B) | |||||||||
| Mistral | 32K | 22.06 | 29.16 | 47.65 | 37.53 | 21.96 | 19.03 | ||
| Infinite | 6K | 18.44 | 30.02 | 39.05 | 32.02 | 22.27 | 15.81 | ||
| Streaming | 6K | 17.92 | 30.05 | 39.09 | 32.18 | 21.83 | 14.71 | ||
| InfLLM | 6K | 22.12 | 29.33 | 47.42 | 36.56 | 22.31 | 17.68 | ||
| InfLLM | 12K | 23.03 | 29.52 | 47.62 | 39.53 | 23.61 | 18.92 | ||
| Llama-3-based Models (8B) | |||||||||
| Llama-3 | 8K | 19.85 | 42.36 | 41.03 | 47.38 | 39.20 | 22.96 | ||
| Infinite | 8K | 19.39 | 42.80 | 40.44 | 43.77 | 37.89 | 18.33 | ||
| Streaming | 8K | 20.05 | 42.46 | 39.54 | 43.69 | 37.89 | 19.68 | ||
| InfLLM | 8K | 22.64 | 43.70 | 49.03 | 49.04 | 35.61 | 26.06 | ||
| Window | GovReport | QMSum | MultiNews | TREC | TQA | SAMSum | |||
| Mistral-based Models (7B) | |||||||||
| Mistral | 32K | 31.12 | 23.87 | 26.62 | 71.00 | 85.97 | 42.29 | ||
| Infinite | 6K | 29.74 | 21.92 | 26.65 | 70.00 | 85.22 | 41.60 | ||
| Streaming | 6K | 29.83 | 21.94 | 26.64 | 70.00 | 85.57 | 41.31 | ||
| InfLLM | 6K | 31.03 | 23.49 | 26.70 | 69.00 | 86.67 | 42.52 | ||
| InfLLM | 12K | 31.37 | 23.77 | 26.66 | 71.00 | 87.34 | 41.80 | ||
| Llama-3-based Models (8B) | |||||||||
| Llama-3 | 8K | 29.94 | 21.45 | 27.51 | 74.00 | 90.50 | 42.30 | ||
| Infinite | 8K | 29.25 | 21.41 | 27.62 | 74.00 | 90.08 | 41.72 | ||
| Streaming | 8K | 29.17 | 21.33 | 27.56 | 73.50 | 90.08 | 41.55 | ||
| InfLLM | 8K | 30.76 | 22.70 | 27.57 | 73.50 | 90.91 | 42.43 | ||
| Window | PsgCount | PsgRetrieval | LCC | RepoBench-P | Avg. | ||||
| Mistral-based Models (7B) | |||||||||
| Mistral | 32K | 3.95 | 86.94 | 57.42 | 54.14 | 43.78 | |||
| Infinite | 6K | 2.08 | 42.80 | 57.12 | 53.43 | 39.07 | |||
| Streaming | 6K | 2.50 | 42.17 | 55.38 | 51.46 | 38.67 | |||
| InfLLM | 6K | 2.87 | 64.00 | 56.67 | 52.97 | 41.90 | |||
| InfLLM | 12K | 3.01 | 87.42 | 56.69 | 52.09 | 44.02 | |||
| Llama-3-based Models (8B) | |||||||||
| Llama-3 | 8K | 8.50 | 62.50 | 60.83 | 49.14 | 44.73 | |||
| Infinite | 8K | 4.50 | 50.00 | 60.12 | 48.62 | 43.03 | |||
| Streaming | 8K | 5.00 | 49.00 | 60.35 | 48.95 | 42.99 | |||
| InfLLM | 8K | 7.17 | 84.00 | 59.88 | 46.48 | 46.95 | |||
分析:
InfLLM的持续优越性:InfLLM在LongBench的各种任务中,普遍优于其他能够处理流式输入的模型(Infinite,Streaming)。这再次证明了context memory提供的上下文信息能够有效增强模型性能。- 滑动窗口方法的不足: 当以
Llama-3作为基础模型时,StreamingLLM和LM-Infinite的性能与原始Llama-3相当甚至更差。这进一步强调了滑动窗口注意力机制虽然能扩展上下文窗口,但由于丢弃了长距离上下文,无法实现有效的长序列理解。 - 对噪声的过滤能力:
Mistral能够处理长达32K的文本,覆盖了LongBench的大部分实例。相比之下,InfLLM在仅12K的窗口大小下,平均性能与Mistral相当甚至更优(Mistral平均43.78,InfLLM平均44.02)。这表明InfLLM能够有效过滤长上下文中的噪声,从而实现更好的长序列理解。
6.4. 在 Vicuna 上的实验 (附录 C.3)
以下是原文 Table 6 的结果,展示了 Vicuna-based 模型的实验结果。
以下是原文 Table 6 的结果:
| R.PK | R.Num | R.KV | Math.F | |
| Vicuna | 5.08 | 4.41 | 1.40 | 11.71 |
| InfLLM | 99.15 | 81.69 | 0.60 | 11.14 |
分析:
- 显著的上下文扩展:
InfLLM成功地将Vicuna(最大长度仅4K)的上下文长度扩展到128K。在Retrieve.Passkey和Retrieve.Number任务上,InfLLM取得了显著的性能提升,准确率分别从5.08提升到99.15和从4.41提升到81.69。这再次证明了InfLLM能够有效地将短上下文模型泛化到长上下文场景。 - 局限性: 然而,在
Retrieve.KV和Math.Find任务上,InfLLM未能显示性能增益,甚至在R.KV任务上略有下降。 - 潜在原因: 论文解释说,这可能是因为
Vicuna的隐藏向量(hidden vectors)在处理极长文本中的噪声时,过滤能力有限。这使得context memory在更复杂的任务(如Retrieve.KV和Math.Find)中难以有效定位相关信息。这暗示了基础模型的内在能力对InfLLM的有效性有一定影响,也为未来的研究指明了方向:设计更强大的记忆机制来应对基础模型在噪声过滤方面的局限性。
7. 总结与思考
7.1. 结论总结
本文提出了 InfLLM,一个无需训练(training-free)的方法,旨在显著提升大型语言模型(LLMs)处理长序列的能力。InfLLM 将 sliding window attention 机制与一个高效的上下文记忆(context memory)模块相结合。该记忆模块通过将远距离上下文存储在块级记忆单元中,并利用代表性词元进行相关性查找,从而允许 LLMs 在有限的上下文窗口内,高效且精确地捕获长距离依赖。
实验结果在 ∞-Bench 和 LongBench 这两个广泛使用的长文本基准测试上证明了 InfLLM 的卓越有效性。它使在数千个词元序列上预训练的 LLMs 达到了与在长序列上持续训练的竞争基线相当甚至更优的性能,同时大幅降低了计算和内存成本。尤其在 1,024K 词元的超长序列上,InfLLM 仍能保持 100% 的准确率,展现了其在处理极端长序列时的强大能力。此外,InfLLM 在上下文检索任务中也优于 Retrieval-Augmented Generation (RAG) 方法,凸显了其训练无关和广泛适用性的优势。
InfLLM 的成功揭示了 LLMs 在无需额外训练的情况下,理解和推理极长序列的内在潜力。
7.2. 局限性与未来工作
论文作者指出了 InfLLM 存在的局限性,并提出了未来的研究方向:
-
CPU 内存使用:
InfLLM将大量过去key-value (KV)缓存存储在CPU memory中,这增加了CPU memory的使用量。- 未来工作: 可以通过集成
KV cache quantization等技术来减少CPU memory需求,进一步优化资源利用。
- 未来工作: 可以通过集成
-
推理速度提升空间: 尽管
InfLLM显著降低了处理长文本的计算开销,但仍有进一步加速的潜力。- 未来工作: 可以将其与
llama.cpp和vLLM(Kwon et al., 2023) 等高效推理框架结合,以进一步增强InfLLM的推理速度。
- 未来工作: 可以将其与
-
记忆模块的训练: 当前
InfLLM的记忆模块是完全训练无关的。- 未来工作: 探索对
context memory module进行高效训练,以进一步提升模型的性能。这可能包括学习更优的单元表示或查找策略。
- 未来工作: 探索对
-
KV cache压缩: 结合key-value cache压缩方法,可以进一步降低计算和内存成本。 -
更强大的单元表示: 实验中发现,代表性词元数量过多可能引入噪声,且平均表示也能取得竞争性结果。
- 未来工作: 探索更有效和强大的记忆单元表示方法,以进一步提高模型性能。
-
动态上下文分割: 记忆单元的最佳大小可能因任务而异,启发式分割可能导致次优性能。
- 未来工作: 研究如何动态地分割上下文以形成记忆单元,以适应不同任务和文本特性的需求。
-
基础模型能力的依赖: 在
Vicuna上的实验表明,基础模型的隐藏向量过滤噪声能力有限时,InfLLM在某些复杂任务上可能无法获得性能增益。- 未来工作: 设计更强大的记忆机制来克服基础模型在噪声过滤方面的局限性。
7.3. 个人启发与批判
7.3.1. 个人启发
- “内在能力”的挖掘: 这篇论文最引人注目的一点是其核心发现——“揭示了
LLMs在无需任何微调的情况下理解极长序列的内在能力”。这提供了一个非常积极的信号:现有的大型、昂贵的LLMs可能比我们想象的更有潜力,我们需要的不是无止境地扩大训练数据和模型规模,而是更智能、更高效地“引导”或“激活”它们已有的能力。这种训练无关(training-free)的思路,对于资源有限的研究者和开发者来说,具有巨大的吸引力。 - 工程与算法的精妙结合:
InfLLM巧妙地结合了sliding window attention的工程效率和context memory的算法智能。通过块级记忆单元、代表性词元选择和LRU缓存管理,它在计算成本、内存效率和性能之间找到了一个非常好的平衡点,为LLM的流式长上下文处理提供了实用的解决方案。 - 对
RAG的挑战与互补:InfLLM在某些任务上优于RAG的结果很有趣。它表明,对于上下文理解任务,模型内部的“自检索”机制可能比外部检索器更直接、更高效,因为无需处理检索器本身的训练、out-of-distribution和LLM与检索器结果的适配问题。这并非意味着RAG无用,而是提示我们,两种机制可能各有侧重,甚至可以互补:InfLLM专注于利用模型内在能力处理核心上下文,RAG则可用于引入外部知识。 - 对
LLM训练的启示: 如果LLMs真的具备这种内在的长上下文理解能力,那么未来的LLM训练策略是否可以更多地关注于如何更好地编码和提取这种能力,而不是仅仅通过简单地增加上下文窗口长度来训练?这可能会导致更高效、更具泛化性的LLM。
7.3.2. 批判与潜在改进
- 代表性词元选择的局限性: 论文中用于选择代表性词元的
representative score公式 依赖于局部窗口内的注意力分数。虽然它是训练无关的,但这种启发式方法可能无法捕获复杂的语义重要性。例如,一个在局部窗口内不显眼的词元,可能在全局语境中具有关键的意义。在Vicuna实验中的局限性也可能部分源于此。- 改进方向: 可以探索更复杂的、基于语义嵌入的聚类方法,或者使用小型、轻量级的编码器来学习更鲁棒的块级表示,即使这需要微量的额外训练。
- 位置编码的简化处理: 将所有远距离记忆单元的位置编码统一设置为 是一种有效的简化,但这也意味着模型失去了这些远距离记忆单元之间更精细的相对位置信息。对于需要精确时间或顺序推理的长文本任务,这可能是一个潜在的瓶颈。
- 改进方向: 可以在不引入
out-of-domain问题的前提下,探索一种分层或稀疏的位置编码策略,以在记忆单元内部或记忆单元之间保留更多的相对位置信息。
- 改进方向: 可以在不引入
- “注意力分散问题”的根本解决: 尽管
InfLLM通过选择性地加载相关记忆单元来缓解distraction issue,但其内部的sliding window attention仍然可能在局部窗口内受到噪声干扰。- 改进方向: 可以考虑将
InfLLM与更先进的稀疏注意力机制或动态注意力机制结合,以在局部和全局层面都实现更精细的注意力控制。
- 改进方向: 可以考虑将
- 记忆查找的开销: 尽管采用了块级查找和
LRU卸载机制,但在极端长序列下,每次都要计算与所有(或大量)记忆单元的相关性分数,这仍然可能带来可观的计算开销,尤其是在CPU memory和GPU memory之间的数据传输。- 改进方向: 引入分层索引或更高效的相似性搜索算法(如
approximate nearest neighbor, ANN)来加速记忆查找过程。
- 改进方向: 引入分层索引或更高效的相似性搜索算法(如
- 评估任务的代表性: 尽管
∞-Bench和LongBench是很好的基准,但长文本理解的复杂性远不止这些任务。例如,涉及复杂推理、多跳问答或需要理解多文档之间关系的任务,对context memory的语义组织和推理能力提出了更高的要求。-
改进方向: 在更具挑战性、更复杂的长文本推理任务上进行评估,以更全面地验证
InfLLM的鲁棒性和能力。总而言之,
InfLLM提供了一个兼顾实用性和高性能的training-free解决方案,为LLM的长上下文处理开辟了新的思路。尽管存在一些可改进的方面,但其核心思想和实验结果无疑为该领域的发展贡献了重要力量。
-
相似论文推荐
基于向量语义检索推荐的相关论文。