LinearARD: Linear-Memory Attention Distillation for RoPE Restoration

2026-03-09

arXiv.org

AI Summary

Extending the context window of large language models often degrades performance on short sequences. This work proposes a self-distillation approach that freezes a base RoPE-equipped teacher model and guides a scaled-RoPE student model to preserve attention structure by aligning the row-wise distributions of their query-query, key-key, and value-value autocorrelation matrices. To achieve this efficiently, the authors introduce a linear-memory kernel based on per-token log-sum-exp statistics, which integrates logit recomputation during backpropagation to enable accurate KL divergence computation and gradient estimation without incurring quadratic memory overhead. When scaling LLaMA2-7B’s context length from 4K to 32K, the method recovers 98.3% of the original short-context performance using only 4.25 million training tokens—significantly fewer than the 256 million tokens required by baseline approaches.

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.

Problem

• context window extension
• positional encoding
• performance degradation
• short-text benchmarks
• RoPE

Innovation

• LinearARD
• attention distillation
• RoPE restoration
• linear-memory attention
• self-distillation