缩放RoPE提升LLM外推能力探索 - YaRN NTK-aware

原生多头注意力在极长上下文中的位置外推能力有限。直接将RoPE从4K扩展至32K甚至更长模型在未见过长距离位置上会崩溃位置编码外推失效。这里参考网络资料探索分析YaRN、NTK-aware等更稳定的缩放RoPE编码和注意力沉淀方法。1 从NTK-aware到YaRN1.1 位置插值局限早期的位置插值(Position Interpolation)将所有位置索引等比例缩小至训练长度内。这种方法虽然简单但代价是丢失了区分相邻词的高频细节导致局部失明性能下降。1.2 NTK-aware缩放为了保留局部细节NTK-aware缩放提出了更精细的策略。NTK-aware缩放采用的方法是高频外推低频内插。不再统一缩放所有频率而是对负责局部细节的高频信息缩放小一点对负责长距离依赖的低频信息缩放大一点。具体操作是调整RoPE的基数(Base)即b常见值为10000。根据Neural Tangent Kernel (NTK) 理论将基数修改为其中s是缩放倍数|D|是模型维度。这样做的好处是通常无需微调即可实现一定长度的扩展。Code Llama用此方法实现了100k上下文窗口。1.3 YaRNYaRN(Yet another RoPE extensioN)可以看作是NTK-aware方法的增强版。YaRN在多个方面进行了关键改进旨在使扩展过程更平滑、结果更稳定1更精细的分段处理通过NTK-by-parts根据每个维度的波长(与上下文长度的比例)进行分类对不同频段采用不同策略处理更精细。2动态温度系数引入动态温度缩放来调整注意力分布。在长序列下这能有效防止注意力过度集中在某些token上起到类似抑制过热的效果使注意力分数更平滑从而稳定多轮对话中的位置偏差。3极佳效率与效果YaRN扩展能力很强仅需极少的微调数据(约400步)即可将LLaMA 2的上下文窗口从4K高效地扩展至128K效果显著。1.4 对比分析1核心原理NTK-aware 缩放采用非均匀缩放不同频率的旋转速度即“高频外推低频内插”。YaRN在NTK-aware基础上引入更精细的分段处理NTK-by-parts和动态温度系数。2扩展能力NTK-aware效果良好通常支持到64K-128K。YaRN效果更优可稳定支持至128K甚至更长。3微调需求NTK-aware通常无需微调即可外推。YaRN推荐少量微调(约400-1000步)以达到最佳效果资源消耗极低。4位置稳定性RTK-Aware位置稳定性良好保留了短距离区分能力。YaRN位置稳定性更好通过温度系数和分段策略长序列下的注意力分布更平滑稳定尤其有利于多轮对话。5实现复杂度NTK-Aware实现复杂度较低主要修改RoPE的基数计算引擎支持广泛。YaRN复杂度中等需要配置额外超参数(如分段阈值、温度系数)但主流框架已集成。目前主流推理框架均已集成这些技术通常通过设置rope_scaling参数即可调用。# 以 HuggingFace Transformers 为例 from transformers import AutoModelForCausalLM, LlamaConfig # 1. NTK-aware 缩放 config_ntk LlamaConfig.from_pretrained(meta-llama/Llama-2-7b-hf) config_ntk.rope_scaling {type: ntk, factor: 8.0} # 4K * 8 32K model_ntk AutoModelForCausalLM.from_pretrained( meta-llama/Llama-2-7b-hf, configconfig_ntk, torch_dtypetorch.float16, ).to(cuda) # 2. YaRN 缩放 config_yarn LlamaConfig.from_pretrained(meta-llama/Llama-2-7b-hf) config_yarn.rope_scaling {type: yarn, factor: 8.0} # 4K * 8 32K model_yarn AutoModelForCausalLM.from_pretrained( meta-llama/Llama-2-7b-hf, configconfig_yarn, torch_dtypetorch.float16, ).to(cuda)2 注意力沉淀除了扩展上下文窗口稳定超长多轮对话的质量还需解决另一个问题。随着对话增长如何确保模型不会“忘记”开头的关键信息注意力沉淀(Attention Sink)技巧便是一个关键的方法适应于多轮对话。2.1 StreamingLLMStreamingLLM最初发现模型天然倾向于为序列开头的几个token分配极高的注意力权重即使这些token在语义上不重要。一旦在滑动窗口推理中丢弃这些开头的token模型性能会迅速崩溃。1.2 Attention Sink基于此StreamingLLM框架利用了这一特性提出了一个简洁而强大的解决方案1锚定关键信息在推理时强制保留对话开头的几个token(如系统提示词)在KV Cache中永不丢弃。这种技巧即注意力沉淀(Attention Sink)这使得它们成为稳定的“注意力锚点”模型在生成任何后续内容时都能通过它“看到”整个对话的起点从而有效防止上下文漂移稳定了位置偏差。2显著提升性能该策略使模型能处理高达400万个token的流式输入且无需微调直接兼容主流LLM。SwiftInfer等优化版本还能带来最高46%的额外推理速度提升。1.3 总结YaRN确保模型能看懂长文本而注意力沉淀则保证模型在多轮对话中“记住”起点。两者结合可高效构建强大的长上下文LLM应用。reference---LLM YaRN详解、代码实现与应用扩展LLM上下文窗口https://awesomeml.com/llm-yarn/大模型的长上下文扩展位置插值、NTK感知缩放与上下文窗口外推https://aladdinedu.blog.csdn.net/article/details/159960976论文总结YaRN——大语言模型上下文窗口的高效扩展方法https://juejin.cn/post/7549878186416521258大语言模型结构之RoPE位置外推https://zhuanlan.zhihu.com/p/670244026LLM学习笔记-长度外推技术https://www.cnblogs.com/laozhanghahaha/p/18345815说明YaRN和NTK Aware外推的异同和优劣https://blog.csdn.net/sinat_37574187/article/details/150921299StreamingLLM让AI模型无限期平稳运行的一种方法https://www.chinaz.com/2023/1008/1564845.shtml社区供稿 | 多轮对话推理速度提升 46%开源方案打破 LLM 多轮对话的长度限制https://hub.baai.ac.cn/view/34743