SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Vaishnavh Nagarajan
vaishnavh@google.com
Aditya Krishna Menon
adityakmenon@google.com
Srinadh Bhojanapalli
bsrinadh@google.com
Hossein Mobahi
hmobahi@google.com
Sanjiv Kumar
sanjivk@google.com
Knowledge distillation (KD) has been widely used to improve the test accuracy of a "student" network, by training it to mimic the soft probabilities of a trained "teacher" network. Yet, it has been shown in recent work that, despite being trained to fit the teacher's probabilities, the student may not only significantly deviate from the teacher probabilities, but may also outdo than the teacher in performance. Our work aims to reconcile this seemingly paradoxical observation. Specifically, we characterize the precise nature of the student-teacher deviations, and argue how they can co-occur with better generalization. First, through experiments on image and language data, we identify that these probability deviations correspond to the student systematically exaggerating the confidence levels of the teacher. Next, we theoretically and empirically establish another form of exaggeration in some simple settings: KD exaggerates the implicit bias of gradient descent in converging faster along the top eigendirections of the data. Finally, we tie these two observations together: we demonstrate that the exaggerated bias of KD can simultaneously result in both (a) the exaggeration of confidence and (b) the improved generalization of the student, thus offering a resolution to the apparent paradox. Our analysis brings existing theory and practice closer by considering the role of gradient descent in KD and by demonstrating the exaggerated bias effect in both theoretical and empirical settings.
In knowledge distillation (KD) [6, 17], one trains a small "student" model to match the predicted soft label distribution of a large "teacher" model, rather than the one-hot labels that the training data originally came with. This has emerged as a highly effective model compression technique, and has inspired an actively developing literature that has sought to explore applications of distillation to various settings [41, 13, 52], design more effective variants [44, 3, 38, 5], and better understand theoretically when and why distillation is effective [32, 40, 35, 2, 8, 34, 10, 43, 25, 14, 39].
On paper, distillation is intended to help by transferring the soft probabilities of the (one-hot loss trained) teacher over to the student. Intuitively, we would desire this transfer to be perfect: the more a student fails to match the teacher's probabilities, the more we expect its performance to suffer. After all, in the extreme case of a student that simply outputs uniformly random labels, the student's performance would be as poor as it can get.
(a) Exaggeration of confidence
Figure 1: (a): Distilled student exaggerates confidence of one-hot-loss trained teacher. For each training sample  $(x,y)$ , we plot  $X = \phi(p_{y^{\mathrm{te}}}^{\mathrm{te}}(x))$  versus  $Y = \phi(p_{y^{\mathrm{te}}}^{\mathrm{st}}(x))$ , which are the teacher and student probabilities on the teacher's predicted label  $y^{\mathrm{te}}$ , transformed monotonically by  $\phi(u) = \log [u / (1 - u)]$ . Note that this is a density plot where higher the brightness, higher the number of datapoints with that  $X$  and  $Y$  value. We find that the distilled student predictions deviate from the  $X = Y$  line by either underfitting teacher's low confidence points (i.e., we find  $Y \leq X$  for small  $X$ ) and/or overfitting teacher's high confidence points (i.e.,  $Y \geq X$  for large  $X$ ). See §3 for details. (b) Distillation exaggerates implicit bias of one-hot gradient descent training. We consider an MLP trained on an MNIST-based dataset. Each plot shows the time-evolution of the  $\ell_2$  norm of the first layer parameters projected onto two randomly picked eigendirections; the  $\star$ 's correspond to the final parameters. First observe that the one-hot-trained teacher moves faster towards its final  $X$  axis value than its final  $Y$  axis value; this corroborates the well-known implicit bias of standard GD training. But crucially, we find that distillation exaggerates this bias: the student moves even faster towards its final  $X$  axis value. In §5 we argue how this exaggerated bias manifests as the exaggerated confidence in Fig 1a.
(b) Exaggeration of implicit bias
However, recent work by Stanton et al. [47] has challenged our presumptions underlying what distillation supposedly does, and how it supposedly helps. First, they show in practice that the student does not adequately match the teacher probabilities, despite being trained to fit them. Secondly, students that do successfully match the teacher probabilities, may generalize worse than students that show some level of deviations from the teacher [47, Fig. 1]. Surprisingly, this deviation occurs even in self-distillation settings [13, 54] where the student and the teacher have identical architecture and thus the student has the potential to fit the teacher's probabilities to full precision. Even more remarkably, the self-distilled student not only deviates from the teacher's probabilities, but also supercedes the teacher in performance.
How is it possible for the student to deviate from the teacher's probabilities, and yet counter-intuitively improve its generalization beyond the teacher? Our work aims to reconcile this paradoxical behavior. On a high level, our answer is that while arbitrary deviations in probabilities may indeed hurt the student, in reality there are certain systematic deviations in the student's probabilities. Next, we argue, these systematic deviations and improved generalization co-occur because they arise from the same effect: a (helpful) form of regularization that is induced by distillation. We describe these effects in more detail in our list of key contributions below:
(i) Exaggerated confidence: Across a wide range of architectures and image & language classification data (spanning more than 20 settings in total) we demonstrate  $(\S 3)$  that the student exaggerates the teacher's confidence. Most typically, on low-confidence points of the teacher, the student achieves even lower confidence than the teacher; in other settings, on high-confidence points, the student achieves even higher confidence than the teacher (Fig 1a). Surprisingly, we find such deviations even with self-distillation, implying that this cannot be explained by a mere student-teacher capacity mismatch. This reveals a systematic unexpected behavior of distillation.
(ii) Exaggerated implicit bias: Next, we demonstrate another form of exaggeration: in some simple settings, self-distillation exaggerates the implicit bias of gradient descent (GD) in converging faster along the top data eigendirections. We demonstrate this theoretically (Thm 4.1) for linear regression, as a gradient-descent counterpart to the seminal non-gradient-descent result of Mobahi et al. [35] (see §1.1 for key differences). Empirically, we provide the first demonstration of this effect for the cross entropy loss on neural networks (Fig 1b and §4.1).
(iii) Reconciling the paradox: Finally, we tie the above observations together to resolve our paradox. We empirically argue how the exaggerated bias towards top eigenvectors causes the student to both (a) exaggerate confidence levels and (b) outperform the teacher (see §5). This presents a resolution for how deviations in probabilities can co-occur with improved performance.
1.1 Bridging key gaps between theory and practice
The resolution above helps paint a more coherent picture of theoretical and empirical studies in distillation that were otherwise disjoint. Mobahi et al. [35] proved that distillation exaggerates the bias of a non-gradient-descent model, one that is picked from a Hilbert space with explicit  $\ell_2$  regularization. It was an open question as to whether this bias exaggeration effect is indeed relevant to settings we care about in practice. Our work establishes its relevance to practice, ultimately also drawing connections to the disjoint empirical work of Stanton et al. [47].
In more explicit terms, we establish relevance to practice in the following ways:
(1) We provide a formal proof of the exaggerated bias of distillation for a (linear) GD setting, rather than a non-GD setting (Theorem 4.1).
(2) We empirically verify the exaggerated bias of KD in more general settings e.g., a multi-layer perceptron (MLP) and a convolutional neural network (CNN) with cross-entropy loss (§4.1). This provides the first practical evidence of the bias exaggeration effect of [35].
(3) We relate the above bias to the student-teacher deviations in Stanton et al. [47]. Specifically, we argue that the exaggerated bias manifests as exaggerates student confidence levels, which we report on a wide range of image and language datasets.
(4) Tangentially, our findings also help clarify when to use early-stopping and loss-switching in distillation (§5.2).
As a more general takeaway for practitioners, our findings suggest that not matching the teacher probabilities exactly can be a good thing, provided the mismatch is not arbitrary. Future work may consider devising ways to explicitly induce careful deviations that further amplify the benefits of distillation e.g., by using confidence levels to reweight or scale the temperature on a per-instance basis.
Our interest in this paper is multiclass classification problems. This involves learning a classifier  $h\colon \mathcal{X}\to \mathcal{Y}$  which, for input  $\mathbf{x}\in \mathcal{X}$ , predicts the most likely label  $h(\mathbf{x})\in \mathcal{Y} = [K]\doteq \{1,2,\dots ,K\}$ . Such a classifier is typically implemented by computing logits  $\mathbf{f}\colon \mathcal{X}\rightarrow \mathbb{R}^{K}$  that score the plausibility of each label, and then computing  $h(\mathbf{x}) = \operatorname{argmax}_{y\in \mathcal{Y}}f_y(\mathbf{x})$ . In neural models, these logits are parameterised as  $\mathbf{f}(\mathbf{x}) = \mathbf{W}^{\top}\mathbf{Z}(\mathbf{x})$  for learned weights  $\mathbf{W}\in \mathbb{R}^{D\times K}$  and embeddings  $\mathbf{Z}(\mathbf{x})\in \mathbb{R}^{D}$ . One may learn such logits by minimising the empirical loss on a training sample  $S\doteq \{(x_{n},y_{n})\}_{n = 1}^{N}$ :
$$
R _ {\mathrm {e m p}} (\mathbf {f}) \doteq \frac {1}{N} \sum_ {n \in [ N ]} \mathbf {e} \left(y _ {n}\right) ^ {\top} \ell \left(\mathbf {f} \left(\mathbf {x} _ {n}\right)\right), \tag {1}
$$
where  $\mathbf{e}(y) \in \{0,1\}^K$  denotes the one-hot encoding of  $y$ ,  $\ell(\cdot) \doteq [\ell(1,\cdot), \ldots, \ell(K,\cdot)] \in \mathbb{R}^K$  denotes the loss vector of the predicted logits, and each  $\ell(y,\mathbf{f}(x))$  is the loss of predicting logits  $\mathbf{f}(x) \in \mathbb{R}^K$  when the true label is  $y \in [K]$ . Typically, we set  $\ell$  to be the softmax cross-entropy  $\ell(y,\mathbf{f}(\mathbf{x})) = -\log p_y(\mathbf{x})$ , where  $\mathbf{p}(\mathbf{x}) \propto \exp(\mathbf{f}(\mathbf{x}))$  is the softmax transformation of the logits.
Equation 1 guides the learner via one-hot targets  $\mathbf{e}(y_n)$  for each input. Distillation [6, 17] instead guides the learner via a target label distribution  $\mathbf{p}^{\mathrm{te}}(\mathbf{x}_n)$  provided by a teacher, which are the softmax probabilities from a distinct model trained on the same dataset. In this context, the learned model is referred to as a student, and the training objective is
$$
R _ {\text {d i s t}} (\mathbf {f}) \doteq \frac {1}{N} \sum_ {n \in [ N ]} \mathbf {p} ^ {\mathrm {t e}} \left(\mathbf {x} _ {n}\right) ^ {\top} \ell \left(\mathbf {f} \left(\mathbf {x} _ {n}\right)\right). \tag {2}
$$
One may also consider a weighted combination of  $R_{\mathrm{emp}}$  and  $R_{\mathrm{dist}}$ , but we focus on the above objective since we are interested in understanding each objective individually.
(a) Test data
Figure 2: Exaggeration of confidence in other settings. There are settings where even on test data, and for cross-architecture distillation settings, where the student exaggerates the teacher's confidence (here specifically on low-confidence points).
(b) Cross-architecture
Compared to training on  $R_{\mathrm{emp}}$ , distillation often results in improved performance for the student [17]. Typically, the teacher model is of higher capacity than the student model; the performance gains of the student may thus informally be attributed to the teacher transferring rich information about the problem to the student. In such settings, distillation may be seen as a form of model compression. Intriguingly, however, even when the teacher and student are of the same capacity (a setting known as self-distillation), one may see gains from distillation [13, 54]. The questions we explore in this paper are motivated by the self-distillation setting; however, for a well-rounded analysis, we empirically study both the self- and cross-architecture-distillation settings.
3 A fine-grained look at student-teacher probability deviations
To analyze student-teacher deviations, Stanton et al. [47] measured the disagreement and the KL divergence between the student and teacher probabilities, in expectation over all points. They found these quantities to be non-trivially large, contrary to the premise of distillation. To probe into the exact nature of these deviations, our idea is to study the per-sample relationship between the teacher and student probabilities.
Setup. Suppose we have teacher and distilled student models  $\mathbf{f}^{\mathrm{te}}$ ,  $\mathbf{f}^{\mathrm{st}}$ :  $\mathcal{X}\to \mathbb{R}^{K}$  respectively. We seek to analyze the deviations in the corresponding predicted probability vectors  $\mathbf{p}^{\mathrm{te}}(\mathbf{x})$  and  $\mathbf{p}^{\mathrm{st}}(\mathbf{x})$  for each  $(\mathbf{x},y)$  in the train and test set, rather than in the aggregated sense as in Stanton et al. [47]. To visualize the deviations, we need a scalar summary of these vectors. An initial obvious candidate is the probabilities assigned to the ground truth class  $y^{\star}$ , namely  $(p_{y^{\star}}^{\mathrm{te}}(\mathbf{x}),p_{y^{\star}}^{\mathrm{st}}(\mathbf{x}))$ . However, the student does not have access to the ground truth class, and is only trying to mimic the teacher. Hence, it is more meaningful and valuable to focus on the teacher's predicted class, which the student can infer i.e., the class  $y^{\mathrm{te}}\doteq \operatorname{argmax}_{y'\in [K]}p_{y'}^{\mathrm{te}}(\mathbf{x})$ . Thus, we examine the teacher-student probabilities on this label,  $(p_{y^{\mathrm{te}}}^{\mathrm{te}}(\mathbf{x}),p_{y^{\mathrm{te}}}^{\mathrm{st}}(\mathbf{x}))$ . Purely for visual clarity, we further perform a monotonic logit transformation  $\phi (u) = \log [u / (1 - u)]$  on these probabilities to produce real values in  $(- \infty , + \infty)$ . Thus, we compare  $\phi (p_{y^{\mathrm{te}}}^{\mathrm{te}}(\mathbf{x}))$  and  $\phi (p_{y^{\mathrm{te}}}^{\mathrm{st}}(\mathbf{x}))$  for each train and test sample  $(\mathbf{x},y)$ . For brevity, we refer to these values as confidence values for the rest of our discussion. It is worth noting that these confidence values possess another natural interpretation. For any probability vector  $\mathbf{p}(\mathbf{x})$  computed from the softmax of a logit vector  $\mathbf{f}(\mathbf{x})$ , we can write  $\phi (p_y(\mathbf{x}))$  in terms of the logits as  $f_{y}(\mathbf{x}) - \log \sum_{k\neq y}\exp (f_{k}(\mathbf{x}))$ . This can be interpreted as a notion of multi-class margin for class  $y$ .
To examine point-wise deviations of these confidence values, we consider scatter plots of  $\phi(p_{y^{\mathrm{te}}}^{\mathrm{te}}(\mathbf{x}))$  ( $X$ -axis) vs.  $\phi(p_{y^{\mathrm{te}}}^{\mathrm{st}}(\mathbf{x}))$  ( $Y$ -axis). We report this on the training set for some representative self-distillation settings in Figures 1a, cross-architecture distillation settings in Fig 2b and test set in Fig 2a. In all plots, the dashed line indicates the  $X = Y$  line. All values are computed at the end of training. The tasks considered include image classification benchmarks, namely CIFAR10, CIFAR-100 [26], Tiny-ImageNet [27], ImageNet [45] and text classification tasks from the GLUE benchmark (e.g., MNLI [51], AGNews [55]). See §C.1 for details on the experimental hyperparameters. The plots for all the settings (spanning more than 25 settings in total) are consolidated in §C.2.
Distilled students exaggerate one-hot trained teacher's confidence. If the objective of distillation were minimized perfectly, the scatter plots would simply lie on  $X = Y$  line. However, we find that across all our settings, the scatter shows stark deviations from  $X = Y$ . Importantly, these deviations reveal a characteristic pattern. In a vast majority of our settings, the confidence values of the student are an "exaggerated" version of the teacher's. This manifests in one of two ways. On the one hand, for points where the teacher attains low confidence, the student may attain even lower confidence — this is true for all our image settings, either on training or on test data. In the scatter plots, this can be inferred from the fact that for a large proportion of points where the  $X$  axis value is small (low teacher confidence), it is also the case that  $Y \leq X$  (even lower student confidence). As a second type of exaggeration, for points where the teacher attains high confidence, the student may attain even higher confidence — this is true for a majority of our language settings, barring some cross-architecture ones. To quantify these qualitative observations, in §C.4 we report the slope  $m$  of the best linear fit  $Y = mX + c$  over the low and high-confidence points and find that they correspond to  $m > 1$  in the corresponding scenarios above.
There are two reasons why these findings are particularly surprising. First, we see these deviations in both self-distillation settings (Fig 1a) and cross-architecture settings (see Fig 2b). It is remarkable that this should occur in self-distillation given that the student has the capacity to match the teacher probabilities. Next, these deviations can occur on both training and test data (Fig 2a). Here, it is surprising that there is deviation on training data, where the student is explicitly trained to match teacher probabilities. However, we must note that there are a few exceptions where it only weakly appears in training data e.g., CIFAR-10, but in those cases it is prominent on test data.
We also conduct ablation studies in §C.5 showing that this observation is robust to various hyperparameter changes (batch size, learning rate, length of training), and choices of visualization metrics. In §C.3, we explore further patterns underlying the student's underfit points. In §C.2, we discuss the exceptions where these characteristic deviations fail to appear even if deviations are stark e.g., in cross-architecture language settings.
Thus, our point-wise visualization of deviations clarifies that the student's mismatch of the teacher's probabilities in Stanton et al. [47] stems from a systematic exaggeration of the teacher's confidence. How do we reconcile this deviation with the student outperforming the teacher in self-distillation? In the next section, we turn our attention to a different type of exaggeration exerted by distillation that will help us resolve this question.
4 Distillation exaggerates implicit bias of GD
While the optimal solution to the KD loss (Eq 2) is for the student to replicate the teacher's probabilities, in practice we minimize this loss using gradient descent (GD). Thus, to understand why the student exaggerates the teacher's confidence in practice, it may be key to understand how GD interacts with the distillation loss. Indeed, in this section we analyze GD and formally demonstrate that, for linear regression, distillation exaggerates a pre-existing implicit bias in GD: the tendency to converge faster along the top eigendirections of the data. We will also empirically verify that our insight generalizes to neural networks with cross-entropy loss. In a later section, we will connect this exaggeration of bias back to the exaggeration of confidence observed in §3.
Concretely, we analyze gradient flow in a linear regression setting with early-stopping. Note that linear models have been used as a way to understand distillation even in prior work [40, 35]; and early-stopping is a typical design choice in distillation practice [12, 7, 23]. Consider an  $n \times p$  dataset  $\mathbf{X}$  (where  $n$  is the number of samples,  $p$  the number of parameters) with target labels  $\mathbf{y}$ . Assume that the Gram matrix  $\mathbf{XX}^{\top}$  is invertible. This setting includes overparameterized scenarios ( $p > n$ ) such as when  $\mathbf{X}$  corresponds to the linearized (NTK) features of neural networks [20, 28]. Then, a standard calculation spells out the weights learned at time  $t$  under GD on the loss  $(1/2) \cdot \| \mathbf{X}\boldsymbol{\beta} - \mathbf{y} \|^2$  as:
$$
\boldsymbol {\beta} (t) = \mathbf {X} ^ {\top} \left(\mathbf {X X} ^ {\top}\right) ^ {- 1} \mathbf {A} (t) \mathbf {y} \tag {3}
$$
$$
\text {w h e r e} \mathbf {A} (t) := \mathbf {I} - e ^ {- t \mathbf {X X} ^ {\top}}. \tag {4}
$$
Thus, the weights depend crucially on a time-dependent matrix  $\mathbf{A}$ . Intuitively,  $\mathbf{A}$  determines the convergence rate independently along each eigendirection: at any time  $t$ , it skews the weight assigned
to an eigendirection of eigenvalue  $\lambda$  by the value  $(1 - e^{-\lambda t})$ . As  $t \to \infty$  this factor increases all the way to 1 for all directions, thus indicating full convergence. But for any finite time  $t$ , the topmost direction would have a larger factor than the rest, implying a bias towards that direction. Our argument is that, distillation further exaggerates this implicit bias. To see why, consider that the teacher is trained to time  $T^{\mathrm{te}}$ . Through some calculation, the student's weights can be similarly expressed in closed-form, but with  $\mathbf{A}$  replaced by the product of two matrices:
$$
\tilde {\boldsymbol {\beta}} (\tilde {t}) = \mathbf {X} ^ {\top} \left(\mathbf {X} \mathbf {X} ^ {\top}\right) ^ {- 1} \tilde {\mathbf {A}} (\tilde {t}) \mathbf {y} \tag {5}
$$
$$
\text {w h e r e} \tilde {\mathbf {A}} (\tilde {t}) := \mathbf {A} (t) \mathbf {A} \left(T ^ {\mathrm {t e}}\right). \tag {6}
$$
One can then argue that the matrix  $\tilde{\mathbf{A}}$  corresponding to the student is more skewed towards the top eigenvectors than the teacher:
Theorem 4.1. (informal; see §B for full version and proof) Let  $\beta_{k}(t)$  and  $\tilde{\beta}_{k}(t)$  respectively denote the component of the teacher and student weights along the  $k$ 'th eigenvector of the Gram matrix  $\mathbf{XX}^{\top}$ , at any time  $t$ . Let  $k_{1} < k_{2}$  be two indices for which the eigenvalues satisfy  $\lambda_{k_1} > \lambda_{k_2}$ . Consider any time instants  $t$  and  $\tilde{t}$  at which both the teacher and the student have converged equally well along the top direction  $k_{1}$ , in that  $\beta_{k_1}(t) = \tilde{\beta}_{k_1}(\tilde{t})$ . Then along the bottom direction, the student has a strictly smaller component than the teacher, as in,
$$
\left| \frac {\tilde {\beta} _ {k _ {2}} (\tilde {t})}{\beta_ {k _ {2}} (t)} \right| <   1. \tag {7}
$$
The result says that the student relies less on the bottom eigendirections than the teacher, if we compare them at any instant when they have both converged equally well along the top eigendirections. In other words, while the teacher already has an implicit tendency to converge faster along the top eigendirections, the student has an even stronger tendency to do so. In the next section, we demonstrate that this insight generalizes to more practical non-linear settings.