arXiv:2604.00004cs.CLcs.AI2026-04

用轻量方法恢复长文本模型的短文本能力

LinearARD: Linear-Memory Attention Distillation for RoPE Restoration

论文配图:LinearARD: Linear-Memory Attention Distillation for RoPE Restoration
图 1 · 摘自论文原文
  • 通过注意力结构一致性自蒸馏恢复旋转位置编码
  • 32K上下文扩展后短文本性能恢复98.3%
  • 仅需425万训练样本,远低于其他方法

大语言模型扩展上下文窗口通常依赖位置编码缩放和轻量持续预训练,但会破坏原有模型在标准短文本任务上的性能。本文提出LinearARD,一种自蒸馏方法,通过与固定原生RoPE教师模型保持注意力结构一致,恢复经RoPE缩放的学生模型。该方法不匹配隐藏状态,而是对齐$Q/Q$、$K/K$、$V/V$自关系矩阵的行分布,直接监督注意力动态。为解决$n \times n$关系图的二次内存瓶颈,引入线性内存核,利用每标记的log-sum-exp统计并融合梯度重构,实现精确KL散度与梯度计算。在将LLaMA2-7B从4K扩展至32K时,LinearARD恢复了98.3%的基准短文本性能,并在长文本任务上超越现有方法。关键优势在于仅需425万训练样本,显著低于LongReD和CPT所需的2.56亿。

原文摘要 · Abstract (English)

The extension of context windows in Large Language Models is typically facilitated by scaling positional encodings followed by lightweight Continual Pre-Training (CPT). While effective for processing long sequences, this paradigm often disrupts original model capabilities, leading to performance degradation on standard short-text benchmarks. We propose LinearARD, a self-distillation method that restores Rotary Position Embeddings (RoPE)-scaled students through attention-structure consistency with a frozen native-RoPE teacher. Rather than matching opaque hidden states, LinearARD aligns the row-wise distributions of dense $Q/Q$, $K/K$, and $V/V$ self-relation matrices to directly supervise attention dynamics. To overcome the quadratic memory bottleneck of $n \times n$ relation maps, we introduce a linear-memory kernel. This kernel leverages per-token log-sum-exp statistics and fuses logit recomputation into the backward pass to compute exact Kullback-Leibler divergence and gradients. On LLaMA2-7B extended from 4K to 32K, LinearARD recovers 98.3\% of the short-text performance of state-of-the-art baselines while surpassing them on long-context benchmarks. Notably, our method achieves these results using only \textbf{4.25M} training tokens compared to the \textbf{256M} tokens required by LongReD and CPT. Our code is available at https://github.com/gracefulning/LinearARD.

位置编码自蒸馏长文本轻量化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。