Towards Understanding Knowledge Distillation

In plain words

Knowledge distillation trains a small model not on the hard right-or-wrong answers, but on the probabilities another model assigns. This paper asks why that works, using the simplest possible setting: a straight-line rule as the teacher and a straight-line rule as the student.

In that setting the student can copy the teacher exactly once it has seen as many examples as there are input dimensions. With fewer examples, each new example moves the student closer to the teacher. How close it gets depends on how the data sits relative to the teacher's boundary: data that stays away from the boundary is easier. The training method itself also helps, because plain gradient descent from a near-zero start settles on a particularly good answer.

TL;DR

• Teacher and student are linear classifiers (a deep linear variant is also covered); the student is trained on the teacher's soft outputs §3.
• With at least d examples, the student's weight vector equals the teacher's and its transfer risk is zero [thm. 1] [thm. 3].
• With fewer examples, a bound on expected transfer risk depends on the angle distribution of the data; it gives exponential rates for angular-margin tasks and polynomial rates otherwise [thm. 3] [cor. 1] [cor. 2].
• Results are proven for linear models only; the conclusion treats the nonlinear case as open §6.

Problem

Distillation is widely used, but earlier explanations were qualitative §1. The paper asks what the student learns, how fast its risk falls with more transfer data, and which properties of the task and the training procedure matter §1. Standard learning-theory bounds for hard labels say little when the number of examples is below the dimension; the paper argues that distillation behaves differently in that regime §4.2.

Method

Setting §3

1. Inputs x in R^d are drawn i.i.d. from a density P_x. The teacher is a fixed linear classifier h*(x) = 1{w*^T x >= 0} with w* != 0.
2. The transfer set holds n inputs and soft labels y_i = sigmoid(w*^T x_i). There are no hard labels and no label noise in the setup.
3. The student is h(x) = 1{w^T x >= 0}, with w^T = W_N ... W_1. N = 1 is a direct linear classifier; N >= 2 is a deep linear network.
4. The student minimises the normalised cross-entropy against the soft labels. A constant L* makes the minimum zero; it does not change the optimum or the gradient §3 app. "Properties of the Cross-Entropy Loss".
5. Training is modelled as gradient flow, the limit of gradient descent with an infinitesimal step size; the student is trained until tau -> infinity §3.
6. The transfer risk is R(h) = P_x[h(x) != h*(x)] §3.

Closed-form limit §4.1

7. For N = 1 with w(0) = 0, the weight vector converges to w_hat, given by the formula below [thm. 1]. Its proof uses the fact that the gradient always lies in the span of the data, and strong convexity on that span app..
8. For N >= 2, zero initialisation is a stationary point, so the weights are initialised small but non-zero. Under three conditions (a norm bound on the initial weights, an initial loss below the loss at zero, and balancedness between layers) the end-to-end weight vector approaches w_hat for n >= d, and stays within epsilon of it for n < d [thm. 2].

\hat w = \begin{cases} w^* & n \ge d \\ X (X^\top X)^{-1} X^\top w^* & n < d \end{cases} [thm. 1]

9. For n < d, w_hat is the orthogonal projection of the teacher onto the span of the training inputs. It is the closest point to w* that gradient descent can reach, since all updates stay in that span §4.1.

Transfer-risk bound §4.2

10. For a task, p(theta) is the reverse CDF of the unsigned angle between w* and a random input x §4.2:

    p(\theta) = P_x[\bar\alpha(w^*, x) \ge \theta], \quad \theta \in [0, \pi/2] §4.2

    where bar alpha(u, v) = arccos(|u^T v| / (|u| |v|)) §4.2.

11. Theorem 3: for n >= d the expected risk is zero. For n < d and any beta in [0, pi/2]:

    \mathbb{E}[R(\hat h_X)] \le p(\beta) + p(\pi/2 - \beta)^n [thm. 3]

    Proof idea: one training point on the far side of the angle cone already pulls the student's angle below pi/2 - beta, and a random training set has such a point with high probability §4.3.

12. Corollary 1 (angular margin): if some beta gives p(beta) = 0 and gamma = p(pi/2 - beta) < 1, then the expected risk is at most gamma^n [cor. 1].

13. Corollary 2 (polynomial tail): if p(theta) <= c (1 - (2/pi) theta)^kappa for all theta, then the expected risk is at most c (1 + (log n)^kappa) / n^kappa [cor. 2]. The corollary's proof sets beta = (pi/2) n^(-1/n) [cor. 2].

Strong monotonicity §4.3

14. Lemma (strong monotonicity): for nested full-rank training sets X_- ⊆ X_+, the distillation solution for X_+ has an angle to w* no larger than the one for X_- [lemma, §4.3]. Adding training data never hurts the closed-form student in this sense.

Extended result app., "Theorem 3 for Approximate Distillation"

15. If the student's weight vector is within epsilon of w_hat (with epsilon <= |w_hat| / 2), define delta = sqrt(2 pi epsilon / |w_hat|). Then, for beta in [0, pi/2 - delta], the expected risk is at most p(beta) + p(pi/2 - delta - beta)^n app.. As epsilon -> 0, delta -> 0 and the bound returns to Theorem 3.

Pseudocode of the closed form [thm. 1]

# Student for N = 1 with zero initialisation  [§4.1] [thm. 1]
# Input:  inputs (d, n) i.i.d. from P_x; teacher weights (d,)
# Output: student weights (d,)

def distilled_student(inputs, teacher_weights):
    # Regime is decided by the number of examples  [thm. 1]
    num_dims, num_examples = inputs.shape
    if num_examples >= num_dims:
        # Enough data: the teacher is recovered exactly  [thm. 1]
        return teacher_weights
    # Otherwise project the teacher onto the data span  [thm. 1]
    gram = inputs.T @ inputs                         # (n, n)
    projector = inputs @ inverse(gram) @ inputs.T    # (d, d)
    return projector @ teacher_weights

Pseudocode of the bound [thm. 3]

# Upper bound on expected transfer risk for n < d  [§4.2] [thm. 3]
# Input:  p, the reverse CDF of the teacher angle; n; beta in [0, pi/2]
# Output: scalar bound

def transfer_risk_bound(p, num_examples, beta):
    # Points away from the boundary cut off large angles  [thm. 3]
    near_term = p(beta)
    # Every training point must fall in the wrong cone for an error  [§4.3]
    far_term = p(pi / 2 - beta) ** num_examples
    return near_term + far_term

Hyperparameters

Theory values are symbols; the experiment values are the ones stated in the text.

Results

Author-reported. Theory covers linear teachers and students with exact soft labels. Experiments are synthetic (§5.1, §5.3) and MNIST digits 0 vs 1 (§5.2), with linear learners only.

The angular schematic (Tasks A, B, C) is fig. 1; the example margin and polynomial distributions are fig. 2. The paper gives no headline benchmark numbers, and no comparison with other methods.

Paper vs code

No official code.

Search done on the card's compile date:

• The arXiv metadata has no comments field and no project-page link. The LaTeX sources contain no URLs.
• GitHub repository search for towards understanding knowledge distillation returned only unrelated repositories (zihuixue/MFH, songmzhang/NMT-KD).
• GitHub repository search for linear distillation returned no repository by these authors. postmachines/Linear-Distillation-Learning (MIT) is a third-party repository and is not used.
• GitHub user search for christoph lampert ist austria returned no accounts matching either author. An IST Austria GitHub organisation was not searched.

Pitfalls

• Exact soft labels. The analysis assumes the labels are exactly the sigmoid of the teacher's logit, and the teacher is exactly linear. Noisy labels, nonlinear teachers and a bias term are outside the stated theory §3 §6.
• Gradient flow, not gradient descent. Every theorem is for the infinitesimal-step limit §3. The paper does not analyse finite step sizes.
• Deep-student initialisation. Theorem 2 needs a norm bound on the initial weights, an initial loss below the loss at zero, and balancedness [thm. 2]. The paper calls these practical, and notes that a sign flip can satisfy the loss condition §4.1. Balancedness is only approximated by near-zero initialisation; the paper cites evidence that approximate balancedness may suffice, without proving it §4.1.
• Corollary 2 in the statement and proof. The statement has a constant c and no restriction to n < d. The displayed proof bound has no c, and the source does not reconcile the two [cor. 2]. Check the constant before using the rate.
• Corollary 2 choice of beta. The bound in Theorem 3 holds for any beta. Making it explicit needs a beta chosen from the task, here (pi/2) n^(-1/n) [cor. 2].
• p is unknown in practice. Every rate needs the angle distribution p, which depends on the unknown teacher and the input distribution. The paper gives no way to estimate it from data §4.2.
• Two loss normalisations. The main text uses one normalisation with a constant L* §3. The appendix uses per-example constants [app. "Properties of the Cross-Entropy Loss"]. The constants do not change the optimum or the gradient, so the closed form is the same under both.
• Strong monotonicity is for the closed-form solution. The lemma is stated for w_hat(X), the distillation solution. For finite-step deep networks it is not established §4.3 §5.3.
• Monotonicity is only correlational. The monotonicity index is binary and is measured, not manipulated; the paper says its effect on risk is unknown §5.3.
• Experiment details are missing. Learning rate, step count, held-out set size, the kappa grid, the teacher for §5.1, and the depth used in §5.2 and §5.3 are not stated. The results are qualitative trends read from plots.
• Not claimed for nonlinear models. The conclusion states that the nonlinear case is a hypothesis and the main direction for future work §6.

Lineage

• Builds on: Hinton, Vinyals and Dean (2014) for distillation itself §1 (no arXiv id in the paper's bibliography, so no card link).
• Deep linear networks as a model: Saxe et al. (2014), Kawaguchi (2016), Hardt and Ma (2017) §3.
• Balancedness and gradient-flow analysis of deep linear networks: Arora, Cohen and Hazan (2018) and Arora, Cohen, Golowich and Hu (2019) §4.1.
• Earlier theory: Urner et al. (2011), who view distillation as a semi-supervised method §2. Other theory on learning from privileged information and teaching sets is discussed in §2.
• Compared against: hard-label learning, where standard bounds need n >> d §4.2 and the paper's strong-monotonicity property does not hold §4.3.
• Superseded by: not known from content on this wiki.