大语言模型100K上下文窗口的秘诀

上下文窗口（context window）是指语言模型在进行预测或生成文本时，所考虑的前一个token或文本片段的大小范围。

在语言模型中，上下文窗口对于理解和生成与特定上下文相关的文本至关重要。较大的上下文窗口可以提供更丰富的语义信息、消除歧义、处理上下文依赖性，并帮助模型生成连贯、准确的文本，还能更好地捕捉语言的上下文相关性，使得模型能够根据前文来做出更准确的预测或生成。

最新发布的语言大模型的上下文窗口越来越大。本文详细探讨了大型上下文窗口的技术可能性，尤其分析了将上下文长度增加到100K背后的六大优化技巧。

最近有几个新的语言大模型（LLM）发布，这些模型可以使用非常大的上下文窗口，例如65K 个tokens（MosaicML的MPT-7B-StoryWriter-65k+）和100K个tokens的上下文窗口（Antropic）。在Palm-2技术报告中，谷歌并没有透露具体上下文大小，但表示他们“显著增加了模型的上下文长度”。

相比之下，当前GPT-4模型可以使用32K个输入tokens的上下文长度，而大多数开源LLM的上下文长度为2K个tokens。

如此大的上下文长度意味着提示（prompt）可以达到一本书的大小。《了不起的盖茨比》有72K个tokens，210页，按1.7分钟/页的阅读速度计算，需要6小时的阅读时间。因此，模型可以扫描并保留此数量的“自定义”信息来处理查询！

我想要弄清楚大型上下文窗口的技术可能性。本文搜集了一些零散信息，内容如下：

为何上下文长度如此重要，且能在LLM中起到举足轻重的作用？
处理大型上下文长度时，原始Transformer架构的主要局限性是什么？
Transformer架构的计算复杂度
目前有哪些可以加速Transformer并将上下文长度增加到100K的优化技术？

重点概览

本文将互换使用“上下文长度”、“上下文窗口”和“输入tokens数量”，并用 $n$ 来表示。

鉴于文章内容较长，以下为要点总结：

第一个问题是注意力层（attention layer）计算的二次方时间（Quadratic time）和空间复杂度，即输入tokens数量 $n$ 。
当embedding大小 $d>n$ 时，第二个问题是embedding大小 $d$ 的线性层的二次方时间复杂度。
第三个问题是原始架构中使用的位置正弦嵌入（Positional Sinusoidal Embedding ）。
在Transformer架构中，可学习（learnable）矩阵权重的形状与输入tokens个数 $n$ 无关。
因此，在2K上下文长度中训练的Transformer可以使用任意长度的tokens，甚至是100K个tokens。但如果不是在100K个tokens上训练出来的，那么该模型在100K个tokens的推理过程中不会产生有意义的推理结果。
由于 $n$ 、 $d$ 相关的二次复杂度，在大型语料库上训练Vanilla Transformer，并且只在较大的上下文长度上训练是不可行的。据估计，在2K上下文长度上训练LLaMA的费用约为300万美元，因此，100K的花费约为1.5亿美元。
一种选择是，可以在2K个tokens上下文中训练模型，然后在更长的上下文tokens（例如65K）中微调。但由于位置正弦编码（Positional Sinusoidal Encoding）的存在，这不适用于原始Transformer模型。
[技巧1] 为解决此问题，可删除位置正弦编码并使用ALiBi，参考苏神博客，这一简单位置嵌入不会影响准确性。然后可以在2K个tokens上训练，在100K个tokens上微调。
[技巧2] 无需计算所有tokens间的注意力分数（attention scores）。某些tokens比其他tokens更重要，因此可使用稀疏注意力。这将提升训练和推理速度。
[技巧3] Flash Attention有效地实现了GPU的注意力层。它使用切片（tiling）技术，避免生成不适合GPU SRAM容量的大型中间矩阵 $（n，n）$ 。这将提升训练和推理速度。
[技巧4] 选择多查询注意力（Multi-Query attention），而非多头注意力。这意味着线性投影 $K$ 和 $V$ 时，可在跨所有注意力头（head）中共享权重。这极大地加快了增量（incremental）推理速度。
[技巧5] 条件计算（Conditional computation）避免将所有模型参数应用于输入序列中的所有tokens。CoLT5仅对最重要的tokens应用重量级计算，并使用较轻量级的层处理其余tokens。这将加速训练和推理。
[技巧6] 为适应大型上下文，需要GPU中有大量RAM，因此人们使用80GB的A100 GPU。

总之，训练和推理速度越快，可使用的上下文长度就越大。

以下是上述要点的详细内容。

为何上下文长度如此重要？

上下文长度是LLM的关键限制之一，将其增加到现在的100K是一项难以置信的成就。

对于语言大模型，其中一个重要用例是人们想要“将大量自定义数据输入LLM”（与公司或特定问题相关的文档，各种异构文本等），并询问有关此特定数据的问题，而不是LLM在训练期间接入一些来自互联网的抽象数据。

为了克服这一局限性，人们做了以下尝试：

尝试总结技巧和复杂的链式提示。
维护向量数据库以保留自定义文档的嵌入，然后通过相似性指标在它们之间展开“搜索”。
尽可能使用自定义数据微调LLM（并非所有商业LLM都允许自定义微调，对开源LLM进行自定义微调并不常见）。
为特定数据开发定制小型LLM（同样，这并非常规任务）

较大的上下文长度能够让已经很强大的LLM（用整个互联网数据训练的模型）查询用户的上下文和数据，以更强的个性化在完全不同的层面与你交互。所有这些都无需更改模型权重并能够“在内存中”即时“训练”。

总体而言，大型上下文窗口可让模型更加准确、流畅，提升模型创造力。

这就好比是计算机的RAM，操作系统保留了所有应用程序的实时上下文，由于上下文长度充足，LLM可以像“推理计算机”一样，保留大量用户上下文。

原始Transformer和上下文长度

需要注意的是，在Transformer架构中，所有可学习矩阵权重的形状与输入tokens数量 $n$ 无关。所有可训练参数（嵌入查找、投影层、softmax层和注意力层）都不依赖于输入长度，并且必须处理可变长度（variable-length）的输入。该架构具有的开箱即用的特性非常不错。

这意味着，如果你用2K的上下文长度训练了一个Transformer模型，可以对任意大小的tokens序列进行推断，唯一的问题在于，如果模型没有在上下文长度为100K的情况下进行训练，那么它在推断过程中将无法对100K个tokens产出有意义的结果。这种情况下，训练数据的分布与推断过程中的分布相差很远，模型的表现就像任何其他机器学习模型一样，面临失败风险。

为训练具有较大上下文长度Transformer，我们的解决方案是将其分为两个阶段进行训练：首先在2K个tokens的上下文长度上训练基本模型，然后继续在更长的上下文中进行训练（微调），例如65K或100K。MosaicML就采用这种方法。但问题是，原始的Transformer架构无法直接实现这一点，因此需要使用一些技巧（请参阅后文的技巧1）。

多头注意力回顾

大型上下文长度所面临的挑战与Transformer架构的计算复杂度有关。为讨论复杂度，我们首先回顾一下注意力层的工作原理。

$Q$ - 查询（query）， $K$ - 键（key）， $V$ - 值（value），这些符号是论文中与信息检索相关的符号表示法。在信息检索中，你可以将一个“查询”输入系统，并搜索与之最接近的“键”。

$n$ - 输入的tokens数量

$d$ - 文本嵌入维度

$h$ - 注意力头的数量

$k$ - $Q$ 和 $K$ 的线性投影大小

$v$ -$ V$的线性投影大小

多头注意力（Multi-Head Attention）

我们有一个Embedding层，用于接收tokens作为输入，并返回大小为 $（1，d）$ 的向量。因此，对于一个由 $n$ 个tokens组成的序列，我们得到大小为 $（n，d）$ 的文本嵌入矩阵 $X$ ，然后将其与位置正弦嵌入相加。
多头注意力层旨在为tokens序列计算新的嵌入表示，该tokens序列可以被视为对原始文本编码 $X$ ，但需要，（1）根据tokens间相对于上下文的重要性进行加权，（2）根据tokens的相对位置进行加权。
我们使用 $h$ 个注意力头对嵌入矩阵 $X（n×d）$ 进行并行处理。为了使所有的注意力头都得到 $Q$ 、 $K$ 和 $V$ ，我们需要对 $X$ 进行线性投影，将其分别投影到 $k$ 、 $k$ 和 $v$ 维度。为此，可以通过将 $X$ 分别与形状为 $（d，k）$ 、 $（d，k）$ 和 $（d，v）$ 的 $h$ 个矩阵相乘来实现。你可将其理解为用 $（n，d）$ 乘以 $（h，d，k）$ 、 $（h，d，k）$ 和 $（h，d，v）$ 。
注意力头返回大小为 $（n，v）$ 的 $h$ 个注意力分数矩阵。然后，我们将来自所有注意力头 $（n，h*v）$ 的片段进行连接，并对其进行线性投影，为后续步骤做准备。

缩放点积注意力（Scaled Dot-Product Attention）

现在详细讨论一个注意力头。

$Q$ 、 $K$ 、 $V$ 是 $X$ 的3个线性投影，大小分别为 $（n，k）$ 、 $（n，k）$ 和 $（n，v）$ ，通过乘以每个注意力头的可学习权重（learnable weight）获得。
通过计算 $Q$ 和 $K$ （转置）之间的距离（点积），我们得到了注意力分数。将矩阵 $（n，k）$ 与 $（k，n）$ 相乘，得到矩阵 $（n，n）$ ，然后我们将其与掩码矩阵相乘，以将一些tokens置零（在解码器中需要）。接下来，我们对其进行缩放，并应用softmax函数，使注意力分数范围在0到1之间。这样，我们就得到一个形状为 $（n，n）$ 的矩阵，其中 $n_ij$ 表示第 $i$ 个和第 $j$ 个token之间的相对注意力分数（0-1之间），这展示了这些tokens在给定长度为n的特定上下文中有多“接近（close）”。
然后，我们将这个注意力分数矩阵 $（n，n）$ 乘以大小为 $（n，d）$ 的“值（value）” $V$ ，以获得由这些相对注意力分数加权得到的文本嵌入。

下图是Multi-Query注意力论文的代码片段，展示了如何使用批（batching）处理计算多头注意力，并且在每一步都清晰地给出了形状信息。代码里还包括在解码过程中使用的掩码乘法操作。

Transformer的复杂度和上下文长度

2个矩阵乘法 $（a，b）\star（b，c）$ 的复杂度为 $O（a \star b \star c）$ 。

为简单起见，我们假设 $k \star h = O（d）$ ，并利用这个假设来推导注意力机制的复杂度。
注意力层的复杂度由两部分组成：

线性投影得到 $Q，K，V$ ：大小为 $（n，d）$ 的嵌入矩阵乘以 $h$ 个可学习矩阵 $（d，k）$ ， $（d，k）$ 和 $（d，v）$ 。因此，复杂度约为 $O(nd^2)$
将 $Q$ 与变换后的 $K$ 相乘，然后再乘以 $V$ ： $（n，k）\star（k，n）=（n，n）$ ，以及 $（n，n）\star（n，v）=（n，v）$ 。复杂度约为 $O(n^2d)$ 。

因此，注意力层的复杂度为 $O（n^2d + nd^2）$ ，其中 $n$ 是上下文长度（输入tokens的数量）， $d$ 是嵌入大小。从这里我们可以看出，注意力层计算的复杂度与输入tokens数 $n$ 和嵌入大小 $d$ 相关，分别是二次方关系。

当 $d>n$ 时（例如，在LLaMa中， $n=2K$ ， $d=4K$ ）， $O(nd^2)$ 这个术语非常重要。

当 $n>d$ 时（例如，在使用 $n=65K$ 和 $d=4K$ 进行MosaicML训练时）， $O(n^2d)$ 这个术语非常重要。

提醒一下，二次方增长的情况有多糟糕：

$2000^2=4000000, 100000^2=10000000000$

举例说明一下二次方复杂度是如何影响模型训练成本的。LLaMa模型的训练估计价格约为300万美元 ，具有650亿个参数，2K的上下文长度和4K的嵌入大小。预估时间大部分是GPU训练时间。如果我们将上下文长度从2K增加到100K（增加了50倍），训练时间也会增加大约50倍（由于上下文更大，迭代次数较少，但每次迭代的时间更长）。因此，以100K上下文训练LLaMa模型的成本约为1.5亿美元。

对该计算稍作详细说明：

假设tokens数量为n时，注意力的复杂度为 $O(n^2d + nd^2)$ ，需要进行M次迭代来进行训练。如果我们将上下文长度从 $n$ 增加到 $p \star n$ ，由于上下文长度变大，所需的迭代次数将变为 $\frac{M}{p}$ （这里简单假设它是线性的，实际情况可能会高点或低点，具体取决于任务）。现在我们有两个方程式：

（1） $n$ 的复杂度为 $M \star (n^2d + nd^2)$

（2） $p \star n$ 的复杂度为 $\frac{M}{p} \star ((pn)^2d + (pn)d^2)$

经过一系列简化和除法，得到比值 $\frac{(2)}{(1)}$ 的近似为$ \frac{(d + p \star n)}{(d + n)}$。

如果 $d << n$ ，将 $n$ 增加 $p$ 倍将导致迭代次数增加约 $p$ 倍。

如果 $d ~ n$ ，将 $n$ 增加 $p$ 倍将导致迭代次数增加约 $\frac{p}{2}$ 倍。

Transformer训练阶段和推理阶段的区别

在深入研究优化技术之前，最后需要讨论的是训练和推理过程中计算的差异。

在训练过程中，你可以并行计算；而在推理过程生成文本时，你需要按顺序逐步生成，因为下一个token依赖于前面的token。实现推理的直接方式是逐步计算注意力分数，并缓存以前的结果供未来的token使用。

这种区别导致了加速训练和推理具有不同方法。因此，下面的一些技巧既可以优化训练阶段，也可以优化推理阶段，但也有一些只能优化推理阶段。

增加上下文长度的优化技术

接下来谈谈研究人员是如何克服所有这些挑战，并能够训练具有较大上下文长度的语言模型。

[技巧1] 更好的位置编码——ALiBi

为训练具有较大上下文长度Transformer，我们的解决方案是将其分为两个阶段进行训练：首先在2K个tokens的上下文长度上训练基本模型，然后在更长的上下文（例如65K）上进行微调。但是之前我们提到原始的Transformer架构不适用于这种方法，为什么？

**这是因为位置正弦编码没有“外推（extrapolation）”能力。**在ALiBI[4]论文中，作者表明，在推理过程中，位置正弦编码对于上下文窗口的扩展不具有健壮性，在增加了一些tokens后，性能开始下降。因此，缺乏“外推”能力基本上意味着在推理/微调过程中不能使用比训练时更大的上下文长度。关于“外推”的概念和各种位置编码的比较详见[4]。

在原始Transformer论文中，位置正弦嵌入与底层架构中的tokens嵌入相加，以添加关于单词顺序的信息。如果你想了解位置正弦嵌入的计算方式，推荐观看这个视频，其中对其进行了直观且详细的解释。

因此，第一个技巧是移除位置正弦嵌入，并由另一种位置嵌入来替代，即线性偏置注意力（ALiBI）。

它应用于注意力头部（而非网络底部），并通过与其距离成比例的惩罚来偏置查询键的注意力分数（在softmax之前）。

这一技巧能够加速训练进程。

计算每个注意力头的注意力分数时，ALiBi为每个注意力分数 $(q_i \dot k_j，左侧)$ 添加了一个常数偏置（右侧）。与未修改的注意力子层一样，之后对这些分数用softmax函数进行转化，其余计算保持不变。m是一个特定于注意力头的标量，在训练期间为定值，且不进行学习。（摘自ALiBi论文）

[技巧2] 稀疏注意力机制

在大小为100K的上下文中，并非所有tokens之间都存在相关性。为了减少计算量，一种方法是在计算注意力分数时仅考虑部分tokens。添加稀疏性的目的是使计算复杂度与 $n$ 呈线性关系，而非二次方关系。有多种方法可以选择tokens之间的连接方式，这篇Google博客文章中有出色的示例。

例如，滑动窗口注意力（Sliding Window Attention ，也称局部注意力）在每个token周围采用了固定大小的窗口注意力。在这一注意力机制中，给定一个固定的窗口大小 $w$ ，每个token会关注其两侧的 $\frac{w}{2}$ 个tokens。这种注意力机制的计算复杂度为 $O(n \star w)$ ，与输入序列的长度 $n$ 成线性关系。为提高计算效率， $w$ 应相对于 $n$ 较小。技巧在于注意力信息在相邻的tokens中“流动（flows）”，近似完全的图。

BigBird注意力分数方法结合了全局、局部和随机机制。在这篇论文中，作者展示了一个重要的观察结果，即在计算相似性分数和不同节点间的信息流动之间存在固有的张力（tension）关系（即一个token对其他tokens的影响能力）。

这一技巧可加快训练和推理。

[技巧3] FlashAttention——用于GPU的注意力层高效实现

在注意力层中，有几个计算操作会反复执行：

$S = Q \star K$
$P = softmax(S)$
$O = P \star V$

请记住 $P$ 、 $S$ 和 $O$ 结果的概念，稍后将用到。FlashAttention的作者“融合”了这些操作：他们实现了一个能有效利用GPU内存，并计算准确注意力的注意力层算法。

为使GPU执行一个运算，输入数据必须在名为SRAM的“快速（quick）”内存中。数据从“慢速”的HBM（高带宽内存）复制到SRAM中，并在计算完成后返回到HBM。SRAM内存的速度比HBM快得多，但容量小得多（例如，A100 40GB GPU中的SRAM为20MB，而HBM为40GB）。

因此，访问HBM的运算成本很高。

就GPU内存利用而言，注意力层面临的主要问题是“中间（intermediate）”乘法结果** $P$ 、 $S$ 和 $O$ 的大小 $(n，n)$ **，需要将它们保存至HBM中，并在注意力运算之间再次读取。将 $P$ 、 $S$ 和 $O$ 从HBM移动到SRAM，以及反向移动是瓶颈所在，作者在论文中解决了这一问题。

FlashAttention算法的主要思路是将输入的 $Q$ 、 $K$ 和 $V$ 矩阵划分成块（block），将这些块从HBM加载至SRAM中，然后根据这些块来计算注意力输出。这个过程被称为“切片（tiling）”。

左图：FlashAttention使用切片技术，防止将大型n \time n注意力矩阵（虚线框内）存储到HBM中。在外部循环（红色箭头）中，FlashAttention循环遍历 $K$ 和 $V$ 矩阵的块，并将它们加载到SRAM中。在每个块中，FlashAttention循环遍历 $Q$ 矩阵的块（蓝色箭头），将它们加载到SRAM中，并将注意力计算的输出写回至HBM。