如何在推理阶段确保模型能处理远超预训练时的文本长度，已成为当前大型模型面临的核心问题之一，我们将此问题视为大模型的长度外推挑战。因为我们总希望模型能够处理任意长的文本，但又不可能把训练样本的长度拉到任意长。

本文从位置编码（Position Encoding, PE）的角度出发来学习 Transformer 模型在长度外推方面的研究进展，研究各种旨在增强 Transformer 长度外推能力的方法，主要包括可外推的位置编码和基于这些位置编码的拓展方法。之前关于 RoPE 的介绍可以参考 rope。

问题背景

Transformer 自诞生以来就席卷了 NLP 领域。随着 LLM 能力的增长，我们对它们的期望也在增长，比如希望模型可以处理更长的文本，因为理解和扩展 LLM 的上下文长度对于提高其在各种 NLP 应用程序中的性能至关重要。

然而，增加 LLM 的上下文窗口并不是那么简单，因为 Transformer 的优势容量是以相对于输入序列长度的二次计算和内存复杂度为代价的。这导致了 Transformer 及在其基础之上的 LLM 都不具备有效长度外推（Length Extrapolation）的能力。这意味着，受限于其训练时预设的上下文长度限制，大模型无法有效处理超过该长度限制的序列。当输入超过该限制时，由于模型没有在预训练中见过超出上下文窗口的新的 token 位置，其性能会显著下降。

因此，如何解决长度泛化问题成为了 LLM 的一项主要挑战。

解决思路

为了实现更长文本的支持，当前的解决思路主要可以分为几个策略：

  • 在预训练阶段尽可能支持更长的文本长度。为实现这一阶段目标，通常采用并行化方法将显存占用分摊到多个 device，或者改造 attention 结构，避免显存占用与文本长度成二次关系。
  • 进行微调。比如在相对较小的窗口（例如 4K 令牌）上使用大量数据训练模型，然后在较大的窗口（例如 64K 令牌）上对其进行微调。
  • 在推理阶段尽可能外推到更大长度。为实现这一阶段目标，通常需要在两个方面进行考虑：对位置编码进行外推，优化 Attention 机制。

微调的挑战

因为微调和预训练本质类似，而微调难度远逊于预训练，所以我们来看看微调的挑战。

LLM背景下的微调代表了 NLP 领域的复杂演变。 这个过程涉及专门完善模型的现有功能，通过微调，LLM可以理解而且可以准确生成超出其初始训练数据参数的文本，在适应新的内容类型和结构方面表现出非凡的灵活性。微调外推侧重于通过额外的、有针对性的训练来提高模型的熟练程度。 然而，进一步扩展上下文窗口（微调）则存在以下几个主要挑战：

  • 高微调成本：扩展预训练的大型语言模型（LLMs）的上下文窗口到更长的文本时，通常需要在相应长度的文本上进行微调。但是由于 attention 的空间复杂度是 $O(n^2)$，这导致计算资源和时间上成本很高。随着上下文窗口的继续扩展，模型的计算量和内存需求将显著增加，带来极其昂贵的微调时间成本和 GPU 资源开销。
  • 长文本稀缺：微调通常需要相应长度的长文本，但当前训练数据中长文本数量有限。在当前的数据集中，尤其是超过1000k的长文本非常有限，这限制了通过微调来扩展上下文窗口的方法。
  • 新位置引入的灾难性值：首先，未经训练的新位置索引引入了许多异常值，使得微调变得困难。例如，当从 4k tokens 扩展超过1000k时，会引入超过90%的新位置。这些位置引入了许多灾难性值，导致分布外问题，使得微调难以收敛。
  • 注意力分散：当扩展到超长的上下文窗口后，由于引入众多新位置信息，大模型的注意力会分散在大量的 token 位置上，从而降低了大模型在原始短上下文窗口上的性能。 尽管上下文长度不会影响模型权重的数量，但它确实会影响这些权重如何编码令牌的位置信息。即使在微调之后，这也会降低模型适应较长上下文窗口的能力，从而导致性能不佳。

因此，人们普遍认为，用更长的上下文窗口对现有模型进行微调要么是有害的，要么是昂贵的。

长度外推的必要性

为了实现更长文本的支持，当前的解决思路主要可以分为几个策略：

  • 在预训练阶段尽可能支持更长的文本长度。为实现这一阶段目标，通常采用并行化方法将显存占用分摊到多个 device，或者改造 attention 结构，避免显存占用与文本长度成二次关系。
  • 进行微调。比如在相对较小的窗口（例如 4K 令牌）上使用大量数据训练模型，然后在较大的窗口（例如 64K 令牌）上对其进行微调。
  • 在推理阶段尽可能外推到更大长度。为实现这一阶段目标，通常需要在两个方面进行考虑：对位置编码进行外推，优化 Attention 机制。

长度外推

定义

外推概念的提出，最早可以追溯到 ALiBi 的论文中。如果模型在不经微调的情况下，在超过训练长度的文本上测试，依然能较好的维持其训练效果，我们就称该模型具有长度外推能力（extrapolation，也称 length extrapolation）。后来这种任务也被称为「上下文窗口拓展」(Context Window Extension)，目的依旧是用已经训好的模型来生成更大的文本，只是不再强调方法是外推。

顾名思义，免训练长度外推，就是不需要用长序列数据进行额外的训练，只用短序列语料对模型进行训练，就可以得到一个能够处理和预测长序列的模型，即“Train Short, Test Long（短训练，长推理）”。

  • train short：大部分文本的长度不会特别长，特别长的输入只是长尾情况，因此训练时的使用特别长的文本其实意义不大。再加上受限于训练成本，因此人们通常使用短序列训练，这样即符合实际情况，也可以显著降低训练开销。
  • test long：这里long是指推理时候的文本长度比训练时的最大文本长度还要长，希望不用微调就能在长文本上也有不错的效果。

Yarn

参考资料