[Paper Review] Towards Understanding Knowledge Distillation

Mary Phuong, Christoph H. Lampert|arXiv (Cornell University)|May 27, 2021

Machine Learning and Algorithms133 citations

TL;DR

The paper provides a theoretical analysis of knowledge distillation for linear and deep linear models, proving fast generalization and identifying three key factors—data geometry, optimization bias, and strong monotonicity—that drive transfer performance.

ABSTRACT

Knowledge distillation, i.e., one classifier being trained on the outputs of another classifier, is an empirically very successful technique for knowledge transfer between classifiers. It has even been observed that classifiers learn much faster and more reliably if trained with the outputs of another classifier as soft labels, instead of from ground truth data. So far, however, there is no satisfactory theoretical explanation of this phenomenon. In this work, we provide the first insights into the working mechanisms of distillation by studying the special case of linear and deep linear classifiers. Specifically, we prove a generalization bound that establishes fast convergence of the expected risk of a distillation-trained linear classifier. From the bound and its proof we extract three key factors that determine the success of distillation: * data geometry -- geometric properties of the data distribution, in particular class separation, has a direct influence on the convergence speed of the risk; * optimization bias -- gradient descent optimization finds a very favorable minimum of the distillation objective; and * strong monotonicity -- the expected risk of the student classifier always decreases when the size of the training set grows.

Motivation & Objective

• Motivate and analyze knowledge distillation beyond empirical observations.
• Derive a generalization bound showing fast convergence of distillation-trained linear classifiers.
• Identify and explain three factors that determine the success of distillation: data geometry, optimization bias, and strong monotonicity.
• Show that distillation can recover the teacher’s weights with finite samples when n >= d.

Proposed method

• Model the distillation setup with a linear teacher and a linear student (direct or deep linear networks).
• Train the student on soft labels produced by the sigmoid of the teacher’s output using infinitesimal gradient flow.
• Derive closed-form asymptotic solutions for the student’s end-to-end weights under gradient flow.
• Prove a transfer risk bound showing zero risk for n >= d and distribution-dependent bounds for n < d.
• Introduce geometric quantities (angles between w* and data) to bound transfer risk.
• Discuss how data geometry, optimization bias, and monotonicity influence learning dynamics and transfer efficiency.

Experimental results

Research questions

• RQ1 Under what conditions can a distillation-trained linear student recover the teacher’s weights with finite samples?
• RQ2 How fast does the student learn from soft labels, and how does data geometry affect the transfer risk?
• RQ3 What are the roles of optimization dynamics and data distribution in distillation success?
• RQ4 How does increasing training data influence the transfer risk (monotonicity) in linear distillation?

Key findings

• If n >= d, the student perfectly identifies the teacher’s weight vector with probability 1 (almost surely).
• When n < d, the student learns the projection of the teacher’s weights onto the data span, i.e., the best subspace-constrained approximation.
• The transfer risk decays to zero for n >= d, and for n < d it is bounded by a distribution-dependent expression involving angular geometry between w* and data.
• For large-margin or well-aligned data distributions, the transfer risk decays exponentially or at rates characterized by a polynomial bound in n (Corollaries 1 and 2).
• The results reveal three key factors: data geometry (class separation and alignment with w*), optimization bias (gradient descent converges to favorable minima), and strong monotonicity (adding data never increases transfer risk).
• The theory provides non-vacuous, finite-sample guarantees contrasting with classical hard-label learning, including fast convergence and explicit risk bounds.