Combining On-Policy Optimization and Distillation for Long-Context Reasoning in Large Language Models

TL;DR

This paper introduces a new training method combining on-policy reinforcement learning and distillation to improve long-context reasoning in large language models, supported by a synthetic dataset.

Contribution

It proposes Distilled Group Relative Policy Optimization (dGRPO), integrating dense guidance from a teacher with on-policy optimization for better long-context performance.

Findings

dGRPO outperforms off-policy methods in long-context tasks.

The LongBlocks dataset enables effective evaluation of long-context reasoning.

Combining policy optimization with distillation improves stability and effectiveness.

Abstract

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.