StreamingLLM论文精读为什么只保留开头几个Token就能让模型持续生成数百万Token在普通大语言模型中KV Cache会随着对话或文本长度不断增长。最直接的解决办法是只保留最近一段Token也就是滑动窗口。但论文发现了一个令人意外的现象一旦把最开头的几个Token从KV Cache中删除模型困惑度会突然爆炸只要把这些开头Token重新放回缓存即使其余大量历史内容已经删除模型又能恢复正常生成。作者将这些持续吸收大量注意力、但语义上未必重要的初始Token称为Attention Sinks注意力汇聚Token。基于这一现象论文提出StreamingLLM固定保留少量初始Attention Sink再配合一个滚动的最近Token窗口使模型能够在固定显存和固定单步计算量下持续生成。一、论文基本信息项目内容论文题目Efficient Streaming Language Models with Attention Sinks方法名称StreamingLLM作者Guangxuan Xiao、Yuandong Tian、Beidi Chen、Song Han、Mike Lewis会议ICLR 2024首次公开2023年9月主要贡献Attention Sink现象、固定大小滚动KV Cache、专用Sink Token主要模型LLaMA-2、MPT、Falcon、Pythia是否需要微调已训练模型不需要是否扩展模型上下文窗口否官方实现MIT HAN Lab发布的StreamingLLM代码论文正式发表于ICLR 2024并提供了公开实现。(ICLR 会议论文集)二、StreamingLLM要解决什么问题论文关注的是持续运行的语言模型例如长时间运行的聊天助手不断接收日志的分析系统持续输入字幕或语音转录长时间代码补全连续新闻或社交媒体流永不重置的多轮交互服务。这些应用并不一定要求模型记住从开始到现在的全部内容但要求模型能持续接收新Token不因输入总长度增加而耗尽显存不需要周期性清空KV Cache始终能够基于最近上下文稳定生成。论文指出传统大语言模型在这种场景中面临两个问题KV Cache随输入长度线性增长输入长度超过预训练窗口后模型本身也可能出现位置外推失败。(arXiv)三、三种基础推理方式为什么都不理想3.1 完整注意力完整注意力保存此前所有Token的Key和Value。它的优点是历史内容都可以被访问但代价是KV Cache持续增长单步注意力需要访问越来越多历史Token长时间运行最终会显存不足输入超过训练窗口后位置编码也可能失效。因此完整注意力不能支持真正持续运行的流式服务。3.2 普通滑动窗口普通窗口注意力只保存最近固定数量的Token。例如缓存大小为1024时始终保留最近1024个Token每加入一个新Token就删除最早的一个KV显存保持固定单步注意力计算量保持固定。从系统角度看这似乎是最理想的方案。但论文发现当序列长度刚刚超过窗口大小、最初Token第一次被删除时模型困惑度就会突然爆炸。也就是说问题并不是逐渐遗忘旧内容而是整个注意力分布突然失去稳定性。(arXiv)3.3 滑动窗口加重新计算另一种方法是保留最近一段原始文本但每生成一个新Token都重新将窗口内全部文本输入模型重新计算其KV。这种方式通常能够保持较好的语言建模质量因为窗口内Token每次都会获得连续、正确的位置。但它存在严重的重复计算窗口越大每一步重新计算的Token越多解码成本随窗口大小近似二次增加无法充分复用已有KV。论文把它作为高质量参考基线但认为其速度不适合实际流式应用。(arXiv)四、什么是Attention Sink作者分析LLaMA-2等模型的注意力图后发现最底部少数层通常更关注最近Token从较深层开始大量注意力头持续关注序列开头的几个Token即使这些Token与当前预测没有明显语义关系它们仍然获得很高注意力。作者把这种现象称为Attention Sink某些Token充当注意力分数的“汇聚池”吸收模型暂时没有必要分配给其他内容的注意力。论文观察到通常保留最初4个Token就足以恢复稳定而只保留1个或2个初始Token在部分模型上仍然不够。(arXiv)五、为什么初始Token会成为注意力汇聚点5.1 论文给出的Softmax解释注意力经过Softmax归一化后所有历史Token的注意力权重必须加起来等于1。有些情况下当前Query可能只需要自身或少量局部信息但注意力仍必须将全部概率质量分配出去。作者认为模型会把部分“不必要的注意力”分配给某些固定Token这些Token便逐渐成为Attention Sink。(arXiv)可以将其理解为模型并不一定需要从这些Token中读取重要语义而是需要一个稳定的位置来容纳多余注意力权重。5.2 为什么偏偏是开头Token在自回归注意力中最开始的Token对后面几乎所有Token都可见。相比之下第一个Token能被整个序列后续位置访问中间Token只能被它后面的部分位置访问最后几个Token仅被极少数后续位置访问。因此在预训练过程中开头Token最容易被大量Query反复利用最终形成稳定的注意力汇聚位置。(arXiv)5.3 语义真的不重要吗作者进行了一个关键实验。他们把原始文本最初4个Token替换成4个换行符但仍将这些位置的KV保留在缓存中。结果显示保留原始前4个Token时LLaMA-2-13B困惑度为5.40替换成4个换行Token后困惑度为5.60完全删除初始Token、只保留最近1024个Token时困惑度达到5158.07。这说明恢复性能的关键更接近于初始位置及其训练形成的注意力模式而不是最初几个Token具体表达了什么内容。(arXiv)不过“语义不重要”不能理解为这些Token在所有场景下都没有内容价值。它只是说明在这个语言建模实验中Attention Sink的主要作用不是携带原始文本开头的语义。六、StreamingLLM的核心设计StreamingLLM将固定大小的KV Cache分成两部分第一部分Attention Sink永久保留最开始的少量Token默认是4个。这些Token用于稳定Softmax注意力分布。第二部分Rolling Cache保存最近的一段连续Token。这些Token提供最近语义局部语法当前对话状态最近实体当前推理步骤。因此其缓存结构可以概括为最初4个Token 最近若干Token例如缓存总预算为1024可以使用4个初始Token 1020个最近Token中间那些既不属于最初位置、也不属于最近窗口的Token会被删除。(arXiv)七、StreamingLLM完整执行流程假设模型总缓存预算为8个Token其中前2个位置作为Attention Sink后6个位置作为滚动窗口。随着文本不断到来缓存可能经历初始阶段[0, 1, 2, 3, 4, 5, 6, 7]继续生成后[0, 1, 4, 5, 6, 7, 8, 9]再继续生成[0, 1, 6, 7, 8, 9, 10, 11]其中Token 0和1始终保留最近6个Token不断向前滚动中间旧Token被淘汰。因此无论模型总共处理了1万、100万还是400万Token缓存大小都不会继续增长。八、为什么位置编号必须重新排列只删除KV还不够位置编码也必须同步处理。假设当前缓存保留的原始Token编号是[0, 1, 2, 3, 6, 7, 8]现在准备处理Token 9。如果仍然按照原始文本位置计算位置编号将是[0, 1, 2, 3, 6, 7, 8, 9]其中存在一个从3跳到6的断层。StreamingLLM不使用这种原始全局位置而是按照它们在当前缓存中的相对位置重新编号[0, 1, 2, 3, 4, 5, 6, 7]这样最近窗口中的Token始终位于模型训练时熟悉的位置范围内而不会不断进入数十万或数百万的位置编号。(arXiv)8.1 RoPE如何处理对于使用RoPE的模型论文缓存旋转位置编码处理之前的Key。每次滚动窗口更新后再按照新的缓存位置对Key施加RoPE变换。这样可以避免将已经带有旧全局位置的Key直接用于新的连续窗口。(arXiv)8.2 ALiBi如何处理对于使用ALiBi的MPT模型不需要重新旋转Key。算法只需要根据当前缓存中的连续相对位置重新计算线性距离偏置而不是保留原始文本中跳跃的位置距离。(arXiv)九、为什么只保留最近Token还不够而加入4个Sink就有效普通滑动窗口删除初始Token后模型不仅失去了几个旧Token还改变了整个注意力归一化环境。模型原本已经习惯一部分注意力可以被分配给初始Sink位置。删除这些位置后原本流向Sink的权重必须重新分配给剩余Token导致局部Token获得异常高的注意力注意力分布偏离训练状态隐藏表示逐层累积偏差模型最终输出失去稳定性。StreamingLLM重新加入初始Token相当于恢复了训练时熟悉的注意力“锚点”。论文在400K Token文本上测试不同初始Token数量模型纯窗口保留1个保留2个保留4个保留8个Falcon-7B17.9012.1212.1212.1212.12MPT-7B460.2914.9915.0014.9914.98Pythia-12B21.6211.9512.0912.0912.02LLaMA-2-7B3359.9511.8810.519.599.54结果说明不同模型需要的Sink数量略有差异1个Token有时已经有效但不够稳定4个初始Token通常足够继续增加到8个收益很小。(arXiv)十、StreamingLLM真的扩展了上下文窗口吗没有。这是理解这篇论文最重要的一点。StreamingLLM可以持续处理无限增长的输入流但模型在任意时刻实际看到的只有最初几个Sink Token最近固定数量的Token。中间已经被删除的内容无法再次访问。官方代码说明也明确指出StreamingLLM没有扩大模型原始上下文窗口模型只识别当前最近窗口输入一本完整书时它可能只能依据结尾部分生成摘要它适合持续对话和流式生成不适合要求记住全部历史细节的长文理解。(GitHub)因此应该区分两个概念。能力StreamingLLM是否提供连续运行数百万Token而不重置缓存是固定KV Cache显存是保持基于最近上下文的流畅生成是访问数百万Token之前的任意信息否将4K模型变成真正4M上下文模型否总结一本完整超长书籍通常不能保证所谓“无限长度”指的是输入流可以无限持续而不是可访问记忆范围无限扩大。十一、为什么论文能测试到400万Token论文在拼接后的PG-19长文本上让多个模型持续进行语言建模。实验覆盖LLaMA-2 7B、13B、70BFalcon 7B、40BPythia 2.8B、6.9B、12BMPT 7B、30B。在超过400万Token的连续文本中StreamingLLM的困惑度保持相对稳定。(arXiv)但这个结果说明的是模型不会因为累计处理长度超过训练窗口而突然数值崩溃。它并不证明模型仍然记得数百万Token之前发生的内容。事实上论文的StreamEval附录显示当问题所需答案超出当前缓存范围后准确率会快速下降并最终降到零。(arXiv)十二、流式问答实验作者把ARC-Easy和ARC-Challenge中的问题连续拼接成一个长数据流模拟不断进行的多轮问答。模型使用1024个KV缓存位置。结果如下模型方法ARC-EasyARC-ChallengeLLaMA-2-7B-Chat独立逐题71.2553.16LLaMA-2-7B-Chat普通窗口3.581.39LLaMA-2-7B-ChatStreamingLLM71.3455.03LLaMA-2-13B-Chat独立逐题78.1663.31LLaMA-2-13B-Chat普通窗口0.250.34LLaMA-2-13B-ChatStreamingLLM80.8965.61LLaMA-2-70B-Chat独立逐题91.2978.50LLaMA-2-70B-Chat普通窗口0.120.32LLaMA-2-70B-ChatStreamingLLM91.3780.20完整Dense Cache在连续数据流中因显存增长而OOM普通窗口在初始Token被删除后接近随机输出StreamingLLM则接近逐题独立推理的结果。(arXiv)这一实验非常适合StreamingLLM因为每个新问题的答案主要依赖最近的题目而不是很久之前的问答。十三、StreamEval进一步说明了什么作者设计了StreamEval模拟持续输入信息并周期性提问。默认设置中每输入10行新信息提出一次问题答案位于之前约20行的位置问题依赖相对较近的历史内容。StreamingLLM在输入总长度接近120K Token时仍保持较合理的准确率而Dense和普通窗口分别受训练长度和缓存淘汰影响。(arXiv)但附录的距离实验揭示了它的边界当答案与问题之间的距离超过当前Rolling Cache容量后准确率会逐渐下降最终变为零。例如在缓存为42044时答案距离达到约2300 Token后准确率已经降为0增加缓存可以推迟崩溃位置但不能消除这一限制。(arXiv)十四、LongBench结果揭示的局限论文还在LongBench上测试长文问答和摘要。使用4个Sink加3496个最近Token时StreamingLLM通常低于“保留开头1750个和结尾1750个Token”的普通截断基线。例如方法NarrativeQAQasperHotpotQA2WikiMQA开头1750结尾175018.719.225.432.84个Sink3496最近Token11.616.921.628.2原因是长文问答中的重要指令、背景或事实可能位于文档开头而4个Attention Sink只负责稳定注意力并不能保存开头的实际语义内容。当StreamingLLM也保留开头1750个Token时表现才恢复到与截断基线接近。(arXiv)这进一步证明Attention Sink提供的是数值稳定性而不是长期语义记忆。十五、专用Sink Token15.1 为什么现有模型需要多个初始Token现有大模型在预训练时通常没有一个在所有训练样本中都稳定出现在第一个位置的特殊Token。因此模型可能将注意力汇聚功能分散到开头若干Token上而不是集中到一个位置。这也是为什么现有LLaMA、MPT和Pythia模型通常需要保留4个左右的初始Token。(arXiv)15.2 作者提出专门训练一个Sink Token作者提出在每个预训练样本最前面增加一个可学习的占位Token。这个Token不需要表达实际语义它的作用是专门吸收模型没有必要分配给其他内容的注意力权重。这样部署时只需要永久保留这一个Sink Token而不必保留多个普通文本Token。(arXiv)15.3 预训练实验作者从头训练了160M参数模型对比普通Softmax模型使用Zero Sink的模型加入可学习Sink Token的模型。普通模型需要保留多个初始Token才能稳定而经过专用Sink Token预训练的模型只保留一个Sink即可恢复正常困惑度。同时加入Sink Token并没有损害普通零样本任务性能。七项任务中Sink模型的结果与普通模型相近部分指标还略高。(arXiv)不过这项结论只在160M小模型上通过从头训练验证并没有在7B或70B模型上重新预训练验证。十六、实际效率结果StreamingLLM与滑动窗口重新计算方法具有相似的固定KV显存但运行过程不同重新计算方法每一步都重新处理整个窗口StreamingLLM直接复用已有KV只计算新Token。论文在单张NVIDIA A6000上测试LLaMA-2-7B和13B。随着缓存窗口变大重新计算方法的每Token延迟近似二次增长StreamingLLM的延迟增长更加缓慢最大报告加速达到22.2倍两种方法的KV内存占用处于相近水平。(arXiv)16.1 怎样理解22.2倍加速22.2倍是相对于每生成一个Token都重新计算最近整个窗口KV的滑动窗口基线。它不是相对于高效的普通KV Cache解码。普通KV Cache只计算新Token速度本来就很快但缓存会无限增长StreamingLLM的价值在于同时获得类似普通KV复用的单步效率与固定窗口类似的恒定内存比纯滑动窗口稳定得多的生成质量。因此22.2倍不能理解为StreamingLLM让标准大语言模型推理普遍加速22倍。更准确地说在必须保持固定缓存且不能接受普通窗口崩溃时StreamingLLM比窗口重新计算方案最高快22.2倍。十七、缓存越大效果一定越好吗不一定。论文测试不同StreamingLLM缓存大小后发现增加最近Token数量并不总能降低困惑度。例如MPT-7B缓存结构困惑度425214.12450814.254102014.334204414.99LLaMA-2-7B则从较小窗口逐渐改善到42044附近最好但继续扩大到44092后又略有恶化。(arXiv)这说明模型未必能够充分利用全部上下文更长缓存不一定带来更好预测位置外推和训练分布仍会影响结果StreamingLLM主要解决稳定性不解决长上下文利用效率。十八、与H₂O的区别上一篇H₂O和StreamingLLM都压缩KV Cache但保留历史Token的规则完全不同。对比维度StreamingLLMH₂O长期保留Token固定的最初几个Token动态累计注意力最高的Token最近Token保留保留是否依赖输入内容Sink基本固定是是否统计注意力重要性否是是否动态改变历史保留集合否是核心目标保持注意力数值稳定保留具有长期内容贡献的Token计算开销非常低需要累计注意力与动态淘汰长期语义保留能力很有限相对更强StreamingLLM认为不需要寻找语义上重要的旧Token只要保留注意力汇聚点再保留最近窗口即可稳定生成。H₂O则认为除最近Token外还应动态保留历史上累计注意力较高的重要Token。H₂O的历史集合随内容变化而StreamingLLM的Sink位置固定。(GitHub)18.1 两者谁更适合什么场景StreamingLLM更适合持续聊天实时日志或字幕流只依赖最近上下文希望实现尽可能简单不愿引入动态排序和注意力统计。H₂O更适合历史中存在可能长期有用的信息希望保留远距离关键词或实体能够接受动态缓存管理希望根据输入内容选择历史Token。两种方法也可以组合保留固定Attention Sink、最近窗口再从中间历史中选择少量Heavy Hitters。十九、它和真正的长上下文扩展有什么区别真正的上下文扩展方法通常希望模型接收超过原训练窗口的长输入访问窗口内所有Token利用跨越数万Token的远距离信息完成长文问答、摘要和检索。StreamingLLM并不完成这些目标。它只是让模型在处理无限输入流时始终使用有限的局部上下文。官方实现将两类技术描述为正交关系长上下文扩展扩大模型一次能访问的窗口StreamingLLM让这个有限窗口能够在无限输入流上持续滚动。因此StreamingLLM可以应用于32K模型此时它能够保存更大的最近窗口但仍然不能记住已被滚动窗口淘汰的全部历史。(GitHub)二十、这篇论文真正证明了什么论文并没有证明语言模型只需要最初4个Token和最近Token就能理解任意长度文本。它真正证明了两个现象。第一普通滑动窗口失败的原因不只是丢失语义即使被删除的最初Token没有重要语义移除它们仍会破坏模型注意力分布。保留少量初始位置可以恢复稳定。第二持续生成长度和有效记忆长度可以分离模型可以在固定大小缓存下持续生成数百万Token但它实际利用的仍然只是有限最近上下文。也就是说运行时长度可以无限有效上下文长度仍然有限。这一区分是理解StreamingLLM最重要的结论。二十一、方法优点21.1 无需微调对已经训练好的LLaMA-2、MPT、Falcon和Pythia模型只需要改变KV Cache管理方式不需要重新训练。(arXiv)21.2 缓存大小固定无论总输入长度增长到多少KV Cache只包含Sink和最近窗口显存不会无限增长。21.3 实现简单它不需要计算Token重要性累积注意力动态排序训练压缩器保存额外摘要。只需要永久保留少量初始KV并滚动最近窗口。21.4 支持多种位置编码论文在RoPE和ALiBi模型上都进行了验证并设计了对应的位置重排方法。(arXiv)21.5 提供真实流式任务实验论文不仅测量PG-19困惑度也测试连续ARC问答、StreamEval以及实际解码延迟。二十二、方法局限22.1 不是真正的无限记忆中间历史Token被删除后无法访问。模型只能记住Sink和最近窗口。因此它不能代替长上下文模型外部记忆RAG历史摘要检索式对话记忆。22.2 初始Sink不保存长期内容Attention Sink主要维持注意力分布稳定并不自动保存开头的任务要求或事实。若系统提示或重要说明位于开头仅保留4个Sink位置可能会删除其大部分实际语义。LongBench结果已经显示保留4个Sink不如保留完整开头文本。(arXiv)22.3 对依赖远距离信息的任务无效一旦答案或证据超出最近窗口StreamEval准确率会逐渐下降并最终变为零。(arXiv)因此它更适合局部依赖的流式任务而不是跨越整个输入历史的推理。22.4 22.2倍基线需要谨慎理解其对比对象是昂贵的窗口重新计算而不是标准全KV解码。与其他现代推理框架相比实际加速幅度可能不同。22.5 Sink数量具有模型依赖性论文默认保留4个初始Token但Falcon可能1个已经足够LLaMA-2保留1个时仍明显较差新模型可能具有不同注意力结构专门训练Sink Token后只需要1个。因此4并不是普遍最优常数。(arXiv)22.6 专用Sink Token只在小模型上验证作者从头训练的模型规模为160M。论文没有证明在数十亿参数模型中加入专用Sink Token后也一定能获得相同效果和训练稳定性。(arXiv)22.7 Attention Sink的理论解释仍偏经验性论文将其主要归因于Softmax归一化和初始Token的全局可见性。这些解释与实验现象一致但并不是对Attention Sink产生机制的完整严格证明。二十三、这篇论文最重要的价值StreamingLLM最重要的贡献并不是简单提出“保留第一个Token”。它揭示了一个此前容易被忽视的问题KV Cache淘汰不仅会删除历史信息还可能破坏模型在预训练阶段形成的注意力数值结构。这意味着设计KV Cache压缩方法时不能只考虑哪些Token语义重要哪些Token距离最近哪些Token注意力最高。还要考虑删除某些位置后模型的Softmax注意力分布是否仍然处于熟悉的状态。StreamingLLM进一步区分了两种能力无限运行模型可以永远继续接收和生成Token无限记忆模型可以访问从开始到现在的全部历史。StreamingLLM解决了前者但没有解决后者。二十四、一句话总结StreamingLLM发现自回归大语言模型会将最初几个Token作为Attention Sink用来吸收大量与语义无关但Softmax必须分配的注意力普通滑动窗口一旦删除这些初始KV注意力分布便会发生剧烈偏移并导致困惑度崩溃。因此StreamingLLM固定保留约4个初始Token并配合滚动的最近Token窗口和重新排列的位置编码使LLaMA-2、MPT、Falcon和Pythia能够以恒定KV Cache持续处理超过400万Token并相对窗口重新计算最高加速22.2倍但它没有扩大模型真正可访问的上下文已被窗口删除的中间历史无法恢复因此“无限流式生成”不等于“无限长上下文理解”。下一篇仍可保持这一模板并继续重点区分论文宣称的理论能力与真实部署边界。