1 引言
论文链接：https://arxiv.org/pdf/2309.00071v3
随着大模型长文本应用场景持续普及，超长文档摘要、整库检索、超长多轮对话等任务对模型上下文窗口的要求持续提升。当前绝大多数开源基座模型均基于旋转位置编码（Rotary Position Embedding，RoPE）[1]训练，原生上下文窗口固定，直接外推会出现位置分布偏移、注意力距离失真、长文本性能断崖下跌等问题。为低成本突破预训练窗口限制，业界陆续提出了位置插值（Position Interpolation, PI）[2]、NTK-aware[3] 全局基底缩放、NTK-by-parts[4] 分段插值等一系列 RoPE[1] 扩窗方案。
其中，PI[2] 简单粗暴压缩全局频率，严重损伤短文本局部细节；NTK-aware[3] 通过全局非线性缩放实现高低频差异化外推，大幅缓解短文本退化，但存在中频频段耦合、大倍率扩窗精度不足的缺陷；后续迭代的 NTK-by-parts[4] 引入波长分段思想，对高、中、低频进行分区加权过渡，实现了更精细的频段调控，成为工业界效果最优的零微调扩窗范式。
但即便 NTK-by-parts[4] 已经大幅平衡长短文本性能，其仍存在一处核心短板：固定的中频过渡曲线无法适配不同扩展倍率与模型特性，大倍数扩窗下远距离注意力衰减、检索准确率下降的问题依旧存在。为彻底解决分段插值的过渡生硬、长距离依赖建模不足问题，ICLR 2024 提出了 YaRN[5] 长上下文扩展算法。
YaRN[5] 在 NTK-by-parts[4] 分段波长架构的基础上，引入斜坡式动态缩放策略与注意力温度补偿机制，让中频过渡不再是固定线性加权，而是随扩展倍率自适应平滑变化，同时通过温度参数修复长距离注意力分布偏差。相比前代方案，YaRN[5] 既完整保留高频局部细节、规避 PI[2] 的全局压缩缺陷，又进一步抹平了大倍率扩窗下的频段断层，实现零微调、长短通吃、超大倍率稳定外推的最优效果，成为目前开源社区落地最广、综合性能最强的 RoPE[1] 原生扩窗方案。本文将从前代算法缺陷、YaRN[5] 核心原理、动态缩放公式、温度补偿机制，完整拆解 YaRN[5] 的技术细节与工程价值。
2 YaRN
RoPE[1] 回顾：Roformer: Enhanced transformer with rotary position embedding-CSDN博客
NTK-by-parts[4]回顾：https://arxiv.org/pdf/2309.00071v3
YaRN[5]在NTK-by-parts[4]基础上引入了与扩展倍率相关的logit温度缩放机制：在logits送入 Softmax之前，引入温度系数t，形式如下：
这等价于缩放 RoPE[1] 嵌入,而无需更改注意力内核。由于 RoPE[1] 嵌入预先生成、在所有前向传播中复用，该方案在训练与推理阶段均不会引入额外开销。将这一机制与 NTK-by-parts[4] 插值框架结合，便构成完整的 YaRN 算法。
这等价于直接缩放RoPE[1]嵌入，而无需改动注意力内核。由于RoPE[1]嵌入可预先生成，并在全部前向传播过程中复用，该机制在训练与推理阶段均不带来额外计算开销。将该温度补偿机制与 NTK-by-parts[4]插值框架相结合，就构成了完整的YaRN[5]算法。
对于LLaMA[6]和LLaMA 2[7]模型，t推荐以下值:
3 动态NTK插值
静态NTK插值在整个推理周期中，嵌入层是固定的，包括缩放因子:
其中 为预训练上下文窗口大小，
为目标上下文窗口大小。使用静态NTK插值时，模型在输入序列长度小于
时可能出现性能折扣，当序列长度大于
时可能出现突然退化。
对此，故提出了动态NTK插值[5]，在推理阶段每次前向传递中，位置嵌入更新缩放因子：
其中是当前输入序列的序列长度，当序列长度不超过上下文窗口大小时，不触发任何插值；仅当序列更长时，才自适应开启RoPE[1]缩放。该方法最早诞生于社区讨论，YaRN[5]将其规范化，并基于它构建了零样本分支Dynamic-YaRN[5]。
4 总结
本文梳理了RoPE[1]系列上下文扩展方案的演进脉络：从全局压缩的PI[2]插值、全局非线性缩放的 NTK-aware[3]，再到采用分段斜坡加权的NTK-by-parts[4]，各方案在长短文本性能之间不断取舍，但大倍率扩窗下远距离注意力衰减、检索精度下滑的问题始终存在。YaRN[5]在 NTK-by-parts[4]分段插值架构之上，新增了依赖扩展倍率的注意力温度补偿机制，借助RoPE[1]可拆分为二维旋转矩阵的特性，将温度缩放等价融合进预计算的RoPE[1]嵌入，无需修改注意力内核，训练与推理均无额外开销。与此同时，YaRN[5]还规范化了源自社区的动态NTK插值策略[5]，衍生出 Dynamic-YaRN[5]零样本分支：推理时根据当前序列长度动态调整缩放因子，短序列不触发插值，保护原生短文本能力，长序列自适应开启RoPE[1]扩展。凭借这套组合设计，YaRN[5]实现了兼顾长短文本、支持超大倍率稳定外推的长上下文扩展能力，也是当前开源生态落地广泛的RoPE[1]扩窗方案。
参考文献
[1] Su J, Ahmed M, Lu Y, et al. Roformer: Enhanced transformer with rotary position embedding[J]. Neurocomputing, 2024, 568: 127063.
[2] Chen S, Wong S, Chen L, et al. Extending context window of large language models via positional interpolation[J]. arXiv preprint arXiv:2306.15595, 2023.
[3] Peng B, Quesnelle J. Ntk-aware scaled rope allows llama models to have extended (8k+) context size without any fine-tuning and minimal perplexity degradation[EB/OL].(2023)
[4] bloc97. Add NTK-Aware interpolation "by parts" correction, 2023. URL https://github. com/jquesnelle/scaled-rope/pull/1.
[5] Peng B, Quesnelle J, Fan H, et al. Yarn: Efficient context window extension of large language models[C]//International Conference on Learning Representations. 2024, 2024: 31932-31951.
[6] Touvron H, Lavril T, Izacard G, et al. Llama: Open and efficient foundation language models[J]. arXiv preprint arXiv:2302.13971, 2023.
[7] Touvron H, Martin L, Stone K, et al. Llama 2: Open foundation and fine-tuned chat models[J]. arXiv preprint arXiv:2307.09288, 2023.
转载自 CSDN-专业IT技术社区
原文链接：https://blog.csdn.net/qq_52771580/article/details/162263692
