在流式应用中部署大型语言模型（LLMs），例如多轮对话这样的长时间交互场景，是迫切需要的，但存在两大挑战。首先，在解码阶段，缓存先前tokens的键和值状态（KV）会消耗大量内存。其次，流行的LLMs无法泛化到超出训练序列长度的更长文本。窗口注意力是一种自然的方法，它只缓存最近的KVs，但我们发现，当文本长度超过缓存大小时，它会失败。我们观察到一个有趣的现象，即attention sink（注意力汇聚点），保留初始tokens的KV在很大程度上可以恢复窗口注意力的性能。在本文中，我们首先证明attention sink的出现是由于对初始tokens的强烈注意力得分，即使它们在语义上并不重要，也会作为“汇聚点”。 基于上述分析，我们引入了StreamingLLM，这是一个高效的框架，可以使得用有限长度注意力窗口训练的LLMs泛化到无限序列长度，而无需任何微调。我们展示了StreamingLLM能够使Llama-2、MPT、Falcon和Pythia等模型在高达400万个tokens甚至更多的情况下，进行稳定且高效的语言建模。此外，我们发现在预训练过程中添加一个占位符token作为专用的 attention sink 可以进一步改善流式部署。在流式设置中，StreamingLLM比滑动窗口重计算基线快高达22.2倍。代码和数据集在：https://github.com/mit-han-lab/streaming-llm。
1. 引言
大型语言模型（LLMs）正变得无处不在，推动了许多自然语言处理应用的发展，如对话系统、文档摘要、代码完成和问题回答。为了充分发挥预训练LLMs的潜力，它们应能有效准确地进行长序列生成。例如，理想的聊天机器人助手可以稳定地处理最近一天长时间对话的内容。然而，对于LLM来说，泛化到比它们预训练的序列更长的长度是非常具有挑战性的，例如，Llama-2的4K。
原因在于LLMs在预训练期间受到注意力窗口的限制。尽管已经做了大量努力来扩展这个窗口大小并改善长输入的训练和推理效率，可接受的序列长度本质上仍然是有限的，这不允许持久部署。
在本文中，我们首先介绍了LLM流应用的概念，并提出了一个问题：
我们能否在不牺牲效率和性能的情况下，将LLM部署到无限长度的输入上？
图1是StreamingLLM与现有方法的对比说明。该语言模型在长度为 L 的文本上进行了预训练，预测第 T 个token（T \gg L）。
- (a) 密集注意力（Dense Attention）具有 O(T^2) 的时间复杂度和不断增加的缓存大小。当文本长度超过预训练文本长度时，其性能会下降。
- (b) 窗口注意力（Window Attention）缓存最近的 L 个tokens的KV。虽然在推理中效率高，但一旦开始tokens的键和值被驱逐，性能就会急剧下降。
- (c) 带重计算的滑动窗口（Sliding Window with Re-computation）为每个新token重建来自 L 个最近tokens的KV状态。虽然它在长文本上表现良好，但由于在上下文重计算中的二次注意力，其 O(TL^2) 的复杂度使其相当缓慢。
- (d) StreamingLLM 保留了用于稳定注意力计算的attention sink（几个初始tokens），并结合了最近的tokens。它高效并且在扩展文本上提供稳定的性能。使用Llama-2-13B模型在PG-19测试集的第一本书（65K tokens）上测量了困惑度。
当将LLMs应用于无限输入流时，会出现两个主要挑战：
1. 在解码阶段，基于Transformer的LLMs缓存所有先前tokens的Key和Value状态（KV）见图1a，这可能导致过度的内存使用和增加的解码延迟。
2. 现有模型的长度外推能力有限，即当序列长度超出预训练期间设置的注意力窗口大小时，它们的性能会下降。
一种直观的方法，称为窗口注意力，只在最近的tokens的KV状态上维护一个固定大小的滑动窗口（图1b）。尽管它确保了初始填充缓存后的恒定内存使用和解码速度，但一旦序列长度超过缓存大小，模型就会崩溃，即使只是驱逐第一个token的KV（图3）。另一种策略是带有重计算的滑动窗口（图1c），它为每个生成的token重建最近tokens的KV状态。虽然它提供了强大的性能，但由于其窗口内的二次注意力计算，这种方法显著更慢，使得这种方法不适用于实际的流式应用。
为了理解窗口注意力的失败，我们发现了自回归LLMs的一个有趣现象：无论它们与语言建模任务的相关性如何，初始tokens都被分配了惊人的大量注意力得分，见图2。我们将这些tokens称为“attention sinks”。尽管它们缺乏语义意义，但它们收集了重要的注意力得分。我们将原因归因于Softmax操作，它要求所有上下文tokens的注意力得分总和为一。因此，即使当前查询在许多先前的tokens中没有强烈的匹配，模型仍然需要将这些不需要的注意力值分配到某个地方，以使其总和为一。将初始tokens作为sink tokens的原因是直观的：由于自回归语言建模的特性，初始tokens对几乎所有后续tokens都是可见的，使它们更容易被训练以用作attention sinks。
基于上述见解，我们提出了StreamingLLM，一种简单而高效的框架，使得用有限注意力窗口训练的LLMs能够在不进行微调的情况下处理无限长度的文本。StreamingLLM利用了attention sinks具有高注意力值的事实，并保留它们可以保持注意力得分分布接近正常的特点。因此，StreamingLLM只保留attention sink tokens的KV（只需4个初始tokens）以及滑动窗口的KV，以锚定注意力计算并稳定模型的性能。有了StreamingLLM，包括Llama-2-[7, 13, 70]B、MPT-[7, 30]B、Falcon-[7, 40]B和Pythia-[2.9,6.9,12]B在内的模型可以可靠地建模400万个tokens，甚至更多。与唯一可行的基线相比，带有重计算的滑动窗口，StreamingLLM实现了高达22.2倍的加速，实现了LLMs的流式使用。
最后，我们确认了我们的attention sink假设，并展示了语言模型可以通过预训练来只需要一个attention sink token就可以进行流式部署。具体来说，我们建议在所有训练样本的开头添加一个额外的可学习token，作为指定的attention sink。通过从头开始预训练160百万参数的语言模型，我们展示了添加这个单一的sink token可以保持模型在流式情况下的性能。这与普通模型形成对比，后者需要重新引入多个初始tokens作为attention sinks才能达到相同的性能水平。
2. 相关工作
在将LLMs应用于长文本方面已经进行了大量研究，主要集中在三个领域：**长度外推**、**上下文窗口扩展**和**提高LLMs对长文本的利用率**。虽然这些领域看似相关，但值得注意的是，一个方向上的进展并不一定导致另一个方向上的进展。例如，扩展LLMs的上下文大小并不会提高模型超出上下文大小的性能，这两种方法都不能确保有效利用长上下文。我们的StreamingLLM框架主要属于第一类，即将LLMs应用于远远超出预训练窗口大小的文本，甚至可能是无限长度的文本。我们没有扩大LLMs的注意力窗口大小，也没有增强模型在长文本上的记忆和使用率。最后两个类别与我们的重点正交，可以与我们的技术整合。
长度外推：旨在使训练有素的短文本语言模型在测试期间能够处理更长的文本。研究的一个主要方向是为Transformer模型开发相对位置编码方法，使它们能够超越其训练窗口运作。其中一个倡议是Rotary Position Embeddings (RoPE)，它在每个注意力层中转换查询和键，以整合相对位置。尽管它很有前景，但随后的研究表明，它在超出训练窗口的文本上表现不佳。另一种方法，ALiBi，根据查询-键注意力得分之间的距离偏置它们，从而引入相对位置信息。虽然这表现出了改善的外推，但我们在MPT模型上的测试突显了当文本长度远大于训练长度时的崩溃。然而，当前的方法尚未实现无限长度的外推，导致没有现有的LLMs适合流应用。
上下文窗口扩展：集中在扩展LLMs的上下文窗口，使其能够在一个前向传递中处理更多的tokens。主要的工作线解决了训练效率问题。鉴于训练期间注意力计算的二次复杂性，开发长上下文LLM既是计算挑战也是内存挑战。解决方案包括从系统聚焦的优化，如FlashAttention，它加速了注意力计算并减少了内存占用，到近似注意力方法，这些方法以效率为代价交换模型质量。最近，有大量的工作致力于使用RoPE扩展预训练的LLMs，涉及位置插值和微调。然而，所有上述技术只在有限的范围内扩展了LLMs的上下文窗口，这远远不能满足我们论文的主要关注点，即处理无限输入。
提高LLMs对长文本的利用率：旨在优化LLMs，以更好地捕获和利用上下文中的内容，而不仅仅是将它们作为输入。正如刘先生和长聊天所强调的，前两个方向的成功并不一定转化为长上下文的有效利用。在LLMs中有效使用长时间上下文仍然是一个挑战。我们的工作集中在稳定地利用最近的tokens，实现LLMs的无缝流应用。
3. StreamingLLM 方法
3.1 窗口注意力的失败和 attention sink
虽然窗口注意力技术在推理过程中提供了效率，但它导致了极高的语言建模困惑度。因此，模型的性能不适合在流应用中部署。在本节中，我们使用“ attention sink ”的概念来解释窗口注意力的失败，这是StreamingLLM背后的灵感。
确定困惑度激增的点
图3显示了一个包含20K tokens文本的语言建模的困惑度。很明显，当文本长度超过缓存大小时，由于排除了初始tokens，困惑度激增。这表明，无论这些初始tokens与被预测的tokens的距离如何，它们对于维持LLMs的稳定性都至关重要。
为什么在移除初始tokens的KV时LLMs会崩溃？
我们从Llama-2-7B和模型的所有层和头部可视化注意力图，在图2中可以找到。我们发现，除了最底下的两层之外，模型在所有层和头部始终关注初始tokens。这意味着：移除这些初始tokens的KV将从注意力计算中的SoftMax函数（方程）的分母中移除相当大的一部分。这种改变导致注意力得分分布的显著偏移，偏离了正常推理设置中的预期。
\text{SoftMax}(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, \dots, N对于语言建模中初始tokens的重要性有两种可能的解释：(1) 它们的语义至关重要，或者 (2) 模型学习了对它们的绝对位置的偏见。为了区分这些可能性，我们进行了实验（表1），在其中用换行符"\n"替换了前四个tokens。观察结果表明，模型仍然显著强调这些初始换行符tokens。此外，重新引入它们将语言建模困惑度恢复到与拥有原始初始tokens相当的水平。这表明，起始tokens的绝对位置而不是它们的语义价值具有更大的重要性。
LLMs将初始Tokens视为 attention sink
为了解释为什么模型会不成比例地关注初始tokens —— 不考虑它们对语言建模的语义相关性，我们引入了“ attention sink ”的概念。SoftMax函数（方程）的性质阻止了所有被关注的tokens具有零值。这需要从所有层的所有头中的其他tokens聚集一些信息，即使当前的嵌入对其预测已经有足够的自我包含信息。因此，模型倾向于将不必要的注意力值丢弃到特定的tokens。在量化异常值的领域中也做出了类似的观察，这导致提出了SoftMax-Off-by-One作为一种潜在的补救措施。
为什么各种自回归LLMs，如Llama-2、MPT、Falcon和Pythia，始终将它们的 attention sink 集中在初始tokens上，而不是其他tokens？我们的解释很简单：由于自回归语言建模的顺序性质，初始tokens对所有后续tokens都是可见的，而稍后的tokens只对有限的一组后续tokens可见。因此，初始tokens更容易被训练为 attention sink ，捕获不必要的注意力。
在表1中，当我们重新引入最初的四个 token 以及最近的1020个 token（4+1020）时，困惑度得到了恢复。用换行符"\n"替换原始的四个初始 token（4"\n"+1020），也实现了类似的困惑度恢复。缓存配置 x+y 表示添加 x 个初始 token 和 y 个最近的 token。困惑度是在 PG19 测试集的第一本书（65K token）上测量的。
我们注意到，LLMs通常被训练为使用多个初始tokens作为 attention sink ，而不仅仅是一个。如图2所示，引入四个初始tokens作为 attention sink 足以恢复LLM的性能。相比之下，只添加一个或两个并不能实现完全恢复。我们认为这种模式的出现是因为这些模型在预训练期间没有在所有输入样本中包括一个一致的起始token。尽管Llama-2确实在每个段落之前加上了一个“<s>”token，但它是在文本分块之前应用的，导致一个大多数随机的token占据了第零位置。这种缺乏统一起始token导致模型使用几个初始tokens作为 attention sink 。我们假设，通过在所有训练样本的开始处加入一个稳定的可学习token，它可以单独地充当一个承诺的 attention sink ，消除了确保一致流式传输所需的多个初始tokens的需要。我们将在部分中验证这一假设。
3.2 带有 attention sink 的滚动KV缓存
为了在已经训练好的LLMs中实现流式处理，我们提出了一种简单的方法，可以在不进行任何模型微调的情况下恢复窗口注意力的困惑度。除了当前的滑动窗口tokens，我们在注意力计算中重新引入了一些起始tokens的KV。在StreamingLLM中，KV缓存可以在概念上分为两部分，如图4所示：（1） attention sink （四个初始tokens）稳定了注意力计算；（2）滚动KV缓存保留了对语言建模至关重要的最新tokens。StreamingLLM的设计具有多功能性，可以无缝集成到任何使用相对位置编码的自回归语言模型中，如RoPE和ALiBi。
在确定相对距离并向tokens添加位置信息时，StreamingLLM侧重于缓存中的位置，而不是原始文本中的位置。这种区别对StreamingLLM的性能至关重要。例如，如果当前缓存有tokens [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这样的编码，我们在引入旋转变换之前缓存tokens的Keys。然后，我们在每个解码阶段对滚动缓存中的keys应用位置变换。另一方面，与ALiBi的集成更为直接。在这里，连续的线性偏差被应用于注意力得分，而不是'跳跃'偏差。在缓存中分配位置嵌入的这种方法对于StreamingLLM的功能至关重要，确保模型即使超出其预训练的注意力窗口大小也能高效运行。
3.3 使用 attention sink 预训练LLMs
正如在部分中详细说明的那样，模型对多个初始tokens过度关注的一个重要原因是缺乏指定的汇聚点token来卸载过多的注意力得分。因此，模型无意中指定了全球可见的tokens，主要是初始tokens，作为 attention sink 。一个可能的补救措施是故意包括一个全球可训练的 attention sink token，称为“汇聚点Token”，它将作为不必要的注意力得分的存储库。或者，用SoftMax-off-by-One这样的变体替换传统的SoftMax函数，这种方法不需要所有上下文tokens的注意力得分总和为一，也可能是有效的。请注意，这种SoftMax替代方案相当于在注意力计算中使用一个具有全零Key和Value特征的token。我们将这种方法称为“零汇聚点”，以使其在我们的框架中保持一致。
为了验证，我们在相同的设置下从头开始预训练了三个具有1.6亿参数的语言模型。第一个模型使用标准的SoftMax注意力（Vanilla），第二个用\text{SoftMax}_1（Zero Sink）替换了常规的注意力机制，并且一个在所有训练样本中前置了一个可学习的占位符token（Sink Token）。如表3所示，虽然零汇聚点在一定程度上缓解了 attention sink 问题，但模型仍然依赖于其他初始tokens作为 attention sink 。引入汇聚点token在稳定注意力机制方面非常有效。仅仅将这个汇聚点token与最近的tokens配对就足以锚定模型的性能，而且结果评估的困惑度甚至略有提高。鉴于这些发现，我们建议在所有样本中使用汇聚点token来训练未来的LLMs，以优化流式部署。
4. 实验
我们使用四个近期著名的模型家族来评估StreamingLLM：Llama-2、MPT、PyThia和 Falcon。值得注意的是，Llama-2、Falcon 和 Pythia 都采用了RoPE，而MPT采用了ALiBi— 这是近期研究中最有影响力的两种位置编码技术。我们多样化的模型选择确保了我们发现的有效性和稳健性。我们将StreamingLLM与已建立的基线进行对比，如密集注意力、窗口注意力和带有重新计算的滑动窗口方法。在所有后续的StreamingLLM实验中，除非另有说明，我们默认使用四个初始tokens作为 attention sink 。
4.1 在LLM家族和规模上对长文本进行语言建模
我们首先使用拼接的PG19测试集来评估StreamingLLM的语言建模困惑度，该测试集包含100本长书。对于Llama-2模型，缓存大小设置为2048，而对于Falcon、Pythia和MPT模型，它设置为1024。这是为了增强可视化清晰度而选择的预训练窗口大小的一半。
图3显示，StreamingLLM可以在跨越20K tokens的文本上匹配oracle基线（带有重新计算的滑动窗口）的困惑度。与此同时，当输入长度超过其预训练窗口时，密集注意力技术会失败，当输入长度超过缓存大小时，窗口注意力技术会遇到困难，导致初始tokens被驱逐。 在图5中，我们进一步证实，StreamingLLM可以在一系列的模型家族和规模上可靠地处理特别长的文本，包括超过400万个tokens。这包括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可以改善流式LLMs，我们在相同条件下训练了两个各具160百万参数的语言模型。其中一个模型坚持原始的训练设置，另一个在每个训练样本的开始处加入了一个sink token。我们的实验使用了Pythia-160M代码库，并遵循了其训练配方。我们使用去重的Pile数据集，在一台配备8xA6000 NVIDIA GPU的服务器上训练模型。除了将训练批次大小减少到256外，我们保留了所有Pythia训练配置，包括学习率计划、模型初始化和数据集排列。两个模型都训练了143,000步。
收敛性和正常模型性能。 在预训练期间加入sink token对模型的收敛性和后续在一系列NLP基准测试上的性能没有负面影响。 如图6所示，带有sink token的训练模型展现了与其原版对应物相似的收敛动态。我们在七个不同的NLP基准测试上评估了这两个模型，包括ARC-[Challenge, Easy]、LAMBADA、OpenbookQA、PIQA和Winogrande。如表4所示，使用sink token预训练的模型的性能与使用原版方法训练的模型相似。
注意力可视化
图7对比了有无sink token的预训练模型的注意力图。没有sink token的模型，类似于Llama-2-7B（图2），显示早期层的局部注意力和更深层对初始tokens的关注。相比之下，带有sink token的模型在各层和头之间始终集中在sink上，表明了有效的注意力卸载机制。这种对sink的强烈关注，以及对其他初始tokens的关注减少，解释了sink token在增强模型流式性能方面的有效性。
4.3 基于指令调整模型的流式问答结果
为了展示StreamingLLM在现实世界的适用性，我们模拟了使用在现实场景中常用的基于指令调整的LLMs进行的多轮问答。
我们首先将ARC-[Challenge, Easy]数据集中的所有问题-答案对连接起来，将连续流输入到Llama-2-[7,13,70]B-Chat模型中，并使用精确匹配标准在每个答案位置评估模型完成情况。正如表5所示，密集注意力导致内存溢出（OOM）错误，表明它不适合这种设置。虽然窗口注意力方法工作效率高，但当输入长度超过缓存大小时，由于随机输出，其准确性较低。相反，StreamingLLM通过有效处理流式格式，与一次性、逐样本基线准确性保持一致，表现出色。
为了更加贴合StreamingLLM的场景，我们引入了一个受LongEval基准测试启发的数据集StreamEval。如图8所示，与LongEval的单一查询在长跨度设置上不同，我们每10行新信息就查询一次模型。每个查询的答案始终是之前的20行，这反映了现实世界中的问题通常与最近的信息有关。如图9所示，即使输入长度接近120K tokens，使用StreamingLLM的LLMs也能保持合理的准确性。相比之下，密集注意力和窗口注意力分别在预训练文本长度和KV缓存大小上失败。此外，我们使用了两个上下文扩展模型，LongChat-7b-v1.5-32k和Llama-2-7B-32K-Instruct，以显示StreamingLLM可以补充上下文扩展技术。在StreamingLLM中，上下文扩展意味着扩大流式LLMs的最大缓存大小，从而能够捕获更广泛的局部信息。
4.4 消融研究
初始Tokens的数量。 在表2中，我们消融了添加不同数量的初始tokens与最近tokens对流式困惑度的影响。结果显示，仅引入一个或两个初始tokens是不够的，而四个初始tokens似乎足够了，随后的添加只有边际效应。这个结果证明了我们在StreamingLLM中引入4个初始tokens作为attention sinks的选择是正确的。
缓存大小。 在表6中，我们评估了缓存大小对StreamingLLM困惑度的影响。与直觉相反，增加缓存大小并不一致地降低语言建模困惑度。这种不一致性显示了一个潜在的局限性，即这些模型可能没有最大化利用它们接收到的整个上下文。未来的研究工作应该致力于增强这些模型更好地利用广泛上下文的能力。
4.5 效率结果
我们将其解码延迟和内存使用情况与滑动窗口重新计算进行了基准测试，这是唯一具有可接受性能的基线。这两种方法都使用Huggingface Transformers库实现，并在单个NVIDIA A6000 GPU上使用Llama-2-7B和Llama-2-13B模型进行了测试。 如图10所示，随着缓存大小的增加，StreamingLLM的解码速度呈线性增长。而滑动窗口重新计算基线在解码延迟上呈二次方增长。因此，StreamingLLM实现了令人印象深刻的加速，每个token达到高达22.2倍。尽管延迟减少了，但StreamingLLM保持了与重新计算基线一致的内存占用。
5. 结论
在流式应用中部署LLMs是迫切需要的，但由于效率限制和更长文本的性能降低而带来挑战。窗口注意力提供了一个部分解决方案，但当排除初始tokens时，其性能骤降。认识到这些tokens作为“attention sinks”的作用，我们引入了StreamingLLM——一个简单而高效的框架，使LLMs能够在不进行微调的情况下处理无限的文本。通过添加attention sinks和最近的tokens，StreamingLLM可以有效地对多达400万个tokens的文本进行建模。我们进一步表明，使用专用的sink token进行预训练模型可以改善流式性能。StreamingLLM首先解耦了LLM的预训练窗口大小和其实际文本生成长度，为LLMs的流式部署铺平了道路。
