中文

LinearARD:用于恢复RoPE的线性内存注意力蒸馏

计算与语言 2026-04-02 v1 人工智能

摘要

大语言模型上下文窗口的扩展通常通过扩展位置编码后进行轻量级持续预训练(CPT)来实现。虽然这种方法对处理长序列有效,但会破坏原始模型能力,在标准短文本基准测试上导致性能下降。我们提出LinearARD,这是一种自蒸馏方法,通过注意力结构一致性恢复RoPE缩放的学生模型。LinearARD不匹配晦涩的隐藏状态,而是对稠密Q/K/V自关系矩阵的行级分布进行对齐,以直接监督注意力动态。为克服n×n关系图的二次内存瓶颈,我们引入线性内存内核。该内核利用每个标记的log-sum-exp统计信息,并将logit重计算融入反向传播,以计算精确的KL散度和梯度。在将LLaMA2-7B从4K扩展至32K时,LinearARD恢复了98.3%的短文本性能,同时在长上下文基准测试中超越现有最先进方法。值得注意的是,我们的方法仅需约4.25M个训练标记,即可实现这些成果,而LongReD和CPT需要256M个标记。我们的代码已公开:https://github.com/gracefulning/LinearARD。

关键词

引用

@article{arxiv.2604.00004,
  title  = {LinearARD: Linear-Memory Attention Distillation for RoPE Restoration},
  author = {Ning Yang and Hengyu Zhong and Wentao Wang and Baoliang Tian and Haijun Zhang and Jun Wang},
  journal= {arXiv preprint arXiv:2604.00004},
  year   = {2026}
}