PaperlayerView exact revision on arXiv

Source-linked AI summary

ROSD: Reflective On-Policy Self-Distillation for Language Model Reasoning across Domains

Ziqi Zhao, Xinyu Ma, Liu Yang, Yujie Feng, Daiting Shi, Jingzhou He, Xin Xin, Zhaochun Ren, Xiao-Ming Wu

arXiv:2605.28014v1cs.CLcs.LG

TL;DR

Existing OPSD methods provide dense supervision but can imitate training-domain reference trajectories and overwrite valid reasoning prefixes, limiting generalization. ROSD uses a self-reflector to generate corrective guidance and localize the first error, then distills only where correction is needed. Across benchmarks and backbone models, it strengthens in-domain reasoning and substantially improves out-of-domain generalization over standard OPSD, though it remains below GRPO out of domain.

  • Problem

    Existing OPSD methods can encourage reference-solution imitation and overwrite valid reasoning prefixes, limiting in-domain gains and out-of-domain generalization.

  • Method

    ROSD uses reflection to extract a corrective idea and locate the first erroneous span, then applies self-distillation only from that localized error onward.

  • Results

    ROSD delivers stronger in-domain reasoning and substantially better out-of-domain generalization than standard OPSD across benchmarks and backbone models.

  • Takeaways & Limitations

    Error-aware, selectively applied supervision can repair faulty reasoning while preserving valid prefixes during reasoning post-training.

  • Takeaways & Limitations

    ROSD still falls short of GRPO in out-of-domain settings, and broader evaluation is needed to establish its generality.

Abstract

from arXiv · show

On-policy self-distillation (OPSD) improves the reasoning performance of large language models (LLMs) by providing dense token-level supervision for on-policy rollouts. However, existing OPSD methods often yield limited gains on in-domain reasoning and generalize poorly to out-of-domain problems. We identify two key causes: conditioning the self-teacher on a verified solution encourages imitation of training-domain reference trajectories rather than error-specific correction, and applying distillation to the full response can overwrite valid reasoning prefixes and reinforce overfitting. We propose Reflective On-policy Self-Distillation (ROSD), a framework that turns reference-solution imitation into targeted reasoning correction through reflection-guided, error-localized distillation. For each rollout, ROSD uses a self-reflector to extract a corrective idea and locate the first erroneous span. The corrective idea guides the self-teacher toward targeted supervision, while the localized error span restricts distillation to where correction is needed. This design corrects flawed reasoning while preserving valid prefixes. Experiments on multiple in-domain and out-of-domain reasoning benchmarks show that ROSD yields stronger in-domain reasoning performance overall and substantially better out-of-domain generalization than standard OPSD. Code is available at https://github.com/ZiqiZhao1/ROSD.

1 Introduction

Existing post-training methods provide either sparse outcome feedback or dense supervision that can imitate reference trajectories and overwrite valid reasoning. ROSD instead uses reflection-guided, error-localized distillation to correct flawed reasoning while preserving valid prefixes, improving in-domain reasoning and out-of-domain generalization.

  • Motivation: RLVR provides outcome-level feedback, while OPSD supplies dense token-level guidance along on-policy rollouts.OPSD conditions a self-teacher on a correct solution and computes token-wise divergence along the sampled rollout.
  • Motivation: Standard OPSD can become unstable in-domain and rapidly degrade out-of-domain because it imitates reference trajectories rather than correcting rollout-specific errors.Its verified solution does not identify where the student fails or what local correction is needed.
  • Motivation: Full-response distillation can overwrite correct reasoning prefixes and inject training-domain preferences into otherwise valid reasoning.The underlying error may be local, but standard OPSD applies its loss across the entire response.
  • ROSD: ROSD uses a self-reflector to extract a corrective idea and an exact error quote identifying the first erroneous span.The corrective idea guides targeted teacher supervision, while the error quote localizes the update region.
  • Results: ROSD improves in-domain reasoning while better preserving out-of-domain generalization across multiple benchmarks and model backbones.The framework shifts self-distillation from reference-solution imitation toward selective reasoning correction.

2 Preliminaries and Related Work

The related work contrasts response-level RLVR with on-policy self-distillation, which provides denser token-level supervision but raises questions about teacher context and where to apply the loss. ROSD addresses both choices by constructing informative context and selecting tokens for distillation.

  • RL-Based Post-Training: GRPO is a representative RLVR algorithm that estimates group-relative advantages from sampled responses without a separate critic model.Its policy objective uses importance-sampling ratios and clipping to constrain updates.
  • RL-Based Post-Training: Because GRPO assigns one response-level advantage to every token, it cannot distinguish valid portions of incorrect trajectories from erroneous portions.This limits token-specific correction during optimization.
  • On-Policy Self-Distillation: On-policy self-distillation samples rollouts from the student policy and uses an auxiliary-context teacher to provide token-level supervision along those rollouts.The teacher predicts next-token distributions conditioned on extra information and the response prefix, while the student matches them with token-level divergence.
  • On-Policy Self-Distillation: Prior self-distillation methods can use successful rollouts as teacher context, with KL, reverse KL, or JSD objectives and fixed or EMA-updated teachers.These design choices vary the divergence and teacher-update mechanism without changing the on-policy setup.
  • ROSD: ROSD differs by constructing more informative distillation context and applying the loss only to selected tokens.Its comparison with prior methods spans training setting, supervision source, and evaluation scope.

3 Proposed Method: ROSD

ROSD uses reflection to turn reference-solution supervision into targeted correction, then localizes distillation to the rollout segment requiring revision. It preserves valid reasoning prefixes while guiding the student with corrective information.

  • Error-Focused Self-Reflection: For each problem, ROSD separates correct and wrong rollouts and pairs each wrong rollout with the shortest correct rollout from the same group.The selected correct rollout provides a contrastive basis for identifying where the wrong rollout deviates.
  • Error-Focused Self-Reflection: ROSD generates error-focused reflections containing a corrective idea and an exact quote marking the first erroneous span.For wrong rollouts, the quote is later aligned to the student rollout for localization.
  • Error-Focused Self-Reflection: Correct-rollout reflections summarize why the reasoning is valid and which key idea makes the final answer correct, avoiding direct full-solution imitation.This preserves useful reasoning information from successful student rollouts while avoiding reuse of the entire rollout as teacher context.
  • Quote-Localized Self-Distillation: The corrective idea conditions the self-teacher, while the error quote remains a localization anchor and the reflector reuses the self-teacher’s model weights.The quote is aligned back to the student rollout before selective distillation is applied.
  • Quote-Localized Self-Distillation: ROSD restricts distillation to the suffix beginning at the identified error, preventing unnecessary updates to valid prefixes.If the quote cannot be matched, the method falls back to full-response distillation for that rollout.
  • Quote-Localized Self-Distillation: The masked token-level objective compares the student with a reflection-conditioned self-teacher, using all response tokens for correct rollouts and error-starting tokens for wrong rollouts.The objective uses Jensen–Shannon divergence, and the mask determines which tokens contribute to the loss.

4 Experiments

ROSD is evaluated against GRPO and SDPO across model scales, datasets, in-domain and out-of-domain settings, and training dynamics. It achieves stronger in-domain averages, preserves substantially better cross-domain performance than SDPO, and combines fast convergence with more stable training.

  • In-Domain Results: ROSD reaches the best in-domain averages: 72.83% on Qwen3-4B and 73.45% on Qwen3-8B.These averages improve over GRPO and SDPO across both model scales.
  • Out-of-Domain Results: ROSD outperforms SDPO on the out-of-domain average across all training datasets and both model scales.GRPO remains strongly robust out of domain, while standard self-distillation can impose stronger token-level constraints and encourage overfitting.
  • Out-of-Domain Results: After ToolUse training on Qwen3-4B, ROSD maintains 41.31% OOD average performance while SDPO drops to 2.88%.The comparison uses science-question-answering and mathematical-reasoning test domains that differ substantially from ToolUse.
  • Training Dynamics: ROSD converges faster than the baselines and maintains in-domain performance more stably, whereas SDPO often improves early before declining.Figure 3 reports mean@16 test scores on Material and Chemistry, with shaded variance over 16 sampled responses.
  • Training Dynamics: Dense self-distillation accelerates early rollout-accuracy optimization, but SDPO’s rising rollout accuracy can diverge from declining test performance on Material.ROSD retains the fast-convergence benefit while achieving better in-domain generalization.
  • Ablation Study: Removing reflection improves OOD performance over SDPO, while removing localization improves both settings; full ROSD performs best across scales and evaluation settings.The ablation indicates that reflection and localized distillation provide complementary benefits.
  • Error Localization: ROSD’s error-quote match rate stays around 0.5, while normalized error position moves later as training progresses.The later error positions help explain why full-response distillation can become harmful late in training.
  • Efficiency: ROSD matches SDPO in average per-step training time, becomes more efficient later, and produces shorter, more stable responses.Further training-efficiency and response-length details are provided in Appendix C.

5 Conclusion

ROSD reframes self-distillation as targeted error correction rather than imitation of entire reference solutions. Reflection supplies corrective guidance and localizes supervision, producing stronger in-domain reasoning and substantially better out-of-domain generalization than standard OPSD.

  • Conclusion: ROSD uses reflection to identify corrective guidance and localize supervision to erroneous spans.This repairs faulty reasoning steps while preserving valid prefixes.
  • Conclusion: Across benchmarks and backbone models, ROSD achieves stronger in-domain reasoning and substantially better out-of-domain generalization than standard OPSD.The conclusion emphasizes error-aware, selectively applied supervision for reasoning post-training.

Limitations

The paper identifies broader evaluation and cross-domain robustness as remaining limitations. ROSD still trails GRPO in out-of-domain settings despite improving substantially over standard OPSD.

  • Limitations: ROSD requires evaluation in broader settings to establish its generality beyond the reported benchmarks and backbone scales.The authors report consistent gains across multiple benchmarks and scales but identify broader evaluation as future work.
  • Limitations: ROSD remains below GRPO in out-of-domain performance, indicating that cross-domain generalization still needs improvement.This limitation persists despite stronger overall in-domain performance than GRPO and substantial OOD gains over standard OPSD.

Ethical Considerations

The paper notes that stronger reasoning capabilities may be misused for harmful or malicious objectives. It therefore emphasizes deployment with rigorous safety alignment mechanisms.

  • Ethical Considerations: Stronger reasoning capabilities may be exploited for harmful or malicious objectives, so deployment should include rigorous safety alignment mechanisms.The experiments use publicly available open-source datasets under their corresponding licenses and intended purposes.

A Dataset Statistics

The experiments use ScienceQA and ToolUse for training and evaluation, while AIME2024 serves only as an additional held-out reasoning benchmark.

  • ScienceQA and ToolUse are used for both training and evaluation, whereas AIME2024 is held out and not used for training.

B Implementation Details

Experiments use the verl distributed actor-rollout pipeline with FSDP actor training and vLLM rollout generation on eight NVIDIA A800 80G GPUs. Training runs for 10 epochs on science question-answering tasks and 5 epochs on ToolUse, with a batch size of 32.

  • Experiments use verl with FSDP actor training and vLLM-based rollout generation on 8 NVIDIA A800 80G GPUs.
  • Training lasts 10 epochs for science question-answering tasks and 5 epochs for ToolUse, using a training batch size of 32.

C Additional Results

ROSD produces shorter responses and maintains a more stable response-length trend during training than GRPO and SDPO.

  • ROSD produces relatively shorter responses with a more stable response-length trend, while GRPO and SDPO responses lengthen during training.

C.2 Training efficiency

ROSD has training time comparable to SDPO and lower than GRPO, becoming more efficient later as baseline rollout lengths increase while its response length remains stable.

  • Training efficiency: ROSD's overall training time is comparable to SDPO and lower than GRPO despite introducing an additional reflector.
  • Training efficiency: ROSD becomes more efficient later in training because baseline rollout lengths increase while its response length remains stable.
Loading 2605.28014v1…