When Top-K Misses the Decision: Tool-Call Drift in Multi-Teacher On-Policy Distillation

2026/07/29 by Jiabin Shen, Guang Chen, Chengjun Mao
Computer Science · #cs.CL #cs.LG

paper · pdf

33 pages, 5 figures. Code and aggregate artifacts: https://github.com/shen-jiabin/topk-support-opd

arxiv created 2026/08/02 · arxiv updated 2026/08/04

Abstract

Top-K teacher logits make on-policy distillation tractable, but probability mass is not decision support. In a two-teacher tool-use setting, vanilla generalized knowledge distillation raises tool-call recall but also over-calls on direct-answer examples. With Qwen3.5-9B, the response teacher's top-32 retains 99.99% of its probability mass yet contains the behavior-switch tool-call token on only 0.4% of 500 unique response prompts; even top-256 covers only 52.2%. Student logits for tokens omitted from the truncated support receive zero direct gradient, so the tool teacher reinforces entry while the response teacher usually cannot oppose it; a wrong entry then amplifies divergence along the generated trajectory. Matched restoration connects this blind spot to behavior: restoring the tool-call token at every supervised response position lowers full-generation over-calling from 14.2+/-2.1% to 3.7+/-0.5% across three seeds, but also lowers call recall by 12.4 points. A non-tool placebo changes over-calling by only 0.95 points, whereas a teacher/student top-32 support-union intervention reaches 7.4+/-0.6% over-calling with 87.0+/-2.0% call recall. The pattern recurs with Llama-3.1-8B's native JSON protocol: the response teacher omits the JSON-entry token through top-256 despite 99.977% retained mass, whereas the tool teacher ranks it first on every audited tool prompt; the matched support-union intervention lowers over-calling from 28.8+/-0.8% to 11.1+/-1.2%. These results identify decision-critical support omission as a causal mechanism in the primary Qwen setting and support it across model families. They also trace a restraint-capability trade-off among support-, loss-, and decoding-level interventions, motivating support-aware audits of compressed distillation. Code and aggregate artifacts are available at https://github.com/shen-jiabin/topk-support-opd.

Citations

• On Group Relative Policy Optimization Collapse in Agent Search: The Lazy Likelihood-Displacement
• LLM-Oriented Token-Adaptive Knowledge Distillation
• SimpleTIR: End-to-End Reinforcement Learning for Multi-Turn Tool-Integrated Reasoning
• ToDi: Token-wise Distillation via Fine-Grained Divergence Control
• When2Call: When (not) to Call Tools
• APIGen-MT: Agentic Pipeline for Multi-Turn Data Generation via Simulated Agent-Human Interplay
• Search-R1: Training LLMs to Reason and Leverage Search Engines with Reinforcement Learning
• DistiLLM-2: A Contrastive Approach Boosts the Distillation of LLMs
• SMART: Self-Aware Agent for Tool Overuse Mitigation
• On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes
• MiniLLM: On-Policy Distillation of Large Language Models
• Multi-Task Learning Using Uncertainty to Weigh Losses for Scene Geometry and Semantics

Related

• Revisiting On-Policy Distillation: Empirical Failure Modes and Simple Fixes
• CADENCE: Closing the Reasoning Gap via Coverage-Adaptive On-Policy Distillation
• On the Generalization vs Fidelity Paradox in Knowledge Distillation
• Prefix-Guided On-Policy Distillation: Mining Golden Trajectories from Rollouts
• SODA: Semi On-Policy Black-Box Distillation for Large Language Models