如何突破技术瓶颈，扩展LLM的上下文长度？

语言模型的上下文长度现在已经达到了惊人的100k级别，我希望理解这在技术上如何实现，所以在这篇文章中，我收集了一些相关的信息，涵盖了以下内容：

为什么上下文长度很重要
目前有哪些优化技术可以加速Transformer并将上下文长度增加到100K

文章导航

一、为什么上下文长度很重要

上下文长度是大型语言模型（LLMs）的关键限制之一。而将上下文长度增加到100K是一个了不起的成就。

人们应用LLMs的一个重要场景是“将大量自定义数据输入到LLM中”（与公司或特定问题相关的文件，各种异构文本等），并询问关于这些特定数据的问题，而不是LLM在训练期间从互联网上看到的一些抽象数据。

为了克服这个限制，人们现在做了各种事情：

尝试摘要技术和复杂的链式提示；
维护向量数据库以保存自定义文档的嵌入，然后通过某些相似性度量进行“搜索”；
在可能的情况下使用自定义数据微调LLM（并非所有商业LLM都允许这样做，这对开源LLM来说也不是一项简单的任务）；
为这些特定数据开发自定义的小型LLM（同样，这也不是一项简单的任务）。

拥有较大的上下文长度可以让一个已经很强大的LLM分析你的上下文和数据，并在一个完全不同的层次上与你互动，实现更高的个性化。而且所有这些都不需要改变模型的权重，并且可以在上下文中进行“训练”(Few-shot learning)。总体而言，大的上下文窗口为模型带来了更多的准确性、流畅性和创造力。

一个类比可能是计算机的RAM，操作系统在其中保持所有应用程序的实时上下文。具有大量上下文长度的LLM可以像一个“推理计算机”一样，保持大量的用户上下文。

二、有哪些技巧可以增大上下文窗口

在下文中，我们使用n表示“上下文长度”。

注意力层计算的时间和空间复杂度相对于输入标记数n是平方级的。
当embedding尺寸d > n时，第二个问题是线性层相对于embedding尺寸d的时间复杂度是平方级的。
第三个问题是原始架构中使用的正余弦位置编码。
在Transformer架构中，可学习矩阵权重的形状与输入标记数n无关。因此，训练过的Transformer在2K上下文长度时可以处理任何长度的标记，甚至是100K。但如果模型没有在100K的上下文长度上进行训练，它在推理时不会产生有意义的结果。
由于相对于n和d的平方复杂度，在巨大的语料库上仅针对大上下文长度训练原始Transformer是不可行的。估计在2K上下文长度上训练LLaMA的费用约为300万美元。因此，在100K上下文长度上训练LLaMA的费用约为1.5亿美元。
一种选择是在2K标记上下文上训练模型，然后在更长的上下文上进行微调（例如65K）。但由于正余弦位置编码，这在原始Transformer中不起作用。
- [技巧 #1] 为了解决这个问题，移除正余弦位置编码，使用ALiBi或者ROPE，一种简单而优雅的位置嵌入，不会影响准确性。然后你可以在2K上进行训练并在100K上进行微调。
- [技巧 #2] 你不需要计算所有标记之间的注意力得分。有些标记比其他标记更重要，因此可以使用稀疏注意力。它将加速训练和推理。
- [技巧 #3] Flash Attention有效地实现了GPU上的注意力层。它使用tiling并避免了不适合GPU SRAM的大中间矩阵（n, n）的实现。它将加速训练和推理。
- [技巧 #4] MQA代替MHA。这意味着在线性投影K和V时，你在所有head之间共享权重。它显著加速增量推理。
- [技巧 #5]重计算避免将所有模型参数应用于输入序列中的所有标记。它将加速训练和推理。
- [技巧 #6] 要适应大上下文，你需要大量的GPU RAM，所以人们使用80GB A100 GPU。