长度外推为在不需要对模型进行额外训练的情况下,模型可以处理更长的序列。本篇文章主要介绍目前大模型用到的一些长度外推技术,包括以RoPE为基础进行位置插值、NTK-aware、动态NTK、NTK-by-parts 和YaRN。关于RoPE,可参见我的上一篇博客LLM学习笔记-位置编码篇
位置插值
回想一下Transformer结构中使用RoPe进行query, key, value(下文用q、k、v表示)的计算
xm表示m位置处的词嵌入, θd是d维的旋转角度。为了在预测阶段扩展窗口长度,一个非常朴素的方法就是在预测阶段给长文本的位置乘上一个缩放因子1/s,将长文本的位置缩放到训练阶段窗口长度之内。缩放因子1/s一般为训练阶段文本长度L/预测阶段文本长度L', 即 1/s=L/L'
优点
- 实现简单
- 兼容性强,与当前的位置编码技术,比如正余弦编码、RoPE高度兼容
缺点
- 仍需要进行微调。根据Meta的论文,需要大约1000步长文本数据微调才能得到一个能处理长文本的模型
- 效果有限。位置插值均等压缩长文本中的每个位置,扰乱了模型的局部分辨率,邻近token的位置关系被影响,导致模型无法捕捉长文本中相邻token的关系,进而影响最终效果
NTK-aware RoPE
根据NTK理论,如果输入维度较低且相应的嵌入缺乏高频分量,深度神经网络将难以学习高频信息。位置插值利用缩放因子s对RoPE embedding的所有维度进行均等缩放会损害embedding中的高频信息。NTK-aware RoPE的提出者认为位置插值效果有限的原因就在于此。NTK-aware RoPE对所有维度采用非均匀缩放,对于高频(低维)部分缩放程度小一点,对于低频(高维)部分缩放程度大一点。总结来说就是‘高频外推,低频内插’,具体通过对θ中的base,10000进行缩放来完成。
原来, 现在改为。为实现低频内插,则需当i = d/2-1时,有,解得。 将该结果带入到表达式中可得。 当i为0时,最右边式子的第二项为1,不会进行缩放;当i为d/2-1时,最右边式子的第二项为1/s,从而实现了高频外推,低频内插。
优点
- 实现免训练外推
- 实现简单
- 兼容性强,与当前的位置编码技术,比如正余弦编码、RoPE高度兼容
缺点
- 外推长度有限
NTK-by-parts
NTK-by-parts在NTK-aware基础上对embedding的不同维度进一步的区别对待, 不对高频做任何改动,只内插低频部分。 NTK-by-parts引入了波长的概念,[关于波长,可参见我的上一篇博客LLM学习笔记-位置编码篇]。对于RoPE embedding而言,给定一个窗口长度L,embedding中某些维度的波长要大于L,这意味着该维度可能训练的不够充分,关于该部分苏剑林博客里解释的很好,我参考他的博客内容简单解释一下。
加入RoPE之后的qk内积可以用复数表示为:
其中表示为单位圆上的一个点,当m-n增大时,该点开始旋转,转速与θ有关。θ越大转速越快,因此m-n在增大到L-1的过程中该点会转很多圈,即该单位圆上的所有点都被训练到。θ越小转速越慢,m-n从0增大到L-1时,可能一个单位圆还没有转完只转了部分弧度,该单位圆上的点训练不充分,因此当m-n超出窗口大小时,该点可能会落到没有经过训练的点上,从而有无法预估的表现,这时就需要用内插法把该点压缩回训练的弧里。而该点在单位圆上转的圈数与维度对应的波长有关 。 Ti为周期,计算方法与波长一样。
NTK-by-parts的做法就是将窗口长度L与维度对应的波长进行比较,如果波长远小于窗口长度,说明该维度对应的点转了很多圈,不用进行改动;如果波长大于等于窗口长度,则需要进行内插;如果波长与窗口长度的关系在前两者之间,则兼而有之。具体操作上,NTK-by-parts引入了两个超参α和β,对于某个维度而言,如果该维度对应的转数r(d) > β,则无需对改动,如果该维度对应的转数r(d) < α,则需要进行内插,实验中α=1, β=32。定义斜坡函数γ如下
对频率进行如下缩放
Dynamic NTK
Dynamic NTK主要提出了动态的缩放因子。对于自回归模型而言,其所处理的序列是不断变化的,每一步比前一步增加一个token。 在这种情况下有两种插值的方法
- 在整个推理阶段,使用固定的缩放因子s=L'/L, 其中L'是预测阶段文本长度
- 在每轮前向传播中,位置编码更新缩放因子s = max(1, l'/L), 其中l'是当前序列的长度
方案1的问题在于,当序列长度小于L'时,模型性能可能会有一点小折扣;而当序列长度大于L'时,模型性能则会突然下降,而使用方案2可以使模型性能下降的更少。论文中把方案二这种在推理阶段不断调整缩放因子的方式称为Dynamic 缩放,再结合NTK感知方法,论文称之为‘动态NTK’插值
YaRN
Yarn实际上是将attention scale 和 NTK-by-parts结合到一起。具体来说,是在attention 计算的时候引入了一个温度因子t,
。
t的值是根据实验得出, 实验结果显示引入这个缩放因子后,ppl有所下降。
优点
- YaRN可以直接和修改注意力机制的库兼容,比如Flash attention
- 效果拔群,在不到原始训练数据0.1%进行微调,YaRN在上下文窗口的扩展上优于其他测试的方法
总结
本篇博客主要总结了基于RoPE的长度外推技术,主要采取缩放的方法,1. 按token位置进行缩放,例如位置插值,2. 按embedding维度进行缩放,例如NTK-aware, NTK-by-parts, Dynamic-NTK等,采用不同的方法对embeeding的维度进行非均匀缩放,保证高频不缩放或者缩放程度小一点,低频进行缩放
Ref
YaRN: Efficient Context Window Extension of Large Language Models
Transformer升级之路:16、“复盘”长度外推技术
大模型长度扩展综述:从直接外推ALiBi、插值PI、NTK-aware插值(对此介绍最详)、YaRN到S2-Attention