Reproduce “When Top-K Misses the Decision: Tool-Call Drift in Multi-Teacher On-Policy Distillation” on OpenResearch

About this paper

Jiabin Shen, Guang Chen, Chengjun Mao

This research reveals how low-probability, decision-critical tokens omitted from top-K teacher logits during on-policy distillation cause student large language models (LLMs) to exhibit excessive tool-calling behavior, termed "tool-call drift." The study establishes a causal link for this omission and evaluates various support, loss, and decoding-level interventions to mitigate this drift while balancing a student model's required tool-calling capabilities.

Key insights

• High probability mass retention in top-K logits is an insufficient diagnostic for distillation fidelity, as decision-critical, low-probability tokens can be causally linked to severe student behavioral drift.
• Decision-critical support omission exhibits directional asymmetry in multi-teacher setups, where a response teacher might omit a tool-call token (low probability in natural responses), while a tool teacher provides attenuated gradients for the same token.
• Interventions at different layers (support, loss, decoding) offer distinct trade-offs in mitigating tool-call drift, impacting over-calling, required-call recall, and multi-turn success differently, necessitating careful selection based on deployment costs and objectives.

Summary by alphaXiv.