论文

LinearARD:用于 RoPE 恢复的线性记忆注意力 蒸馏

LinearARD: Linear-Memory Attention Distillation for RoPE Restoration

摘要

大语言模型 中上下文窗口的扩展通常通过缩放位置编码和轻量级连续 预训练 (CPT) 来实现。虽然对于处理长序列有效,但这种范例通常会破坏原始模型的功能,导致标准短文本基准的性能下降。我们提出了 LinearARD,一种自我 蒸馏 方法,通过与冻结的本地 RoPE 教师的注意力结构一致性来恢复旋转位置嵌入 (RoPE) 规模的学生。 LinearARD 不是匹配不透明的隐藏状态,而是对齐密集 $Q/Q$、$K/K$ 和 $V/V$ 自关系矩阵的行分布,以直接监督注意力动态。为了克服 $n \times n$ 关系图的二次内存瓶颈,我们引入了线性内存内核。该内核利用每个词元的 log-sum-exp 统计数据,并将 logits 重新计算融合到向后传递中,以计算精确的 Kullback-Leibler 散度和梯度。在从 4K 扩展到 32K 的 LLaMA2-7B 上,LinearARD 恢复了最先进基线的 98.3% 的短文本性能,同时在长上下文基准上超越了它们。值得注意的是,与 LongReD 和 CPT 所需的 \textbf{256M} 词元相比,我们的方法仅使用 \textbf{4.25M} 训练词元即可实现这些结果。我们的代码可在 https://github.com/gracefulning/LinearARD 获取。