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

TL;DR

OPSDL introduces an on-policy self-distillation technique that enhances long-context understanding in large language models by leveraging their own short-context capabilities, leading to improved performance without sacrificing short-term skills.

Contribution

This paper presents OPSDL, a novel self-distillation method that uses the model's own short-context abilities as a teacher to improve long-context performance efficiently and stably.

Findings

OPSDL outperforms standard post-training methods like SFT and DPO.

The method improves long-context performance across models from 7B to 32B parameters.

OPSDL maintains short-context performance while enhancing long-context understanding.

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.