PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Siyan Zhao, Zhihui Xie, Mengchen Liu, Jing Huang, Guan Pang, Feiyu Chen, Aditya Grover

arXiv:2601.18734v3cs.LGcs.CL

TL;DR

Reasoning post-training methods either use off-policy data, sparse rewards, or a separate teacher, leaving a gap for self-contained dense supervision on student trajectories. OPSD uses one LLM in teacher and student roles with different contexts, and it performs on par with or better than GRPO while being more sample-efficient and outperforming off-policy distillation or SFT. Its evidence is limited to models up to 8B parameters and does not establish whether the trend continues at larger scales.

  • Problem

    Off-policy distillation suffers distribution mismatch, while on-policy distillation typically requires a separate teacher and does not explicitly exploit ground-truth reasoning solutions.

  • Method

    OPSD uses a single LLM as teacher and student: the teacher sees privileged answers, while the student samples trajectories from only the question and receives per-token divergence supervision.

  • Results

    OPSD performs on par with or better than GRPO, outperforms off-policy distillation and SFT, and has significantly better sample efficiency than GRPO.

  • Takeaways & Limitations

    Self-distillation can use privileged reasoning information to provide dense supervision on a model’s own rollouts without an external teacher.

  • Takeaways & Limitations

    Experiments are limited to models up to 8B parameters, leaving continuation of the trend beyond that scale open.

Abstract

from arXiv · show

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, we introduce On-Policy Self-Distillation (OPSD), a learning algorithm where a single LLM acts as both teacher and student with 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 superior token efficiency compared to reinforcement learning methods and better performance over off-policy distillation methods. Code repo: https://github.com/siyan-zhao/OPSD.

1. Introduction

Existing reasoning post-training methods trade off on-policy relevance, dense supervision, and teacher-model requirements. OPSD addresses this trade-off by using one LLM as both teacher and student, with privileged answer context guiding dense supervision on student-generated trajectories.

  • Motivation: RLVR uses costly sampling and sparse sequence-level rewards, while SFT and traditional distillation face distribution mismatch or off-policy data limitations.Traditional knowledge distillation provides dense token-level supervision but trains on fixed trajectories; SFT suffers exposure bias and weaker generalization.
  • Motivation: OPSD asks whether a capable LLM can teach itself by rationalizing correct solutions that are unavailable to the student during generation.The teacher receives privileged information such as a ground-truth answer or reference chain-of-thought, while the student sees only the problem.
  • Method: At each student-generated step, OPSD minimizes divergence between teacher and student next-token distributions using dense supervision that backpropagates only through student logits.The divergence may be forward KL, reverse KL, or JSD.
  • Method: OPSD instantiates teacher and student policies from a single LLM, preserving on-policy training by sampling trajectories exclusively from the student.The teacher conditions on both the problem and privileged answer, whereas the student conditions only on the problem.
  • Contributions and evaluation: OPSD introduces per-token pointwise KL clipping to stabilize training and evaluates the method on three competition-level mathematical reasoning tasks.The evaluation reports matching GRPO with improved token efficiency and outperforming supervised fine-tuning.

2. Background

Autoregressive distillation can be off-policy or on-policy, while RLVR optimizes student-generated responses with sparse sequence-level rewards. OPSD is positioned as combining on-policy trajectories with dense feedback without an external teacher.

  • Knowledge distillation: Traditional supervised distillation averages teacher–student token-distribution divergence over fixed reference outputs.Both autoregressive teacher and student define token-level distributions over the vocabulary.
  • Knowledge distillation: Off-policy distillation suffers distribution mismatch because inference exposes the student to partial sequences different from those used during training.This mismatch can produce compounding errors during autoregressive generation.
  • Method comparison: Table 1 characterizes OPSD as combining on-policy training with dense feedback without requiring an external teacher model.This is the method-level contrast highlighted in the table caption.
  • On-policy distillation: On-policy distillation trains on sequences sampled from the student, obtaining dense token-level feedback from a teacher on those samples.This connects distillation with imitation learning and combines on-policy relevance with dense supervision.
  • RLVR: GRPO samples groups of responses, assigns binary correctness rewards, and gives every token in each response the same group-normalized advantage.Its objective also uses clipped policy updates and a reverse KL penalty.
  • RLVR: RLVR provides only sequence-level feedback, and identical rewards across all samples make advantages zero, preventing policy updates despite sampling cost.These limitations motivate denser token-level supervision.

3. Methods

OPSD uses one language model as both teacher and student by varying their conditioning contexts: the teacher accesses a reference solution, while the student generates on-policy trajectories from the problem alone. Training distills the teacher’s token distributions along those trajectories, with optional clipping or sampled-token policy-gradient objectives.

  • Teacher and student policies: OPSD instantiates teacher and student policies from the same model, eliminating the need for a separate teacher model.The policies share parameters θ and differ only in conditioning context.
  • Teacher and student policies: The teacher conditions on the problem and reference solution, whereas the student observes only the problem statement at inference-matched context.The reference solution may include chain-of-thought reasoning and serves as privileged information during training.
  • On-policy sampling from the student: The student samples on-policy responses, and both policies evaluate each student-generated prefix to produce next-token distributions for distillation.The teacher evaluates the student trajectory while conditioning additionally on the reference solution.
  • Full-vocabulary logit distillation: Training minimizes token-wise divergence between teacher and student distributions along student rollouts, using full-vocabulary measures such as forward KL, reverse KL, or JSD.Gradients backpropagate only through the student logits, while the teacher acts as a fixed target.
  • Per-token pointwise divergence clipping: Pointwise clipping limits the influence of highly divergent stylistic vocabulary tokens so training signals are less dominated by style than by mathematically meaningful tokens.The imbalance is attributed to skewed divergence contributions across vocabulary entries.
  • Alternative objective: An alternative sampled-token objective uses reverse-KL-derived advantages in a policy-gradient update, providing dense trajectory-level shaping without matching the full vocabulary.The advantage is treated as constant with respect to θ, yielding the usual policy-gradient form.

4. Experiments

Experiments evaluate OPSD against SFT and GRPO, then test divergence objectives, generation styles, KL clipping, generation length, and distillation variants. OPSD matches or exceeds GRPO while using substantially fewer sampled tokens, and ablations identify forward KL, TM-off students with TM-on teachers, clipping, and full-vocabulary divergence as effective choices.

  • Main Results: OPSD matches or exceeds GRPO and outperforms SFT across competition-level mathematical reasoning benchmarks.The comparison covers Qwen3-1.7B, Qwen3-4B, and Qwen3-8B settings.
  • Main Results: 1024 sampled tokens per problem and one rollout suffice for OPSD, whereas GRPO uses 8 rollouts of 16k tokens each.OPSD checkpoints converge within 100 steps, while GRPO can degrade later from entropy collapse.
  • Effect of Divergence Objective: Forward KL raises AIME25 Avg@12 from 36.7 to 43.9 at step 50, while reverse KL and JSD provide limited or negative gains.All divergence objectives use the same pointwise clipping scheme, and forward KL is used thereafter.
  • Generation Styles and KL Clipping: A TM-off student paired with a TM-on teacher gives the largest KL on math tokens and achieves the best downstream performance across model sizes.Stylistic tokens can contribute disproportionately large KL values, motivating pointwise clipping.
  • Generation Styles and KL Clipping: Per-token pointwise KL clipping stabilizes training and prevents performance degradation on AIME24 for Qwen3-1.7B.Clipping controls heavy-tailed divergence contributions from stylistic tokens.
  • Effect of Generation Length: Increasing student generation length from 1024 to 4096 tokens does not consistently improve performance on AIME25 or AIME24.The analysis attributes this to later tokens becoming more predictable and receiving less penalty.
  • Learning Objective Comparison: Full-vocabulary divergence consistently improves over sampled-token distillation when Qwen3-4B uses a 2048-token distillation budget.The full-vocabulary variant matches complete token distributions, whereas sampled-token distillation uses a teacher-shaped policy-gradient objective.

5. Related Work

Prior work improves LLM reasoning through self-training, on-policy distillation, SFT, and RL, while OPSD extends self-training by using privileged information to guide the student on its own trajectories.

  • LLM Self-Training: Context distillation uses one underlying model as teacher and student, giving the teacher privileged context and training the context-free student on generated outputs.
  • LLM Self-Training: Reasoning self-training methods such as ReST and STaR generate rationales with hints or answers, filter successful trajectories, and fine-tune on them as hard distillation.
  • On-Policy Distillation: On-policy distillation trains students on their own sampled trajectories while a teacher provides per-token guidance through KL-based or related objectives.
  • On-Policy Distillation: OPSD explores whether privileged answer information lets an LLM use its reasoning capability to guide a weaker version of itself toward improved reasoning.
  • On-Policy Distillation: On-policy training paradigms also appear in robotics and deep reinforcement learning, including DAgger, where humans supervise states visited by the student policy.
  • Improving LLM Reasoning through SFT and RL: SFT on reasoning traces performs strongly but may rely on memorization rather than robust generalization, whereas RL optimizes outcome-based objectives.

6. Conclusion

The paper presents OPSD as a framework for post-training LLMs on reasoning tasks. Experiments report better performance than off-policy distillation and SFT, performance on par with or better than GRPO, and significantly better sample efficiency than GRPO.

  • OPSD is presented as a simple framework for post-training large language models on reasoning tasks.
  • OPSD achieves better performance than off-policy distillation and SFT, and performs on par with or better than GRPO.
  • OPSD exhibits significantly better sample efficiency than GRPO.

7. Impact Statement

The paper aims to advance machine learning by improving the efficiency of training language models for reasoning tasks. It states that the method reduces computational costs compared with existing reinforcement learning approaches and anticipates no specific negative societal consequences.

  • The work aims to advance the field of machine learning.
  • The method improves training efficiency for language models on reasoning tasks and reduces computational costs compared with existing reinforcement learning approaches.
  • The authors do not foresee specific negative societal consequences from the method.

A. Limitations and Future Directions

The experiments are limited to models up to 8B parameters, and the paper identifies verification, problem difficulty, and curriculum design as directions for extending OPSD.

  • Limitations: Experiments are limited to models up to 8B parameters because of computational constraints.
  • Limitations: The paper leaves open whether the observed trend continues at scales beyond 8B parameters.
  • Future Directions: The framework does not explicitly use correctness verification of generated answers, leaving room for additional learning objectives beyond distribution matching.
  • Limitations: If reasoning problems exceed the model’s comprehension threshold, the teacher cannot provide meaningful supervision even with ground-truth solutions.
  • Future Directions: The paper suggests curriculum learning and adaptive curricula that track the frontier of model capabilities for more challenging reasoning tasks.

B. Experimental Details

The experiments compare SFT, GRPO, and OPSD under documented training and evaluation configurations, using Thinking-Mode-off students with Thinking-Mode-on teachers for the main OPSD runs.

  • Configuration selection: The student-generation setup with thinking mode off and teacher thinking mode on produces the highest KL signal on math-related tokens.This finding motivates the configuration selected for the main OPSD experiments.
  • Experimental setup: Main OPSD experiments use a Thinking-Mode-off student and Thinking-Mode-on teacher configuration.The paper reports this configuration as the main experimental setup.
  • Configurations: Training and evaluation configurations for SFT, GRPO, and OPSD are provided in Tables 7, 6, and 8, respectively.These tables document the configurations used for the experiments.
  • Hyperparameters: The clipping parameter τ was not tuned, leaving potential performance gains within the same 100-step budget.The paper notes that optimizing τ may improve performance.
  • Infrastructure: All experiments use 8 A100 or H100 GPUs, gradient checkpointing, Flash Attention 2, AdamW, and bfloat16 precision.For OPSD, full-vocabulary logit distillation is used unless otherwise stated.

C. Token Category Definitions

The analysis divides generated tokens into stylistic and mathematical categories using predefined keyword lists, then examines their per-token KL divergence.

  • Category construction: Tokens are categorized into style and math groups using predefined keyword lists.These categories support the per-token KL divergence analysis in Section 4.3.1.
  • Style Tokens: Style tokens include discourse, reasoning, transition, modality, and evaluation terms such as “because,” “therefore,” “first,” and “check.”The supplied style list also includes modal verbs, evaluative adjectives, and procedural terms.
  • Math Tokens: Math tokens include technical terms for quantities, operations, algebraic objects, and optimization, such as “equation,” “factor,” “integer,” and “bound.”The supplied math list covers concepts including logarithms, roots, coefficients, variables, and inequalities.
  • Analysis use: The objective analysis interprets OPSD as a dense, token-level reward method and contrasts it with sequence-level STaR.This connects token categorization to the paper’s broader analysis of token-level learning signals.

D.1. STaR as Sequence-Level Policy-Gradient

STaR can be interpreted as a policy-gradient method whose binary answer reward supervises entire sampled trajectories. OPSD instead assigns dense token-level rewards along student rollouts, enabling finer credit assignment even when the final answer is wrong.

  • D.1. STaR as Sequence-Level Policy-Gradient: STaR models a joint distribution over a latent rationale r and final answer y, sampling the rationale before predicting the answer.The model’s rationale is latent in the policy-gradient formulation.
  • D.1. STaR as Sequence-Level Policy-Gradient: STaR uses the indicator reward R(y) = 1(y = y⋆), and the log-derivative trick converts its expected return into a policy gradient.This formulation underlies the filtering step used in STaR.
  • D.1. STaR as Sequence-Level Policy-Gradient: The indicator reward discards gradient contributions from sampled rationales that do not produce the correct answer y⋆.When all sampled trajectories are incorrect, the learning signal vanishes.
  • D.2. OPSD as Dense-Reward Policy Gradient: OPSD samples a trajectory ˆy from the student and defines a per-token reward as log pT(ˆyn | x, y⋆, ˆy<n) − log pS(ˆyn | x, ˆy<n).The reward compares the privileged teacher’s preference for the sampled token with the student’s preference.
  • D.2. OPSD as Dense-Reward Policy Gradient: OPSD’s reward is dense, providing a learning signal at every token position regardless of final-answer correctness.Unlike STaR’s sequence-level indicator, this enables fine-grained credit assignment when the final answer is wrong.
Loading 2601.18734v3…