Search papers, labs, and topics across Lattice.
This paper investigates knowledge distillation as a method for transferring long-context retrieval capabilities to smaller student models without long-context pre-training. They show that logit-based knowledge distillation effectively transfers positional information, particularly when using phase-wise RoPE scaling. Analysis of packed repeated token sequences reveals how positional perturbations propagate through transformer layers, influencing the teacher's output distribution and guiding the student's learning.
Forget expensive long-context pre-training: knowledge distillation can transfer long-context retrieval skills to smaller models using only short-context data.
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.
