# UNDO: Understanding Distillation as Optimization

arXiv: https://arxiv.org/abs/2504.02521 · HTML: https://mlsanity.com/paper/2504.02521/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 30 of its 36 references that are in the map.
- First posted on arXiv 3 Apr 2025.

## Builds on (most cited first, 30 of 30)

- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Chain-of-Thought Prompting Elicits Reasoning in Large Language Models](https://mlsanity.com/paper/2201.11903/) (2022)
- [Distilling the Knowledge in a Neural Network](https://mlsanity.com/paper/1503.02531/) (2015)
- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter](https://mlsanity.com/paper/1910.01108/) (2019)
- [Let's Verify Step by Step](https://mlsanity.com/paper/2305.20050/) (2023)
- [Qwen Technical Report](https://mlsanity.com/paper/2309.16609/) (2023)
- [MMLU-Pro: A More Robust and Challenging Multi-Task Language Understanding Benchmark](https://mlsanity.com/paper/2406.01574/) (2024)
- [Knowledge Distillation: A Survey](https://mlsanity.com/paper/2006.05525/) (2020)
- [MiniLM: Deep Self-Attention Distillation for Task-Agnostic Compression of Pre-Trained Transformers](https://mlsanity.com/paper/2002.10957/) (2020)
- [STaR: Bootstrapping Reasoning With Reasoning](https://mlsanity.com/paper/2203.14465/) (2022)
- [On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes](https://mlsanity.com/paper/2306.13649/) (2023)
- [Are NLP Models really able to Solve Simple Math Word Problems?](https://mlsanity.com/paper/2103.07191/) (2021)
- [Did Aristotle Use a Laptop? A Question Answering Benchmark with Implicit Reasoning Strategies](https://mlsanity.com/paper/2101.02235/) (2021)
- [Distilling Step-by-Step! Outperforming Larger Language Models with Less Training Data and Smaller Model Sizes](https://mlsanity.com/paper/2305.02301/) (2023)
- [WizardMath: Empowering Mathematical Reasoning for Large Language Models via Reinforced Evol-Instruct](https://mlsanity.com/paper/2308.09583/) (2023)
- [SmolLM2: When Smol Goes Big -- Data-Centric Training of a Small Language Model](https://mlsanity.com/paper/2502.02737/) (2025)
- [Teaching Small Language Models to Reason](https://mlsanity.com/paper/2212.08410/) (2022)
- [TheoremQA: A Theorem-driven Question Answering dataset](https://mlsanity.com/paper/2305.12524/) (2023)
- [Distilling Reasoning Capabilities into Smaller Language Models](https://mlsanity.com/paper/2212.00193/) (2022)
- [Symbolic Chain-of-Thought Distillation: Small Models Can Also "Think" Step-by-Step](https://mlsanity.com/paper/2306.14050/) (2023)
- [Baby Llama: knowledge distillation from an ensemble of teachers trained on a small dataset with no performance penalty](https://mlsanity.com/paper/2308.02019/) (2023)
- [Survey on Knowledge Distillation for Large Language Models: Methods, Evaluation, and Application](https://mlsanity.com/paper/2407.01885/) (2024)
- [Distilling Mathematical Reasoning Capabilities into Small Language Models](https://mlsanity.com/paper/2401.11864/) (2024)
- [Democratizing Reasoning Ability: Tailored Learning from Large Language Model](https://mlsanity.com/paper/2310.13332/) (2023)
- [FIRST: Teach A Reliable Large Language Model Through Efficient Trustworthy Distillation](https://mlsanity.com/paper/2408.12168/) (2024)
- [SIKeD: Self-guided Iterative Knowledge Distillation for mathematical reasoning](https://mlsanity.com/paper/2410.18574/) (2024)
- [First-Step Advantage: Importance of Starting Right in Multi-Step Math Reasoning](https://mlsanity.com/paper/2311.07945/) (2023)
- [SMART: Self-learning Meta-strategy Agent for Reasoning Tasks](https://mlsanity.com/paper/2410.16128/) (2024)

## Built on by (most cited first, 1 of 1)

- [Recursive Meta-Distillation: An Axiomatic Framework for Iterative Knowledge Refinement](https://mlsanity.com/paper/2601.13100/) (2026)

## Abstract (arXiv)

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.