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

2026-05-08 · Qiyong Zhong, Mao Zheng, Mingyang Song, Xin Lin, Jie Sun, Houcheng Jiang, Xiang Wang, Junfeng Fang arxiv

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.

📄 PDF Abstract BibTeX arXiv:2605.07725

Code (0)

등록된 구현이 없습니다.

Tasks

Reinforcement Learning

Improving Reasoning Capabilities in Small Models through Mixture-of-Layers Distillation with Stepwise Attention on Key Information

2026-04-17 · Yao Chen, Jiawei Sheng, Wenyuan Zhang, Tingwen Liu arxiv

The significant computational demands of large language models have increased interest in distilling reasoning abilities into smaller models via Chain-of-Thought (CoT) distillation. Current CoT distillation methods mainl…

Trust Region Policy Distillation

2026-07-06 · Zhengpeng Xie, Li Lyna Zhang, Zeke Xie, Mao Yang arxiv

Big goals are hard to achieve all at once; breaking them into small steps is wiser. We present Trust Region Policy Distillation (TOP-D), which transforms the notoriously unstable, high-variance On-Policy Distillation (OP…

Mathematical Reasoning

PaD: Program-aided Distillation Can Teach Small Models Reasoning Better than Chain-of-thought Fine-tuning

2023-05-23 · Xuekai Zhu, Biqing Qi, Kaiyan Zhang, Xinwei Long 외

While large language models (LLMs) excel in various natural language processing tasks, their huge size and the inaccessibility of parameters present challenges for practical deployment. Previous studies try to distill ta…

Arithmetic Reasoning GSM8K Hallucination

SelecTKD: Selective Token-Weighted Knowledge Distillation for LLMs

2025-10-28 · Haiduo Huang, Jiangcheng Song, Yadong Zhang, Pengju Ren arxiv

Knowledge distillation (KD) is a standard route to compress Large Language Models (LLMs) into compact students, yet most pipelines uniformly apply token-wise loss regardless of teacher confidence. This indiscriminate sup…

Knowledge Distillation Mathematical Reasoning Instruction Following Code Generation

Anti-Self-Distillation for Reasoning RL via Pointwise Mutual Information

2026-05-12 · Guobin Shen, Xiang Cheng, Chenxiao Zhao, Lei Huang 외 arxiv

On-policy self-distillation, where a student is pulled toward a copy of itself conditioned on privileged context (e.g., a verified solution or feedback), offers a promising direction for advancing reasoning capability wi…