论文状态:已完成

Efficient Streaming Language Models with Attention Sinks

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

TL;DR 精炼摘要

在流式应用中部署大型语言模型(LLMs),传统方法面临内存消耗大和泛化能力弱的问题。论文发现“注意力汇聚”现象,即保留初始词元的KV状态可以提高性能。基于此,提出StreamingLLM框架,允许模型无微调处理无限序列,显著提升了处理速度和效率。

摘要

Deploying Large Language Models (LLMs) in streaming applications such as multi-round dialogue, where long interactions are expected, is urgently needed but poses two major challenges. Firstly, during the decoding stage, caching previous tokens' Key and Value states (KV) consumes extensive memory. Secondly, popular LLMs cannot generalize to longer texts than the training sequence length. Window attention, where only the most recent KVs are cached, is a natural approach -- but we show that it fails when the text length surpasses the cache size. We observe an interesting phenomenon, namely attention sink, that keeping the KV of initial tokens will largely recover the performance of window attention. In this paper, we first demonstrate that the emergence of attention sink is due to the strong attention scores towards initial tokens as a "sink" even if they are not semantically important. Based on the above analysis, we introduce StreamingLLM, an efficient framework that enables LLMs trained with a finite length attention window to generalize to infinite sequence lengths without any fine-tuning. We show that StreamingLLM can enable Llama-2, MPT, Falcon, and Pythia to perform stable and efficient language modeling with up to 4 million tokens and more. In addition, we discover that adding a placeholder token as a dedicated attention sink during pre-training can further improve streaming deployment. In streaming settings, StreamingLLM outperforms the sliding window recomputation baseline by up to 22.2x speedup. Code and datasets are provided at https://github.com/mit-han-lab/streaming-llm.

思维导图

论文精读

中文精读

1. 论文基本信息

1.1. 标题

具有注意力汇聚的高效流式语言模型 (Efficient Streaming Language Models with Attention Sinks)

1.2. 作者

Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, Mike Lewis

隶属机构:

  • Guangxuan Xiao, Song Han: 麻省理工学院 (Massachusetts Institute of Technology)
  • Yuandong Tian, Mike Lewis: Meta AI
  • Beidi Chen: 卡内基梅隆大学 (Carnegie Mellon University)
  • Song Han: 英伟达 (NVIDIA)

1.3. 发表期刊/会议

预印本 (arXiv)。

声誉与影响力: arXiv 是计算机科学和物理学等领域重要的预印本平台,许多前沿研究成果在此首发,具有较高的学术关注度。该论文被 NVIDIA TensorRT-LLM、Intel Extension for Transformers、HuggingFace Transformers 等多个主流 LLM 服务解决方案采纳,显示了其在工业界和学术界的显著影响力。

1.4. 发表年份

2023年

1.5. 摘要

部署大型语言模型 (Large Language Models, LLMs) 于多轮对话等流式应用中是当前急需解决的问题,但面临两大挑战:

  1. 在解码阶段,缓存先前词元 (tokens) 的键和值 (Key and Value, KV) 状态会消耗大量内存。

  2. 流行的 LLMs 无法泛化 (generalize) 到超出其训练序列长度的文本。

    窗口注意力 (Window attention) 是一种自然的方法,它只缓存最新的 KV 状态,但论文发现当文本长度超过缓存大小时(即初始词元被逐出时),其性能会急剧下降。研究观察到一个有趣的现象,即“注意力汇聚 (attention sink)”:保留初始词元的 KV 状态能在很大程度上恢复窗口注意力的性能。

本文首先证明了注意力汇聚的出现是由于即使初始词元在语义上不重要,它们也作为“汇聚点 (sink)”吸引了强大的注意力分数。基于此分析,论文引入了 StreamingLLM,一个高效的框架,使经过有限长度注意力窗口训练的 LLMs 能够泛化到无限序列长度,且无需任何微调 (fine-tuning)。实验表明,StreamingLLM 能够使 Llama-2、MPT、Falcon 和 Pythia 等模型稳定高效地进行语言建模,处理多达 400 万甚至更多的词元。此外,研究发现,在预训练 (pre-training) 期间添加一个占位符词元 (placeholder token) 作为专用的注意力汇聚 (dedicated attention sink) 可以进一步改善流式部署。在流式设置中,StreamingLLM 的速度比滑动窗口重计算 (sliding window recomputation) 基线快 22.2 倍。

1.6. 原文链接

原文链接: https://arxiv.org/abs/2309.17453 PDF 链接: https://arxiv.org/pdf/2309.17453v4.pdf

2. 整体概括

2.1. 研究背景与动机

大型语言模型 (LLMs) 已广泛应用于对话系统、文档摘要、代码补全和问答等自然语言处理 (Natural Language Processing, NLP) 领域。为了充分发挥其潜力,LLMs 应该能够高效准确地生成长序列文本。例如,一个理想的聊天机器人助手应该能够稳定地处理持续一整天的对话内容。然而,当前的 LLMs 在泛化到比它们预训练时更长的序列长度时面临巨大挑战。例如,Llama-2 模型的预训练长度通常为 4K 个词元。

这种限制源于 LLMs 在预训练阶段受到的注意力窗口 (attention window) 约束。尽管研究人员在扩展窗口大小、提高训练和推理效率方面做了大量工作,但可接受的序列长度仍是有限的,这不适用于需要持续运行的流式部署。

在将 LLMs 应用于无限输入流时,会遇到两个主要挑战:

  1. 内存消耗和解码延迟: 在解码阶段,基于 Transformer 的 LLMs 需要缓存所有先前词元 (tokens) 的键和值 (Key and Value, KV) 状态。随着序列长度的增加,这会导致内存使用量过大和解码延迟不断增加。

  2. 长度外推能力有限: 现有模型的长度外推能力有限,即当序列长度超出预训练时设定的注意力窗口大小时,其性能会显著下降。

    直观的解决方案,如窗口注意力 (window attention),只维护一个固定大小的滑动窗口来缓存最新词元的 KV 状态。虽然这能确保在缓存填满后内存使用和解码速度保持恒定,但一旦序列长度超出缓存大小(即最早的词元被逐出),模型性能便会崩溃。另一种策略是滑动窗口重计算 (sliding window with re-computation),它为每个生成的词元重新构建最近词元的 KV 状态。这种方法虽然性能强劲,但由于窗口内的注意力计算复杂度是二次方的,导致其速度非常慢,不适合实际的流式应用。

因此,论文的核心动机是解决 LLMs 在流式应用中面临的性能下降和效率瓶颈问题,尤其是在处理超长上下文时。

2.2. 核心贡献/主要发现

本文的主要贡献和关键发现总结如下:

  1. 揭示“注意力汇聚 (attention sink)”现象: 论文首次观察并证明了自回归 LLMs 中存在“注意力汇聚”现象。即,即使初始词元在语义上不重要,模型也会将不成比例的大量注意力分数分配给它们,将其用作注意力值的“倾倒 (dump)”点。这是由于 SoftMax 函数要求注意力分数总和为一的特性,以及初始词元对所有后续词元都可见的性质。
  2. 提出 StreamingLLM 框架: 基于对注意力汇聚现象的深入理解,论文提出了 StreamingLLM 框架。该框架通过保留少数初始词元的 KV 状态(通常仅需 4 个)以及一个滚动窗口中的最新词元 KV 状态,有效地稳定了注意力计算,从而使 LLMs 能够在不进行任何微调的情况下,从有限的预训练注意力窗口泛化到无限长的序列。
  3. 广泛的性能验证: StreamingLLM 能够使 Llama-2 (7B, 13B, 70B)、MPT (7B, 30B)、Falcon (7B, 40B) 和 Pythia (2.9B, 6.9B, 12B) 等主流 LLM 系列模型稳定高效地处理多达 400 万甚至更长的词元,其语言建模困惑度与需要昂贵重计算的滑动窗口基线相当。
  4. 效率显著提升: 在流式设置中,StreamingLLM 相较于滑动窗口重计算基线,实现了高达 22.2 倍的解码速度提升,同时保持了相似的内存占用。
  5. 预训练优化建议: 论文发现,在预训练阶段为所有训练样本添加一个专用的可学习占位符词元 (learnable placeholder token) 作为注意力汇聚,可以进一步简化和改善流式部署的性能。这种方法使得模型仅需一个汇聚词元即可实现稳定性能,而无需依赖多个初始内容词元。

3. 预备知识与相关工作

3.1. 基础概念

理解本文需要以下基本技术、理论和模型:

  • 大型语言模型 (Large Language Models, LLMs):
    • 概念定义: LLMs 是基于深度学习,特别是 Transformer 架构,通过在海量文本数据上进行预训练而获得的巨型神经网络模型。它们能够理解、生成人类语言,执行多种自然语言处理任务,如问答、文本生成、摘要、翻译等。LLMs 的强大能力来源于其庞大的参数量和训练数据规模,使其能够捕捉语言的复杂模式和知识。
    • Transformer 架构: LLMs 的核心架构,由编码器 (Encoder) 和解码器 (Decoder) 组成,但现代 LLMs(如 GPT 系列和本文讨论的 Llama-2 等)多采用仅解码器 (Decoder-only) 架构。Transformer 引入了自注意力 (Self-Attention) 机制,使其能够并行处理序列,并捕获序列中任意位置之间的依赖关系。
  • 解码阶段 (Decoding Stage):
    • 概念定义: 指的是语言模型根据给定的输入(提示或已生成的序列)逐步生成新的词元(单词、子词或字符)以构建完整输出序列的过程。这是一个自回归 (autoregressive) 过程,即模型在生成每个新词元时,会考虑此前已生成的所有词元作为上下文。
  • 键和值状态 (Key and Value States, KV States):
    • 概念定义: 在 Transformer 的自注意力机制中,每个输入词元都会被投影成三个不同的向量:查询 (Query, Q)键 (Key, K)值 (Value, V)。在生成序列的解码过程中,为了避免重复计算已生成词元的 K 和 V 向量,模型会将它们缓存起来。这个缓存被称为 KV 缓存 (KV Cache)。当生成新的词元时,只需要计算新词元的 Q 向量,然后将其与 KV 缓存中的 K 向量进行点积计算注意力分数,再用这些分数加权 KV 缓存中的 V 向量,从而高效地获得上下文表示。KV 缓存是提高解码效率的关键,但其内存消耗随序列长度线性增长。
  • 注意力窗口 (Attention Window):
    • 概念定义: 在 Transformer 模型中,注意力窗口是指在计算当前词元注意力时,模型能够“看到”或考虑的上下文词元的最大数量。在预训练时,模型通常被设定一个固定的最大序列长度,这个长度就是其注意力窗口大小。超出这个窗口的词元将不被考虑。
  • 困惑度 (Perplexity, PPL):
    • 概念定义: 困惑度是衡量语言模型预测文本质量的重要指标。它表示一个语言模型在给定文本上对下一个词元的预测有多“困惑”或不确定。困惑度越低,说明模型对文本的建模能力越好,预测准确性越高,反之亦然。它通常被认为是交叉熵损失的指数形式。
  • 窗口注意力 (Window Attention):
    • 概念定义: 窗口注意力是一种优化技术,旨在减少 Transformer 在处理长序列时的内存和计算开销。它通过限制每个词元计算注意力时只能关注其附近固定大小的“窗口”内的词元来实现。在自回归解码中,这意味着只缓存最近 LL 个词元的 KV 状态,形成一个滑动窗口。优点是内存和计算效率高,但会丢失窗口外的早期上下文信息。
  • 相对位置编码 (Relative Positional Encoding):
    • 概念定义: 传统的绝对位置编码 (Absolute Positional Encoding) 为序列中的每个位置分配一个唯一的嵌入,而相对位置编码则关注词元之间的相对距离。它允许模型根据两个词元之间的相对位置来计算注意力分数,而不是它们的绝对位置。这使得模型理论上能够更好地泛化到训练时未见过的更长序列。
    • 旋转位置嵌入 (Rotary Position Embeddings, RoPE) (Su et al., 2021):
      • 概念定义: RoPE 是一种通过对查询 (Q) 和键 (K) 向量应用旋转变换来整合相对位置信息的方法。它将位置信息编码到 Q 和 K 向量中,使得点积注意力能够自然地计算词元间的相对位置关系。RoPE 已被 Llama 系列模型广泛采用,以支持长序列处理。
    • 带线性偏差的注意力 (Attention with Linear Biases, ALiBi) (Press et al., 2022):
      • 概念定义: ALiBi 通过直接向注意力分数 (在 SoftMax 之前) 添加一个与词元相对距离成比例的线性偏差来引入位置信息。距离越远,负偏差越大,从而降低了远距离词元的注意力分数。这种方法被认为在长度外推方面表现良好,并被 MPT 等模型采用。
  • SoftMax 函数:
    • 概念定义: SoftMax 函数是一种常用的激活函数,它将一个包含任意实数的向量转换为一个概率分布向量。即,输出向量中的每个元素都在 0 到 1 之间,并且所有元素的总和为 1。在自注意力机制中,SoftMax 函数用于将计算出的注意力对数 (logits) 转换为注意力权重,这些权重表示每个上下文词元对当前词元的重要性,总和为 1。
    • 数学公式: SoftMax(x)i=exij=1Nexj \mathrm { S o f t M a x } ( x ) _ { i } = \frac { e ^ { x _ { i } } } { \sum _ { j = 1 } ^ { N } e ^ { x _ { j } } }
      • 符号解释:
        • SoftMax(x)i\mathrm{SoftMax}(x)_i: SoftMax 函数对输入 xx 的第 ii 个元素的输出,表示第 ii 个词元的概率或注意力权重。
        • ee: 自然对数的底数。
        • xix_i: 输入向量 xx 中的第 ii 个元素,在注意力机制中通常是注意力对数 (attention logit)。
        • j=1Nexj\sum_{j=1}^{N} e^{x_j}: 输入向量 xx 中所有元素的指数化之和,作为归一化因子。
        • NN: 输入向量 xx 的维度或序列中的词元总数。

3.2. 前人工作

关于将 LLMs 应用于长文本的研究主要集中在三个方面:长度外推 (Length Extrapolation)、上下文窗口扩展 (Context Window Extension) 和提高 LLMs 对长文本的利用率 (Improving LLMs' Utilization of Long Text)。

  • 长度外推 (Length Extrapolation):

    • 目标是使在短文本上训练的语言模型在测试时能够处理更长的文本。
    • 主要研究方向是开发 Transformer 模型的相对位置编码 (Relative Positional Encoding) 方法,使其能够超越训练窗口长度进行工作。
    • 旋转位置嵌入 (RoPE) (Su et al., 2021): 通过对每个注意力层的查询 (Q) 和键 (K) 进行变换来整合相对位置。尽管有潜力,但后续研究 (Press et al., 2022; Chen et al., 2023) 表明其在超出训练窗口的文本上表现不佳。
    • 带线性偏差的注意力 (ALiBi) (Press et al., 2022): 基于词元距离对 Q-K 注意力分数进行偏差处理,引入相对位置信息。虽然显示出改进的外推能力,但本文测试表明在文本长度远超训练长度时仍会崩溃。
    • 现状: 当前方法尚未实现无限长度外推,导致没有现有 LLMs 适用于流式应用。
  • 上下文窗口扩展 (Context Window Extension):

    • 目标是扩展 LLM 的上下文窗口,使其在一次前向传播中能处理更多词元。
    • 主要挑战是注意力计算的二次方复杂性,导致训练长上下文 LLM 在计算和内存上都非常困难。
    • 系统级优化: 例如 FlashAttention (Dao et al., 2022; Dao, 2023),加速注意力计算并减少内存占用。
    • 近似注意力方法: 例如 Longformer (Beltagy et al., 2020)BigBird (Zaheer et al., 2020b) 等,通过牺牲模型质量来换取效率。
    • 扩展预训练 LLM: 最近工作多集中于通过位置插值 (position interpolation) 和微调 (fine-tuning) 扩展基于 RoPE 的预训练 LLM (Chen et al., 2023; kaiokendev, 2023; bloc97, 2023; Peng et al., 2023)。
    • 现状: 所有上述技术都只能在有限程度上扩展 LLM 的上下文窗口,未能解决无限输入处理的问题。
  • 提高 LLMs 对长文本的利用率 (Improving LLMs' Utilization of Long Text):

    • 目标是优化 LLMs,使其能更好地捕获和利用上下文信息,而不仅仅是将其作为输入。
    • 现状: Liu et al. 和 Li et al. 指出,前两类方向的进展不一定能转化为对长上下文的有效利用。LLMs 在有效利用长上下文方面仍面临挑战。

3.3. 差异化分析

本文提出的 StreamingLLM 框架主要属于长度外推类别,但其目标和方法与现有工作存在显著差异:

  • 无限长度泛化 vs. 有限上下文扩展: 现有长度外推和上下文扩展方法都致力于将 LLMs 的上下文窗口扩展到有限的更长长度(如 8K、32K 词元),但 StreamingLLM 的目标是实现对无限长度文本的稳定处理,这对于流式应用至关重要。
  • 无需微调 vs. 微调: 许多上下文扩展方法(如位置插值)需要对预训练模型进行额外的微调。StreamingLLM 能够在不进行任何微调的情况下,使现有预训练 LLMs 具备流式处理能力,这大大降低了部署成本和复杂性。
  • 注意力机制的稳定 vs. 上下文扩展: StreamingLLM 的核心在于通过保留“注意力汇聚”来稳定模型在长序列中的注意力机制,从而防止性能崩溃。它并不直接扩展 LLMs 的上下文窗口大小,也不直接增强模型对长文本的记忆和利用能力。
  • 正交性: 论文强调 StreamingLLM 可以与上下文窗口扩展方法相结合。这意味着 StreamingLLM 解决的是如何在有限缓存下维持模型稳定性和外推能力的问题,而其他方法则可能在此基础上进一步增加模型可以“看到”的最新上下文的容量。
  • 发现新现象: 本文的核心创新之一是发现了“注意力汇聚”这一内在现象,并基于此现象设计解决方案,而不仅仅是基于经验或系统优化。

4. 方法论

本章节将详细拆解 StreamingLLM 的技术方案,从发现“注意力汇聚 (attention sink)”现象开始,到提出 StreamingLLM 框架,再到优化预训练过程以支持更高效的流式部署。

4.1. 窗口注意力 (Window Attention) 的失败与注意力汇聚 (Attention Sinks)

4.1.1. 识别困惑度飙升点

描述: 窗口注意力 (window attention) 技术在推理 (inference) 阶段虽然提供了效率,但在长文本处理中会导致极高的语言建模困惑度 (language modeling perplexity),使其不适合流式部署。论文通过实验观察到,困惑度会在特定时间点急剧升高。

实验观察: 图 3 展示了在 20K 词元文本上语言建模的困惑度。当文本长度超过 KV 缓存 (KV cache) 大小,特别是当初始词元 (initial tokens) 被逐出时,困惑度会急剧飙升。这表明,无论初始词元与预测词元距离多远,它们对维持 LLMs 的稳定性都至关重要。

Figure 3: Language modeling perplexity on texts with 20K tokens across various LLM. Observations reveal consistent trends:(1) Dense attention fails once the input length surpasses the pre-training attention window s Windowattenti colapses nce the iput lnth exceds the cache ize, ie the iiial tokens e evicted. (3) StreamingLLM demonstrates stable performance, with its perplexity nearly matching that of the sliding window with re-computation baseline. 该图像是一个图表,展示了不同语言模型在处理20K tokens文本时的语言建模困惑度(log perplexity)。观察到,当输入长度超过预训练的注意力窗口时,密集注意力表现出明显的崩溃,而StreamingLLM则保持稳定,几乎与滑动窗口的重新计算基线相当。

图 3: 不同 LLM 在 20K 词元文本上的语言建模困惑度。观察显示出一致的趋势:(1) 当输入长度超过预训练注意力窗口大小时,密集注意力 (Dense attention) 失败。(2) 当输入长度超过缓存大小,即初始词元被逐出时,窗口注意力 (Window attention) 崩溃。(3) StreamingLLM 表现稳定,其困惑度几乎与滑动窗口重计算 (sliding window with re-computation) 基线匹配。

4.1.2. LLMs 移除初始词元 KV 状态时为何崩溃?

描述: 为了理解为什么移除初始词元的 KV 状态会导致模型崩溃,论文可视化了 Llama-2-7B 模型所有层和注意力头 (attention heads) 的注意力图 (attention maps)。

可视化分析: 图 2 展示了 Llama-2-7B 模型在 256 个句子上(每个句子长度为 16 个词元)的平均注意力对数 (average attention logits)。

Figure 2: Visualization of the average attention logits in Llama-2-7B over 256 sentences, each with a length o.Observations include 1) Theattention maps in the first two layers (layers 0 and ) exhibit he "l pattn, wi reokes ei otten.Beyond he bottwyers, thede heavie to the initial token across all layers and heads. 该图像是一个示意图,展示了Llama-2-7B模型在不同层和头部的平均注意力权重。左侧的图表为层0和头0及层1和头0的注意力分布,而右侧的图表显示了层2和头0、层9和头0、层16和头0、层23和头0、层31和头0的注意力情况。图中可以观察到初始令牌在多个层和头部的注意力集中现象。

图 2: Llama-2-7B 在 256 个句子上(每个句子长度为 16 个词元)的平均注意力对数可视化。观察结果包括:1) 前两层(第 0 层和第 1 层)的注意力图呈现“局部”模式,注意力集中在邻近词元。2) 在底部层之上,注意力明显偏向所有层和头部中的初始词元。

观察结果: 发现除了最底部的两层之外,模型在所有层和头部都持续强烈地关注初始词元。

解释: 移除这些初始词元的 KV 状态将从注意力计算的 SoftMax 函数(见下文公式)的分母中移除相当大的一部分。这种改变会导致注意力分数分布发生显著偏移,使其偏离正常推理设置下的预期分布,从而引发模型性能下降。

SoftMax 函数的数学公式 (Equation 1): SoftMax(x)i=exiex1+j=2Nexj,x1xj,j2,,N \mathrm { S o f t M a x } ( x ) _ { i } = \frac { e ^ { x _ { i } } } { e ^ { x _ { 1 } } + \sum _ { j = 2 } ^ { N } e ^ { x _ { j } } } , \quad x _ { 1 } \gg x _ { j } , j \in 2 , \ldots , N

  • 符号解释:
    • SoftMax(x)i\mathrm{SoftMax}(x)_i: SoftMax 函数对输入 xx 的第 ii 个元素的输出,表示第 ii 个词元的注意力权重。
    • ee: 自然对数的底数。
    • xix_i: 第 ii 个词元的注意力对数 (attention logit)。
    • ex1e^{x_1}: 初始词元的注意力对数指数化结果。
    • j=2Nexj\sum_{j=2}^{N} e^{x_j}: 从第二个词元到第 NN 个词元注意力对数指数化结果的总和。
    • NN: 序列中词元的总数。
    • x1xj,j2,,Nx_1 \gg x_j, j \in 2, \ldots, N: 表示初始词元 x1x_1 的注意力对数远大于其他词元 xjx_j 的注意力对数,这是论文观察到的“注意力汇聚”现象。移除 ex1e^{x_1} 会显著改变分母,进而扭曲整个注意力分布。

4.1.3. 语义重要性或绝对位置偏置?

描述: 为了区分初始词元的重要性是因为其语义关键,还是模型学习了一种对其绝对位置的偏置,论文进行了一项对照实验。

实验设计 (Table 1): 将前四个初始词元替换为换行符词元 "\n""\n",然后观察模型行为。 以下是原文 Table 1 的结果:

Llama-2-13BPPL (↓)
0 + 1024 (Window)5158.07
4 + 10205.40
4"\n"+10205.60
  • 缓存配置 (X+YX+Y): 表示添加 XX 个初始词元和 YY 个最新词元。
  • 评估: 困惑度在 PG-19 测试集的第一本书 (65K 词元) 上测量。

观察结果:

  • “0 + 1024 (Window)”配置(即纯窗口注意力,无初始词元)的困惑度高达 5158.07,表明其性能极差。
  • “4 + 1020”配置(保留 4 个原始初始词元和 1020 个最新词元)的困惑度显著降低至 5.40,性能得到极大恢复。
  • “4 "\n""\n" + 1020”配置(用 4 个换行符词元替换原始初始词元,并保留 1020 个最新词元)的困惑度为 5.60,与“4 + 1020”的困惑度相当。

结论: 即使初始词元没有语义上的重要性(被替换为无意义的换行符),模型仍然显著强调它们。重新引入它们可以使语言建模困惑度恢复到与具有原始初始词元相似的水平。这强烈表明,起始词元的绝对位置,而非其语义价值,在模型稳定性中扮演更重要的角色。

4.1.4. LLMs 将初始词元作为注意力汇聚 (Attention Sinks)

描述: 为了解释模型为何不成比例地关注初始词元(无论其语义相关性),论文引入了“注意力汇聚 (attention sink)”的概念。

解释原理:

  • SoftMax 函数的性质(公式 1)决定了所有被关注词元的注意力分数不能都为零。这意味着即使当前查询 (query) 已经包含了足够的信息进行预测,并且在许多先前词元中没有强匹配,模型仍然需要将这些不需要的注意力值分配到某个地方,以使总和为一。

  • 为什么是初始词元? 在自回归语言建模的顺序特性下,初始词元对几乎所有后续词元都是可见的。相比之下,后面的词元只能被有限的后续词元看到。因此,初始词元更容易被训练成接收这些“不必要”注意力值的“汇聚点”。

  • 多个注意力汇聚: 论文观察到 LLMs 通常会使用多个初始词元作为注意力汇聚,而不仅仅是一个 (Table 2)。例如,只引入一个或两个初始词元并不能完全恢复模型困惑度,而引入四个初始词元通常就足够了,再增加收益递减。这可能是因为在预训练期间,模型没有一个始终如一的特定起始词元(例如,Llama-2 虽然会为每个段落添加 <s><s> 词元,但它在文本分块前应用,导致第 0 位词元通常是随机的),这使得模型需要将几个初始词元作为注意力汇聚。

    以下是原文 Table 2 的结果:

    Cache Config 0+20481+20472+2046 4+2044 8+2040
    Falcon-7B17.9012.1212.1212.12 12.12
    MPT-7B460.2914.9915.0014.99 14.98
    Pythia-12B21.6211.9512.0912.09 12.02
    Cache Config0+40961+40952+40944+4092 8+4088
    Llama-2-7B3359.9511.8810.519.59 9.54
  • 缓存配置 (X+YX+Y): 表示添加 XX 个初始词元和 YY 个最新词元。

  • 评估: 困惑度在连接的 PG19 测试集中的 400K 词元上评估。

4.2. 带有注意力汇聚的滚动 KV 缓存 (Rolling KV Cache with Attention Sinks) - StreamingLLM 框架

描述: 为了使已训练的 LLMs 能够进行流式处理,论文提出了 StreamingLLM,这是一种直接的方法,可以在不进行任何模型微调 (fine-tuning) 的情况下,恢复窗口注意力 (window attention) 的困惑度表现。该方法在当前的滑动窗口词元之外,重新引入了少量起始词元 (starting tokens) 的 KV 状态进行注意力计算。

KV 缓存结构: StreamingLLM 中的 KV 缓存可以概念上分为两部分 (Figure 4):

  1. 注意力汇聚 (Attention Sinks): 由少数(论文中默认使用 4 个)初始词元的 KV 状态组成,用于稳定注意力计算。

  2. 滚动 KV 缓存 (Rolling KV Cache): 保留最近的词元 KV 状态,这对于语言建模至关重要。

    Figure 4: The KV cache of StreamingLLM 该图像是一个示意图,展示了StreamingLLM中的KV缓存机制。图中显示了在生成第7、8和9个token时的不同状态,包括注意力汇聚(Attention Sinks)、被驱逐的token和滚动KV缓存。该图旨在说明如何利用初始token来改进窗口注意力的性能。

图 4: StreamingLLM 的 KV 缓存。该图示意了 StreamingLLM 如何通过结合注意力汇聚和滚动 KV 缓存来处理连续的词元流。当生成新的词元时,滚动 KV 缓存会更新,而注意力汇聚(即最早的几个词元)的 KV 状态则始终保留。

相对位置编码 (Relative Positional Encoding) 的处理:

  • StreamingLLM 在确定相对距离和添加位置信息时,关注缓存内的位置,而非原始文本中的位置。这对于 StreamingLLM 的性能至关重要。

  • 示例: 如果当前缓存包含词元 [0, 1, 2, 3, 6, 7, 8],并且正在解码第 9 个词元,那么分配给它们的内部位置将是 [0, 1, 2, 3, 4, 5, 6, 7],而不是它们在原始文本中的位置 [0, 1, 2, 3, 6, 7, 8, 9]

  • 对于 RoPE (Rotary Position Embeddings): 模型缓存词元的 Key (K) 向量在引入旋转变换之前。然后在每个解码阶段,将位置变换应用于滚动缓存中的 Key 向量。

  • 对于 ALiBi (Attention with Linear Biases): 集成更为直接。应用连续的线性偏差 (linear bias) 而非“跳跃式 (jumping)”偏差到注意力分数上。

    这种在缓存内分配位置嵌入的方法是 StreamingLLM 功能的关键,确保模型即使在超出其预训练注意力窗口大小的情况下也能高效运行。

4.3. 使用注意力汇聚预训练 LLMs (Pre-Training LLMs with Attention Sinks)

描述: 正如 4.1.4 节所述,模型对多个初始词元过度关注的一个重要原因是缺乏一个指定的汇聚词元 (sink token) 来卸载多余的注意力分数。模型因此无意中使用了全局可见的词元(主要是初始词元)作为注意力汇聚。论文提出了两种可能的补救措施:

  • 引入可学习汇聚词元 (Learnable Sink Token): 有意地在所有训练样本的开头包含一个全局可训练的注意力汇聚词元,作为不必要注意力分数的一个存储库。
  • 替代 SoftMax 函数: 将传统的 SoftMax 函数替换为 SoftMax-off-by-One (Miller, 2023) 的变体。

SoftMax-off-by-One (SoftMax_1) 数学公式 (Equation 2): SoftMax1(x)i=exi1+j=1Nexj, \mathrm { S o f t M a x } _ { 1 } ( x ) _ { i } = \frac { e ^ { x _ { i } } } { 1 + \sum _ { j = 1 } ^ { N } e ^ { x _ { j } } } ,

  • 符号解释:
    • SoftMax1(x)i\mathrm { S o f t M a x } _ { 1 } ( x ) _ { i }: SoftMax-off-by-One 函数对输入 xx 的第 ii 个元素的输出。
    • ee: 自然对数的底数。
    • xix_i: 第 ii 个词元的注意力对数。
    • 1: 分母中额外增加的常数 1,使得注意力分数不再强制总和为 1,为注意力分数提供了一个“卸载点”。
    • j=1Nexj\sum _ { j = 1 } ^ { N } e ^ { x _ { j } }: 所有词元注意力对数指数化结果的总和。
    • NN: 序列中词元的总数。
  • 等价性: 注意到 SoftMax1\mathrm { S o f t M a x } _ { 1 } 等价于在注意力计算中预置一个 Key 和 Value 特征均为零的词元。论文将此方法称为 “零汇聚 (Zero Sink)”,以适应其框架。

验证实验: 为了验证这些假设,论文从头开始预训练了三个 1.6 亿参数的语言模型,在相同设置下进行:

  1. Vanilla: 使用标准的 SoftMax 注意力。

  2. Zero Sink: 将常规注意力机制替换为 SoftMax1\mathrm { S o f t M a x } _ { 1 }

  3. Learnable Sink: 在所有训练样本的开头预置一个可学习的占位符词元。

    以下是原文 Table 3 的结果:

    Cache Config 0+1024 1+1023 2+1022 4+1020
    Vanilla27.8718.4918.0518.05
    Zero Sink2921419.9018.2718.01
    Learnable Sink123518.0118.0118.02
  • 缓存配置 (X+YX+Y): 表示添加 XX 个初始词元和 YY 个最新词元。
  • 评估: 困惑度在 PG-19 测试集的第一个样本上评估。

实验结果:

  • 零汇聚 (Zero Sink): 尽管在一定程度上缓解了注意力汇聚问题,但模型仍然依赖其他初始词元作为注意力汇聚。当没有初始词元时(0+1024),困惑度仍然很高 (29214)。

  • 可学习汇聚词元 (Learnable Sink Token): 在稳定注意力机制方面非常有效。只需将这个汇聚词元与最新词元配对,就能充分锚定模型的性能,甚至评估困惑度略有改善。当只有一个汇聚词元时(1+1023),困惑度为 18.01,与使用 4 个初始词元的情况几乎一致。

    结论: 这些发现表明,在所有训练样本的开头包含一个稳定的可学习词元,可以使其作为专用的注意力汇聚,从而优化流式部署。论文建议未来的 LLMs 训练时应包含一个汇聚词元,以优化流式部署。

5. 实验设置

本节详细描述了 StreamingLLM 的实验设置,包括使用的数据集、评估指标、对比基线模型以及采用的 LLM 家族和规模。

5.1. 数据集

实验使用了多个数据集来全面评估 StreamingLLM 的性能:

  • PG19 (Rae et al., 2020) 测试集:
    • 特点: 包含 100 本长篇书籍的测试集。
    • 用途: 主要用于评估语言模型的困惑度 (perplexity),尤其是在超长文本(如 400 万词元)上的性能稳定性。论文通过连接 (concatenated) 整个 PG19 测试集来创建超长序列。
  • ARC-[Challenge, Easy] (Clark et al., 2018):
    • 特点: 问答数据集,分为“简单 (Easy)”和“挑战 (Challenge)”两部分,考察常识推理能力。
    • 用途: 用于评估 StreamingLLM 在多轮问答场景下,结合指令微调 (instruction-tuned) LLMs 的准确率。
  • StreamEval (本文引入):
    • 特点: 受 LongEval (Li et al., 2023) 基准测试启发而设计。与 LongEval 的单次长跨度查询不同,StreamEval 模拟了真实世界中常见的流式问答场景:每接收 10 行新信息就进行一次查询,且查询的答案始终位于 20 行之前。每行包含 23 个词元。

    • 用途: 专门用于评估 StreamingLLM 在连续流式输入中的问答准确率,尤其是在输入长度逐渐增加(例如达到 120K 词元)时的性能。

    • 样本示例 (原文 Figure 8):

      Figure 8: The first sample in StreamEval. 该图像是一个示意图,展示了输入内容与期望输出之间的关系。输入内容呈现为多行线索记录,其中包括特定行号的查询。每个查询都涉及相应行的注册内容,并提供了一些省略的行信息。最终的期望输出展示了一个包括选定内容的简洁表示。

      图 8: StreamEval 中的第一个样本。该图展示了 StreamEval 数据集的一个示例。输入由一系列带行号的文本组成,每隔一定行数(例如 10 行)会提出一个问题,该问题通常指向之前某个特定行号的内容(例如 20 行之前)。模型需要根据不断更新的上下文来回答这些问题。

  • LongBench (Bai et al., 2023):
    • 特点: 一个双语、多任务基准测试,用于评估长上下文理解能力。包含多种 NLP 任务,如单文档问答 (NarrativeQA (Koisky et al., 2017), Qasper (Dasigi et al., 2021))、多文档问答 (HotpotQA (Yang et al., 2018), 2WikiMQA Ho et al. (2020)) 和摘要 (GovReport (Huang et al., 2021), MultiNews (Fabbri et al., 2019))。
    • 用途: 用于与 LongBench 提供的默认截断基线进行比较,评估 StreamingLLM 在各种长文本理解任务中的表现,并探讨注意力汇聚数量对性能的影响。
  • Pile (Gao et al., 2020) 数据集:
    • 特点: 一个 800GB 的多样化文本数据集,用于语言建模的预训练。论文使用了其去重版本。
    • 用途: 用于预训练 1.6 亿参数的语言模型,以验证引入专用汇聚词元 (sink token) 对流式部署性能的影响。

5.2. 评估指标

论文中使用的评估指标及其详细说明如下:

  • 困惑度 (Perplexity, PPL):

    1. 概念定义: 困惑度是衡量语言模型预测下一个词元 (token) 准确性的常用指标。它量化了模型对给定文本序列的“不确定性”程度。具体来说,困惑度是模型在测试集上平均每个词元预测概率的几何平均的倒数。困惑度值越低,表示模型对文本的预测能力越好,对文本序列的建模效果越准确。
    2. 数学公式: 对于一个包含 NN 个词元的序列 W=(w1,w2,,wN)W = (w_1, w_2, \ldots, w_N),其困惑度 PPL(W) 计算公式为: PPL(W)=exp(1Ni=1NlogP(wiw1,,wi1)) PPL(W) = \exp\left(-\frac{1}{N} \sum_{i=1}^{N} \log P(w_i | w_1, \ldots, w_{i-1})\right)
    3. 符号解释:
      • PPL(W): 序列 WW 的困惑度。
      • exp()\exp(\cdot): 自然指数函数。
      • NN: 序列 WW 中词元的总数。
      • P(wiw1,,wi1)P(w_i | w_1, \ldots, w_{i-1}): 在给定前 i-1 个词元的情况下,语言模型预测第 ii 个词元 wiw_i 的概率。
      • log\log: 自然对数。
  • 准确率 (Accuracy):

    1. 概念定义: 准确率是最直接的性能衡量指标之一,特别适用于分类任务或精确匹配的问答任务。它表示模型正确预测的样本数量占总样本数量的比例。在问答场景中,如果模型生成的答案与参考答案完全一致,则视为正确。
    2. 数学公式: Accuracy=Number of Correct PredictionsTotal Number of Predictions Accuracy = \frac{\text{Number of Correct Predictions}}{\text{Total Number of Predictions}}
    3. 符号解释:
      • Number of Correct Predictions\text{Number of Correct Predictions}: 模型做出正确预测的样本数量。
      • Total Number of Predictions\text{Total Number of Predictions}: 模型进行预测的总样本数量。
  • 解码延迟 (Decoding Latency):

    1. 概念定义: 解码延迟衡量的是语言模型在生成一个新词元或整个序列时所花费的时间。在流式应用中,通常更关注生成每个词元的平均延迟 (per-token decoding latency),因为这直接影响用户体验的响应速度。
    2. 数学公式: 论文中未给出具体公式,但通常以单位时间 (例如毫秒) 每个词元来表示。
    3. 符号解释: 无需特定符号,通常直接报告时间值。
  • 内存使用 (Memory Usage):

    1. 概念定义: 内存使用量衡量的是语言模型在运行过程中所需的计算内存资源,特别是 KV 缓存所占用的显存。在部署 LLMs 时,内存效率是关键考量因素,因为它直接影响可部署的模型规模和批处理大小。
    2. 数学公式: 论文中未给出具体公式,通常以 GB 或 MB 为单位报告。
    3. 符号解释: 无需特定符号,通常直接报告内存值。

5.3. 对比基线

论文将 StreamingLLM 与以下几种主流的注意力机制和流式处理方法进行了比较:

  • 密集注意力 (Dense Attention):
    • 描述: 这是 Transformer 模型中标准的自注意力机制,它缓存所有先前词元 (tokens) 的键和值 (KV) 状态。
    • 代表性: 它是最基础、最直接的实现方式,但其 KV 缓存的内存消耗和计算复杂度随序列长度呈二次方增长,导致在长序列下效率低下,最终可能导致内存溢出 (Out-of-Memory, OOM)。
  • 窗口注意力 (Window Attention):
    • 描述: 一种优化策略,它只维护一个固定大小的滑动窗口来缓存最近 LL 个词元的 KV 状态。当新的词元加入时,最旧的词元 KV 状态会被逐出。
    • 代表性: 旨在解决密集注意力的内存效率问题,提供恒定的内存使用和解码速度。
  • 滑动窗口重计算 (Sliding Window with Re-computation):
    • 描述: 这种方法在生成每个新词元时,会重新计算最近 LL 个词元的 KV 状态。
    • 代表性: 它能提供与密集注意力相近的性能质量,因为每次都“重新看到”了完整的近期上下文。然而,由于每次都需要在滑动窗口内进行二次方复杂度的注意力计算,其速度非常慢,不适用于实时流式应用。在论文中,它被视为“oracle baseline (神谕基线)”,代表了在给定窗口大小下的最佳性能。

5.4. 模型和硬件

  • 模型家族:
    • Llama-2 (Touvron et al., 2023b): 包括 Llama-2-[7, 13, 70]B 模型。Llama-2 使用 RoPE (Rotary Position Embeddings)
    • MPT (Team, 2023): 包括 MPT-[7, 30]B 模型。MPT 使用 ALiBi (Attention with Linear Biases)
    • Pythia (Biderman et al., 2023): 包括 Pythia-[2.9, 6.9, 12]B 模型,以及用于预训练实验的 Pythia-160M。Pythia 使用 RoPE
    • Falcon (Almazrouei et al., 2023): 包括 Falcon-[7, 40]B 模型。Falcon 使用 RoPE
  • 硬件:
    • 性能和效率测试: 在单个 NVIDIA A6000 GPU 上进行,用于衡量解码延迟和内存使用。
    • 预训练带汇聚词元模型: 在一个 8xA6000 NVIDIA GPU 服务器上进行。

6. 实验结果与分析

本节将深入解读论文的实验结果,分析 StreamingLLM 在不同场景下的性能表现,并与基线方法进行对比。

6.1. 语言建模在超长文本上的表现 (Language Modeling on Long Texts)

6.1.1. 困惑度在 20K 词元文本上的表现

描述: 论文首先在 20K 词元文本上评估了 StreamingLLM 的语言建模困惑度,与密集注意力、窗口注意力以及滑动窗口重计算基线进行对比。

结果分析 (原文 Figure 3):

Figure 3: Language modeling perplexity on texts with 20K tokens across various LLM. Observations reveal consistent trends:(1) Dense attention fails once the input length surpasses the pre-training attention window s Windowattenti colapses nce the iput lnth exceds the cache ize, ie the iiial tokens e evicted. (3) StreamingLLM demonstrates stable performance, with its perplexity nearly matching that of the sliding window with re-computation baseline. 该图像是一个图表,展示了不同语言模型在处理20K tokens文本时的语言建模困惑度(log perplexity)。观察到,当输入长度超过预训练的注意力窗口时,密集注意力表现出明显的崩溃,而StreamingLLM则保持稳定,几乎与滑动窗口的重新计算基线相当。

图 3: 不同 LLM 在 20K 词元文本上的语言建模困惑度。观察显示出一致的趋势:(1) 当输入长度超过预训练注意力窗口大小时,密集注意力 (Dense attention) 失败。(2) 当输入长度超过缓存大小,即初始词元被逐出时,窗口注意力 (Window attention) 崩溃。(3) StreamingLLM 表现稳定,其困惑度几乎与滑动窗口重计算 (sliding window with re-computation) 基线匹配。

核心发现:

  • 密集注意力 (Dense Attention) 的失败: 当输入长度超过模型的预训练注意力窗口大小时,密集注意力模型迅速崩溃,困惑度飙升。这印证了其固有的长度限制。
  • 窗口注意力 (Window Attention) 的崩溃: 窗口注意力在输入长度超过其固定缓存大小(即初始词元被逐出)时,性能急剧下降并崩溃。这再次验证了论文关于“注意力汇聚”现象以及初始词元重要性的论断。
  • StreamingLLM 的稳定性: StreamingLLM 展示了稳定的性能,其困惑度几乎与滑动窗口重计算基线 (sliding window with re-computation baseline) 相匹配。滑动窗口重计算由于每次都重新计算上下文,提供了高质量的性能,因此 StreamingLLM 能够达到这一水平,表明其在保持上下文质量方面的有效性。

6.1.2. 困惑度在 400 万词元文本上的表现

描述: 论文进一步将评估扩展到更长的文本序列,即 400 万词元,以验证 StreamingLLM 在极其长的流式输入下的稳定性。

结果分析 (原文 Figure 5):

Figure 5: Language modeling perplexity of StreamingLLM on super long texts with 4 million tokens across varius LLM families and scales. The perplexity remains stable throughout. We use the concatenated test set of PG19 (100 books) to perform language modeling, with perplexity fluctuations due to book transitions. 该图像是一个图表,展示了StreamingLLM在超长文本上进行语言建模时的困惑度(perplexity)变化。图中分别列出了Llama-2、Pythia、Falcon和MPT这几种LLM的表现,困惑度在不同输入长度下保持稳定。

图 5: StreamingLLM 在 400 万词元超长文本上的语言建模困惑度,涵盖了多种 LLM 家族和规模。困惑度在整个过程中保持稳定。实验使用连接的 PG19 (100 本书) 测试集进行语言建模,困惑度波动是由于书籍切换。

核心发现:

  • 超长文本的泛化能力: StreamingLLM 能够可靠地处理超过 400 万词元的超长文本,涵盖了 Llama-2 (7B, 13B, 70B)、Falcon (7B, 40B)、Pythia (2.8B, 6.9B, 12B) 和 MPT (7B, 30B) 等多个 LLM 家族和不同规模的模型。
  • 持续的性能稳定性: 在如此长的序列中,模型的困惑度保持稳定,仅在书籍切换时出现轻微波动,这表明 StreamingLLM 在实际的流式应用中具有极强的鲁棒性。

6.2. 预训练带注意力汇聚词元的结果 (Results of Pre-Training with A Sink Token)

6.2.1. 收敛和正常模型性能 (Convergence and Normal Model Performance)

描述: 论文研究了在预训练阶段引入专用注意力汇聚词元 (sink token) 是否会影响模型的收敛行为和在标准 NLP 基准测试上的性能。

结果分析 (原文 Figure 7 和 Figure 6):

Figure 6: Pre-training loss curves of models w/ and w/o sink tokens. Two models have a similar convergence trend. 该图像是训练损失曲线图,展示了模型在使用和不使用 sink token 时的收敛趋势。蓝色线条表示未使用 sink token 的模型,橙色线条表示添加了 sink token 的模型,二者在 k 步骤的训练损失趋于相似。

图 7: 平均注意力对数可视化,比较了未经汇聚词元 (左侧) 和使用汇聚词元 (右侧) 预训练的模型在 256 个句子上 (每个句子长度 16 个词元) 的注意力分布。关键观察:(1) 未经汇聚词元训练的模型在低层显示局部注意力,在深层对初始词元注意力增加。 (2) 使用汇聚词元训练的模型在各层和头部持续集中于汇聚词元,并且对其他初始词元的注意力减少,支持了指定汇聚词元以增强流式性能的益处。

Figure 7: Visualization of average attention logits over 256 sentences, each 16 tokens long, comparing models pre-trained without (left) and with (right) a sink token. Both maps show the same layers and heads. Key observations: (1) Without a sink token, models show local attention in lower layers and increased attention tkes deyersWi tke the ee e , eivel ceanatenion.Wit he ec the tokn attenion ive initial tokens, supporting the benefit of designating the sink token to enhance the streaming performance. 该图像是一个比较图,显示了在使用和不使用池化标记(sink token)预训练的模型在不同层和头上的平均注意力日志的可视化。左侧为未使用池化标记的模型,其低层显示局部注意力,而右侧则是使用池化标记的模型,初始标记的注意力增强了流式性能。

图 6: 带有/不带有汇聚词元模型的预训练损失曲线。两个模型具有相似的收敛趋势。

核心发现:

  • 收敛性: 如图 6 所示,预训练损失曲线显示,无论是否包含汇聚词元,模型的收敛动态都非常相似,表明引入汇聚词元对模型训练的稳定性没有负面影响。

  • 通用任务性能: 论文在 ARC-[Challenge, Easy]、HellaSwag、LAMBADA、OpenbookQA、PIQA 和 Winogrande 这七个 NLP 基准测试上评估了两个 1.6 亿参数模型(一个传统预训练,一个带汇聚词元预训练)的零样本 (zero-shot) 准确率。 以下是原文 Table 4 的结果:

    MethodsARC-c ARC-eHS LBDOBQAPIQA WG
    Vanilla18.645.229.439.616.062.250.1
    +Sink Token19.645.629.839.916.662.650.8
  • 结果: 如表 4 所示,预训练时包含汇聚词元并没有损害模型的性能,反而略有提升,这表明引入汇聚词元并不会牺牲模型的通用语言理解能力。

6.2.2. 流式性能 (Streaming Performance)

描述: 论文比较了传统预训练模型和带汇聚词元预训练模型在流式场景下的困惑度表现。

结果分析 (原文 Table 3): 以下是原文 Table 3 的结果:

Cache Config 0+1024 1+1023 2+1022 4+1020
Vanilla27.8718.4918.0518.05
Zero Sink2921419.9018.2718.01
Learnable Sink123518.0118.0118.02
  • 缓存配置 (X+YX+Y): 表示添加 XX 个初始词元和 YY 个最新词元。
  • 评估: 困惑度在 PG-19 测试集的第一个样本上评估。

核心发现:

  • 普通模型的依赖性: 传统模型 (Vanilla) 需要添加多个初始词元(如 4 个)作为注意力汇聚才能保持稳定的流式困惑度 (从 27.87 降到 18.05)。
  • 带汇聚词元模型的效率: 预训练时带有汇聚词元 (Learnable Sink) 的模型,仅使用该汇聚词元本身(1+1023 配置)即可实现令人满意的流式性能(困惑度为 18.01),这与普通模型需要 4 个初始词元才能达到的性能相当。
  • 注意力可视化 (原文 Figure 7): 右侧图显示,带汇聚词元预训练的模型在所有层和注意力头中都一致地将注意力集中到汇聚词元上,而对其他初始词元的注意力则显著减少,这解释了汇聚词元在增强模型流式性能方面的功效。

6.3. 流式问答结果 (Results on Streaming Question Answering)

6.3.1. 多轮问答 (Multi-round Question Answering)

描述: 为了展示 StreamingLLM 的实际应用性,论文使用指令微调 (instruction-tuned) 的 LLMs 模拟多轮问答场景,这在现实世界中非常常见。

实验设计: 将 ARC-[Challenge, Easy] 数据集中的所有问答对连接起来,形成一个连续的输入流,输入到 Llama-2-[7, 13, 70]B-Chat 模型中。在每个答案位置评估模型的补全结果,使用精确匹配 (exact match) 作为评估标准。

以下是原文 Table 5 的结果:

Llama-2-7B-ChatLlama-2-13B-ChatLlama-2-70B-Chat
MethodsARC-c ARC-eARC-c ARC-eARC-c ARC-e
One-Shot22.5 45.128.3 54.434.1 61.2
Dense AttentionOOM OOMOOM OOMOOM OOM
Window Attention (2048)0.0 0.00.0 0.00.0 0.0
StreamingLLM (4+2044)22.1 45.028.0 54.333.9 61.1
  • One-Shot: 表示单次问答基线,即每个问答对独立处理,不考虑流式上下文。
  • OOM: 内存溢出 (Out-of-Memory)。

核心发现:

  • 密集注意力 (Dense Attention) 的局限性: 密集注意力在所有模型上都导致了内存溢出 (OOM),表明其不适用于流式场景。
  • 窗口注意力 (Window Attention) 的失败: 窗口注意力方法虽然高效,但由于在输入长度超过缓存大小时产生随机输出,导致准确率非常低(0.0),无法实际应用。
  • StreamingLLM 的高效和准确性: StreamingLLM 能够高效处理流式格式,其准确率与单次问答基线 (one-shot baseline) 基本持平。这表明 StreamingLLM 在保持性能质量的同时,成功解决了流式应用的挑战。

6.3.2. StreamEval 基准测试结果

描述: 论文引入了 StreamEval 数据集,以更好地模拟 StreamingLLM 的应用场景。与 LongEval 的单次长查询不同,StreamEval 每 10 行新信息查询一次,且答案始终在 20 行之前,反映了问题通常与最近信息相关的现实情况。

结果分析 (原文 Figure 9):

Figure 9: Performance on the StreamEval benchmark. Accuracies are averaged over 100 samples. 该图像是一个性能比较图,展示了不同模型(如Llama-2系列)在StreamEval基准测试上的精度表现。通过密集注意力、窗口注意力和StreamingLLM三种方法,随着输入长度增加,精度变化趋势有所不同,StreamingLLM在长输入下显示出更稳健的性能。

图 9: StreamEval 基准测试上的性能。准确率平均自 100 个样本。

核心发现:

  • StreamingLLM 的稳健性: LLMs 采用 StreamingLLM 后,即使输入长度接近 120K 词元,仍能保持合理的准确率。
  • 基线的崩溃: 密集注意力和窗口注意力分别在达到预训练文本长度和 KV 缓存大小时失败。
  • 与上下文扩展方法的互补性: 论文还展示了 StreamingLLM 可以与上下文扩展模型(如 LongChat-7b-v1.5-32k 和 Llama-2-7B-32K-Instruct)相结合。在这种情况下,上下文扩展意味着拓宽 StreamingLLMs 的最大缓存大小,从而能够捕获更广泛的局部信息。

6.3.3. StreamEval 随查询-答案距离增加的准确率

描述: 为了评估 StreamingLLM 处理更长距离上下文的能力,论文在 StreamEval 上测试了 Llama-2-7B-32K-Instruct 模型,评估不同查询-答案行距离下的准确率,以及不同缓存配置的影响。

结果分析 (原文 Table 7): 以下是原文 Table 7 的结果:

Llama-2-7B-32K-InstructCache Config
Line DistancesToken Distances4+20444+40924+81884+16380
2046085.8084.6081.1577.65
4092080.3583.8081.2577.50
60138079.1582.8081.5078.50
80184075.3077.1576.4073.80
10023000.0061.6050.1040.50
15034500.0068.2058.3038.45
20046000.000.0062.7546.90
40092000.000.000.0045.70
600138000.000.000.0028.50
800184000.000.000.000.00
1000230000.000.000.000.00
  • 行距离 (Line Distances): 查询与答案之间的行数距离。
  • 词元距离 (Token Distances): 行距离乘以每行 23 个词元。
  • 缓存配置 (X+YX+Y): 表示添加 XX 个初始词元和 YY 个最新词元,总缓存大小为 X+YX+Y

核心发现:

  • 性能受限于缓存大小: StreamingLLM 在查询与答案之间的词元距离在缓存大小范围内时能保持准确率。然而,一旦这个距离增加并超出缓存容量,准确率就会迅速下降并最终降至零。
  • 非上下文扩展: 这些结果再次证实,StreamingLLM 的作用是基于近期上下文生成连贯文本,它不能扩展语言模型的上下文长度。它仅在自身缓存的范围内提供稳定性能。
  • 模型对上下文利用的局限性: 这些发现也强调了当前语言模型的一个更广泛的挑战:它们可能无法充分利用其缓存中的所有上下文信息。

6.4. LongBench 长距离基准测试评估 (Long-Range Benchmark Evaluation)

描述: 论文使用 Llama-2-7B-chat 模型(最大上下文长度 4K)在 LongBench 上评估了 StreamingLLM,LongBench 包含单文档 QA、多文档 QA 和摘要等任务。LongBench 默认将输入截断到 3,500 个词元,其中保留开头 1,750 个和结尾 1,750 个词元。

结果分析 (原文 Table 8): 以下是原文 Table 8 的结果:

Llama2-7B-chatSingle-Document QA NarrativeQA QasperMulti-Document QA HotpotQA 2WikiMQASummarization GovReport MultiNews
Truncation 1750+175018.719.225.432.827.325.8
StreamingLLM 4+349611.616.921.628.223.925.5
StreamingLLM 1750+175018.219.724.932.026.325.9
  • Truncation 1750+1750: 基线方法,将输入截断为开头 1750 词元和结尾 1750 词元。
  • StreamingLLM 4+3496: 使用 4 个注意力汇聚词元和 3496 个最新词元。
  • StreamingLLM 1750+1750: 使用 1750 个初始词元(作为注意力汇聚)和 1750 个最新词元。

核心发现:

  • 初始提示信息的重要性: StreamingLLM 在 4+34964+3496 配置下表现不如截断基线 (Truncation 1750+1750),这很可能是因为丢失了输入提示中关键的初始信息(StreamingLLM 4+34964+3496 只保留了 4 个初始词元,而基线保留了 1750 个)。
  • 与截断基线性能相当: 当将注意力汇聚的数量调整为 1750 个(即 StreamingLLM 1750+1750 配置)时,StreamingLLM 的性能恢复到与文本截断基线相当的水平。
  • 结论的强化: 这些结果进一步证实了第 6.3.3 节的发现:StreamingLLM 的有效性取决于其缓存内包含的信息。在缓存内部,其性能与文本截断基线表现相当。这再次强调了 StreamingLLM 并非魔术般地扩展了上下文理解,而是在有限的缓存容量内优化了模型的稳定性。

6.5. 效率结果 (Efficiency Results)

描述: 论文将 StreamingLLM 的解码延迟和内存使用与唯一的具有可接受质量的基线——滑动窗口重计算 (sliding window with re-computation)——进行了基准测试。

实验设置: 两种方法均使用 Huggingface Transformers 库 (Wolf et al., 2020) 实现,并在单个 NVIDIA A6000 GPU 上使用 Llama-2-7B 和 Llama-2-13B 模型进行测试。

结果分析 (原文 Figure 10):

Figure 10: Comparison of per-token decoding latency and memory usage between the sliding window approach with re-computation baseline and StreamingLLM, plotted against the cache size (attention window size) on the \(\\mathrm { X }\) -axis. StreamingLLM delivers a remarkable speedup of up to \(2 2 . 2 \\times\) per token and retains a memory footprint similar to the re-computation baseline. 该图像是图表,比较了 Llama-2-7B 和 Llama-2-13B 在不同缓存大小下,滑动窗口重新计算方法与 StreamingLLM 的每个 token 解码延迟及内存使用。结果显示,StreamingLLM 在延迟和内存使用方面均优于滑动窗口方案。

图 10: 滑动窗口重计算基线与 StreamingLLM 在不同缓存大小(注意力窗口大小)下的每个词元解码延迟和内存使用比较。StreamingLLM 实现了每个词元高达 22.2 倍的显著加速,并保持了与重计算基线相似的内存占用。

核心发现:

  • 解码速度 (Decoding Speed):
    • 随着缓存大小 (cache size) 增加,StreamingLLM 的解码速度呈线性增长
    • 滑动窗口重计算基线的解码延迟则呈二次方增长
    • 显著加速: 因此,StreamingLLM 实现了令人印象深刻的加速,每个词元高达 22.2 倍。这使得 StreamingLLM 在实时流式应用中具有极高的实用性。
  • 内存使用 (Memory Usage): 尽管延迟大幅降低,StreamingLLM 仍保持了与重计算基线一致的内存占用。这表明 StreamingLLM 在效率提升的同时,没有引入额外的内存负担。

6.6. 消融实验/参数分析 (Ablation Studies)

6.6.1. 初始词元数量的影响 (Numbers of Initial Tokens)

描述: 论文通过消融实验 (ablation study) 探讨了保留不同数量的初始词元作为注意力汇聚对流式困惑度的影响。

结果分析 (原文 Table 2): 以下是原文 Table 2 的结果:

Cache Config 0+20481+20472+2046 4+2044 8+2040
Falcon-7B17.9012.1212.1212.12 12.12
MPT-7B460.2914.9915.0014.99 14.98
Pythia-12B21.6211.9512.0912.09 12.02
Cache Config0+40961+40952+40944+4092 8+4088
Llama-2-7B3359.9511.8810.519.59 9.54
  • 缓存配置 (X+YX+Y): 表示添加 XX 个初始词元和 YY 个最新词元。
  • 评估: 困惑度在连接的 PG19 测试集中的 400K 词元上评估。

核心发现:

  • 纯窗口注意力失效: 0+Y0+Y 配置(纯窗口注意力,无初始词元)的困惑度急剧增加,再次验证了初始词元的重要性。
  • 少量初始词元的不足: 引入一个或两个初始词元并不能完全恢复模型的困惑度。例如,Llama-2-7B 从 0+40960+4096 的 3359.95 降到 1+40951+4095 的 11.88 和 2+40942+4094 的 10.51,虽有大幅改善但仍未达最优。
  • 四个初始词元的阈值: 引入四个初始词元通常足以恢复 LLM 的性能。例如,Llama-2-7B 在 4+40924+4092 配置下困惑度为 9.59,而 8+40888+4088 配置为 9.54,收益递减。
  • 结论: 这一结果证明了论文选择使用 4 个初始词元作为注意力汇聚在 StreamingLLM 中的合理性。

6.6.2. 缓存大小的影响 (Cache Sizes)

描述: 论文评估了 StreamingLLM 中缓存大小对困惑度的影响。

结果分析 (原文 Table 6): 以下是原文 Table 6 的结果:

Cache4+2524+5084+1020 4+2044
Falcon-7B13.6112.8412.34 12.84
MPT-7B14.1214.2514.33 14.99
Pythia-12B13.1712.5212.08 12.09
Cache4+5084+10204+2044 4+4092
Llama-2-7B9.739.329.08
  • 缓存配置 (X+YX+Y): 表示添加 XX 个初始词元和 YY 个最新词元,总缓存大小为 X+YX+Y
  • 评估: 困惑度在连接的 PG19 测试集中的 400K 词元上评估。

核心发现:

  • 非线性提升: 增加缓存大小(即保留更多近期词元)并不总是能一致地降低语言建模困惑度。例如,Falcon-7B 在 4+10204+1020 时困惑度为 12.34,但在 4+20444+2044 时却略有上升到 12.84。MPT-7B 也呈现出类似的不一致趋势。
  • 模型利用上下文的局限性: 这种不一致性表明,这些模型可能未能最大化利用它们所接收到的全部上下文信息。即使提供了更多的上下文,模型也可能无法有效地将其整合到预测中。
  • 结论: 这一发现指出了未来研究的方向:应着力于增强模型更好地利用其广泛上下文的能力。

7. 总结与思考

7.1. 结论总结

部署大型语言模型 (LLMs) 于流式应用(如多轮对话)的需求日益增长,但受到效率限制和长文本性能下降的挑战。传统的窗口注意力 (window attention) 方法虽然部分解决了效率问题,但在初始词元 (initial tokens) 被逐出时性能会急剧下降。

本文的核心贡献在于:

  1. 发现了“注意力汇聚 (attention sink)”现象: 即使初始词元在语义上不重要,它们也会作为“汇聚点”吸引大量注意力分数,这是 SoftMax 函数特性和自回归模型训练过程的产物。

  2. 提出了 StreamingLLM 框架: 这是一个简单高效的框架。它通过在滚动窗口缓存中额外保留少量初始词元的 KV 状态,有效稳定了注意力计算,从而使 LLMs 能够在不进行任何微调 (fine-tuning) 的情况下,从有限的预训练注意力窗口泛化到无限长的序列。

  3. 性能和效率的显著提升: StreamingLLM 在 Llama-2、MPT、Falcon 和 Pythia 等主流 LLM 上实现了对高达 400 万词元甚至更长文本的稳定高效语言建模。相比于滑动窗口重计算基线,它带来了高达 22.2 倍的解码速度提升,同时保持了相似的内存占用。

  4. 预训练优化建议: 论文进一步发现,在预训练阶段引入一个专用的可学习汇聚词元 (learnable sink token) 可以进一步优化流式部署,使得模型仅需一个汇聚词元即可稳定性能。

    StreamingLLM 首次成功地将 LLM 的预训练窗口大小与其实际文本生成长度解耦,为 LLM 的流式部署铺平了道路。

7.2. 局限性与未来工作

论文作者指出了 StreamingLLM 的局限性,并提出了未来可能的研究方向:

  • 非上下文扩展: StreamingLLM 提升了 LLMs 在流式上下文中的效率,但它不扩展模型的上下文窗口,也不增强其长期记忆能力。模型的有效上下文长度仍然受限于其 KV 缓存的大小。这意味着 StreamingLLM 不适用于那些需要依赖超长历史数据和全面上下文理解的任务,例如长文档问答 (Long document QA) 或摘要 (Summarization)。它更适合需要短时记忆的任务,如日常对话和短文档问答。
  • 上下文利用率不足: 实验结果(特别是缓存大小的消融实验)表明,即使模型能够访问更大的上下文窗口,它也可能未能充分利用所有可用的上下文信息。这表明当前 LLMs 在高效整合和利用长距离信息方面仍存在潜在限制。
  • 未来工作方向:
    • 增强上下文利用能力: 未来的研究应致力于增强 LLMs 充分利用其缓存内大量上下文信息的能力。这可能涉及新的模型架构、训练目标或推理策略。
    • 结合上下文扩展方法: StreamingLLM 可以与现有的上下文扩展技术(如位置插值和微调)结合使用,以在保持流式效率的同时,进一步扩大模型可关注的最新上下文范围。

7.3. 个人启发与批判

个人启发:

  • 洞察力驱动的创新: StreamingLLM 的核心在于对“注意力汇聚”这一现象的深刻洞察。它没有试图“修正” SoftMax 函数的这一“副作用”,而是巧妙地利用了它,将其从一个潜在的问题转化为解决方案的一部分。这种从模型内在机制中发现并利用“怪异”行为来解决实际问题的思路非常具有启发性。
  • 工程与理论的结合: 论文不仅提出了一个实用的框架,还通过理论分析(SoftMax 函数和词元可见性)和可视化(注意力图)来解释现象的根源,并用预训练实验来验证优化方案。这种理论洞察与高效工程实践相结合的方式,是高质量学术研究的典范。
  • 无需微调的巨大价值: StreamingLLM 能够在不进行微调的情况下,使现有 LLMs 获得处理无限长文本的能力,这对于实际部署具有巨大的经济和时间成本优势。它降低了 LLMs 在新应用场景中落地的门槛。
  • 预训练阶段的优化潜力: 预训练时引入专用汇聚词元的想法非常精妙。它表明我们可以在模型构建的早期阶段就为特定的部署场景(如流式推理)进行优化,这为未来的 LLM 设计提供了新的方向。

批判与潜在改进:

  • “理解”与“稳定”的区分: 尽管 StreamingLLM 解决了 LLMs 在长文本流式处理中的“稳定性”和“效率”问题,但它并未根本上解决模型对超长上下文的“理解”问题。模型在处理超出其滑动窗口的远距离信息时,仍然无能为力。这意味着对于需要长期记忆、整合多处分散信息的复杂任务,StreamingLLM 本身仍有局限。未来的工作可以探索如何将 StreamingLLM 的效率优势与长文本记忆/检索机制(如基于检索增强生成 RAG)相结合,以实现真正的长文本理解。
  • 注意力汇聚的普遍性与理论深度: 论文在附录中提到,“注意力汇聚”现象也存在于编码器 Transformer (BERT) 和 Vision Transformer (ViT) 中。这暗示了 SoftMax 在 Transformer 架构中的一个普遍特性。如果能对这一现象进行更深入的理论分析,例如,从信息论或优化角度解释 SoftMax 为什么会自然地产生这种“倾倒”行为,可能会进一步提升研究的理论价值。
  • 对下游任务的影响: 论文主要评估了语言建模困惑度和一些问答任务。对于其他需要更强上下文理解或生成质量的下游任务(如长篇写作、复杂代码生成等),StreamingLLM 在“固定窗口”内的表现可能仍需更细致的评估。
  • “注意力汇聚”数量的自适应性: 论文通过消融实验发现 4 个初始词元通常足够,但这个数字是否可以根据模型架构、数据特性或任务需求进行自适应调整?或者模型是否可以动态地选择哪些词元作为注意力汇聚?这将是一个有趣的探索方向。

相似论文推荐

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

暂时没有找到相似论文。