Track papers across labs, topics, and trends.
Search papers, labs, and topics across Lattice.
This paper addresses the problem of "confident conflicts" in student-teacher distillation for LLMs, where students are forced to imitate low-probability teacher tokens despite high student confidence in alternatives, leading to generalization degradation. The authors introduce Entropy-Aware Mixing, a dynamic per-token interpolation between student and teacher distributions gated by the student's predictive entropy, to stabilize distillation. Results demonstrate improved in-distribution and out-of-distribution math reasoning and better preservation of general capabilities compared to fixed-teacher supervision, though optimal schedules vary by training source.
Confident conflicts in LLM distillation are mitigated by dynamically mixing student and teacher distributions based on student entropy, boosting reasoning and preserving general capabilities.
Small language models are often post-trained as students on reasoning traces from stronger teacher models to efficiently learn new skills. However, token-level imitation on traces that lie far outside the student's expected distribution often produces \textit{confident conflicts}, whereby the student is required to imitate a continuation that it deems unlikely (i.e., low-probability) despite being confident in a different continuation (i.e., in a low-entropy state). To mitigate the degradation in generalisation and catastrophic forgetting caused by these conflicts, we propose \textbf{Entropy-Aware Mixing}: a dynamic per-token interpolation of the student and teacher distributions, gated by the student's predictive entropy. We implement both convex and geometric interpolations for both offline trace generation (via speculative decoding, then SFT) and on-policy forward-KL distillation. Our results show that entropy-aware mixing stabilises distillation, improving in-distribution and out-of-distribution math reasoning while better preserving general capabilities than fixed-teacher supervision. Nonetheless, the optimal entropy schedule depends on the training source, with offline-generated traces favouring concave schedules (greater overall teacher influence) and on-policy training favouring linear or convex schedules (teacher concentrated in high-entropy states).
AI2