Short Data, Long Context: Distilling Positional Knowledge in Transformers

Patrick Huber; Ernie Chang; Chinnadhurai Sankar; Rylan Conway; Igor Fedorov; Md Rifat Arefin; Adithya Sagar

cs.CLApr 7, 2026

Short Data, Long Context: Distilling Positional Knowledge in Transformers

Authors: Patrick Huber, Ernie Chang, Chinnadhurai Sankar, Rylan Conway, Igor Fedorov, Md Rifat Arefin, Adithya Sagar

Organizations: Meta Reality Labs

Abstract

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.

Figures & tables

Figure 1: The three dimensions we explore in this work: RoPE embeddings, the objective function (knowledge distillation), and the underlying data. By combining these components, the student model acquires long-context abilities despite being trained only on short-context data.

Panel of Figure 2 — Figure 2: (Top) Each dimension pair undergoes incremental rotation as position $m$ increases, encoded by a unique rotation angle $m\theta_{i}$ . (Center) Different dimension pairs rotate at different frequencies: low-index pairs distinguish nearby tokens, while high-index pairs encode long-range positional relationships. (Bottom) The base parameter $\theta$ controls how quickly angles accumulate: a smaller $\theta$ causes faster wrapping (limiting context length), while a larger $\theta$ spreads rotations across longer sequences.

	Student	Teacher
Model	Random Init	Llama-4 Scout
Layers	30	48
Attn Heads	20	40
KV Heads	4	8
Dimension	1,280	5,120
Hidden Dim	6,144	16,384

Table 1: Student and Teacher Model Specifications

Percentile	10	20	30	40	50	60	70	80	90	99
Word Count per Sample	154	220	293	375	466	571	710	935	1,441	4,955

Table 2: Word count percentiles of our pre-training dataset.

Figure 3: Schematic visualization of scaled and unscaled RoPE $\theta$ in the short-context (Phase 1) and long-context (Phase 2) pre-training stages.

Panel of Figure 4 — Figure 4: Performance for different RoPE $\theta$ values and training methods (CE vs. KD) across Needle in a Haystack (left) and RULER (right).

Panel of Figure 5 — Figure 5: Comparison of the average pre-training performance (left, in %) and training loss across different $\theta$ settings (right) for KD and CE objectives.

Figure 6: Schematic visualization of the experimental approach. We repeat a fixed sequence of $n=2{,}048$ tokens $N=64$ times to generate an analysis sequence of shape $(bsz,64,dim)$ , filling the complete context of $131{,}072$ tokens. By using identical content at every position, we isolate positional perturbations from semantic and syntactic differences.

Panel of Figure 7 — Figure 7: Distance between position index 0 and subsequent positions for a single hidden-dimension element in the query hidden state, measured before (Pre-RoPE) and after (Post-RoPE) applying rotary positional embeddings.

Panel of Figure 8 — Figure 8: Cosine similarity between the query hidden state at position index 0 and subsequent positions, measured before (Pre-RoPE) and after (Post-RoPE) applying rotary positional embeddings.

Figure 9: Per-layer cosine similarity between position index 0 and subsequent position indices (offset by 2,048 tokens each). Each curve represents a different layer, showing progressive amplification of positional divergence through the model.

Figure 10: Output layer top-10 token ranking and probability mass distribution across sequence positions. Letters represent token indices and their rank.

Panel of Figure 11 — Figure 11: Euclidean distance between Phase 1 and Phase 2 pre-trained hidden states per hidden-state element for Layer 0 (top) and Layer 30 (bottom).

Panel of Figure 12 — Figure 12: Euclidean distance between Phase 1 and Phase 2 pre-trained hidden states per sequence element for Layer 0 (top) and Layer 30 (bottom).

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

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.
Position Encoding in Transformers: From Absolute and Relative Methods to Rotary Position Embeddings and Long-Context Scaling
Aug 9, 2026Jiguo LiRotary Positional Embeddings Long-Context Modeling
RoPE Distinguishes Neither Positions Nor Tokens in Long Contexts, Provably
May 15, 2026Yufeng Du, Phillip Harris, Minyang Tian +5Rotary Positional Embeddings Self-Attention
University of Illinois at Urbana-Champaign · USA · University of Bonn +3

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

Long-Context Language Modeling

Explore similar work

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

Position Encoding in Transformers: From Absolute and Relative Methods to Rotary Position Embeddings and Long-Context Scaling

RoPE Distinguishes Neither Positions Nor Tokens in Long Contexts, Provably