EFFICIENT STREAMING LANGUAGE MODELS WITH ATTENTION SINKS翻译

摘要

迫切需要在流应用（例如多轮对话）中部署大型语言模型（LLM），这些应用需要长上下文交互，但会带来两个主要挑战。首先，在解码阶段，缓存先前token的key和value状态（KV）会消耗大量内存。 其次，流行的LLM不能推广到比训练长度序列更长的文本。窗口注意力机制（仅缓存最近的 KV）是一种自然的方法，但我们表明，当文本长度超过缓存大小时，它会失败。 我们观察到一个有趣的现象，即attention sink，保持初始 token 的 KV 将在很大程度上恢复窗口注意力的性能。在本文中，我们首先证明attention sink的出现是由于将初始token具有较大的注意力分数，尽管它们在语义上并不重要。基于上述分析，我们引入了 StreamingLLM，这是一个高效的框架，能够使用有限长度的注意力窗口训练 LLM，并泛化到无限序列长度，而无需任何微调。我们证明 StreamingLLM 可以使 Llama-2、MPT、Falcon 和 Pythia 能够使用多达 400 万个甚至更多的token执行稳定且高效的语言建模。此外，我们发现在预训练期间添加占位符token作为专用注意力池可以进一步改进流部署。在流设置中，StreamingLLM 的性能比滑动窗口重新计算基线高出 22.2 倍。链接中提供了代码和数据集（https://github.com/mit-han-lab/streaming-llm）。

1.介绍

　　大型语言模型 (LLM) 正变得无处不在，为许多自然语言处理应用提供支持，例如对话系统、文档摘要、代码补全和问答。为了充分发挥预训练LLM的潜力，他们应该能够高效、准确地执行长序列生成。例如，理想的 ChatBot 助手可以稳定地处理最近一整天的对话内容。然而，LLM 很难推广到比预训练更长的序列长度，例如 Llama-2的 4K。
　　原因是LLM在预训练期间受到的注意力窗口的限制。尽管付出了大量努力来扩大此窗口大小并提高长输入的训练和推理效率，但可接受的序列长度本质上仍然是有限的，这不允许持久部署。
　　在本文中，我们首先介绍LLM流式应用的概念并提出以下问题：
　　我们能否在不牺牲效率和性能的情况下为无限长度的输入部署LLM？
　　当将LLM应用于无限的输入流时，会出现两个主要挑战：

1. 在解码阶段，基于 Transformer 的 LLM 会缓存所有先前token的key和value状态 (KV)，如图 1 (a) 所示，这可能会导致内存使用过多并增加解码延迟。
2. 现有模型的长度外推能力有限，即当序列长度超出预训练期间设置的注意力窗口大小时，其性能会下降。

一种直观的方法，称为窗口注意力（图 1 b），仅在最近token的 KV 状态上维护固定大小的滑动窗口。尽管它在缓存被填满后确保了恒定的内存使用和解码速度，但一旦序列长度超过缓存大小，模型就会崩溃，如图 3 所示。另一种策略是带有重计算的滑动窗口（如图1c所示），它为每个生成的token重建最近token的KV状态。 虽然它提供了强大的性能，但由于在其窗口内计算二次注意力，这种方法的速度明显慢，使得这种方法对于现实世界的流应用来说不切实际。
　　为了理解窗口注意力的失败，我们发现了自回归 LLM 的一个有趣现象：大量的注意力分数被分配给初始token，而不管它们与语言建模任务的相关性如何，如图 2 所示。我们将这些称为token的“注意力池”。尽管它们缺乏语义意义，但它们收集了大量的注意力分数。我们将原因归因于 Softmax 操作，该操作要求所有上下文token的注意力分数总和为 1。因此，即使当前query在许多先前的token中没有很强的匹配，模型仍然需要在某个地方分配这些不需要的注意力值，因此它的总和为1。初始token作为注意力池的背后原因很直观：初始token对几乎所有后续token都是可见的，因为自回归语言建模的性质，使它们更容易被训练为注意力池。
　　基于上述见解，我们提出了 StreamingLLM，这是一个简单而高效的框架，使在有限注意力窗口下训练的 LLM 能够处理无限长度的文本，而无需进行微调。StreamingLLM 利用了注意力池具有高注意力值的事实，保留它们可以保持注意力分数分布接近正常。因此，StreamingLLM 只是将注意力集中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 倍的加速，实现了 LLM 的流式使用。

2.RELATED WORK

3.STREAMINGLLM

3.2 ROLLING KV CACHE WITH ATTENTION SINKS

　　为了在已经训练好的LLM中启用LLM流，我们提出了一种简单的方法，可以恢复窗口注意力的困惑，而无需任何模型微调。除了当前的滑动窗口token之外，我们在注意力计算中重新引入了一些起始token的 KV。StreamingLLM中的KV缓存在概念上可以分为两部分，如图4所示：（1）注意力池（四个初始token）稳定注意力计算；2）滚动KV缓存保留最新的token，这对于语言建模至关重要。StreamingLLM 的设计用途广泛，可以无缝地融入任何采用相对位置编码的自回归语言模型，例如 RoPE 和 ALiBi。
　　在确定相对距离并向token添加位置信息时，StreamingLLM 重点关注缓存中的位置而不是原始文本中的位置。这种区别对于 StreamingLLM 的性能至关重要。例如，如果当前缓存有token [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 这样的编码，我们在引入旋转变换之前缓存token的key。然后，我们在每个解码阶段对滚动缓存中的key应用位置变换。另一方面，与ALiBi的集成更加直接。在这里，对注意力分数应用连续线性偏差，而不是“跳跃”偏差。这种在缓存中分配位置嵌入的方法对于 StreamingLLM 的功能至关重要，确保模型即使超出其预训练注意窗口大小也能高效运行。

3.3 PRE-TRAINING LLMS WITH ATTENTION SINKS

　　正如第 3.1 节所述，模型过度关注多个初始token的一个重要原因是缺乏指定的接收器token来卸载过多的注意力分数。因此，该模型无意中将全局可见的token（主要是初始标记）指定为注意力池。一种潜在的补救措施是有意包含一个全局可训练的注意力池token，表示为“Sink Token”，它将用作不必要的注意力分数的存储库。或者，用 SoftMax-off-by-One 等变体替换传统的 SoftMax 函数，
$$SoftMa{x}_1(x{)}_i=\frac{e^{x_i}}{1+{\sum }_{j=1}^N{e}^{x_j}}\text{\hspace{1em}(2)}$$
　　它不需要所有上下文token的注意力分数总和为 1，也可能是有效的。请注意，此 SoftMax 替代方案等效于在注意力计算中使用具有全零 Key 和 Value 特征的token。我们将此方法称为“Zero Sink”，以使其始终适合我们的框架。
　　为了进行验证，我们在相同的设置下从头开始预训练具有 1.6 亿个参数的三种语言模型。 第一个模型利用标准 SoftMax 注意力（Vanilla），第二个模型用 $SoftMa{x}_1$（Zero Sink）替换常规注意力机制，以及在所有训练样本中prepend可学习的占位符token（Sink Token）。如表3所示，虽然Zero Sink在一定程度上缓解了注意力池问题，但模型仍然依赖其他初始标记作为注意力池。引入sink token对于稳定注意力机制非常有效。简单地将这些token与最近的token配对就足以满足模型的性能，并且由此产生的评估困惑度甚至略有改善。鉴于这些发现，我们建议在所有样本中使用接收器token来训练未来的LLM，以优化流式部署。

4.实验