English

LinearARD: Linear-Memory Attention Distillation for RoPE Restoration

Computation and Language 2026-04-02 v1 Artificial Intelligence

Abstract

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/QQ/Q, K/KK/K, and V/VV/V self-relation matrices to directly supervise attention dynamics. To overcome the quadratic memory bottleneck of n×nn \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.

Keywords

Cite

@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}
}
R2 v1 2026-07-01T11:46:49.811Z