在长上下文的大语言模型中Transformer 架构：综述

24年2月南京大学的论文“Advancing Transformer Architecture in Long-Context Large Language Models: A Comprehensive Survey”。

目前的LLM主要是针对短文本片段进行预训练，这损害了它们处理实际场景中经常遇到长上下文提示的有效性。本文对基于 Transformer 的 LLM 架构最新进展进行了全面综述，旨在增强 LLM 在整个模型生命周期（从预训练到推理）中的长上下文功能。首先描述并分析当前基于 Transformer 的模型处理长上下文输入和输出的问题。然后，提供 Transformer 架构的分类和升级方案来解决这些问题。随后，针对长上下文 LLM 广泛使用的评估方法进行了调查，包括数据集、指标和基线模型，以及库、框架和编译器等优化工具包，提高 LLM 在这些不同阶段的效率。运行。最后，讨论未来研究的挑战和潜在途径。

Transformer 是一种复杂的深度神经网络模型，它集成了多种先前的设计 [11,13,86] 和新组件来支持序列-到-序列的语言建模，最初用于机器翻译 [225]。当代LLM大多采用 Transformer 架构，利用其模块 [61,178,181]，其中的成功主要归功于其精心设计的注意机制，该机制捕获整个输入中每对tokens的全局依赖关系，从而使模型成为可能处理有复杂关系的序列。然而，其时间和空间二次复杂性带来了巨大的计算资源挑战，限制了训练期间的输入文本长度和推理期间的有效上下文窗口。此外，位置嵌入（PE）缺乏稳健且可泛化的机制，会导致推理过程中的性能下降和波动，特别是在相关信息较长的序列或位置漂移的情况下[139]。

随着LLM根深蒂固于需要长上下文理解[114, 248]和生成[89, 142]的各种应用中，对能够有效且高效地理解和生成极长序列的长上下文LLM的需求变得越来越紧迫。因此，研究人员投入了大量精力来增强 Transformer 架构，以解决LLM中的长上下文问题，包括注意效率的优化、具有额外记忆机制的上下文窗口扩展、外推 PE 的有效长度泛化、上下文预处理/后处理和其他杂项方法，例如特定的预训练目标、专家混合 (MoE)、量化、并行性等。

如图是该综述概述：（a）当代基于 Transformer 的解码器LLM的典型架构图，图例位于最右上角；(b) 增强 Transformer 架构模块的方法分类（按颜色对应（a））：高效注意（注意核的子模块）、长期记忆（针对 KV 缓存）、外推 PE（针对位置嵌入模块） 、上下文处理（与上下文预处理/后处理相关）和杂项（一般用于整个解码器块以及损失模块）。

简而言之，（神经）语言建模旨在近似任何给定文本出现的对数概率，表示为 logP(𝑋1:𝐿; 𝜃)，其中 𝜃 代表要学习的网络参数，𝑋1:𝐿 包括 长度为 𝐿 的序列代表自然语言，包括单词、标点符号、数学符号等。语言建模的一个重大实际障碍是维数灾难，即随 𝐿 的增加，概率分布的支持呈指数级增长。LLM采用掩码语言建模 (MLM) 和因果语言建模 (CLM) 等变型。

MLM建模目标：

普通的 Transformer 架构[225]主要包括一个编码器和一个解码器，每个编码器都堆叠有多个相同的块。每个块的骨架大部分与图（a）中的骨架兼容。一般来说，第一个块采用由词嵌入层编码的token化序列，然后是 多头自注意（MHA） 层，其具有与特定语言建模目标相对应的注意掩码和 前馈网络（ FFN） 层。MHA 和 FFN 层在块的每个入口/出口处都进行 层归一化（BN） [11] 和残差连接 [86] 。然后，每个更高层的块将前一个块的输出隐藏状态作为输入，用 MHA 和 FFN 层表示它们，并将它们提供给下一个块。最后一个块输出的最终隐态被馈送到称为语言建模头的线性层，并且输出 logits 通过 softmax 运算转换为目标词汇上的概率分布。请注意，编码器和解码器块之间的细微差别在于，后者在输入 FFN 层之前还通过交叉注意 (CA) 层与编码器的输出进行交互。

这种二进制结构最初是为机器翻译任务中的序列-到-序列建模而设计的。随后，针对更通用的语言建模目标（例如 MLM 和 CLM）提出了几种变型。BERT 系列 [61, 141] 仅利用带有 MLM 的编码器来增强双向信息，作为判别模型。相反，GPT系列[23,178,179]仅使用带有CLM的解码器，专注于单向生成模型。然而，T5 [181] 和 BART [125] 将每个 NLP 任务视为文本-到-文本的转换，同时利用编码器和解码器。仅解码器生成模型架构最近已成为当前LLM的主要选择。值得注意的例子包括 Llama [221, 222]、OPT [262]、Bloom [236]、GLM [66, 258] 和 Mistral [4, 100] 等。

注意机制 [13] 作为 MHA 层中Transformer 的核心设计，根据输入序列中每个token与其他tokens的相关性来计算其加权表示。而位置编码（PE）是为了保持Transformer中序列的次序信息。

从狭义上讲，KV缓存是一个张量列表，用于存储每个块的注意层中所有先前tokens的 k、v 嵌入，在因果 LLM 的自回归生成过程中使用和更新。如图 (a) 所示，在生成第一个 token 之前，所有 KV 缓存都初始化为空，并在使用 𝐿 query和 𝐿 key进行提升注意计算后填充 𝐿（key，value）对。然后，第一个生成的token也将被视为输入，将整个序列扩展到 𝐿 + 1 个tokens。为了避免冗余计算，实际输入将仅包含最新生成的token，从而派生出一个新的三元组（query，key，value）。但为了进行等效计算，新Q必须参与并应用于所有 𝐿 + 1 个先前的K和V。因此，新的（K，V）必须与存储在 KV 缓存中的过去 𝐿 对进行连接，并将其自身更新到其中，以便下一个生成的token参与。然而，从广义上讲，可以将KV缓存视为LLM的记忆存储，其占用量随着生成tokens的增加而线性增长。这直接导致缺乏有效记忆的限制之一，需要提出增强LLM长期记忆机制的方法。

关于高效注意，第一类方法是优化注意机制，特别是作为 Transformer 计算瓶颈的核操作。这种方法在预训练阶段直接增加超参𝐿𝑚𝑎𝑥 来扩展推理过程中 LLM 的上下文长度边界。这些方法分为五种不同的策略，每种策略都有一个特定的重点：局部注意、分层注意、稀疏注意、近似注意和 IO- 觉察注意。

1 局部注意

传统的注意机制的特点是其全局性和充分的注意性质，其中每个token都应该关注其他每个token，从而导致时间和空间复杂度呈二次方。考虑到局部上下文在某些应用中的重要性[246]，近年来已经引入了各种方法来实现局部注意机制。这些机制将每个token的注意仅限于其相邻的token。这些方法的变体，源于确定什么是token邻居的启发式方法，如图所示。

受卷积神经网络 (CNN) [116, 122] 的启发，另一种方法是使用滑动窗口技术，如 Longformer [17] 中所示。在这种方法中，每个token被分配一个连续的固定窗口，并且只允许关注之前相邻的 𝑤 ≪ 𝐿 tokens 作为其邻居。

ETC [6]和LongT5 [78]采用全局-局部注意机制，它们显式或隐式地构造辅助全局tokens，用全局注意表示片段信息，同时仅将局部注意应用于源tokens。为了避免调整，LongLM [103] 对窗口外tokens采用分组注意，对相邻窗口内的tokens采用标准注意。

除了直接位置邻接之外，Reformer [110]还利用基于kNN和LSH算法的一个邻居token选择机制 [96]。

2 分层注意

全局token技术[6,17,78]和块间注意[198]本质上是在自注意中引入层次特征，补偿来自高层注意的更多全局信息，同时保持来自低层局部注意较低的计算成本。从这个角度来看，更多的工作探索各种分层机制，将结构化层次结构引入自注意中，利用更高层的全局信息和低层的局部注意来实现多尺度的上下文感受野。

一般来说，分层注意机制，源自自然语言中与局部注意相同的上下文局部性原理。然而，它们采用了更复杂的结构，通常是启发式设计的，在捕获远程上下文依赖关系和保持低级计算效率之间取得平衡。

3 稀疏注意

虽然一些方法引入启发式方法来实现自注意内的局部性和分层结构，但另一个方向探索全注意矩阵中固有的稀疏模式 [43, 254]。这些方法旨在引入稀疏注意掩码，表示为 𝑀S，其中每行 𝑖 分配第 𝑖 个token所关注的一组稀疏索引 S𝑖 ⊆ {𝑗|𝑗 < 𝑖}。这些基于稀疏性的注意机制提供计算效率和捕获全局上下文信息的能力[151]。如图提供这些稀疏注意机制的可视化。

一些方法不是仅根据位置设置固定的稀疏索引，而是以可学习的方式自适应地寻求稀疏性，同时考虑嵌入值。

其他一些方法将完全注意视为全连接图，其中节点表示每个token的嵌入，边表示通过注意的连接。这些方法将稀疏性视为图稀疏化问题。

4 近似注意

除了限制完全注意计算的启发式方法之外，一些研究还基于具有线性复杂度的注意矩阵稀疏性或低秩来近似注意，尽管是以精度为代价的。

Linformer [229] 采用SVD近似注意矩阵为一个低秩矩阵。Luna [144] 将注意核解耦为两种嵌套的注意方法，二者均具有 L 的线性复杂度。除了低秩先验之外，有些工作是基于广义可核化注意，将其视为kernel machine来提高注意效率，催生进一步的核策略，例如Fourier注意[161]和基础（primal）注意[39]。

此外，受 Robust-PCA [28] 的启发，结合基于 LSH 的稀疏矩阵𝑆（如 Reformer 的 [110]）和具有随机特征图（如 Performer 的 [110]）的低秩核分解，提供了更准确且有效的近似。

5 IO-觉察注意

上述所有追求高效注意的方法都可以被认为是基于注意矩阵和 NLP 任务的一些理论或经验特性，包括局部性、稀疏性、低秩性和其他启发式或数学技巧，以高注意质量换取低计算复杂度。相比之下， IO -觉察注意机制整体代表了考虑内存瓶颈来优化注意计算的努力，同时保持注意核计算的准确性。

Flash Attention [54, 55] 在 GPU 高带宽内存 (HBM) 和 GPU 片上 SRAM 之间保证 IO 觉察，成功地减少时间和内存消耗，同时应用精确的注意计算。它已经得到广泛采用，并已直接纳入 Pytorch v2.0 [175]。更具体地说，对于前向传输，它利用平铺（tiling）技术[244]将大型softmax注意分解为更小的块，从HBM[99]逐块加载到SRAM，在片上执行所有注意计算步骤以减少 HBM 访问，并利用在线 softmax 技术 [156] 通过重新缩放逐步将输出更新回 HBM。

虽然 Flash Attention 有效解决了训练内存瓶颈，但 LLM 仍然面临与推理过程中 KV 缓存的内存消耗相关的挑战，该内存消耗会随着批量请求而动态增长。认识到碎片和冗余造成的内存浪费，vLLM 提出 Paged Attention 机制 [119]。该技术从虚拟内存操作系统 [57] 中的内存分页（memory paging）技术 [107] 中汲取灵感，有效管理 KV 缓存内存以最大程度地减少浪费，并允许在批请求中灵活共享。

由于上下文的工作记忆，Transformer 架构经常难以捕获长期依赖关系。有两种主要途径来应对这一挑战，同时又不损害全注意机制的优势。受 RNN 的启发，一些人通过合并可通过注意层访问的内部记忆缓存，将循环机制引入注意中。

这种方法使模型能够在较长的序列上维护和检索信息，从而弥补内置长期记忆的固有缺陷。另一种方法涉及利用现有模型作为外部知识库（例如特定文档或数据集）的接口。在推理过程中，模型可以从这些知识库中读取内容以丰富其上下文输入，并从用户的响应中写入这些知识库以刷新其长期记忆。通过以这种方式集成外部知识，模型可以获得更广泛的上下文，从而增强其有效处理长期依赖关系的能力。

1 内部MemoryCache

回顾自然语言表示的时间性，引入基于递归机制的内部MemoryCache的概念。它将长文本分成固定长度的片段流，并通过更多上下文信息 𝐾𝑛𝑡 , 𝑉𝑛𝑡 增强第 𝑛 层中当前第 𝑡 片段的查询 𝑄𝑡𝑛。该上下文信息是从先前片段的缓存或提取信息中获得的，存储在MemoryCache中。假设每个段具有相同的长度𝑙，并且模型由𝑁层Transformer块组成。运算符[°]表示沿长度维度的串联操作。值得注意的是，MemoryCache中的变量通常与计算图分离，消除了梯度计算的需要。

分段递归最初从 Transformer-XL [53] 引入到 Transformer 中。它缓存最后一层𝑚个之前连续段的输出，并将它们连接到当前层的当前段中，扩展当前查询的上下文。

类似于Feedback Transformer[69]，ERNIE-Doc[64]提出一种增强的递归机制，连接同一层中先前段的输出隐态而不是最后一层，进行直接替换。通过这种方式，不仅可以隐式扩展最大有效上下文长度，而且还可以利用过去的高层表示来丰富未来的低层表示。此外，它采用了回馈机制，馈送片段两次，其中第一次仅浏览每个片段，而第二次回顾以实现双向信息流，这类似于 READTWICE [257] 中的机制。

一种连续信号表示允许使用固定内存存储进行无界上下文表示，与上下文长度无关，类似于 LSTM。然而，由于MemoryCaches被存储为连续信号，它不能简单地添加到当前段之前，而是必须通过连续注意机制转换回嵌入。

RMT [25] 将Memory Cache形式化为特殊的 [mem] tokens，在每个段的开头和结尾处都预先考虑。处理完每个段后，读/写tokens将从输出嵌入中分离出来，写tokens将作为下一个段的 [mem]tokens。利用这种带有全局记忆tokens的循环机制，RMT 被证明可以将有效上下文大小扩展到 1M tokens [26]。Memorizing Transformer [238]仅将（Key，Value）Memory Caches应用于顶部注意层，但具有较大的缓存大小且未压缩。此外，不是使用简单的 FIFO 缓存来读取内存，而是使用 k-NN 算法来检索每个查询的前 k 个最相似的（key，value）对，将其添加到本地查询中。相比之下，Memformer [237] 充分利用自注意的变体和遗忘机制来读取和写入Memory Cache，通过长范围的时间步骤，检索和保留最重要的信息。

2 外部记忆库

所讨论的机制通过从内部MemoryCache中添加额外的隐态来增强普通无状态 Transformer 的顺序循环。然而，它们也存在缺点。首先，记忆机制的微小变化可能需要完整的模型重训练，未充分利用已经熟练捕获上下文依赖关系的预训练 LLM，尽管时间不够长。其次，正如 Memorizing Transformer [238] 中所指出的，它们经常面临内存泄气的挑战，其中内存缓存中较旧的隐态在训练过程中可能与最新的隐态分布不同，从而阻碍内存增强的有效性。

作为解决方案，另一种检索增强机制RAG[126]，用上下文信息编码器将长序列作为分段嵌入向量存储到外部记忆库，将模型与其长期记忆存储解耦 。在查询期间，模型根据特定标准从该记忆库中检索信息，并将其动态连接以形成上下文工作记忆。

关于位置编码外展，认识到需要将推理长度 𝐿𝑚𝑎𝑥 外推，研究界在这个方向上做出了重大努力。值得注意的是，根据[8]，干扰因素是在parity任务中长度泛化失败的主要原因。然而，这些问题可以通过scratchpad提示等方法得到很大程度的缓解[163]。

1 增强理解

在 Transformers 时代之前，基于 RNN 的语言模型是在较短的序列上进行训练的，但预计可以有效地泛化到较长的上下文，这种现象称为长度外推或长度泛化 [154, 155]。不幸的是，最近的研究 [8,35,173] 强调了基于 Transformer 语言模型的长度外推能力存在显着缺陷。当应用于实际应用程序时，这会导致推理过程中上下文长度范围不足。

在最初的 Transformer 论文 [225] 中，几乎没有讨论其 SinPE 的设计见解或理论解释。这导致许多研究人员质疑其必要性和有效性，尤其是对外推缺陷的指责，这也指向相同基于三角几何的 RoPE [207]。为了理解当前三角 PE 引起的糟糕外推，有以下两个见解。

1. 从数学角度来看，正如 Su [206] 在他的博客中解释的那样，外推法（涉及从局部信息推断整体）取决于函数的高阶平滑度。然而，这些 PE 被设计为高频振荡三角基函数的组合，以容纳足够的位置信息。这种选择使得模型在训练阶段没有特定学习的情况下很难进行泛化。

2. 从训练的角度来看，由于基函数的波长或周期呈指数增加，和 {𝛽𝑖 }成比例，受当前支持𝐿𝑚𝑎𝑥 约束的训练样本对后部低频维度来说通常太短，无法跨越整个周期。这表明只有少数几个维度感知完整的周期性信息，从而接受足够的外推训练，并且边界在[140]被定义为临界维度（例如，对于Llama2-4k [222]，临界维度仅为92）。因此，当依赖于这些学习不佳的低频分量时，直接外推法很容易失败。