LinearARD: Linear-Memory Attention Distillation for RoPE Restoration

Ning Yang; Hengyu Zhong; Wentao Wang; Baoliang Tian; Haijun Zhang; Jun Wang

cs.CLMar 9, 2026

LinearARD: Linear-Memory Attention Distillation for RoPE Restoration

Authors: Ning Yang, Hengyu Zhong, Wentao Wang, Baoliang Tian, Haijun Zhang, Jun Wang

Organizations: Institute of Automation, Chinese Academy of Sciences, Beijing, China · Southwest University, Chongqing, China · Dalian University of Technology, Dalian, Liaoning, China · Tianjin University, Tianjin, China ·

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 4.25M training tokens compared to the 256M tokens required by LongReD and CPT. Our code is available at https://github.com/gracefulning/LinearARD.

Figures & tables

Figure 1: LinearARD pipeline. A frozen native-RoPE teacher provides final-layer dense row-wise relation distributions ( $Q/Q$ , $K/K$ , and $V/V$ ), and the RoPE-scaled student is restored by minimizing relation KL with an exact linear-memory kernel.

Figure 2: Final-layer standard attention vs. QKV self-relations. (a) Attention matrix $\mathbf{A}$ computed from $\mathbf{Q}\mathbf{K}^{\top}$ (shown for reference), and (b–d) Q/Q, K/K, and V/V relation matrices computed from $\mathbf{Q}\mathbf{Q}^{\top}$ , $\mathbf{K}\mathbf{K}^{\top}$ , and $\mathbf{V}\mathbf{V}^{\top}$ (Eq. 4 ). LinearARD distills the row-wise relation distributions in (b–d) from the final Transformer layer by aligning teacher vs. student rows via forward KL (Eq. 3 ) under the same mask $\mathbf{M}$ .

Figure 3: Memory scaling of the linear-memory KL distillation Kernel in Sec. 3.3 as a function of sequence length.

				Short-text accuracy (%)									Summary
Model	CW	PE	Method	MMLU	LAMB.	MathQA	BoolQ	OBQA	PIQA	SIQA	ARC-C	Avg.	Rec.	RULER	Tokens
LLaMA2-7B	4K	–	–	45.99	71.18	29.92	78.23	44.00	78.73	40.28	49.23	54.69	100.0	–	–
	32K	PI $8\times$	–	24.65	6.18	19.73	58.29	36.20	69.31	34.60	24.91	34.23	62.6	–	–
	32K	PI $8\times$	CPT	37.35	67.30	27.60	77.52	44.00	78.73	36.18	48.89	52.20	95.4	59.6	256M
	32K	PI $8\times$	LongReD	38.40	67.57	27.60	78.35	44.00	78.73	36.39	50.94	52.75	96.4	59.7	256M
	32K	PI $8\times$	LinearARD	39.54	67.82	29.31	76.82	44.20	78.40	39.82	51.11	53.38	97.6	65.1	256M

Table 1: Main results on short-text benchmarks and long-context robustness. Accuracy in % is reported on eight short-text benchmarks and their mean, denoted Avg. CW denotes context window length, PE denotes the RoPE scaling configuration, Rec. denotes recovery relative to the native model, and RULER is averaged over 8K/16K/32K evaluations. Near-zero RULER values for 4.25M-token CPT/LongReD indicate unstable or unparseable long-context outputs. LAMB. denotes LAMBADA, OBQA denotes OpenBookQA, and ARC-C denotes ARC-Challenge.

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

PE	Variant	MMLU	LAMB.	BoolQ	OBQA	Avg.	Rec.
–	Native	45.99	71.18	78.23	44.00	59.85	100.00
PI $8\times$	w/o restore	24.65	6.18	58.29	36.20	31.33	52.35
	ARD (attention)	31.63	65.17	76.76	41.80	53.84	89.95
	ARD (QKV rel.)	36.81	64.72	79.05	43.60	56.05	93.64
	ARD (+logit KL)	36.38	68.85	79.17	43.60	57.00	95.24
	ARD (full FT QKV)	30.42	58.66	75.38	43.60	52.02	86.91

Appendix

Table 2: ARD component ablation on LLaMA2-7B with PI

8\times

and no CPT. Avg. is the mean of MMLU, LAMBADA, BoolQ, and OpenBookQA. Rec. is the recovery ratio relative to the native model. This diagnostic table compares relation targets under the same ablation protocol; the main LinearARD setting in Table 1 uses final-layer relation alignment for efficiency.

		Short-text accuracy (%)
PE	Method	MMLU	LAMB.	MathQA	BoolQ	OBQA	PIQA	SIQA	ARC-C	Avg.
–	Native	46.0	71.2	29.9	78.2	44.0	78.7	40.3	49.2	54.7
PI $8\times$	ARD (no CPT)	36.3	63.2	27.2	79.4	43.8	78.1	36.3	49.4	51.7
PI $8\times$	ARD + CPT	36.8	64.7	27.9	77.5	43.8	78.6	36.0	49.6	51.9

Appendix

Table 3: Short-context side of the CPT ablation on LLaMA2-7B with PI

8\times

. ARD denotes relation distillation on PG19, and CPT denotes the subsequent lightweight SlimPajama stage. LAMB., OBQA, and ARC-C denote LAMBADA, OpenBookQA, and ARC-Challenge.

Method	R@8K	R@16K	R@32K	Avg.
ARD (no CPT)	46.5	42.5	21.4	36.8
ARD + CPT	74.9	67.4	47.3	63.2

Appendix

Table 4: Long-context side of the CPT ablation on LLaMA2-7B with PI

8\times

. Avg. is the mean over 8K, 16K, and 32K.

Seq. Len.	Naive (ms)	Kernel (ms)	Speedup	Naive	Kernel
1K	5.30	3.32	1.60 $\times$	OK	OK
2K	20.74	11.83	1.75 $\times$	OK	OK
4K	81.25	43.82	1.85 $\times$	OK	OK
8K	326.93	172.56	1.90 $\times$	OK	OK
16K	OOM	683.54	–	OOM	OK
32K	OOM	2549.54	–	OOM	OK

Appendix

Table 5: Single-operator latency for the naive dense implementation and the IO-aware kernel. The naive implementation runs out of memory at 16K and above.

Context	Step Time (s)						Throughput (tokens/s)
Context	All-32	Last-16	Last-8	Last-4	Last-2	Last-1	All-32	Last-16	Last-8	Last-4	Last-2	Last-1
1K	0.77	0.62	0.54	0.50	0.49	0.48	1327.2	1655.6	1882.2	2034.6	2070.9	2144.0
2K	1.84	1.30	1.02	0.88	0.81	0.78	1113.8	1581.4	2005.7	2320.1	2519.2	2629.0
4K	5.50	3.47	2.45	1.94	1.68	1.56	744.7	1181.9	1672.5	2108.4	2431.5	2628.1
8K	19.08	11.09	7.11	5.11	4.11	3.61	429.4	738.4	1152.9	1602.2	1991.3	2267.1
16K	76.32	38.35	22.53	14.64	10.69	8.70	214.7	427.2	727.2	1119.3	1532.6	1882.9

Appendix

Table 6: End-to-end ARD training step time and throughput under different layer scopes. All-32 denotes all-layer relation distillation, while Last-

k

distills only the final

k

layers.

Figure 4: Layerwise contribution of ARD on LLaMA2-7B with PI

8\times

and no CPT. Later layers dominate the restoration signal, while early layers contribute little or can hurt recovery.

PE	CW	Method	MMLU	LAMB.	MathQA	BoolQ	OBQA	PIQA	SIQA	ARC-C	Avg.	Rec.
–	4K	Native	45.99	71.18	29.92	78.23	44.00	78.73	40.28	49.23	54.69	100.0
PI $8\times$	32K	w/o restore	24.65	6.18	19.73	58.29	36.20	69.31	34.60	24.91	34.23	62.60
PI $8\times$	32K	ARD	36.81	64.72	27.94	77.54	43.80	78.62	35.98	49.57	51.87	94.85
PI $16\times$	64K	w/o restore	23.22	0.66	19.53	44.04	32.60	63.71	34.54	24.83	30.39	55.57
PI $16\times$	64K	ARD	32.72	59.89	26.26	77.55	42.20	77.53	36.30	49.15	50.20	91.79
PI $32\times$	128K	w/o restore	22.98	0.35	18.99	38.44	32.40	61.86	34.08	24.32	29.18	53.35

Appendix

Table 7: Short-context stress test under larger PI scaling factors on LLaMA2-7B. The 64K/128K ARD runs use 288 ARD steps, corresponding to 4.5M distillation tokens, and do not include CPT.

Backbone	Method	8K	16K	32K	Avg. all
LLaMA2-7B	CPT (4.25M)	0.0625	0.0192	0.0383	0.040
LLaMA2-7B	LongReD (4.25M)	0.0458	0.0292	0.0300	0.035
LLaMA3-8B	CPT (4.25M)	0.0475	0.0442	0.0375	0.043
LLaMA3-8B	LongReD (4.25M)	0.0408	0.0400	0.0508	0.044
Mistral-7B-v0.1	CPT (4.25M)	0.0100	0.0108	0.0100	0.010
Mistral-7B-v0.1	LongReD (4.25M)	0.0150	0.0075	0.0075	0.010

Appendix

Table 8: Per-length RULER averages for budget-matched CPT and LongReD controls. Values remain near zero under the same 4.25M-token budget used by LinearARD.

Task	CPT (256M)			LongReD (256M)			LinearARD+CPT (2M)			LinearARD (no CPT)
	8K	16K	32K	8K	16K	32K	8K	16K	32K	8K	16K	32K
NIAH-S1	95.0	91.0	82.0	95.0	94.0	88.0	100.0	100.0	100.0	97.0	90.0	76.0
NIAH-S2	95.0	96.0	78.0	95.0	97.0	86.0	100.0	100.0	96.0	80.0	74.0	30.0
NIAH-S3	91.0	84.0	55.0	95.0	91.0	60.0	88.0	88.0	43.0	45.0	52.0	15.0
NIAH-K1	93.0	90.0	62.0	95.0	94.0	63.0	96.0	96.0	76.0	66.0	58.0	25.0
NIAH-K2	76.0	68.0	23.0	76.0	60.0	20.0	74.0	70.0	29.0	47.0	43.0	13.0

Appendix

Table 9: Per-task RULER scores in % for LLaMA2-7B with PI scaling factor

8\times

, evaluated at 8K, 16K, and 32K. NIAH denotes Needle-In-A-Haystack variants. VT, CWE, FWE, and QA(SQuAD) follow RULER task naming. Avg. averages over subtasks at a fixed length. Avg. all averages over all subtasks and lengths and matches the aggregate RULER reported in Table 1 . LinearARD denotes our method. LinearARD without CPT removes the lightweight CPT stage.

Task	CPT (256M)			LongReD (256M)			LinearARD+CPT (2M)
	8K	16K	32K	8K	16K	32K	8K	16K	32K
NIAH-S1	99.0	100.0	99.0	89.0	89.0	91.0	99.0	99.0	99.0
NIAH-S2	100.0	100.0	99.0	86.0	84.0	62.0	96.0	96.0	88.0
NIAH-S3	100.0	98.0	94.0	81.0	73.0	63.0	66.0	71.0	64.0
NIAH-K1	99.0	95.0	93.0	84.0	82.0	56.0	98.0	90.0	73.0
NIAH-K2	99.0	80.0	60.0	92.0	83.0	66.0	95.0	96.0	61.0

Appendix

Table 10: Per-task RULER scores in % for LLaMA3-8B with PI scaling factor

4\times

, evaluated at 8K, 16K, and 32K. Avg. averages over subtasks at a fixed length. Avg. all averages over all subtasks and lengths and matches the aggregate RULER reported in Table 1 . LinearARD denotes our method.

Task	CPT (256M)			LongReD (256M)			LinearARD+CPT (2M)
	8K	16K	32K	8K	16K	32K	8K	16K	32K
NIAH-S1	100.0	100.0	73.0	94.0	88.0	50.0	93.0	91.0	51.0
NIAH-S2	100.0	100.0	60.0	91.0	90.0	63.0	96.0	89.0	64.0
NIAH-S3	89.0	68.0	53.0	85.0	70.0	46.0	73.0	49.0	53.0
NIAH-K1	98.0	93.0	62.0	95.0	90.0	45.0	92.0	88.0	54.0
NIAH-K2	80.0	40.0	3.0	96.0	93.0	49.0	94.0	92.0	54.0

Appendix

Table 11: Per-task RULER scores in % for Mistral-7B-v0.1 with PI scaling factor

4\times

Figure 5: Training dynamics of distillation objectives. The plot shows the optimization target ( QKV-loss , blue) alongside monitored but non-optimized metrics ( Attention-loss , gold; Logit-loss , maroon). The synchronous decline of all three curves suggests a strong causal link between internal attention drift and output distribution shift.

Length	Forward Error	Backward Error
Length	( $\times 10^{-7}$ )	Mean ( $\times 10^{-4}$ )	Max ( $\times 10^{-2}$ )
256	$4.9$	$1.8$	$0.6$
512	$4.9$	$1.7$	$0.7$
1024	$4.7$	$1.5$	$0.8$
2048	$4.6$	$1.2$	$0.9$
4096	$4.9$	$1.0$	$1.0$

Appendix

Table 12: Kernel numerical verification. Relative deviations from a materialized dense reference where memory permits are reported. Forward error is the FP32 loss mean absolute error divided by the mean output magnitude. Backward errors are computed from BF16 gradients using the same normalization, reported as mean and maximum relative deviations.

Attention-Guided Distillation

See all 11 papers

Shuffle the Context: RoPE-Perturbed Self-Distillation for Long-Context Adaptation
Apr 15, 2026Zichong Li, Chen Liang, Liliang Ren +3Long-Context Language Modeling Language Model Robustness
Work is done during internship at Microsoft. · Georgia Institute of Technology · Microsoft.
AdaRoPE: Not All Attention Heads Should Rotate and Scale Equally
Jun 5, 2026Shaowen Wang, Yuke Zheng, Tansheng Zhu +4Rotary Positional Embeddings Self-Attention
Tsinghua University · Hunyuan Team, Tencent · Xiongan AI Institute +1
Short Data, Long Context: Distilling Positional Knowledge in Transformers
Apr 7, 2026Patrick Huber, Ernie Chang, Chinnadhurai Sankar +4Long-Context Language Modeling LLM Context Management
Meta Reality Labs

Apr 15, 2026J/K move · Enter open · S save

Shuffle the Context: RoPE-Perturbed Self-Distillation for Long-Context Adaptation

Zichong Li, Chen Liang, Liliang Ren, Tuo Zhao+2

Work is done during internship at Microsoft. · Georgia Institute of Technology · Microsoft.

Long-Context Language Modeling Language Model Robustness Self-Distillation

Open paper AskPDF arXiv

Large language models (LLMs) increasingly operate in settings that require reliable long-context understanding, such as retrieval-augmented generation and multi-document reasoning. A common strategy is to fine-tune pretrained short-context models at the target sequence length. However, we find that standard long-context adaptation can remain brittle: model accuracy depends strongly on the absolute placement of relevant evidence, exhibiting high positional variance even when controlling for task format and difficulty. We propose RoPE-Perturbed Self-Distillation, a training regularizer that improves positional robustness. The core idea is to form alternative "views" of the same training sequence by perturbing its RoPE indices -- effectively moving parts of the context to different positions -- and to train the model to produce consistent predictions across views via self-distillation. This encourages reliance on semantic signals instead of brittle position dependencies. Experiments on long-context adaptation of Llama-3-8B and Qwen-3-4B demonstrate consistent gains on long-context benchmarks, including up to 12.04% improvement on RULER-64K for Llama-3-8B and 2.71% on RULER-256K for Qwen-3-4B after SFT, alongside improved length extrapolation beyond the training context window.

Abstract

Attention-Guided Distillation

Explore similar work

Shuffle the Context: RoPE-Perturbed Self-Distillation for Long-Context Adaptation

AdaRoPE: Not All Attention Heads Should Rotate and Scale Equally

Short Data, Long Context: Distilling Positional Knowledge in Transformers