UNDO: Understanding Distillation as Optimization

2025-04-03

AI Summary

Standard one-step knowledge distillation suffers from performance limitations due to misalignment between teacher explanations and student learning needs. Method: This paper proposes an iterative knowledge distillation framework that reformulates distillation as a teacher-student co-optimization process. It dynamically identifies student errors to guide the teacher in generating targeted, enhanced reasoning justifications; integrates iterative prompt refinement, error-driven teacher re-prompting, reasoning path alignment, and evaluation-guided distillation cycles to progressively refine teacher explanations. The method generalizes across student models without requiring teacher retraining. Results: On mathematical and commonsense reasoning benchmarks, our approach achieves up to a 20% absolute improvement over standard single-step distillation. Moreover, the distilled high-quality explanations demonstrate strong transferability across diverse student architectures, validating the framework’s generality and practical utility.

Technology Category

• Search and Optimization: Learning to Search
• Knowledge Representation and Reasoning: Computational Complexity of Reasoning
• Machine Learning: Calibration & Uncertainty Quantification

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• User Modeling, Personalization and Recommendation: Explainable and interpretable methods for personalization

Abstract

Knowledge distillation has emerged as an effective strategy for compressing large language models' (LLMs) knowledge into smaller, more efficient student models. However, standard one-shot distillation methods often produce suboptimal results due to a mismatch between teacher-generated rationales and the student's specific learning requirements. In this paper, we introduce the UNDO: UNderstanding Distillation as Optimization framework, designed to bridge this gap by iteratively identifying the student's errors and prompting the teacher to refine its explanations accordingly. Each iteration directly targets the student's learning deficiencies, motivating the teacher to provide tailored and enhanced rationales that specifically address these weaknesses. Empirical evaluations on various challenging mathematical and commonsense reasoning tasks demonstrate that our iterative distillation method, UNDO, significantly outperforms standard one-step distillation methods, achieving performance gains of up to 20%. Additionally, we show that teacher-generated data refined through our iterative process remains effective even when applied to different student models, underscoring the broad applicability of our approach. Our work fundamentally reframes knowledge distillation as an iterative teacher-student interaction, effectively leveraging dynamic refinement by the teacher for better knowledge distillation.

Problem

• Optimizing knowledge distillation for better student model performance
• Addressing mismatch between teacher rationales and student learning needs
• Iterative refinement of teacher explanations to target student weaknesses

Innovation

• Iterative distillation targeting student errors
• Dynamic teacher refinement of explanations
• Broad applicability across student models