💬 NLP

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

This paper introduces Reflective On-policy Self-Distillation (ROSD), a framework that enhances language model reasoning across domains by using a self-reflector to localize errors and guide targeted, error-specific distillation, thereby overcoming the overfitting and poor generalization limitations of existing on-policy self-distillation methods.

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

Published 2026-05-28

📖 4 min read☕ Coffee break read

CC BY 4.0

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

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are teaching a student how to solve complex math problems. You have a smart teacher (the AI model) who is trying to learn by practicing on its own.

The Problem: The "Copycat" Teacher

In the past, when this student made a mistake, the standard way to teach them was to show them the perfect, final answer from a textbook and say, "Look at this! You need to write your whole answer exactly like this."

The paper calls this On-Policy Self-Distillation (OPSD). The problem is that the student starts acting like a "copycat."

They memorize the style, not the fix: Instead of learning why they made a mistake, they just copy the teacher's specific words and formatting.
They break what was already right: If the student got the first half of the answer correct but messed up the second half, the teacher forces them to rewrite the entire thing to match the textbook. This accidentally overwrites the good parts they already knew, making them forget their own valid reasoning.

This works okay for the specific practice problems, but when the student faces a new type of problem (a different domain), they fail because they just memorized the "look" of the answers, not the logic.

The Solution: ROSD (The "Reflective Tutor")

The authors propose a new method called ROSD (Reflective On-Policy Self-Distillation). Think of ROSD not as a teacher who hands you a finished essay, but as a critical editor who uses a magnifying glass.

Here is how ROSD works, step-by-step:

1. The "Self-Reflector" (The Detective)
Instead of just showing the perfect answer, the system first acts like a detective. It looks at the student's wrong answer and the correct answer side-by-side.

It asks: "Where exactly did the logic break?"
It finds: The specific sentence or step where the student went off the rails.
It creates a "Corrective Idea": A short note explaining how to fix that specific mistake (e.g., "You forgot to convert units here," rather than "Here is the whole essay").

2. The "Error Quote" (The Highlighter)
The system doesn't just guess; it physically highlights the exact chunk of text in the student's answer where the error happened. It's like a teacher using a red pen to circle only the wrong sentence, leaving the rest of the page untouched.

3. Targeted Correction (The Surgery)
Now, the "Self-Teacher" steps in. But instead of making the student rewrite the whole essay, the teacher only gives guidance on the highlighted error.

The student keeps their correct introduction and valid reasoning steps (the "valid prefix").
They only re-write the part that was wrong, guided by the "Corrective Idea."

The Analogy: Fixing a Broken Car

Old Method (OPSD): Your car has a flat tire. The mechanic takes the whole car apart, throws away the engine, the seats, and the wheels, and builds a brand new car from scratch based on a blueprint. It works, but you lost everything that was already working.
ROSD: The mechanic looks at the car, finds the flat tire (the error), and says, "Okay, the engine and seats are fine. Let's just swap out this one tire and make sure it's inflated correctly."

The Results

The paper tested this on various "subjects" (like Physics, Chemistry, and using computer tools).

Better at the subject: The students learned the material better than before.
Better at new subjects: Because they learned the logic of fixing errors rather than memorizing the style of the answer, they were much better at solving problems they had never seen before.
Stability: The old method often got worse over time as the student got confused by trying to copy too much. ROSD stayed stable and kept getting better.

In short: ROSD teaches AI to fix its own mistakes surgically, preserving the good parts of its thinking, rather than forcing it to copy a perfect model from start to finish.

Technical Summary: Reflective On-Policy Self-Distillation (ROSD)

Problem Statement

While Reinforcement Learning with Verifiable Rewards (RLVR) has proven effective for optimizing Large Language Models (LLMs) in verifiable domains, standard algorithms like GRPO rely on sparse, outcome-level rewards. This lack of token-level supervision for intermediate reasoning steps limits the model's ability to refine its reasoning quality.

On-Policy Self-Distillation (OPSD) addresses this by providing dense token-level supervision from a "self-teacher" conditioned on a correct solution. However, existing OPSD methods suffer from two critical limitations that hinder both in-domain stability and out-of-domain (OOD) generalization:

Imitation over Correction: Conditioning the self-teacher on a verified solution encourages the student to imitate the reference trajectory's specific patterns and domain-specific expressions rather than learning to correct its own specific errors.
Overwriting Valid Reasoning: Applying distillation loss to the entire response forces the model to align with the teacher distribution across all tokens. This can overwrite valid reasoning prefixes and inject training-domain preferences into correct steps, leading to overfitting and poor generalization.

Methodology: ROSD

The authors propose Reflective On-Policy Self-Distillation (ROSD), a framework designed to shift self-distillation from reference-solution imitation to targeted reasoning correction. ROSD operates through two primary mechanisms:

1. Error-Focused Self-Reflection

Instead of directly using a correct solution as the teacher's context, ROSD employs a self-reflector (using the same model weights as the student/teacher) to analyze the discrepancy between a student's rollout and a correct solution.

For Wrong Rollouts: The reflector compares a wrong rollout ( $y^-$ $y^{-}$ ) with the shortest correct rollout ( $y^*$ $y^{*}$ ) from the same group. It generates two outputs:
- Corrective Idea ( $e$ ): An explanation of why the rollout failed and how to repair the logic.
- Error Quote ( $q$ ): An exact substring from the wrong rollout marking the first span where reasoning goes wrong.
For Correct Rollouts: The reflector summarizes the validity of the reasoning and the key idea behind the correct answer, avoiding direct full-solution imitation.

2. Quote-Localized Self-Distillation

The self-teacher is conditioned on the corrective idea ( $e$ ) rather than the full reference solution. This ensures the teacher provides error-specific guidance.

Localization: The error quote ( $q$ ) is aligned back to the original student rollout to identify the starting token index ( $k$ ) of the error.
Selective Loss: A mask is applied to the distillation loss function. The model is only trained to match the teacher distribution from token $k$ onwards (the erroneous suffix). The valid reasoning prefix (tokens $< k$ ) is preserved and excluded from the loss calculation.
Objective: The training objective minimizes the token-level Jensen-Shannon Divergence (JSD) between the student and the masked teacher distribution, ensuring updates are restricted to the segments requiring correction.

Key Contributions

ROSD Framework: A novel reflective on-policy self-distillation method that transitions from imitating reference solutions to performing selective reasoning corrections.
Dual-Technique Design:
- An error-focused self-reflector that extracts a corrective idea and precisely locates the first erroneous span.
- Error-localized distillation that restricts supervision updates to the suffix starting from the identified error, preserving valid reasoning prefixes.
Empirical Validation: Extensive experiments demonstrating that ROSD improves in-domain reasoning performance while significantly mitigating the OOD generalization collapse observed in standard OPSD methods.

Experimental Results

The authors evaluated ROSD on Qwen3-4B and Qwen3-8B backbones across five datasets: four science QA benchmarks (Chemistry, Physics, Biology, Materials Science) and a Tool-Use dataset.

In-Domain Performance: ROSD achieved the highest average scores across both model scales (72.83% on 4B and 73.45% on 8B), outperforming both GRPO and the standard SDPO baseline. Improvements were consistent across most domains, indicating no significant trade-off in specific areas.
Out-of-Domain Generalization: ROSD demonstrated substantially better OOD robustness compared to SDPO.
- While SDPO suffered severe degradation when training and testing domains differed significantly (e.g., training on ToolUse and testing on Science QA, where SDPO dropped to ~~2.88% average), ROSD maintained strong performance (~~41.31%).
- This suggests ROSD effectively reduces the noise and overfitting associated with full-response distillation.
Training Dynamics: ROSD converged faster than baselines and maintained stable performance throughout training. In contrast, SDPO often showed early gains followed by performance decline, indicating overfitting to training-domain solution styles.
Ablation Studies: Removing either the reflection component or the localized distillation mechanism resulted in performance drops, confirming that both error-focused guidance and selective supervision are necessary and complementary.

Significance and Claims

The paper claims that ROSD highlights the importance of error-aware, selectively applied supervision for effective reasoning post-training. By correcting flawed reasoning steps while preserving valid prefixes, ROSD avoids the pitfalls of imitating domain-specific reference trajectories.

The authors position ROSD as a step toward more robust reasoning models that can generalize across domains. However, they modestly note that while ROSD improves OOD performance over standard OPSD, it still does not fully match the OOD robustness of GRPO in all settings, indicating that further advances are needed to fully strengthen cross-domain generalization. The work emphasizes that strong reasoning capabilities require rigorous safety alignment to prevent potential misuse.