Long Context 是一份面向研究工程的技术参考,主题是扩展 Transformer 模型的上下文窗口,核心方法包括 RoPE(旋转位置嵌入)、YaRN、ALiBi(带线性偏置的注意力)以及位置插值(Position Interpolation)。
【用途】 资料给出的适用场景包括:用 Transformer 模型处理 32k、64k、128k+ token 的长文档;对 LLaMA、Mistral 等预训练模型扩展上下文窗口;实现高效的位置编码(RoPE、ALiBi);训练具备长度外推能力的模型;部署可高效处理可变长输入的模型;以及以较低算力微调既有模型以适应更长上下文。
【主要能力与方法】
- RoPE:通过旋转矩阵编码绝对位置,并在注意力中提供相对位置依赖。资料给出公式 q_m=(W_q*x_m)e^(imθ)、k_n=(W_kx_n)*e^(inθ),其中 θ_j=base^(-2j/d)。优点包括随距离衰减的 token 间依赖、兼容线性注意力、外推能力优于绝对位置编码。提供了 RotaryEmbedding 模块、rotate_half 与 apply_rotary_pos_emb 的代码。
- YaRN:采用 NTK-aware 插值与注意力温度缩放,宣称可用约为基线少 10 倍 token 高效扩展上下文,可将 LLaMA 扩展到 128k token,训练步数比基线少 2.5 倍。资料给出 scale、original_max_position、extrapolation_factor、attn_factor、beta_fast、beta_slow 等参数配置。
- ALiBi:不向 token 添加位置嵌入,而是直接对注意力分数施加与 key-query 距离成正比的惩罚,公式为 attention_bias[i,j]=-m*|i-j|,m 为每个注意力头的斜率。资料列出 11% 训练加速、11% 显存减少、较强的长度外推(训练 1k、测试 2k+)以及近期性归纳偏置等优点,并给出 get_alibi_slopes 与 create_alibi_bias 代码。
- 位置插值:线性缩小位置索引,在已训练范围内插值而非外推。公式为 scaled_position[i]=i/extension_factor,可将 LLaMA 7B–65B 扩展到 32k token,资料称约 1000 步微调即可,稳定性比外推好 600 倍。
【方法对比】 资料以表格列出:RoPE 最大上下文 8k–32k、需完整预训练、显存中等、外推良好;YaRN 32k–128k、需极少训练、外推优秀;ALiBi 上下文不受限、需完整预训练、显存更低(-11%)、外推优秀,适合从零训练;位置插值 32k+、需极少训练(约 1k 步)、外推较差(设计使然),适合快速扩展。
【使用方式】 安装命令:pip install transformers torch;pip install einops;pip install rotary-embedding-torch;可选 pip install flash-attn --no-build-isolation。资料给出 HuggingFace 中通过 config.rope_scaling 配置 linear、dynamic、yarn 等类型与 factor 的示例,并附 LLaMA 位置插值、自定义 LongContextAttention、用 Trainer 进行最少步数微调(max_steps=1000 等)以及 YaRN 微调脚本命令的示例。
【最佳实践与限制】 资料建议:新模型从零训练用 ALiBi;扩展已有 RoPE 模型用 YaRN;快速低算力扩展用位置插值。缩放因子从 2.0 到 16.0 递增,因子越大所需微调步数越多。微调数据应使用与目标长度匹配的长文档(如 32k、24k、16k token),并提到 PG-19、arXiv 论文、长对话、GitHub 仓库等数据集;避免只使用短文档。常见陷阱包括:不微调直接套用位置插值会导致表现差;在缺乏大量微调的情况下做过于激进的扩展(如直接到 1M token)不可行,建议按 8k→16k→32k→64k 逐步扩展并每步微调。
【生产部署】 资料给出加载 32k 上下文模型(如 togethercomputer/LLaMA-2-7B-32K)进行推理的示例,以及内存优化手段:梯度检查点、Flash Attention 2(宣称 2–3 倍加速)、vLLM 的 paged attention(max_model_len=32768)。
【相关资源】 引用的论文与实现包括 RoFormer(arXiv 2104.09864)、YaRN(arXiv 2309.00071)、ALiBi(arXiv 2108.12409)、Position Interpolation(arXiv 2306.15595),以及 HuggingFace RoPE 工具、jquesnelle/yarn 实现与 Together AI 博客。
关于收费、价格、官方支持范围等信息,资料暂未提供。