YaRN论文解读

论文标题：YaRN: Efficient Context Window Extension of Large Language Models

论文地址：https://arxiv.org/abs/2309.00071

论文发布时间：2023-08-31

本篇论文所提出的算法 YaRN，被 Qwen2.5 在外推拓展方面使用。

Abstract

旋转位置嵌入（RoPE）无法在超出其训练的序列长度上进行泛化。

论文提出了 YaRN（Yet another RoPE extensioN method，另一种 RoPE 扩展方法），一种计算高效的方法来扩展此类模型的上下文窗口，所需的标记数量比以前的方法少 10 倍，训练步骤少 2.5 倍。

Introduction

目前 Decoder Only 结构的大模型，在扩展长上下文能力上的主要限制来自于位置编码。想办法通过少量微调（或不进行微调）动态扩展上下文窗口是目前研究的重点。

最初的 Transformer 使用绝对位置编码（正余弦），后来改进为可学习的绝对位置编码（Embedding）。目前通用的位置编码是相对位置编码，RoPE。后续很多的工作，如 NTK、Dynamic NTK，基本上都是基于 RoPE 的改进。

包括本篇文章提出的 YaRN。

YaRN 在对不到 0.1% 的原始预训练数据进行微调后，在上下文窗口扩展中达到了最先进的性能。同时，通过与称为 Dynamic Scaling 的推理时间技术相结合，Dynamic-YaRN 允许 2 倍以上的上下文窗口扩展，无需任何微调。

Background and Related Work

旋转位置编码

RoPE 是老生常谈了，就不介绍了。

在实际坐标中，可以使用以下函数表示 RoPE

$f_{W}\left(x_{m}, m, \theta_{d}\right)=\left(\begin{array}{ccccccc} cos m \theta_{1} & -sin m \theta_{1} & 0 & 0 & \cdots & 0 & 0 \\ sin m \theta_{1} & cos m \theta_{1} & 0 & 0 & \cdots & 0 & 0 \\ 0 & 0 & cos m \theta_{2} & -sin m \theta_{2} & \cdots & 0 & 0 \\ 0 & 0 & sin m \theta_{2} & cos m \theta_{2} & \cdots & 0 & 0 \\ 0 & 0 & 0 & 0 & \cdots & cos m \theta_{l} & -sin m \theta_{l} \\ 0 & 0 & 0 & 0 & \cdots & sin m \theta_{l} & cos m \theta_{l} \end{array}\right) W x_{m}$

位置插值

给定一个是用 RoPE 做预训练的大语言模型，通过以下方式修改 RoPE：

$f_{W}'\left(x_{m}, m, \theta_{d}\right)=f_{W}\left(x_{m}, \frac{m L}{L'}, \theta_{d}\right)$

其中 $L^{'} > L$ 是超出预训练限制的新上下文窗口。利用原始的预训练模型和修改后的 RoPE 公式，可以在减少几个数量级的 token 上进一步微调语言，实现上下文窗口扩展。

原理就是“抽象的拉抻”，把原来训练好的 L 个窗口大小的输入长度，通过加因子变长。（比如原长度为 1000，除以 0.8，就变成 1250。）

为了方便介绍后续方法，重写并简化为以下一般形式：

$f_{W}'\left(x_{m}, m, \theta_{d}\right)=f_{W}\left(x_{m}, g(m), h\left(\theta_{d}\right)\right)$

因此对于 $f_{W}'$ ， $g (m) = m / s$ ， $h(\theta_{d})=\theta_{d}$

Methodology

高频信息丢失——“NTK 感知”插值

为了解决在插值 RoPE 嵌入时丢失高频信息的问题，提出了 “NTK-aware” 插值，NTK 感知插值不是将 RoPE 的每个维度均等地缩放一个因子 s ，而是通过减少高频和增加低频来将插值压力分散到多个维度。可以通过多种方式获得这样的变换，但最简单的是对 θ 的值进行更改：

$\quad$

$h\left(\theta_{d}\right)=b^{\prime-2 d /|D|}, \quad$

$\cdot s^{\frac{|D|}{|D|-2}}$

$b$ 是常数，通常会给个固定值，如 32k。

与 PI 相比，这种方法在扩展非微调模型的上下文大小方面表现得要好得多。这种方法的一个主要缺点是，鉴于它不仅仅是一个插值方案，某些维度会略微外推到“越界”值，因此使用 “NTK 感知” 插值进行微调会产生不如 PI 的结果。此外，由于“越界”值，理论比例因子 s 并不能准确描述真实的上下文扩展比例。在实践中， scale 值 s 必须设置为高于给定上下文长度扩展的预期 scale 。

没太理解为什么会越界，改造后的 $b^{'}$ 不应该比原来的 $b$ 小吗？

相对局部距离丢失 - “NTK-by-parts” 插值

因为 $\lambda_{d}=\frac{2 \pi}{\theta_{d}}=2 \pi b^{\frac{2 d}{|D|}}$ ，所以作者打算抛开辐角不谈，从波长的角度理解位置编码。作者提出两个观点：

给定上下文大小 L ，有一些维度 d 的波长长于预训练期间看到的最大上下文长度 $(\lambda>L)$ ，这表明某些维度的嵌入在旋转域中分布不均匀。
当以更小的 $b^{'}$ 拉伸所有 RoPE 维度时，辐角更小，波长更大，所有 token 都彼此更接近，两个向量旋转量较少的点积更大。这种缩放严重损害了 LLM 理解其内部嵌入之间的小关系和局部关系的能力。