PaperlayerView exact revision on arXiv

Source-linked AI summary

SOD: Step-wise On-policy Distillation for Small Language Model Agents

Qiyong Zhong, Mao Zheng, Mingyang Song, Xin Lin, Jie Sun, Houcheng Jiang, Xiang Wang, Junfeng Fang

arXiv:2605.07725v3cs.CLcs.AI

TL;DR

Small-model tool-integrated reasoning is challenging because long-horizon interactions and sparse rewards can destabilize learning, while erroneous tool calls can make teacher supervision unreliable. SOD adjusts distillation strength by step-level divergence, and experiments report improvements across math, science, and code benchmarks.

  • Problem

    Small-model tool-integrated reasoning is difficult to train stably: long-horizon tool interactions meet sparse outcome-level rewards, and erroneous tool calls can make later teacher supervision unreliable.

    Source: §1 · p. 1

  • Method

    SOD adaptively reweights on-policy distillation at each reasoning step using student–teacher divergence, preserving guidance in aligned states and attenuating it in high-divergence states.

    Source: §1 · p. 2, §4.2 · p. 5

  • Results

    Up to 20.86% improvement over the second-best baseline is reported across challenging math, science, and code benchmarks.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The findings suggest that agent distillation benefits from supervision mechanisms that adapt to evolving state distributions.

    Source: §7 · p. 10

    Experiments use the Qwen3 model family and Python code interpreter, so validation on other model families and agentic settings remains future work.

    Source: §6 · p. 9

Abstract

from arXiv · show

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.

1 Introduction

Small language models are important for deploying agentic reasoning in constrained settings, but stable tool-integrated reasoning remains challenging. SOD adapts on-policy distillation to step-level divergence, and experiments report improvements over baselines.

1 Introduction

  • Transferring agentic capabilities to on-device small language models matters in latency-, resource-, and privacy-sensitive settings where large models incur deployment costs and complexity.The passage identifies inference cost, deployment overhead, and system complexity as concerns for large-scale models.Source: §1 · p. 1
  • For small-model TIR, long-horizon tool interactions and sparse outcome-level rewards can exacerbate exploration failure and leave policies in a cold-start regime.OPD instead provides dense token-level supervision on student-generated trajectories, alleviating sparse-reward credit assignment and improving sample efficiency and training stability.Source: §1 · p. 1, Figure 1 · p. 2
  • Erroneous tool calls can corrupt later reasoning states, amplify student–teacher divergence, and make teacher token-level supervision unreliable or misleading.The resulting discrepancy can produce unstable gradients and even training collapse.Source: §1 · p. 2
  • SOD adapts distillation strength to student–teacher divergence, preserving dense guidance in aligned steps and attenuating potentially misleading signals at high divergence.Across math, science, and code benchmarks, it improves over the second-best baseline by up to 20.86%.Source: §1 · p. 2
  • A 0.6B student achieves 26.13% accuracy on AIME 2025, and SOD improves over the second-best baseline by up to 20.86%.The authors describe this as the first sub-billion-parameter model to reach that level on AIME 2025.Source: §1 · p. 2

2 Related Work

Prior work applies reinforcement learning and on-policy distillation to language agents, addressing tool use and supervision challenges through different training signals.

Reinforcement Learning for Agents

  • Agent RL research extends from RLHF with PPO to scalable GRPO, while language-agent approaches include tool-use paradigms and RL across code, tools, GUIs, and web navigation.ReAct, Toolformer, and FireAct rely on demonstrations rather than online optimization.Source: §2 · pp. 2–3
  • Agentic RL faces sparse, delayed feedback and credit assignment challenges, while distribution shift and compounding errors can amplify optimization instability.Broader frameworks still rely on trajectory-level rewards without dense supervision.Source: §2 · p. 3

On-policy Distillation

  • OPD trains on student-generated trajectories with token-level supervision; prior work studies its divergence objectives, reverse-KL formulation, implicit rewards, compatibility, and self-distillation extensions.Recent work also explores OPD to mitigate RL limitations.Source: §2 · p. 3

3 Preliminaries

The preliminaries define multi-turn tool-integrated trajectories and introduce GRPO and on-policy distillation as post-training approaches with different supervision signals.

3.1 Multi-turn Tool-Integrated Reasoning

  • A multi-turn TIR trajectory comprises the input, successive model responses and any tool observations, and a final response; tool observations condition later generations.The policy generates model tokens, while the environment supplies observations.Source: §3.1 · p. 3

3.2 Group Relative Policy Optimization

  • GRPO samples trajectory groups from the old policy and uses relative rewards for on-policy learning, but supplies only sparse outcome-level rewards.The reward is defined on complete trajectories.Source: §3.2 · p. 3

3.3 On-policy Distillation

  • OPD provides dense token-level supervision on student-generated trajectories by aligning the student policy with a teacher distribution.Its objective is a sampled estimator of reverse KL divergence on student-visited states.Source: §3.3 · pp. 3–4

4 Methodology

The methodology explains why standard OPD becomes unreliable under tool-induced state drift, then introduces a divergence-based step-wise weighting scheme and a combined training objective. SOD reduces distillation strength in high-mismatch steps while preserving or restoring guidance where alignment is stronger.

4 Methodology

  • The methodology proceeds from diagnosing OPD unreliability under tool-induced drift to defining a divergence score and formulating step-wise reweighted OPD.The overview identifies step-level divergence as the basis for adaptive reweighting.Source: §4 · p. 4, Figure 2 · p. 5

4.1 Failure of On-policy Distillation under Tool-Induced State Drift

  • Tool errors can amplify student–teacher divergence and push trajectories into low-overlap states where OPD gradients become high-variance and uninformative.The propositions characterize divergence amplification and gradient signal-to-noise degradation; the authors say reweighting suppresses the weighted second moment under monotonically increasing divergence.Source: §4.1 · p. 4

4.2 Student-teacher Divergence and Step-wise Reweighting

  • SOD uses step-level divergence as a proxy for teacher-supervision reliability and adjusts distillation weights as divergence changes across steps.When divergence rises, weights attenuate supervision; when it falls, weights can increase again, subject to an upper bound of 1 + δ.Source: §4.2 · p. 5

4.3 Training Objective

  • The training objective combines sparse-reward GRPO with dense token-level OPD whose strength is reweighted by student–teacher divergence.The trajectory is partitioned into model-generated reasoning steps, excluding tool-observation tokens from the model-token sets used for step-level treatment.Source: §4.2 · p. 4, §4.3 · pp. 5–6

5 Experiment

SOD is evaluated against supervised, reinforcement-learning, and distillation baselines on math, science, and code benchmarks at two student scales. It leads all baselines, while ablations and training analyses examine its reweighting, component roles, scalability, dynamics, and distillation patterns.

5.1 Experimental Setup

  • The 3k SFT corpus and 30k RL dataset support evaluation on AIME 2024/2025, GPQA-Diamond, and LiveCodeBench, using Qwen3-0.6B and Qwen3-1.7B students.Evaluation uses temperature 1.0, top_p=0.6, and average@32; the default teacher is Qwen3-4B optimized with GRPO.Source: §5.1 · p. 6

5.2 Overall Performance

  • Across all four tasks, SOD achieves the highest scores at both model scales, exceeding strongest baseline OPD by relative average improvements of 20.86% at 0.6B and 18.50% at 1.7B.At 0.6B, SFT and GRPO both fail to outperform Vanilla.Source: §5.2 · pp. 6–7

5.3 Ablation Study

  • 34.70% average with uniform weighting, 37.14% with heuristic decay, 31.85% with hard masking, and 38.10% without clipping show that the tested alternatives underperform adaptive reweighting.The authors attribute the shortfall to equal weighting of corrupted states, fixed schedules missing recovered alignment, discarding later supervision, and instability from unbounded amplification.Source: §5.3 · p. 7, Table 2 · p. 7
  • 40.78% average without GRPO versus 25.39% without step-wise OPD indicates that both objective components contribute, with step-wise OPD providing the primary fine-grained guidance.GRPO broadens trajectory-level exploration beyond the teacher distribution, while step-wise OPD supplies fine-grained guidance.Source: §5.3 · p. 7, Table 2 · p. 7

5.4 Scalability and Generalization

  • SOD with a 14B teacher outperforms its 4B-teacher counterpart at both student scales, whereas OPD’s 0.6B average accuracy declines when its teacher grows from 4B to 14B.The authors attribute OPD’s decline to increased capacity-gap distribution mismatch and describe SOD as suppressing harmful signals while using richer teacher knowledge.Source: §5.4 · p. 8, Figure 3 · p. 8

5.5 Dynamic training analysis

  • SOD achieves higher accuracy than GRPO across both scales; GRPO’s entropy collapse accompanies repetitive outputs and abandonment of multi-step tool interactions.Compared with OPD, SOD sustains improvement without degradation and uses fewer tool-calling turns.Source: §5.5 · pp. 8–9, Figure 4 · p. 8

5.6 Three Distillation Patterns of SOD

  • SOD’s stable, erroneous, and recovery patterns correspond to high weights under low divergence, attenuation under persistent divergence, and restored weights after alignment returns.This permits supervision to remain useful after recovery rather than continuing to penalize temporary errors.Source: §5.6 · p. 9, Figure 5 · p. 9

6 Limitations

The study’s scope is limited to python-interpreter tool-integrated reasoning and the Qwen3 model family, leaving other agent settings and model families for future validation.

6 Limitations

  • Experiments focus on agentic TIR with a python code interpreter, so findings do not establish how other tool settings behave.The interpreter provides a clean, reproducible environment for isolating tool-induced state drift from complex API configurations.Source: §6 · p. 9
  • All experiments use Qwen3 models; validation on other model families remains an open question.Qwen3 was selected for performance across scales, availability, and adoption in OPD studies; resource constraints deferred broader validation.Source: §6 · p. 9
  • Web browsing and API calls may exhibit different drift patterns, and the authors identify them as settings for future exploration.The paper does not test these settings.Source: §6 · p. 9

7 Conclusion

SOD adapts step-level distillation to student–teacher divergence, and the paper’s conclusion reports improvements across math, science, and code benchmarks. Its training combines on-policy tool rollouts, divergence-based weighting, and GRPO with step-wise OPD.

7 Conclusion

  • SOD improves training stability and final performance across math, science, and code benchmarks, and frames adaptive supervision as a general principle for robust agentic training.The conclusion attributes unreliable teacher supervision to tool-induced transitions pushing students into out-of-distribution regions.Source: §7 · pp. 9–10

A Algorithmic Details of SOD

  • SOD generates on-policy tool-interaction trajectories, computes step-level student–teacher divergence, and derives adaptive distillation weights for each reasoning step.The algorithm partitions model-generated tokens into steps, excluding tool observations, and sets the first step’s weight to 1.Source: A Algorithmic Details of SOD · p. 16

Stage III: Combined Optimization

B.1 Training Datasets.

  • Training uses a 3k high-quality SFT corpus and a diverse 30k RL dataset spanning mathematical, coding, and scientific problems.The RL sources include 17k DAPO-Math problems, 4,902 math and 3,586 code tasks from Skywork-or1, and 3k MegaScience problems.Source: §B.1 · pp. 16–17, Table 3 · p. 17

B.2 Benchmarks

The evaluation uses AIME, GPQA-Diamond, and LiveCodeBench benchmarks, compares supervised, RL, and distillation baselines, and applies a shared implementation setup across methods.

B.2 Benchmarks

  • The four benchmarks are AIME 2024/2025, GPQA-Diamond, and LiveCodeBench, covering advanced mathematics, scientific reasoning, and coding.AIME uses recent competition problems; GPQA-Diamond is an expert-validated subset; LiveCodeBench v6 evaluates recent coding problems across several capabilities.Source: §B.2 · p. 17

B.3 Baselines

  • Baselines span supervised learning, reinforcement learning, and distillation, including SFT, GRPO, OPD, and two on-policy self-distillation variants.OPSDgt uses full ground-truth solutions as privileged information, while OPSDhint uses a hint; GRPO supplies sequence-level rewards and OPD token-level teacher supervision.Source: §B.3 · p. 18

B.4 Implementation Details.

  • Experiments use shared training infrastructure, optimizer, and core hyperparameters for fair comparison; RL and distillation baselines start from SFT checkpoints, while OPD variants receive token-level supervision from a GRPO-optimized Qwen3-4B teacher along students’ on-policy rollouts.RL and distillation methods use AdamW at 1e-6, batch size 64, and mini-batch size 16; they train for at most 1 epoch, except teacher models at most 2 epochs, and support up to 16 tool-call turns.Source: §B.4 · pp. 18–19

C.1 Training Cost and Practicality Analysis

SOD’s training-cost comparison reports substantial performance gains over OPD with modest overhead for the 1.7B student, while 0.6B training is slightly faster per step.

C.1 Training Cost and Practicality Analysis

  • 3.5% faster per step for the 0.6B model: SOD takes 1052.3s versus OPD’s 1090.5s.The authors attribute this to suppressing learning from erroneous tool-call patterns, yielding fewer failed retries and shorter responses.Source: §C.1 · p. 19
  • +18.50% average improvement over OPD accompanies +4.9% overhead for the 1.7B model.Per-step times are 1105.4s for SOD and 1053.6s for OPD.Source: §C.1 · p. 20
  • GRPO’s lower per-step time reflects training collapse, with progressively less effective tool use and shorter responses rather than faster equivalent work.OPD and SOD maintain longer reasoning trajectories and active tool use, which requires more generation time per step.Source: §C.1 · p. 20

C.2 Quantitative Analysis of Distillation Pattern Distribution

The analysis tracks how stable, recovery, and erroneous trajectories change during training, then examines why on-policy distillation can become unreliable in tool-integrated reasoning. It connects cascading tool errors and weak teacher-student overlap to noisy supervision, motivating adaptive step-wise weighting.

C.2 Quantitative Analysis of Distillation Pattern Distribution

  • For both model scales, stable trajectories increase and erroneous trajectories decrease; recovery rises and stabilizes, while 1.7B reaches a higher stable proportion faster than 0.6B.The 0.6B model starts with more erroneous and fewer recovery trajectories and converges more slowly.Source: §C.2 · pp. 20–21, Figure 6 · p. 20

C.3 The Performance of Naive On-policy Distillation

  • 33.79 vs. 40.78: SOD without GRPO exceeds naive OPD by 6.99 points (+20.69%), while naive OPD averages 33.79 versus OPD’s 36.27.The authors attribute the improvement to step-wise reweighting, with GRPO providing complementary gains.Source: §C.3 · p. 21, Table 5 · p. 21

D Proofs for Section 4.1

  • The proofs define student and teacher next-token distributions on student-generated prefixes and express token-level OPD loss using student-sampled tokens.The step-level mismatch and overlap terms follow the main-text definitions.Source: D Proofs for Section 4.1 · p. 21

D.1 Proof of Proposition 1

  • Consecutive erroneous tool calls compound divergence super-linearly because later failures further corrupt prefixes and increase subsequent per-token distribution shifts.The accumulated multi-error context becomes less familiar to the teacher, making its conditional predictions increasingly unreliable.Source: §D.1 · pp. 22–23

D.2 Proof of Proposition 2

  • As student-teacher overlap approaches zero, OPD’s gradient signal-to-noise ratio tends to zero because informative teacher-supported tokens become rare.Low-overlap tokens can produce large but uninformative loss contributions.Source: §D.2 · p. 24

D.3 Discussion: Failure Cascade and Design Implications

  • The failure cascade runs from tool-error-induced divergence through low-overlap, high-variance supervision to uniform aggregation that gives corrupted steps equal weight.The proposed design principle is to preserve dense supervision in aligned steps and attenuate it when divergence signals teacher miscalibration.Source: §D.3 · p. 24

D.4 Variance Reduction under Step-wise Reweighting

Step-wise weighting suppresses high-divergence contributions while preserving useful signal from aligned steps, and its divergence proxy avoids the cost of computing a full KL. The included prompt sections describe task-specific prompts and how the OPSDhint baseline creates partial guidance.

D.4 Variance Reduction under Step-wise Reweighting

  • Under monotonically increasing divergence, step-wise reweighting keeps weighted gradient contributions bounded and restores positive aggregate SNR by suppressing high-divergence noise.The factor (d1 + ϵ)^2/(dk + ϵ)^2 attenuates later contributions while leaving well-aligned early-step signal intact.Source: §D.4 · pp. 25–26

D.5 Justification of dk as a Proxy for ∆k

  • The proxy dk uses the student-sampled token’s teacher log-probability and preserves the ordering of divergence needed for step-wise attenuation and recovery.It avoids evaluating the teacher’s full output distribution at each token position.Source: §D.5 · p. 26

E.1 Datasets Prompt Template

  • The dataset includes mathematics, programming, and scientific problems, with specialized prompts for each task type.The prompts are presented following Yu et al. [52].Source: §E.1 · p. 26

E.2 Baseline Prompt Template

  • OPSDhint conditions on a concise hint distilled from the ground-truth solution, intended to provide partial guidance without revealing the answer or full derivation.The hint-generation instructions request a key insight or intermediate step, such as a substitution, theorem, or decomposition.Source: §E.2 · p. 27, Hint Generation Prompt · p. 28

Hint Generation Prompt

  • The hint-generation prompt asks for a brief directional suggestion and explicitly excludes the final answer, a complete derivation, and tool-use instructions.It can request a useful substitution, relevant theorem, or strategic decomposition.Source: Hint Generation Prompt · p. 28

F Visualization and Analysis of Token Entropy from The Teacher

The section uses teacher conditional entropy as an indicator of confidence in supervision along student-generated trajectories. Its examples contrast a stable, successful tool-use trajectory with repeated tool failures accompanied by sharply rising entropy.

Entropy interpretation

  • Teacher conditional entropy reflects confidence: low entropy indicates concentrated, reliable guidance, while high entropy signals uncertainty under the student’s context.The paper treats high entropy as a hallmark of out-of-distribution states where distillation becomes unreliable.Source: F Visualization and Analysis of Token… · p. 28

Erroneous trajectory

  • In the erroneous example, mean teacher entropy rises from 0.85 to 2.14 across steps, with over 78% of final-step tokens above H=1.0.The authors interpret high entropy as teacher uncertainty under corrupted context and say SOD attenuates distillation in these regions.Source: F Visualization and Analysis of Token… · pp. 28–29

Stable trajectory

G Case Study

The case studies illustrate how adaptive weights respond to stable, erroneous, and recovering trajectories. They show high weights retained for aligned reasoning, reduced weights after repeated failures, and weight recovery after a student corrects course.

Stable pattern

  • In a stable trajectory, moderate divergence and low teacher entropy accompany weights at or near 1.2 after the first step.The authors contrast this with uniform weighting and trajectory-level filtering, which may fail to preserve guidance for each reliable step.Source: G Case Study · p. 30

Erroneous pattern

  • After three repeated numpy.roots failures, divergence rises from 0.22 to 1.28, teacher entropy from 0.32 to 2.38, and the adaptive weight falls from 1.00 to 0.17.The student then guesses 15 instead of the correct 1536.Source: G Case Study · p. 31

Recovery pattern

  • After three indentation failures, the student rewrites the code successfully; weights recover from 0.41 to 1.20 by the final step.The divergence and entropy fall after the rewrite, while the trajectory-level alternative would weight all seven steps uniformly.Source: G Case Study · p. 32
Loading 2605.07725v3…