Knowledge Distillation Performs Partial Variance Reduction
Authors: Mher Safaryan
Alexandra Peste
Dan Alistarh
Venue: NeurIPS 2023 poster (2023)
Keywords: knowledge distillation
stochastic optimization
variance reduction
TL;DR: We show that knowledge distillation acts as a form of partial variance reduction, which can reduce the stochastic gradient noise, but may not eliminate it completely, depending on the properties of the “teacher” model.

Abstract: Knowledge distillation is a popular approach for enhancing the performance of "student" models, with lower representational capacity, by taking advantage of more powerful "teacher" models. Despite its apparent simplicity, the underlying mechanics behind knowledge distillation (KD) are not yet fully understood. In this work, we shed new light on the inner workings of this method, by examining it from an optimization perspective. Specifically, we show that, in the context of linear and deep linear models, KD can be interpreted as a novel type of stochastic variance reduction mechanism. We provide a detailed convergence analysis of the resulting dynamics, which hold under standard assumptions for both strongly-convex and non-convex losses, showing that KD acts as a form of \emph{partial variance reduction}, which can reduce the stochastic gradient noise, but may not eliminate it completely, depending on the properties of the ``teacher'' model. Our analysis puts further emphasis on the need for careful parametrization of KD, in particular w.r.t. the weighting of the distillation loss, and is validated empirically on both linear models and deep neural networks.

Knowledge Distillation Performs Partial Variance Reduction
Mher Safaryan
Alexandra Peste
IST Austria
IST Austria
mher. mher. atber.safaryaa@istaaac...
alexardra..esteeiista.cc
Dan Alistarh
IST Austria
dan www.aaaaaiiittaaaccoooom
Abstract
Knowledge distillation is a popular approach for enhancing the performance of "student" models, with lower representational capacity, by taking advantage of more powerful "teacher" models. Despite apparent simplicity and widespread use, the underlying mechanics behind knowledge distillation (KD) are still not fully understood. In this work, we shed new light on the inner workings of this method, by examining it from an optimization perspective. We show that, in the context of linear and deep linear models, KD can interpreted of stochastic variance reduction mechanism. We provide detailed convergence analysis of the resulting dynamics, which hold under standard assumptions for both strongly-conver and non-convex losses, showing that KD acts as a form of partial variance reduction, which can reduce the stochastic gradient noise, but may not eliminate it completely, depending on the properties of the "teacher" model. Our analysis puts further emphasis on the need for careful parametrization of KD, in particular w.r.t. the weighting of the distillation loss, and is validated empirically both linear models and deep neural networks.
novel type
Introduction
Knowledge Distillation (KD) [13, 3] is standard tool for transferring information between machine learning model lower representational sapacity-usually called the student-and more accurate and powerful teacher model. In the context of classification using neural networks, it is common consider the student to be a smaller network [2], whereas the teacher is a network that is larger but also more accurate. Assuming supervised classification task, distillation consists in training the student minimize the cross-entropy with respect to the teacher's 's logits on every given sample, in addition to minimizing the standard cross-entropy loss with respect the ground truth labels.
and more computationally-heaw
Since its introduction [3], distillation has been developed and applied in a wide variety of settings, from obtaining compact high-accuracy encodings of model ensembles [14], to boosting the accuracy of compressed models [50, 39, 32], to reinforcement learning [51, 43, 36, 5, 7, 46] and learning with privileged information [52]. Given its apparent simplicity, there has been significant interest in finding explanations for effectiveness of distillation [2, 14, 38). For instance, one hypothesis [2, 14] is that the smoothed labels resulting from distillation present the student with a decision surface that is easier to learn than the one presented by the categorical (one-hot) outputs. Another hypothesis [2, 14, 52] starts from the observation that the teacher's outputs have higher entropy than the ground truth labels, and therefore, higher information content. Despite this work, we still have a limited analytical understanding regarding why knowledge distillation is so effective [38]. Specifically, very little is
37th Conference on Neural Information Processing Systems (NeurIPS 2023).

---

known about the interplay between distillation and stochastic gradient descent (SGD), which is the standard optimization setting in which this method applied.
1.1. Contributions. In this paper, we investigate the impact of knowledge distillation on the convergence of the "student" model when optimizing via SGD. Our approach starts from simple re-formulation of distillation in the context of gradient-base based optimization, which allows us to connect KD to stochastic variance reduction techniques, such SVRG [16], which are popular in stochastic optimization. Our results apply both to self-distillation, where KD is applied while training relative to an earlier version of the same model, well as distillation for compression, where compressed model leverages outputs from an uncompressed one during training.
In nutshell, in both cases, we show that SGD with distillation preserves the convergence speed of vanilla SGD, but that the teacher's outputs serve to reduce the gradient variance term, proportionally to the distance between the teacher model and the true optimum. Since the teacher model may not be at an optimum, this means that variance reduction is only partial, as distillation may not complctely eliminate noise, and in fact may introduce bias or even increase variance for certain parameter values. Our analysis precisely characterizes this effect, which can be controlled by the weight of the distillation loss, and is validated empirically for both linear models and deep networks.
1.2. Results Overview. To illustrate our results, consider the case of self-distillation [59], which is a popular supervised training technique in which both the student model x E I and the teacher model 0 E Rd have the same structure and dimensionality. The process starts from teacher model 0 trained using regular cross-entropy loss, and then trains the student model with respect to a weighted combination of cross-entropy w.r.t. the data labels, and cross-entropy w.r.t. the teacher's outputs. This process is often executed over multiple iterations, in the sense that the student at iteration becomes the teacher for iteration l: + 1, and so on.
Our first observation is that, in the case of self-distillation with teacher weight 2 À 2 0, the gradient of the distilled student model on a sample i, denoted by VzJ((1) 0, at a given iteration can simply be written
Vafil = -
where (z) and (() are the student's and teacher's standard gradients on sample respectively, This expression is exact for linear models and generalized linear networks, and we provide evidence that it holds approximately for general classifiers. With this re-formulation, self-distillation can be interpreted as a truncated form of the classic SVRG iteration [16], which never employs full (non-stochastic) gradients.
Our main technical contribution is in providing convergence guarantees for iterations of this form, for both strong quasi-convex, and non-convex functions under the Polyak-Lojasiewicz (PL) condition. Our analysis covers both self-distillation (where the teacher is a partially-traines version of the same model), and more general distillation, in which the student is compressed version of the teacher model. The convergence rates we provide are similar in nature to SGD convergence, with one key difference: the rate dependence the gradient variance x² is dampened by term depending on the gap between the teacher model and an optimum. Intuitively, this says that, if the teacher's 's accuracy poor, then distillation will not have any positive effect. However, the better-trained the teacher is, the more it can help reduce the student's 's variance during optimization. Importantly, this effect occurs even if the teacher is not trained to near-zero loss, thus motivating the usefulness of the teacher in self-distillation. Our analysis highlights the importance of the distillation weight, as a means to maximize the positive effects of distillation: for linear models, we can even derive a closed-form solution for the optimal distillation weight. We validate our findings experimentally for both linear models and deep networks, confirming the effects predicted by the analysis.
2 Related Work
We now provide overview for some of the relevant related work regarding KD. Knowledge distillation, in its current formulation, was introduced in the seminal work of [13], which showed that the predictive performance of a model can be improved if it is trained match the soft targets produced by a large and accurate model. This observation has motivated the adoption of KD as a standard mechanism to enhance the training of neural networks in wide range of settings, such as compression [39], learning with noisy labels [26]. and has also become an essential tool in training accurate compressed versions of large models [39, 49, 54].

---

Despite these important practical advantages, providing a thorough theoretical justification for the mechanisms driving the success of KD has so far been elusive. Several works have focused on studying KD from different theoretical perspectives. For example, Lopez et al. [28] connected distillation with privileged information [52] by proposing the notion of generalized distillation. and presented an intuitive explanation for why generalized distillation should allow for higher sample efficiency, relative to regular training. Phuong and Lampert [38] studied distillation from the perspective of generalization bounds in the case of linear and deep linear models. They identify three factors which influence the success of KD: (I) the geometry of the data; (2) the fact that the expected risk of the student always decreases with more data; (3) the fact that gradient descent finds a favorable minimum of the distillation objective. By contrast to all these previous references, our work studies the impact of distillation on stochastic (SGD-based) optimization.
More broadly, there has been extensive work on connecting KD with other areas in learning. Dao et al. [6] examined links between KD and semi-parametric Inference, whereas Li et al. [61] performed an empirical study on KD in the context of learning with noisy labels. Yuan et al. [58] and Sultan et al. [48] investigated the relationships between KD and label smoothing. a popular heuristic for training neural networks, showing both similarities and substantive differences. In this context, our work the first to signal the connection between KD and variance-reduction, as well as investigating the convergence of KD in the context of stochastic optimization.
3 Knowledge Distillation
1,2,..., NJ. where inputs On E A (c.g., vectors from Id and outputs b, E B (c.g., categorical labels or real numbers). Consider a set of models F = { A + B E P _ Rd with fixed neural network architecture parameterized by vector x. Depending on the supervised learning task and the class F of models, we define a loss function é B B + R4 in order to measure the performance of the model. In particular, the loss associated with a data point (an, by) and model Or E J would be (()x(an), bw)) In this framework, the standard Empirical Risk Minimization (ERM) takes the following form:
3.1. Background. Assume we are given finite dataset {(an; bn)
min aERd
E ((po(an);
In the objective above, the model oz is trained to match the true outputs b, given in the training dataset. Suppose that in addition to the true labels b,, we have access to sufficiently well-trained and perhaps more complicated teacher model's 's outputs (a ) E B for each input on E AA. Similar to the student model 6u. the teacher model & maps A - B but can have different architecture, more layers and parameters. The fundamental question is how to exploit the additional knowledge of the teacher te to facilitate the training of a more compact student model Φz with lower representational capacity. Knowledge Distillation with parameter À E [0, 1] from teacher model 49 to student model 6₂ is the following modification to the objective (1):
min ZERM N
-m=1 (1 ())() bw.) + Xe(o)( M())a(().),
(o(a.)))
(2)
Here we customize the loss penalizing dissimilarities from the teacher's feedback ( ) in addition to the true outputs 6n In case of is linear in the second argument (e.g., cross-entropy loss), the problem simplifies into
min 1 EN= ((ox(an), (1 - A)b,, + (a PERd N
(3)
which is a standard ERM (I) with modified "soft" labels Sn = (1- A)bn + (an)))) ) as the target.
3.2. Seff-distillation. As already mentioned, the teacher's model dp can have more complicated neural network architecture and potentially larger parameter space 0 E Q C RDD In particular, & does not have to be from the same set of models F as the student model Φz. The special case when both the student and the teacher share the same structure/architecture called self-distillation
[33, 60). which is the key setup for our work. In this case, the teacher model & = 6o E F with
E T (i.e., Q = P, D d) and the corresponding distillation objective would be
#ERd N + (1 (an)) - ((()x(aa)) (a())
min 1
(4)
3

---

Algorithm 1 Knowledge Distillation via SGD
1: Input: learning rate y > O, initial student model xd P S R
2: for each distillation iteration m do
3: choose teacher model TTh E S 12 and distillation weight A [0, (e.g., see Sec. 5)
4: for each training iteration t do
5: sample an unbiased mini-batch { ~ DD form the train set
6: compute distillation gradient V (x' 0 { = V + (1- XT V."''" 0
7: update the student model via x'+1 = x² Veee(z Am)
8: end for
9: end for
Our primary focus in this work would be the objective mentioned above of self-distillation. For convenience, let (z) == ((0x(an) bx)) be the prediction loss with respect to the output b, fn (x 8) == (an) 00 (an)) be the loss with respect the teacher's output probabilities and fn A) = (z) + () () be the distillation loss. See Algorithm for an illustration.
3.3. Distillation Gradient. As the first step towards understanding how self-distillation affects the training procedure, we analyze the modified loss landscape (4) via stochastic gradients of (1) and (4). In particular, we put forward the following proposition regarding the form of distillation gradient in terms of gradients of (1).
Proposition (Distillation Gradient). For a student model x Rdd teacher model 0 Rd and distillation weight X, the distillation gradient corresponding to self-distillation (4) is given by
V=Sa(o 0, 1) Vfa(z) AV
Before justifying this proposition formally, let us provide some intuition behind the expression (5) and its connection to distillation loss (4). First, the gradient expression (5) suggests that the teacher has little no effect on the data points classified correctly with high confidence (i.e. those for which V (o, is close to 0 or ( (a.) is close to b,)) In other words, the more accurate the teacher is, the less it can affect the learning process. In the extreme case, a perfect or overfitted teacher (one that (O) = 0 or (a))n) = b, for all n) will have no effect. In fact, this is expected since, in this case, problems (1) and (4) coincide. Alternatively, if the teacher is perfect, then the modified objective (4) intuitively suggests that the learning dynamics of a student model is adjusted based on the teacher's knowledge. As we can see in (5), the adjustment from the teacher is enforced by -XV fn (f) term. is worth mentioning that the direction of distillation gradient fn 0, A) can be different from the usual gradient's s direction V fz (z:) due the influence of the teacher. Thus, Proposition explicitly shows how the teacher guides the student by adjusting its stochastic gradient.
As we will show later, distillation gradient (5) leads to partial variance reduction because of the additional (a (0) term. When chosen properly (distillation weight À and proximity optimal solution x²)) this additional stochastic gradient is capable of adjusting the student's stochastic gradient since both are computed using the same batch from the train data. In other words, both gradients have the same source of randomness which makes partial cancellations feasible.
to the
Linear regression. To support Proposition 1 rigorously, consider the simple setup of linear regression. Let A = Rd, P = Rd+1 Φx (a) = x a E RR where a = (a T] E RRd+1 is the input vector in the lifted space (to include the bias term), B R. and the loss defined by E(t, t') (t-
for all t, t' E R. Thus, based on (4), we have
fn( A,
(1 1)- T an - bn,) Mix on - -8 aTn))²
from which we compute its gradient straightforwardly as
Vafa(( (, = 2(1 - xx an bn) + 22(x T ān - 00 ān
= 2(x an by,) - 22(0 an - bn) = fz) - AV (Q).
Hence, the distillation gradient for linear regression tasks has the form (5).
Classification with a single hidden layer. We can extend the above argument and derivation for a K-class classification model with one hidden layer that has soft-max as the last layer, i.. Φ (an) o(X an) K where X [2] D2 ... ØK] dd are the model parameters,

---

1-1.1
38
Cosine Similarity
3.32
1-13
2.08
23
0.22
J+B1
J-033
3.M
1-0.7
3.00
91
26
66
226
20
42
100 20
60
120
Epoch
Epach
Epoch
Figure 1: Cosine similarity, l2 distance and SNR (i.e., 27 distance over the gradient norm of standard KD) statistics between the true and approximated distillation gradient for a neural network during training. As predicted, larger À leads to larger differences, although gradients remain well-correlaed.
an E A = Rdd is the input data and a is the soft-max function. Then, we show in the Appendix B.2 that for all = 1.2.... K holds
V ae fx = V
where θ [01 02 0k] e Raxk are the teacher's parameters.
Generic classification. Proposition 1 will not hold precisely for arbitrary deep non-lincat neural networks. However, careful calculations reveal that, in general, distillation gradient takes form similar to (5). Detailed derivations are deferred to Appendix B.3, here we provide the sketch.
Consider an arbitrary neural network architecture for classification that ends with soft-max layer, i.e. (a.) (x)) where an AA is the input data, th, (x) are the produced logits with respect to the model parameters x, and a is the soft-max function. Denote ypn (=) == ((a(z), ) b,) the loss associated with logits = and the true label bn In words, th,, gives the logits from the input data, while &Pn computes the loss from given logits. Then, the representation for the loss function is (x) (P((V(()))) We show in Appendix B.3 that the distillation gradient can be written as
Vafn(1 0) ())n(z) (VPP AVipn(Us (v))
(.).(a)) (z) Ox (en(2)
(V. (5,
de
(0)
where (v) (v)(w) I Venn2(0).... ... (x), Raxk is the Jacobian of the
vector-valued function vhx for logits. Notice that the first term (v)(2) (5)(w)
on (v..(a) coincides with the
student's gradient (x) (w) 8r However, the second term (x),.2) (f)
0%%, (0) differs from the
teacher's gradient as the partial derivatives of logits are with respect to the student model.
Despite these differences in the case of deep non-linear models, we observe that the distillation gradient defined by Equation 5 can approximate well the true distillation gradient from Equation 3. Specifically, we consider fully connected neural network with one hidden layer and ReLU activa- tion [34]. trained on the MNIST dataset [24], using regular self-distillation, from SGD-teacher, with fixed learning rate, and SGD with weight decay and no momentum. At each training iteration we compute the cosine similarity between the gradient of the distillation loss and the approximation from Equation 5, and we average the results across each epoch. The results presented Figure show that the distillation gradient approximates well the true distillation gradient. Moreover, the behavior is monotonic the distillation weight À (higher similarity for smaller A), as predicted by the analysis above, and it stabilizes as training progresses. The decrease of cosine similarity can be explained as follows: at the beginning the cosine similarity is high (and SNR is low) since we start from the same model. Then, initial perturbations caused by either the KD or modified KD gradient don't cause big shifts (the teacher has enough confidence and small gradients). These perturbations accumulate over the training leading to decreased cosine similarity and eventually stabilize.
4 Convergence Theory for Self-Distillation
4.1. Optimization Setup and Assumptions. We abstract the standard ERM problem (1) into stochastic optimization problem of the form
{ f(x) = EE-D l
min zERd
(6)
5

---

where fa( is the loss associated with data sample & ~ D given model parameters E R. For instance, E = (an, bn) a single data point, then the corresponding loss is fe(z) f (on),.)). b,). The goal is to find parameters I minimizing the risk fzr). To solve the problem (6), we employ Stochastic Gradient Descent (SGD). Based on Section 3, applying SGD the problem (6) with self-distillation amounts to the following optimization updates in the parameter space:
1+ = x' > (V - (e)).
with initialization z0 E IR step size learning rate O, teacher model's parameters E R and distillation weight X. To analyze the convergence behavior of iterates (7), we need to impose some assumptions in order to derive reasonable convergence guarantees. First, we assume that the problem (6) has a non-empty solution set X o 0 and f = f(x") for some minimizer x" E EXX
Assumption (Strong quasi-convexity.. The function f: R + R is differentiable and p-strongly quasi-convex some constant μ > 0, i.e., for any E2 it holds
f(x)) f (x) + (VI()), 2)
M 2 1 x1|2
(8)
Strong quasi-convexity [9] is weaker version of strong convexity [35], which assumes that the quadratic lower bound above holds for at every point y E Id instead of x E X. Notice that strong quasi-convexity implies that the minimizer x" is unique'. A more relaxed version of this assumption (PL) condition [40].
is the PPolykk-Lojaiewvin condition). Function f: F.R.d - R is differentiable and satisfies
Assumption
PL condition with parameter > 0, if for any E R2 it holds
|V
2u(f(x)-f')
(9)
Note that the requirement imposed by the PL condition above weaker than by strong convexity. Functions satisfying PL. condition do not have to be convex and can have multiple minimizers [17]. We make use of the following form of smoothness assumption on the stochastic gradient commonly referred to as expected smoothness in the optimization literature [11, 10, 19].
Assumption 3 (Expected Smoothness). Functions fę(x) are differentiable and L-smooth in expecta- tion with respect to subsampling & ~ D, i.e., for any x E IRd it holds
READ [1] [[) - VJe(2 )13} < 2C([(x)- - 1r)
(10)
for some constant L C C(S,).
The expected smoothness condition above is a joint property of loss function ff and data subsampling strategy from the distribution D. In particular, it subsumes the smoothness condition for f(x since (10) also implies [(x) V[(=")]1 < 2C(((x) f*)) for any x E Rd. We denote by I. the smoothness constant of (a:) and notice that L.S.
4.2. Convergence Theory and Partial Variance Reduction. Equipped with the assumptions described in the previous part, we now present our convergence guarantees for the iterates (7) for both strong quasí-convex and PL loss functions.
1 and properly Theorem 1 (See Appendix C.2). Let Assumptions and 3 hold. For any 7 S sC.
chosen distillation weight X, the iterates (7) of SGD with self-distillation using teacher's 's parameters 0 converge as
E [1 - x x01 S (1
x=1 20 min(y: ((f()) - )),,,
(11)
where o3 = Bl (x 12] is the stochastic noise the optimum.
Theorem 2 (See Appendix C.3). Let Assumptions and 3 hold. For any 7 S - t and properly chosen distillation weight X, the iterates (7) of SGD with self-distillation using teacher's parameters O converge as
- (1 (1- (f(x=)- + Loi Iff f(x") fx**). then $1100 "0*1| f(x') f(x*)) 0. Hence, x² x"'
{*)),
(12)

---

Proof overview. Both proofs follow similar steps and can be divided into three logical parts.
Part 1 (Descent inequality). Generally, integral part of essentially any convergence theory for optimization methods is a descent incquality quantifying the progress of an algorithm in one iteration. Our theory is not an exception: we first define our "potential" e' = =Iæz* quasi-convex setup, and e' = f(x*) f for the PL setup. Then, we start our derivations by bounding E (1 - yu)e'. Here, E, is the conditional expectation with respect the randomness of previous iterate x². Specifically, up to constants, both setups allow the following bound:
x=1 f for the strongly
E [e S (1 7u)e (()) - (())))((()))(()))))))))) -() (())
(13)
Choosing the learning rate y to be small enough, we ensure that the second term is non-positive and hence negligible in the upper bound.
where ((A) = I + B
Part (Optimal distillation weight). Next, we focus our attention on the third term in (13) involving
the iteration- intrrendenddee neighborhood term N (A). Note that the (()) factor next to (A) will be
absorbed once we unfold the recursion (13) up to initial iterate. Hence, the convergence neighborhood
is proportional to (A). Now the question is how small this term can get if we properly tune
the parameter A. Notice that N (() 7 corresponds the neighborhood size for plain SGD without any distillation involved. Luckily, due to the quadratic dependence, we can minimize N analytically with respect to À and find the optimal value
([VT)("") B|I ((0))) 31V(0)
A
(14)
Consequently, the analysis puts further emphasis on the need for careful parametrization with respect to the weighting À of the distillation loss as there exists a particularly privileged value X,
Part 3 (Impact of the teacher). In the final step, we quantify the impact of the teacher the reduction in the neighborhood term (Aa) compared to the plain SGD neighborhood N (0). Via algebraic transformations, we show that
2 O(
1²(x", 0) 1- 2(0) (-300)
(15)
where ((0) [VV E [o, 1] is the signal-to-noise ratio, and p(z",0) E 1-1,
is the correlation coefficient between stochastic gradients Vfe(2*) and ((() This representation gives analytical means to measure the impact of the teacher. For instance, the optimal teacher 0 = xx satisfics p(xx, 8) = 1 and ((0) = O. and thus ((A)) = 0. In general, if the teacher is not optimal, then the reduction (15) is of order 7 (f - f*) and N(A.) = o min(7,0(f(0) - -}))
We discuss these results highlighting the key aspects and significance.
Structure of the rates. The structure of these two convergence rates is typical in gradient- based stochastic optimization literature: linear convergence up to some neighborhood controlled by the stochastic noise term of and learning rate y. In fact, these results (including learning rate restrictions) are identical to ones for SGD [8] except the non- vanishing terms in (11) and (12) include an additional O(f(f) f*) factor due to distillation and proper selection of weight parameter X. For both setups, the rate of SGD is the same (11) or (12) with only one difference: min(y, (f f*)) term is replaced with y. So, ((f(0) f*) is the factor that makes our results better compared to SGD in terms of optimization performance.
Importance of the results. First, observe that in the worst case when the teacher's parameters are trained inadequately (or not trained at all), that is Off(0) - 2 % then the obtained rates recover the known results for plain SGD. However, the crucial benefit of these results is to show that a sufficiently well-trained teacher, i.c. O(f(0) f+) SGD without slowing down the speed of convergence. In the best case scenario, when the teacher's model is perfectly trained, namely f(0) = f then the neighborhood term vanishes, and the method converges to the exact solution (see SGD-star Algorithm 4 in [9]). Thus, self-distillation in the form of iterates (7) acts as form of partial variance reduction, which can reduce the stochastic gradient noise, but may not eliminate it completely, depending on the properties of the teacher model.
7 provably reduces the neighborhood size of
Choice of distillation weight X. As we discussed in the proof sketch above, our analysis reveals that the performance of distillation is optimized for a specific value (14) of distillation weight X, depending on the teacher model. One way to interpret the expression (14) for weight parameter
7

---

0.00
Train Stateetee
0:20
SGD teacher
teacher
0.2
0.6
0.0
0.4
0.0 0.0
02
5.5 0.8
0.03
5.
0.3 04 0.5 Lambda
0.7 0.0 100
Lambda
(a) Self-distillation on MNIST.
(b) Self-distillation on CIFAR-10.
The minimum training loss, and average the last 10 epochs, for models trained with SGD and with self-disillation, using different values the distillation parameter A, on MNIST and CIFAR-10 10. SGD is equivalent to À O. The curves for the SGD- teachers (which do not have zero loss) reflect our analysis, corroborating the existence of the "optimal" distillation weight. By contrast, in the L-BFGS teacher, higher distillation weight always leads to lower loss.
Figure 2:
intuitively is that the better the teacher's model 0 is, the bigger X S gets. In other words, X. quantifies the quality of the teacher: X, indicates poor teacher model (f(0) > (*) and X. for the optimal teacher (f(0)= f*).
4.3. Experimental Validation. In this section we illustrate that our theoretical analysis in the convex case also holds empirically. Specifically, consider classification problems using linear models in two different setups: training a linear model on the MNIST dataset [24] and linear probing on the CIFAR-11 10 dataset [23], using ResNet50 model [12]. pre-trained the ImageNet dataset [42]. For the second setup we train a linear classifier on top of the features extracted from a ResNet50 model pre-trained ImageNet. This is standard setting, commonly used in the transfer learning literature, see e.g. (21, 45, 15]. In both cases we train using SGD without momentum and regularization, with a fixed learning rate and mini-batch of size 10, for total of 100 epochs. The models trained with SGD are compared against self-distillation (Equation 7), using the same training hyper-pone parameters, where also consider the
the teacher is the model trained with SGD. In the case of CIFAR- 10 features, "optimal" teacher, which is a model trained with L-BFGS [27]. We perform all experiments using multiple values of the distillation parameter À and measure the cross entropy loss between student and true labels. At each training epoch we computing the running average over all mini-batch losses seen during that epoch.
The results presented in Figure 2 show the minimum cross entropy train loss obtained over 100 epochs, as well as the average over the last 10 epochs, for models trained with SGD, as well as with self-distillation, with À [0.01, 1]. We observe that when the teacher the model trained with SGD (A = 0), there exists À > O which achieves a lower training loss than SGD, which is in line with our statement from Theorem 1. Furthermore, when the teacher is very close the optimum, À closer to reduces the training loss the most compared to SGD, which is also in line with the theory (see Theorem 1). This behavior is illustrated in Figure 2b, when using an L-BFGS teacher.
5 Removing Bias and Improving Variance Reduction
In this section, we investigate the cause of having variance reduction only partially and suggest possible workaround obtain complete variance reduction. In brief, the potential source partial variance reduction is the biased nature of distillation. Essentially, distillation bias is reflected in the
iterates (7) since the expected update direction E [ [e(((²)) AV (0) x²) Vf(x²) - XV (0)
can be different from V)(a=). This comes from the fact that distillation loss (4) modifies the initial
loss (1) composed true outputs. To make our argument compelling, next we correct the bias by
adding AVf(O) to iterates (7) and analyze the following dynamics:
++1 x² (V fe +V
(16)
Besides making the estimate unbiased, the advantage of this adjustment is that no tuning is required for the distillation weight X; we may simply set 1. The obvious disadvantage is that ff(@) is the batch gradient over the whole train data that can be very costly to compute. However, we could compute it once and reuse it for all further iterates. The resulting iteration is similar to the popular and well-studied SVRG [16, 57, 41, 20, 25, 22, 29] method, and therefore iterates (16) will enjoy full variance reduction.
8

---

Thaiing
E3353
Grudfenn
Epoch
Epech
Figure 3: (Lefft plot) The train loss of self-distillation, unbiased self-distillation and vanilla SGD training. (Right plot) The progress of gradient variances (averaged over the iterations within each epoch) for the same setup
Theorem 3 (See Appendix D). Let Assumptions 2 and 3 hold. Then for any 7 5 # (16) with X converge as
3LE the iterates
-alf(=)-ff1e - (10 -f) + 31LL+2)
(f(0)) ff)
(17)
The key improvement that bias correction brings in (17) is the convergence up to neighborhood (()(J(0) - (*)) in contrast to min(r, ((S(0) - -)) as (11) and (12). The multiplicative dependence of learning rate and the quality of the teacher leads the method (16) to full variance reduction. Indeed, if we choose the teacher model as = 2., then the rate (17) becomes
E [f(x²) - sr1 s [(1 - (1-> +y 3L(L ()(x) - "*) (f(x0)) - f³),
provided sufficiently small step-size 7 S 12ZC and enough training iterations = (1/z)) Hence, following SVRG and updating the teacher model in every ((+/wx) training iterations, that is choosing 8t n! as the teacher at the m distillation iteration (see line of Algorithm 1), we have
[f(x"") - [ S ((((x(mm - (*) S 2m (f(x²) - f*).
Thus, we need Ollog !) bias-corrected distillation iteration phases, each with = ((/rn) training iterates, get accuracy in function value. Overall, this amounts to O 2111 log }) iterations of (16).
7ps
Experimental Validation. Similarly to the previous section, we further validate empirically the result from Theorem 3. Specifically, we consider the convex setup described before, where we train linear models on features extracted on the CIFAR-10 dataset. Based on Figure 2b, we select À = 0.4 achieving the largest reduction in train loss, compared SGD, and we additionally perform unbiased self-distillation (Equation 16), using the same training hyperparameters. Similar to the setup from Figure 2, we measure the cross entropy train loss of the student and with the labels, which is computed at each epoch by averaging the mini-batch losses. The results are averaged over three runs and presented in Figure 3. The first plot on the left shows that, indeed, the unbiased self-distillation update further reduces the training loss, compared to the update from Equation 7, The second plot explicitly tracks gradient variance (averaged over the iterations within each epoch) for the same setup. As expected, both variants of KD (biased and unbiased) have reduced gradient variance compared plain SGD. The plot also highlights that both variants KD have similar variance reduction properties, while the unbiasedness of unbiased KD amplifies the reduction of train loss.
6 Convergence for Distillation of Compressed Models
So far, the theory we have presented is for oiff-ditillatton, i.e the teacher's and student's 's architectures are identical, To understand the impact of knowledge distillation, we relax this requirement and allow the student's model be sub-network of the larger and more powerful teacher's model. Our approach to model this relationship between the student and the teacher is to view the student as masked or, in general, compressed version of the teacher. Hence, as an extension (7) we analyze the following dynamics of distillation with compressed iterates:
++ = C(x² - (V (V))
(18)
where student's parameters are additionally compressed each itcration using an unbiased compres- sion operator defined below.

---

Assumption 4. The compression operator C: B + R is unbiased and there exists finite w 2 0 bounding the compression variance variance, i.e., for all x E IR we have
B[C(+)] - ElIC(a) - =1 S w ||:l|
(19)
Typical examples of compression operators satisfying conditions (19) are sparsification [55, 47] and quantization [1. 56], which are heavily used in the context of communication efficient distributed optimization and federated learning [30, 44, 37, 53]. In this context, we obtain the following:
Theorem 4 (See Appendix E). Let smoothness A Assumption hold and I be p--ttrongyy convex. Choose any Y S 162 and compression operator with variance parameler properly selecting distillation weight X, the iterates (18) satisfy
(u//c). Then,
E |100 1oof-=11'
[²] (0w++1) [11 (000-1))110 x²1 of -1 min (7. ((1(0) r1]]
Clearly, there are several factors influencing the speed of the rate and the neighborhood of the convergence that require some discussion. First of all, choosing the identity map compression operator (C(x) = x for all x E Rd), we recover the same rate (11) as before (w = in this case). Next, consider the case when the stochastic noise at the optimum vanishes (o² 0) and distillation is switched off (A = 0) in (18). In this case, the convergence is still up to some neighborhood proportional to ||=*1 since compression is applied to the iterates. Intuitively, the neighborhood term
O(I|==|| I|²) corresponds to the compression noise at the optimum x" ((19) when x x²") Also note that the presence of this non- non-vanishing term O[|æ=" and the variance restriction (+/cc) is consistent with the prior work [18].
So, the convergence neighborhood of iterates (18) has two terms, one from each source of randomness: compression noise/variance O[||=* )I²) at the optimum and stochastic noise/variance O(a)) at the optimum. Therefore, this case as well, distillation with a properly chosen weight parameter (partially) reduces the stochastic variance of sub-sampling
7 Discussion and Future Work
Our work has provided a new interpretation of knowledge distillation, examining this mechanism for the first time from the point of view of optimization. Specifically, we have shown that knowledge distillation acts as a form of partial variance reduction, whose strength depends on the characteristics of the teacher model. This finding holds across several variants of distillation, such as self-distillation and distillation of compressed models, as well as across various families of objective functions.
Prior observations showed that significant capacity gap between the student and the teacher may in fact lead to poorer distillation performance [31]. To reconcile the issue of large capacity gap our results, notice that, in our case "better teacher" means better parameter (i.e.. weights and biases) values, evaluated in terms of training loss. In particular, in the case of self-distillation, covered in Sections and 5, the teacher and student architectures are identical, and hence they have the same capacity. In our second regime, distillation for compressed models (Section 6), we actually consider the case when the student network is subnetwork of the teacher, we consider a sparsification compression operator that selects i parameters for the student out of d parameters of the teacher. Then, clearly. the teacher has larger capacity with capacity ratio d/k 2 1. However, our result in this direction (Theorem 4) does not allow the capacity ratio to be arbitrarily large. Indeed, the constraint w = O(C/H) on compression variance implies a constraint on capacity ratio since w d/k 1 for the sparsification operator. Thus, our result holds when the teacher's size is not significantly larger than the student's size, which is in line with the prior observations on large capacity gap.
As we mentioned, our Proposition 1 does not hold precisely for arbitrary deep non-linear neural networks. However, we showed that this simple model (5) of distillation gradient approximates the true distillation gradient reasonably well both empirically (see Figure 1) and analytically (see Appendix B.3). There is much more to investigate for the case of non-convex deep networks where exact tracking of teacher's 's impact across multiple layers of non-linearities becomes harder. We see our results as a promising first step towards a more complete understanding of the effectiveness of distillation. One interesting direction of future work would be to construct more complex models for distillation gradient and to investigate further connections with more complex suriance-reduction methods, c.g. [4], which may yield even better-performing variants of KD.
10

---

Acknowledgements
MS has received fundiig from the European Union's Horizon 2020 research and innovation pro- gramme under the Marie Sklodowska-Curie grant agreement No 101034413.
References
[1] Dan Alistarh, Demjan Grubic, Jerry Li, Ryota Tomioka, and Milan Vojnovic. QSGD: Communicationn-eiicciennt SGD via gradient quantization and encoding. In Advances Neural Information Processing Systems 30, pages 1709-1720, 2017.
[2] Lei Jimmy Ba and Rich Caruana. Do deep nets really need to be deep? CoRR, abs/ 1312.618488 2013.
[3] C. Buciluă, R. Caruana, and A. Niculescu-Mizie Model compression. Proceedings of the 12th ACM SIGKDD International Conference Knowledge Discovery and Data Mining, KDD '06, pages 535-541, New York, NY, USA, 2006.
[4] Ashok Cutkosky and Francesco Orabona. Momentum baseed variance reduction in non-convex SGD. Advances in Neural Information Processing Systems, 32, 2019.
[5] Wojeiech Czarnecki, Siddhant Jayakumar, Max Jaderberg. Leonard Hasenclever, Yee Whye Teh, Nicolas Heess, Simon Osindero, and Razvan Pascanu. Mix & match agent curricula for reinforcement learning. In Proceedings of the 35th International Conference on Machine Learning, 2018.
[6] Tri Dao, Govinda M Kamath, Vasilis Syrgkanis, and Lester Mackey. Knowledge Distillation as Semiparametric Inference. International Conference on Learning Representations (ICLR), 2021.
[7] Alexandre Galashov, Siddhant Jayakumar, Leonard Hasenclever, Dhruva Tirumala, Jonathan Schwarz, Guillaume Desjardins, Wojtek M. Czarnecki, Yee Whye Teh, Razvan Pascanu, and Nicolas Heess. Information asymmetry in KL-regularized RL. In International Conference on Learning Representatiinss 2019.
[8] Guillaume Garrigos and Robert M Gower. Handbook of convergence theorems for (stochastic) gradient methods. arXiv preprint arXiv:2301.11235, 2023.
[9] Eduard Gorbunov, Filip Hanzely, and Peter Richtárik. Unified Theory of SGD: Variance Reduction, Sampling, Quantization and Coordinate Descent. 23rdInternational Conference Artificial Intelligence and Statistics (AISTATS), 2020.
[10] R.M. Gower, P. Richtárik, and F. Bach. Stochastic via Jacobian sketching. Math. Program. 188, pages 135-192, 2021.
methods: variance reduction
[11] Robert Mansel Gower, Nicolas Loizou, Xun Qian, Alibek Sailanbayev, Egor Shulgin, and Peter Richtárik. SGD: General Analysis and Improved Rates. Proceedings of the 36th International Conference on Machine Learning, 2019.
[12] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. Conference on Computer Vision and Patlern Recognition, 2016.
[13] G. Hinton, O. Vinyals, and J. Dean. Distilling the knowledge neural network. Deep Learning Workshop al NIPS, 2014.
[14] G. Hinton, O. Vinyals, and J. Dean. Distilling the Knowledge in Neural Network. ArXiv e-prints, March 2015.
[15] Eugenia Iofinova, Alexandra Peste, Mark Kurtz, and Dan Alistarh. How Well Do Sparse Imagenet Models Transfer? Conference on Computer Vision and Pattern Recognition, 2022.
[16] Ric Johnson and Tong Zhang. Accelerating stochastic gradient descent using predictive variance reduction. Advances in Neural Information Processing Systems, 26, 2013.
11

---

[17] Hamed Karimi, Julie Nutini, and Mark Schmidt. Linear Convergence of Gradient and Proximal- Gradient Methods Under the
Condition. arXiv:1608.0466: 2020.
[18] Ahmed Khaled and Peter Richtárik. Gradient descent with compressed iterates. arXiv preprint WMXiv:1909.0777 2019.
[19] Ahmed Khaled and Peter Richtárik. Better Theory for SGD in the Nonconvex World. Transac- tions on Machine Learning Research, 2023.
[20] Jakub Koněcný and Peter Richtárik. Semi-Stochastic Gradient Descent Methods. Frontiers in Applied Mathematics and Statistics 3:9, 2017.
[21] Simon Korblitth, Jonathon Shlens, and Quoe V. Le. Do Better ImageNet Models Transfer Better? Conference on Computer Vision and Pattern Recognition, 2019.
[22] Dmitry Kovalev, Samuel Horváth, and Peter Richtárik. Don't t jump through hoops and re- move those loops: SVRG and Katyusha are better without the outer loop. 31 International Conference on Learning Theory (ALT), 2020.
[23] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. CiteSeer, 2009.
[24] Yann LeCun and Corinna Cortes. 2010.
MNIST
handwritten
digit database.
[25]
Lihua Lei, Cheng Ju, Jianbo Chen, and Michael I. Jordan. Non-Convex Finite-Sum Optimization Via SCSG Methods. 31st Conference on Neural Information Processing Systems (NIPS), 2017.
[26] Yuncheng Li, Jianchao Yang, Yale Song, Liangliang Cao, Jiebo Luo, and Li-Jia Li. Learning from noisy labels with distillation. In 2017 IEEE International Conference on Computer Vision (ICCV), pages 1928-1933 IEEE, 2017.
[27] Dong C Liu and Jorge Nocedal. On the limited memory bfgs method for large scale optimization. Mathematical programming, 45(1)500--528, 1989.
[28] David Lopez-Paz, Léon Bottou, Bernhard Schölkopf, and Vladimir Vapnik. Unifying distillation and privileged information. arXiv preprint arXiv:1511.03643, 2015.
[29] Grigory Malinovsky, Alibek Sailanbayev, and Peter Richtárik. Random reshuffling with variance reduction: new analysis and better rates. 39th Conference on Uncertainty in Artificial Intelligence (UAI), 2023.
[30] Ilia Markov, Adrian Vladu, Qi Guo, and Dan Alistarh. Quantized Distributed Training of Large Models with Convergence Guarantees. arXiv preprint arXiv:2302.023390 2023.
[31] Seyed Iman Mirzadeh, Mehrdad Farajtabar, Ang Li, Nir Levine, Akihiro Matsukawa, and Hassan Ghasemzadeh. Improved knowledge distillation via teacher assistant. Proceedings of the AAAI conference artificial intelligence, volume 34, pages 519/-5198, 2020.
[32] Asit Mishra and Debbie Marr. Apprentice: Using knowledge distillation techniques to improve low-precisio precisiin network accuracy. arXiv preprint arXiv:1711.0588 2017.
[33] Hossein Mobahi, Mehrdad Farajtabar, and Peter L. Bartlett. Delf-Distillation Amplifies Regular- ization in Hilbert Space. 34th Conference on Neural Information Processing Systems (NeurIPS), 2020.
[34] Vinod Nair and Geoffrey E Hinton. Rectified linear units improve restricted boltzmann machines. In International Conference on Machine Learning, 2010.
[35] Yurii Nesterov. Introductory lectures on convex optimization: A basic course, volume 87.
Springer Science & Business Media, 2013.
[36] Emilio Parisotto, Jimmy Lei Ba, and Ruslan Salakhutdinov. Actor-mimic: Deep multitask and transfer reinforcement learning. In International Conference on Learning Representations, 2016.
12

---

[37] Constantin Philippenko and Aymeric Dieuleveut. Preserved central model for faster bidirectional compression in distributed settings. 35th Advances in Neural Information Processing Systems, 2021.
[38] Mary Phuong and Christoph Lampert. Towards understanding knowledge distillation. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 5142-5151, Long Beach, California, USA, 09-15 Jun 2019. PMLR.
quantization. International Conference on Learning Representations (ICLR), 2018. [39] Antonio Polino, Razvan Pascanu, and Dan Alistarh. Model compression via distillation and
[40] B. T. Polyak. Gradient methods for minimizing functionals (in Russian). Zh. Vychisl. Mat. Mat. Fiz. pages 643-653 1963.
[41] Sashank J. Reddi, Ahmed Hefny, Suvrit Sra, Barnabas Poczos, and Alex Smola. Stochastic Variance Reduction for Nonconvex Optimization. Proceedings of the 33rd International Conference on Machine Learning, 2016.
[42] Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, et al. ImageNet large scale visual recognition challenge. ICV, 115(3)2211-2555 2015.
[43] Andrei A Rusu, Sergio Gomez Colmenarejo, Caglar Gulcehre, Guillaume Desjardins, James Kirkpatrick, Razvan Pascanu, Volodymyr Mnih, Koray Kavukcuoglu, and Raia Hadsell. Policy distillation. International Conference on Learning Representations, 2016.
[44] Mher Safaryan, Filip Hanzely, and Peter Richtárik. Smoothness Matrices Beat Smoothness Constants: Better Communication Compression Techniques for Distributed Optimization. In 35th Conference on Neural Information Processing Systems, 2021.
[45] Hadi Salman, Andrew Ilyss, Logan Engstrom, Ashish Kapoor, and Aleksander Madry. Do Adversarially Robust ImageNet Models Transfer Better? Advances in Neural Information Processing Systems, 2020.
[46] Simon Schmitt, Jonathan Hudson, Augustin Zidek, Simon Osindero, Carl Doersch, Woj- ciech M Czarnecki, Joel Z Leibo, Heinrich Kuttler, Andrew Zisserman, Karen Simonyan, et al. Kickstarting deep reinforcement learning. arXiv preprint arXiv:1803. 038333 2018.
[47] Sebastian U Stich, Jean-Baptiste Cordonnier, and Martin Jaggi. Sparsified SGD with memory. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems 31, pages 1452-4463 Curran Associates, Inc., 2018
[48] Md Arafat Sultan, Knowledge Distillation x Label Smoothing: Fact or Fallacy? arXiv preprint arXiv:2301.12609 2023.
[49] Zhiqing Sun, Hongkun Yu, Xiaodan Song, Renjie Liu, Yiming Yang, and Denny Zhou. Mo- bilebert: a compact aask-agnostic bert for limited devices. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 2158-21 2020.
[50] Hokchhay Tann, Soheil Hashemi, R Iris Bahar, and Sherief Reda. Hardware-software codesign of accurate. multiplier-fe free deep neural networks. In Design Automation Conference (DAC), 2017 54th ACM/EDAC/EEE, pages 1-6. IEEE, 2017.
[51] Yee Teh, Victor Bapst, Wojeiech M Czarnecki, John Quan, James Kirkpatrick, Raia Hadsell, Nicolas Heess, and Razvan Pascanu. Distral: Robust multitask reinforcement learning. In Advances in Neural Information Processing Systems, pages 4496-4506, 2017.
[52] Vladimir Vapnik and Rauf Izmailov. Learning using privileged information: similarity control and knowledge transfer. Journal of machine learning research, 16(20232049):55, 2015.
[53] Thijs Vogels, Sai Praneeth Karimireddy, and Martin Jaggi. PowerSGD: Practical Low-Rank Gradient Compression for Distributed Optimization. 33th Advances in Neural Information Processing Systems, 2019.
13

---

[54] Wenhui Wang, Furu Wei, Li Dong, Hangbo Bao, Nan Yang. and Ming Zhou. Minilm: Deep self-attention distillation for task-agnostic compression of pre-trained transformers. Advances in Neural Information Processing Systems, 33:5776-5785 -7788, 2020.
[55] Jianqiao Wangni, Jialei Wang. Ji Liu, and Tong Zhang. Gradient sparsification for caamumication--fffcieee distributed optimization. Advances in Neural Information Pro- cessing Systems, pages 1306-1316, 2018.
[56] Wei Wen, Cong Xu, Feng Yan, Chunpeng Wu, Yandan Wang. Yiran Chen, and Hai Li. Terggrad: Ternary gradients reduce communication distributed deep learning. In Advances Neural Information Processing Systems, page 1509-15199 2017.
[57] Lin Xiao and Tong Zhang. A Proximal Stochastic Gradient Method with Progressive Variance Reduction. SIAM Journal on Optimization, 2014.
[58] Li Yuan, Francis EH Tay, Guilin Li, Tao Wang, and Jiashi Feng. Revisiting Knowledge Distillation via Label Smoothing Regularization. Conference Computer Vision and Pattern Recognition, 2020.
[59] Linfeng Zhang, Jiebo Song, Anni Gao, Jingwei Chen, Chenglong Bao, and Kaisheng Ma. Be your own teacher: Improve the performance of convolutional neural networks via self distillation. In Proceedings of the IEEE/CVF International Conference on Computer Vision, 2019.
pages
[60] Linfeng Zhang, Jiebo Song, Anni Gao, Jingwei Chen, Chenglong Bao, and Kaisheng Ma. Be Your Own Teacher: Improve the Performance of Convolutional Neural Networks via Self Distillation. ICCV, 2019.
[61] Xinshao Wang Ziyun Li, Di Hu, Neil M. Robertson, David A. Clifton, Christoph Meinel, and Haojin Yang. Not All Knowledge Is Created Equal: Mutual Distillation of Confident Knowledge. NeurlPS 2022 (rrnshooprrssworrtty and Socially Responsible Machine Learning). 2022.
14

---

Appendix
Contents
Introduction
2 Related Work
3 Knowledge Distillation
4 Convergence Theory for Sell-Distillation
5 Removing Bias and Improving Variance Reduction
Convergence for Distillation of Compressed Models
7 Discussion and Future Work
10
Basic Facts and Inequalities
16
Proofs for Section
17
B.1
Binary Logistic Regression
17
B.2 Multi-class Classification with Soft-max
18
B.3 Generic Non-linear Classification.
18
Proofs for Section
20
C.1 Key lemma
20
C.2 Proof of Theorem
21
C.3 Proof of Theorem 2
23
Proofs for Section
23
D.I Proof of Theorem 3
23
Proofs for Section
24
E.1 Five Lemmas
24
E.2 Proof of Theorem 4
26
Additional Experimental Validation
28
F.1
The impact of the learning hyperparameters
28
F.2
The impact of the teacher's quality
29
F3 The impact of knowledge distillation compression
29
Limitations
29
15

---

A Basic Facts and Inequalities
To facilitate the reading of technical part of the work, here we present several standard inequalities and basic facts that we are going to use in the proofs.
Usually we need to bound the sum of two error terms by individual errors using the following simple bound
|| + b12 - 2 |all + [161] A direct generalization of this inequality for arbitrary number of summation terms is the following one:
(20)
Σ d2
|0: 1
(21)
This inequality can be seen as a special case of Jensen's 'ss inequality for convex functions h: I2 - R.
h Σ
a,h(z,),
i=1
1=1
where E Rd are any vectors and E [0, 1] with Ei= = 1. Then, (20) follows from Jensen's s
inequality when ((x) || and α1 ... = an = A more general version of the Jensen's inequality, from the perspective of probability theory, is
h(Ez) $ Ez(2)
(22)
for any convex function h and random vector E Rdd.
Another extension of (20), which we actually apply in that form, is Peter-Paul inequality given by
(a, b) - 2x
||||| 2 + l1bl
(23)
for any positive > 0. Then, (20) is the special case of (23) with s =
Typically, function f is called L-smooth if its gradient is Lipschitz continuous with Lipschitz constant L 2 0, namely
[V/(x) V L || ,
(24)
In particular, smoothness inequality (24) implies the following quadratic upper bound [35]]
f(v) < f(x) + (V/(=), - x) + Ill - x0|2,
(25)
for all points y, E Rd. If the function f is additionally convex, then the following lower bound holds too:
1
f(y) 2 f()) + (Vf(x), -x)
(Vf()) - Vv(+))
(26)
2L
As we mentioned in the main part of the paper, function f is called j-strongly convex if the following inequality holds
f(g) 2 f(x) + (Vf(z). x)
/ lly x
(27)
2
for all points x. y Rdd. Recall that strong quasi-convexity assumption (1) the special case of (27) when y = x=. In other words, strong convexity (27) implies strong quasi-convexity (1). Continuing this chain of conditions, strong quasi consi-convexity (1) implies PL condition (9), which, in turn, implies the so-called Quadratic Growth condition given below
f(x) ff 1
(28)
for all & Rd, Derivations of these implications and relationship with other conditions can be found in [17].
As many analyses with linear exponential convergence speed, analysis also uses a very standard transformation from a single-step recurrence relation to convergence incquality. Specifically, assume we have the following recursion for ex, 0, 0,,,
16

---

e+ S <1 (]- + +N,
with constants E (0, 1] and N 2 0. Then, repeated application of the above recursion gives
e' 5 - n)- + +11 (1- n) (1- -1 (1-n)'-23) n + .....
30
N 7)c (1--))
s (1 ') 9 N
Σ
n)²
(29)
n
j=0
If E Rd is a random vector and x E IRdd is fixed (or has randomness independent of 2), then the following decomposition holds:
E +1 + =12] = ||=|| + B [1=13]
(30)
B Proofs for Section 3
For the sake of presentation, we first consider binary logistic regression problem special case of multi-class classification.
B.1 Binary Logistic Regression
In this case we have d-dimensional input vectors on E Rd with their binary true labels b, E B [0, 1). predictor (a) = (x a) E (0, 1) with parameters x E P = Rd+ and lifted input vector á = 1] IRd+1 (to avoid additional notation for bias terms), where o(I) = Ite-7, E
the sigmoid function. Besides, the loss is given by the cross entropy loss below
((p.g) H (1 ) (1=)]) (
-q log (1 9) log(1 -p), P.g [0.
Based (2) and (3), have
Jn 1 A)
= (1 A)fn(z) + Afn(I 0)
(1- (a)) bn) +
((n)) (1 2)b, + ))o(an))
= E(o(x T an), (1 - A)bn + -oa an)) = E(o(xx T an),
8n log + 1 To.) + - 8n) log (1 t Tan)))
where 8n = (1 - 2)bn + (0 an) are the soft labels. Notice that fn 0, = 0) = fa(r). Next, we derive expression for the stochastic gradient for the above objective, namely the gradient V Va 2) of the loss associated with th data point (an, bn).
a
V x 8n log ( (1+e==1am)
-Sn
On -,,0(-x' an) an.
1 x- an
V (1 sn) log (I (l+eeaa)))
e= ān + c=ta, an on = - sn)o [ T an) an.
(1- sn)
(1-3p)o
Hence, using the identity a(1) + + = 1 for the sigmoid function, we get
Vafa( 0,1)
[-s - o(x an)) [xx)) + (1
sn) an.] ān
10xx a.) -s,] ān
[o(x (n)) -(1 (1- a)bn Tan)) on
(o(x (a)) - bm) ān - > (a(0 an) bm) an
(O).
Thus, the distillation gradient for binary logistic regression tasks the same form (5).
17

---

B.2 Multi-class Classification with Soft-max
Now we extend the above steps for multi-class problem. Here we consider a K-classification model with one hidden layer that has soft-max the last layer, i.e., the forward pass has the following steps:
an aTa, + (x = o an) E
where X = [21 32 ... FK] Rdd are the student's model parameters, on A = IRd is the input data and a is the soft-max function (as a generalization of sigmoid function). Then we simplify the loss
(x ()
= (1 ())
(1 -1- X))) ba) (nn) (o
= f(ox(an); (1 N)6n - (d)((n))
= f(o(X) an), an), sn)
K
K
c an
= Σ
sn,k log o (X an)x =
kp,k log
2)11
655
an
Jmm
K
K
K
K
K
Σ
Sn,k
log
e" Or
log en an
Σ
Sn,k log
c
Sn,k log 2
j=1
j=1
k
K
K
K
K
= log
k
Sn,k log c loge
a. = log
c
Sn Ox.
k=1
k=1
k=1
where 8n (1 A)b, + (an) R are the soft labels. Next, we derive an expression for the stochastic gradient for the above objective, namely the gradient V. fn 1 A) of the loss associated with nh data point (an, ba).
K
K Σ S
c an
V Vx 0,A)
Vze
zz
log
c
=
Σx=1 60T0n On - -n
k=1
k=1
(a(XT an)
sn) On
(a(XT ann - -bn) an (a(0T an) - Dn)) an
Vx.f((X)
(v))
Again, we get the same expression (5) for the distillation gradient
VxJa(X A2
B.3 Generic Non-linear Classification
Finally, consider arbitrary classification model that ends with linear layer and soft-max the last layer, ie., the forward pass has the following steps
an - vn(x) + &z (2(an) = (v)) E R
where On E A = Rd is the input data and (x(x)) E R are the logits with respect to the model parameters x. Denote (2()) = (() b,) the loss associated with logits and the true label 6 In words, 2n gives the logits from the input data, while pn gives the loss from given logits. Then, clearly we have the following representation for the loss function fx(z) (z)) Next, let us
18

---

simplify the distillation loss as
fn(x 2) = (1 -
+ +
(1 -1- ()))) b,) +
(e))
E(or(an), (1 1)bn + (an((an)
(1)(vi(a);; - 2)b, + (1i
((o(v)(1)) 8n)
K
elna(z))
Σ
Sn,k log o (1)(())))
E n log
2j=1 eVn (a)
K
Σ
8a,k
log
Mn)(())
Uinnk (2t)
1=1
K
K
K
Σ
Sa,k
log
e (nn
(n,ken)) R=1
k=
==1
log
k=1
where 8n =1 X)b, + (an)(an). Now we need to differentiate obtained expression and derive an expression for the stochastic gradient for the above objective, namely the gradient V fz 9, A) of the loss associated with th data point (an, ba)- Applying the gradient operator, we get
Vafn(1 10,
V
log
(x)
Sn (n
k=1
K
K
Va)nn(())
Σ 1=
sn, Velnn(()
j=1 elns(())
(x)) -8n) (aVin(a)
k=1
= (x) (o(w)(z)) - sn),
where
(x) Oy (x)
Vinn, (x) V Kx),k())) e gaxk
dr
is the Jacobian of vector-valued function vn: Rd + RRKK From the derivation so far we imply that
(v)((x)) sn (o((,,))) bw.) ((o(()()) b..) = (v)(((())) X (v))
Taking into account that fa(x) = (x)(v(())) we show the following form for the distilled gradient
Vafn(z 0, = (1)((z) (((()((V(pp) - -VVp,
(z)(()) (z) - @v)(()) (f)
Or Oih
Ox (1)
19

---

C Proofs for Section
Before we proceed to the proofs of Theorems and we prove a key lemma that will be useful in both proofs. The lemma we are about to present covers Part (Optimal distillation weight) and Part 3 (Impact of the teacher) the proof overview discussed in the main content.
C.1
Key lemma
To simplify the expressions in our proofs, let us introduce some notation describing stochastic gradients. Denote the signal-to-noise ratio with respect to parameters 0 by
IV ((O))22 E (0) E [0. 1],
((0)
(31)
and correlation coefficient between stochastic gradients f(e(()) and Vfe(y) by
Cov(2, y)
p(x,z)= = Var(x)Var()) E [[1,1]]
(32)
where
Vf(x), Vie(v) Vs(II))
Var(o) (C,((2,
the covariance and variance respectively.
Lemma 1. Let N(A) (ZI eyE [[I { for some constant 2 0.
Then, the optimal X that minimizes (A) is given by
E (Vfg(=)), V
(33)
E +
Moreover,
(Aa)) N(0)
,²(=", 1- 3(0
S min (1.0 (1)))))))))
1+ 1 3(0)
Proof. Notice that ((A) is quadratic and using the first-order optimality condition, we conclude
d (X))==2/V/(0111 (-22 (VSg(=") Vfc(O))] 2XE [IV]{(())1)) 0,
d
we get (33). Furthermore, plugging the expression A. into (A), we get N(A.) = cr
X2
V =
E + (I
E[IVS( + ب, ))(O))))
cr
[I [ - 2X. V + X
E
E V
[
E
E +
(VV + (E (e(I
c7 [1V
( |
Note that ((0 c70 From E [[ v/(²") 0. we imply that
(")) - -1000-1001177700000000000000-1-- -
and therefore
E[V
plx,00 V
20

---

Now we can simplify the expression for (A.) as follows
((a))
[E Vfe())))
1
N(0)
E V (E[Vft(==', V (E
++01+++
E -
1 8
1 1- (x" 0)
1+ m4 (0)
Here, p(=", 0) is the correlation coefficient between stochastic gradients at x¹ and 0. Hence, showed with tuned distillation weight the neighborhood can shrink by some factor depending on the teacher's parameters. In the extreme case when the teacher 0 = x² is optimal, we have p(x*, 0) 1, (()) = and, thus, no neighborhood N(2.) = This hints us on the fact that the reduction factor N (A.)/N(()) of the neighhorhood is controlled by the "quality" of the teacher.
To make this argument rigorous, consider the teacher's model to be away from the optimal solution xx within the limit described by the following inequality
f(0)- - 0 (x=))(
(34)
C
where d²(x) E [[()))] is the second moment of the stochastic gradients. Without loss of generality, we assume that (z > for all parameter choices æ E Rd. otherwise we have a = and even plain SGD ensures full variance reduction. Then, simplify the reduction factor
1 3())
(E((V (((()))
1
1 0² 8
1+ 300)
E (
1 1 5, 3(0)
0²(0) =000)+ E I (x()) 2
(x)
27 3(0)
2
o2(0) + 0z)))
E [[I [1
(0)o((x**
(0)
1 0r 3(0)
(10)+)33)
L()(0) m3)
1
a(0)o(z*)
1
2L f1o--"
o7 200
2L o(x')
(((())
2
e7C (0))
((0))(2*)
2C
ey (2)) (f(0) - - - (-
2b
aax))a(0))
where the last inequality used 1 (1-u)² 1+a0
< (2 + v)u for all u, 2 0.
C.2 Proof of Theorem
Denote by [ = E[ xx]] the conditional expectation with respect to x'. Then, we start bounding the error using the update rule (7).
21

---

Er [| ++ - x²1
|x - x² z¹113 - 27 (x² x², Vf(a') - V())))--NVVV0)--7((((())7)))771
1| - z¹11 =x'12 2y (x' x,, Vf((z)) (V((=)) +2 (1V(())))))
(1+
(1 7/121100 10
+ B, [1V.f(('') -
(1 7x)lla - 1- [[V -V
(28)+223) S
7/1 1| 12 7/1
|| ==11== +
81 X² (
8
P
+ B, [1V (x)) 22 (VVe((=') (()((())1
(3)
87 {1)
(1 - -"'
+ + (V
f(x)))
-
+ " x -
87 _ + 2
(1 (1---)ll''-""
- -
where we used Peter-Paul inequality (23) with parameter s _ and the step-size bound ? S 42 in the last inequality. Applying full expectation and unrolling the recursion, we get
87
|1100 - 1²]
<
(1 - --/))E
+
x2 ((()) + 2 [[V - AV))((
(29)
812
|V + 2B
(I - 11)))/11100 - x=1
- -
n22
(1 - (--))))(a>>> 11 -1=0
1012
8
(A),
(35)
where N1 = X2I [V( (()))1 Applying Lemma I with 4/4,
we imply that for some À = X the neighborhood size
Lemma (1.) S (1) min (1,0 (((10)))))) r1)) no? min (y,0(f(0) - -f))
Plugging the above bound of N2 into (35) completes the proof.
22

---

C.3 Proof of Theorem 2
We start the recursion from the L-smoothness condition of f. As before, E, denotes conditional expectation with respect to x².
E. [ - |*]
(25)
(f(x') - M* - (VS('), V((")) + L-2
2 T. -VV
(20)
(f(x') - f*)
E, [1V]'((') - V -VV
(9423)
(1 - 7p) (f(x") /r) - 2 (f(x*) f4) -
-
S
+ E, [IV -
(0)
(
(1 - (1) (J(x") -)) TPE
- ( -
2
+ 2 (((x')17) + B
(1 - (2) (f(=*) - 7/4 1 1-7
4LC
(f(=9)-')
-
1
- +
+ > +
S (1- - 7p) (f(=") + + L,
-
where in the last inequality we used step-size bound 7 iz +. Applying full expectation and unrolling the recursion, we get
[[((a')-
S (1 TR [5 [1] + >²2 (0)) I, (IV])((=*) - XV)
(29)
S (1 mu) (f(x*)-13)) - -5") + X² /
+ 17 E [[V] AV))(
(1 (1- 7w)² (f(z9)- f²) + 1 N2)),
(36)
where N₂(A)= = f(0)I EA [g(x)) - AV (0)) Similar to the previous case, we applying Lemma 1 with c = L and conclude that for some X = X. the neighborhood size is
Lemma 1 S
(0) min (1. (4 (1
(2,).)
Loo min (7, (((00) (*)).
Plugging the above bound of N₂ into (36) completes the proof.
D Proofs for Section
D.1 Proof of Theorem 3
Again, we start the recursion from the smoothness condition of f.
E. [[xxtt]] - [*]
(25)
(f(x') - 1*) - (Vf(x') V.(()) + Ln2
E.
[[V - V1e(0)+ V/())
2
(21)
(f(x') - f*) ->(()(((')1
3
3
3 -. T>>
+ #2 2R [IV V(e(2')) [13] + I 2 [[V])(()) V +
2
(91(10) S
(1 - 7P) (f(x*) - f*) - -1 (f(2)))) - )
- + 3 (f(x')- - -f) + +3LC? (f(0) +3' f")
<
(1-7w) (f(z")- fr) + +LL 22² ff)
23

---

where we used size bound y - 4 in the last step. Therefore, 3LC
[ - M
(1 - 74)E T)))E(((a') - 1"] 31 1)7 - 1)) (1 ())=())) 5r) + 3 + L) 1 (5(0) 1),
which concludes the proof.
E Proofs for Section 6
First, we break the update rule (18) into two parts by introducing an auxiliary model parameters y² e gd:
y g+ x' -
x 2+1 C(y't).
Without loss of generality, we assume that initialization satisfies x. = yo = C(y)). Then, for all t 2 0 we have x² = C(y') and the update rule of y can be written recursively without x² via
y + C(y') -
(O)).
(37)
Using unbiasedness of the compression operator C, we decompose the error
E ||
x 11
(30)
E || - y 112] + EE |1 - x |1²]
E |||cu')) - 2/ 11]]] + E 'lll x= 11²]
WE ||0511 + E [lle* x= 11]]
(2w + 1) ||I* x=" [1]] 2w ||x="1 11?
(38)
Thus, our goal would be to analyze iterates (37) and derive the rate for x², In fact, the special case of (37) was analyzed by [18] in the non-stochastic case and without distillation (A 0). The analysis we provide here for (37) is based on [18], and from this perspective, our analysis can be seen as an extension of their analysis.
To avoid another notation for the expected smoothness constant (see Assumption 3), analogous to (24) we assume that L also satisfies the following smoothness inequality:
E [VVee())
V < < y11
(39)
E.1 Five Lemmas
First, we need to upper bound the compression error by function suboptimality. Denote 8(x) =
C(x) ((z)-z. z. Let Es and Be be the expectations with respect the compression operator C. and sampling & respectively.
Lemma (Lemma [18]). Let a = 2n and v = 2w [ax0>>1]] For all x E Rd we have
Es 11((x)111 < 2a (f(x) - f(x))) + +v,
(40)
Proof. From (20) we imply ||a||2 < 2| 10 + 2|| 1|² and from p-strong convexity condition (27) we get || 111 3/2 (f(x) f(x")). Putting these inequalities together, we arrive at
4w (f(x) f(x") ) + 2w |
E, 1 < 2|| - 20|l x²122 2w| 10 11
P
24

---

Lemma 3. For all E R we have
E V + (
- V & 6² (1x- yll + BB1((()))))
(41)
sff + (()) - f()) + (Vi(v), x- - y) + 2 | - vlI + I. (1)0())
(42)
Proof. Fix a and let 5 = (z). Inequality (41) follows from Lipschitz continuity of the gradient, applying expectation and using (19):
E V + 8) - Vfe(ii)l C3 + 8 y/l² (19)+ C (1x2- y/1² + Eslldl²).
The first inequality in (42) follows by applying Jensen's inequality (22) and using (19). Since f is L-smooth, have
(25)
L 1 + 5 yl1
Ef(x + 6)
ES(a) (VI()), x + 0 y) +
(19++330)
f(:) + (VI()), x .) +
L
|
ul1 + 2 11101
Lemma 4. For all y E Ted holds
8(x)
B5
+
7
2 AV 2C || +2
+2 (c
1
Es|l(())1? (43)
72
Proof.
Fix and let 5 = 8(x). Then for every E IRd we can write
E,
Vfg(I + 5) + (0)
6
E,
Vie(v) fe 8) +
7
8
2E.
V(e(v) + (V) + Vee(x V
Ba|18112
Ed. (v AV) + [VV((()) AV + 2c (lx=-vl u12 + Esllo1?)
222
(19)
2
2
Es|l8112
+ - 2c² (1K-~1² v1 Es|ld|²))
The next lemma generalizes the strong convexity inequality (27) (special case of 6(x) 0).
Lemma 5. If s-moooth and y-strongly convex, then for all x. y E Rdd, holds
L -
f(v) 2 f(x) ( V fx + 8)] x² 11 x
E11((1)))
(44)
2
Proof. Fix x and let 5 = 5(z). Using (27) with r x 8, we get
f()) 2 / (x + 6) + (V () 8) + #
|l x 5112
2
Applying expectation and (30), we get
f() 2 Ef(r + E(V +
11 1ly 2012 + 1
1|8811
2
2
25

---

It remains to bound the term -E(Vf(x + 6), 6), which can be done using L-smoothness and applying expectation as follows:
-V( 00,
29)
2 E [(x) f(x +0) [|]1² f(x) EEf( + 8)
L
B1881
2
2C + (C² + } and B= (c+ 2/2) v. where a, v are defined in
Lemma Denote
Lemma 2. Then
8(x)
E
+ + (V
+
$ 4x) - f(xx)) + +B V[((') AV)((())?? (45)
Proof. Using (43) with y x, we get
8(x)
E
+
- +2 C² 1
B110(z)11
72
(20++
+4 Ve vf AV + -2 ( ( f(xx)) +v)
S
SC(f(x)
4 2c (2²
a (f(x)-
- f(x.)) +2 C2
1 (2))
v 4E -
E.2 Proof of Theorem 4
Denoting 8¹ = -(yy)) we have C(r*)= 36 + 5¹). Then
11/2 - z¹11 -x112
I[c(2") [V(e((()))) + +>A[(((0) - =01
110' x² + - +
1|o' - 211 -=-1=
2 VV + AV -z) =110* VV +80 (x)
26

---

Taking conditional expectation E,
:0 E y" we get
Er || - x=
118 - x²11² + +2 (Be [Vf(y' + +8'] - XVf( x- -1¹) + Es [110 -VV +8 + +
(44) S
110/ - x²12 + 2 f(xx) - f(v") - 5" 'lel - =x-1/
2
L L- "E, [|0-13]
+ + ------- x²)
2
8¹
+ B
Vie(V' + +00) + (O
(1 7/4 11
112 1|1 2-((('') - f(=')) - W)), 11011 +2 -x) + (1 - 7141 1// x012 - (f((') - f(x))) - 7044
8' Vieee' + + (0)
B,
(27)+(23)
=01
110' -y'-2"1 =211 + 13 |811
2
+ 2/ 11 - 111 + 271²
| 1 Veel + +0²) + xx
2
11
(1 (041 /21 1= - +( 0'1| + 27 |
+
+4> - --00-//()))---------1/711+)7>))-
(1 - 7/4) 1100 -1 7(L 14) |10º11
=
2y1²
1
(40)
(1 7/4 //1 - x01 - 1)((())) - f(x')) + 2
(T 11)
- (2()((((==) +D)
+ 2712
|] + = AV
1
(1
7/1 101 - + 22 - μ4) - - +2 +7(L - -1)x
2
4 -VV
P
where a and v in Lemma 2 and A and B are defined Lemma 6. Next, we show that the bounds
on y and w lead to A + a(L p)) S 1/2. Plugging the expression of A, we the former inequality becomes
2
2C
L2 +
a
+ (L- 1)) = 2
Rearranging terms, we get
2w S
1/2 476
p 2 +
2 7
+ 1
For y S 162. the above inequality holds i = O(4)).. If = zēž the condition on variance
1
parameter w becomes w = (v//c). Thus, by assumption on w and y, we have 2 aa p)) S 1/2. and hence
Ee [11 41 =1]]] s 711) 1100 x x01 + D,
where
2712
D = 2²²B + 7(L - 1)v
[V)) + +4,7B -
2Y
= 22²B + 7 - +1v== - Na(A),
P
27

---

Validation Loss Statistics
0.55
0.45
0.40
0.35
0.0
0.05
0.1
0.3 0.4 0.5
0.7 0.8
1.0
Lambda
Figure 4: Validation loss statistics for the same setup as in Figure 2a.
with (N) = = ((x fc()l Applying Lemma with 2p get (x((*) (0) min (1, (1)( 2nos min (v. O(f(0) f*)). Taking expecta- tion, unrolling the recurrence, and applying the tower property, we get
E ||v - x 11] s (1-74)' 11 x² z¹1
D
7p
where the neighborhood is given by
D 1 + + 44v 2 N(A))
1 (2 + 1////v + + -4) + 27 N(Aa))
7/
7/
P 122722 + 2/7 L pi) + n
2
N(X*)
" (2 C2 + 2/7 + L p)) + 4o
min (7. ((f(0) - f*)).
Ignoring the absolute constants and using bounds
S I6C and w O(rr),, we have
E Ilv' x017
=01 s 740 1| 2²1' +
O
-
|
40 min (7, 0(f(0) f*)).
Applying this inequality in (38) conclude the proof.?
F Additional Experimental Validation
In this section we provide additional experimental validation for our theory.
E1
The impact of the learning hyperparameters
We begin by measuring the impact that the optimization parameters, in particular the step size, have on the convergence of SGD and KD. For this, we perform experiments linear models trained on the MNIST dataset, without momentum and regularization, using a mini-batch size of 10, for a total of 100 epochs. We compute the cross entropy train loss between the student and true labels, measured as a running average over all iterations within an epoch, similar to Figure 2. We also compute the minimum train loss, as well as the average and standard deviation across the last 10 epochs. The results in Figure show the impact that different learning rates have on the overall training dynamics of self-distillation. In all cases, the teacher was trained using the same hyperparamters as the self-dit distillld models (A in the plot). We can see that using higher learning rate introduces more variance in the SGD update, and KD would have more pronounced variance reduction effect. In all cases, however, we can find an optimum achieving lower train loss compared to SGD.
The rate in Theorem 4 uses 7 Tez value for the learning rate.
28

---

1440
Waln Lags Statistis
8375
Train Loss Statsticss
0.30
9.355
3.30
0.05 6.1
0.5 0.6
0.7
0.0 0.01 0.05
0.3 04 0.5
0.9
1.0
Lambda
Lambda
(a) The impact of the step size on the learning dy- namics of SGD and KD. The teacher was trained dynamics of self-distillation. The same teacher with the same hyperparameters the correspond- was used in both setups. ing self-distilled models.
(b) The impact of better teacher on the learning
Figure 5: Ablation study on the training loss of self-distillation and SGD, when taking into account the training hyperparameters (learning rate) and quality of the teacher.
Sparsit
Train Loss Statistics
Random, Sparsity
0.34
0.33
0.32
0.31
0.30
0.29
0.28
0.5
Lambda
Figure 6: The impact of self-distillation on the train loss of compressed models. Here a random mask computed at initialization, and kept fixed throughout training.
F22 The impact of the teacher's quality
Next, we quantify the impact that better trained teacher could have on self-distillation. Using the same setup of convex MNIST above, we perform self-distillation using a better teacher, i.e. one that achieves 93.7% train accuracy and 92.5% test accuracy. (In comparison, the teacher trained using step size of 0.05 achieved 92% train accuracy and 90.9% test accuracy) We can see in Figure 5b that this better-trained teacher has more substantial impact on the models which inherently have higher variance, i.e. those trained with a higher learning rate; in this case, the optimal value of À is closer to which is also suggested by the theory (see Equation 14). We note that a similar behavior was also observed on the CIFAR-10 features linear classification task presented in Figure 2b.
F.3
The impact of knowledge distillation on compression
Now we turn our attention towards validating the theoretical results developed in the context of knowledge distillation for compressed models, presented in Section 6. We consider again the convex MNIST setting, as described in the previous section, and we perform self-distillation from the better trained teacher. We prune the weights at initialization, using random mask at chosen sparsity level, and we apply this fixed mask after each parameter update. The results presented in Figure 6 show that self-distillation can indeed reduce the train loss, compared to SGD, even for compressed updates. Moreover, we observe that with increased sparsity the impact of self-distillation is less pronounced, as also suggested by the theory.
G Limitations
Lastly, following our discussion from Section 7, we discuss some limitations of our work.
As a theoretical paper, we used several assumptions to make our claims rigorous. However, one can always question each assumption and extend the theory under certain relaxations.
29

---

Our theoretical claims are based on strong (quasi) convexity or Polyak-Lojasiewiice condition, which are standard assumptions in the optimization literature.
Another limitation concerning the "distillation for compression" part of our theory is the unbiasedness condition E [C(x)] = in Assumption Ideally, we would utilize any "biased" compression operator, such TopK, with similar convergence properties. However, it is known that biased estimators (c.g.. biased compression operators or biased stochastic gradients) are harder to analyze in theory.
30


Reviews:
[{"review_id": "vWp3FIL0z7", "reviewer": "Reviewer_1HYD", "summary": "This work analyzes Knowledge Distillation (KD) using an optimization point of view.\nBy 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.\nUsing some approximations, the authors argue that their analysis should hold for deep networks too.\n\nThe paper complements its analysis by proposing a technique to reduce the bias of KD and extend the result to generic KD with smaller students.\nThe proposes analysis is supported by experiments for linear models.", "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?\n\nI 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.\nI 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.", "limitations": "The authors have addressed the limitations of their work properly.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is clear and easy to read even for non-optimization experts.\nThe whole idea of analyzing KD as a variance-reducing algorithm is novel and interesting\nAlthough the assumptions are limited to linear models, the empirical experiments show that such assumptions approximately could hold deep networks is convincing.\nThe 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.\nWhile 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.\nI don't think the authors could either add some empirical evidence on this matter or focus on the linear/convex case.\nI any case, I encourage the authors to further investigate KD as variance-reductions on deep networks, even if only empirically in another paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your comments and positive evaluation of our work.\n\n\n- *The biggest weakness of this work is its ambivalence on the validity of the results for deep learning. ...*\n\n***Response.*** We acknowledge the fact that the focus of our work is analytical: we identify a first non-trivial interpretation of KD from the optimization perspective, and characterize cases when KD can lead to better rates. Our interpretation appears to match experiments precisely in the convex case. Our further experiments are meant to show that our “distillation gradient” interpretation can also be relevant in the case of SGD-based optimization of DNNs. We indeed plan to investigate the challenging DNN / non-convex case more precisely in future work, for which both new techniques and further experimental validation will be needed.\n\n- *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?*\n\n***Response.*** Yes, we train a linear classifier on top of the features extracted from a ResNet50 model pre-trained on ImageNet. This is a standard setting, commonly used in the transfer learning literature, see e.g. [1], [2], [3].\n\n[1] Simon Kornblith, Jonathon Shlens, Quoc V. Le, *Do Better ImageNet Models Transfer Better?* CVPR, 2019 (https://arxiv.org/abs/1805.08974)\n\n[2] Hadi Salman, Andrew Ilyas, Logan Engstrom, Ashish Kapoor, Aleksander Madry, *Do Adversarially Robust ImageNet Models Transfer Better?* NeurIPS, 2020 (https://arxiv.org/abs/2007.08489)\n\n[3] Eugenia Iofinova, Alexandra Peste, Mark Kurtz, Dan Alistarh, *How Well Do Sparse Imagenet Models Transfer?* CVPR, 2022 (https://arxiv.org/abs/2111.13445)\n\n\n- *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.*\n\n***Response.*** We acknowledge that only provide partial evidence for the validity of our “distillation gradient” assumption for deep models. We provided evidence that the assumption holds approximately for deep networks, and show varying behaviour of the alignment during training. \nFurther, we emphasize that our Theorem 2 is designed to hold specifically for self-distillation in the non-convex case, which approximates very popular heuristics in practice, e.g. [55].   \n\nWe agree that there is much more to investigate for the case of deep networks where exact tracking of teacher’s impact across multiple layers of non-linearities becomes harder. We see our results as a promising first step towards a more complete understanding of the effectiveness of distillation, and will provide a more complete discussion towards applicability in the non-convex case in the next revision.  \n\n\n"}, {"review_id": "JmD79asQ4m", "reviewer": "Reviewer_6fjL", "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.", "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?\n\n2. 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.\n\n3. Beyond empirical evidence, is there an intuitive reason to believe that Prop. 1 would approximately apply to deep neural networks?\n", "limitations": "Yes.", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "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.\n* The connection between the distillation gradient and variance reduction methods like SVRG is interesting and novel to the best of my knowledge.\n* The presented analyses and results seem sound.\n* The authors present some empirical evidence that support the claims of the paper (Figs. 2-5).\n", "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).\n* The presented results, including the empirical ones, apply to training loss, rather than test error.\n* 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.\n\n[1] https://arxiv.org/pdf/1902.03393.pdf \n\n[2] https://arxiv.org/pdf/2202.03680.pdf \n\n[3] https://arxiv.org/pdf/2206.06067.pdf ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your comments and evaluation of our work.\n\n\n- *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).*\n\n***Response.*** We acknowledge the fact that the focus of our work is analytical: we identify a first non-trivial interpretation of KD from the optimization perspective, and characterize cases when KD can lead to better rates. Our interpretation appears to match experiments precisely in the convex case. Our further experiments are meant to show that our “distillation gradient” interpretation can also be relevant in the case of SGD-based optimization of DNNs. We indeed plan to investigate the challenging DNN / non-convex case more precisely in future work, for which both new techniques and further experimental validation will be needed.\n\n- *The presented results, including the empirical ones, apply to training loss, rather than test error.*\n\n***Response.*** As we also mentioned in the abstract, in this work, we investigate KD specifically from an optimization perspective. Of course, investigating test errors is also important. However, it is not the goal of this work. Nonetheless, we added a plot (Figure 8 in the response PDF file) showing validation loss for the same setup as Figure 2(a) with very similar behavior.\n\n\n- *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?*\n\n***Response.*** To reconcile the mentioned prior observations with our results, notice that the notion of a “better teacher” has different meanings. In the papers you mentioned, a better teacher means a teacher with a much larger capacity (for instance wider and/or deeper architecture) which can have higher performance than the student model. However, in our case “better teacher” means better parameter (i.e., weights and biases) values, evaluated in terms of training loss.\n\nIn particular, in the case of self-distillation, covered in Sections 4 and 5, the teacher and student architectures are identical, and hence they have the same capacity. Here, a better teacher means better parameter values within the same architecture. \n\nIn our second regime, distillation for compressed models (Section 6), we actually consider the case when the student network is a subnetwork of the teacher; we consider a sparsification compression operator that selects $k$ parameters for the student out of $d$ parameters of the teacher. Then, clearly, the teacher has a larger capacity with a capacity ratio $d/k\\ge1$. However, our result in this direction (Theorem 4) does not allow the capacity ratio to be arbitrarily large. Indeed, the constraint $\\omega = \\mathcal{O}(\\mu/\\mathcal{L})$ on compression variance implies a constraint on capacity ratio since $\\omega = d/k-1$ for the sparsification operator. Thus, our result holds when the teacher’s size is not significantly larger than the student’s size, which does not contradict the observations from prior work noted by the reviewer. \n\n\n- *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.*\n\n***Response.*** Please see our PDF response where we added the plot (Figure 6) with a larger 100 epoch training supporting the claim on stability. The decrease of cosine similarity can be explained as follows: at the beginning the cosine similarity is high (and SNR is low) since we start from the same model. Then, initial perturbations caused by either the KD or modified KD gradient don’t cause big shifts (the teacher has enough confidence and small gradients). These perturbations accumulate over the training leading to decreased cosine similarity and eventually stabilize.\n\n- *3. Beyond empirical evidence, is there an intuitive reason to believe that Prop. 1 would approximately apply to deep neural networks?*\n\n***Response.*** Yes, there is. Intuitively, distillation has essentially negligible effect on the data points for which the teacher classifies correctly with high confidence (since the distillation gradient almost equals the regular gradient). In such cases, the feedback from the teacher is similar to or the same as from the true labels. This intuition is reflected in Proposition 1 since high confidence of teacher’s classification means that teacher’s gradient $\\nabla f_n(\\theta)$ with respect to those data points is close to zero and does not affect the student’s gradient much.\n"}, {"review_id": "OzjANRr3Wb", "reviewer": "Reviewer_AWoB", "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.", "questions": "In addition to the \"weaknesses\" above, here are some questions/suggestions.\n\n- 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?\n\n- It would be great if there is a plot that explicitly tracks the gradient variance.", "limitations": "The limitations are given in the appendices, which I think is okay.", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "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.\n\n- 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.\n\n- The contributions and the proposed unbiased KD are novel, as far as I know.\n\n- 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.\n\n- 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.\n\n- 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)?\n\n- 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.\n\n- 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.\n\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your comments and positive evaluation of our work\n\n- *It would have been better if ...*\n\n***Response.*** We acknowledge this suggestion, and in fact this is one of the key aspects of our theory which is discussed in lines 269-277 (Importance of the results) after the proof overview of the theorems. The comparison of the convergence rates between theorems 1 & 2 and SGD can be done as follows. For both setups, the rate of SGD is the same (11) or (12) with only one difference: $\\min(\\gamma, \\mathcal{O}(f(\\theta) - f^*))$ term is replaced with $\\gamma$. So, $\\mathcal{O}(f(\\theta) - f^*)$ is the factor that makes our results ***better*** compared to SGD in terms of optimization performance. \n\nWe thank you for this suggestion, and will add further clarifying discussion on this point.\n\n- *Two logical connections could be made a little bit clearer...*\n\n***Response.*** Thank you for the suggestion, we will make them clearer in the revision.\n\n*(1)* Exactly! Equation (5) leads to partial VR because of the additional $ -\\lambda\\nabla f_{\\xi}(\\theta)$ term. When chosen properly (distillation weight $\\lambda$ and proximity of $\\theta$ to the solution $x^*$), this additional stochastic gradient is capable of adjusting the student’s stochastic gradient since both are computed using the same batch from the train data. In other words, both gradients have the same source of randomness which makes partial cancellations feasible.\nAs a matter of fact, we do have direct empirical corroboration for this in Figure 3 (see also the new plot, Figure 7 in our response PDF, tracking gradient variance for the same setup). The Blue line is pure SGD, without any distillation involved. The Green line employs distillation with distillation weight $\\lambda=0.4$ and achieves uniformly lower train errors (and gradient variance).\n\n*(2)* First of all, in the optimization perspective we adopt, “better training” means precisely smaller train error. The analysis of test error is a different aspect which we do not consider here. The standard/complete variance reduction mechanism (such as SVRG) is capable of removing the stochastic neighborhood term completely, and guarantees that the convergence is to the exact minimizer, as opposed to SGD. To give a specific example, in the strongly convex and smooth regime SGD with constant learning rate still converges up to some neighborhood of the minimizer, while SVRG converges to the exact solution with constant learning rate. See reference [4] below for a complete discussion. \n\n[4] Robert M. Gower, Mark Schmidt, Francis Bach, Peter Richtarik, *Variance-Reduced Methods for Machine Learning*, IEEE, 2020 (https://arxiv.org/abs/2010.00892).\n\nWe will provide a clarifying discussion on both points in the next version of our paper. \n\n- *In Figure 1, I am not sure whether the cosine similarity is the right metric to use. ...*\n\n***Response.*** We believe that cosine similarity is a relevant metric since it measures the alignment of the gradients. Nevertheless, to fully address this concern, we additionally provided plots tracking $\\ell_2$ distance and SNR for the same gradients in the rebuttal PDF (Figure 6). As we can see, they follow exactly the same trend as the plot of cosine similarity.\n\n\n- *The empirical results on the proposed unbiased KD are given in the form of \"training loss\" only, rather than test loss ...*\n\n***Response.*** Please note that the purpose of proposing / analyzing unbiased KD is to support our claim that the potential source of ***partial*** variance reduction in KD is the biased nature of distillation. Indeed, we removed the biased and improved variance reduction both analytically and empirically. Analyzing test loss/accuracy of unbiased KD would be interesting in itself, but is irrelevant for our argument related to variance reduction. In fact, unbiased KD closely relates to the popular and well-studied SVRG algorithm (please see the references in the paper), applied in the context of deep networks.  \n\n- *A mild suggestion on notations*\n- *A minor mistake*\n\n***Response.*** Thank you for the detailed comments on improving the presentation, which we will take into account for the next version of the work. \n\n- *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?*\n\n***Response.*** Yes, we believe it is possible to extend the current method of analysis to general knowledge distillation under certain restrictions. The main issue there is obtaining reasonably accurate analytical expressions for the distillation gradients. If the student’s architecture is not the same as the teacher’s nor is a sub-network of the teacher’s architecture, then it may be tricky to derive a closed form for the distillation gradient. Technically, the gradient space of the teacher could be very different from the gradient space of the student; dimensions may not match and the student's gradient space is not a subspace of the teacher's gradient space. One possible workaround could be to impose additional assumptions on how the student’s gradient space can be embedded reasonably into the teacher’s gradient space. (For instance, this would be the case with pruning, quantization, or structured compression/neural architecture search.) In such cases, we will indeed be able to extend Proposition 1.\n\n- *It would be great if there is a plot that explicitly tracks the gradient variance.*\n\n***Response.*** Please find the plot in our PDF response (Figure 7). The plot explicitly tracks gradient variance (averaged over the iterations within each epoch) for the same setup as Figure 3. As expected, both variants of KD (biased and unbiased) have reduced gradient variance compared to plain SGD. The plot also highlights that both variants of KD have similar variance reduction properties, while the unbiasedness of unbiased KD amplifies the reduction of train loss in Figure 3."}, {"review_id": "ui23gpUlAT", "reviewer": "Reviewer_ASQK", "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.", "questions": "N/A", "limitations": "would be interesting if there are also results concerning generalizations.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "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.\nThe 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. \n2. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your comments and positive evaluation of our work.\n\n- *There is no explanation why LBFGS teacher is better than SGD teacher.*\n\n***Response.*** This is simply because LBFGS as an optimization algorithm can train the teacher to achieve almost zero train loss, while the SGD teacher converges only within some neighbourhood of the optimal solution. Throughout, for us, having a better teacher means that the loss $f(\\theta) - f^*$ is smaller.\n\n- *The optimal choice of distillation weight  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.*\n\n***Response.*** This is correct. From the practical perspective, our results could be interpreted in (at least) two ways: \n\n*(1)* Our characterization shows that one could build a “line search”-type procedure to approximate the optimal choice of distillation weight for a given teacher. \n\n*(2)* More broadly, our analysis suggests that careful tuning of the distillation weight parameter can be crucial to obtain the full benefits of KD.\n\n\n- *Would be interesting if there are also results concerning generalizations.*\n\n***Response.*** Indeed, results on generalization would be interesting, and we plan to investigate this in future work. However, in this work we take an optimization view of KD, which means that we focus on training loss."}]

Author rebuttal:
Dear Reviewers and Area Chair,

Thank you for the time and effort you put into evaluating our work. We have responded to all comments in your reviews by providing additional discussions/clarifications and numerical validations (please find the attached PDF response for the additional plots).

Please let us know whether we managed to address your concerns regarding the paper.

Regards,
Authors


Meta-review: This works studies knowledge distillation through the optimization lens and shows that distillation can lead to variance reduction of SGD. The paper presents a theoretical analysis of deep linear models, and empirically verifies that the predictions of the theory carry over to general deep networks. All reviewers find the connection between distillation and variance reduction interesting and novel and they commonly recommend accept. A shared shortcoming of this work mentioned by the reviewers is that the linear network analysis is not enough to guarantee similar results always hold true for general neural networks . However, they also agree that the paper even has enough contribution to be published. In concordance with the reviewers, I recommend accept. Please incorporate the clarifications that you provided during the rebuttal into the final version of the paper.
Decision: Accept (poster)