400 万 token，大模型推理飙升 22 倍！清华校友爆火一作
转自 | 新智元
编辑 | 润
【导读】大模型上下文限制不是事儿？
大语言模型的上下文长度限制从此不存在了？！
Meta、MIT、CMU的研究者最近刚刚发表了一篇论文，提出了一种被称为是「高效流式语言模型」（Efficient Streaming Language Models，ESLM）的方法，可以让有限上下文能力的语言模型能够支持几乎无限的上下文窗口。
https://arxiv.org/pdf/2309.17453.pdf
瞬间登顶Hacker News！
甚至，这一项目在GitHub上狂揽1.8k星。
上个动图大家先简单感受一下效果：
一句话来形容就是——完全停不下来！
因为LLM在预训练期间在有限的注意力窗口的限制下进行训练，尽管可以通过各种方法来扩展模型支持的上下文长度，但是模型可接受的序列长度仍然是有限的。
核心观点
研究人员想要解决的问题是：
能否在不牺牲效率和性能的情况下让大预言模型支持无限长度的上下文？
研究人员在尝试解决无限上下文长度输入时发现，限制模型能力最主要的原因是这两个：
1. 在解码阶段，基于Transformer的 LLM 会缓存之前所有token的键值状态（KV），如下图a所示，这可能会导致内存使用过多并增加解码延迟。
2. 现有模型的长度外推能力有限，当序列长度超出预训练期间设置的注意力窗口大小时，模型性能会严重下降。
解决这两个问题，研究人员尝试了两种方法：
第一种方法，称为窗口注意力（window attention）（上图b）。
这种方法只维护最新的token的KV状态的固定大小的滑动窗口。
但它的问题是虽然在缓存最初填满后确保了恒定的内存使用和解码速度，但一旦序列长度超过缓存大小，即使只是逐出第一个token的KV，模型也会崩溃，如下图所示。
另一种可能的策略是重新计算滑动窗口（如上图c所示）。这种发放为每个生成的token重建最新token的KV状态。
虽然这种方法提供了强大的稳定性，但由于需要在窗口范围内计算二次注意力（quadratic attention），这种方法的速度明显过慢，使得这种方法对于没法满足实际应用的需求。
在探索窗口注意力方法为什么行不通的过程中，研究人员发现了自回归LLM的一个有趣的特征：
大量的注意力分数被分配给初始token，而不管这些token与语言建模任务的相关性如何，如下图所示。
研究人员将这些占用了大量注意力的初始token称为——「注意力池（attention sinks）」
尽管这些初始token很多时候缺乏语义含义，但它们却占用了很高的注意力分数（attention scores）。
研究人员认为主要原因是因为Softmax操作，要求所有上下文token的注意力分数总和为1。
因此，即使当前任务和许多先前的token不匹配，模型仍然需要在某个地方分配这些不需要的注意力分数，使得分数总和为1。
初始token为什么会获得如此之高的注意力，原因也很简单：
初始token对几乎所有后续token都是可见的，因为自回归语言模型的性质，使初始token更容易被训练为「注意力池」。
基于这些原因，研究人员提出了一个新的框架：Streaming LLM。
这是一个简单而高效的框架，使得通过有限注意力窗口训练的LLM能够处理无限长度的文本，而无需进行微调。
Streaming LLM利用了初始token——「注意力池」具有高注意力值的这一特点，保留初始token可以保持注意力分数分布接近正常。
因此，Streaming LLM只是将注意力集中token的KV（只需 4 个初始token就足够了）与滑动窗口的KV保持在一起，以锚定注意力计算并稳定模型的性能。
使用 StreamingLLM，模型包括 Llama-2-[7, 13, 70]B、MPT-[7, 30]B、Falcon-[7, 40]B 和 Pythia
最后，研究人员证实了他们提出的「注意力池」的假设，并证明语言模型可以经过预训练，只需要一个注意力池token即可进行流部署。
具体来说，研究人员建议在所有训练样本的开头添加一个额外的可学习token可以作为指定的注意力接收器。
通过从头开始预训练 1.6 亿个参数语言模型，研究人员证明添加这个单一接收器token可以保持模型在流情况下的性能。这与普通模型形成鲜明对比，普通模型需要重新引入多个初始token作为注意力集中器才能达到相同的性能水平。
Streaming LLM
为什么窗口注意力方法行不通？
虽然窗口注意力技术在推理过程中提供了效率，但它导致了极高的语言建模复杂性。因此，该模型的性能不适合部署在流应用程序中。
在本节中，研究人员使用注意力接收器的概念来解释窗口注意力的失败，作为 StreamingLLM 背后的灵感。
识别困惑度激增点
上图显示了 20K token文本的语言建模的复杂性。很明显，当文本长度超过缓存大小时，由于排除初始token而导致困惑度度激增。这表明初始token，无论与预测token的距离如何，对于维持 LLM 的稳定性至关重要。
为什么LLM在删除初始token的KV时会崩溃？
研究人员将 Llama-2-7B 和模型的所有层和头部的注意力图可视化，如图 2 所示。研究人员发现，除了底部两层之外，模型始终关注所有层和头部的初始token。
含义很明确：删除这些初始token的 KV 将删除注意力计算中 SoftMax 函数（公式 1）中相当一部分分母。这种改变导致注意力分数的分布发生显着偏离正常推理设置中的预期。
对于语言模型中初始token的重要性，有两种可能的解释：（1）它们的语义至关重要，或者（2）模型学习到对其绝对位置的偏见。
为了区分这些可能性，研究人员进行了实验（上表），其中前四个token被换行token「\n」替换。
观察结果表明，模型仍然显着强调这些初始换行token。此外，重新引入它们可以恢复将语言建模的复杂性达到与原始初始token相当的水平。这表明起始token的绝对位置，而不是它们的语义值，具有更重要的意义。
LLM将初始token作为注意力池。为了解释为什么模型不成比例地关注初始token——无论它们与语言建模的语义相关性如何，研究人员引入了「注意力池」的概念。
SoftMax 函数（公式 1）的本质阻止所有参与的token具有零值。即使当前嵌入有足够的自包含信息用于其预测，也需要聚合来自所有层中所有头的其他token的一些信息。
因此，该模型倾向于将不必要的注意力值转储到特定token。类似的观察已在量化异常值领域，导致提出 SoftMax-Off-by-One作为一种潜在的补救措施。
为什么各种自回归LLM（例如 Llama-2、MPT、Falcon 和 Pythia）在注意力集中时始终关注初始token，而不是其他token？
研究人员的解释很简单：由于自回归语言模型的顺序性质，初始token对所有后续token都是可见的，而后面的token仅对有限的一组后续token可见。因此，初始token更容易被训练为注意力集中器，捕获不必要的注意力。
研究人员注意到，LLM通常接受训练，利用多个初始token作为注意力池，而不仅仅是一个。如上图所示，随着注意力的下降，引入四个初始token足以恢复 LLM 的表现。
相反，仅添加一两个并不能实现完全恢复。研究人员认为出现这种模式是因为这些模型在预训练期间没有在所有输入样本中包含一致的起始token。尽管 Llama-2 确实为每个段落添加了“”token作为前缀，但它是在文本分块之前应用的，导致大多数随机token占据第零个位置。
缺乏统一的起始token导致模型使用多个初始token作为注意力池。研究人员假设，通过在所有训练样本的开始处合并一个稳定的可学习token，它可以单独充当承诺的注意力接收器，从而消除对多个初始token来确保一致流的需要。研究人员将在第 3.3 节中验证这一假设。
3.2 使用注意力接收器滚动 KV 缓存
为了在已经训练好的LLM中启用LLM Streaming，研究人员提出了一种简单的方法，可以恢复窗口注意力的困惑度，而无需任何模型微调。
除了当前的滑动窗口token之外，研究人员在注意力计算中重新引入了一些起始token的 KV。
StreamingLLM中的KV缓存在概念上可以分为两部分，如上图所示：
1）注意力接收器（四个初始token）稳定注意力计算；
2）滚动KV缓存保留最新的token，这对于语言建模至关重要。StreamingLLM 的设计用途广泛，可以无缝地融入任何采用相对位置编码的自回归语言模型，例如 RoPE和 ALiBi。
在确定相对距离并向token添加位置信息时，StreamingLLM 重点关注缓存中的位置而不是原始文本中的位置。这种区别对于 StreamingLLM 的性能至关重要。例如，如果当前缓存有token [0, 1, 2, 3, 6, 7, 8] 并且正在解码第 9 个token，则分配的位置为 [0, 1, 2, 3, 4, 5, 6, 7]，而不是原文中的位置，即 [0, 1, 2, 3, 6, 7, 8, 9]。
3.3 使用注意力池预训练 LLMS
正如上文所述，模型过度关注多个初始token的一个重要原因是缺乏指定的接收器token来卸载过多的注意力分数。因此，该模型无意中将全局可见的token（主要是初始token）指定为注意力池。
一种潜在的补救措施是有意包含一个全局可训练的注意力接收器token，表示为「Sink Token」，它将作为不必要的注意力分数的存储库。或者，用 SoftMax-off- 等变体替换传统的SoftMax函数。
不需要所有上下文token的注意力分数总和为 1，也可能是有效的。
请注意，此 SoftMax 替代方案等效于在注意力计算中使用具有全零 Key 和 Value 特征的token。研究人员将此方法称为「零池（Zero Sink）」，以使其始终适合研究人员的框架。
为了进行验证，研究人员在相同的设置下从头开始预训练具有 1.6 亿个参数的三种语言模型。
第一个模型利用标准 SoftMax 注意力（Vanilla），第二个模型用 SoftMax1（零池）替换常规注意力机制，并且在所有训练样本中预先考虑可学习的占位符token（Sink Token）。如下表所示，虽然零汇在一定程度上缓解了注意力池问题，但模型仍然依赖其他初始token作为注意力池。
引入sink token对于稳定注意力机制非常有效。简单地将这个接收器token与最近的token配对就足以锚定模型的性能，并且由此产生的评估困惑度甚至略有改善。
鉴于这些发现，研究人员建议在所有样本中使用接收器token来训练未来的法学硕士，以优化LLM Stream的部署。
作者介绍
Guangxuan Xiao
MIT EECS 博士一年级在读，导师为韩松教授。
2022年毕业于清华大学计算机科学与技术系，师从刘志远教授和翟继东教授。研究方向是机器学习应用的高效算法和系统，特别是基础模型方向。
田渊栋
Meta AI Research (FAIR) 的研究科学家和高级经理，致力于强化学习、表示学习和优化方向的研究。
他是2021年ICML杰出论文荣誉奖和2013年ICCV Marr奖荣誉奖的第一作者获得者，并且是该项目的首席科学家和工程师。
他还获得了 2022 年 CGO 杰出论文奖。
此前，他于2013-2014年在谷歌自动驾驶汽车团队工作。他于2013年在卡内基梅隆大学机器人研究所获得博士学位。
BEIDI CHEN
她目前是 Meta/Facebook AI Research (FAIR) 的客座研究员。之前是斯坦福大学的博士后研究员，与 Chris Ré 博士一起工作。
将于 2023 年秋季开始担任卡内基梅隆大学电气与计算机工程助理教授，她的研究方向是大规模机器学习，设计和优化随机算法（算法-硬件协同设计）来加速大型机器学习系统解决现实问题。
韩松
韩松现任麻省理工学院（MIT）电子工程学院副教授。本科毕业于清华大学，在斯坦福大学取得博士学位。他提出了广泛应用于高效人工智能计算的「深度压缩」技术。
他还率先提出将权重稀疏性引入现代人工智能芯片的EIE（Efficient Inference Engine），并影响了英伟达安培GPU架构。他开创了TinyML研究，将机器学习引入到物联网设备，实现了边缘机器学习。
推荐阅读：