Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

UNDO: Understanding Distillation as Optimization

Kushal Jain, Piyushi Goyal, Kumar Shridhar

TL;DR

The paper addresses the mismatch between teacher rationales and student learning in knowledge distillation by introducing UNDO, an iterative, gap-driven distillation framework. At each cycle, the student’s weaknesses are diagnosed on a validation set, guiding the teacher to regenerate targeted rationales that address those gaps, followed by fine-tuning the student and checking for convergence. Empirical results across multiple math and reasoning tasks show that UNDO consistently improves over standard one-shot distillation, with gains up to 20% on challenging benchmarks and robustness across different student models and cross-teacher data. The approach reframes distillation as a dynamic, iterative teacher–student interaction, with theoretical grounding in KL divergence and practical implications for personalized rationale transfer and improved out-of-domain generalization.

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.

UNDO: Understanding Distillation as Optimization

TL;DR

The paper addresses the mismatch between teacher rationales and student learning in knowledge distillation by introducing UNDO, an iterative, gap-driven distillation framework. At each cycle, the student’s weaknesses are diagnosed on a validation set, guiding the teacher to regenerate targeted rationales that address those gaps, followed by fine-tuning the student and checking for convergence. Empirical results across multiple math and reasoning tasks show that UNDO consistently improves over standard one-shot distillation, with gains up to 20% on challenging benchmarks and robustness across different student models and cross-teacher data. The approach reframes distillation as a dynamic, iterative teacher–student interaction, with theoretical grounding in KL divergence and practical implications for personalized rationale transfer and improved out-of-domain generalization.

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.

Paper Structure

This paper contains 24 sections, 7 equations, 3 figures, 4 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Knowledge distillation
  4. Iterative distillation
  5. Methodology
  6. Problem Setup.
  7. Standard Knowledge Distillation
  8. UNDO: UNderstanding Distillation as Optimization
  9. 1) Identify Student Learning Gaps & Incorporate in Teacher Prompt.
  10. 2) Regenerate from the Teacher with Gap Analysis.
  11. 3) Fine-tune the Student.
  12. 4) Convergence Check.
  13. KL Divergence Perspective.
  14. Experimental Details
  15. Dataset
  16. ...and 9 more sections

Figures (3)

  • Figure 1: We compare UNDO: UNderstanding Distillation as Optimization (bottom) with Standard Distillation (top). The main difference lies in the iterative nature of UNDO where once a knowledge is distilled into the student, the teacher iteratively identifies the learning gap, regenerate its reasoning chains and distill the knowledge again.
  • Figure 2: Accuracy (maj@1) for Qwen and SmolLM2 models on GSM8K, MATH, MMLU Pro and SVAMP datasets for 4 iterations.
  • Figure 3: Accuracy comparison of training standard distillation and UNDO for the equal number of epochs. Note that iteration 1 for UNDO is same as epoch 5 for standard distillation.