跳到正文
arXiv cs.CL· Ning Yang, Hengyu Zhong, Wentao Wang, Baoliang Tian, Yuan Zhou, Haijun Zhang, Jun Wang·· 12 小时前AI 评分40

LinearARD:用线性内存注意力蒸馏恢复 RoPE,仅需 1.6% token 预算

LinearARD: Linear-Memory Attention Distillation for RoPE Restoration

AI 导读

LinearARD 提出一种恢复式蒸馏方法,通过对齐冻结原生 RoPE 教师模型最终注意力层的 Q/Q、K/K、V/V 自关系矩阵行分布,将 RoPE 缩放后的学生模型恢复至长上下文能力。

正文 · AI 翻译

查看 PDF HTML(实验性)

摘要:大语言模型上下文窗口的扩展通常通过缩放位置编码并随后进行持续预训练(CPT)来实现。虽然有效,但这一范式以数据需求量大和计算成本高昂而著称,需要大量长文本语料来使模型重新校准到偏移后的位置分布。我们提出 LinearARD,一种恢复蒸馏方法,通过与冻结的原生 RoPE 教师模型的注意力结构一致性,来恢复经过旋转位置嵌入(RoPE)缩放的学生模型。LinearARD 不采用下一 token 预测或难以解释的隐藏状态匹配,而是对齐最终注意力层中稠密 $Q/Q$、$K/K$ 和 $V/V$ 自关系矩阵的逐行分布,以直接监督注意力动态。为消除 $n \times n$ 关系图的二次内存瓶颈,我们引入一种线性激活内存内核,它仅存储每个 token 的 log-sum-exp 统计量,并在反向传播中重新计算 logits,以获得精确的 Kullback-Leibler 散度梯度。在扩展到 32K 上下文的 LLaMA2-7B、LLaMA3-8B 和 Mistral-7B-v0.1 上,LinearARD 分别恢复了原生短上下文性能的 93.1%/94.2%/94.3%,并且仅使用 4.25M 个 token 就在 RULER 上实现了强大的长上下文鲁棒性——这仅相当于最先进基线所需的 256M token 预算的 1.6%(约 $\sim$60$\times$ 的减少)。在这种严重受限的预算下,CPT 和 LongReD 在 RULER 上仍接近于零,表明关系级恢复在效率与质量之间提供了显著更好的权衡。我们的代码可在 this https URL 获取。
主题: 计算与语言(cs.CL);人工智能(cs.AI)
引用为: arXiv:2604.00004 [cs.CL]
  (或 arXiv:2604.00004v3 [cs.CL] 用于此版本)
  https://doi.org/10.48550/arXiv.2604.00004

arXiv 通过 DataCite 颁发的 DOI

提交历史

来自:Hengyu Zhong [查看邮箱]
[v1] 2026年3月9日星期一 10:51:01 UTC(4,047 KB)
[v2] 2026年7月23日星期四 07:35:13 UTC(3,544 KB)
[v3] 2026年10月26日星期一 06:15:39 UTC(3,553 KB)

来源:arXiv cs.CL · arxiv.org