论文解读: streaming-LLM 使各种模型稳定、高效地处理长达400万tokens的文本

1. 论文题目
Efficient streaming language models with attention sinks

2. 论文链接
https://arxiv.org/pdf/2309.17453v1.pdf

3. 项目链接
https://github.com/mit-han-lab/streaming-llm

友情提示：

为了保持论文中原汁原味的语义信息，保留相关原文，没有刻意翻译成中文，避免造成误解。

Abstract

该论文探讨了大型语言模型(LLMs)在流式应用中的部署挑战。主要问题是在解码时缓存先前tokens的KV状态占用大量内存，并且流行的LLMs不能推广到比训练序列更长的文本。尽管窗口注意力(windows attention)是一种方法，但当文本长度超过缓存大小时，它会失败。通过观察“注意力汇聚 attention sink”现象，保留初始tokens的KV可大大恢复性能。基于此，研究引入了StreamingLLM框架，使得有限长度注意窗口训练的LLMs能够推广到无限序列长度。实验证明，StreamingLLM可使各种模型稳定、高效地处理长达400万tokens的文本，并通过添加占位符token进一步优化性能。

一、Introduction

大型语言模型（LLMs）已广泛应用于多种自然语言处理任务，但长序列生成仍是个挑战。LLMs在预训练时受到注意力窗口的限制，导致它们在超过预训练长度时性能下降。论文首先介绍了LLM流式应用的概念，并探讨了两个主要挑战： 1. 解码阶段的内存使用和延迟； 2. 模型在超过预训练注意力窗口大小时的性能退化。

尽管有滑动窗口等方法，但效率仍然受到挑战。研究发现，模型中存在"注意力汇聚 attention sink "的现象，即初步tokens获得大量的注意力分数，而与其语义重要性无关。基于这一发现，提出了StreamingLLM框架，使得LLMs能处理无限长度的文本。该方法通过保留注意力汇聚tokens的KV来稳定模型性能，与其他方法相比，速度提高了22.2倍。此外，研究还证实了通过在预训练阶段加入一个额外的可学习token (learnable token)，可以进一步优化流式部署的性能。

二、Related Work

在应用大型语言模型(LLMs)于长文本方面已有大量研究，主要集中在三个方向：长度外推、上下文窗口扩展以及提高LLMs对长文本的利用。这些方向虽然相关，但进展并不总是相互促进。例如，扩展LLMs的上下文大小并不一定改善模型性能。StreamingLLM框架主要集中在第一个方向，不扩展LLMs的注意力窗口大小，也不增强其对长文本的记忆和利用。

1. 长度外推：旨在使训练于较短文本的语言模型在测试时处理更长的文本。主要方法是为Transformer模型开发相对位置编码方法。RoPE和ALiBi是此方向的两种方法，但都存在限制。

2. 上下文窗口扩展：重点是扩展LLMs的上下文窗口，以便在一次前向传递中处理更多tokens。主要工作集中在解决训练效率问题，如FlashAttention和近似注意力方法。尽管如此，这些技术仅在有限程度上扩展了LLMs的上下文窗口。

3. 提高LLMs对长文本的利用：优化LLMs以更好地捕获和使用上下文中的内容。尽管在前两个方向取得了进展，但LLMs如何有效地使用长上下文仍是一个挑战。

总的来说，此研究的核心是稳定地利用最近的tokens，使LLMs能够无缝地进行流式应用。

三、Streaming LLM

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

尽管窗口注意力(window attention)在inference期间提供了效率，但它导致了极高的语言建模困惑度，使模型的性能不适于流式应用中的部署。本节通过注意力汇聚(attention sinks)的概念来解释窗口注意力的失败，这也是StreamingLLM背后的灵感来源。

3.1.1 确定困惑度激增的点 当文本长度超过缓存大小时，困惑度会激增，主要是由于初步tokens的排除。这表明，初始tokens对维持LLMs的稳定性至关重要。

3.1.2 为什么移除初始tokens的KV会导致LLMs崩溃？ 观察Llama-2-7B的注意力图，发现模型一致地关注初始tokens。移除这些初始tokens的KV将导致注意力分数的分布发生重大变化。实验结果表明，模型重视初始tokens的绝对位置，而不是它们的语义价值。

3.1.3 LLM将初始tokens作为注意力汇聚(attention sinks) 由于SoftMax函数的性质，模型趋向于将不必要的注意力值丢弃到特定的tokens。各种自回归LLMs，如Llama-2, MPT, Falcon和Pythia，都倾向于关注初始tokens作为它们的注意力汇聚。这是因为在自回归语言建模的顺序性质下，初始tokens对所有后续tokens都是可见的。因此，初始tokens更容易被训练为注意力汇聚。

3.1.4 为何需要多个初始tokens作为注意力汇聚？ 实验表明，引入四个初始tokens作为注意力汇聚可以恢复LLM的性能。可以假设，通过在所有训练样本的开始加入一个稳定的可学习tokens，它可以单独充当一个专注的注意力汇聚(attention sink)，从而无需多个初始tokens来确保一致的流式传输。

3.2 Rolling KV cache with attention sinks

为了在已训练的LLMs中实现流式处理，作者提出了一个直接的方法，无需对模型进行微调即可恢复窗口注意力的困惑度。StreamingLLM的KV缓存分为两部分：1) 注意力汇聚（四个初始tokens）稳定注意力计算；2) 滚动KV缓存保留最近的tokens。此设计可以无缝整合到使用相对位置编码的任何自回归语言模型中。StreamingLLM在确定相对距离和添加位置信息时，关注缓存中的位置而非原文。此策略对于模型效率至关重要，确保其超出预训练注意窗口大小时仍能高效运行。

3.3 Pre-training LLMs with attention sinks

模型过度关注多个初始tokens的一个重要原因是缺少一个专门的sink token来卸载过多的注意力分数。因此，模型无意中指定了主要是初始的、全局可见的token作为注意力汇聚(attention sinks)。 ① 一个可能的解决方法是有意地包括一个全局可训练的注意力汇聚token，称为“sink token"。 ② 或者，用类似SoftMax-off-by-One的变体替换传统的SoftMax函数也可能有效。 为了验证，我们从头开始使用160百万参数预训练了三个语言模型。如表3所示，虽然 Zero Sink 在一定程度上缓解了注意力汇聚的问题，但模型仍然依赖于其他初始tokens作为注意力汇聚。引入sink token在稳定注意机制上非常有效。 简单地将这个sink token与最近的tokens配对足以锚定模型的性能，而且结果评估的困惑度甚至略有改善。鉴于这些发现，我们建议在所有样本中使用一个sink token来训练未来的LLMs，以优化流式部署。

四、Experiments

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

我们首先使用拼接的PG19测试集评估StreamingLLM的语言建模困惑度PPL，该测试集包含100本长书。对于Llama-2模型，缓存大小设置为2048；对于Falcon、Pythia和MPT模型，设置为1024。图-3显示StreamingLLM可以在跨越 20K tokens的文本上匹配oracle基线的困惑度。同时，当输入长度超过其预训练窗口时，dense attention技术失败；当输入长度超过缓存大小时，window attention技术出现困难。图-5进一步证实，StreamingLLM可以可靠地处理超过400万tokens的异常扩展文本。

4.2 使用一个sink token进行预训练的结果

为验证在所有预训练样本中引入sink token能改善流式LLMs的假设，我们训练了两个各有1.6亿参数的语言模型。一个模型遵循原始设置，另一个在每个训练样本开始时加入了sink token。在Pythia-160M代码库的指导下进行实验。两模型在NVIDIA GPU服务器上使用Pile数据集训练，保留了所有Pythia的配置，并训练了143,000步。

4.2.1 收敛性与正态模型性能

在预训练中加入sink token不会对模型的收敛性和后续在多个NLP基准上的性能产生负面影响。 如图6所示，加入sink token的模型与原始模型在收敛动态上表现相似。我们评估了两种模型在七个不同的NLP基准上的性能，包括ARC、HellaSwag、LAMBADA、OpenbookQA、PIQA和Winogrande。

如表4所示，预训练中加入sink token的模型与使用传统方法训练的模型表现相似。

4.2.2 流性能

如表3所示，使用传统方法训练的模型与增加了a sink token的模型在流式困惑度上存在差异。值得注意的是，传统模型需要添加多个tokens作为attention sinks来保持稳定的流式困惑度。相比之下，使用a sink token训练的模型只需使用该sink token就能实现令人满意的流式性能。

4.2.3 注意力可视化

图7展示了带有和不带有a sink token的预训练模型的注意力图。不带sink token的模型，类似于Llama-2-7B（图2），显示了早期层次的局部注意力和对初始tokens的深层关注。与此相反，带有a sink token的模型始终在各层和头部集中注意力于sink，这表明了有效的注意力卸载机制。这种对sink的强烈关注，减少了对其他初始tokens的关注，解释了sink token在增强模型的流式性能中的有效性。

4.3 使用指令微调模型进行流式问答的结果

为展示StreamingLLM的实际适用性，我们模拟了多轮问答场景，使用了ARC数据集。与传统方法相比，StreamingLLM表现出色，能有效处理流式格式，与单次测试精度一致。我们还引入了一个新数据集StreamEval，模拟现实中的查询场景，其中问题通常与最近的信息相关。结果显示，StreamingLLM即使在长文本上也能保持合理的精度，而其他方法则在特定长度上失败。此外，结合其他上下文扩展技术，StreamingLLM表现得更好，可以捕获更多的局部信息。

4.4 消融研究

论文探讨了StreamingLLM中初始tokens数量和缓存大小的影响。结果显示，仅引入1或2个初始tokens是不够的，但使用4个初始tokens作为关注点是足够的。另外，增加缓存大小并不总是降低语言建模的困惑度，暗示模型可能未充分利用其接收到的整个上下文，需要进一步研究。

4.5 效率结果

论文使用Huggingface Transformers库在NVIDIA A6000 GPU上测试StreamingLLM的解码延迟和内存使用情况，与滑动窗口重新计算基线进行比较。结果显示，随着缓存大小的增加，StreamingLLM的解码速度呈线性增长，而滑动窗口的解码延迟呈二次增长。因此，StreamingLLM实现了高达22.2倍的速度提升，同时其内存占用与重新计算基线保持一致。

五、Conclusion

为了在流应用中部署LLMs，我们面临着效率和长文本性能下降的挑战。虽然窗口注意力提供了部分解决方案，但当排除初始tokens时，其性能急剧下降。识别到这些tokens作为“注意力汇聚”的角色，我们引入了StreamingLLM——一个简单且高效的框架，能够处理长达400万tokens的文本，无需微调。我们进一步展示，通过预训练模型，使用专用的sink token可以提高流性能。StreamingLLM首次将LLM的预训练窗口大小与其实际文本生成长度解耦，为LLMs的流式部署铺平了道路。