Surgical Post-Training: Proximal On-Policy Distillation for Reasoning with Knowledge Retention
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
| Model | In-domain Reasoning | OOD Reasoning | General | Avg | ||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|
| AIME24 | AIME25 | AMC23 | Math500 | Minerva | Olympia | Avg | GPQA-D | Connect4 | Avg | IFEval | ||
| 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 5.8 | 39.0 | 12.0 | 25.5 4.4 | 79.6 3.4 | 41.8 5.3 |
| Qwen3-8b+RFT | 27.3 | 20.0 | 59.0 | 83.0 | 40.4 | 53.8 | 47.3 0.5 | 49.1 | 3.1 | 26.1 3.8 | 81.5 1.5 | 46.4 0.7 |
| Qwen3-8b+ SFT+ | 29.3 | 24.7 | 67.3 | 84.8 | 42.5 | 54.5 | 50.5 3.7 | 47.3 | 14.1 | 30.7 0.8 | 80.0 3.0 | 49.4 2.3 |
| Qwen3-8b + SPoT | 28.0 | 27.3 | 71.5 | 87.4 | 39.7 | 58.5 | 52.1 5.3 | 46.8 | 36.0 | 41.4 11.5 | 84.8 1.8 | 53.3 6.2 |
| Model | In-domain Reasoning | OOD Reasoning | General | Avg | ||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|
| AIME24 | AIME25 | AMC23 | Math500 | Minerva | Olympia | Avg | GPQA-D | Connect4 | Avg | IFEval | ||
| SPoT-BCO | 28.0 | 27.3 | 71.5 | 87.4 | 39.7 | 58.5 | 52.1 5.3 | 46.8 | 36.0 | 41.4 11.5 | 84.8 1.8 | 53.3 6.2 |
| SPoT-BCE | 33.3 | 25.3 | 67.0 | 87.6 | 40.1 | 55.4 | 51.5 4.7 | 49.5 | 31.1 | 40.3 10.4 | 85.8 2.8 | 52.8 5.7 |
| DPO | 22.0 | 19.3 | 62.5 | 85.2 | 40.8 | 51.6 | 46.9 0.1 | 48.9 | 31.8 | 40.4 10.5 | 84.7 1.7 | 49.6 2.5 |
| Reward-SFT | 20.0 | 16.7 | 60.5 | 83.0 | 39.0 | 52.3 | 45.3 1.5 | 49.2 | 35.2 | 42.2 12.3 | 83.7 0.7 | 48.9 1.8 |
| DFT | 24.0 | 23.3 | 70.5 | 82.4 | 40.4 | 52.4 | 48.8 2.0 | 45.6 | 3.7 | 24.7 5.2 | 79.5 3.5 | 46.9 0.2 |
| Model | In-domain Reasoning | OOD Reasoning | General | Avg | ||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|
| AIME24 | AIME25 | AMC23 | Math500 | Minerva | Olympia | Avg | GPQA-D | Connect4 | Avg | IFEval | ||
| 2k direct data | 20.0 | 14.7 | 63.0 | 81.6 | 37.5 | 47.6 | 44.1 2.7 | 42.6 | 17.2 | 29.9 0.0 | 82.6 0.4 | 45.2 1.9 |
| 2k rectified data, | 26.0 | 19.3 | 67.0 | 87.8 | 41.2 | 52.4 | 49.0 2.2 | 49.8 | 26.0 | 37.9 8.0 | 84.1 1.1 | 50.4 3.3 |
| 4k rectified data, | 27.3 | 22.7 | 64.0 | 87.6 | 41.5 | 58.2 | 50.2 3.4 | 49.2 | 33.3 | 41.3 11.4 | 83.6 0.6 | 51.9 4.8 |
| 4k rectified data, | 28.0 | 27.3 | 71.5 | 87.4 | 39.7 | 58.5 | 52.1 5.3 | 46.8 | 36.0 | 41.4 11.5 | 84.8 1.8 | 53.3 6.2 |
| Method | AIME 2024 | AIME 2025 | AMC 2023 | Connect4 | IFEval | Avg |
|---|---|---|---|---|---|---|
| SPoT | 28.0 6.0 | 27.3 8.0 | 71.5 5.0 | 36.0 25.1 | 84.8 1.8 | 49.5 9.2 |
| GRPO | 36.3 14.3 | 28.7 9.4 | 77.5 11.0 | 7.7 3.2 | 81.2 1.8 | 46.3 6.0 |
| SPoT + GRPO | 41.3 19.3 | 36.7 17.4 | 86.0 19.5 | 29.4 18.5 | 82.1 0.9 | 55.1 14.8 |
Appendix figures & tables13 assets
Supplementary material from the paper’s appendix.
Appendix
| 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 |
| 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 |
| Setting | AIME24 | AIME25 | AMC23 | Connect4 | IFEval |
|---|---|---|---|---|---|
| Base | 22.0 | 19.3 | 66.5 | 10.9 | 83.0 |
| SPoT , (1,998 pairs) | 19.3 | 22.0 | 64.0 | 23.7 | 84.5 |
| SPoT , (3,646 pairs) | 26.0 | 22.0 | 67.0 | 27.9 | 82.3 |
| Component | GRPO | SPoT |
|---|---|---|
| Total time | 13 h | 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 | 104 | 8 |
| Model | In-domain Reasoning | OOD Reasoning | General | ||||||
|---|---|---|---|---|---|---|---|---|---|
| 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 |
| Model | TruthfulQA | MMLU-Pro (500 subset) |
|---|---|---|
| Qwen3-8B Base | 68.39 | 51.2 |
| Qwen3-8B + SPoT | 67.90 | 55.2 |
| 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 |
| 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 |
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