Long-Context LLM综述

前言

题目： Beyond the Limits: A Survey of Techniques to Extend the Context Length in Large Language Models
论文地址：Beyond the Limits: A Survey of Techniques to Extend the Context Length in Large Language Models

1. Preliminary

long context，长上下文技术通过增加模型直接处理的文本长度来维持更多的上下文信息，需要通过模型训练来逐步拉升大模型能够接纳的输入文本长度。一般来说，把接受4K-8K输入token的LLM，算作普通的LLM。能够接受10K~200K甚至数百万的LLM，叫做长上下文大模型。

为什么需要long-context LLM？在许多现实场景中，例如文档摘要和多回合对话，llm需要理解和生成长序列，以便在推理阶段准确地执行任务。这些上下文序列通常比LLM所接受的训练要长得多，这强调了LLM必须具备处理长序列的能力，否则模型的困惑度（perplexity）将急剧上升。然而，序列长度增加不仅增加了计算成本，而且内存需求经常超过高级gpu的容量，从而阻碍了有效的训练。

潜在的应用场景：

1. 长文档摘要：较短的窗口限制了模型对更多细节的疏忽，导致无法对文档进行透彻理解。而使用较长的上下文窗口被证明有助于模型更好地理解细节，减轻歧义。
2. QA任务：当面对多回合对话时，上下文窗口的扩展被证明有助于在连续对话中促进对话主题的连贯跟踪。
3. 翻译：多义性词汇的翻译在翻译领域是一个难以解决的问题。更长的上下文窗口在多义性词汇的语境化中起到了明显的帮助作用。此外，当面对很多技术术语时，长上下文模型往往表现出更高的效率，特别是在适应特定领域的上下文细微差别方面。
4. 指代消解：通过分析文本上下文和语义信息，确定代词、名词短语等所指代的具体对象或实体的过程。指代消解的过程需要建立代词与其先行词之间的联系。更长的上下文窗口促进了对信息的更全面的评估，从而通过包含远参考和上下文相关的细节来帮助精确的代词解析。
5. 对话系统：在对话中，当面临一些幽默、讽刺或微妙的表达时，扩展上下文窗口被证明可以帮助LLM更好理解这些信息。

2. 方法分类

主要分为5类：

1. Length Extrapolation，长度外推，即"Train Short, Test Long"。如位置编码外推和上下文窗口分割。
2. Attention Approximation，注意力近似。旨在降低注意力依赖计算的复杂度，比如低秩分解和稀疏注意力。降低复杂度后，在有限资源的情况下可以输入更多的token。
3. Attention-free Transformer，在不依赖于传统注意力机制的情况下提供token之间依赖信息的计算方法，如基于RNN的状态空间模型SSM，最近比较火的Mamba。
4. 模型压缩：压缩模型减小参数量和计算量，进而可以处理更长的上下文。
5. 硬件感知的Transformer：主要从I/O、资源管理和多设备等方面提高计算效率，如FlashAttention。

总结：除了Length Extrapolation以外，其他方法都是基于提高计算效率、减少计算量和降低复杂度等角度出发。

3. 长度外推

只用短序列语料对模型进行训练，就可以得到一个能够处理和预测长序列的模型，即“Train Short, Test Long”。

主要分为三种方法：位置外推和插值(Positional Extrapolation and Interpolation)、上下文窗口分割与滑动(Context Window Segmentation and Sliding)和提示压缩(Prompt Compression)。

3.1 位置外推和插值

位置外推和插值指的是调整与输入token相关的位置嵌入(PE)的技术，从而修改这些token在模型体系结构中的定位和解释方式。位置编码赋予输入token一种重要的意义，使模型能够识别序列中每个token的特定位置。这确保了模型可以有效地捕获和利用输入数据中固有的顺序信息。

原始Transformer中经典的正余弦位置编码：

对于pos+k位置的位置向量某一维2i或2i+1而言，可以表示为pos位置与k位置的位置向量的2i与2i+1维的线性组合。

如何实现长度外推？如果推理时的长度大于训练时的长度，那么就会出现从未见过的位置编码，虽然该编码可以由之前的编码线性组合得到，但实验证明效果会大打折扣，泛化性不够强。

位置外推

位置外推能力是指模型处理超过其训练序列长度的输入序列的能力，从而能够在扩展序列上保留上下文和一致性。一些经典的里程碑工作如下。

（1）[Train short, test long: Attention with linear biases enables input length extrapolation. In ICLR, 2022.Alibi]

线性偏差注意力ALiBi，属于研究长度外推的开篇工作。方法十分简单，在计算 attention score的时候，会对之前位置的分数按照与当前位置的差距进行不同程度的惩罚，如下图所示。
有效性解释：

• 传统的绝对位置编码（如正弦/余弦编码）在训练时会为每个位置分配一个固定的向量，模型可能会过度拟合这些特定长度的模式。而ALiBi通过在注意力分数计算中直接使用线性偏置，减少了模型对特定序列长度的依赖，从而提高了对未见过的序列长度的泛化能力。
• ALiBi的位置偏差随距离线性增长，这种设计让模型在处理不同长度的序列时，可以自然地根据距离调整注意力权重，无需显式学习位置编码的复杂周期性结构。
• 由于线性偏置的引入直接与序列中元素的位置相关，没有固定大小的编码矩阵限制，理论上模型可以更容易地处理任意长度的序列，从而展现出良好的长度外推性能。
• 通过直接在注意力分数上施加与距离相关的线性惩罚，ALiBi鼓励模型关注更接近的词对，同时不对远处的依赖完全排除，从而在一定程度上平衡了局部和全局依赖的学习，这对于处理长序列尤其有利。

（2）旋转位置编码RoPE [Roformer: Enhanced transformer with rotary position embedding]

RoPE几乎是目前常见的流行大模型的标配位置编码方式。使用旋转矩阵对绝对位置进行编码，同时在自注意力公式中结合了明确的相对位置依赖性。旋转位置嵌入保持了序列长度的灵活性、随相对距离的增加而衰减的token间依赖性，以及为线性自注意力配备的相对位置编码的能力。

相关原理：


简单总结RoPE的self-attention流程：首先计算query矩阵和key矩阵，接着将每个位置的query和key向量的元素按照两两一组应用旋转变换。以位置m的query向量为例，其应用旋转变换后变为：
其中${\theta }_j=10000^{-2j/d}$。接着对位置n的key向量进行同样的操作，最后求它们的内积也就是注意力分数：

最后再利用注意力分数对value矩阵进行加权。需要注意的是，value向量不需要进行旋转变换。

最终计算得到内积注意力分数中包含两个位置的相对位置信息m - n，并且这种分数随着m - n的增大而减小，因此可以建模相对位置信息。

（3）xPOS [A length-extrapolatable transformer. In ACL, 2023]

xPOS在旋转位置编码的基础上做了一些改进。其在旋转角度向量的每个维度上都包含了独特的指数衰减因子，从而改进了长度外推。

一些发现：事实上，通过对比现有的一些位置编码的外推效果，可以发现相对位置编码的长度外推性，平均好于绝对位置编码的；像RoPE这样的函数式相对位置编码，又会比训练式相对位置编码的外推效果好些。

插值

为了解决模型处理更长序列时的问题，研究者们提出了一个策略：将位置索引本身进行下采样或缩放。具体来说，可以将原本在[0, 4096]范围内索引缩放到[0, 2048]的范围内。通过这种方式，所有的位置索引都被映射回了模型预训练时的范围内，从而帮助模型更好地处理这些原本超出其处理能力的输入序列。一些里程碑式的研究如下所示。

（1）[Extending context window of large language models via positional interpolation]

这篇文章中提出了位置插值 (Position Interpolation ，PI)的概念，可以将基于RoPE的预训练LLM的上下文窗口大小扩展到 32768，并只需要进行最小步数的微调（1000 步以内）。

PI将RoPE中的映射函数变为：

该变换将更大的位置索引范围$[0,L^{{}^{\prime }})$缩减至原始的索引范围$[0,L)$。例如训练时长度为2048，当推理长度为4096时，位置0被映射为0, 1被映射成0.5, 2被映射为1，4096被映射为2048。

结合下式：

可以发现，插值实际上是对旋转角度进行了缩放。

（2）[LongRoPE: Extending LLM Context Window Beyond 2 Million Tokens]

LongRoPE中发现，有效的位置编码插值应考虑两种非均匀性：不同的RoPE维度和不同的token位置。低维和初始token位置存储着关键信息，因此需要进行更少程度的插值。相比之下，高维存储的信息相对较为稀疏，可进行较大程度的插值。

LongRoPE具体做法：搜索RoPE每个维度以及不同token位置的旋转角度缩放因子，有效地保留了原始 RoPE 位置编码中的信息。这种方法最大程度地减小了位置插值引起的信息损失，从而为微调提供了更好的初始化。

具体来讲，首先在预训练的大模型上搜索256K上下文窗口对应的位置编码插值方案，并在此长度下进行微调。其次，由于LongRoPE的非均匀插值允许在不微调的情况下进行8倍扩展，所以接着对已扩展的微调后的大模型进行了二次非均匀插值搜索，最终达到了2048K上下文窗口。在将上下文窗口扩展到极长的2048K后，文章中发现原始上下文窗口内的性能出现了下降。这是因为它导致原始上下文窗口内的位置被压缩在更窄的区域内，从而对大模型的性能产生了负面影响。为此，LongRoPE在扩展后的大模型上对8K长度内的RoPE缩放因子进行了重新搜索，旨在引导其在较短长度上进行较少的位置插值，来恢复短上下文窗口的性能。

3.2 上下文窗口分割与滑动

具体思想：将输入划分为段，并应用滑动窗口方法来管理上下文。

（1）[Structured prompting: Scaling in-context learning to 1,000 examples.]和[Parallel context windows for large language models. In ACL, 2023]

二者方法类似：将长文本分割成多个连续的块或窗口，每个窗口都在模型的最大上下文长度范围内。然后，模型可以独立处理这些窗口，然后采用某种方式结合相邻窗口的信息以捕捉跨越窗口边界的依赖关系。

PCW将LLM的输入分为上下文token（上下文文档或者检索到的文档）和任务token（要分类的句子或者问题本身）。如果原始LLM长度为N，任务token长度为T，那么留给上下文的窗口长度就只有C=N-T。然后将长的上下文分成了B块，每块长度都为C，总的长度为BC+T，PCW重新编码如下：

即任务token占据最后的T的位置，B个窗口占据前边大小为N-T的位置。注意力计算：B个窗口内各自执行自回归，然后将结果进行拼接。解码时任务token和所有上下文中的token都计算注意力。

（2）StreamingLLM [Efficient streaming language models with attention sinks]

文中观察到大量的注意力得分被分配给初始标token，即使它们与语言建模任务的相关性不高。这一观察结果表明，通过合并窗口上下文和第一个（论文中为前4个）token，可以得到一个有效的框架，这使得用有限长度的注意窗口训练的llm能够有效地推广到无限长的序列长度，而无需任何微调。

3.3 提示压缩

prompt压缩是指在保留重要信息的同时缩短原始提示的方法。prompt压缩过程通过压缩大量prompt输入或通过学习获得prompt的简明表示来加速LLM输入的处理。

（1）[LLMLingua: Compressing prompts for accelerated inference of large language models. In EMNLP, 2023]

采用如GPT2-small或LLaMA-7B，来识别和消除提示符中的多余标记。这种方法可以实现高达20倍的压缩比，同时最大限度地减少任何可识别的性能下降。

（2）[Compressing context to enhance inference efficiency of large language models. In EMNLP, 2023]

引入了一种称为selective context的新方法，旨在通过系统地识别和修剪输入上下文中的冗余来提高llm的推理效率。主要目标是简化输入，使其更紧凑，从而优化语言模型推断的整体效率。

4. 注意力近似

注意力近似，即Attention Approximation，其目的是降低自注意力的计算和记忆复杂度$O(n^2)$。主要分为：低秩分解(Low-rank Decomposition)、稀疏注意力(Sparse Attention)、Softmax-free Attention。

4.1 低秩分解

低秩分解方法通过减少矩阵中参数的个数，提高了注意力计算的效率。

（1）线性编码器-解码器(Linear encoder - decoder, LED) [Lightweight and efficient end-to-end speech recognition using low-rank transformer. In ICASSP, 2020]

通过在内积计算之前将三个矩阵分解成更小的矩阵，逼近线性参数效率。

（2）Linformer [Linformer: Self-attention with linear complexity]

引入了另一种线性投影机制，即在K和V之前添加两个较小的矩阵，将其投影到较小的尺寸，同时保持Q不变。这两种方法都通过线性逼近优化矩阵计算。

4.2 稀疏注意力

在计算内积时，每个Token并不会与所有其他Token计算内积，进而造成稀疏性，减少计算量，扩大上下文窗口长度。

相关参考文献：

1. Guo Q, Qiu X, Liu P, et al. Star-transformer.
2. Beltagy I, Peters M E, Cohan A. Longformer: The long-document transformer.
3. Ainslie J, Ontanon S, Alberti C, et al. Etc: Encoding long and structured inputs in transformers.
4. Oord A v d, Li Y, Vinyals O. Representation learning with contrastive predictive coding.

4.3 Softmax-free Attention

softmax在时间和空间上都有着较高复杂度，替换softmax操作有可能降低计算复杂性，为提高处理大量token序列的效率铺平道路。这类方法被称为softmax-free attention。

（1）CosFormer [cosformer: Rethinking softmax in attention]通过一个加权余弦距离的线性算子模拟softmax行为。
（2）SOFT [Soft: Softmaxfree transformer with linear complexity. NeurIPS, 2021]采用高斯核函数代替softmax。
（3）SIMA [Sima: Simple softmax-free attention for vision transformers. In WACV, 2024]选择使用简单的L1范数对Q和K矩阵进行规范化。
（4）[A study on relu and softmax in transformer. arXiv:2302.06461, 2023]用ReLU函数代替softmax进行归一化，证明这种替代在保持性能的同时保持了线性可扩展性。

5. Attention-free Transformers

无注意力Transformer是指在不依赖于传统注意力机制的情况下提供token之间依赖信息的计算方法。这些机制为依赖性计算提供了不同的视角，同时保持了小于二次方的复杂度。

5.1 状态空间模型(State Space Model, SSM)

SSM是一种统计seq2seq模型，它使用隐藏状态的线性投影来计算基于输入序列的输出序列。支持并行训练并优化了推理效率。

（1）[Mamba: Lineartime sequence modeling with selective state spaces]

Mamba能够处理长达百万级别的token序列，这对于需要长期记忆的任务至关重要。Mamba的推理速度极快，比Transformer模型快5倍，这意味着它可以在更短的时间内处理更多的数据。Mamba在序列长度上的扩展性是线性的，这意味着随着序列长度的增加，模型的性能不会受到显著影响。

5.2 位置依赖注意力

token间的依赖取决于它们的位置，而不是计算内积等形式的直接交互。

（1）[An attention free transformer]

文中提出了Attention-free Transformer (AFT)，仅使用K和V，同时消除Q及其与K的点积。然后，引入了一种新的可学习矩阵W，它与普通注意中的矩阵不同，其在所有输入序列中保持一致。

（2）RWKV：Rwkv: Reinventing rnns for the transformer era.

6. 模型压缩

模型压缩减小了计算和内存需求，进而可以处理更长的上下文。

1. 通过消除冗余权重来最小化LLM体系结构的大小，从而减少计算和内存需求。
2. 考虑降低计算精度以减轻计算复杂性。
3. 提高内存效率和优化数据存储方法。

7. 硬件感知的Transformer

7.1 I/O

（1）[Flashattention: Fast and memory-efficient exact attention with io-awareness] 利用gpu中不同的流式多处理器(SMs)来并行处理不同的块。
（2）[Blockwise parallel transformers for large context models] 在FlashAttention的基础上，Block-wise Parallel Transformer (BPT)融合了前馈层和自注意力，以进一步减少IO使用，使模型能够处理比FlashAttention长四倍的序列。

7.2 多设备分布式注意力

（1）[Ring attention with blockwise transformers for near-infinite context] Ring attention在多个gpu上分配注意力计算来扩展到非常长的序列。
（2）[Longnet: Scaling transformers to 1,000,000,000 tokens] LongNet具有计算多个注意的能力，每个注意具有不同的扩展稀疏模式。这些计算独立运行，可以分布在多个设备上，每个设备对应一个扩展模式。这种集体方法有利于处理较长的序列。

7.3 内存管理

（1）[Efficient memory management for large language model serving with pagedattention]

PagedAttention在推理阶段采用虚拟内存启发策略来解决顺序生成中固有的内存限制挑战。通过将KV缓存分段为块，这种方法实现了灵活的内存管理，有效地减轻了内部和外部碎片问题。
（2）[Flash-decoding for long-context inference]

Flash-Decoding建立在FlashAttention原理的基础上。为K/V序列长度引入新的并行化维度，即使在小批量大小和大上下文长度的情况下，它也能确保最佳的GPU利用率。