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

Zichong Li; Chen Liang; Liliang Ren; Tuo Zhao; Yelong Shen; Weizhu Chen

cs.CLApr 15, 2026

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

Authors: Zichong Li, Chen Liang, Liliang Ren, Tuo Zhao, Yelong Shen, Weizhu Chen

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

Abstract

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.

Figures & tables

Panel of Figure 1 — Figure 1: (a) Relation between answer position and accuracy on the NIAH multikey-2 task from RULER ( Hsieh et al., 2024 ) with 64k-token sequences. We evaluate the ProLong-64k-base model ( Gao et al., 2025 ) as causal LM baseline and model trained with our method. (b) Overview of our proposed objective.

Figure 2: Analysis of attention distance patterns on a 64k-token sequence from Prolong ( Gao et al., 2025 ) dataset. For both the baseline model and our model, we inspect attention scores from the final 256 query tokens. We consider attention weights above a threshold (

10^{-3}

) and measure the signed distance between query and key indices. The figure shows the histogram of attention distances in the 24th layer.

Method	RULER		HELMET
	32K	64K	RAG	ICL	LongQA	Rerank	Avg.
Standard	85.23	57.25	59.08	86.30	30.30	1.51	44.30
LongCE	86.11	52.51	54.71	85.80	31.64	11.20	45.84
PoSE	87.20	67.47	57.22	86.30	30.23	8.85	45.65
Ours (cyclic shift)	87.33	71.30	58.34	86.80	30.63	8.70	46.12
Ours (skip)	87.87	69.29	59.12	87.65	32.41	12.34	47.88

Table 1: Average performance of trained Llama-3-8B-Instruct under different sequence lengths on RULER and HELMET (64K) benchmarks.

Method	Avg.	S-1	S-2	S-3	MK-1	MK-2	MK-3	MV	MQ	VT	CWE	FWE	QA-1	QA-2
Standard	57.3	76.8	99.4	67.2	84.2	31.0	0.0	85.2	87.1	70.4	7.9	54.7	46.8	33.6
LongCE	52.5	98.2	99.2	74.6	93.2	18.4	0.0	90.3	89.9	4.9	4.9	18.9	54.6	35.6
PoSE	67.5	99.6	100	75.6	90.2	70.6	6.6	86.2	90.5	61.4	22.7	79.8	53.8	40.2
Ours (cyclic shift)	71.3	100	100	81.6	86.8	88.4	27.4	93.1	93.6	67.1	16.1	80.7	51.6	41.0
Ours (skip)	69.3	100	100	75.0	94.0	82.6	20.4	88.6	94.5	43.9	26.1	82.5	53.2	40.0

Table 2: Performance of Llama-3-8B-Instruct experiments on RULER tasks (64K context).

Method	RULER		HELMET
	128K	256K	RAG	ICL	LongQA	Rerank	Avg.
Standard	74.45	67.01	51.54	86.08	56.02	13.57	51.80
LongCE	76.87	68.01	52.88	87.24	56.00	15.53	52.90
PoSE	75.47	65.48	52.04	87.16	55.47	12.51	51.79
Ours (cyclic shift)	78.37	67.91	53.00	86.52	54.48	13.58	51.89
Ours (skip)	77.51	68.10	53.96	87.40	57.23	13.39	52.99

Table 3: Average performance of Qwen-3-4B under different sequence lengths on RULER and HELMET (128K) benchmarks.

Figure 3: Positional sensitivity on RULER NIAH MultiKey-2 for Qwen-3-4B.

Method	128K	256K
Standard	42.03	23.90
LongCE	46.54	26.91
PoSE	58.21	43.18
Ours (cyclic shift)	60.89	45.33
Ours (skip)	59.43	44.74

Table 4: Extrapolation performance of Llama using YaRN at extended context lengths.

Method	512K	1M
Standard	58.11	45.31
LongCE	60.52	47.53
PoSE	58.00	44.62
Ours (cyclic shift)	60.27	49.56
Ours (skip)	60.73	50.41
Ours (skip) + LongCE	60.13	50.02

Table 5: Extrapolation performance of Qwen using YaRN at ultra-long context lengths.

Method	RULER		HELMET	LBv2
	128K	256K	HELMET	LBv2
Standard	75.52	64.44	52.61	27.63
LongCE	76.95	65.92	53.36	31.73
PoSE	76.22	63.88	52.53	28.83
Ours (skip)	77.28	66.86	53.48	29.40
+ LongCE	78.68	67.15	54.02	32.22

Table 6: Average performance of Qwen-3-4B after SFT under different sequence lengths on RULER, HELMET, and LongBenchV2 (LBv2) benchmarks.

Model	Method	MMLU	HellaSwag	Winogrande	OpenBookQA	Avg.
Llama	Base	64.7	75.6	71.4	43.0	63.7
	Standard	61.7	78.7	72.5	43.4	64.1
	LongCE	62.9	78.7	72.2	43.8	64.4
	PoSE	62.2	78.7	72.0	43.6	64.1
	Ours (cyclic shift)	61.4	78.4	72.4	44.2	64.1
	Ours (skip)	61.2	78.3	72.2	43.6	63.8

Table 7: Downstream task performance across models under different methods.

Method	Avg.
Baseline	47.9
Ours	59.2
w/ forward KL	57.8
w/ CLM	56.5
w/o random skip	55.1
Dropout	48.9

Table 8: Ablation of objective design and comparison with generic regularization on Llama-3-8B-Instruct (RULER, 64K).

Skipping Range	RULER (64K)
32K	59.63
64K	59.23
128K	58.88

Table 9: Effect of skipping range

Y

on RULER performance at 64K context (Llama-3-8B-Instruct).

Method	Avg.
Baseline	47.9
Ours	59.2
NoPE	29.5
Random dilation	53.4
Chunked permutation	46.2

Table 10: Ablation of RoPE perturbation schemes on Llama (RULER, 64K).

Method	RULER (64K)
Standard	47.9
Ours (skip, $\lambda=0.2$ )	54.0
Ours (skip, $\lambda=0.5$ )	60.0
Ours (skip, $\lambda=1$ )	59.2
Ours (skip, $\lambda=2$ )	58.9

Table 11: Effect of the KL coefficient

\lambda

on Llama-3-8B-Instruct (RULER, 64K).

Panel of Figure 4 — Figure 4: RULER performance versus wall-clock training time under matched compute.

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Method	Avg.	S-1	S-2	S-3	MK-1	MK-2	MK-3	MV	MQ	VT	CWE	FWE	QA-1	QA-2
Standard	67.01	100	98.8	97.2	85.8	41.2	25.6	78.15	89.4	95.4	8.4	78.2	41.0	32.0
LongCE	68.01	100	97.8	99.0	84.2	45.4	17.2	90.1	89.3	90.8	8.4	78.2	47.4	36.4
PoSE	65.48	100	99.8	97.6	84.8	46.8	23.6	73.3	86.7	86.8	1.7	77.6	41.2	31.4
Ours (cyclic shift)	67.92	100	99.4	97.8	86.6	54.2	58.0	75.2	87.3	78.6	1.2	71.3	39.2	34.2
Ours (skip)	68.10	100	100	98.4	82.6	50.2	39.6	75.15	88.5	93.1	2.5	77.6	43.6	34.0
Ours (skip) + LongCE	68.95	100	97.2	99.4	83.2	45.8	28.4	96.5	90.0	85.8	4.6	83.1	46.4	36.0

Appendix

Table 12: Performance of Qwen-3-4B on RULER tasks at 256K context.

Method	RULER		HELMET
	32K	64K	Avg.
Standard	79.94	52.48	46.73
Ours (skip)	85.27	65.28	48.80

Appendix

Table 13: Average performance of mix-length-trained Llama-3-8B-Instruct under different sequence lengths on RULER and HELMET (64K) benchmarks.

Method	64K NIAH-position
Baseline	44.0
LongCE	56.2
PoSE	48.2
Ours (shift)	42.0
Ours (skip)	59.4

Appendix

Table 14: Accuracy on the 64K NIAH-position benchmark (multi-value ordinal retrieval).

Method	RULER 32K	RULER 64K	HELMET
Standard	85.26 (0.08)	57.32 (0.12)	44.26 (0.17)
Ours (cyclic shift)	87.29 (0.08)	71.25 (0.14)	46.17 (0.09)
Ours (skip)	87.82 (0.09)	69.32 (0.09)	47.81 (0.13)

Appendix

Table 15: Evaluation variance across random seeds on Llama-3-8B-Instruct. We report mean performance with standard deviation in parentheses.

Long-Context Language Modeling

Long-context language modeling studies how language models can use, retain, and efficiently process information across very long sequences.

See all 180 papers

LinearARD: Linear-Memory Attention Distillation for RoPE Restoration
Mar 9, 2026Ning Yang, Hengyu Zhong, Wentao Wang +3Attention-Guided Distillation Memory-Efficient Attention
Institute of Automation, Chinese Academy of Sciences, Beijing, China · Southwest University, Chongqing, China · Dalian University of Technology, Dalian, Liaoning, China +3
EndPrompt: Efficient Long-Context Extension via Terminal Anchoring
May 14, 2026Han Tian, Luxuan Chen, Xinran Chen +10Long-Context Language Modeling LLM Fine-Tuning
1Nankai University · 2Baidu Inc. · 3Shanghai Jiao Tong University +1
Self-Guided Test-Time Training for Long-Context LLMs
Jul 10, 2026Xinyu Zhu, Zhe Xu, Xiaohan Wei +10Long-Context Language Modeling Test-Time Training
Meta AI · University of Virginia

Mar 9, 2026J/K move · Enter open · S save

LinearARD: Linear-Memory Attention Distillation for RoPE Restoration

Ning Yang, Hengyu Zhong, Wentao Wang, Baoliang Tian+2

Institute of Automation, Chinese Academy of Sciences, Beijing, China · Southwest University, Chongqing, China · Dalian University of Technology, Dalian, Liaoning, China+3

Attention-Guided Distillation Memory-Efficient Attention Long-Context Language Modeling Language Model Distillation Knowledge Distillation Efficient Attention+1

Open paper AskPDF arXiv

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.

Abstract

Long-Context Language Modeling

Explore similar work

LinearARD: Linear-Memory Attention Distillation for RoPE Restoration

EndPrompt: Efficient Long-Context Extension via Terminal Anchoring

Self-Guided Test-Time Training for Long-Context LLMs