大模型基础架构的变革：剖析Transformer的挑战者（下）
上一篇文章中，我们介绍了UniRepLKNet、StripedHyena、PanGu-π等有可能会替代Transformer的模型架构，这一篇文章我们将要介绍另外三个有可能会替代Transformer的模型架构，它们分别是StreamingLLM、SeTformer、Lightning Attention-2，想了解之前的内容，请前往《大模型基础架构的变革：剖析Transformer的挑战者（中）》一文查阅。
七、StreamingLLM
StreamingLLM是一种结合了模型固有的「注意力池」token和文本中最近 token 的滚动缓存进行处理无限文本的新型架构。
在流媒体应用程序（如多轮对话）中部署大型语言模型（LLM），迫切需要长时间的交互，但是存在两个主要的挑战。第一个挑战是在解码阶段，缓存以前的token的键值状态（KV）会消耗大量的内存；第二个挑战是主流的LLM不能推广到比训练序列长度更长的文本。窗口注意力只缓存最近的kv，这是一种自然的方法——但是作者发现，当文本长度超过缓存大小时，它就会失败。
因此观察到一个有趣的现象，即注意力池集，即保持初始token的KV将在很大程度上能够恢复窗口注意的性能。在论文中，作者首先证明了注意力池聚的出现是由于初始token作为“汇聚的强注意分数”，即使它们在语义上并不重要。基于上述分析，作者引入了StreamingLLM，这是一个非常有效的方法，使应用了有限长度注意窗口训练的LLM可以推广到无限序列长度，而无需任何微调。
目前业内对将LLM应用于长文本进行了广泛的研究，主要关注三个领域：长度外推、上下文窗口扩展，以及提高LLM对长文本的利用。.
长度外推的目的是使在较短文本上训练的语言模型能够在测试过程中处理较长的文本。一个主要的研究途径是开发Transformer模型的相对位置编码方法，使它们能够在训练窗口之外发挥作用。
上下文窗口扩展集中于扩展LLM的上下文窗口，允许在一次向前传递中处理更多的token。一个主要的工作方向是解决训练效率的问题。考虑到Transformer在训练过程中存在对注意力计算随token增长的二次复杂度，开发一个长上下文的LLM，无论是计算还是存储都面临着巨大的挑战。
对不同LLM中有20K token的文本进行语言建模困惑。观察结果显示出一致的趋势：
(1)一旦输入长度超过训练前的注意窗口大小，密集的注意就会失败。
(2)一旦输入长度超过缓存大小，窗口注意就会崩溃，即初始token被丢弃。
(3)streamingLLM表现出稳定的性能，其困惑几乎与滑动窗口的重新计算基线匹配。
改进LLM对长文本的利用，以更好地捕获和使用上下文中的内容，而不是仅仅将它们作为输入。在长度外推和上下文窗口扩展两个方向上的成功并不一定会转化为对冗长的上下文的有效利用。在LLM中解决这种对长期上下文的有效使用仍然是一个挑战。论文的主要工作集中于稳定地利用最新的token使LLM的无缝流应用。
作者最终证明了StreamingLLM可以使Llama-2、MPT、Falcon和Pythia能够使用多达400万个token来执行稳定和有效的语言建模。此外，作者还发现在训练前添加一个占位符token作为专用的注意力池可以进一步改善流媒体部署。在流媒体设置中，StreamingLLM的性能比滑动窗口重新计算基线高出了22.2倍的加速。
StreamingLLM与现有方法的说明：语言模型对长度为L的文本进行预训练预测token（T≫L）。
(a)密集注意具有O（T 2）时间复杂度和不断增加的缓存大小。当文本长度超过训练前的文本长度时，其性能就会下降。
(b)窗口注意缓存最近的L的token的KV。虽然在推断方面很有效，但一旦启动token的键和值被丢弃，性能就会急剧下降。
(c)滑动窗口具有重新计算的重建能力，它为每个新token的L的最近token重建KV状态。虽然它在长文本上表现良好，但它的O（T L²）复杂性，源于上下文重新计算中的二次注意，使得它的速度相当慢。
(d)StreamingLLM保持注意力池（几个初始token），以实现稳定的注意力计算，并结合最近的token。它非常高效，并且在扩展文本上提供了稳定的性能。在PG-19测试集的第一本书（65Ktoken）上的Llama-2-13B模型测量复杂性。
当对无限输入流应用llm时，出现了两个主要的挑战：
1.在解码阶段，基于Transformer的LLM缓存所有以前token的密钥和值状态（KV），这可能导致过度的内存使用和增加解码延迟。
2.现有模型具有有限的长度外推能力，即当序列长度超过训练前设置的注意窗口大小时，其性能会下降。
为窗口注意力只在最近的token的KV状态上保持一个固定大小的滑动窗口。尽管它确保了在缓存初始填充后恒定的内存使用和解码速度，但一旦序列长度超过缓存大小，模型就会崩溃，甚至会丢弃掉第一个token的KV。另一种策略是带有重新计算的滑动窗口注意力，它为每个生成的token重建最近token的KV状态。虽然它提供了很强的性能，但由于在其窗口内计算二次注意力，该方法明显较慢，这使得该方法在现实流媒体应用中不现实。
为了理解窗口注意的问题，作者发现了一个有趣的自回归LLM现象：大量的注意力分数被分配给初始token，而不管它们是否与语言建模任务有相关性。作者称这些token为“注意力不集中”。尽管他们缺乏语义意义，但他们收集了显著的注意力分数。作者将此原因归因于Softmax的操作，它要求所有上下文token的注意力分数总和为1。
因此，即使当前查询在许多以前的token中没有强匹配，模型仍然需要将这些不需要的注意力值分配到某个token，以便其汇总，让总和为1。初始token作为汇聚token背后的原因是直观的：由于自回归语言建模的性质，初始token对几乎所有后续的token都是可见的，使它们更容易被训练为注意力池。如果删除掉初始的几个token的KV，模型的表达能力会急速下降。
为什么在删除初始token的KV时，LLM会中断？作者可视化了Llama-2-7B和模型的所有层和头部的注意力地图。最终发现，在底部的两层之外，该模型始终关注于所有层和头部的初始token。其含义很清楚：在注意力计算中，去除这些初始token的KV将去除SoftMax函数中相当大的一部分分母。这种变化导致注意力分数分布偏离正常推理设置的预期的显著变化。
对于初始toeken在语言建模中的重要性，有两种可能的解释：
(1)它们的语义是至关重要的
(2)模型学习了对其绝对位置的偏向。
为了区分这些可能性，作者进行了实验，其中前四个token被替换为行划线token“\n”。观察结果表明，该模型仍然显著强调这些初始划线的token。此外，重新引入它们可以将语言建模的困惑恢复到可与拥有原始初始token相媲美的级别。这表明，这表明起始token的绝对位置比它们的语义值更重要。
LLM在注意力集中时关注初始token。为了解释为什么该模型不成比例地关注初始token——而不管它们与语言建模的语义相关性如何，作者引入了“注意力池”的概念。SoftMax函数的性质具备防止所有参与计算的token出现零值，因此会导致在所有层的所有头部中聚合来自其他token的一些信息，即使当前的嵌入有足够的独立信息来进行预测。因此，该模型倾向于将不必要的注意值转储到特定的token上。
在自回归的llm，如Llama-2、MPT、Falcon和Pythia中，关于注意力集中时始终关注最初的token，而不是其他token？更简单的解释是：由于自回归语言建模的顺序性质，初始token对所有后续token都是可见的，而后来的token只对有限的一组后续token可见。因此，最初的token更容易被训练成为注意力池，吸引不必要的注意力。
基于上述见解，作者提出了StreamingLLM，这是一个简单而有效的框架，使应用有限注意窗口训练的LLM能够在不进行微调的情况下处理无限长度的文本。StreamingLLM利用了注意力池具有较高的注意力值这一事实，保留它们可以保持注意力分数分布接近正态分布。
StreamingLLM只是保持注意力池内token的KV（只有4个初始token就足够了）和滑动窗口的KV，以锚定注意力计算和稳定模型的性能。通过StreamingLLM，包括Llama-2-[7,13,70]B、MPT-[7,30]B、Falcon-[7,40]B和Pythia- [2.9,6.9,12]B在内的模型可以高效地创建400万个token，甚至可能更多。与唯一可行的基线，通过重新计算的滑动窗口相比，StreamingLLM实现了高达22.2倍的加速，实现了LLM的流媒体应用。
重新引入初始token对StreamingLLM的影响：
(1)引入一两个初始token通常不足以完全恢复模型的困惑，这说明模型并不仅仅使用第一个token作为注意力池。
(2)引入四个初始token一般就足够了；进一步增加的收益会递减。缓存配置x+y表示将初始token添加到最近的token。复杂性在连接的PG19测试集中的400K token上进行评估。
根据研究发现，LLM通常被训练为使用多个初始token作为注意力池，而不仅仅是一个。引入四个初始token，作为注意力的集合，就足以恢复LLM的性能。相比之下，只添加一两个并不能实现完全恢复。作者认为这种模式的出现是因为这些模型在训练前没有包含所有输入样本的一致的起始token。
尽管Llama-2在每个段落前都有一个“<>”token，但它在文本分块之前应用，导致一个几乎随机的token占据第零位置。由于缺乏统一的起始token，导致模型使用几个初始token作为注意力池。如果通过在所有训练样本开始时合并一个稳定的可学习token，那么它可以奇异地作为一个承诺的注意力池，消除需要多个初始token，以确保一致的streaming。
作者确认了注意力池的假设，并证明了语言模型预训练只需要一个注意力池token便可以来进行流媒体部署。作者建议在所有训练样本开始时，增加一个额外的可学习token来作为一个指定的注意力池。作者通过从头开始对1.6亿个参数语言模型进行预训练，证明了添加这个单一的接收token可以保持模型在流媒体中的性能。这与普通模型形成了鲜明对比，普通模型需要重新引入多个初始token作为注意力池，以达到相同的性能水平。
streamingLLM的设计是多用途的，可以无缝地集成到任何使用相对位置编码的自回归语言模型中。为了在已经训练过的LLM中实现LLM流，作者提出了一种简单的方法，可以在不进行任何模型微调的情况下恢复窗口注意力的困惑度。在当前的滑动窗口的token之外，还在注意力计算中重新引入了一些启动token的KV。
StreamingLLM中的KV缓存可以分为两部分：
(1)注意汇（四个初始token）稳定注意力计算；
(2)滚动KV缓存保留了最新的token，这对语言建模至关重要。
总体来说，StreamingLLM保存了模型固有的「注意力池」（attention sinks）锚定其推理的初始 token。结合最近 token 的滚动缓存，使LLM能够在不进行微调的情况下处理无限的文本。通过在最近的token中增加注意力池，StreamingLLM可以有效地建模多达400万token的文本。带有专用接收器token的预训练模型可以提高流媒体性能。StreamingLLM解耦了LLM的训练窗口大小和实际文本生成长度，为LLM的流媒体部署铺平了道路。
八、SeTformer
SeTformer是一种通过将点积自注意力（DPSA）完全替换为自我最优传输（SeT）来提高性能和计算效率的新型架构。
点积自注意（DPSA）是Transformer模型中一个重要的组成部分。简单来说，它是一种帮助模型聚焦在输入数据中关键部分的机制。然而，使用DPSA在处理长文档或高清图像时，由于一些复杂的运算，比如softmax操作，会导致计算时间和内存的大幅增加，变得非常耗费资源。
为了解决这个问题，科学家们想到了一种新的方法，叫做自优传输Transformer（Self-optimal Transport ：SeT）。在这个新型模型中，他们完全舍弃了原来的DPSA，采用了一种叫做最优传输的方式，以期望在性能和计算效率上都有所提升。SeT基于两个基本的softmax属性：一是保持注意矩阵中的数值都是非负的，二是利用一种非线性的机制来突出显示输入序列中的关键信息。
通过引入最优传输的核代价函数，SeT模型成功地满足了这些特性，从而在处理长序列时更为高效。这种新方法在实际应用中表现得非常出色，特别是在处理小型和基础模型时，SeTrermer在ImageNet-1K数据集上表现出色，top1精度分别达到了84.7%和86.2%，给人留下了深刻的印象。在目标检测领域，SeTformer在mAP上比FocalNet高出了2.2个百分点，并且还成功减少了29%的参数量。在语义分割任务中，基本大小模型相较于NAT，mIoU提高了3.5个百分点，同时参数减少了33%。此外，SeTformer在GLUE基准测试中的语言建模方面也取得了先进的结果。
这些发现凸显了SeTformer在视觉和语言任务中的广泛适用性，表现出其在不同领域的强大性能。这意味着SeTformer可以更高效地处理图像识别、目标检测、语义分割等任务，同时在语言建模方面也取得了先进的成果。SeTformer的成功应用为各种应用场景提供了一种更强大、更高效的解决方案。
在研究过程中发现，虽然softmax计算的成本很高，但它是一个用于精确计算注意力权重的函数。在面临这一挑战时，一个基本问题浮现：“是否存在一种替代DPSA的方法，可以在有效地建模，并且在远程交互时保持softmax的特性呢？”
在softmax注意力中，有两个关键因素影响其性能：注意矩阵中的非负元素、以及增强注意权重稳定性的非线性重加权机制。为了应对这些因素，作者提出了一种新的DPSA变体，名为自最优传输（SeT）。SeT利用最优传输（OT）原理和核方法，以满足softmax注意力的这两个关键特性。
在SeT中，输入特征经过映射，进入一个再生核希尔伯特空间（RKHS）。在该空间中，点评估采用线性函数的形式。与其使用点积来寻找元素之间的相似性得分不同，作者采用了通过计算输入和参考特征之间的最优传输而获得的对齐得分。
OT是一种有效的对齐分布的计算技术，OT的目标是找到将质量从一个位置传输到另一个位置的最有效方法。在Set模型中，通过将输入中的每个特征与参考集对齐，OT有效地计算相似度权重。这种方法在维护softmax特性的同时，提供了一种有效的建模远程交互的途径。
SeT包含两个简单步骤，首先是构建核特征映射，然后是应用OT计算对齐矩阵。详细过程如下：
第一步，先将输入特征向量引入到RKHS的空间中。在这个空间里，使用一个正定核来确保所有特征都是非负的。这有助于模型避免将与任务无关的信息聚合在一起。为了更好地处理相关性，作者采用了一个非线性重加权方案，通过根据标记与参考集的相关性来分配不同的权重。
第二步，应用OT来对齐输入特征向量和参考集，从而增强局部相关性并捕获复杂的依赖关系。这种对齐的过程有助于模型更好地理解不同特征之间的关系，特别是在处理视觉和语言任务时。
为了保持计算的可行性，作者使用核方法来获取有限维度的嵌入，这类似于Swin的方法。通过核的组合，最终构建了跨层次的降采样特征图，使得模型更加高。Setformer通过广泛的实验证明了其在各种任务中的有效性。
在DPSA中，每个输入特征xi都汇聚了所有其他x的特征，形成一个很大的稀疏矩阵。Set模型采用了一种叫做OT（最优传输）的方法，通过基于最佳匹配的x和y特征进行聚合。这意味着模型在整合特征时，通过找到最佳匹配的方式，让模型更智能地学习如何有效地汇聚信息。这样的方法有助于提高模型对输入信息的理解和处理能力。
更进一步地说，OT基于最优传输的思想是通过一种有效的方式将一个特征集从一个地方传输到另一个地方，以实现最佳匹配。这对于模型的理解能力来说是很有帮助的，因为它让模型更聪明地汇总信息，使得在处理各种任务时都能更好地发挥作用。这种智能的信息整合方式有助于提高模型的性能，使其更适用于各种实际应用场景。
Setformer架构包含了四个序列阶段，每个阶段都有多个SeT块。这些连续的阶段之间通过下采样器层连接在一起，这样可以既减小特征空间大小，又增加处理的深度。这个架构的设计就是为了更好地处理数据信息，使得模型能够更有效地学习和理解数据中的特征。
在右边的部分，解释了注意力公式，也就是模型是如何聚合信息的。简而言之，即把数据中的元素x和y映射到一个叫做RKHS的空间，然后通过计算它们之间的最优传输（OT），将它们与相应的参考对齐。这就是说，模型通过找到元素之间最佳的匹配方式，来有效地整合它们的信息。这种注意力机制有助于模型更有针对性地关注数据中重要的部分，提高处理效率。
总体来说Setformer是一种高效的加权注意力模型，它基于自优传输（SeT）的概念。作者t借鉴了内核方法和最优传输的特性，这有助于模型更好地捕捉数据中复杂的关系。这种模型确保了注意力分数是非负的，并采用一种类似于softmax机制的重新加权方案，使得注意力集中在关键的信息上。
通过巧妙地利用这种独特的结构，作者还构建了一系列setromx-x架构，它们在视觉和语言领域一直表现出色，超越了已有的Transformer模型。这些架构不仅提高了模型的性能，而且在处理复杂的任务时也表现得更为出色。
九、Lightning Attention-2（TransNormer）
Lightning Attention-2是OpenNLPLab在2023年7月份提出TransNormerLLM模型之后，再次提出的一种能生成无限序列长度、具恒定算力开销、以及具更高建模精度的新注意力机制。Lightning Attention-1则是在TransNormerLLM模型中提出的一种注意力方式。
关于transformer随着token的增长而出现的计算复杂性的二次方增长问题逐渐成为影响大模型进一步发展的阻力，线性注意力一直被认为是一个不错的替代方案。它能够以线性计算的方式处理token，理论上可以在不减慢速度的情况下处理无限长度的序列，也就是说，在使用一定内存的情况下，可以以相同的速度处理不同长度的序列。然而，目前的线性注意力算法由于在计算中的累积和问题，无法在实际情况中展现其理论上的优势。
需要注意的是，在线性关注中，从O（n^2）到O（n）的复杂性显著减少只是在理论上，实际上可能不能直接转化为硬件上计算效率的显著提高。实际实施中，要实现实质性的加速面临一些挑战，主要有两个问题：(1) GPU上的内存访问可能会对整体关注计算速度产生影响。(2) 线性关注核心技巧所需的累积求和操作阻碍了它在实际应用中达到理论上的训练速度。
上面的第一个问题，已经被Lightning Attention-1成功地解决了。在本文中，将主要介绍Lightning Attention-2是如何来解决第二个问题的。Lightning Attention-2是一个实现线性注意的方法，使其能够充分发挥其理论上的计算优势。为了实现这一点，作者采用了平铺的思想，分别处理线性注意计算中的块内和块间组件。其关键思想是在线性注意计算中分别处理块内和块间分量，利用“分治”的概念。
具体的实现步骤是，对于内块，保持使用传统的注意计算机制来计算QKV的乘积；而对于间块，使用线性注意核技巧。平铺技术在正向和后向过程中实现，以充分利用GPU硬件功能。因此，Lightning Attention-2可以训练具有无限序列长度的LLM，而不需要额外的成本，因为在固定的内存消耗下，其计算速度随着序列长度的增加而保持不变。
对不同序列长度范围的Lightning Attention-2进行了全面评估，以评估其准确性，并将其计算速度和内存利用率与Lightning Attention-2和Lightning Attention-1进行了比较。研究结果表明，Lightning Attention-2在计算速度上具有显著的优势，这是由于其创新的内部分离策略。此外，在不影响性能的情况下，Lightning Attention-2的内存占用减少了。
首先看一下线性关注的公式，然后介绍一下我们新提出的Lightning Attention-2。在TransNormer模型内的NormAttention的情况下，Lightning Attention-2计算采用了一种不同于传统Transformer结构的方法，避免了softmax函数和缩放操作带来的复杂性。正常的注意力机制可以用以下方式来表示：
上式中，Q、K、V∈R^(n×d)分别为查询矩阵、键矩阵和值矩阵，n表示序列长度，d表示特征维数。为了利用右矩阵乘法固有的计算效率，上面的方程可以无缝地等价地转换为其线性变量：
这种线性公式便于循环预测的O（nd^2），使其在训练中相对于序列长度有效。此外，采用线性注意确保了O（d^2）的恒定计算复杂度，而不考虑序列长度，从而能够对无限的长序列进行推理。这一成就是通过反复更新K^T V来实现的，而不需要重复计算整个注意矩阵。相比之下，标准的softmax注意力在推理过程中需要的计算复杂度为O（md^2），其中m表示token索引。
下图详细介绍了Lightning Attention-2的结构框架。在第i次迭代中，把矩阵Qi、Ki、Vi的块从高带宽存储器（HBM）搬移到静态随机存取存储器（SRAM）。在SRAM中，输出Ointra和Ointer分别被独立计算，然后更新到KV矩阵。接着，最终的输出Oi，也就是Ointra和Ointer的总和，被写回到HBM。
Lightning Attention-2中采用的“分治”和平铺技术的思想，巧妙地解决了当前线性关注算法的一些限制，尤其是在处理累积求和的挑战时。通过将计算分为块内和块间两个部分，巧妙地利用了GPU硬件的性能，确保了高效的计算。
通过广泛的实验证明，Lightning Attention-2在不同模型大小和序列长度上都表现出色。它不仅在处理长序列时能够保持一致的训练速度，而且在速度和准确性方面都优于目前最先进的关注力机制。这一技术突破对未来的大型语言模型具有深远的影响，尤其是那些需要处理长序列的模型。
总体来说，Lightning Attention-2是一种创新的线性关注方法，充分发挥了其在理论计算上的优势，特别是在处理因果关系时表现出色。在处理块之间时应用了线性关注核心技巧。通过正向和反向的过程利用平铺技术，充分发挥了GPU硬件的性能。在涉及不同模型大小和序列长度的多种实验中发现，不论输入序列有多长，Lightning Attention-2都能保持一致的训练和推理速度，并且明显快于其他关注机制。
未来有望结合Lightning Attention-2引入序列并行性，旨在有效地克服当前硬件的限制，促进超长序列的训练。这将为语言模型领域带来更大的发展空间。
以上列举出的模型只是可能在未来会替代Transformer架构的基础模型，至于替代的可能性有多大，多久能够替代，这些都跟实际的模型应用情况相关，不仅要看学术界的努力，更要看工业界的行动。关于这一点，在上一篇文章《Transformer架构的局限已凸显，被取代还有多久？》中已有详细论述，可以翻看查阅。
除本系列文章所列出的9种改进或可能替代Transformer的模型之外，还有HGRN、TNN等其他类似的语言基座模型，以及D-iGPT、Conv2Former、SPFormer、QFormer 、Fus-MAE、VisionMamba、FlatFormer、PLG-ViT 、EfficientViT等视觉类基座模型，可以肯定的是未来有可能替代Transformer的基础模型还会继续出现，最终到底哪个模型会更有可能胜出，未来是否会出现下一个王者，就像Transformer一样一统江湖，还是会出现多类基础模型并驾齐驱，百家争鸣的局面，这就要看哪一个模型在设计方式上更能够降低训练和应用成本，提高模型的生产效率和效果了。