大模型位置编码与长文本：算法侧与系统侧优化
Published:
长文本情境下，大模型面临了来自算法侧和系统侧的双重挑战：文本过长时不仅会导致模型生成质量下降，还会导致推理成本急剧增加。这里针对这两个方面的挑战进行了分析，以及记录了一些常见的算法侧优化方法。
1. 长文本模型的综合策略
- 算法侧：
- 位置编码扩展，如插值等
- 注意力偏置，鼓励局部注意力，避免注意力分布的熵增大
- 保留attention sinks确保注意力分布稳定等
- 系统侧：
- 稀疏注意力
- 线性注意力
- attention kernel优化
- GQA等kv cache缩减工作
2. 位置编码基础：RoPE
特性：低维高频，高维低频
较低的维度（$i$较小时）$\theta_i$较大，此时相对位置稍微变化一点就会造成旋转很大角度，因此低维度是高频分量，用于精细区分短依赖关系
较高的维度（$i$较大时）$\theta_i$较小，此时相对位置变化很大才能造成比较小的旋转，因此高维度是低频分量，用于捕捉长依赖关系
为什么RoPE比绝对位置编码更适合长文本/适合推理时长度外推：
绝对位置编码训练的时候序列多长，推理时序列就不能超过这个长度，否则将造成训推不一致导致模型效果变差。而相对位置编码学的是识别相对位置，例如第1和第5个token之间的位置差与第1001和第1005个token之间的位置差是一样的，因此不管序列绝对长度多大时，一定范围内的相对位置关系仍是可以被捕捉的。
3. 长文本外推（算法侧）
长文本外推指的是，给定一个在短文本上预训练的模型，希望在不额外训练（或只添加少量额外微调）的情况下，使得模型可以直接处理更长的序列（或微调后可以处理更长的序列），且长序列任务上性能损失较小。
训练和推理时长度不一致时，如果直接做长文本推理导致的模型算法性能问题：
- 推理时用到了没训练到的位置编码（不管是绝对位置编码还是相对位置编码都涉及这个问题）。可以通过插值的方法来缓解这个问题。
- 推理时做attention的token数量远大于训练时的token数量。虽然attention理论上可以处理任意长度的序列，但训练和推理长度不一致时，推理时attention中对于每个query会有更多的token去平均注意力分数，导致最后的注意力分数分布更“均匀”（注意力的熵更大），使得注意力更分散。而训练长度短则意味着注意力的熵更低，注意力更集中。这种差异会影响推理效果。因此，可以通过鼓励局部注意力来避免关注到过远的tokens，从而减小长文本带来的影响。
以下为RoPE结合插值方法做长文本外推的经典思路，主要解决第一个问题：
位置插值
设预训练长度为$L$，推理时遇到文本长度为$L’>L$的长序列，则对于长序列中的每个token，将它们的RoPE中的位置索引（设某个token实际上在位置$m$）缩放为$m\cdot \frac{L}{L’}$，这样就可以将整个句子在位置编码意义上“压缩”到预训练长度$L$以内。
这种方法需要使用少量长文本数据进行微调，才能得到一个能处理长文本的模型
NTK-aware RoPE
使用位置插值对于所有位置均等缩放会带来一个问题：直接压缩后会使得高频（低维）区域过于“拥挤”，位置稍微变化一点就会产生巨大震荡，这样会损害高频信息。
因此，NTK-aware RoPE提出对于所有维度采用非均匀缩放，对于高频（低维）部分（的位置索引参数$m$）缩放程度小一点，对于低频（高维）部分缩放程度大一点。总结来看也即：“高频外推，低频内插”
可以做到免训练外推，但外推长度有限。
NTK-by-parts
显式地对不同维度进行区别对待：不对高频（低维）做任何改动，只内插低频（高维）部分。
主要出发点在于：通常高频部分训练的比较充分，但低频部分可能训练的不够充分（相对距离增长到预训练长度的过程中，可能高频部分早已经转了很多圈，也即单位圆上所有点都训练充分了，但低频部分还没转完整的一圈，导致单位圆上某些点训练不充分），因此相对距离超过预训练长度后，就会落到这些没有充分训练的点上，导致性能下降。此时就需要使用内插法将其压到训练到的范围内。
Dynamic NTK
在自回归模型处理的序列不断增长的过程中，序列长度会不断变化，因此可以在推理阶段根据当前长度不断调整缩放因子。
YaRN
结合了attention scale和NTK-by-parts。前者是在attention计算时，在scale因子（也即推理长度$L’$和训练长度$L$的比值）上引入一个和缩放因子（也即文本长度）有关的温度因子$t$，来动态调节注意力机制：
\[a=\text{softmax}(\frac{q_mk^T_n}{t\sqrt{d}})\] \[\sqrt{\frac{1}{t}}=0.1\ln(s)+1\] \[s=\frac{L'}{L}\]上述经验公式表示了温度因子和文本长度缩放因子之间的关系（负相关）。
这样做可以更好地处理长序列，是因为模型上下文窗口扩展时，会导致注意力权重分布发生显著变化（一个token需要关注更多的tokens，导致注意力的熵更大，或者说注意力更分散），这种变化会导致ppl上升，因此引入温度因子来精细化控制注意力分布：当上下文窗口较大（使得缩放因子$s$也较大）时，需要降低温度因子$t$（使得注意力分布更加“尖锐”），从而使得注意力分布更类似长尾分布，一方面减小了注意力的熵，防止长文本下注意力分布和预训练时相比变化过大，同时也增强对于长距离信息的捕捉能力。
YaRN相比于普通的位置插值，在训练效率和收敛性上更好。
以下为通过添加注意力偏置来鼓励局部注意力的经典方法，主要解决第二个问题：
ALIBI： 是一种位置编码方法，不希望token数量过多时导致注意力太分散。其做法是：直接在attention score上减一个预设好的偏置，偏置值是两个token之间的距离，也即两个token距离越远则这个偏置值越大，导致二者注意力被削减越多：
\[q_m^Tk_n-\lambda|m-n|\]这样一来，就可以削弱较远的token的注意力贡献，使得每个token关注的范围尽可能保持在（预训练范围）的局部范围内，通过注意力局部化从而减小注意力熵增。
QKV bias： 结合RoPE位置编码进行使用，思路类似ALIBI，也是通过添加偏置来减弱长距离的贡献，从而促进注意力局部化。一个初始思路（GAU）为：
\[q_m^TR_m^TR_nk_n\to q_m^TR_m^TR_nk_n+a^TR_m^TR_nb\]其中$a,b$为两个可学习参数。可以证明$a^TR_m^TR_nb$呈现出关于$|m-n|$递减的趋势，可以起到局部注意力的作用。为了使得公式优雅，避免直接在attention矩阵上进行额外加项，进一步改进为在权重$W_q,W_k$处分别添加一个bias项$a,b$（之前的工作大多认为bias在模型足够大时没什么用，所以主流模型的权重都不带bias），在计算量几乎不变下使得attention变成：
\[q_m^TR_m^TR_nk_n\to (q_m+a)^TR_m^TR_n(k_n+b)\]将后式展开后变成：
\[q_m^TR_m^TR_nk_n+a^TR_m^TR_nk_n+q_m^TR_m^TR_nb+a^TR_m^TR_nb\]其中最后一项就是初始方法中加到attention score上的那一项。而也可以证明中间两项的影响是很小的（因为$q_m,k_n$的分布是“方向均匀的”，接近球面采样，因此二者分别和另一个随机向量$a,b$的夹角接近90度，因此内积接近0）。
这样一来，在预训练中为QKV权重添加bias，训好的模型直接在超过预训练长度的文本上运行效果明显比无bias时更好。但这种方法可复现性不是很稳定，可能造成训练稳定性问题。
4. 长文本下优化推理效率
长文本下系统层面的挑战：
Attention的二次复杂度导致计算压力大增
解决方案：
1、稀疏注意力：不让每个token都关注所有tokens，详见稀疏注意力相关部分
2、线性近似注意力通过改写注意力公式实现线性注意力，破除二次复杂度，详见线性注意力相关部分
3、IO感知的精确注意力
以flash attention为代表，通过tiling技术来在不同层次内存中高效移动数据，减少了GPU显存访存复杂度，详见注意力kernel优化相关部分
4、直接使用O(n)复杂度的SSM模型
KV Cache过大导致显存占用和访存压力大增
- GQA减少head数量
- KV Cache量化/压缩
系统级优化
使用continuous batching等，避免过长的prefill序列对调度带来损害