Surgical Post-Training: Proximal On-Policy Distillation for Reasoning with Knowledge Retention

Wenye Lin; Kai Han

cs.CLMar 2, 2026

Surgical Post-Training: Proximal On-Policy Distillation for Reasoning with Knowledge Retention

Authors: Wenye Lin, Kai Han

Organizations: The University of Hong Kong

Abstract

Injecting new reasoning knowledge into Large Language Models (LLMs) via post-training often induces catastrophic forgetting. Recent studies emphasize the importance of on-policy data but suggest that KL-divergence fails to mitigate forgetting. In contrast, we show, both analytically and empirically, that the KL-constrained reward formulation actually plays a critical role in retaining knowledge during post-training. This motivates our Surgical Post-Training (SPOT), a proximal on-policy distillation framework designed to optimize reasoning efficiently while preserving prior knowledge. SPOT consists of (1) a data rectification pipeline employing an Oracle to surgically correct erroneous steps via minimal edits, generating proximal on-policy data; and (2) a reward-based binary cross-entropy objective essential for enhancing reasoning and mitigating forgetting. Empirically, with only 4k rectified math pairs, SPOT improves Qwen3-8B's accuracy by 6.2% on average across in-domain and out-of-domain tasks, requiring merely 16-minute model training on 8x H800 GPUs. Moreover, SPOT provides a superior initialization for subsequent reinforcement learning, significantly elevating the performance ceiling. Code: https://github.com/Visual-AI/SPoT

Figures & tables

Figure 1: Illustration of SPoT . Left: An Oracle surgically corrects reasoning errors, yielding positive samples proximal to the model’s original distribution. Top-Right: Unlike DPO’s relative ranking, we leverage an explicit classification loss that proves more effective for reasoning tasks. Bottom-Right: The tethering effect inherent in our reward definition substantially mitigates catastrophic forgetting.

Figure 2: IFEval acc results (avg@5). SFT+ forgets instruction following ability, while reward-based methods do not.

Figure 3: Training loss curve. Reward-SFT and DPO converge rapidly; SFT+ remains high due to absolute likelihood maximization.

Figure 4: Evolution of implicit rewards during training. Left: reward scores for chosen responses; Right: reward scores for rejected responses. The reward measures how much the model’s preference for the chosen response has grown relative to the initial state.

Model	In-domain Reasoning							OOD Reasoning			General	Avg
Model	AIME24	AIME25	AMC23	Math500	Minerva	Olympia	Avg	GPQA-D	Connect4	Avg	IFEval	Avg
Qwen3-8b	22.0	19.3	66.5	82.8	37.9	52.3	46.8	48.9	10.9	29.9	83.0	47.1
Qwen3-8b+SFT	15.3	13.3	56.0	78.2	38.6	44.3	41.0 $\downarrow$ 5.8	39.0	12.0	25.5 $\downarrow$ 4.4	79.6 $\downarrow$ 3.4	41.8 $\downarrow$ 5.3
Qwen3-8b+RFT	27.3	20.0	59.0	83.0	40.4	53.8	47.3 $\uparrow$ 0.5	49.1	3.1	26.1 $\downarrow$ 3.8	81.5 $\downarrow$ 1.5	46.4 $\downarrow$ 0.7
Qwen3-8b+ SFT+	29.3	24.7	67.3	84.8	42.5	54.5	50.5 $\uparrow$ 3.7	47.3	14.1	30.7 $\uparrow$ 0.8	80.0 $\downarrow$ 3.0	49.4 $\uparrow$ 2.3
Qwen3-8b + SPoT	28.0	27.3	71.5	87.4	39.7	58.5	52.1 $\uparrow$ 5.3	46.8	36.0	41.4 $\uparrow$ 11.5	84.8 $\uparrow$ 1.8	53.3 $\uparrow$ 6.2

Table 1: Main results for in-domain reasoning, OOD reasoning, and general instruction following (OOD non-reasoning). We compare SPoT (with BCO loss) against standard baselines: SFT (rejection-sampled Gemini 2.5 Pro responses), RFT (rejection-sampled self-generated responses), and SFT+ (rectified

y^{+}

). Direct SFT leads to performance degradation on both in-domain and OOD tasks. RFT fails to yield clear in-domain improvements and causes OOD regression. While SFT+ improves in-domain reasoning with near-distribution data, it still suffers OOD losses. In contrast, SPoT demonstrates superiority in both in-domain and OOD performance. We report avg@16 for AIME24, AIME25, and AMC23, and avg@5 for the others.

Model	In-domain Reasoning							OOD Reasoning			General	Avg
Model	AIME24	AIME25	AMC23	Math500	Minerva	Olympia	Avg	GPQA-D	Connect4	Avg	IFEval	Avg
SPoT-BCO	28.0	27.3	71.5	87.4	39.7	58.5	52.1 $\uparrow$ 5.3	46.8	36.0	41.4 $\uparrow$ 11.5	84.8 $\uparrow$ 1.8	53.3 $\uparrow$ 6.2
SPoT-BCE	33.3	25.3	67.0	87.6	40.1	55.4	51.5 $\uparrow$ 4.7	49.5	31.1	40.3 $\uparrow$ 10.4	85.8 $\uparrow$ 2.8	52.8 $\uparrow$ 5.7
DPO	22.0	19.3	62.5	85.2	40.8	51.6	46.9 $\uparrow$ 0.1	48.9	31.8	40.4 $\uparrow$ 10.5	84.7 $\uparrow$ 1.7	49.6 $\uparrow$ 2.5
Reward-SFT	20.0	16.7	60.5	83.0	39.0	52.3	45.3 $\downarrow$ 1.5	49.2	35.2	42.2 $\uparrow$ 12.3	83.7 $\uparrow$ 0.7	48.9 $\uparrow$ 1.8
DFT	24.0	23.3	70.5	82.4	40.4	52.4	48.8 $\uparrow$ 2.0	45.6	3.7	24.7 $\downarrow$ 5.2	79.5 $\downarrow$ 3.5	46.9 $\downarrow$ 0.2

Table 2: Comparison on loss objectives. We find that SPoT-BCO yields the best overall results. SPoT-BCE is more competitive on keeping general abilities. DPO mitigates catastrophic forgetting but shows no clear improvements for in-domain reasoning. Although DFT Wu et al. (2025) is designed for generalization, we observe a clear degradation on OOD tasks.

Model	In-domain Reasoning							OOD Reasoning			General	Avg
Model	AIME24	AIME25	AMC23	Math500	Minerva	Olympia	Avg	GPQA-D	Connect4	Avg	IFEval	Avg
2k direct data	20.0	14.7	63.0	81.6	37.5	47.6	44.1 $\downarrow$ 2.7	42.6	17.2	29.9 $\downarrow$ 0.0	82.6 $\downarrow$ 0.4	45.2 $\downarrow$ 1.9
2k rectified data, $\gamma=0.6$	26.0	19.3	67.0	87.8	41.2	52.4	49.0 $\uparrow$ 2.2	49.8	26.0	37.9 $\uparrow$ 8.0	84.1 $\uparrow$ 1.1	50.4 $\uparrow$ 3.3
4k rectified data, $\gamma=1$	27.3	22.7	64.0	87.6	41.5	58.2	50.2 $\uparrow$ 3.4	49.2	33.3	41.3 $\uparrow$ 11.4	83.6 $\uparrow$ 0.6	51.9 $\uparrow$ 4.8
4k rectified data, $\gamma=0.6$	28.0	27.3	71.5	87.4	39.7	58.5	52.1 $\uparrow$ 5.3	46.8	36.0	41.4 $\uparrow$ 11.5	84.8 $\uparrow$ 1.8	53.3 $\uparrow$ 6.2

Table 3: Ablation study on different data. “Rectified data” refers to the dataset curated using our surgical rectification pipeline, while “direct data” uses the rejected-sampling direct answers from Gemini 2.5 Pro.

\gamma

defines the threshold of change ratio in Equation 1 .

Method	AIME 2024	AIME 2025	AMC 2023	Connect4	IFEval	Avg
SPoT	28.0 $\uparrow$ 6.0	27.3 $\uparrow$ 8.0	71.5 $\uparrow$ 5.0	36.0 $\uparrow$ 25.1	84.8 $\uparrow$ 1.8	49.5 $\uparrow$ 9.2
GRPO	36.3 $\uparrow$ 14.3	28.7 $\uparrow$ 9.4	77.5 $\uparrow$ 11.0	7.7 $\downarrow$ 3.2	81.2 $\downarrow$ 1.8	46.3 $\uparrow$ 6.0
SPoT + GRPO	41.3 $\uparrow$ 19.3	36.7 $\uparrow$ 17.4	86.0 $\uparrow$ 19.5	29.4 $\uparrow$ 18.5	82.1 $\downarrow$ 0.9	55.1 $\uparrow$ 14.8

Table 4: Comparison with GRPO. SPoT + GRPO: GRPO initialized from SPoT checkpoint.

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: Random Rectification Example 1. Left: answer from Qwen3-8B; Right: rectification by Gemini 2.5 Pro.

Figure 6: Random Rectification Example 2. Left: answer from Qwen3-8B; Right: rectification by Gemini 2.5 Pro.

Figure 7: Random Rectification Example 3. Left: answer from Qwen3-8B; Right: rectification by Gemini 2.5 Pro.

Figure 8: Random Rectification Example 4. Left: answer from Qwen3-8B; Right: rectification by Gemini 2.5 Pro.

Model	AIME24	AIME25	AMC23	Connect4	IFEval
Base	22.0	19.3	66.5	10.9	83.0
SPoT (Qwen3-30B-A3B Oracle)	25.6	22.5	67.0	26.9	82.9

Appendix

Table 5: SPoT with a weaker open-source Oracle ( Qwen3-30B-A3B-2507 ) on Qwen3-8B. Even without a frontier API, the framework yields consistent gains.

Prompt	AIME24	AIME25	AMC23	Connect4	IFEval
Default	26.0	19.3	67.0	26.0	84.1
Simplified	26.5	18.7	66.5	27.0	84.1

Appendix

Table 6: Prompt sensitivity of the Oracle on Qwen3-8B (2k pairs). Performance is robust to prompt variation.

Setting	AIME24	AIME25	AMC23	Connect4	IFEval
Base	22.0	19.3	66.5	10.9	83.0
SPoT , $\tau{=}0.5$ (1,998 pairs)	19.3	22.0	64.0	23.7	84.5
SPoT , $\tau{=}0.7$ (3,646 pairs)	26.0	22.0	67.0	27.9	82.3

Appendix

Table 7: Effect of the LCS filtering threshold on Qwen3-8B. The threshold

0.6

used in the main paper balances the trade-off.

Figure 9: Distributions of change ratio. A higher change ratio indicates that reasoning failures occur earlier in the reasoning chain. The distribution shape is determined by the model ability and the data difficulty together.

Component	GRPO	SPoT
Total time	$\sim$ 13 h	$\sim$ 1 h
Sampling	5.56 h	43 min
Policy training	5.12 h	16 min
Reference log-prob	2.25 h	(in training)
Teacher rectification	—	overlapped with sampling
Total GPU hours	$\sim$ 104	$\sim$ 8

Appendix

Table 8: Wall-clock and GPU-hour for GRPO and SPoT on Qwen3-8B. Teacher rectification is parallel with student sampling and therefore does not appear as a separate term.

Model	In-domain Reasoning						OOD Reasoning		General
Model	AIME24	AIME25	AMC23	Math500	Minerva	Olympia	GPQA-D	Connect4	IFEval
Qwen3-8b	22.0 ±1.7	19.3 ±1.2	66.5 ±1.7	82.8 ±0.4	37.9 ±0.6	52.3 ±0.5	48.9 ±0.7	10.9 ±1.4	83.0 ±0.3
Qwen3-8b+SFT	15.3 ±1.7	13.3 ±1.4	56.0 ±1.7	78.2 ±0.5	38.6 ±0.7	44.3 ±0.6	39.0 ±1.6	12.0 ±1.5	79.6 ±0.4
Qwen3-8b+RFT	27.3 ±1.4	20.0 ±1.6	59.0 ±2.0	83.0 ±0.4	40.4 ±0.6	53.8 ±0.5	49.1 ±0.9	3.1 ±1.1	81.5 ±0.5
Qwen3-8b+ SFT+	29.3 ±0.8	24.7 ±0.8	67.3 ±3.0	84.8 ±0.4	42.5 ±0.5	54.5 ±0.5	47.3 ±1.2	14.1 ±1.3	80.0 ±0.6
Qwen3-8b + SPoT	28.0 ±1.1	27.3 ±1.1	71.5 ±1.3	87.4 ±0.5	39.7 ±0.5	58.5 ±0.7	46.8 ±0.5	36.0 ±1.8	84.8 ±0.4

Appendix

Table 9: Detailed results with standard deviations. This table corresponds to the main results presented, additionally reporting the standard deviations to demonstrate statistical significance. The standard deviations are calculated across independent sampling runs (

k=16

for AIME24, AIME25, and AMC23, and

k=5

for the others). The average summary columns are omitted here for brevity.

Model	TruthfulQA	MMLU-Pro (500 subset)
Qwen3-8B Base	68.39	51.2
Qwen3-8B + SPoT	67.90	55.2

Appendix

Table 10: General-capability evaluations on Qwen3-8B. SPoT preserves truthfulness and improves on MMLU-Pro.

Model	AIME24	AIME25	AMC23	Connect4	IFEval
Base	22.0	19.3	66.5	10.9	83.0
DPO+SFT ( Wang et al., 2024 )	24.6	18.9	66.5	21.9	81.5
SPoT (Ours)	26.0	19.3	67.0	26.0	84.1

Appendix

Table 11: Comparison with DPO+SFT ( Wang et al., 2024 ) on Qwen3-8B. SPoT dominates across in-domain math, OOD reasoning, and instruction following.

Model	AIME24	AIME25	AMC23	IFEval	Connect4
Base	12.67	9.33	41.00	67.64	0.00
SPoT	14.67	13.33	46.00	68.30	0.00

Appendix

Table 12: SPoT on Qwen3-1.7B. The recipe transfers to smaller models.

LLM Post-Training

Large language model post-training adapts pretrained models into more useful, reliable, and specialized assistants through supervised and preference-based learning.

See all 246 papers

Distilled Reinforcement Learning for LLM Post-training
Jul 19, 2026Chen Wang, Zhaochun Li, Jionghao Bai +4RL for Language Model Reasoning Cross-Architecture Knowledge Distillation
College of Elite Engineers, Nankai University · Zhongguancun Academy · Beijing Institute of Technology +4
Reinforcement-aware Knowledge Distillation for LLM Reasoning
Feb 26, 2026Zhaoyang Zhang, Shuli Jiang, Yantao Shen +6Knowledge Distillation RL for Knowledge Distillation
AWS Agentic AI
From Rollouts to Recipes: Self-Contained Post-Training for LLMs
Sep 1, 2026Yifei Li, Lingling Zhang, Muye Huang +3Self-Training for Language Models RL for Language Model Reasoning
School of Computer Science and Technology, Xi’an Jiaotong University, China · Shaanxi Province Key Laboratory of Big Data Knowledge Engineering · Zhongguancun Academy, Beijing, China +1

Jul 19, 2026J/K move · Enter open · S save

Distilled Reinforcement Learning for LLM Post-training

Chen Wang, Zhaochun Li, Jionghao Bai, Yining Zhang+3

College of Elite Engineers, Nankai University · Zhongguancun Academy · Beijing Institute of Technology+4

RL for Language Model Reasoning Cross-Architecture Knowledge Distillation On-Policy Distillation RL Post-Training Knowledge Distillation RL for Knowledge Distillation

Open paper AskPDF arXiv

Large language model (LLM) post-training is essential for improving reasoning, adaptation, and alignment. Existing methods mainly follow two paradigms: reinforcement learning (RL) and on-policy distillation (OPD). However, RL relies on coarse-grained outcome supervision, resulting in difficult credit assignment and limited capability to acquire new knowledge. OPD, meanwhile, unconditionally matches teacher logits through KL divergence, which creates a dilemma: similar teachers provide little new knowledge, while substantially different teachers often yield ineffective guidance, largely restricting OPD to within-family distillation. We propose Distilled Reinforcement Learning (Distilled RL), which integrates teacher supervision into the RL objective to provide fine-grained guidance, selectively transfer new knowledge and avoid unconditional imitation. Distilled RL contains three components: reverse importance sampling with clipping, negative sample reset, and sequence-level geometric normalization. Through a concise and interpretable case study, we demonstrate that Distilled RL can effectively transfer previously unavailable knowledge from a teacher model to a student model. Extensive experiments across both within-family and cross-family distillation settings show that Distilled RL substantially outperforms standard RL and OPD in terms of both pass@1 and pass@k. Our code is available at https://github.com/597358816/Distilled-RL.

Abstract

LLM Post-Training

Explore similar work

Distilled Reinforcement Learning for LLM Post-training

Reinforcement-aware Knowledge Distillation for LLM Reasoning

From Rollouts to Recipes: Self-Contained Post-Training for LLMs