论文状态:已完成

InfLLM: Training-Free Long-Context Extrapolation for LLMs with an Efficient Context Memory

发表:2024/02/07
原文链接PDF 下载
价格:0.100000
价格:0.100000
价格:0.100000
已有 2 人读过
本分析由 AI 生成,可能不完全准确,请以原文为准。

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 1,0241,024K, 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. 发表期刊/会议

arXivarXiv 是一个预印本(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,024KInfLLM 仍能有效地捕获长距离依赖。我们的代码可以在 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 通常在相对较短的序列上进行预训练(例如,几千个词元),这导致它们在处理远超训练长度的序列时面临严峻挑战。

主要的挑战包括:

  1. 域外问题 (Out-of-domain Issue): LLMs 未在长序列上进行训练,导致其在遇到超出其训练范围的上下文长度时,性能会急剧下降。这通常与位置编码(positional encoding)机制的限制有关。

  2. 注意力分散问题 (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 方法,其核心贡献和主要发现如下:

  1. 揭示 LLMs 的内在长序列理解能力 (Unveiling Intrinsic Capacity): 论文的一个关键发现是,LLMs 本身具备在无需微调的情况下理解极长序列的内在能力。InfLLM 通过提供相关的长距离上下文,有效地利用了这种能力。
  2. 提出训练无关的基于记忆的方法 InfLLM (Training-Free Memory-Based Method): InfLLM 是一种无需任何额外训练即可扩展 LLMs 上下文窗口的方法。这显著降低了计算开销,避免了持续训练可能引入的性能下降和模型能力变化。
  3. 结合滑动窗口注意力与高效上下文记忆 (Efficient Context Memory with Sliding Window Attention): InfLLMsliding window attention 的基础上,额外构建了一个高效的 context memory 模块。这使得模型既能处理流式长输入(通过滑动窗口),又能通过记忆模块捕获被滑动窗口忽略的远距离依赖。
  4. 创新性块级记忆单元机制 (Novel Block-Level Memory Units): 为了解决大规模噪声上下文和内存查找效率问题,InfLLM 将历史 key-value 向量组织成块(blocks),每个块作为一个记忆单元。并通过选择块内最具代表性的词元(representative tokens)作为单元表示,来高效计算相关性并进行查找。
    • 有效查找 (Effective Lookup): 块级语义一致性提高了信息检索的准确性。
    • 高效查找 (Efficient Lookup): 避免了逐词元计算,减少了计算开销,并通过连续内存访问提高了效率。
  5. 高效缓存管理与卸载机制 (Efficient Cache Management and Offloading): InfLLM 采用将不常用记忆单元卸载到 CPU memory、只在 GPU memory 中保留常用单元的策略,并结合 LRU(最近最少使用)缓存策略,大大降低了 GPU memory 占用,实现了在有限资源下处理极长序列的能力。
  6. ∞-BenchLongBench 上的卓越性能 (Superior Performance on Benchmarks):
    • 在无需任何训练的情况下,InfLLM 使 Mistral-7BLlama-3-8B 等基础模型在处理几千个词元的训练长度时,达到了与在长序列上持续训练的竞争基线相当甚至更优的性能。
    • InfLLM 成功将 Llama-3 的上下文长度从 8K 扩展到 128K,甚至在 1,024K 的超长序列上也能有效捕获长距离依赖,并在 Retrieve.Passkey 任务上保持 100% 的准确率。
  7. 效率优势 (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(如 LlamaMistral)使用的位置编码,它通过旋转矩阵的方式将位置信息融入到 querykey 向量中,使其能够更好地处理相对位置关系。

  • 自注意力机制 (Self-Attention Mechanism): Transformer 架构的核心组件,允许模型在处理序列中的每个词元时,都能关注到序列中的其他所有词元,并根据它们的重要性分配不同的权重。其计算涉及三个核心向量:

    • 查询向量 (Query, Q): 代表当前词元在寻找相关信息时发出的“查询”。
    • 键向量 (Key, K): 代表序列中所有词元的“索引”或“特征”,用于响应查询。
    • 值向量 (Value, V): 包含序列中所有词元的实际信息,根据注意力权重进行加权求和后输出。 基本的 Self-Attention 计算公式如下: Attention(Q,K,V)=softmax(QKTdk)V \mathrm{Attention}(Q, K, V) = \mathrm{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V 其中,QQKKVV 分别是 querykeyvalue 矩阵,dkd_kkey 向量的维度,用于缩放点积以防止梯度过小。softmax 函数将注意力分数归一化为概率分布。
  • Key-Value Cache (KV Cache):LLMs 的推理(inference)阶段,尤其是在自回归(autoregressive)生成时,为了避免重复计算已经处理过的词元对应的 KeyValue 向量,这些向量会被缓存起来。当生成下一个词元时,只需要计算新词元的 KeyValue,然后与 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 的方法: 针对广泛使用的 RoPESu 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-InfiniteHan et al., 2023)和 Streaming-LLMXiao 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 的长上下文处理能力经历了从固定短上下文到尝试扩展、再到更精细化管理的演进:

  1. 初始阶段: Transformer 模型由于其二次计算复杂度,通常在几千个词元(例如 5124096)的固定上下文窗口内进行训练。超过这个长度,计算和内存成本都难以承受。
  2. 早期扩展尝试: 研究者尝试通过修改位置编码机制(如 RoPE 的缩放和插值方法)来“欺骗”模型,使其能够泛化到更长的位置。同时,稀疏注意力等高效注意力机制开始被提出,以降低计算复杂度,允许在预训练阶段处理更长的序列。
  3. 流式处理与记忆引入: 随着 LLMs 应用场景对实时、流式输入的需求增加,sliding window attention 等机制被提出,以固定有限的窗口处理无限长的流。然而,这种方式丢失了长距离信息。为了弥补这一缺陷,将外部记忆引入 Transformer 成为一个自然的方向。早期的记忆增强 Transformer 模型通常需要专门设计和训练。
  4. 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-domaindistraction issues

具体而言,它构建了一个无需训练的上下文记忆(context memory)模块。当模型处理长序列时,一部分历史 key-value 向量会被“驱逐”并存储到这个记忆中。在每个计算步骤,InfLLM 会智能地从这个记忆中查找并加载与当前词元最相关的记忆单元,然后将这些记忆单元与当前局部上下文拼接起来,形成扩展的上下文窗口,供 LLM 进行注意力计算。

4.2. 核心方法详解

InfLLM 的整体框架和具体组件设计如下:

4.2.1. 整体框架 (Overall Framework)

为了处理远超 LLM 训练长度的输入序列 s={ti}i=1ls = \{ t_i \}_{i=1}^lInfLLM 采取了分块(chunk-by-chunk)编码和逐词元(token-by-token)生成的方式。

  1. 序列分段: 对于每个计算步骤,输入由过去 key-value 向量 P={(kj,vj)}j=1lP\mathbf{P} = \{ (\mathbf{k}_j, \mathbf{v}_j) \}_{j=1}^{l_P} 和当前词元 Xˉ={ti+lP}i=1lˉX\bar{\mathbf{X}} = \{ \mathbf{t}_{i+l_P} \}_{i=1}^{\bar{l}_X} 组成。在编码阶段,lˉX\bar{l}_X 等于块大小(chunk size);在解码阶段,lˉX\bar{l}_X 等于 1。

  2. 过去上下文的划分: 过去 key-value 向量 P\mathbf{P} 根据其与当前词元的距离,被划分为三个部分:

    • 初始词元 (Initial Tokens), I=P[1:lI]\mathbf{I} = \mathbf{P}_{[1:l_I]} 最早的 lIl_I 个词元,通常包含系统提示(system prompts)和任务描述,这些信息通常非常重要,因此被保留。
    • 驱逐词元 (Evicted Tokens), E=P[lI+1:lPlL]\mathbf{E} = \mathbf{P}_{[l_I+1:l_P-l_L]} 介于初始词元和局部词元之间的词元。这些词元由于距离当前词元较远,通常会被滑动窗口注意力机制忽略。InfLLM 将这部分词元存储到 context memory 中。
    • 局部词元 (Local Tokens), L=P[lPlL+1:lP]\mathbf{L} = \mathbf{P}_{[l_P-l_L+1:l_P]} 距离当前词元最近的 lLl_L 个词元,它们构成了 sliding window attention 的核心部分。
  3. 构建当前 Key-Value 缓存: 在每个计算步骤中,InfLLM 会动态地构建一个用于注意力计算的 key-value 缓存 C\mathbf{C}。这个缓存由三部分拼接而成:

    • 初始词元 I\mathbf{I}
    • context memory 中通过查找操作 f()f(\cdot) 获取的相关记忆单元。
    • 局部词元 L\mathbf{L}。 因此,当前 key-value 缓存表示为: C=Concat(I,f(X,E),L) \mathbf{C} = \mathrm{Concat}(\mathbf{I}, f(\mathbf{X}, \mathbf{E}), \mathbf{L}) 其中,f(X,E)f(\mathbf{X}, \mathbf{E}) 表示从 evicted tokens E\mathbf{E} 构成的 context memory 中,根据当前词元 X\mathbf{X} 查找并返回相关的记忆单元。
  4. 注意力计算: 构建好 key-value 缓存 C\mathbf{C} 后,LLM 会使用它与当前词元的 query 向量 QX\mathbf{Q}\mathbf{X} 进行注意力计算,生成输出 O\mathbf{O}O=Attn[QX,Concat(Ck,KX),Concat(Cv,VX)] \mathbf{O} = \mathrm{Attn} \left[ \mathbf{Q}\mathbf{X}, \mathrm{Concat}(\mathbf{C}_k, \mathbf{K}\mathbf{X}), \mathrm{Concat}(\mathbf{C}_v, \mathbf{V}\mathbf{X}) \right]

    • Q\mathbf{Q}, K\mathbf{K}, V\mathbf{V}LLM 注意力层中的参数矩阵,用于将输入特征映射为 querykeyvalue 向量。

    • X\mathbf{X} 是当前需要编码的词元序列。

    • KX\mathbf{K}\mathbf{X}VX\mathbf{V}\mathbf{X} 是通过 K\mathbf{K}V\mathbf{V} 矩阵对当前词元 X\mathbf{X} 线性变换得到的 keyvalue 向量。

    • Ck\mathbf{C}_kCv\mathbf{C}_v 分别表示由 InfLLM 动态构建的 key-value 缓存 C\mathbf{C} 中的 keyvalue 向量部分。

    • Concat 函数用于将 key 向量 (Ck\mathbf{C}_kKX\mathbf{K}\mathbf{X}) 拼接起来,以及将 value 向量 (Cv\mathbf{C}_vVX\mathbf{V}\mathbf{X}) 拼接起来,形成完整的 keyvalue 序列,供 Attn 函数使用。

    • Attn 函数代表标准的自注意力计算,它会计算 QX\mathbf{Q}\mathbf{X} 与拼接后的 key 序列的点积,并通过 softmax 得到注意力权重,然后与拼接后的 value 序列进行加权求和,得到最终的注意力输出 O\mathbf{O}

      特别说明: 如果查找操作 f()f(\cdot) 总是返回空集,那么 InfLLM 将退化为 LM-InfiniteStreaming-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)

InfLLMevicted tokens E\mathbf{E} 划分为多个记忆单元,每个单元包含 lbsl_{bs} 个连续词元。为了高效地进行相关性查找,每个块需要一个能够代表其语义内容的单元表示。直接使用整个块的 key-value 向量作为表示会导致计算量过大。因此,InfLLM 采取了选择块内“代表性词元”的方法。

  1. 代表性词元选择: InfLLM 不引入额外的训练参数,而是利用 LLM 内部已有的注意力机制来识别块内最重要的词元。对于第 mm 个词元,其代表性分数(representative scorermr_m 定义为: rm=1lLj=1lLqm+jkm r_m = \frac{1}{l_L} \sum_{j=1}^{l_L} \mathbf{q}_{m+j} \cdot \mathbf{k}_m

    • qm+j\mathbf{q}_{m+j} 是第 (m+j)(m+j) 个词元的 query 向量。
    • km\mathbf{k}_m 是第 mm 个词元的 key 向量。
    • lLl_L 是局部窗口的大小。 这个公式直观地表示了第 mm 个词元在其对应局部窗口内对其他词元的“影响程度”或“重要性”。分数越高,表示该词元在该局部窗口内越显著。通过计算每个词元的 rmr_m,并选择分数最高的 rkr_k 个词元作为该块的代表性词元。 对于一个记忆块 B\mathbf{B},其包含 lbsl_{bs}key-value{(kˉjB,vjB)}j=1lbs\{ (\bar{\mathbf{k}}_j^B, \mathbf{v}_j^B) \}_{j=1}^{l_{bs}}。通过上述方法,从 B\mathbf{B} 中选出 rkr_k 个代表性词元,形成该块的代表性 key-value 对集合 R(\mathbf{B}) = \{ (\mathbf{k}_{b_j}^B, \mathbf{v}_{b_j}^B) \}_{j=1}^{r_k}。这些代表性词元将用于后续的相关性计算。
  2. 记忆查找与相关性计算: 在记忆查找阶段,InfLLM 需要从所有存储的记忆单元中选择 kmk_m 个与当前词元最相关的单元加载。为了衡量记忆块 B\mathbf{B} 与当前词元 X\mathbf{X} 之间的相关性,InfLLM 计算它们之间的相似度分数(relevance score): sim(X,B)=i=1lXj=1rkqi+lPkbjB \mathrm{sim}(\mathbf{X}, \mathbf{B}) = \sum_{i=1}^{l_X} \sum_{j=1}^{r_k} \mathbf{q}_{i+l_P} \cdot \mathbf{k}_{b_j}^B

    • X\mathbf{X} 是当前需要编码的词元序列(长度为 lXl_X)。
    • B\mathbf{B} 是一个记忆块。
    • qi+lP\mathbf{q}_{i+l_P} 是当前词元序列中第 ii 个词元的 query 向量。
    • kbjB\mathbf{k}_{b_j}^B 是记忆块 B\mathbf{B} 中第 jj 个代表性词元的 key 向量。 这个公式通过计算当前词元的 query 向量与记忆块中所有代表性词元的 key 向量的点积之和,来衡量它们的语义相关性。选择相关性分数最高的 kmk_m 个记忆单元加载到当前 key-value cache 中。

4.2.2.2. 位置编码 (Positional Encoding)

现有的 LLM 训练通常使用有限数量的位置编码,导致直接应用于长序列时出现 out-of-domain 问题。此外,InfLLM 的当前 key-value cache 是由不连续的文本块(初始词元、相关记忆单元、局部词元)组成的,直接分配连续的位置编码会造成信息不匹配。

受先前工作启发,InfLLM 采取了一种简化的位置编码策略:对于所有超出局部窗口大小的词元,它们被赋予相同的(固定的)位置编码。

  • 具体来说,context memory 单元中的词元与当前词元之间的距离被统一设置为 局部窗口大小 lLl_L
  • 解释: 这种策略旨在避免模型因看到训练时未见过的超大位置索引而崩溃。通过将所有远距离的记忆单元都视为与当前词元距离为 lLl_L,模型可以在其熟悉的相对位置范围内进行计算。尽管这种处理方式似乎忽略了远距离记忆单元之间的相对位置信息,但论文在附录中通过实验(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 中的记忆单元会根据其频率分数 sbs_b 进行排序。分数最低的单元会被卸载回 CPU memory。频率分数 sbs_b 的更新公式如下: sb=sbd+j=1lXi=1lbsattention_score(qj+lP,ki) s_b = s_b \cdot d + \sum_{j=1}^{l_X} \sum_{i=1}^{l_{bs}} \mathrm{attention\_score}(\mathbf{q}_{j+l_P}, \mathbf{k}_i)
      • sbs_b:记忆块 bb 的频率分数。
      • dd:衰减系数(decay coefficient),用于反映之前查找的影响。
      • lXl_X:当前计算步骤中涉及的词元数量。
      • lbsl_{bs}:记忆单元的块大小。
      • attention_score(q,k)\mathrm{attention\_score}(\mathbf{q}, \mathbf{k}):在执行注意力计算后,query 向量 q\mathbf{q}key 向量 k\mathbf{k} 之间的注意力分数(值在 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. 评估指标

论文中使用的评估指标并未在正文中详细定义,但从任务类型和结果表格中的列名可以推断。根据 ∞-BenchLongBench 的官方说明以及常见的 LLM 评估实践,以下是对可能使用的评估指标的概念定义、数学公式和符号解释的补充:

5.2.1. 准确率 (Accuracy)

  • 概念定义: 准确率是最常见的分类和检索任务评估指标,它衡量模型正确预测的样本数量占总样本数量的比例。对于 Retrieve.PassKeyRetrieve.NumberRetrieve.KVChoiceMath.Find 等具有明确正确答案的离散型任务,准确率是直接衡量模型性能的指标。
  • 数学公式: Accuracy=Number of Correct PredictionsTotal Number of Predictions \mathrm{Accuracy} = \frac{\text{Number of Correct Predictions}}{\text{Total Number of Predictions}}
  • 符号解释:
    • 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-LROUGE-1/2F1 score
  • 数学公式 (以 ROUGE-N Recall 为例): ROUGENRecall=sentenceReferencen-gramsentenceCountmatch(n-gram)sentenceReferencen-gramsentenceCount(n-gram) \mathrm{ROUGE-N_{Recall}} = \frac{\sum_{\text{sentence} \in \text{Reference}} \sum_{n\text{-gram} \in \text{sentence}} \mathrm{Count_{match}}(n\text{-gram})}{\sum_{\text{sentence} \in \text{Reference}} \sum_{n\text{-gram} \in \text{sentence}} \mathrm{Count}(n\text{-gram})}
    • 其中,Count_match(n-gram) 是指模型生成摘要和参考摘要中共同出现的 n-gram 的数量。
    • Count(n-gram) 是指参考摘要中 n-gram 的数量。 ROUGE-N PrecisionROUGE-N F1-score 也可以类似计算。
  • 符号解释:
    • n-gram:由文本中 nn 个连续词元组成的序列。
    • 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): F1=2PrecisionRecallPrecision+Recall \mathrm{F1} = 2 \cdot \frac{\mathrm{Precision} \cdot \mathrm{Recall}}{\mathrm{Precision} + \mathrm{Recall}} 其中, Precision=Number of Correct Tokens in PredictionTotal Number of Tokens in Prediction \mathrm{Precision} = \frac{\text{Number of Correct Tokens in Prediction}}{\text{Total Number of Tokens in Prediction}} Recall=Number of Correct Tokens in PredictionTotal Number of Tokens in Reference \mathrm{Recall} = \frac{\text{Number of Correct Tokens in Prediction}}{\text{Total Number of Tokens in Reference}}
  • 符号解释:
    • Precision:模型答案中正确词元的比例。

    • Recall:参考答案中被模型捕获的正确词元的比例。

    • Number of Correct Tokens in Prediction:模型预测答案中与参考答案重叠的词元数量。

    • Total Number of Tokens in Prediction:模型预测答案的总词元数量。

    • Total Number of Tokens in Reference:参考答案的总词元数量。

      由于论文中表格直接给出了 R.PKR.NumR.KVChoiceQASumMath.F 的数值,且大部分任务(除 Sum 外)是离散型判断或抽取任务,因此可以合理推断,这些任务的评估指标很可能就是准确率 (Accuracy),而 Sum 任务则会是 ROUGE 分数,QA 可能是 F1EM。在论文中,这些数值通常被统一为百分比或分数,以方便比较。

5.3. 对比基线

论文将 InfLLM 与多种具有代表性的长上下文处理方法进行了比较:

  1. 原始模型 (Original Models):

    • Mistral-7B-inst-v0.2 (最大上下文 32K 词元)
    • Llama-3-8B-Instruct (最大上下文 8K 词元)
    • 这些模型代表了在不进行任何长上下文扩展时 LLM 的基准性能,它们在遇到超长序列时通常表现不佳。
  2. 位置降尺度和重用方法 (Position Downscaling and Reusing):

    • NTK-aware scaled RoPE (NTK) (LocalLLaMA, 2023):通过非线性插值方法调整 RoPE 的旋转基,以扩展上下文窗口。
    • SelfExtend (Jin et al., 2024):通过重用相邻词元的位置 ID,使扩展的相对位置保持在训练上下文窗口的范围内。
    • 这些方法旨在解决 LLM 在未见长度上的“域外问题”。
  3. 滑动窗口方法 (Sliding Window):

    • LM-Infinite (Infinite) (Han et al., 2023):采用滑动窗口机制直接丢弃远距离上下文,流式读取超长序列。
    • StreamingLLM (Stream) (Xiao et al., 2023):同样采用滑动窗口机制,只关注局部上下文。
    • 这类方法效率高,但无法捕获长距离依赖。
  4. Key-Value 淘汰方法 (Key-Value Eviction):

    • H2O (Zhang et al., 2023b):一种广泛使用的 key-value 淘汰方法,旨在通过丢弃无用的 key-value 向量来减少计算复杂度。
    • 论文指出,这类方法主要用于提高计算效率,但由于未见位置编码问题,通常无法在不额外训练的情况下外推上下文窗口。
  5. 持续训练模型 (Models with Continual Training):

    • Llama-3-8B-Instruct-Gradient-1048k (Llama-1M):一个在长文本数据和聊天数据集上进行了额外微调的模型,其上下文窗口扩展到 1048K
    • 这个基线用于与 InfLLM 进行性能和效率的对比,以证明 InfLLM 无需训练即可达到竞争性能的优势。
  6. 检索增强生成 (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 lbsl_{bs}):128 词元。
    • 代表性词元数量(number of representative tokens rkr_k):4
    • 局部窗口大小(local window size lLl_L):4K 词元(对于 Mistral-basedLlama-3-based InfLLM 均是)。
    • 加载的相关记忆单元数量:Mistral-based InfLLM 加载 96 个,Llama-3-based InfLLM 加载 32 个。
    • 初始词元数量(initial tokens):128 词元(LM-Infinite, StreamingLLM, InfLLM 均采用此设置,用于覆盖系统提示和任务描述)。
  • 计算加速: 采用 FlashAttention (Dao, 2023) 加速所有基线模型的实验。
  • 硬件: NVIDIA A100A800 GPU
  • 精度: 所有实验均采用半浮点精度(half-float precision)。
  • 附录中的额外细节:
    • context memory 在所有 LLM 层构建。
    • GPU cache 大小设置为 32(是每个步骤加载单元数量的两倍)。
    • 频率分数衰减系数(frequency score decay coefficient dd):0.1
    • 在扩展到 1,024K 上下文的实验中,编码块大小设置为 2048,代表性词元数量设置为 1,以加速实验。

6. 实验结果与分析

6.1. 核心结果分析

6.1.1. 在 ∞-Bench 上的表现

以下是原文 Table 1 的结果,展示了 InfLLM 和基线模型在 ∞-Bench 上的表现。

以下是原文 Table 1 的结果:

WindowStreamingR.PKR.NumR.KVChoiceQASumMath.FAvg.
Mistral-based Models (7B)
Mistral32K28.828.814.844.512.925.920.625.2
NTK128K×100.086.819.240.216.920.326.944.3
SelfExtend128K100.0100.015.642.817.318.819.144.8
Infinite32K28.828.80.442.811.422.516.321.6
Streaming32K28.828.50.242.411.522.116.921.5
H2032K8.64.82.648.015.624.426.918.7
InfLLM16K100.096.196.843.715.725.825.757.7
Llama-3-based Models (8B)
Llama-38K8.57.86.244.115.524.721.718.4
NTK128K×X0.00.00.00.00.46.42.61.3
SelfExtend128K×100.0100.00.219.78.614.722.638.0
Infinite8K6.87.60.241.514.620.820.616.0
Streaming8K8.58.30.440.614.320.421.416.3
H208K2.52.40.00.00.72.86.02.1
InfLLM8K100.099.05.043.719.524.323.745.0

分析:

  1. InfLLM 的显著优势: 无论是基于 Mistral 还是 Llama-3InfLLM 在总体平均性能 (Avg.) 上都远超所有其他基线。尤其是在 Mistral 上达到了 57.7Llama-3 上达到了 45.0。这表明 InfLLM 能够有效弥补原始 LLM 在处理长上下文时的不足。
  2. 滑动窗口机制的局限性: InfiniteStreaming 模型(如 Mistral-basedInfinite 21.6, Streaming 21.5)虽然能处理长序列,但由于直接丢弃了远距离上下文,其性能甚至低于原始 Mistral 模型(25.2)。这凸显了 InfLLMcontext memory 在捕获长距离依赖方面的关键作用。
  3. 位置编码外推方法的不足: NTKSelfExtend 方法虽然能够将上下文窗口扩展到 128K,但在某些任务上表现不佳,例如 Mistral-basedSelfExtendSum 任务上表现为 18.8,低于原始 Mistral25.9。这可能是因为它们主要解决了 out-of-domain 问题,但未能有效解决噪声上下文引起的 distraction 问题。Llama-3-based NTK 甚至在多个任务上得到 0.0 的极低分数,表明其稳定性可能不足。
  4. KV 淘汰方法的局限性: H2O 的表现普遍非常差(Mistral-based18.7Llama-3-based2.1),这验证了论文的观点,即 KV 淘汰方法虽然旨在提高效率,但不能在不额外训练的情况下有效外推上下文窗口。
  5. InfLLM 在特定任务上的表现: InfLLMRetrieve.PassKeyRetrieve.NumberRetrieve.KV 等检索任务上表现尤为突出,例如 Mistral-based InfLLMR.PKR.KV 上分别达到 100.096.8。这直接证明了其 context memory 能够准确地从海量信息中检索出关键信息。

6.1.2. 与持续训练模型的比较

以下是原文 Table 2 的结果,比较了 InfLLM 与经过持续训练的 Llama-3-8B-Instruct-Gradient-1048k (Llama-1M) 的性能和效率。

以下是原文 Table 2 的结果:

Train-FreeR.PKR.NumR.KVChoiceQASumMath.FVRAMTime
Llama-1M100.099.823.251.513.618.518.376.6G40.4s
InfLLM100.099.05.043.719.524.323.726.3G26.7s
Llama-1M+InfLLM100.0100.055.839.320.317.131.426.3G26.7s

分析:

  1. 性能对比:
    • InfLLM (训练无关) 在 R.PKR.Num 上与 Llama-1M (持续训练) 表现相当,甚至在 QASumMath.F 上超越了 Llama-1M。这有力地支持了论文的核心论点:LLMs 即使未经额外训练,也具有理解极长序列的内在能力,InfLLM 只是有效地激活了这种能力。
    • Llama-1M 需要 512GPU 进行持续训练,这对于大多数研究者来说是无法承受的,而 InfLLM 无需训练,具有极高的实用性。
  2. 效率优势:
    • InfLLMVRAM 使用量仅为 Llama-1M34% (26.3G vs 76.6G)。
    • InfLLM 的时间消耗降低了 34% (26.7s vs 40.4s)。
    • 论文还提到,在 256K 词元长度时,full-attention 基线模型会因为内存不足而失败,而 InfLLM 可以在单个 GPU 上高效处理高达 1024K 词元的序列。这证明了 InfLLM 在资源受限环境下的强大实用性。
  3. InfLLM 与持续训练模型的结合: Llama-1M+InfLLM 的结果表明,InfLLM 也可以与持续训练的模型结合,并且在某些任务(如 QA, Math.F)上带来进一步的性能提升,同时保持了 InfLLM 的效率优势。这表明 InfLLM 还可以作为一种提高推理速度的通用方法。

6.1.3. 与检索增强生成 (RAG) 的比较

以下是原文 Table 3 的结果,比较了 InfLLMRAG 在三个上下文检索任务上的表现。

以下是原文 Table 3 的结果:

Task| R.PKR.NumR.KV
RAG-E589.265.413.2
InfLLM100.096.196.8

分析:

  1. InfLLM 性能优势: InfLLM 在所有三个检索任务上都显著优于 RAG-E5 模型。尤其在 R.PKR.KV 任务上,InfLLM 接近或达到满分(100.096.8),而 RAG-E5 分别只有 89.213.2
  2. InfLLM 的泛化能力: 结果表明,即使没有额外数据或训练,InfLLM 也表现出卓越的泛化能力。
  3. RAG 的局限性:
    • RAG 需要额外的检索数据来训练检索模型,并需要微调 LLM 以适应检索增强的输入,而 InfLLM 是完全训练无关的。
    • RAG 模型的性能受其检索组件的限制,并且可能面临 out-of-distribution 问题,导致在训练分布之外的任务上表现不佳。InfLLM 则无特定任务要求,具有更广泛的适用性。

6.1.4. 扩展到 1,024K 上下文

下图是原文 Figure 3 的结果,展示了 InfLLM 在不同长度序列上的表现。

Figure 3: The results on sequences with different lengths.
分析:

  1. InfLLM 的超长上下文处理能力:Retrieve.PassKey 任务中,即使序列长度扩展到 1,024K 词元,InfLLM 仍能保持 100% 的准确率。这充分证明了 InfLLM 在极端长序列中有效捕获长距离依赖和识别关键信息的能力。
  2. 滑动窗口方法的局限性: 相比之下,LM-Infinite 的性能随着序列长度的增加而迅速下降。这证实了仅依赖局部窗口而丢弃远距离上下文的方法在长距离依赖任务中的根本性缺陷。
  3. 对噪声的鲁棒性: 这个实验也间接表明 InfLLM 能够有效过滤长序列中的噪声,精准定位所需信息。

6.2. 消融实验/参数分析

6.2.1. 记忆设置的影响 (The Impact of Memory Settings)

下图是原文 Figure 2 的结果,展示了 InfLLM 在不同代表性词元数量、选择单元数量和记忆单元大小下的性能。

Figure 2: Extra studies about InfLLM. Here, (a), (b), and (c) investigate the impact of the context memory under different numbers of representative tokens, different numbers of selected units, and memory unit sizes, respectively.
分析:

  • 不同代表性词元数量 (Figure 2a):
    • 随着代表性词元数量的增加(从 14),模型性能呈现上升趋势,表明更多的代表性词元能更好地捕捉记忆单元的语义内容。
    • 当数量达到 8 时,性能略有下降。这可能是由于引入了语义不相关的词元作为单元表示,反而增加了噪声。这暗示了存在一个最优的代表性词元数量,并且未来的工作可以探索更有效和强大的单元表示方法。
  • 不同选择单元数量 (Figure 2b):
    • 随着选择的记忆单元数量增加(从 132),模型性能显著提升。这表明更多的记忆单元能提供更全面的相关上下文,提高了信息召回率。
    • 然而,更多的单元也意味着更高的内存调度时间和注意力计算时间。因此,提升查找准确性以减少所需单元数量,是提高效率的关键。
  • 不同记忆单元大小 (Figure 2c):
    • 记忆单元大小对性能的影响因任务而异。例如,对于 Retrieve.KV 任务,64128 可能是较好的选择;对于 Math.Find 任务,32 可能是最优的。
    • 这说明记忆单元的语义连贯性很重要。过大的单元可能导致查找不精确,过小的单元则增加查找开销。
    • 目前使用启发式规则分割上下文可能导致次优性能,未来研究方向是探索动态分割上下文的方法。

6.2.2. 消融研究 (Ablation Study)

以下是原文 Table 4 的结果,展示了关于动态记忆查找和单元表示的消融研究。

以下是原文 Table 4 的结果:

TaskR.KVMath.FQA
InfLLM96.825.715.7
Decoding-Only85.226.312.0
w/o Lookup0.416.311.4
Mean Repr84.625.114.9

分析:

  • 上下文记忆查找 (Context Memory Lookup):
    • w/o Lookup (无任何记忆查找):性能急剧下降,尤其在 R.KV 任务上仅为 0.4。这验证了 context memory 对于长文本理解的不可或缺性,以及在编码和解码阶段都集成远距离上下文的重要性。模型需要整合长距离上下文来生成连贯的记忆,并从中收集有用信息以生成正确答案。
    • Decoding-Only (仅在解码阶段查找):性能相较于完整 InfLLM 有所下降(R.KV96.8 降至 85.2QA15.7 降至 12.0)。这表明远距离上下文信息对于输入编码和输出生成两个阶段都至关重要,模型需要全面的长距离上下文才能进行有效的理解和推理。
  • 单元表示 (Unit Representation):
    • Mean Repr (使用平均 key 向量作为单元表示):虽然性能略低于 InfLLM 的代表性词元选择方法(R.KV96.8 降至 84.6),但仍能达到竞争性水平。这说明 LLM 原始的注意力向量本身就具有有效进行相关性计算的能力。这也暗示了未来可以进一步探索更高效的单元表示方法来提升模型性能。

6.3. 在 LongBench 上的表现 (附录 C.2)

以下是原文 Table 5 的结果,展示了 InfLLM 和基线模型在 LongBench 上的表现。

以下是原文 Table 5 的结果:

WindowNQAQasperMFQAHQA2WikiMQAMusique
Mistral-based Models (7B)
Mistral32K22.0629.1647.6537.5321.9619.03
Infinite6K18.4430.0239.0532.0222.2715.81
Streaming6K17.9230.0539.0932.1821.8314.71
InfLLM6K22.1229.3347.4236.5622.3117.68
InfLLM12K23.0329.5247.6239.5323.6118.92
Llama-3-based Models (8B)
Llama-38K19.8542.3641.0347.3839.2022.96
Infinite8K19.3942.8040.4443.7737.8918.33
Streaming8K20.0542.4639.5443.6937.8919.68
InfLLM8K22.6443.7049.0349.0435.6126.06
WindowGovReportQMSumMultiNewsTRECTQASAMSum
Mistral-based Models (7B)
Mistral32K31.1223.8726.6271.0085.9742.29
Infinite6K29.7421.9226.6570.0085.2241.60
Streaming6K29.8321.9426.6470.0085.5741.31
InfLLM6K31.0323.4926.7069.0086.6742.52
InfLLM12K31.3723.7726.6671.0087.3441.80
Llama-3-based Models (8B)
Llama-38K29.9421.4527.5174.0090.5042.30
Infinite8K29.2521.4127.6274.0090.0841.72
Streaming8K29.1721.3327.5673.5090.0841.55
InfLLM8K30.7622.7027.5773.5090.9142.43
WindowPsgCountPsgRetrievalLCCRepoBench-PAvg.
Mistral-based Models (7B)
Mistral32K3.9586.9457.4254.1443.78
Infinite6K2.0842.8057.1253.4339.07
Streaming6K2.5042.1755.3851.4638.67
InfLLM6K2.8764.0056.6752.9741.90
InfLLM12K3.0187.4256.6952.0944.02
Llama-3-based Models (8B)
Llama-38K8.5062.5060.8349.1444.73
Infinite8K4.5050.0060.1248.6243.03
Streaming8K5.0049.0060.3548.9542.99
InfLLM8K7.1784.0059.8846.4846.95

分析:

  1. InfLLM 的持续优越性: InfLLMLongBench 的各种任务中,普遍优于其他能够处理流式输入的模型(Infinite, Streaming)。这再次证明了 context memory 提供的上下文信息能够有效增强模型性能。
  2. 滑动窗口方法的不足: 当以 Llama-3 作为基础模型时,StreamingLLMLM-Infinite 的性能与原始 Llama-3 相当甚至更差。这进一步强调了滑动窗口注意力机制虽然能扩展上下文窗口,但由于丢弃了长距离上下文,无法实现有效的长序列理解。
  3. 对噪声的过滤能力: Mistral 能够处理长达 32K 的文本,覆盖了 LongBench 的大部分实例。相比之下,InfLLM 在仅 12K 的窗口大小下,平均性能与 Mistral 相当甚至更优(Mistral 平均 43.78InfLLM 平均 44.02)。这表明 InfLLM 能够有效过滤长上下文中的噪声,从而实现更好的长序列理解。

6.4. 在 Vicuna 上的实验 (附录 C.3)

以下是原文 Table 6 的结果,展示了 Vicuna-based 模型的实验结果。

以下是原文 Table 6 的结果:

R.PKR.NumR.KVMath.F
Vicuna5.084.411.4011.71
InfLLM99.1581.690.6011.14

分析:

  1. 显著的上下文扩展: InfLLM 成功地将 Vicuna(最大长度仅 4K)的上下文长度扩展到 128K。在 Retrieve.PasskeyRetrieve.Number 任务上,InfLLM 取得了显著的性能提升,准确率分别从 5.08 提升到 99.15 和从 4.41 提升到 81.69。这再次证明了 InfLLM 能够有效地将短上下文模型泛化到长上下文场景。
  2. 局限性: 然而,在 Retrieve.KVMath.Find 任务上,InfLLM 未能显示性能增益,甚至在 R.KV 任务上略有下降。
  3. 潜在原因: 论文解释说,这可能是因为 Vicuna 的隐藏向量(hidden vectors)在处理极长文本中的噪声时,过滤能力有限。这使得 context memory 在更复杂的任务(如 Retrieve.KVMath.Find)中难以有效定位相关信息。这暗示了基础模型的内在能力对 InfLLM 的有效性有一定影响,也为未来的研究指明了方向:设计更强大的记忆机制来应对基础模型在噪声过滤方面的局限性。

7. 总结与思考

7.1. 结论总结

本文提出了 InfLLM,一个无需训练(training-free)的方法,旨在显著提升大型语言模型(LLMs)处理长序列的能力。InfLLMsliding window attention 机制与一个高效的上下文记忆(context memory)模块相结合。该记忆模块通过将远距离上下文存储在块级记忆单元中,并利用代表性词元进行相关性查找,从而允许 LLMs 在有限的上下文窗口内,高效且精确地捕获长距离依赖。

实验结果在 ∞-BenchLongBench 这两个广泛使用的长文本基准测试上证明了 InfLLM 的卓越有效性。它使在数千个词元序列上预训练的 LLMs 达到了与在长序列上持续训练的竞争基线相当甚至更优的性能,同时大幅降低了计算和内存成本。尤其在 1,024K 词元的超长序列上,InfLLM 仍能保持 100% 的准确率,展现了其在处理极端长序列时的强大能力。此外,InfLLM 在上下文检索任务中也优于 Retrieval-Augmented Generation (RAG) 方法,凸显了其训练无关和广泛适用性的优势。

InfLLM 的成功揭示了 LLMs 在无需额外训练的情况下,理解和推理极长序列的内在潜力。

7.2. 局限性与未来工作

论文作者指出了 InfLLM 存在的局限性,并提出了未来的研究方向:

  1. CPU 内存使用: InfLLM 将大量过去 key-value (KV) 缓存存储在 CPU memory 中,这增加了 CPU memory 的使用量。

    • 未来工作: 可以通过集成 KV cache quantization 等技术来减少 CPU memory 需求,进一步优化资源利用。
  2. 推理速度提升空间: 尽管 InfLLM 显著降低了处理长文本的计算开销,但仍有进一步加速的潜力。

    • 未来工作: 可以将其与 llama.cppvLLM (Kwon et al., 2023) 等高效推理框架结合,以进一步增强 InfLLM 的推理速度。
  3. 记忆模块的训练: 当前 InfLLM 的记忆模块是完全训练无关的。

    • 未来工作: 探索对 context memory module 进行高效训练,以进一步提升模型的性能。这可能包括学习更优的单元表示或查找策略。
  4. KV cache 压缩: 结合 key-value cache 压缩方法,可以进一步降低计算和内存成本。

  5. 更强大的单元表示: 实验中发现,代表性词元数量过多可能引入噪声,且平均表示也能取得竞争性结果。

    • 未来工作: 探索更有效和强大的记忆单元表示方法,以进一步提高模型性能。
  6. 动态上下文分割: 记忆单元的最佳大小可能因任务而异,启发式分割可能导致次优性能。

    • 未来工作: 研究如何动态地分割上下文以形成记忆单元,以适应不同任务和文本特性的需求。
  7. 基础模型能力的依赖:Vicuna 上的实验表明,基础模型的隐藏向量过滤噪声能力有限时,InfLLM 在某些复杂任务上可能无法获得性能增益。

    • 未来工作: 设计更强大的记忆机制来克服基础模型在噪声过滤方面的局限性。

7.3. 个人启发与批判

7.3.1. 个人启发

  • “内在能力”的挖掘: 这篇论文最引人注目的一点是其核心发现——“揭示了 LLMs 在无需任何微调的情况下理解极长序列的内在能力”。这提供了一个非常积极的信号:现有的大型、昂贵的 LLMs 可能比我们想象的更有潜力,我们需要的不是无止境地扩大训练数据和模型规模,而是更智能、更高效地“引导”或“激活”它们已有的能力。这种训练无关(training-free)的思路,对于资源有限的研究者和开发者来说,具有巨大的吸引力。
  • 工程与算法的精妙结合: InfLLM 巧妙地结合了 sliding window attention 的工程效率和 context memory 的算法智能。通过块级记忆单元、代表性词元选择和 LRU 缓存管理,它在计算成本、内存效率和性能之间找到了一个非常好的平衡点,为 LLM 的流式长上下文处理提供了实用的解决方案。
  • RAG 的挑战与互补: InfLLM 在某些任务上优于 RAG 的结果很有趣。它表明,对于上下文理解任务,模型内部的“自检索”机制可能比外部检索器更直接、更高效,因为无需处理检索器本身的训练、 out-of-distributionLLM 与检索器结果的适配问题。这并非意味着 RAG 无用,而是提示我们,两种机制可能各有侧重,甚至可以互补:InfLLM 专注于利用模型内在能力处理核心上下文,RAG 则可用于引入外部知识。
  • LLM 训练的启示: 如果 LLMs 真的具备这种内在的长上下文理解能力,那么未来的 LLM 训练策略是否可以更多地关注于如何更好地编码和提取这种能力,而不是仅仅通过简单地增加上下文窗口长度来训练?这可能会导致更高效、更具泛化性的 LLM

7.3.2. 批判与潜在改进

  • 代表性词元选择的局限性: 论文中用于选择代表性词元的 representative score 公式 rm=1lLj=1lLqm+jkmr_m = \frac{1}{l_L} \sum_{j=1}^{l_L} \mathbf{q}_{m+j} \cdot \mathbf{k}_m 依赖于局部窗口内的注意力分数。虽然它是训练无关的,但这种启发式方法可能无法捕获复杂的语义重要性。例如,一个在局部窗口内不显眼的词元,可能在全局语境中具有关键的意义。在 Vicuna 实验中的局限性也可能部分源于此。
    • 改进方向: 可以探索更复杂的、基于语义嵌入的聚类方法,或者使用小型、轻量级的编码器来学习更鲁棒的块级表示,即使这需要微量的额外训练。
  • 位置编码的简化处理: 将所有远距离记忆单元的位置编码统一设置为 lLl_L 是一种有效的简化,但这也意味着模型失去了这些远距离记忆单元之间更精细的相对位置信息。对于需要精确时间或顺序推理的长文本任务,这可能是一个潜在的瓶颈。
    • 改进方向: 可以在不引入 out-of-domain 问题的前提下,探索一种分层或稀疏的位置编码策略,以在记忆单元内部或记忆单元之间保留更多的相对位置信息。
  • “注意力分散问题”的根本解决: 尽管 InfLLM 通过选择性地加载相关记忆单元来缓解 distraction issue,但其内部的 sliding window attention 仍然可能在局部窗口内受到噪声干扰。
    • 改进方向: 可以考虑将 InfLLM 与更先进的稀疏注意力机制或动态注意力机制结合,以在局部和全局层面都实现更精细的注意力控制。
  • 记忆查找的开销: 尽管采用了块级查找和 LRU 卸载机制,但在极端长序列下,每次都要计算与所有(或大量)记忆单元的相关性分数,这仍然可能带来可观的计算开销,尤其是在 CPU memoryGPU memory 之间的数据传输。
    • 改进方向: 引入分层索引或更高效的相似性搜索算法(如 approximate nearest neighbor, ANN)来加速记忆查找过程。
  • 评估任务的代表性: 尽管 ∞-BenchLongBench 是很好的基准,但长文本理解的复杂性远不止这些任务。例如,涉及复杂推理、多跳问答或需要理解多文档之间关系的任务,对 context memory 的语义组织和推理能力提出了更高的要求。
    • 改进方向: 在更具挑战性、更复杂的长文本推理任务上进行评估,以更全面地验证 InfLLM 的鲁棒性和能力。

      总而言之,InfLLM 提供了一个兼顾实用性和高性能的 training-free 解决方案,为 LLM 的长上下文处理开辟了新的思路。尽管存在一些可改进的方面,但其核心思想和实验结果无疑为该领域的发展贡献了重要力量。

相似论文推荐

基于向量语义检索推荐的相关论文。

暂时没有找到相似论文。