近日，清华大学孙茂松团队在 AI Open 上发表了题为“Optimal RoPE Extension via Bayesian Optimization for Training-Free Length Generalization”的研究论文。该工作首先提出了一个统一的频率权重（Frequency Weight）框架，将PE、PI、NTK、YaRN等现有RoPE扩展方法统一归纳为频率权重调整问题，然后利用贝叶斯优化（Bayesian Optimization）自动搜索最优频率权重，提出了ORoPE方法。该方法完全无需训练，仅通过在代理任务上优化即可实现长上下文泛化。在Llama-2-7b和MiniCPM-2b两个模型的8K、16K、32K多种上下文长度实验中，ORoPE均显著超越了现有手工设计方法，并可与经过全长度训练的模型相媲美。

研究背景

大语言模型（LLMs）通常使用固定长度的上下文窗口进行训练，但实际应用中往往需要处理远超训练长度的长文本。旋转位置编码（Rotary Position Embedding, RoPE）是Llama、MiniCPM等主流LLM广泛采用的位置编码方式。如何在不重新训练的前提下扩展RoPE以支持更长上下文，是提升LLM实用性的关键挑战。现有方法如Position Interpolation（PI）、NTK-aware Scaling和YaRN等，通过手工设计频率缩放因子来实现长度泛化，但这些手工方法受限于人工先验知识，难以找到最优解。

研究亮点

• 提出统一的频率权重框架，将PE、PI、NTK、YaRN等现有RoPE扩展方法统一表示为 θ_i = w_i · b^{-2i/d} 的形式，揭示了不同方法本质上仅是频率权重w_i的差异。

• 利用贝叶斯优化自动搜索最优频率权重，完全无需训练或微调即可扩展LLM的上下文长度，大幅降低了长度泛化的计算成本。

• 在LongBench、Retrieve等多个基准测试上，ORoPE在Llama-2-7b和MiniCPM-2b两个模型的8K/16K/32K长度上均达到最优性能，超越所有无训练基线方法。

• 通过分段搜索、约束设置和代理任务设计，将搜索空间从d/2维降低到少量分段，使贝叶斯优化能够高效收敛。

图文导读

为了理解现有方法在长度泛化中失败的原因，作者首先分析了不同RoPE扩展方法的注意力输出标准差。图1展示了在Retrieve.PassKey任务（8K长度）上的分析结果。可以观察到，PI和YaRN在某些token位置出现标准差接近零的现象，说明这些位置的注意力输出几乎失去了区分能力，导致模型无法正确提取目标信息。NTK虽然标准差分布相对正常，但在生成质量上仍有不足。相比之下，ORoPE在所有位置都保持了良好的标准差分布。

表1进一步展示了不同方法在Retrieve.PassKey检索任务上的生成文本案例。PI方法出现了直接复制输入前缀的问题，NTK输出了混乱的数字序列，YaRN及其变体的输出也存在不同程度的错误。而ORoPE和其变体ORoPE-s能够正确定位并输出目标信息（PassKey: 48 820 和 99 408），验证了该方法在长文本理解上的优越性。

基于上述分析，作者提出了统一的频率权重（FW）框架。在RoPE中，位置编码由旋转频率 θ_i 决定，作者将其统一表示为 θ_i = w_i · b^{-2i/d}，其中 w_i 称为频率权重。在此框架下，PE对应 w_i=1（不做缩放），PI对应 w_i=1/s（统一缩小），NTK对应 w_i=s^{-2i/d}（按维度渐变缩放），YaRN则使用分段线性插值。不同方法的本质区别仅在于频率权重的取值方式，因此可以通过优化频率权重来寻找更优的RoPE扩展方案。

作者采用贝叶斯优化来自动搜索最优频率权重。为了评估不同频率权重配置的效果，需要设计合适的代理任务（Proxy Task）。图2展示了Retrieve.PassKey代理任务的格式：在一段长文本中随机插入一个密钥数字，要求模型从中检索出该数字。作者还使用了Retrieve.Number和Retrieve.KV等多种检索任务作为代理任务。通过分段搜索策略，将搜索空间从原始的d/2维（如Llama-2的64维）降低到8个分段，大幅提高了搜索效率。

表3展示了Llama-2-7b-chat在8K上下文长度下的性能对比。评估涵盖了建模任务（GovReport困惑度）、检索任务（Retrieve.KV/Number/PassKey）以及LongBench理解任务（2WikiMultiHopQA、GovReport摘要、HotpotQA、Musique、NarrativeQA）共9个基准测试。ORoPE-s以28.09的平均得分和2.42的平均排名位居第一，显著超越了所有手工设计的基线方法。值得注意的是，原始PE（直接外推）几乎完全失效，平均得分仅为0.69。

表4展示了MiniCPM-2b在8K上下文长度下的对应结果。ORoPE-s同样以31.24的平均得分领先，特别是在检索任务（Retrieve.KV: 13.54, Retrieve.Number: 98.44）上大幅超越了NTK等方法。这说明ORoPE方法具有良好的模型通用性，不依赖于特定的模型架构。

表9将ORoPE与经过全长度语料库训练的长上下文模型进行对比。这些模型包括CodeLlama（将RoPE基数从10,000提高到1,000,000）、LongLoRA（使用32K和100K长度训练）以及YaRN（使用64K和128K长度训练）。尽管ORoPE完全不需要任何训练，但其性能（Avg score: 28.09，Avg rank: 2.89）仍然优于大部分经过大规模训练的模型，仅略低于YaRN-64K（21.24）和YaRN-128K（20.58）在某些特定任务上的表现。总体而言，ORoPE以极低的计算成本实现了与训练方法相当甚至更优的效果。

图3和图4分别展示了Llama-2-7b和MiniCPM-2b上，各方法的频率权重分布。可以看到，PI的频率权重为常数（所有维度统一缩放），NTK呈指数递增趋势，YaRN则呈分段线性形态。ORoPE通过贝叶斯优化搜索得到的频率权重具有最宽的动态范围，在不同频率维度上采取了差异化的缩放策略：低频维度（高索引）的权重较大，高频维度的权重较小，但整体分布与手工方法有着本质区别。这说明贝叶斯优化能够发现人工难以设计的最优频率配置。

图5展示了消融实验结果，分析了三个关键因素的影响。在分段数量方面，8段搜索在效率与性能之间取得最佳平衡。在代理任务方面，Retrieve.PassKey容易饱和，Retrieve.KV则更具区分度，组合多个代理任务能获得更稳健的结果。在搜索边界方面，更大的上界能够获得更好的结果，但也需要更长的搜索时间；过大的下界会排除最优解，过小则增加搜索时间。总体而言，ORoPE的超参数设置较为鲁棒，在合理范围内都能获得良好性能。

资助信息

该工作得到国家重点研发计划项目（No.2022ZD0116312）、国家自然科学基金项目（No.62236004）和清华大学国强研究院的资助。

原文信息

DOI：https://doi.org/10.1016/j.aiopen.2025.01.002

Optimal RoPE extension via Bayesian Optimization for training-free length generalization

期刊介绍

AI Open是一本致力于分享人工智能及其应用理论的英文国际期刊，期刊侧重人工智能领域可操作知识层面及具有前瞻性观点的研究。期刊主编由清华大学计算机与科学技术系唐杰教授担任。

AI Open欢迎人工智能及其应用相关领域的文章。

期刊收录的所有文章都经过严格的同行评审，并发表在月活用户超过2000万的ScienceDirect平台，供领域内的学者、及全球读者免费阅读、下载及引用。

目前，期刊已被ESCI、Ei Compendex、Scopus、DOAJ、dblp computer science bibliography、EBSCOhost等权威数据库收录。在COMPUTER SCIENCE, INTERDISCIPLINARY APPLICATIONS学科175种国际期刊中位列第3位（Q1区），在COMPUTER SCIENCE, ARTIFICIAL INTELLIGENCE学科204种国际期刊中位列第5位（Q1区）。

主编

• 唐杰，清华大学

• Björn W. Schuller， Imperial College London

副主编

• Wendy Hall， University of Southampton

• Michalis Vazirgiannis，Ecole Polytechnique

• Jose A. Lozano，University of the Basque Country UPV/EHU

• Esma Aïmeur，University of Montreal, Canada

• 刘知远，清华大学

• 张静，中国人民大学

• 东昱晓，清华大学

• 吴乐，合肥工业大学

• 马家祺，University of Illinois Urbana-Champaign, USA

• 何向南，中国科学技术大学

• 邱锡鹏，复旦大学

转载本文请联系原作者获取授权，同时请注明本文来自科爱KeAi科学网博客。

链接地址：https://wap.sciencenet.cn/blog-3496796-1525333.html?mobile=1

收藏

分享