YaRN — 大模型上下文窗口高效扩展

概述

YaRN（Yet another RoPE extensioN method）是一种 RoPE 位置编码扩展方法，将仅在 4K 上下文训练的 Llama 2 等模型扩展到 64K/128K，论文发表于 ICLR 2024。

核心定位：它不是插件，而是一套训练脚本——修改 RoPE 位置编码、校正注意力温度并轻量微调，产出长上下文的新模型权重。

核心原理

1. NTK-aware 缩放：不同维度用不同缩放因子。低频维度（远距离）大幅拉伸，高频维度（近距离）几乎不变，通过调整 RoPE 的 base 参数实现：

base' = base · s^(D/(D-2))    （s = L'/L，D = 维度数）

2. 温度校正：缩放后注意力 softmax 熵变大，在 softmax 前引入温度系数 t 校正：

attn = softmax(q·k / t)

t 的取值由"让扩展后注意力熵回到训练时水平"推导得出。

3. 轻量微调：少量数据（约 400 步）在扩展后长度上微调，适应新位置分布。

仓库结构

支持模型

Llama 2（7B/13B），覆盖 Meta 官方及 CodeLlama 等衍生权重。

使用流程

git clone https://github.com/jquesnelle/yarn.git
cd yarn

# 按 README 指定缩放因子与目标长度修改配置
# 在扩展后的长度上用少量数据微调
# 产出新权重后用 vLLM / HF Transformers 部署

关键要点

• 产物是新模型权重，不可挂载到现成服务化模型
• 原论文扩展上限为 64K/128K，1M 上下文需 LongRoPE 或原生长上下文模型
• 权重下载与复现指南见仓库 README 与 paper/ 目录

来源

• GitHub: https://github.com/jquesnelle/yarn
• 论文: arXiv:2309.00071（ICLR 2024）
• 许可证: MIT