SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Mher Safaryan
mher.safaryan@ista.ac.at
Alexandra Peste
alexandra.peste@ista.ac.at
dan.alistarh@ista.ac.at
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 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 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 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 (KD) [13, 3] is a standard tool for transferring information between a machine learning model of lower representational capacity—usually called the student—and a more accurate and powerful teacher model. In the context of classification using neural networks, it is common to consider the student to be a smaller network [2], whereas the teacher is a network that is larger and more computationally-heavy, but also more accurate. Assuming a supervised classification task, distillation consists in training the student to minimize the cross-entropy with respect to the teacher's logits on every given sample, in addition to minimizing the standard cross-entropy loss with respect to the ground truth labels.
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 the 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
known about the interplay between distillation and stochastic gradient descent (SGD), which is the standard optimization setting in which this method is 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 a simple re-formulation of distillation in the context of gradient-based optimization, which allows us to connect KD to stochastic variance reduction techniques, such as 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, as well as distillation for compression, where a compressed model leverages outputs from an uncompressed one during training.
In a 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 completely 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, we consider the case of self-distillation [59], which is a popular supervised training technique in which both the student model  $x \in \mathbb{R}^d$  and the teacher model  $\theta \in \mathbb{R}^d$  have the same structure and dimensionality. The process starts from a teacher model  $\theta$  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  $k$  becomes the teacher for iteration  $k + 1$ , and so on.
Our first observation is that, in the case of self-distillation with teacher weight  $1 \geq \lambda \geq 0$ , the gradient of the distilled student model on a sample  $i$ , denoted by  $\nabla_{x}f_{i}(x \mid \theta, \lambda)$  at a given iteration can simply be written as
$$
\nabla_ {x} f _ {i} (x \mid \theta , \lambda) \simeq \nabla f _ {i} (x) - \lambda \nabla f _ {i} (\theta),
$$
where  $\nabla f_{i}(x)$  and  $\nabla f_{i}(\theta)$  are the student's and teacher's standard gradients on sample  $i$ , 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-trained version of the same model), and more general distillation, in which the student is a 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 on the gradient variance  $\sigma^2$  is dampened by a term depending on the gap between the teacher model and an optimum. Intuitively, this says that, if the teacher's accuracy is 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 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.
We now provide an 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 to 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 a 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: (1) 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 is 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
3.1. Background. Assume we are given a finite dataset  $\{(a_n, b_n) \mid n = 1, 2, \ldots, N\}$ , where inputs  $a_n \in \mathcal{A}$  (e.g., vectors from  $\mathbb{R}^d$ ) and outputs  $b_n \in \mathcal{B}$  (e.g., categorical labels or real numbers). Consider a set of models  $\mathcal{F} = \{\phi_x : \mathcal{A} \to \mathcal{B} \mid x \in \mathcal{P} \subseteq \mathbb{R}^d\}$  with fixed neural network architecture parameterized by vector  $x$ . Depending on the supervised learning task and the class  $\mathcal{F}$  of models, we define a loss function  $\ell: \mathcal{B} \times \mathcal{B} \to \mathbb{R}_+$  in order to measure the performance of the model. In particular, the loss associated with a data point  $(a_n, b_n)$  and model  $\phi_x \in \mathcal{F}$  would be  $\ell(\phi_x(a_n), b_n)$ . In this framework, the standard Empirical Risk Minimization (ERM) takes the following form:
$$
\min  _ {x \in \mathbb {R} ^ {d}} \frac {1}{N} \sum_ {n = 1} ^ {N} \ell \left(\phi_ {x} \left(a _ {n}\right), b _ {n}\right). \tag {1}
$$
In the objective above, the model  $\phi_{x}$  is trained to match the true outputs  $b_{n}$  given in the training dataset. Suppose that in addition to the true labels  $b_{n}$ , we have access to sufficiently well-trained and perhaps more complicated teacher model's outputs  $\Phi_{\theta}(a_n)\in \mathcal{B}$  for each input  $a_{n}\in \mathcal{A}$ . Similar to the student model  $\phi_{x}$ , the teacher model  $\Phi_{\theta}$  maps  $\mathcal{A}\rightarrow \mathcal{B}$  but can have different architecture, more layers and parameters. The fundamental question is how to exploit the additional knowledge of the teacher  $\Phi_{\theta}$  to facilitate the training of a more compact student model  $\phi_{x}$  with lower representational capacity. Knowledge Distillation with parameter  $\lambda \in [0,1]$  from teacher model  $\Phi_{\theta}$  to student model  $\phi_{x}$  is the following modification to the objective (1):
$$
\min  _ {x \in \mathbb {R} ^ {d}} \frac {1}{N} \sum_ {n = 1} ^ {N} \left[ (1 - \lambda) \ell \left(\phi_ {x} \left(a _ {n}\right), b _ {n}\right) + \lambda \ell \left(\phi_ {x} \left(a _ {n}\right), \Phi_ {\theta} \left(a _ {n}\right)\right) \right]. \tag {2}
$$
Here we customize the loss penalizing dissimilarities from the teacher's feedback  $\Phi_{\theta}(a_n)$  in addition to the true outputs  $b_{n}$ . In case of  $\ell$  is linear in the second argument (e.g., cross-entropy loss), the problem simplifies into
$$
\min  _ {x \in \mathbb {R} ^ {d}} \frac {1}{N} \sum_ {n = 1} ^ {N} \ell \left(\phi_ {x} \left(a _ {n}\right), (1 - \lambda) b _ {n} + \lambda \Phi_ {\theta} \left(a _ {n}\right)\right), \tag {3}
$$
which is a standard ERM (1) with modified "soft" labels  $s_n \coloneqq (1 - \lambda)b_n + \lambda \Phi_\theta(a_n)$  as the target.
3.2. Self-distillation. As already mentioned, the teacher's model  $\Phi_{\theta}$  can have more complicated neural network architecture and potentially larger parameter space  $\theta \in \mathcal{Q} \subseteq \mathbb{R}^{D}$ . In particular,  $\Phi_{\theta}$  does not have to be from the same set of models  $\mathcal{F}$  as the student model  $\phi_{x}$ . The special case when both the student and the teacher share the same structure/architecture is called self-distillation [33, 60], which is the key setup for our work. In this case, the teacher model  $\Phi_{\theta} \equiv \phi_{\theta} \in \mathcal{F}$  with  $\theta \in \mathbb{R}^{d}$  (i.e.,  $\mathcal{Q} = \mathcal{P}$ ,  $D = d$ ) and the corresponding distillation objective would be
$$
\min  _ {x \in \mathbb {R} ^ {d}} \frac {1}{N} \sum_ {n = 1} ^ {N} \left[ (1 - \lambda) \ell \left(\phi_ {x} \left(a _ {n}\right), b _ {n}\right) + \lambda \ell \left(\phi_ {x} \left(a _ {n}\right), \phi_ {\theta} \left(a _ {n}\right)\right) \right]. \tag {4}
$$
Algorithm 1 Knowledge Distillation via SGD
1: Input: learning rate  $\gamma > 0$ , initial student model  $x^0 \in \mathcal{P} \subseteq \mathbb{R}^d$   
2: for each distillation iteration  $m$  do  
3: choose a teacher model  $\theta^m \in \mathcal{Q} \subseteq \mathbb{R}^D$  and distillation weight  $\lambda^m \in [0,1]$  (e.g., see Sec. 5)  
4: for each training iteration  $t$  do  
5: sample an unbiased mini-batch  $\xi \sim \mathcal{D}$  form the train set  
6: compute distillation gradient  $\nabla f_{\xi}(x^t \mid \theta^m, \lambda^m) = \lambda^m\nabla f_{\xi}(x^t) + (1 - \lambda^m)\nabla f_{\xi}(x^t \mid \theta^m)$   
7: update the student model via  $x^{t+1} = x^t - \gamma\nabla f_{\xi}(x^t \mid \theta^m, \lambda^m)$   
8: end for  
9: end for
Our primary focus in this work would be the objective mentioned above of self-distillation. For convenience, let  $f_{n}(x) \coloneqq \ell(\phi_{x}(a_{n}), b_{n})$  be the prediction loss with respect to the output  $b_{n}$ ,  $f_{n}(x \mid \theta) \coloneqq \ell(\phi_{x}(a_{n}), \phi_{\theta}(a_{n}))$  be the loss with respect to the teacher's output probabilities and  $f_{n}(x \mid \theta, \lambda) \coloneqq \lambda f_{n}(x) + (1 - \lambda) f_{n}(x \mid \theta)$  be the distillation loss. See Algorithm 1 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 1 (Distillation Gradient). For a student model  $x \in \mathbb{R}^d$ , teacher model  $\theta \in \mathbb{R}^d$  and distillation weight  $\lambda$ , the distillation gradient corresponding to self-distillation (4) is given by
$$
\nabla_ {x} f _ {n} (x \mid \theta , \lambda) = \nabla f _ {n} (x) - \lambda \nabla f _ {n} (\theta). \tag {5}
$$
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 or no effect on the data points classified correctly with high confidence (i.e. those for which  $\nabla f_{n}(\theta)$  is close to 0 or  $\phi_{\theta}(a_n)$  is close to  $b_{n}$ ). 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  $\nabla f_{n}(\theta) = 0$  or  $\phi_{\theta}(a_n) = b_n$  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 not 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  $-\lambda \nabla f_{n}(\theta)$  term. It is worth mentioning that the direction of distillation gradient  $\nabla_{x}f_{n}(x\mid \theta ,\lambda)$  can be different from the usual gradient's direction  $\nabla f_{n}(x)$  due to the influence of the teacher. Thus, Proposition 1 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  $-\lambda \nabla f_{\xi}(\theta)$  term. When chosen properly (distillation weight  $\lambda$  and proximity of  $\theta$  to the 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.
- Linear regression. To support Proposition 1 rigorously, consider the simple setup of linear regression. Let  $\mathcal{A} = \mathbb{R}^d$ ,  $\mathcal{P} = \mathbb{R}^{d + 1}$ ,  $\phi_x(a) = x^\top \bar{a} \in \mathbb{R}$ , where  $\bar{a} = [a 1]^\top \in \mathbb{R}^{d + 1}$  is the input vector in the lifted space (to include the bias term),  $\mathcal{B} = \mathbb{R}$ , and the loss is defined by  $\ell(t, t') = (t - t')^2$  for all  $t, t' \in \mathbb{R}$ . Thus, based on (4), we have
$$
{f _ {n} (x \mid \theta , \lambda)} = {(1 - \lambda) (x ^ {\top} \bar {a} _ {n} - b _ {n}) ^ {2} + \lambda (x ^ {\top} \bar {a} _ {n} - \theta^ {\top} \bar {a} _ {n}) ^ {2},}
$$
from which we compute its gradient straightforwardly as
$$
\begin{array}{l} \nabla_ {x} f _ {n} (x \mid \theta , \lambda) = 2 (1 - \lambda) (x ^ {\top} \bar {a} _ {n} - b _ {n}) \bar {a} _ {n} + 2 \lambda (x ^ {\top} \bar {a} _ {n} - \theta^ {\top} \bar {a} _ {n}) \bar {a} _ {n} \\ { = } { 2 ( x ^ { \top } \bar { a } _ { n } - b _ { n } ) \bar { a } _ { n } - 2 \lambda ( \theta ^ { \top } \bar { a } _ { n } - b _ { n } ) \bar { a } _ { n } = \nabla f _ { n } ( x ) - \lambda \nabla f _ { n } ( \theta ) . } \\ \end{array}
$$
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.e.  $\phi_X(a_n) = \sigma(X^\top a_n) \in \mathbb{R}^K$ , where  $X = [x_1 x_2 \ldots x_K] \in \mathbb{R}^{d \times K}$  are the model parameters,
Figure 1: Cosine similarity,  $l_{2}$  distance and SNR (i.e.,  $l_{2}$  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  $\lambda$  leads to larger differences, although gradients remain well-correlated.
$a_{n}\in \mathcal{A} = \mathbb{R}^{d}$  is the input data and  $\sigma$  is the soft-max function. Then, we show in the Appendix B.2 that for all  $k = 1,2\ldots ,K$  it holds
$$
\nabla_ {x _ {k}} f _ {n} (X \mid \Theta , \lambda) = \nabla_ {x _ {k}} f _ {n} (X) - \lambda \nabla_ {\theta_ {k}} f _ {n} (\Theta),
$$
where  $\Theta = [\theta_{1}\theta_{2}\dots \theta_{K}]\in \mathbb{R}^{d\times K}$  are the teacher's parameters.
- Generic classification. Proposition 1 will not hold precisely for arbitrary deep non-linear neural networks. However, careful calculations reveal that, in general, distillation gradient takes a 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.  $\phi_x(a_n) = \sigma(\psi_n(x))$ , where  $a_n \in \mathcal{A}$  is the input data,  $\psi_n(x)$  are the produced logits with respect to the model parameters  $x$ , and  $\sigma$  is the soft-max function. Denote  $\varphi_n(z) \coloneqq \ell(\sigma(z), b_n)$  the loss associated with logits  $z$  and the true label  $b_n$ . In words,  $\psi_n$  gives the logits from the input data, while  $\varphi_n$  computes the loss from given logits. Then, the representation for the loss function is  $f_n(x) = \varphi_n(\psi_n(x))$ . We show in Appendix B.3 that the distillation gradient can be written as
$$
\nabla_ {x} f _ {n} (x \mid \theta , \lambda) = J \psi_ {n} (x) (\nabla \varphi_ {n} (\psi_ {n} (x)) - \lambda \nabla \varphi_ {n} (\psi_ {n} (\theta))) = \frac {\partial \psi_ {n} (x)}{\partial x} \frac {\partial f _ {n} (x)}{\partial \psi_ {n} (x)} - \lambda \frac {\partial \psi_ {n} (x)}{\partial x} \frac {\partial f _ {n} (\theta)}{\partial \psi_ {n} (\theta)},
$$
where  $J\psi_{n}(x)\coloneqq \frac{\partial\psi_{n}(x)}{\partial x} = [\nabla \psi_{n,1}(x)\nabla \psi_{n,2}(x)\dots \nabla \psi_{n,K}(x)]\in \mathbb{R}^{d\times K}$  is the Jacobian of the vector-valued function  $\psi_{n}$  for logits. Notice that the first term  $\frac{\partial\psi_n(x)}{\partial x}\frac{\partial f_n(x)}{\partial\psi_n(x)}$  coincides with the student's gradient  $\nabla_xf_n(x) = \frac{\partial f_n(x)}{\partial x}$ . However, the second term  $\frac{\partial\psi_n(x)}{\partial x}\frac{\partial f_n(\theta)}{\partial\psi_n(\theta)}$  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 a fully connected neural network with one hidden layer and ReLU activation [34], trained on the MNIST dataset [24], using regular self-distillation, from an SGD-teacher, with a 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 in Figure 1 show that the distillation gradient approximates well the true distillation gradient. Moreover, the behavior is monotonic in the distillation weight  $\lambda$  (higher similarity for smaller  $\lambda$ ), 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 a stochastic optimization problem of the form
$$
\min  _ {x \in \mathbb {R} ^ {d}} \left\{f (x) := \mathbb {E} _ {\xi \sim \mathcal {D}} \left[ f _ {\xi} (x) \right] \right\}, \tag {6}
$$
where  $f_{\xi}(x)$  is the loss associated with data sample  $\xi \sim \mathcal{D}$  given model parameters  $x \in \mathbb{R}^d$ . For instance, if  $\xi = (a_n, b_n)$  is a single data point, then the corresponding loss is  $f_{\xi}(x) = \ell(\phi_x(a_n), b_n)$ . The goal is to find parameters  $x$  minimizing the risk  $f(x)$ . To solve the problem (6), we employ Stochastic Gradient Descent (SGD). Based on Section 3, applying SGD to the problem (6) with self-distillation amounts to the following optimization updates in the parameter space: