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 · University of Science and Technology Beijing, Beijing, China · University College London, London, United Kingdom
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×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 A computed from QK⊤ (shown for reference), and (b–d) Q/Q, K/K, and V/V relation matrices computed from QQ⊤ , KK⊤ , and VV⊤ (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 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×
–
24.65
6.18
19.73
58.29
36.20
69.31
34.60
24.91
34.23
62.6
–
–
32K
PI 8×
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×
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×
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×
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× 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×
ARD (no CPT)
36.3
63.2
27.2
79.4
43.8
78.1
36.3
49.4
51.7
PI 8×
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× . 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× . 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 ×
OK
OK
2K
20.74
11.83
1.75 ×
OK
OK
4K
81.25
43.82
1.85 ×
OK
OK
8K
326.93
172.56
1.90 ×
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)
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× 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×
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×
32K
ARD
36.81
64.72
27.94
77.54
43.80
78.62
35.98
49.57
51.87
94.85
PI 16×
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×
64K
ARD
32.72
59.89
26.26
77.55
42.20
77.53
36.30
49.15
50.20
91.79
PI 32×
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× , 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× , 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× , 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.
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
( ×10−7 )
Mean ( ×10−4 )
Max ( ×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.
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.
Zichong Li, Chen Liang, Liliang Ren +3
Work is done during internship at Microsoft. · Georgia Institute of Technology · Microsoft.
Rotary Position Embedding (RoPE) is widely adopted in Transformers to encode positional information, yet standard implementations enforce a uniform frequency schedule and scaling across all attention heads. Using simplified retrieval tasks and length generalization scenarios, we show -- both empirically and theoretically -- that heads with different functional roles require distinct frequency ranges and attention scaling factors to operate effectively. Ignoring this structure leads to suboptimal utilization of embedding dimensions and degraded performance, particularly under long-context settings. To address these limitations, we propose AdaRoPE, which equips each attention head with learnable rotation frequencies and attention scaling factors. Pretrained LLMs with AdaRoPE consistently outperform existing RoPE variants, including partial RoPE and NoPE baselines. For context extension, we further show that uniform frequency and attention scaling, used in methods such as YaRN, are suboptimal. By applying head-specific scaling, AdaRoPE enables better context extension while better preserving short-context performance in both the extrapolation setting and the long-context continued pretraining setting. These results highlight the importance of optimizing rotary position embedding at the level of individual attention heads.
Shaowen Wang, Yuke Zheng, Tansheng Zhu +4
Tsinghua University · Hunyuan Team, Tencent · Xiongan AI Institute +1
Extending the context window of language models typically requires expensive long-context pre-training, posing significant challenges for both training efficiency and data collection. In this paper, we present evidence that long-context retrieval capabilities can be transferred to student models through logit-based knowledge distillation, even when training exclusively on packed short-context samples within a long-context window. We provide comprehensive insights through the lens of Rotary Position Embedding (RoPE) and establish three key findings. First, consistent with prior work, we show that phase-wise RoPE scaling, which maximizes rotational spectrum utilization at each training stage, also achieves the best long-context performance in knowledge distillation setups. Second, we demonstrate that logit-based knowledge distillation can directly enable positional information transfer. Using an experimental setup with packed repeated token sequences, we trace the propagation of positional perturbations from query and key vectors through successive transformer layers to output logits, revealing that positional information systematically influences the teacher's output distribution and, in turn, the distillation signal received by the student model. Third, our analysis uncovers structured update patterns in the query state during long-context extension, with distinct parameter spans exhibiting strong sensitivity to long-context training.
Patrick Huber, Ernie Chang, Chinnadhurai Sankar +4