cs.LGJun 8, 2026

Escaping the KL Agreement Trap in On-Policy Distillation

Authors: Haoran XinAnhao ZhaoYing SunJin LiXiaoyu ShenHui Xiong

Organizations: Thrust of Artificial Intelligence, The Hong Kong University of Science and Technology (Guangzhou) · The Hong Kong Polytechnic University · Eastern Institute of Technology, Ningbo · Department of Computer Science and Engineering, The Hong Kong University of Science and Technology, Hong Kong SAR, China

Abstract

On-policy distillation (OPD) provides dense token-level supervision by asking a teacher to score student-generated rollouts. However, when the student drifts into an unrecoverable prefix, the teacher may locally agree with the degraded state, producing low reverse KL but little corrective training signal. We identify this persistent regime as a low-KL agreement trap. Further analyses show that tokens during and after such traps produce less useful supervision signals. We propose KAT (KL Agreement Trap Termination), an online OPD termination rule that detects persistent low-KL agreement with a dynamic training-adaptive threshold. By filtering weak supervision from degenerate agreement, KAT improves avg@k accuracy by 2.66% and pass@k by 3.43% across four mathematical benchmarks, while reducing average rollout length by 59.73%.

Explore similar work

CardsList