OPSDL: On-Policy Self-Distillation for Long-Context Language Models

Xinsen Zhang; Zhenkai Ding; Tianjun Pan; Run Yang; Chun Kang; Xue Xiong; Jingnan Gu

Source-linked AI summary

OPSDL: On-Policy Self-Distillation for Long-Context Language Models

Xinsen Zhang, Zhenkai Ding, Tianjun Pan, Run Yang, Chun Kang, Xue Xiong, Jingnan Gu

arXiv:2604.17535v1Apr 19, 2026cs.CLcs.AI

TL;DR

OPSDL targets the difficulty of extending effective context length when existing methods require high-quality supervision or sparse rewards. It uses the model’s short-context behavior as an on-policy token-level teacher for long-context generations, and consistently improves long-context benchmarks across model scales and context lengths without degrading short-context performance.

Problem
Long-context optimization is limited by dependence on high-quality training data or sparse sequence-level rewards, which can make training unstable and sample-inefficient.
Method
OPSDL uses the model’s own short-context capability as a self-teacher, applying on-policy token-level reverse-KL supervision to long-context generations.
Results
OPSDL consistently improves RULER and LongBench V2 performance across 7B, 14B, and 32B models, with particularly strong gains at extended context lengths.
Takeaways & Limitations
On-policy token-level self-distillation provides a lightweight alternative to dedicated long-context training while preserving short-context capabilities.

Abstract

from arXiv · show

Extending the effective context length of large language models (LLMs) remains a central challenge for real-world applications. While recent post-training methods have made progress in long-context scaling, they either rely on high-quality supervision data or sparse sequence-level rewards, leading to unstable and inefficient optimization. We propose OPSDL, an On-Policy Self-Distillation method for enhancing the Long-context capabilities of LLMs. Unlike other recent self-distillation methods that inject privileged information and rely on the model's in-context learning ability to act as a teacher, OPSDL leverages the model's own inherently strong short-context capability as a self-teacher to supervise its own generation in long-context scenarios. The model first generates responses conditioned on the full long-context, then the self-teacher provides per-token supervision signals via point-wise reverse KL divergence under the relevant extracted short-context. This dense token-level signal encourages faithful use of relevant evidence and mitigates hallucinations induced by irrelevant context. We evaluate OPSDL on long-context benchmarks across a range of models from 7B to 32B parameters. Results show consistent and substantial improvements across varying context lengths, outperforming standard post-training approaches such as SFT and DPO with higher sample efficiency. Notably, these gains are achieved without degrading general short-context performance. These findings highlight the effectiveness of OPSDL as a scalable and stable approach for long-context learning.

1 Introduction

OPSDL addresses the gap between maximum and effective context length by replacing data- or reward-dependent optimization with on-policy self-distillation. It uses short-context behavior as a token-level teacher and improves long-context performance across models and context lengths.

Motivation: Effective context length remains a central challenge despite models accepting increasingly long inputs.The gap affects long-document comprehension, repository-level code understanding, and multi-hop reasoning over scattered evidence.
Motivation: Existing long-context methods depend on high-quality data or sparse sequence-level rewards, making optimization unstable and sample-inefficient.They may also require frozen verifiers or learned reward models, increasing system complexity.
Approach: OPSDL distills a model’s short-context behavior into its long-context behavior using on-policy token-level reverse KL divergence.The model serves as both student and teacher, without external supervision or reward models.
Results: Across models and context lengths, OPSDL consistently outperforms standard instruction-tuned models, with larger gains as context length grows.It also reaches performance comparable to, and sometimes exceeding, a model explicitly trained for million-token contexts.
Contributions: The method provides dense supervision, avoids external auxiliary components, and improves long-context performance without degrading short-context performance.Evaluations cover models from 7B to 32B parameters.

2 Related Work

OPSDL is presented against post-training and on-policy self-distillation methods that use supervision, rewards, or privileged information. Its framework instead uses the model’s own short-context behavior as the teacher signal.

Post-Training for Long-Context Modeling: Long-context post-training includes supervised fine-tuning and preference optimization, but these approaches face data-quality, distribution-shift, or reward-construction constraints.LongReward uses multidimensional LLM feedback with DPO, while LongPO uses short-context generations as positive samples.
On-Policy Distillation and Self-Distillation: On-policy distillation emphasizes token-level supervision over student-generated trajectories rather than only sequence-level feedback.This line of work includes generalized knowledge distillation and subsequent self-distillation methods.
OPSDL Positioning: OPSDL differs from related self-distillation methods by using the model’s short-context capability as its teacher rather than injecting privileged context or feedback.The supplied framework overview depicts long-context generation supervised by the same model under an extracted short context.

3 Method

OPSDL constructs long-context, short-context, and query triplets, then aligns on-policy long-context generations with short-context predictions through token-level reverse-KL signals. The teacher shares the student’s parameters and co-evolves during training.

Problem Formulation: OPSDL uses a long context CL, an extracted short context CS containing core information, and a query Q answerable from both contexts.Queries are generated from CS so the short-context teacher and long-context student can be meaningfully compared.
Self-Alignment: The same policy produces a more calibrated short-context distribution, which serves as an anchor for self-aligning long-context behavior.The short context avoids distracting noise and fits within the model’s well-trained context window.
Token-Level Self-Distillation: OPSDL samples responses from the long-context policy and computes token-level advantages from short-context teacher and long-context student log probabilities.The advantage is defined as their log-probability ratio and estimates the gradient of point-wise reverse KL divergence.
Token-Level Self-Distillation: Positive advantages raise tokens favored by the short-context teacher, while negative advantages identify tokens associated with hallucination or irrelevant long-context content.Near-zero advantages indicate agreement and produce negligible gradient signal.
Policy Optimization: Only tokens where long- and short-context policies diverge receive substantial gradient signal, avoiding noise from uniform supervision.The policy-gradient objective is optimized over responses sampled from the long-context distribution.
Training Procedure: Unlike a frozen-teacher distillation setup, OPSDL uses the same parameters for teacher and student, so the teacher co-evolves with the policy.Training iterates through triplet construction, on-policy rollout, token-level advantage estimation, and policy updates.

4 Experiment

The experiments evaluate OPSDL against long-context post-training baselines across model scales, benchmarks, and context lengths. OPSDL delivers stronger long-context performance, remains stable across realistic and synthetic evaluations, narrows the gap with dedicated 1M-token models, and largely preserves short-context capabilities.

Experiment Setup: OPSDL is trained from raw long documents using (CL, CS, I) triplets without human annotation, and is compared with Long-SFT and LongPO on Qwen2.5-Instruct models.The evaluation spans 7B, 14B, and 32B backbones, plus officially long-context-optimized 1M variants.
Experiment Setup: The evaluation uses RULER for synthetic long-context capacity and LongBench V2 for reasoning over long documents across domains and context lengths.Table 1 reports results across different model sizes; LongBench V2 cells average four runs at sampling temperature 0.1.
Main Results: At 7B, OPSDL raises performance from 51.68 to 56.61, achieving the largest improvement over the base instruct model among trainable methods.The method consistently achieves the largest improvement at each evaluated scale.
Main Results: At 128K tokens, OPSDL improves over the base instruct model by +48.70, +34.25, and +30.29 points at 7B, 14B, and 32B, respectively.The gains become increasingly important beyond 64K tokens, where base instruct models degrade substantially.
Main Results: OPSDL concentrates optimization on tokens where long-context behavior deviates from the short-context anchor, producing more targeted learning than uniform sequence-level supervision.The reported improvements are consistently larger than those achieved by Long-SFT and LongPO.
Main Results: OPSDL narrows the RULER-average gap to the 1M variant from 13.10 to 3.94 at 7B and from 10.57 to 3.28 at 14B.The officially optimized 1M variants still achieve the highest absolute scores.
Main Results: OPSDL improves both RULER and LongBench V2, while training stably across model sizes where LongPO failed to converge at 14B and 32B.This pattern covers synthetic long-context capability and realistic long-document reasoning.
Main Results: Across MMLU, ARC-C, Hellaswag, and Winogrande, OPSDL’s average drop is about 1.3 points, while MT-Bench remains virtually unchanged.Long-SFT shows a larger average degradation of 3–4 points.

5 Conclusion

OPSDL uses a model’s short-context capability as a dynamically co-evolving teacher for long-context generation. Across RULER and LongBench V2, it outperforms Long-SFT and LongPO, improves extended-context performance, narrows the gap with 1M-token models, and minimally degrades short-context capabilities.

Conclusion: OPSDL uses on-policy token-level reverse-KL supervision from a model’s short-context capability to guide its long-context generation.The framework removes the need for external reward models, human-annotated data, and offline preference pairs.
Conclusion: Across RULER and LongBench V2 at 7B, 14B, and 32B, OPSDL consistently outperforms Long-SFT and LongPO, with especially large gains at extended context lengths.It also substantially narrows the gap with officially long-context-optimized Qwen2.5-Instruct-1M variants.
Conclusion: OPSDL preserves short-context capabilities with minimal degradation while providing a simple, scalable, and effective paradigm for long-context post-training.The conclusion frames the method as both effective and scalable within the evaluated setting.