๐ค AI Summary
This work addresses the challenges of training instability and inefficiency when scaling the context length of large language models, which existing approaches mitigate using high-quality supervised data or sparse rewards. The authors propose an in-policy self-distillation method that leverages the modelโs own strong short-context capabilities as a โself-teacherโ to generate dense supervision signals via token-wise reverse KL divergence. This guides the model to attend to relevant evidence and suppress hallucinations during long-context reasoning, without requiring external data or privileged information, and integrates seamlessly into standard language model training pipelines. Experiments across 7B to 32B parameter models demonstrate that the method significantly outperforms post-training strategies such as supervised fine-tuning (SFT) and direct preference optimization (DPO), achieving higher sample efficiency while preserving short-context general performance.
๐ Abstract
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.