Combining On-Policy Optimization and Distillation for Long-Context Reasoning in Large Language Models
For researchers and practitioners working on long-context LLM alignment, this work provides a practical recipe that mitigates key limitations of existing post-training methods.
The paper addresses the challenge of adapting LLMs to long-context reasoning tasks, where existing methods suffer from exposure bias, instability, or inefficiency. The proposed dGRPO method combines on-policy optimization (GRPO) with on-policy distillation (OPD) to achieve more stable and effective long-context alignment, outperforming off-policy and sparse-reward baselines while preserving short-context performance.
Adapting large language models (LLMs) to long-context tasks requires post-training methods that remain accurate and coherent over thousands of tokens. Existing approaches are limited in several ways: 1) off-policy methods such as supervised fine-tuning (SFT) and knowledge distillation (KD) suffer from exposure bias and limited recovery from model-generated errors over long horizons; 2) on-policy reinforcement learning methods such as Group Relative Policy Optimization (GRPO) better align training with model-generated states, but are unstable and sample-inefficient due to sparse rewards; 3) on-policy distillation (OPD) provides dense token-level guidance, but does not directly optimize arbitrary reward signals. In this paper, we propose Distilled Group Relative Policy Optimization (dGRPO), a method for long-context reasoning that augments GRPO with dense guidance from a stronger teacher via OPD. We also introduce LongBlocks, a synthetic long-context dataset spanning multi-hop reasoning, contextual grounding, and long-form generation. We conduct extensive experiments and ablations comparing off-policy training, sparse-reward GRPO, and our combined approach, leading to an improved recipe for long-context alignment. Overall, our results show that combining outcome-based policy optimization with knowledge distillation in a single objective provides a more stable and effective path to long-context reasoning, while preserving short-context capabilities.