【论文笔记】Efficient Streaming Language Models with Attention Sinks

原文

摘要

在流式应用中部署大型语言模型（LLMs），例如多轮对话，是迫切需要的，但存在两大挑战。首先，在解码阶段，缓存先前token的Key和Value状态（KV）会消耗大量内存。其次，流行的LLMs无法泛化到比训练序列长度更长的文本。窗口注意力，即只缓存最近的KVs，是一种自然的方法——但我们展示了当文本长度超过缓存大小时它就会失败。我们观察到一个有趣的现象，即attention sink，即使初始token在语义上不重要，保留KV也能在很大程度上恢复窗口注意力的性能。在本文中，我们首先展示了注意力sink的出现是由于对初始token的强烈注意力分数作为一个“sink”，即使它们在语义上不重要。基于以上分析，我们引入了StreamingLLM，这是一个高效的框架，使得训练时具有有限长度注意力窗口的LLMs能够泛化到无限序列长度，无需任何微调。我们展示了StreamingLLM可以使Llama-2、MPT、Falcon和Pythia等模型在4百万个token以上的表现稳定和高效。此外，我们发现在预训练期间添加一个占位符token作为专用的注意力汇可以进一步改善流式部署。在流式设置中，StreamingLLM的性能超过了滑动窗口重计算基线，速度提升了高达22.2倍。代码和数据集提供了链接。

1 引言

大型语言模型（LLMs）如雨后春笋般涌现，为许多自然语言处理应用提供动力，例如对话系统、文档摘要、代码补全和问答系统。为了充分发挥预训练LLMs的潜力，它们应该能够高效、准确地执行长序列生成。例如，一个理想的ChatBot助手可以稳定地处理整天的对话内容。然而，对于LLM来说，泛化到比它们预训练过的序列长度更长的序列是非常具有挑战性的，例如Llama-2的4K。原因是LLMs在预训练期间受到注意力窗口的限制。尽管有大量努力扩大这个窗口大小并提高长输入的训练和推理效率，但可接受的序列长度仍然是有限的，这不允许持久部署。

在本文中，我们首先介绍了LLM流式应用的概念，并提出了一个问题：我们能否在不牺牲效率和性能的情况下部署一个LLM来处理无限长度的输入？

在应用LLMs于无限输入流时，出现了两个主要挑战：

1. 在解码阶段，基于Transformer的LLMs缓存了所有先前token的Key和Value状态（KV），如图1(a)所示，这可能导致内存使用过多和解码延迟增加。

2. 现有模型的序列长度外推能力有限，即当序列长度超出预训练期间设定的注意力窗口大小时，它们的性能会下降。

一种直观的方法，称为窗口注意力（Beltagy et al., 2020）如图1(b)所示，只维护最近token的KV状态的固定大小滑动窗口。尽管这确保了在缓存最初填充后内存使用和解码速度保持恒定，但一旦序列长度超过缓存大小时，模型就会崩溃，即甚至只是逐出第一个token的KV，如图3所示。另一种策略是滑动窗口与重计算（如图1©所示），它为每个生成的token重建最近token的KV状态。虽然这种方法性能强劲，但由于其在窗口内计算二次方的注意力，使得这种方法显著变慢，使得这种方法不适用于实际的流式应用。

为了理解窗口注意力的失败，我们发现了一个有趣的现象：自回归LLMs分配了相当大的注意力分数给初始token，无论它们与语言建模任务的相关性如何，我们都将这些token称为“attention sink”。尽管它们在语义上不重要，但它们却收集了大量的注意力分数。我们将其原因归因于Softmax操作，该操作要求所有上下文token的注意力分数之和必须为一。因此，即使当前查询与许多先前的token没有强烈的匹配，模型仍然需要在某个地方分配这些不需要的注意力值，以便它们的总和为一。初始token作为汇token的原因是直观的：由于自回归语言建模的特性，初始token几乎对所有后续token都是可见的，这使得它们更容易被训练为注意力sink。

基于上述洞察，我们提出了StreamingLLM，一个简单高效的框架，使得训练时具有有限注意力窗口的LLMs能够处理无限长度的文本，无需微调。StreamingLLM利用了attention sink具有高注意力值的事实，保留它们可以保持注意力分数分布接近正常。因此，StreamingLLM简单地保留了attention sink token的KV（只需4个初始token就足够了）以及滑动窗口的KV，以锚定注意力计算并稳定模型的性能。有了StreamingLLM，包括Llama-2-[7, 13, 70]B、MPT-[7, 30]B、Falcon-[7, 40]B和Pythia[2.9,6.9,12]B在内的模型可以可靠地建模400万个token，甚至更多。与唯一可行的基线相比，即滑动窗口与重计算，StreamingLLM实现了高达22.2倍的速度提升，实现了LLMs的流式使用。

此外，我们证实了我们的注意力汇假设，并证明语言模型可以预训练，以便在流式部署中仅需要一个注意力sink token。具体来说，我们建议在所有训练样本的开头添加一个额外的可学习token，作为指定的注意力sink。通过从头开始预训练1.6亿参数的语言模型，我们展示了添加这个单一汇token可以保持模型在流式情况下的性能。这与普通模型形成对比，普通模型需要重新引入多个初始token作为注意力sink，以达到相同的性能水平。

最后，我们强调StreamingLLM能够高效地从KV缓存中的token生成连贯的文本，而不需要扩展LLMs的上下文长度。它适合于连续操作需求，使用最小的内存和对过去数据的依赖。此外，StreamingLLM可以补充上下文扩展方法，以增加可关注的近期上下文。

2 相关工作

大量研究已经集中在将大型语言模型（LLMs）应用于长文本上，主要有三个研究领域：长度外推、上下文窗口扩展和提高LLMs对长文本的利用。虽然看似相关，但值得注意的是，一个方向上的进步并不一定导致另一个方向上的进步。例如，扩展LLMs的上下文大小并不提高模型的性能超过上下文大小，这两种方法都不能确保有效地使用长上下文。我们的StreamingLLM框架主要属于第一类别，即将LLMs应用于明显超出预训练窗口大小的文本，甚至可能是无限长度。我们不扩展LLMs的注意力窗口大小或增强模型对长文本的内存和使用。最后两个类别与我们的重点正交，可以与我们的技术集成。

长度外推

长度外推的目标是使语言模型能够处理比它们在训练时见过的序列更长的文本。一个主要的研究途径是开发相对位置编码方法，使Transformer模型能够超越其训练窗口。例如，旋转位置编码（RoPE）通过在每个注意力层中转换查询和键来集成相对位置信息。尽管有希望，但后续研究表明，当文本超出训练窗口时，其性能不佳。另一种方法，ALiBi，根据它们之间的距离对查询-键注意力分数进行偏差调整，从而引入相对位置信息。虽然这显示了改进的外推，我们在MPT模型上的测试突出了当文本长度远远大于训练长度时的崩溃。目前的方法还没有实现无限长度的外推，导致没有现有的LLMs适合流式应用。

上下文窗口扩展

上下文窗口扩展的重点是扩展LLMs的上下文窗口，使得在一个前向传递中能够处理更多的token。鉴于注意力计算的二次复杂度，开发长上下文LLM既是计算也是内存挑战。解决方案包括系统级优化，如FlashAttention，它加速了注意力计算并减少了内存占用，以及近似注意力方法，它们牺牲了模型质量以换取效率。最近，有关通过RoPE扩展预训练LLMs的工作，包括位置插值和微调。然而，所有上述技术只将LLMs的上下文窗口扩展到有限的程度，这不足以处理无限输入。

提高LLMs对长文本的利用

提高LLMs对长文本的利用优化了LLMs，使其更好地捕获和使用上下文中的内容，而不仅仅是将它们作为输入。如Liu等人和Li等人所强调的，前两个方向的成功并不一定转化为对长上下文的熟练利用。解决LLMs在长上下文中的有效使用仍然是一个挑战。我们的工作集中在稳定地利用最近的token，使LLMs能够无缝地应用于流式应用。

3 StreamingLLM

3.1 窗口注意力的失败和attention sink

虽然窗口注意力技术在推理期间提供了效率，但它导致了极高的语言模型困惑度。因此，模型的性能不适合部署在流式应用中。在本节中，我们使用注意力汇的概念来解释窗口注意力的失败，作为StreamingLLM背后的灵感来源。

识别困惑度激增的点。图3显示了在20K token文本上的语言模型困惑度。显然，当文本长度超过缓存大小时，困惑度激增，这导致排除了初始token。这表明，无论初始token与预测token的距离如何，初始token对于维持LLMs的稳定性至关重要。

为什么当移除初始token的KV时LLMs会崩溃？我们从所有层和头的视觉化注意力图中发现，如图2所示，对于Llama-2-7B和模型，超过底部两层后，模型在所有层和头中始终关注初始token。这意味着，移除这些初始token的KV将从SoftMax函数中移除相当一部分分母，这将导致注意力分数分布从正常推理设置中预期的分布发生显著变化。
$SoftMax(x{)}_i=\frac{e^{x_i}}{(e^{x_1}+{\sum }_{j=2}^N(e^{x_j}))}$，其中$x_1\gg x_j,j\in 2,...,N$

初始token在语言建模中的重要性有两个可能的解释：(1) 它们的语义至关重要，或者 (2) 模型学会了对它们的绝对位置的偏好。为了区分这些可能性，我们进行了实验（表1），将前四个token替换为换行符“\n”。观察结果表明，模型仍然显著强调这些初始换行符token。此外，重新引入它们将语言模型困惑度恢复到与拥有原始初始token相当的水平。这表明，起始token的绝对位置，而不是它们的语义价值，更为重要。

LLMs将初始Token作为attention sink。为了解释为什么模型会不成比例地关注初始token——无论它们与语言建模的语义相关性如何，我们引入了“attention sink”的概念。SoftMax函数的性质（方程1）阻止了所有被关注的token的值都为零。这要求在所有层的所有头中，从其他token中聚合一些信息，即使当前的嵌入已经包含了足够的信息来进行预测。因此，模型倾向于将不必要的注意力值转储到特定的token上。在量化异常值（Xiao et al., 2023; Bondarenko et al., 2023）的领域中也观察到了类似的模式，这导致了SoftMax-Off-by-One（Miller, 2023）作为潜在解决方案的提出。

3.2 带有注意力sink的滚动KV缓存

为了使已经训练好的LLMs能够进行LLM流式处理，我们提出了一个简单的方法，可以在不进行任何模型微调的情况下恢复窗口注意力的困惑度。除了当前滑动窗口token之外，我们重新引入了一些起始token的KV在注意力计算中。StreamingLLM中的KV缓存可以概念上分为两部分，如图4所示：（1）注意力汇（四个初始token）稳定注意力计算；（2）滚动KV缓存保留了最近的token，这对于语言模型至关重要。StreamingLLM的设计是通用的，可以无缝集成到任何使用相对位置编码的自回归语言模型中，如RoPE和ALiBi。

当确定token之间的相对距离并添加位置信息时，StreamingLLM关注的是缓存内的位置，而不是原始文本中的位置。对于StreamingLLM的性能来说，这种区分至关重要。例如，如果当前缓存（见图4）包含token [0, 1, 2, 3, 6, 7, 8]，并且正在解码第9个token，那么分配的位置是[0, 1, 2, 3, 4, 5, 6, 7]，而不是原始文本中的位置，那将是[0, 1, 2, 3, 6, 7, 8, 9]。

对于RoPE这样的编码，我们会在引入旋转变换之前缓存token的Keys。然后，在每个解码阶段对滚动缓存中的keys应用位置变换。对于ALiBi的集成则更直接，这里应用的是连续的线性偏差，而不是“跳跃”偏差到注意力分数上。在缓存内分配位置嵌入的方法对于StreamingLLM的功能至关重要，确保模型即使在其预训练注意力窗口大小之外也能高效运作。

3.3 预训练带有注意力sink的LLMs

如第3.1节所述，模型对多个初始token过度关注的一个重要原因是缺乏一个指定的sink token来卸载过多的注意力分数。因此，模型无意中使用全局可见的token，主要是初始的token，作为注意力sink。一个可能的补救措施是有意地包含一个全局可训练的注意力sink token，这将作为一个存储不必要注意力分数的仓库。或者，用一个变体替换传统的SoftMax函数，如SoftMax-off-by-One，可能也是有效的。注意，SoftMax1等同于在注意力计算中添加一个带有全零Key和Value特征的token。我们将这种方法称为“Zero Sink”以适应我们的框架。

为了验证我们的假设，即在所有预训练样本中引入一个汇token可以改善流式LLMs，我们从零开始训练了三个具有1.6亿参数的语言模型，设置完全相同。第一个模型使用标准的SoftMax注意力（普通），第二个模型将常规注意力机制替换为SoftMax1（Zero Sink），还有一个在所有训练样本的开头添加了一个可学习的占位符token（Sink Token）。如表3所示，虽然零汇在一定程度上缓解了注意力sink问题，但模型仍然依赖其他初始token作为注意力sink。引入一个sink token在稳定注意力机制方面非常有效。简单地将这个汇token与最近的token配对就足以锚定模型的性能，并且得到的评估困惑度甚至略有提高。鉴于这些发现，我们建议未来的LLMs在所有样本中训练时都包含一个sink token，以优化流式部署。

4 实验

我们使用四个著名的最新模型系列评估StreamingLLM：Llama-2（Touvron et al., 2023b）、MPT（Team, 2023）、PyThia（Biderman et al., 2023）和Falcon（Almazrouei et al., 2023）。特别是，Llama-2、Falcon和Pythia包含RoPE（Su et al., 2021），而MPT使用ALiBi（Press et al., 2022）——这是最近研究中最有影响力的两种位置编码技术。我们多样化的模型选择确保了我们发现的有效性和鲁棒性。我们将StreamingLLM与建立的基线进行基准测试，如密集注意力、窗口注意力和滑动窗口方法与重新计算。在所有后续的StreamingLLM实验中，我们默认使用四个初始token作为注意力sink，除非另有说明。

4.1 跨LLM家族和规模的长文本语言建模

我们首先评估StreamingLLM在连接的PG19测试集上的语言建模困惑度，该测试集包含100本长书。对于Llama-2模型，缓存大小设置为2048，而对于Falcon、Pythia和MPT模型，设置为1024。这是为了增强可视化清晰度而选择的一半预训练窗口大小。

图3显示了StreamingLLM在20K token文本上的语言模型困惑度，可以与滑动窗口与重计算的基线相匹配。与此同时，密集注意力技术在输入长度超过其预训练窗口时失败，窗口注意力技术在输入长度超过缓存大小时挣扎，导致初始token被逐出。在图5中，我们进一步证实了StreamingLLM可以可靠地处理超过400万个token的异常长文本，涵盖了各种模型家族和规模。这包括Llama-2-[7,13,70]B、Falcon-[7,40]B、Pythia-[2.8,6.9,12]B和MPT-[7,30]B。

4.2 引入 Sink Token 预训练的实验结果

为了验证我们关于在所有预训练样本中引入 sink token 可以提升流式大型语言模型（StreamingLLM）性能的建议，我们在相同条件下训练了两个具有1.6亿参数的语言模型。一个模型使用了原始的训练设置，另一个则在每个训练样本的开头加入了一个 sink token。我们的实验采用了 Pythia-160M（Biderman等，2023）的代码库，并按照其训练方案进行。训练在一台配备8块 NVIDIA A6000 GPU 的服务器上进行，使用了去重后的 Pile 数据集（Gao等，2020）。除将训练批量大小减少到256外，我们保留了所有 Pythia 的训练配置，包括学习率调度、模型初始化和数据集的排列顺序。两个模型都训练了143,000步。

收敛性和常规模型性能

在预训练中引入sink token并不会对模型的收敛性或在多种自然语言处理基准测试上的性能产生负面影响。如图6所示，包含sink token的模型在收敛动态上与其原始模型（vanilla model）相似。我们在七个不同的自然语言处理基准上对这两种模型进行了评估，包括ARC-[Challenge, Easy]、HellaSwag、LAMBADA、OpenbookQA、PIQA和Winogrande。表4显示，经过sink token预训练的模型在性能上与使用原始方法训练的模型相当。

流式性能

如表3所示，传统方法训练的模型与加入 sink token 的模型在流式困惑度上存在差异。值得注意的是，标准模型需要加入多个 token 作为注意力汇聚点（attention sinks）以维持稳定的流式困惑度，而经过 sink token 训练的模型则仅需使用该 token 即可实现令人满意的流式性能。

Attention可视化

图7对比了经过和未经过sink token预训练的模型的注意力图。未包含sink token的模型（如图2中Llama-2-7B所示）在初始层表现出“局部”注意力，而在较深层则集中关注初始token。相比之下，预训练中加入sink token的模型在所有层和头上都专注于sink token，这表明sink token作为一种有效的注意力分散机制。sink token的这种集中作用减少了对其他初始token的关注，从而支持了sink token在增强模型流式性能方面的效果。

4.3 流式问答与指令调整模型的结果

为了展示StreamingLLM在现实世界中的应用性，我们模拟了多轮问答，使用指令调整的LLMs，这在现实世界场景中常用。

我们首先将ARC-[Challenge, Easy]数据集中的所有问答对连接起来，将连续的流输入到Llama-2-[7,13,70]B-Chat模型，并使用精确匹配标准在每个答案位置评估模型完成情况。如表5所示，密集注意力导致内存溢出（OOM）错误，显示出在这种设置中的不适用性。虽然窗口注意力方法工作高效，但由于在输入长度超过缓存大小时随机输出，其准确性较低。相反，StreamingLLM在处理流式格式方面表现出色，与逐个样本的一次性基线准确性一致。

为了更贴合StreamingLLM的场景，我们引入了一个受LongEval（Li et al., 2023）基准启发的数据集StreamEval。如图8所示，与LongEval的单次查询长跨度设置不同，我们每10行新信息查询模型一次。每个查询的答案始终是20行之前的内容，反映了现实世界中问题通常涉及最近信息的情况。如图9所示，使用StreamingLLM的LLMs即使在输入长度接近120K个token时也能保持合理的准确性。相比之下，密集和窗口注意力都在预训练文本长度和KV缓存大小时失败。此外，我们使用了两个上下文扩展模型，LongChat-7b-v1.5-32k（Li et al., 2023）和Llama-2-7B-32KInstruct（Together, 2023），以展示StreamingLLM可以补充上下文扩展技术。在StreamingLLM中，上下文扩展意味着扩大流式LLMs的最大缓存大小，使模型能够捕获更广泛的局部信息。

4.4 消融实验

初始 Token 的数量

在表2中，我们测试了引入不同数量的初始 token 与最近 token 对流式困惑度的影响。实验结果表明，仅引入一两个初始 token 是不够的，而达到四个初始 token 时效果显著，之后增加初始 token 的作用微乎其微。因此，我们在 StreamingLLM 中选择了4个初始 token 作为注意力汇聚点的合理配置。

缓存大小

在表6中，我们评估了缓存大小对 StreamingLLM 困惑度的影响。出乎意料的是，增加缓存大小并不总是能显著降低语言建模的困惑度。这种不一致性揭示了一个潜在的局限性，即模型可能无法最大限度地利用其接收到的全部上下文信息。未来的研究应致力于增强这些模型有效利用大范围上下文的能力。

4.5 效率结果

我们对StreamingLLM的解码延迟和内存使用与滑动窗口方法与重新计算进行了基准测试，这是唯一一个质量可接受的基线。这两种方法都使用Huggingface Transformers库实现，并在单个NVIDIA A6000 GPU上使用Llama-2-7B和Llama-2-13B模型进行测试。如图10所示，随着缓存大小的增加，StreamingLLM的解码速度线性增长。滑动窗口与重新计算基线的解码延迟呈二次方增长。因此，StreamingLLM实现了显著的速度提升，每个token高达22.2倍。尽管延迟降低，StreamingLLM保持了与重新计算基线一致的内存占用。

5 结论

在流式应用中部署LLMs是迫切需要的，但由于效率限制和长文本性能下降而面临挑战。窗口注意力提供了部分解决方案，但当初始token被排除时，其性能会急剧下降。认识到这些token作为“注意力sink”的作用，我们引入了StreamingLLM——一个简单高效的框架，使LLMs能够处理无限文本，无需微调。通过添加注意力汇和最近的token，StreamingLLM可以高效地建模高达400万个token的文本。我们进一步展示了预训练模型与专用汇token可以改善流式性能。StreamingLLM首先将LLMs的预训练窗口大小与其实际文本生成长度解耦，为LLMs的流式部署铺平了道路。