PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Wenye Lin, Kai Han

arXiv:2603.01683v2cs.CLcs.AI

TL;DR

Post-training can inject reasoning knowledge while causing catastrophic forgetting, and existing on-policy distillation may require inaccessible teacher distributions. SPOT uses minimal-edit Oracle rectification and a reward-based binary objective; with 4k rectified pairs, it delivers gains across in-domain and OOD benchmarks and improves later GRPO initialization.

  • Problem

    Existing on-policy distillation mitigates distribution shifts but requires identical tokenizers and white-box teachers, while data proximity alone does not prevent forgetting.

  • Method

    SPOT combines minimal-edit Oracle correction of erroneous student outputs with a reward-based binary objective that preserves reference-model tethering.

  • Results

    SPOT outperforms baselines on average across in-domain and OOD metrics, with 4k rectified pairs and gains reported across Qwen3-8B and Llama-3.1-8B-Instruct.

  • Takeaways & Limitations

    SPOT provides a lightweight bridge between SFT and RL and substantially raises the performance ceiling of subsequent GRPO.

  • Takeaways & Limitations

    Experiments focus on models up to 8B and text-based mathematical reasoning, and performance depends on the Oracle’s capability.

Abstract

from arXiv · show

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

1 Introduction

The paper addresses how to inject external reasoning knowledge while mitigating catastrophic forgetting. SPOT combines proximal surgical rectification with a reward-based binary objective to preserve prior knowledge and improve reasoning.

  • Motivation: On-policy data can reduce distribution shifts, but existing on-policy distillation requires identical tokenizers and white-box teachers.This limits cross-family distillation and use of proprietary models.
  • Motivation: Data proximity alone is insufficient: SFT on proximal data still forgets, whereas reward-based methods retain prior knowledge through KL-constrained tethering.Reward-SFT isolates the reward formulation and confirms its role in retention.
  • Motivation: Positive-only training can raise likelihood for erroneous responses, while DPO’s relative ranking is inadequate for verifiable reasoning.These are identified as the pull-up effect and the insufficiency of relative ranking.
  • SPOT: SPOT combines minimal-edit data rectification with a reward-based binary cross-entropy objective that maximizes rectifications while suppressing errors.The objective preserves the reference model’s tethering effect.
  • Contributions: SPOT establishes a stronger initialization for later GRPO, raising performance ceilings by 7.2% on in-domain math and 21.7% on Connect4.It also provides better general instruction-following capabilities.

2 Data Rectification Pipeline

The data rectification pipeline constructs proximal contrastive examples by minimally correcting erroneous policy outputs. It preserves shared generation structure while concentrating differences at decision-critical reasoning steps.

  • Motivation: Offline SFT uses divergent data, while fully on-policy sampling is inefficient for hard reasoning problems.At a 1% pass rate, one correct response requires roughly 100 rollouts in expectation.
  • Pipeline: The pipeline constructs contrastive triples (x, y−, y+) where y+ minimally corrects y− while preserving its style and lexical structure.The positive response is intended to remain proximal to the model’s generation distribution.
  • Error Elicitation: Error elicitation samples responses from the current policy and retains incorrect final answers for rectification.The retained pair is (x, y−).
  • Oracle-Guided Surgical Rectification: An Oracle modifies only incorrect reasoning steps and filters out responses with incorrect final answers, producing y+ as a nearest valid neighbor to y−.The Oracle may be a human expert or teacher model and is queried once per failure in the offline setting.
  • LCS Filtering: LCS filtering constrains structural changes by rejecting pairs whose change ratio R_LCS exceeds γ, with γ = 0.6.The threshold is selected to retain training samples while preserving downstream performance.
  • Resulting Dataset: Each rectified pair shares most token trajectories and diverges mainly at decision-critical parts, focusing gradients on divergent reasoning tokens.This shared structure supports the pipeline’s proximal-data objective.

3 The Reward Is Secretly a Regularizer

The section shows that the optimization objective, not data proximity alone, provides intrinsic regularization against forgetting. Reward-based training limits updates as the policy becomes sufficiently separated from the reference model.

  • Empirical Observation: Proximal data alone does not prevent forgetting, so the learning objective must provide an intrinsic regularizer.The analysis compares SFT with reward-based methods on the same proximal dataset.
  • Reward-SFT: Reward-SFT classifies the chosen response y+ as correct under an implicit reward, unlike SFT’s direct token-level likelihood maximization.It isolates the implicit reward’s effect without using negative samples.
  • Empirical Observation: SFT+ loses instruction-following ability monotonically, while DPO and Reward-SFT remain stable on IFEval without negative data for Reward-SFT.The result attributes resistance to forgetting to implicit regularization in the reward formulation.
  • Gradient Analysis: Standard SFT applies constant optimization pressure even to high-probability samples, causing continued updates and distribution shift.For p(y+|x) ≈ 0.99, SFT still pushes probability toward 1.0.
  • Gradient Analysis: Reward-SFT uses λ(x, y+) = 1 − σ(rθ(x, y+)) as an instance-dependent scaling coefficient that adapts optimization pressure to reward.When the policy is close to the reference, λ ≈ 0.5 and adaptation remains rapid.
  • Gradient Analysis: At rθ = 10, λ ≈ 4.5 × 10^-5, reducing the gradient by over 2000-fold relative to SFT’s fixed coefficient of 1.0.This suppresses over-optimization and limits drift from pretrained knowledge.
  • Gradient Analysis: The Elastic Tether acts as sample-wise early stopping by suppressing updates on well-learned samples and preserving knowledge in the reference policy.Optimization halts once the policy has sufficiently diverged from the reference distribution.

4 The Surgical Optimization Objective

Positive-only training can preserve knowledge yet still raise the likelihood of erroneous responses, while relative ranking may optimize suppression rather than correct reasoning. SPOT addresses these issues with decoupled reward-based BCE objectives and adaptive reward shifting, trading off saturation, optimization, and retention.

  • 4.1 The “Pull-Up” Effect: SFT+ and Reward-SFT inadvertently increase the likelihood of rejected responses because rectified positive and negative samples share most token trajectories.The shared prefix receives positive-sample gradient updates, preventing a sharp decision boundary at the critical divergence point.
  • 4.2 The Insufficiency of Relative Ranking: DPO can minimize its relative-ranking loss by decreasing rejected rewards while leaving chosen rewards stagnant or reduced, limiting reinforcement of correct reasoning.The paper reports that later DPO training shows sharply decreasing rejected rewards without sustained chosen-reward growth.
  • 4.3 Reward-based Binary Cross Entropy Optimization: SPoT-BCE decouples supervision into independent terms that maximize confidence in y+ and minimize confidence in y−.The implicit reward serves as the classification logit, providing dense supervision while retaining pre-trained knowledge.
  • 4.3 Reward-based Binary Cross Entropy Optimization: SPoT-BCO adds an adaptive reward shift δ, computed from batch statistics with an exponential moving average, to the BCE objective.The shift functions as a proxy for the intractable partition term β log Z(x).
  • 4.3 Reward-based Binary Cross Entropy Optimization: Table 1 reports SPOT superiority across in-domain reasoning, OOD reasoning, and general instruction following, whereas baseline methods show degradation, regression, or unclear gains.The table uses avg@16 for AIME24, AIME25, and AMC23, and avg@5 for the other evaluations.
  • 4.4 Theoretical Analysis of δ: SPoT-BCO continues optimization beyond SPoT-BCE’s saturation but permits larger KL divergence from the reference model, so SPoT-BCE retains prior knowledge more faithfully.SPoT-BCE’s early saturation limits deviation from πref, whereas adaptive shifting enables further updates and higher rewards.

5 Experiments

Experiments show that SPOT combines surgical data rectification with a regularized reward-based objective to improve reasoning while retaining prior capabilities. Across comparisons and ablations, proximity alone is insufficient, whereas SPOT also provides a stronger starting point for subsequent reinforcement learning.

  • 5.1 Results: SFT damages general capabilities and out-of-distribution reasoning, including an 11.5-point IFEval drop on Llama-3.1-8B-Instruct.Qwen3-8B in-domain reasoning also drops from 46.8% to 41.0% under off-policy SFT.
  • 5.1 Results: Proximal or on-policy data alone does not prevent forgetting when optimized with SFT.SFT+ improves in-domain reasoning but still forgets on general tasks, while RFT produces negligible in-domain gains and notable forgetting.
  • 5.1 Results: SPOT outperforms all baselines on average across in-domain and out-of-distribution metrics for both Qwen3-8B and Llama-3.1-8B-Instruct.SPOT also boosts general instruction following on Qwen3-8B.
  • 5.2 Comparison on Loss Objectives: SPoT-BCO achieves the best overall average and in-domain reasoning accuracy, while SPoT-BCE better preserves general capabilities.SPoT-BCE reaches 85.8% on IFEval and 49.5% on GPQA-D; DPO mitigates forgetting but does not improve in-domain reasoning.
  • 5.3 The Effect of the Surgical Rectification Pipeline: +5.2%: rectified data outperforms direct data in the data-source ablation.With SPoT-BCO, γ = 0.6 reaches a 53.2% average score, while the same training set size and γ = 1 is lower.
  • 5.4 Superior Initialization for Reinforcement Learning: +7.2% in-domain and +21.7% OOD reasoning: SPOT + GRPO outperforms GRPO alone across all benchmarks.GRPO from the SPOT checkpoint improves reasoning further but causes some forgetting, including a 1.8% absolute IFEval decrease.

6 Related Work

Related work emphasizes on-policy data for reducing catastrophic forgetting and identifies KL-constrained rewards as important for knowledge preservation, while noting that DPO remains sub-optimal for reasoning.

  • 6 Related Work: On-policy data is linked to better knowledge retention, while KL-constrained reward formulations support preservation of prior knowledge.The paper positions these findings alongside work seeking to improve DPO for reasoning tasks.

7 Conclusion and Discussion

SPOT combines minimal-edit rectification with a binary optimization objective to inject reasoning knowledge while preserving prior knowledge. The paper attributes retention to KL-constrained rewards and reports gains across in-domain and out-of-distribution benchmarks, while identifying scope limits for current experiments.

  • Conclusion and Discussion: SPOT combines surgical rectification with binary optimization to inject reasoning knowledge without catastrophic forgetting.The pipeline corrects only erroneous steps, preserving the model’s prior token distribution.
  • Conclusion and Discussion: KL-constrained reward formulation acts as a sample-wise early stopping mechanism that retains knowledge during post-training.The formulation maintains the reference model’s tethering effect.
  • Conclusion and Discussion: Binary optimization is critical for rigid reasoning tasks where verifiable truth makes pairwise preferences ill-suited.The objective maximizes rectifications while suppressing errors, unlike relative ranking in DPO.
  • Limitations and Future Work: Experiments focus on models up to 8B and text-based mathematical reasoning, while Oracle capability remains a performance dependency.The paper leaves larger-scale studies across diverse Oracles and extensions to multimodal, code, and agentic domains for future work.
  • Conclusion and Discussion: SPOT remains effective with a weaker open-source Oracle, and its performance is robust to simplified prompt formulations.The weaker Oracle result is reported in Table 5, while nearly identical benchmark performance follows prompt simplification.

F Rationale for the LCS Filtering Threshold

The LCS filter balances curation efficiency against fidelity to the on-policy constraint. The paper selects 0.6 as the practical sweet spot and validates LCS as a proxy for student-side perplexity.

  • Filtering Trade-off: A strict LCS threshold retains more on-policy data but discards many candidate pairs, increasing valid-pair gathering costs.A loose threshold retains more pairs but risks admitting off-policy data.
  • Threshold Selection: 0.6 is the empirically selected sweet spot between curation efficiency and on-policy fidelity.The ablation starts from 6,171 Qwen3-8B error samples.
  • Proxy Validation: Mean PPL is 1.60 for 1k pairs with change ratio < 0.6 and 1.95 for 1k pairs with change ratio > 0.6.The lower-PPL group is more on-policy according to LCS.
  • Proxy Validation: The PPL relationship confirms LCS as an effective proxy for student-side perplexity during curation.Computing PPL directly would require a full forward pass per candidate.
  • Change Ratio: Change ratio measures the proportion of the reasoning chain modified during rectification, with higher values indicating earlier reasoning failures.Its distribution depends jointly on model ability and data difficulty.

H Cost and Efficiency Analysis

The cost analysis compares SPOT and GRPO under identical hardware while accounting for Oracle curation and sampling. SPOT is designed for low training overhead, using one rollout per sample and parallelizing teacher rectification.

  • Oracle Cost: Gemini 2.5 Pro curation of 4k rectified pairs costs approximately $278 and takes roughly 40 minutes when parallelized with sampling.The stated API price is $10 per 1M output tokens.
  • Training Efficiency: SPOT requires only one rollout per sample in the end-to-end comparison with GRPO on Qwen3-8B.The comparison uses identical 8× H800 80GB hardware.
  • Cost Accounting: Teacher rectification runs in parallel with student self-sampling and is excluded as a separate term in the wall-clock comparison.This accounting choice is stated in the Table 8 description.
  • Statistical Reporting: Generative-evaluation standard deviations use 16 independent runs for AIME24, AIME25, and AMC23, and 5 runs for the remaining datasets.The paper reports these deviations to support statistical significance.

J Experiment Setting

Experiments use instruction-tuned policy models, controlled data-generation and training settings, and evaluations spanning reasoning tasks and dynamically constructed Connect4 states. Connect4 supplies deterministic, solver-verified out-of-distribution reasoning tasks.

  • Models and Data: Qwen3-8B and Llama3.1-8B-Instruct serve as policy models because instruction-tuned models are more susceptible to catastrophic forgetting.The experiments use the English subset of DAPO-Math-17k.
  • Models and Data: SPOT uses 4k generated contrastive pairs for Qwen3-8B, while Llama3.1-8B-Instruct uses 1.5k pairs matched to its inherently correct responses.Methods are compared using the same amount of training data within the relevant comparison.
  • Hyperparameters: Training uses two epochs, context length 8,192, β = 0.1, temperature 0.7, top-p 0.8, and a maximum of 32,768 tokens.Learning rates are separately specified for SPOT and baseline methods.
  • RL Setting: RL experiments use GRPO with 4k DAPO-Math-17k prompts, n=8 rollouts per prompt, and 8× H800 80GB GPUs.The RL runs train for 10 epochs with batch size 256 and mini-batch size 64.
  • Connect4 Task: Connect4 provides a deterministic 6 × 7 environment requiring board parsing, lookahead planning, and winning-topology identification.Its state space contains approximately 4.5 × 10^12 legal positions.
  • Connect4 Dataset: Connect4 states are dynamically generated, deduplicated, and balanced to keep empty-answer cases at no more than 20%.A perfect GAMEBoT solver supplies ground truth, which is checked against parsed final answers.
  • Evaluation Protocol: Each Connect4 query asks for winning moves for both the model and its opponent.The evaluation extracts final answers from the model’s chain of thought.

L Additional Evaluation

Additional evaluations indicate that SPOT preserves general competencies and outperforms DPO+SFT across the tested benchmarks on Qwen3-8B.

  • SPOT preserves truthfulness and improves performance on MMLU-Pro in general-capability evaluations.
  • SPOT outperforms DPO+SFT on every benchmark under the same Qwen3-8B setting.The comparison covers in-domain math, out-of-distribution reasoning, and instruction following.

N Experiment on Qwen3-1.7B

Applying SPOT to Qwen3-1.7B yields consistent gains across in-domain math and IFEval, while Connect4 remains too difficult for both Base and SPOT.

  • SPOT yields consistent gains across in-domain math and IFEval on Qwen3-1.7B.This indicates that the rectification-based recipe transfers to smaller model capacities.
  • Connect4 is too difficult for both Base and SPOT at the Qwen3-1.7B scale.The passage attributes this outcome to insufficient reasoning capacity rather than a failure of the training recipe.
Loading 2603.01683v2…