大家好,今天我们将讨论来自麻省理工学院、Meta和卡内基梅隆大学提出的论文《具有注意力同步的高效流式语言模型》。这篇论文做了什么?根据他们GitHub仓库的描述,他们在不牺牲效率或性能的情况下部署了支持无限长度输入的大语言模型。

从左侧可以看到,这种方法也被称为流式LLM(Streaming LLM)。左边没有开启流式语言模型,右边开启了。一段时间后,普通模型的性能会崩溃并耗尽内存;而流式LLM则可以持续运行下去。他们指出,部署用于多轮对话等预期长交互的流式应用是迫切需要的,但也面临一些挑战。

这种方法是针对此类应用场景设计的。例如,你正在与聊天机器人对话并希望长时间持续交互。如果你的模型是在固定数量的序列(如4000个token)上训练的,你不希望在达到4000个token后停止,而是希望继续无限延伸。这就是所谓的“流式”处理。

他们为此开发了一种方法,使在有限长度上训练的LLM能够泛化到无限序列长度,且无需任何微调。例如,如果模型是在4000个token的长度上训练的,该方法允许它在无限序列长度或任何长度的输入上工作,且无需微调。

他们表示,能够将此集成到流行的LLM(如Llama、MPT、Falcon和Pia)中,以执行高达400万token及更多的建模。这是一个巨大的突破:这些模型原本在更短的序列长度上训练,但他们利用这种方法将其扩展到大量token(如400万)。

他们通过观察和利用一种称为“注意力同步”的现象来实现这一点,我们稍后会详细讨论。他们说,在将LLM应用于无限输入流时主要有两个挑战:首先是内存使用增加和延迟升高。

如果在执行密集注意力机制(即关注所有之前的token),其复杂度将是序列长度的平方级。如果大幅增加序列长度,比如达到400万token,“爆炸”就会发生。其次是现有模型的长度外推能力有限,即当序列长度超过预训练期间的序列长度时,性能会下降。

他们的意思是,如果模型是在长度为N的序列上训练的,然后在2N的长度上进行推理,性能就会下降。首先,是否可能将训练长度为N的模型用于2N?这实际上是可能的,因为注意力是一种通用操作,可以应用于任何序列长度。

如果你有多个输入并以此训练(长度为N),想在2N上应用,这并没有真正的问题。因为你所做的只是将这些矩阵(W_key, W_Q, W_V)通过它们,然后通过Softmax计算注意力分数。没有什么能阻止你在架构中扩展它,无需任何更改即可做到这一点。(如果你使用具有固定序列长度参数的位置编码,那可能会有问题;但存在不受长度限制的位置编码。)

Alii论文做了类似的事情。我认为旋转位置编码(RoPE)也可以扩展,虽然我不确定具体细节,但你可以使用那些设计用于允许延长序列长度的位置编码。然而问题是,如果你这样做,性能会下降。

密集注意力(Vanilla Dense Attention)的复杂度为O(T*S),其中T是序列长度。它效率低下,且性能由困惑度(Perplexity)衡量,越低越好,越高越差。因此,原始的密集注意力行不通。即使进行长度外推,它在性能和效率上也会失败。

那么如何解决呢?有一种称为“窗口注意力”的方法可以解决至少性能问题——抱歉,是效率问题。在窗口注意力中,你在最近token的KV状态上维护一个固定大小的滑动窗口,并确保恒定的内存使用和解码速度。

一旦序列长度超过缓存大小,模型就会崩溃。他们的意思是:让我举个例子。X1, X2... 假设序列长度为3。当你处理到第4个token时,你需要为X2, X3和X4创建键值。但你真的需要计算X3的值吗?因为你已经在该窗口中处理过X3,并且已经拥有它的键和值。

因此,你实际上可以缓存这些先前值,只需为新token计算新的键和值。在这种情况下,如果上下文长度为L,你不再对每个token执行平方级操作(L^2),而是只执行线性操作。对于序列长度T,复杂度为O(T*L)。

这就是滑动窗口方法。这种窗口注意力具有良好的效率但性能较差。原因是:假设你在X8,并使用缓存的X6值(该值是以前在包含X4, X5和X6的上下文窗口中计算的),但现在你在不同的上下文中使用它。这种不匹配会导致性能下降。

因此,这种方法效率好但困惑度差(即性能差)。另一种策略是“带重计算的滑动窗口”,在这种情况下,当你处理X8时,你基本上忘记之前发生的所有事情,并重新进行整个注意力计算(这是平方级的)。

这种方法具有好的困惑度但效率差,因为复杂度为O(T*L^2)。而他们的方法是两全其美:你获得了线性复杂度和良好的困惑度(即良好性能)。以下是不同语言模型(Llama, Falcon等)的经验证据:随着输入长度增加,蓝色线表示密集注意力。在你超过预训练长度后……

困惑度会急剧上升,或者性能下降。而在窗口注意力机制中,当你达到KV缓存的大小并开始从初始令牌中进行驱逐时,性能也会下降。相比之下,他们提出的红色方法则能保持稳定的表现。为什么会发生这种情况呢?他们认为这是因为一种被称为“注意力同步”的现象,即大量的注意力分数被分配给初始令牌,无论它们与语言建模任务的相关性如何。

这是他们观察到的结果。如果你查看前两层注意力的情况,你会看到一个合理的注意力模式:所有令牌都具有一定的局部性。例如,某个令牌主要关注其附近的令牌,另一个令牌也关注一些邻近的令牌。这意味着第14个令牌主要关注周围的令牌,而对初始令牌的关注度并不高。但这仅适用于最初的几层(如前两层)。

然而,随着层级加深会出现差异。如果你看到那条深红色的线,这意味着大部分注意力分数都被分配给了第0个令牌。对于所有令牌(即使是第14或15个令牌),大部分注意力分数都流向初始令牌。这表现为一种“注意力同步”现象,这种情况出现在第2、9、16、23、31层等。模型 heavily 关注初始令牌,原因可能与Softmax有关:因为Softmax要求注意力分数之和为一。

当应用Softmax时,所有注意力分数的总和必须为1。因此,即使当前的查询在许多先前令牌中没有强烈的匹配项,模型仍然需要将这些不必要的注意力值分配到某个地方。这意味着,假设我们正在处理第14个令牌并与所有前14个初始令牌进行注意力计算。有可能这个第14个令牌与大多数令牌并不匹配,但由于注意力分数之和必须为1,它不得不给出相应的值。

模型学到的技巧是:如果匹配度不高,它就会将大部分注意力分数分配给初始令牌,以确保总和为1。例如X1、X2...X6,当处于X6时,所有令牌都会有一些注意力分数。即使它们并不真正匹配,所需的分数之和也必须为1。因此模型可能会给出类似0.9、0.01、0.02等数值,这些加起来必须等于1。

对于不同的上下文长度也是如此,注意力分数之和仍需为1。模型学到的“捷径”是:给前几个令牌(不一定是第一个,而是前面几个)赋予极高的注意力分数,从而使总和为1。至于为什么是初始令牌获得所有权重(即使它们本身的匹配度不高),是因为这些初始令牌对所有后续令牌都是可见的,使模型更容易训练出这种倾向。

因为如果存在上下文窗口,这些初始令牌始终存在于每个窗口中:在这个窗口里、在那个窗口里……它们是出现在每次注意力计算中的共同令牌。因此,模型学到的策略是:既然这个令牌在我进行注意力分数计算时始终存在,为了将总和凑为1,我会给它一个高分,并根据其他令牌的重要性给予较低的分数。这就像是为模型提供了一个归一化常数。

问题是,如果你使用窗口注意力机制并驱逐这些令牌,就会破坏这种计算结构,导致性能下降。而他们的方法本质上是:既然我们知道前几个初始令牌很重要,我们就保留这些令牌而不驱逐它们。例如在X9处使用窗口时,我会使用该窗口以及这两个初始部分。我不会忘记这些作为“注意力同步”作用的初始令牌。

在他们的案例中,他们保留了前4个初始令牌。这4个令牌固定不变,即使你在计算过程中使用后面的窗口,也会同时包含这4个初始令牌。因为它们充当“注意力同步”的作用。这就是他们论文及方法的简要总结:观察到前几个令牌具有高分数的“注意力同步”现象,因此保留这些令牌。

他们进一步调查了这种“注意力同步”现象的两种可能解释:一是语义至关重要(即前几个令牌在语义上有用);二是模型倾向于它们的绝对位置(即它们是初始令牌这一事实)。是为了区分这两种可能性,他们将前4个原始替换为换行符令牌(该令牌实际上没有语义意义)。

观察结果表明,无论初始部分是否为换行符令牌,“注意力同步”现象依然存在:即使只是垃圾令牌位于初始位置,它们依然获得高注意力分数。此外,如果保留这些令牌,模型困惑度恢复到良好水平。这意味着语义价值并非关键:即使初始阶段是垃圾令牌,只要它们位于该位置就能帮助模型。

这表明模型偏向于这些起始令牌的绝对位置,而非其语义价值。这再次解释了为何需要这样做:因为Softmax之和必须为1,模型需要一个地方来“倾倒”所有不必要的注意力值。而它们为何关注初始令牌而非其他令牌,是因为初始令牌在所有上下文窗口中都是可见的。

他们指出,预训练LLM倾向于利用多个初始令牌作为“注意力同步”机制,而不仅仅是一个。这些模型(如X1...)会反复利用初始令牌。

仅仅使用 X1 作为上下文窗口,或者所谓的注意力同步机制并不会有帮助。他们观察到的是,对于这些预训练模型,你需要大约四个初始令牌作为“感知”部分;只使用其中一两个并没有帮助。但这之所以发生,是因为你使用的是预训练模型,并试图将这种方法整合进去。他们指出的是,如果你真正考虑到注意力同步现象,并在开头有一个特殊的令牌——即同步令牌——那么你可以只使用其中一个。

如果在预训练期间拥有这个始终存在于注意力窗口中的同步令牌,即使在进行计算时也有这个同步令牌存在;如果在预训练阶段就进行这种操作,那么你可以只使用一个令牌。这一个令牌将充当注意力同步令牌。但是,如果你使用的是像 Llama 这样没有使用此方法训练的预训练模型,那么你将需要大约四个初始令牌。

在这里他们解释了如何实现这一点。其中一个重要点是:在确定相对距离并向令牌添加位置信息时,它关注的是缓存内的位置,而不是原始文本。这意味着当你在 X9 的位置时,位置不会是 7、8、9;而零的位置将是 0、1、2 和 3。因此,你关注的是缓存内的位置,而不是绝对位置。

然后他们解释了如何将此整合到不同的定位方法中,比如 RoPE 和 Alibi。这些都是主要的实现细节。接着他们描述了如果将注意力整合到预训练本身中(正如我之前提到的初始同步令牌),你可以只保留一个令牌并获得相似的性能。

在这里他们展示了其方法的有效性:对于非常长的文本(如 300 万或 400 万个令牌),困惑度保持稳定。即使将序列扩展到数百万,模型性能也不会下降。因此他们的方法有效。他们还表示,在预训练期间拥有同步令牌对模型收敛和性能没有负面影响;添加该同步令牌不会改变性能和收敛性。

此外,他们观察到使用该方法使流式处理速度快得多:实现了 22 倍的速度提升。但在这种情况下,他们将其与滑动窗口和重新计算的方法进行了比较,因为那是唯一具有相似性能(如困惑度)的基线。滑动窗口方法的复杂度为 T 乘以 L 平方,而他们的方案是 T 乘以 L,因此有 22 倍的速度提升。

总的来说,这是一篇非常酷的文章,因为它效果很好,并且可以整合到不同的预训练模型中。他们已经将其整合到了 Llama、Pia Falcon 和 Mistral(Mt)等模型中。代码是公开的,人们正在基于它制作演示版本,这很酷。他们还有一个很好的常见问题解答部分。

一个有趣的事实是:他们实际上并没有改变上下文窗口的大小。如果 Llama 2 是在 4096 上预训练的,那么其最大上下文窗口也是 4096。模型的上下文并没有真正改变;这并不意味着你有一个长度为 400 万的序列,当你在第 400 万的位置时,模型会记住在 200 万位置发生的事情。上下文不会转移到这里。

这并不意味着上下文窗口增加了,模型并不拥有所有 400 万个令牌的上下文。它只有非常局部的上下文,只有关于周围这 4096 个令牌的上下文。如果在第 200 万位置出现了一个名字,它在第 400 万的位置可能已经忘记了。正如他们所说:虽然你可以输入很长的文本,但模型只会识别最新的令牌。如果整本书被输入流式 LLM 中,它可能只总结最后的段落,这可能并不很有洞察力。

他们既没有扩展上下文窗口,也没有增强长期记忆;它只是使流式处理更快。即使只有局部上下文,事情也能以 T 乘以 L(线性复杂度)的速度进行;他们的工作是在不降低性能的情况下实现这种线性复杂性。本质上,这并不是说你神奇地扩展了 Llama 的上下文窗口至 400 万并记住第 200 万位置发生的事情。

事实是,你现在可以在流式处理的意义上对 Llama 用于长度为 400 万的输入:你可以保持生成,拥有局部上下文,并以线性复杂度高效地进行处理,而不是二次方复杂度。在上下文窗口大小方面不是长度;无论如何总长度是线性的。这就是视频的全部内容,谢谢观看。