SOD: Step-wise On-policy Distillation for Small Language Model Agents
SOD uses step-wise reweighted OPD to lift a 0.6B agent to 26.13% on AIME 2025.
Key Findings
Methodology
SOD decomposes multi-turn TIR trajectories into reasoning steps, estimates student-teacher drift with a step-level divergence score d_k, and uses it to adapt the distillation weight at each step. When alignment is good, it preserves dense OPD token supervision; when tool errors induce state drift, Eq. (7) attenuates the weight w_k to suppress misleading teacher signals. The final objective combines trajectory-level exploration from GRPO with step-wise OPD: L = L_GRPO + L_stepOPD.
Key Results
- On AIME 2024/2025, GPQA-Diamond, and LiveCodeBench, the 0.6B student with SOD scores 20.84/26.13/22.19/27.72, averaging 24.22. This beats the second-best OPD baseline at 20.04 by a relative 20.86%, and yields a notable 26.13% on AIME 2025, a strong result for a sub-billion agent.
- For the 1.7B student, SOD reaches 50.83/41.72/38.72/40.63, with a 42.98 average versus 36.27 for OPD. The paper also reports that the 1.7B student recovers 69.8% of the 4B teacher’s performance, compared with 58.9% for OPD, showing better distillation efficiency at a larger scale.
- Ablation confirms that the adaptive weighting is essential: uniform weighting drops to 34.70, heuristic decay to 37.14, masking after the first wrong tool call to 31.85, and removing weight clipping to 38.10. Removing GRPO gives 40.78, while removing step-wise OPD collapses performance to 25.39.
Significance
The paper addresses a central bottleneck in agentic small-model training: once a tool call goes wrong, the resulting corrupted observation changes the state distribution, making later teacher supervision increasingly unreliable. SOD turns this failure mode into a design principle by conditioning distillation strength on local reliability. This matters academically because it explains why vanilla OPD can fail under long-horizon tool use, and practically because it offers a stable route for transferring tool-using reasoning to resource-constrained, privacy-sensitive, on-device agents.
Technical Contribution
Technically, SOD introduces a zero-extra-label, zero-extra-inference-cost proxy for local mismatch: d_k is computed as the mean absolute log-probability gap between student and teacher on the current step. The weighting rule uses ratios of consecutive divergence scores and is clipped by 1+δ, allowing the model to downweight corrupted regions while preserving recovery when the student re-aligns. The authors also provide a theoretical analysis showing super-linear divergence accumulation under erroneous tool observations and gradient SNR degradation for vanilla OPD in low-overlap regions. SOD directly counteracts this failure cascade.
Novelty
Unlike standard OPD, GRPO, or self-distillation variants, SOD is the first to make step-level reliability the central control variable for distilling TIR agents. It is also distinct from hard masking or fixed exponential decay because it can represent non-monotonic trajectories: if the student recovers after an earlier mistake, the distillation weight can rise again. That makes the method both more faithful to agentic reasoning and more forgiving of partial recovery.
Limitations
- SOD relies on student-teacher log-probability differences as a proxy for reliability. If the teacher is itself weak on a tool pattern, or if teacher and student differ substantially in style, the divergence score may partially reflect teacher error rather than student drift, which can distort weighting.
- The evaluation is centered on the Qwen3 family and on math, science, and code benchmarks. Broader generalization to web navigation, GUI interaction, memory-heavy long-horizon tasks, or multi-tool ecosystems remains unproven, especially where state transitions are more heterogeneous.
- Although clipping improves stability, SOD still combines GRPO and OPD, making the training pipeline more complex than plain SFT. The robustness of δ, ϵ, and teacher scale choices deserves deeper systematic study.
Future Work
Future work could extend the step-wise divergence idea to finer-grained tool substeps, function-call boundaries, or multimodal environments, and learn the weighting function automatically rather than using a hand-designed ratio rule. Another promising direction is to combine reliability-aware weighting with uncertainty estimation, verifiers, or counterfactual trajectories, so that small agents can distinguish student drift from teacher unreliability in longer open-ended settings.
AI Executive Summary
Deep Dive
Abstract
Tool-integrated reasoning (TIR) is difficult to scale to small language models due to instability in long-horizon tool interactions and limited model capacity. While reinforcement learning methods like group relative policy optimization provide only sparse outcome-level rewards. Recently, on-policy distillation (OPD) has gained popularity by supplying dense token-level supervision from a teacher on student-generated trajectories. However, our experiments indicate that applying OPD to TIR leads to a critical failure mode: erroneous tool calls tend to cascade across subsequent reasoning steps, progressively amplifying student-teacher divergence and rendering the teacher's token-level supervision increasingly unreliable. To address this, we propose SOD, a step-wise on-policy distillation framework for small language model agents, which adaptively reweights distillation strength at each step based on step-level divergence. Therefore, SOD can attenuate potentially misleading teacher signals in high-divergence regions while preserving dense guidance in well-aligned states. Experiments on challenging math, science, and code benchmarks show that SOD achieves up to 20.86% improvement over the second-best baseline. Notably, our 0.6B student achieves 26.13% on AIME 2025, demonstrating effective transfer of agentic reasoning to lightweight models. Our code is available at https://github.com/YoungZ365/SOD.