大家好，我是Manish Gupta。在本视频中，我将探讨“注意力汇点”（Attention Sinks）这一有趣的概念，并讲解如何利用它进行语言建模，从而处理长达400万以上的token输入序列。如今的大型语言模型常以能处理极长序列为荣，它们是如何做到的呢？那么，“注意力汇点”究竟是什么呢？让我们谈谈什么是注意力汇点，它们为何会出现，以及为什么在这一意义上它们是成问题的。
大家在上图中看到的是自回归模型的注意力热力图，具体为Llama 2（70亿参数模型），输入序列长度为16。我们获得了不同层和不同头的注意力热力图，例如第0层的头0、第1层的头0……直到第31层的头0。观察发现，除了前两层（层0和层1）对初始token的注意力较低且模式更局部外，随着层级加深（如第2、9、16、23、31层），第一个token呈现出极高的注意力值（红色）。
在高层中，局部token的注意力不再高，反而第一个token获得了极高的跨头注意力。这种对初始token的高强注意力分数使得这些初始token被称为“注意力汇点”。这是因为在自回归建模中，某些词本身已足够充分，无需关注其他token。但Softmax函数强制注意力总和为1，且由于初始token在自回归过程中被频繁观察，大部分未分配的注意力质量便集中分配给了第一个或前几个token。
接下来谈谈这对各种注意力机制的影响。众所周知，密集注意力（Dense Attention）让每个token关注过去的所有token，效率低下。若输入序列长度为1024，需进行约百万次计算。在文本长度超过模型预训练序列时，性能会显著下降。相比之下，窗口注意力（Window Attention）通过仅缓存L个token并基于此计算前L个窗口的注意力，非常高效。
还有一种滑动窗口注意力配合重计算机制。这意味着不使用之前token的键值缓存，而是仅基于当前滑动窗口重新计算这些缓存。这与普通窗口注意力不同：后者可能使用了已不可用的先前上下文，而滑动窗口仅依赖当前窗口的token。有趣的是，如果窗口足够长，带重计算的滑动注意力实际上比普通窗口注意力能带来更低的困惑度（即更高的准确性）。
当第一个或前几个token被移出缓存时，普通窗口注意力的困惑度会显著上升；而重计算机制则能保持较低的困惑度。基于这一发现，流式LLM（Streaming LLM）成为一种利用注意力汇点的方法：既然前几个token至关重要，我们不仅关注局部窗口，也计算对初始少数token的注意力。这使得流式LLM具有较低的延迟（复杂度与序列长度乘以窗口大小成正比，而非平方），且准确性很高。
在下方的图表中，展示了不同输入序列长度（最高20k）下的复杂度与困惑度对比。图中绘制了密集注意力、窗口注意力和流式LLM的数据线，其中困惑度越低越好。观察发现，流式LLM是一个简单的想法：不仅关注局部上下文，也关注初始token。
关于输入序列长度中的前几个标记。正如你所观察到的，流式大语言模型（Streaming LLM）的基本困惑度值保持平稳且较低，即使序列长度达到两万。然而，如果你注意到“dension”（注：原文口误或特定术语），其困惑度会在输入序列长度超过预训练长度的瞬间急剧升高。同时，窗口注意力机制在KV缓存大小小于输入序列长度时也会产生极高的困惑度，这是不利的。在这种情况下，初始的几个标记会被从缓存中驱逐出去。
这是一张Llama 2 7B模型的图表，论文中也包含了其他模型的数据，但不如这张图直观。由于Llama 2使用RoPE嵌入，这种效应甚至更加显著。那么流式LLM如何解决注意力汇聚（Attention Sinks）问题呢？如前一张幻灯片所述，流式LLM的核心思路是保留四个注意力汇聚标记的键值缓存。
所谓前几个标记被称为“注意力汇聚标记”，除了保留滑动窗口的缓存外，还要专门保存这些标记的缓存。例如，在生成第9个标记时，不仅要关注6、7、8，还要始终关注最初的0、1、2、3。因此，本质上你总是会对这些标记施加注意力。
有趣的是，流式LLM的计算复杂度与带有重新计算（recomp）的滑动窗口几乎相同。然而，带重新计算的滑动窗口比流式LLM慢22倍，因此流式LLM更具优势。在执行此操作时，使用注意力汇聚标记和滚动KV缓存需要确保重新索引位置ID。
例如，在这种情况下，你不会将位置称为0、1、2、3、6、7、8，而是将其重新索引为0、1、2、3、4、5、6。本质上，你根据所使用的实际位置重新分配ID，让模型感觉这些注意力汇聚标记并非遥不可及，而是直接前置在当前滚动KV缓存之前。
你可以观察到，如果采用这种流式LLM方法——它不需要任何预训练或微调的改变，仅在推理时前置注意力汇聚标记——在不同模型（如Llama、Pythia、Falcon和MPT）中，困惑度都保持较低水平。即使序列长度达到百万级别（约400万标记），困惑度依然很低，表明流式LLM保留了这些模型的准确性和低困惑度。
这里展示的是如果改变注意力汇聚标记的数量会发生什么，而不仅仅是使用第一个。为什么要用四个？为什么不用一个、两个或三个呢？下表显示了不同模型的结果：使用零个同步标记（即窗口注意力）时困惑度很高；使用一个同步标记时，困惑度显著降低但未饱和。
如果使用两个标记，困惑度会进一步下降；而在某些情况下，使用四个标记时（至少对于Llama 2 7B），复杂度趋于饱和。这就是为什么他们建议必须使用四个注意力汇聚标记，这已经足够好了。此外，窗口注意力的大小并不那么重要。
事实上，他们观察到对于Llama模型而言，将4000大小的窗口缩小到2000甚至1000时，困惑度反而更低且并未显著升高。这意味着即使窗口尺寸较小而非极长，注意力汇聚的公式依然有效。
虽然使用四个标记效果很好，但能否在预训练阶段仅保留第一个注意力汇聚标记呢？因此他们还进行了一项有趣的实验：在设计上，将预训练时的第一个标记指定为注意力汇聚标记。这一概念类似于在视觉Transformer模型预训练中使用的寄存器（register）标记。
在这种情况下，他们仅使用一个注意力汇聚标记（即第一个标记），并对Pythia 1.6亿参数模型进行预训练。他们将第一个标记指定为全局可训练的专用注意力汇聚标记，本质上是一个通过反向传播随机初始化和训练的可嵌入向量。
观察预训练结果：不使用同步标记时，在较低层（如第0、2层）观察到大量的局部注意力；随着层级加深（如第10层），对初始标记的注意力显著增加。而使用同步标记时，从第零层开始注意力就指向该汇聚标记。
你会发现对同步标记的注意力在所有层级中都很高，而对其他初始标记的关注逐渐减少。
注意力机制，这正是你预期它在某种意义上所具备的特性。好的。现在想法是，如果你使用这个同步令牌（sync token），本质上你不再需要在缓存中保留最初的几个令牌。你可以基本上使用稀疏注意力或基于窗口的注意力，只需一个同步令牌即可。如果你不使用同步令牌，显然困惑度（perplexity）会非常高。但如果你使用同步令牌，
从第一个令牌开始困惑度值就开始下降。你并不需要更多的令牌来获得较低的困惑度值，对吧？所以如果你基本上是在进行预训练而不知道这一点的话……如果你是使用原始的注意力机制进行预训练，也就是原始流式语言模型（streaming LLM），你会观察到困惑度值虽然较低且合理，但当你使用可学习的同步令牌时，
在某种意义上你的困惑度值可以非常低。而且这种带有注意力同步机制的学习技术并不会导致下游任务性能的下降。所以基本上，如果你在各种自然语言处理任务上评估这些模型的准确性，你会观察到使用同步令牌时没有任何任务受到影响。所以在某种意义上这完全不是问题。好的，这就是基本内容。
那么流式语言模型在流式问答（streaming QA）和流式基准测试中表现如何呢？基本上，什么是流式问答？流式问答主要指的是这两个基准测试：streaming QA 和 streaming eval，它们在某种意义上属于长序列基准。那么什么是流式问答？基本上意味着你取这两个数据集：ARC Challenge 或 ARC Easy，
对于每个数据集，你要做的是将所有问答对连接起来。基本上它可能包含几个问答对。你将这些数百个问答对连接成一个单一的序列，然后将其交给模型，并在答案位置预测答案。在每个答案位置，你使用精确匹配准确率来评估模型的表现如何。好的。由于序列非常长，你知道密集注意力基本上意味着这些模型——Llama 2 70亿参数、130亿参数，
以及 LLaMA-2 700亿参数的图表——基本上都会给出内存溢出异常。好的，窗口注意力给出的准确率很差，正如你在这里看到的。但流式语言模型基本上给出了非常好的准确率，正如你在这里看到的。好的，事实上，这里显示的这些准确率等同于单样本（oneshot）准确率，当你基本上将这些示例逐个提供给
这些模型时。所以如果你逐个给出示例，序列长度很小，因此这些模型在 ARC Challenge 和 ARC Easy 基准测试中表现非常好。即使在将所有内容连接起来并作为单个序列给出时，这些准确率也会保留下来，如果你进行流式语言模型类型的评估。好的，他们进行的另一种评估是在这个 StreamEval 基准上
进行的。所以 StreamEval 基准基本上是这样的：你知道你创建这样的样本。这是 StreamEval 基准中的其中一个样本，其中基本内容是说“以下是我想让你记住的行记录”。你给出大约20行，然后在第20行之后的某个地方，
给它一些数值让它记住这20行中的内容。好的，在第20行之后的基本内容是问“第零行的内容是什么？”所以你询问它在时间上倒退20行时的内容。好的。然后你基本上再给出10行，并要求模型告诉你第10行的内容，这又是时间上倒退20行的位置。是的。在第30行之后，第10行是时间上倒退20行的位置。同样地，在第40行之后的第20行
是时间上倒退20行的位置，依此类推。好的。模型应该记住哪个数字是在时间上倒退20行时提到的。好的。这是一个期望的输出。好的。所以在这个基准测试中，正如你可以观察到的那样，这些是使用不同模型的结果，比如 Llama 2 模型、Llama 2 Chat 70亿和130亿参数，还有这些长模型
它们是在非常长的序列长度上进行预训练的。好的，那么你观察到什么？你观察到在这些不同的基准测试和这些不同的模型中，流式语言模型基本上保留了非常高的准确率分数。你看到很高的准确率分数。然而，密集注意力和窗口注意力在准确率方面突然下降。
你知道，一旦你的输入长度超过预训练长度，密集注意力的准确率就会迅速下降；而窗口注意力在输入长度超过 KV 缓存大小、即最初的几个令牌从 KV 缓存中被驱逐时，准确率开始下降。好的。所以就是这样。简而言之，在这段视频中我们讨论了注意力同步机制以及窗口注意力的低效性或者说是
在某种意义上较差的性能，特别是针对长序列的窗口注意力。当初始令牌被排除在外时，即当你的序列长度基本增加并超过 KV 缓存长度时，窗口注意力的表现会变差。好的，带有重计算的窗口注意力可以保留那种准确率，然而重计算是非常昂贵的，它需要大量时间
特别是当窗口大小很大时。流式语言模型本质上是将注意力同步机制添加到最近的令牌中。这基本上是一个非常简单的想法，但其效果非常棒。它将语言模型的预训练窗口大小与其实际文本生成长度解耦了。它在某种意义上甚至能以400万个令牌表现出非常出色的性能。
使用专用的同步令牌进行预训练基本上能提高性能，对吧？虽然我的意思是我认为在这篇论文中关于预训练的想法似乎有点过于强调。我是说，我不知道是否真的有人会 bother 去保留四个注意力同步令牌 versus 一个注意力同步令牌。可能四个也是同样可以的。所以我不确定是否真的需要预训练，但一般来说
使用最初的几个注意力同步令牌以及最近的基于窗口的注意力的想法是一个很棒的想法，这基本上帮助他们建模非常长的序列。好的，就这样了。这就是这段视频的全部内容。你知道的，
谢谢观看。希望你喜欢这个视频。在我的 LinkedIn 上联系我或查看我主页上的研究成果。