LinearARD: Linear-Memory Attention Distillation for RoPE Restoration

2026-07-25 · Source: cs.CL updates on arXiv.org · Field: Technology & Digital — Artificial Intelligence & Machine Learning, Software Development & Engineering · Depth: Expert, extended

Summary

LinearARD is a novel self-distillation method designed to restore the original capabilities of Large Language Models (LLMs) whose context windows have been extended using Rotary Position Embeddings (RoPE) scaling. This approach addresses the common issue where RoPE scaling and subsequent Continual Pre-Training (CPT) degrade short-text performance. LinearARD achieves this by enforcing attention-structure consistency, aligning the row-wise distributions of dense Q/Q, K/K, and V/V self-relation matrices with a frozen native-RoPE teacher model. To overcome the quadratic memory bottleneck associated with n x n relation maps, it introduces a linear-memory kernel that computes exact Kullback-Leibler divergence and gradients. On LLaMA2-7B extended from 4K to 32K, LinearARD recovers 98.3% of short-text performance and surpasses state-of-the-art baselines on long-context benchmarks, using only 4.25M training tokens, a 60x reduction compared to 256M tokens required by methods like LongReD and CPT. It also recovers 94.8%, 94.2%, and 95.0% of native short-context average accuracy on LLaMA2-7B, LLaMA3-8B, and Mistral-7B-v0.1, respectively, while achieving strong RULER scores.

Key takeaway

For Machine Learning Engineers extending LLM context windows with RoPE scaling, you should consider LinearARD to efficiently restore model capabilities. This method significantly reduces training costs, requiring 60x fewer tokens (4.25M vs. 256M) than CPT or hidden-state distillation, while recovering over 94% of native short-text performance and improving long-context robustness. Implement its two-stage approach, focusing on QKV relation distillation, to maintain structural integrity and avoid catastrophic forgetting.

Key insights

RoPE-scaled LLM capabilities are restored by directly aligning internal QKV self-relation distributions using a linear-memory KL distillation kernel.

Principles

• Direct attention-structure supervision is superior to hidden-state matching for RoPE restoration.
• Forward KL divergence offers robust gradients for sparse attention distributions.
• Aligning internal relational structure implicitly restores output proficiency.

Method

LinearARD employs a two-pass tiled kernel: first, compute and store row-wise log-sum-exp statistics; then, iterate query/key tiles, recompute logits, reconstruct probabilities, and accumulate exact KL gradients with O(n) memory.

In practice

• Prioritize Q/Q, K/K, and V/V self-relation targets for attention distillation.
• Apply LoRA/QLoRA to attention projection matrices for parameter efficiency.
• Separate short-context restoration from long-context adaptation with lightweight CPT.

Topics

• Long-context LLMs
• Rotary Position Embeddings
• Attention Distillation
• Linear-Memory Kernel
• Self-Distillation
• LLaMA2-7B

Code references

• gracefulning/LinearARD

Best for: Research Scientist, AI Engineer, AI Scientist, Machine Learning Engineer, NLP Engineer

Related on AIssential

• Towards Generalization of Block Attention via Automatic Segmentation and Block Distillation — Semantic segmentation and an efficient distillation framework enable practical, scalable block attention for long-context LLMs.
• Shuffle the Context: RoPE-Perturbed Self-Distillation for Long-Context Adaptation — RoPE-Perturbed Self-Distillation improves LLM long-context understanding by reducing reliance on absolute positional encoding.
• DepthWeave-KV: Token-Adaptive Cross-Layer Residual Factorization for Long-Context KV Cache Compression — Token-adaptive cross-layer residual factorization compresses KV caches, improving long-context LLM inference efficiency and quality without retraining.
• Beyond Softmax: The Future of Attention Mechanisms — Linear attention offers constant memory complexity by replacing softmax with a state matrix and recurrence relation.
• Formalizing and Mitigating Structural Distortion in LLM Attention for Zero-Shot Graph Reasoning — LLM performance on graphs is bottlenecked by structural distortion from sequence linearization, correctable by attention alignment.
• YaRN: Extending RoPE Without Breaking It — YaRN (Yet another RoPE extensioN) is a highly efficient method for extending the context window of pre-trained Large Language Models, such as LLaMA 2, from 4,096 to 128,000 tokens with significantly r…

See Counsel's argued verdicts on the open AI decisions leaders are weighing →

Open in AIssential →

Editorial summary, takeaway, and curation by AIssential. Original article published by cs.CL updates on arXiv.org.