LLM推理技术之StreamingLLM：如何拥有无限长生成能力
来源：AI大模型面试宝典 作者：宁静致远 发布时间：2023年10月22日
一、前言
1.1 大型语言模型（LLM）存在什么问题？
当前，大型语言模型（LLM）在推理时只能记住有限的上下文。例如，LLaMA-2只能处理4K的上下文，这不仅导致其无法记住超过最近4K上文的内容，而且在生成文本达到4K时就会停止。理想的AI对话助手应该不受输出长度的限制，并且需要记住历史的对话内容。
📝通俗解释：就像人的工作记忆有限一样，LLM也有一个"记忆容量"。LLaMA-2的4K上下文意味着它只能同时"看到"约4000个词的对话内容，超过这个长度它就会"忘记"之前的内容。这就像一个只能记住最近对话的聊天机器人，无法记住很久以前的聊天内容。
1.2 StreamingLLM 背景介绍
MIT、Meta AI、CMU的研究人员最近提出了一种StreamingLLM，声称可以使得经过有限序列长度训练的大型语言模型能够在无需任何微调的情况下，推广到无限序列长度的输入和输出。
⚠️这里值得强调的是：这个方法并没有增加LLM对上文内容的记忆能力，只是让它能够实现无限长的输入和输出。
一个显而易见的好处是，在对话机器人生成一个很长的回答时，用户不需要再点击“继续”按钮了。
📝通俗解释：想象你在看一部很长的电视剧，但电视只能连续播放1小时。之前的模型就像这台电视，看到1小时就自动停止，你需要说"继续"才能看下一小时。StreamingLLM就像一台不会停止的电视，可以一直播放下去，但它的"记忆力"其实没有变好，只是学会了如何流畅地处理很长的内容。
图1. 对话机器人长文本生成截断示意图 图片描述：图片展示了一段古文《出师表》的生成内容，文本显示到"愿陛下托臣以讨贼兴复之效，不效则治臣之罪，以"处截断，右下角有一个蓝色的"继续"按钮。
1.3 StreamingLLM 核心问题？
核心问题：能否在不牺牲效率和性能的情况下，部署一个能处理无限输入的LLM？
实现流式LLM应用部署的效果，即可以不受长度限制不停地输出。这个需要说明：无限长输入和无限长上下文是不同的，前者不需要对所有输入有记忆。
📝通俗解释：这里的"无限长输入"就像一条流动的河流，模型不需要记住整条河流的水，只需要处理当前流过来的水。而"无限长上下文"则要求模型记住整条河流的所有历史。这两个概念完全不同，StreamingLLM实现的是前者。
1.4 StreamingLLM 存在哪些挑战？
模型性能瓶颈：在解码阶段，由于KV Cache存在导致内存使用或延迟增加，内存上限和推理服务SLA存在限制，导致KV Cache不能无限大，这是性能瓶颈。
模型能力瓶颈：现有模型的外推（extrapolation）能力有限，也就是说当序列长度超过预训练时设定的注意力窗口大小时，它们的表现会下降，这是模型能力的瓶颈。如下图所示，Dense Attention具有$O(T^2)$的时间和内存复杂度。当文本长度超过预训练文本长度时，其运行性能会下降。
📝通俗解释：
- 性能瓶颈就像电脑内存不够——KV Cache是存放"记忆"的地方，但内存有限，不能无限存放
- 能力瓶颈就像让一个只学过小学数学的人做高等数学——模型在训练时只见过最多4K长度的内容，超过这个长度它就不会"思考"了
1.5 目前主流的增加输入文本长度的方法有哪些？
方法一：长度外推（Length Extrapolation）
该方法让训练在较短文本上的LLM能够在推理时处理较长的文本。比如，大家经常听到的编码方法RoPE、ALiBi等都归于此类。然而，目前尚未有方法实现无限长度的外推，还无法满足流式应用的需求。
📝通俗解释：长度外推就像一个学会了骑自行车的小孩，虽然没骑过三轮车，但凭借骑自行车的经验也能很快学会。模型通过学习短文本的处理方式，尝试推广到长文本。但目前这种方法还有局限性，无法做到真正的"无限长"。
方法二：上下文窗口扩展（Context Window Extension）
该方法实打实地去扩大LLM的上下文窗口长度，也就是序列长度。由于Attention的计算量和内存需求都随着序列长度增加而成平方增长，所以增加序列长度很难。具体的实现方法包括：
- a. 训练时用FlashAttention等工程优化，以打破内存墙的限制；
- b. 一些approximate attention方法，比如Longformer这种Window Attention方法。如图1所示，Window Attention缓存最近的L个token的KV。虽然在推理过程的效率高，但一旦开头的token的KV被驱逐出Cache，模型推理的表现就会急剧下降（PPL越高模型表现越差）。在图2中，橙色PPL曲线在token数目超过KVCache Size后出现跃升。
- c. 降低内存需求的优化：让Window Attention重新计算从每个新令牌的L个最近令牌中重建KV Cache。虽然它在长文本上表现良好，但由于上下文重新计算中的二次注意力导致的$O(T*L^2)$复杂性，使其相当慢。
📝通俗解释：
- FlashAttention就像优化了自行车的传动系统，让骑车更省力
- Window Attention就像只记住最近发生的事情，忘记 older things
- 重新计算方法就像每次需要回忆时，把最近的事情重新想一遍——虽然记得牢，但太慢了
图2. 四种注意力机制的对比 图片描述：图1展示了四种注意力机制的对比。(a Dense Attention: 计算复杂度$O(T^2)$，PPL: 5641X，效率低且长文本表现差。(b) Window Attention: 计算复杂度$O(TL)$，PPL: 5158X，当初始token被驱逐时失效。(c) Sliding Window w/ Re-computation: 计算复杂度$O(TL^2)$，PPL: 5.43，需要为每个传入token重新计算cache。(d) StreamingLLM (ours): 计算复杂度$O(TL)$，PPL: 5.40，具有Attention Sink，能在长文本上实现高效的语言建模。)
图2. 四种注意力机制的对比
图3. 不同模型在不同输入长度下的PPL变化 图片描述：图2展示了四个模型（Llama-2-7B, Pythia-12B, Falcon-7B, MPT-7B）在不同输入长度下的PPL变化。蓝色线代表Dense Attention，橙色线代表Window Attention，绿色线代表Sliding Window w/ Re-computation，红色线代表StreamingLLM。可以看出Window Attention在超过一定长度后PPL急剧上升，而StreamingLLM保持平稳。
图3. 将输入的文本长度增加到20K进行推理时的困惑度（PPL）
通常，使用这些技术后，大型语言模型（LLM）的推理输入长度会受到一定的限制。然而，这篇论文通过使用approximate attention的方法，放松了对全部输入记忆的限制，仍然只记住最近的上下文，但实现了处理无限输入并获得无限输出的效果。可以说没有和长度外推法硬钢，而是另辟蹊径。
二、StreamingLLM 的思路是什么？
对于Window Attention在超长文本输入时失败的原因，需要深入研究。如图2的橙色曲线所示，即使窗口大小只比KVCache Size大1，也就是说，注意力计算只减少了第一个token，模型推理的PPL值却会急剧上升。
这个现象确实令人感到诧异，因为直觉告诉我们，随着窗口大小的增大，模型推理的表现应该逐渐变差。然而，仅仅少输入一个token，模型的性能就一触即溃，这似乎暗示着开头的第一个token可能具有关键的作用。事出反常必有妖！
📝通俗解释：这就像一个很奇怪的现象——让一个记忆力不好的人记住最近发生的10件事，他做得还行；但如果让他忘记最早的那件事，只记住9件事，他的表现就会急剧下降。这说明"第一件事"可能具有某种特殊的重要性。
于是乎，作者们把attention每一层每一个Head经过softmax输出后的logits值翻出来观察。这一看，不得了，果然发现了问题。如图3所示，作者们发现：
- 第一和第二layer（0和1 layer）的注意力图展示了"local"模式：离当前处理token最近的token收到了更多的attention，即attention矩阵对角线位置值相对更大。
- 除了最网络前面的两层外，模型在所有layer和head都重点对于initial token（开头的几个tokens）给予更多的attention值。
📝通俗解释：研究者发现模型有两个"癖好"：一是"喜新厌旧"，更关注最近的内容（对角线模式）；二是"不忘初心"，始终关注开头的那几个token。这两个发现成为了StreamingLLM的关键洞见。
图4. Llama-2-7B模型的注意力热力图 图片描述：图3展示了Llama-2-7B模型不同层和头的注意力热力图。Layer 0 Head 0, Layer 1 Head 0, Layer 2 Head 0 等显示出明显的对角线模式（local attention）以及左侧边缘的高亮区域（initial tokens attention）。
图4. 使用Llama-2-7B时，对256个句子的平均注意力logits进行的可视化，每个句子的长度为16
2.1 Attention Sink 概念的提出
给予如上观察，作者提出了**"attention sink"概念来解释Window Attention失败的原因。输入给LLM推理开头的几个initial tokens是非常特殊的，仿佛水池（sink）中的排水口一样，吞噬了大量的attention。而且initial tokens与被预测token的距离如何，语义信息如何都不重要，重要的只是它的绝对位置。也就是说前几个位置上的token不管是啥，对维持LLMs推理的稳定性都很关键**。
📝通俗解释：可以把这些initial tokens想象成厨房的排水口——无论你往水槽里倒多少水，最后都会从排水口流走。模型就像一个有"排水口依赖症"的人，无论处理什么内容，都习惯性地"看向"开头的那几个位置，这似乎是一个根深蒂固的"本能"。
2.2 Attention Sink 的原因分析
那么Attention sink是什么原因造成的呢？作者尝试给出一些解释，原来这和前一段时间沸沸扬扬的一段公案有关。
不知道读者们是否收到过如下公众号文章的推送，其标题很耸人听闻：
文章：Attention机制竟有bug，Softmax是罪魁祸首，影响所有Transformer 地址：https://zhuanlan.zhihu.com/p/645844743
这件事简单讲是：高通AI Research的人研究LLM量化方法时发现Attention Head激活张量里有一些值异常突出（常被称为outliner），追查发现是Softmax引发的。这个问题引起了网红程序员Evan Miller的注意，他研究发现softmax函数存在Bug，并发表了一篇博客《Attention Is Off By One》。
个人理解：在Attention机制中，Softmax的输出代表了key/query的匹配程度的概率。因此，如果softmax在某个位置的值非常大，那么在反向传播时，这个位置的权重就会被大幅度地更新。然而，有时候attention机制并不能确定哪个位置更值得关注，但由于Softmax需要所有位置的值的总和为1，因此必须"表态"给某些位置较大的权重，这就可能导致错误的权重更新，而这个错误在后续的过程中很难被纠正。
📝通俗解释：Softmax就像一个"必须表态"的投票系统——它要求所有选项的得票率加起来必须等于100%。即使某个位置没什么值得关注的，Softmax也不得不"强行分配"一些注意力过去。这就导致模型会把一些"无处安放"的注意力错误地投给某些位置，而initial tokens就成了这些"错误注意力"的"垃圾回收站"。
如下是Miller的原话：
The problem with using softmax is that it forces each attention head to make an annotation, even if it has no information to add to the output vector
于是乎，他改进了一下Softmax，也就是把softmax的分母加了个1仅此而已，这样所有位置值可以加和不为1，这样Attention就有了可以不对任何位置"表态"的权利。
$$ (\text{softmax}_1(x))_i = \frac{\exp(x_i)}{1 + \sum_j \exp(x_j)} $$
StreamingLLM的作者采用了类似的观点解释attention sink现象。SoftMax函数的性质使得所有经过attention结构的激活张量不能全部为零，虽然有些位置其实不需要给啥注意力。因此，模型倾向于将不必要的注意力值转嫁给特定的token，作者发现就是initial tokens。
📝通俗解释：这就像一个必须选班长的投票——即使全班同学都觉得谁当班长都行，投票结果也必须有人当选。模型也是如此，即使某些位置没什么值得关注的，Softmax也必须"选出"一些位置来承载注意力，而开头的那几个位置就成了"冤大头"。
2.3 StreamingLLM 的解决方案
有了这个洞见，作者设计Window Attention的改进版。思路也是很直接，在当前滑动窗口方法基础上，重新引入了一些initial tokens的KV在注意力计算中使用。
StreamingLLM中的KV缓存可以概念上分为两部分，如图5所示：
- （1）Attention Sink：4个initial tokens，稳定了注意力计算
- （2）Rolling KV缓存：保留了最近的token，这个窗口值是固定的，图中为3
📝通俗解释：这个设计就像一个"双保险"——既保留了最近的记忆（滑动窗口），又永远保留开头的那几个"元老"token。这样模型既能正常处理新内容，又不会因为"忘记"开头而崩溃。
图5. StreamingLLM的KV Cache机制 图片描述：图4展示了StreamingLLM的KV cache机制，分为三行，分别对应生成Token 7、Token 8和Token 9的过程。Attention Sinks (黄色块 ：始终保留最开始的4个token (0, 1, 2, 3)。Evicted Tokens (灰色虚线块)：随着生成进行，中间部分的token（如4, 5）被移出缓存。Rolling KV Cache (蓝色/红色块)：保留最近的几个token（窗口大小为3），随着新token生成而滑动更新。)
图5. The KV cache of StreamingLLM
2.4 位置编码的适配
还需要有些小改动来给attention注入位置信息，StreamingLLM就可以无缝地融入任何使用相对位置编码的自回归语言模型，如RoPE和ALiBi。方法很直接，大家去看原文。
📝通俗解释：为了让这个方案适用于不同的位置编码方式（如RoPE、ALiBi），需要进行一些适配工作，确保模型能够正确理解token的位置信息。这就像给不同的汽车配上合适的钥匙，让它们都能正常启动。
到目前为止，StreamingLLM不需要做训练，把initial tokens数目设置为4就可以获得不错的长输入下的推理表现了。
2.5 可训练版本（进阶）
如果我们解除不能训练模型的限制，可以通过Pre-training LLMs with attention sinks获得更好的表现。作者提出两种方法：
指定一个全局可训练的attention sink token，称之为"Sink Token"，它将作为不必要的注意力的存储库，从而把initial tokens作为attention sink的作用转移到sink token上。
用类似Miller提出的SoftMax-off-by-One的变体替换传统的SoftMax函数，作者称之为Zero Sink方法。
📝通俗解释：
- 第一种方法就像专门设立一个"垃圾桶"岗位，让专人负责接收那些"无处安放"的注意力
- 第二种方法就像修改投票规则，允许投"弃权票"，让注意力可以不被强制分配到任何位置
三、总结
StreamingLLM 的最大好处是实现无限长的输出，但不能记忆超长的输入。
📝通俗解释：StreamingLLM就像一个"流式播放器"——可以连续不断地播放视频（输出），但它并没有因此变成"过目不忘"的天才。它的"记忆力"（能记住的上下文长度）并没有变长，只是学会了如何优雅地处理很长的内容流，不会因为内容太长而"崩溃"。这对于需要长时间对话或生成长文本的场景非常有用。