Self-Distilled Reasoner: On-Policy Self-Distillation for Large Language Models

TL;DR

OPSD addresses the distribution mismatch and data efficiency challenges in LLM reasoning by enabling a single model to teach itself using privileged information about ground-truth reasoning traces. It employs two policy-conditioned variants from the same model and minimizes token-level divergence along on-policy rollouts, avoiding the need for an external teacher. Across mathematical reasoning benchmarks, OPSD delivers strong performance with 4 to 8x token efficiency relative to GRPO and SFT, particularly at 4B–8B scales. The findings indicate that longer reasoning generations and full-vocabulary distillation provide richer supervision, and they point to future benefits from larger models and curriculum-based training.

Abstract

Knowledge distillation improves large language model (LLM) reasoning by compressing the knowledge of a teacher LLM to train smaller LLMs. On-policy distillation advances this approach by having the student sample its own trajectories while a teacher LLM provides dense token-level supervision, addressing the distribution mismatch between training and inference in off-policy distillation methods. However, on-policy distillation typically requires a separate, often larger, teacher LLM and does not explicitly leverage ground-truth solutions available in reasoning datasets. Inspired by the intuition that a sufficiently capable LLM can rationalize external privileged reasoning traces and teach its weaker self (i.e., the version without access to privileged information), we introduce On-Policy Self-Distillation (OPSD), a framework where a single model acts as both teacher and student by conditioning on different contexts. The teacher policy conditions on privileged information (e.g., verified reasoning traces) while the student policy sees only the question; training minimizes the per-token divergence between these distributions over the student's own rollouts. We demonstrate the efficacy of our method on multiple mathematical reasoning benchmarks, achieving 4-8x token efficiency compared to reinforcement learning methods such as GRPO and superior performance over off-policy distillation methods.

Figures (4)

• Figure 1: Overview of On-Policy Self-Distillation (OPSD): Given a reasoning dataset $\mathcal{S} = \{(x_i, y_i^\star)\}_{i=1}^N$, we instantiate two policies from the same LLM: a student policy$p_S(\cdot \mid x)$ and a teacher policy$p_T(\cdot \mid x, y^\star)$. The student generates an on-policy response $\hat{y} \sim p_S(\cdot \mid x)$. Both policies then evaluate this trajectory to produce next-token distributions $p_S(\cdot \mid x, \hat{y}_{<n})$ and $p_T(\cdot \mid x, y^\star, \hat{y}_{<n})$ at each step $n$. The learning objective minimizes the per-token divergence $D(p_T \| p_S)$ along the student's rollout. Crucially, gradients backpropagate only through the student's logits, allowing the model to self-distil.
• Figure 2: Prompt example for student and teacher policies. Both policies share the same parameters $\theta$ but differ in conditioning context. The teacher receives the ground-truth solution $y^\star$ as privileged information before generation. To ensure a natural transition before evaluating the student's rollout, the teacher is prompted to rationalize and generate its own solution.
• Figure 3: Token Efficiency of OPSD. We compare OPSD and GRPO on Qwen3-4B under the same effective training batch size, reporting average@16 performance as a function of gradient update steps and total generated tokens. Both methods are trained with the same effective batch size in terms of sampled generations per update, but differ in generation length: each generation is capped at 2048 tokens for OPSD and 16384 tokens for GRPO. OPSD achieves comparable or better performance with substantially fewer generated tokens, resulting in lower sampling cost and reduced training time. In this experiment, OPSD can be 4-8$\times$ more token-efficient than GRPO.
• Figure 4: Pass@K performance averaged across four mathematical reasoning benchmarks for Qwen3-4B. We study the effect of the generation length of on-policy sampled student responses in OPSD, comparing 1024, 2048, and 4096 tokens. Longer generations provide more teacher signals. Increasing the generation length from 1k to 2k and 4k consistently improves pass@K, with both 2k and 4k substantially outperforming the 1k setting.