大型语言模型（LLM）已经成为自然语言处理领域不可或缺的一部分，但它们在处理长文本时却面临着巨大的挑战。传统的 LLM 在解码过程中需要缓存所有先前解码的 token 的键值对（KV），这会导致内存占用量随着对话长度的增加而线性增长，并且模型的性能也会受到限制。
StreamingLLM 应运而生，它是一种高效的框架，可以使 LLM 在不进行任何微调的情况下，将有限长度的注意力窗口推广到无限序列长度，从而实现高效的流式解码。
窗口注意力的瓶颈：为什么需要 StreamingLLM？
为了解决传统 LLM 解码过程中的内存占用和性能下降问题，研究人员提出了窗口注意力机制，它只缓存最近的 KV，从而减少了内存占用。然而，当文本长度超过缓存大小时，窗口注意力机制就会失效。
StreamingLLM 的核心思想是利用注意力汇聚（Attention Sink）现象。研究人员发现，即使一些初始 token 在语义上并不重要，但模型仍然会对它们保持很强的注意力。这种现象被称为注意力汇聚。
Attention Sink 的发现：解开窗口注意力的秘密
通过分析 LLM 在处理长文本时的注意力机制，研究人员发现了一个有趣的现象：模型会将大量的注意力集中在初始 token 上，即使这些 token 在语义上并不重要。他们将这种现象称为“注意力汇聚”。
为什么会出现注意力汇聚？ 这是因为 Softmax 函数的特性，它要求所有上下文 token 的注意力分数之和为 1。即使当前查询在许多之前的 token 中没有强匹配，模型仍然需要将这些不必要的注意力值分配到某个地方，以使它们加起来为 1。
初始 token 成为注意力汇聚的原因也很直观： 由于自回归语言模型的特性，初始 token 对所有后续 token 都可见，而后续 token 只能对有限的后续 token 可见。因此，初始 token 更容易被训练成注意力汇聚，从而捕获不必要的注意力。
StreamingLLM 的工作原理：高效利用注意力汇聚
StreamingLLM 框架基于以下几个关键原则：
- 保留注意力汇聚： 保留初始 token 的 KV，以弥补窗口注意力机制的不足。
- 滑动窗口： 使用滑动窗口机制，只缓存最近的 token 和注意力汇聚，从而减少内存占用。
- 预训练优化： 在预训练过程中添加一个占位符 token 作为专门的注意力汇聚，进一步提高流式部署的性能。
StreamingLLM 的优势：高效、稳定、无需微调
- 高效的流式解码： StreamingLLM 可以实现高效的流式解码，即使在处理非常长的文本时，也能保持良好的性能。
- 无限序列长度： StreamingLLM 可以处理无限长度的文本，突破了传统 LLM 的限制。
- 无需微调： StreamingLLM 不需要对 LLM 进行任何微调，就可以实现高效的流式解码。
实验结果：实证验证 StreamingLLM 的有效性
研究人员对多个流行的 LLM 家族（包括 Llama-2、MPT、Falcon 和 Pythia）进行了实验，结果表明：
- StreamingLLM 在处理长文本时，其性能与滑动窗口重新计算基线相当，而传统的窗口注意力机制则会失效。
- StreamingLLM 可以处理超过 400 万个 token 的文本，并且在处理超长文本时仍然保持稳定。
- 在预训练过程中添加一个专门的注意力汇聚 token 可以进一步提高 StreamingLLM 的性能。
StreamingLLM 的应用场景：无限对话、实时翻译、文本摘要
StreamingLLM 在许多场景中都有广泛的应用，例如：
- 多轮对话： StreamingLLM 可以用于构建高效的聊天机器人，实现流畅的对话交互。
- 文本摘要： StreamingLLM 可以用于对长文本进行摘要，并保留关键信息。
- 实时翻译： StreamingLLM 可以用于实时翻译，提供流畅的翻译体验。
未来展望：更强大、更智能的流式解码
StreamingLLM 的出现为 LLM 的流式应用开辟了新的可能性。未来，研究人员将继续探索如何进一步优化 StreamingLLM 框架，使其在更多场景中发挥更大的作用。
https://github.com/mit-han-lab/streaming-llm