Knowledge Distillation Performs Partial Variance Reduction

Summary

This work explores knowledge distillation, a technique used to improve the performance of smaller "student" models by leveraging the knowledge of more powerful "teacher" models. The authors analyze knowledge distillation from an optimization perspective and reveal that it can be seen as a stochastic variance reduction mechanism, reducing stochastic gradient noise and acting as a partial variance reduction technique. However, complete noise elimination depends on the characteristics of the "teacher" model. The study highlights the importance of careful parameterization, particularly regarding the weighting of the distillation loss. Empirical experiments with linear models and deep neural networks validate the findings, providing further insight into knowledge distillation.

Strengths

This paper took a common empirical practice in the field, formalized the problem and provided analytical explanation from a unique perspective. The theoretical conclusions match with what the community observes in the real life, which provide valuable insights and guidance. The experiments also perfectly corroborated the theoretical results, which makes the conclusions very convincing.

Weaknesses

1. There is no explanation why LBFGS teacher is better than SGD teacher. 2. The optimal choice of distillation weight $\lambda$ does reflect the quality of the teacher, however, the selection itself is rather post hoc, is there any other more practical indicator of the quality of the teacher models that is more informative before running the optimization process.

Questions

N/A

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

would be interesting if there are also results concerning generalizations.

Summary

This paper gives a new interpretation of the logit-based knowledge distillation algorithm. In particular, for simple one-layer models, authors show theoretically that the knowledge distillation can be viewed as diminishing the scale of the student model gradient by some amount proportional to the teacher model gradient. For deeper models, one gets a slightly different quantity, which is empirically well-correlated with the single-layer version. Building on this observation, the authors give the convergence guarantee of the knowledge distillation algorithm, in terms of the parameter and the risk (Theorem 1 and 2, respectively). Authors also propose a slight refinement of the KD, which performs a more proper form of variance reduction.

Strengths

- This paper gives one of the most pleasantly simple yet insightful interpretations of the KD objective. There have been many theoretical studies that attempted to explain why knowledge distillation helps, but I do not think I have seen any explanation as simple as this. As it offers parameter-level insights, I believe that this observation can lead to many algorithmic consequences in future work. - The theoretical analysis also provides some idea about determining the distillation weight $\lambda$, which is not practical in its current form but could give some inspiration nevertheless. - The contributions and the proposed unbiased KD are novel, as far as I know. - The paper is written quite clearly, and the contributions are easy to understand.

Weaknesses

- It would have been better if the paper provides a direct head-to-head comparison of theorems 1 & 2 to the convergence guarantees one could get from the vanilla SGD procedure. The provided "convergence bound" type of results does not really give too much insight into how the knowledge distillation (or variance reduction) leads to a **better** convergence than the vanilla SGD. - Two logical connections could be made a little bit clearer. (1) How does eq(5) lead to standard variance reduction? Are we saying that KD is a variance reduction, only because there are some negative terms? Do we have any direct empirical corroborations of this claim? (2) Please explain how variance reduction leads to better training, e.g., by citing prior work. - In Figure 1, I am not sure whether the cosine similarity is the right metric to use. The magnitude should be a very critical issue, especially because this paper is claiming that KD is about variance reduction. Could you provide additional plots on the $\ell_2$ difference, or maybe the SNR-like result (i.e., the ratio between the difference and the original distillation gradient)? - The empirical results on the proposed unbiased KD are given in the form of "training loss" only, rather than test loss or test accuracy. The results on the test/validation dataset may be very useful in understanding the strengths and limitations of the proposed algorithm. - A mild suggestion on notations---please consider changing the symbol for the student model parameters. In most machine learning papers, $x \in \mathcal{X}$ denotes the input feature (and $y \in \mathcal{Y}$ denotes the label). Many KD papers use $\theta_t$ for teacher parameters and $\theta_s$ for student parameters. Although the current notation is okay, using more common notations may help the readers greatly. - A minor mistake---the example on "classification with single hidden layer" may not really be the case with one "hidden layer." The one-hidden-layer network is the same as a two-layer neural network.

Questions

In addition to the "weaknesses" above, here are some questions/suggestions. - Is there any way one could extend the proposed method of analysis to the distillation cases where the student model is not necessarily a "compressed" version of the teacher model? - It would be great if there is a plot that explicitly tracks the gradient variance.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

3 good

Contribution

3 good

Limitations

The limitations are given in the appendices, which I think is okay.

Summary

This paper examines the benefits of Knowledge Distillation (KD) from an optimization perspective. They show that, under certain assumptions, KD performs partial variance reduction on SGD noise and that the amount of reduction depends on the quality of the teacher model. Their analysis suggests that the distillation weight used in the KD loss should be appropriately tuned, and the authors provide a closed-form solution for the optimal weight in the case of linear models. Even though their core result does not directly apply to deep networks, they present some empirical evidence supporting that it remains a reasonable approximation.

Strengths

* Understanding the underlying mechanics of KD and the reason for its benefit is highly relevant to the ML community given the widespread use of KD. A deeper understanding may also lead to better distillation algorithms, as exemplified with the closed-form optimal distillation weight in this paper. * The connection between the distillation gradient and variance reduction methods like SVRG is interesting and novel to the best of my knowledge. * The presented analyses and results seem sound. * The authors present some empirical evidence that support the claims of the paper (Figs. 2-5).

Weaknesses

* The core proposition of the paper (Prop. 1) does not apply to deep neural networks, and the empirical evidence presented to support the claim that it’s a good approximation is quite limited (one scenario on MNIST with one hidden layer FFN and fixed learning rate). * The presented results, including the empirical ones, apply to training loss, rather than test error. * The presented theory states that a higher performing teacher leads to higher variance reduction. However, prior works in distillation have shown that significantly better teachers often lead to worse students [1-3]. Please see the questions section for specific questions on this. [1] https://arxiv.org/pdf/1902.03393.pdf [2] https://arxiv.org/pdf/2202.03680.pdf [3] https://arxiv.org/pdf/2206.06067.pdf

Questions

1. Continuing from the comment from the Weakness section: how does the presented theory reconcile with the prior observations that better teachers may in fact lead to poorer students due to, e.g., the large capacity gap between the student and the teacher? Does this suggest that variance reduction is not the end of the story or that it is solely responsible for the success of distillation? 2. Why does the cosine similarity between the approximation of Prop 1 and the true gradient in deep networks decrease as the training progresses? It would be helpful to show the behavior until a larger epoch like 100 instead of 50 to support the claim that the behavior stabilizes as training progresses. 3. Beyond empirical evidence, is there an intuitive reason to believe that Prop. 1 would approximately apply to deep neural networks?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

3 good

Contribution

3 good

Limitations

Yes.

Summary

This work analyzes Knowledge Distillation (KD) using an optimization point of view. By recasting the KD problem as a standard learning problem with a custom loss, the authors analyze the convergence of SGD on such loss and identify the variance-reducing properties of KD, through the bias induced by the teacher in the loss. Using some approximations, the authors argue that their analysis should hold for deep networks too. The paper complements its analysis by proposing a technique to reduce the bias of KD and extend the result to generic KD with smaller students. The proposes analysis is supported by experiments for linear models.

Strengths

The paper is clear and easy to read even for non-optimization experts. The whole idea of analyzing KD as a variance-reducing algorithm is novel and interesting Although the assumptions are limited to linear models, the empirical experiments show that such assumptions approximately could hold deep networks is convincing. The paper proposes a good balance of novelty, clarity and significance and it should be relevant to anybody interested in KD.

Weaknesses

The biggest weakness of this work is its ambivalence on the validity of the results for deep learning. While I do understand this is not the goal of the paper, I do not understand why the authors claim that their results could be approximately true for deep networks, show some empirical evidence of this and then, drop the topic for the rest of the paper. I don't think the authors could either add some empirical evidence on this matter or focus on the linear/convex case. I any case, I encourage the authors to further investigate KD as variance-reductions on deep networks, even if only empirically in another paper.

Questions

My doubts are about the most empirical aspects of the paper. The "Experimental validation" in sec 4.3 mentions "linear probing" on CIFAR-10, is this just taking the last features produced by a neural network and training a linear classifier on top of those? I also want to ask to the authors if they tried to replicate experiments with deep networks to show that empirically, the same behaviours can be obtained even for non-linear models. I don't understand why you write a paragraph to claim that the fundamental assumption for your results can hold approximately for deep learning and then don't elaborate on the matter any further, either with experiments or theoretical analysis.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The authors have addressed the limitations of their work properly.