LinearARD: Linear-Memory Attention Distillation for RoPE Restoration

TL;DR

LinearARD is a novel self-distillation method that restores Rotary Position Embeddings in large language models with linear memory, achieving high performance on long-context tasks using minimal training data.

Contribution

The paper introduces LinearARD, a memory-efficient self-distillation technique that aligns attention-structure distributions to restore RoPE in LLMs, outperforming existing methods on long-context benchmarks.

Findings

Recovers 98.3% of short-text performance on LLaMA2-7B with extended context.

Uses only 4.25 million training tokens, significantly less than prior methods.

Surpasses state-of-the-art long-context benchmarks with minimal training data.

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/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.