用短输入模拟长样本，高效拓展LLM上下文窗口，北大联合MSRA提出PoSE

©PaperWeekly 原创 · 作者 | 朱大卫

单位 | 北京大学计算语言所博士生

研究方向 | 长上下文建模

论文题目：

PoSE: Efficient Context Window Extension of LLMs via Positional Skip-wise Training

论文链接：

https://arxiv.org/abs/2309.10400

代码链接：

https://github.com/dwzhu-pku/PoSE

研究简介

大型语言模型（LLMs）通常有一个预定义的上下文窗口大小，这限制了它们在长输入的场景中的使用。为了使 LLMs 适应更长的输入，通常需要用目标长度的样本对其进行微调（全长微调），由此导致训练成本十分昂贵。

举例来说，在 Positional Interpolation [1] 这份工作中，将 LLaMA 的上下文窗口从 2048 拓展到 8192 使用了 32 张 A100，对于更大的上下文窗口则使用了 128 张 A100。

为了将训练长度与目标长度解耦合，以实现高效的上下文窗口扩展，我们提出了一种称为位置跳跃式训练（Positional Skip-wisE training, PoSE）的方法，在原始的上下文窗口中模拟更长的训练样本。

如下图所示，我们将原始的上下文窗口分成几块，然后引入不同的 bias 项来调整每个块的位置编码。对于每一条训练样本，这些 bias 项和块的长度都会发生变化，因此通过大量的训练，模型能适应目标长度内的所有位置。

实验结果表明，PoSE 有以下三方面的优势：

• 训练的时空效率：由于只需要按照原始的上下文长度进行训练，PoSE避免了由目标上下文长度增加带来的平方级别的计算复杂度，使得训练对于内存和时间的开销都大大减小。

• 能支持极长的上下文：通过解耦合训练长度和目标长度，我们仅使用2k的训练窗口就成功将 LLaMA 拓展到 128k。

• 兼容所有基于 RoPE 的模型和位置插值策略：PoSE 的有效性在 LLaMA、GPT-J、Baichuan 等多种基础模型，和 Linear、NTK、YaRN 等多种插值策略上得到了验证。

技术背景

旋转位置编码 RoPE：RoPE 是当下主流的位置编码方式，被 LLaMA、GPT-J 等大语言模型所采用。给定一个 维的隐向量  和位置 ，RoPE 通过如下方式编码位置信息：