SlowGuess/NeurIPS_2023_OCR huggingface dataset:
vasilis@cs.utexas.edu
Fotis Iliopoulos
fotisi@google.com
Khoa Trinh
khoatrinh@google.com
Cenk Baykal
baykalc@google.com
Gaurav Menghani
gmenghani@google.com
Erik Vee
erikvee@google.com
Knowledge distillation with unlabeled examples is a powerful training paradigm for generating compact and lightweight student models in applications where the amount of labeled data is limited but one has access to a large pool of unlabeled data. In this setting, a large teacher model generates "soft" pseudo-labels for the unlabeled dataset which are then used for training the student model. Despite its success in a wide variety of applications, a shortcoming of this approach is that the teacher's pseudo-labels are often noisy, leading to impaired student performance. In this paper, we present a principled method for knowledge distillation with unlabeled examples that we call Student-Label Mixing (SLaM) and we show that it consistently improves over prior approaches by evaluating it on several standard benchmarks. Finally, we show that SLaM comes with theoretical guarantees; along the way we give an algorithm improving the best-known sample complexity for learning halfspaces with margin under random classification noise, and provide the first convergence analysis for so-called "forward loss-adjustment" methods.
While good quality human-labeled data are often hard to obtain, finding huge amounts of unlabeled data is relatively easy. Therefore, in modern machine learning applications, we often face the situation where we have a small "golden" dataset with human labels and a large unlabeled dataset. In Distillation with Unlabeled Examples [12, 26, 16] a large teacher model is first trained (or fine-tuned) on the human-labeled data and is then used to generate "soft" pseudo-labels for the unlabeled dataset. Then the (typically smaller) student model, i.e., the model that will be deployed for the purposes of the application, is trained on the combined dataset that contains both the labels generated by humans and the pseudo-labels generated by the teacher model. This general-purpose training paradigm has been applied in a wide variety of contexts [16, 46, 53, 54, 58, 57] including but not limited to distilling knowledge from large-scale foundational models like BERT [18] and GPT-3 [11]. We remark that in such settings one does not have access to the teacher model but only on its pseudo-labels (which were generated during some previous "bulk-inference" phase). This "bulk-inference" step is typically computationally expensive and happens once: one cannot modify the teacher network (or even use it for inference) during the training process of student.
Despite its widespread success in practice, the effectiveness of this powerful approach generally depends on the quality of the pseudo-labels generated by the teacher model. Indeed, training the student model on noisy pseudo-labels often leads to significant degradation of its generalization
performance, and this is a well-known phenomenon that has been observed and studied in a plethora of papers in the literature, e.g., [6, 36, 44, 51, 53, 8, 27].
In this work, we propose Student-Label Mixing (SLaM), a principled method for knowledge distillation with unlabeled examples that accounts for the teacher's noise and consistently improves over prior approaches. At the heart of our method lies the observation that the noise introduced by the teacher is neither random nor adversarial, in the sense that it correlates well with metrics of "confidence" such as the margin score or the entropy of the teacher's predictions. We exploit this empirical fact to our benefit in order to introduce a model for the teacher's noise, which we use to appropriately modify the student's loss function. At a high level, for any given example during the student's training process, we evaluate the student's loss function on a convex combination of the student's current prediction and another (soft-)label that we estimate using our model for the teacher's noise (hence the name "student-label mixing").
1. We propose SLaM: a principled method for improving knowledge distillation with unlabeled examples. The method is efficient, data-agnostic and simple to implement.
2. We provide extensive experimental evidence and comparisons which show that our method consistently outperforms previous approaches on standard benchmarks. Moreover, we show that SLaM can be combined with standard distillation techniques such as temperature scaling and confidence-based weighting schemes.
3. We give theoretical guarantees for SLaM under standard assumptions. As a byproduct of our analysis we obtain a simple "forward loss-adjustment" iteration that provably learns halfspaces with  $\gamma$ -margin under Random Classification Noise with  $O(1 / (\epsilon^2\gamma^2))$  samples improving over prior works that had worse dependence on either the margin  $\gamma$  or the generalization error  $\epsilon$  (see Theorem 5.1 and Remark 5.2).
Knowledge Distillation. Most of the literature on knowledge distillation has been focused on the fully supervised/labeled setting, i.e., when distillation is performed on the labeled training data of the teacher model rather than on new, unlabeled data — see e.g. the original paper of [26]. Naturally, in this setting the pseudo-labels generated by the teacher are almost always accurate and so many follow-up works [2, 14, 15, 41, 52] have developed advanced distillation techniques that aim to enforce greater consistency between the teacher's and the student's predictions, or even between the intermediate representations learned by the two models. Applying such methods in our setting where the training dataset contains mainly unlabeled examples is still possible but, in this case, it is known [51, 27] that fully trusting the teacher model can be actually harmful to the student model, making these methods less effective. (In fact, when the teacher is highly noisy these methods even underperform vanilla distillation with unlabeled examples.) In Section 4.2 we present results that show the improved effectiveness of SLaM relative to the state-of-the-art supervised knowledge distillation methods like the Variational Information Distillation for Knowledge Transfer (VID) framework [2]. Moreover, in Appendix D.5 we show that our method can be combined with (i.e., provide an additional improvement) the most simple, yet surprisingly effective, methods of improving knowledge distillation, namely the temperature-scaling idea introduced by [26].
For distillation with unlabeled examples, many approaches [17, 33, 29] propose filtering-out or reweighting the teacher's pseudo-labels based on measures of teacher's uncertainty, such as dropout variance, entropy, margin-score, or the cut-statistic. These methods are independent of the student model and can be synergistically combined with our technique. For instance, in Section D.4 we demonstrate that combining our method with teacher-uncertainty-based reweighting schemes leads to improved student performance relative to applying the reweighting scheme alone.
Much more closely related to our approach is the recently introduced approach of [27]. There, the authors design a model for the teacher's noise and utilize it in order to modify the student's loss function so that, in expectation, the loss simulates the loss with respect to noise-free pseudo-labels. One of the main advantages of our method compared to that of [27] is that our model for the teacher's noise is more structured and easier to learn, which — as our experiments in Section 4.2 show — leads to consistently better student performance.
Learning From Noisy Labels. Learning from noisy labels is an important and well-studied problem with a vast literature [7, 21, 23, 28, 31, 37, 40, 42, 45, 47] — see [50] for a recent survey. The fundamental difference between our setting and papers in this literature is that the noise introduced by the teacher is structured, and this is a crucial observation we utilize in our design. Specifically, our approach is inspired by the so-called forward loss-adjustment methods, e.g. [43], but it is specifically tailored to the structure of the distillation with unlabeled examples setting. Indeed, forward methods typically attempt to estimate a noise transition matrix whose  $(i,j)$  entry is the probability of the true label  $i$  being flipped into a corrupted label  $j$ , which can be rather problematic when dealing with general, instance specific noise like in the case of distillation with unlabeled examples. On the other hand, we exploit that (i) we have access to confidence metrics of the teacher's predictions; and (ii) that often times, when the teacher model's top-1 prediction is inaccurate the true label is within its top- $k$  predictions for some appropriate  $k$ , to design and estimate a much more refined model for the teacher's noise that we use to inform the design of the student's loss function.
Another related technique for dealing with noisy data is using "robust" loss functions [4, 20, 24, 35, 56] such that they achieve a small risk for new clean examples even under the presence of noise in the training dataset. In Section 4.2 we compare our method with the general framework of [20] for designing robust loss functions and we show that our approach, when applied to the standard cross-entropy loss, consistently outperforms [20] in the setting of distillation with unlabeled examples. That said, we stress that our method is not tied to the cross-entropy loss and, in fact, it often gives better results when combined with more sophisticated loss functions. We demonstrate this in Appendix D.6 where we apply our method in cases where the student loss function comes from the families of losses introduced in [20] and [35].
Semi-Supervised Learning. Akin to our setting, in semi-supervised learning (SSL) (see e.g. [55] for a recent survey) the learner is presented with a small labeled dataset  $A$  and a typically much larger unlabeled dataset  $B$ . Unlike to our setting though, there is typically no distinction between the student and teacher: the model of interest generates pseudo-labels on  $B$  which are utilized by using appropriate loss functions or preprocessing procedures (e.g. "filtering" or "correcting") — often times in an iterative fashion with the goal of improving the quality of the newly-generated pseudo-labels. It is also worth noting that in many real-world applications of distillation with unlabeled examples either the teacher model is unavailable or it is too expensive to retrain it and create fresh pseudo-labels on the data (e.g., when we request labels from a pretrained large language model). Therefore, SSL approaches that either (i) update the "teacher" model (e.g., [34]), or (ii) require several fresh teacher-generated pseudo-labels (e.g., by requesting teacher-predictions on random data-augmentations or perturbed version of the unlabeled examples of  $B$  e.g., [9]) are not applicable in our setting. We implement the recent SSL technique of [48] and show that our method outperforms it in the context of distillation with unlabeled examples. Besides performing on par with state-of-the-art SSL approaches like [9], the method of [48] is free of inherent limitations like using domain-specific data augmentations — which is also an important feature of our approach.
Learning Halfspaces with Random Classification Noise. The theoretical study of classification with Random Classification Noise (RCN) was initiated by [5]. For the fundamental class of linear classifiers (halfspaces) the first polynomial time algorithms for the problem where given in [13] and [10]. The iteration proposed in [13] is a "backward loss-adjustment" method [43] for which it is known that resulting optimization landscape is convex (for linear classifiers). In [19] an improved analysis of the method of [13] was given, showing that SGD on this convex loss learns  $\gamma$ -margin halfspaces with RCN with  $\widetilde{O}(1 / (\gamma^4 \epsilon^2))$  samples. On the other hand, forward loss-adjustment methods for dealing with RCN are known to result in an inherently non-convex landscape, see [38] and Figure 9). Our theoretical result for SLaM (see Theorem 5.1) is the first convergence result for a "forward loss-adjustment" method and, at the same time, achieves a sample complexity of  $O(1 / (\gamma^2 \epsilon^2))$  improving over the prior work.
3 SLaM: Student-Label Mixing Distillation
In this section, we describe our distillation with unlabeled examples setting and present SLaM. In what follows, we assume that examples are represented by feature-vectors in some space  $\mathcal{X}$ . We shall denote by  $X$  the distribution over examples. We consider multi-class classification with  $L$  classes and assume that the ground-truth label of an example  $x$  is represented by a one-hot vector in  $\mathcal{Y} = \{0,1\}^L$  given by some unknown function  $g(x):\mathcal{X}\mapsto \mathcal{Y}$ . In multi-class classification the learning algorithm typically
optimizes a parametric family of classification models  $\mathcal{F} = \{f(\cdot ;w):\mathcal{X}\mapsto \mathbf{R}^L:w\in \mathcal{W}\}$ , i.e., for every parameter  $w\in \mathcal{W}$ ,  $f(x;w)$  is an  $L$ -dimensional "score vector", where  $f(x;w)_i$  corresponds to the probability that the model assigns to the class  $i$  for the example  $x$ . We shall denote by  $\ell (\cdot ,\cdot):\mathbf{R}^L\times \mathbf{R}^L\mapsto \mathbf{R}$  the classification loss function used by the learning algorithm. During training the algorithm considers a set of labeled examples  $S = \{(x^{(1)},g(x^{(1)})),\ldots ,(x^{(n)},g(x^{(n)}))\}$  and optimizes the loss  $\ell (\cdot ,\cdot)$  over  $S$ , i.e., solves the problem  $\min_{w\in \mathcal{W}}\frac{1}{|S|}\sum_{(x,g(x))\in S}\ell (g(x),f(x;w))$ . For two vectors  $v,u\in \mathbf{R}^L$  we denote by  $\mathrm{err}(v,u) = 1\{\mathrm{argmax}(v)\neq \mathrm{argmax}(u)\}$  the indicator of the event that the positions of the maximum elements of  $v,u$  agree. Similarly, for two classifiers  $h(x),f(x):\mathbf{R}^d\mapsto \mathbf{R}^L$  we can use  $\mathrm{err}(h(x),f(x))$  to denote whether their top-1 predictions for the example  $x$  agree. Our goal is to train a classifier over the sample  $S$  so that its generalization error, i.e.,  $\mathbf{E}_{x\sim X}[\mathrm{err}(f(x;w),g(x))]$ , is small.
Distillation with Unlabeled Examples. We assume that we are given a (usually small) dataset  $A$  of correctly labeled examples  $(x,g(x))$  and a set of unlabeled data  $U$ . A "teacher" model  $y_{s}(\cdot):\mathcal{X}\mapsto \mathbf{R}^{L}$  is first trained on the labeled dataset  $A$  and then provides soft-labels for the examples of dataset  $U$ , i.e., we create a dataset  $B = \{(x,y_s(x)):x\in U\}$  containing examples labeled with the corresponding probability distribution over classes (soft-labels) of the teacher model. We then train a (typically smaller) student model using both the original labeled data  $A$  and the teacher-labeled dataset  $B$ , i.e.,  $\min_{w\in \mathcal{W}}\frac{1}{|A\cup B|}\sum_{(x,z)\in A\cup B}\ell (z,f(x;w))$ . In what follows, we shall call the above training procedure as "vanilla-distillation".
Remark 3.1 ("Hard-" vs "Soft-" Distillation). We remark that the process where instead of using the soft-labels provided by the teacher model on the unlabeled dataset U, we use one-hot vectors representing the class with maximum score according to the teacher, is known as hard-distillation. We will denote by  $y_{s}(x)$  the soft-label of the teacher and by  $y(x)$  the corresponding hard-label, i.e.,  $y(x)$  is the one-hot representation of  $\operatorname{argmax} y_{s}(x)$ . When it is clear from the context we may simply write  $y$  instead of  $y(x)$ .
Modelling the Teacher as a "Noisy" Label Oracle. In the distillation setting described in the previous paragraph, it is known [51, 27, 8, 44] that the teacher model often generates incorrect predictions on the unlabeled examples, impairing the student's performance. Given any  $x \in U$ , we model the teacher's prediction  $y$  as a random variable. Similarly to [27] we assume that, for every unlabeled datapoint  $x \in U$ , the provided teacher label  $y$  is correct with probability  $\alpha(x)$  and incorrect with probability  $1 - \alpha(x)$ . However, in contrast with [27], our noise model prescribes a non-adversarial (semi-random) behavior of the teacher when its top-1 prediction is incorrect.
A first step towards more benign noisy teachers is to assume that, conditionally on being wrong, the teacher label is a uniformly random class of the remaining  $L - 1$  classes. We remark that this model is already enough to give improvements in datasets with a moderately large number of classes (e.g., up to 100). In particular, it perfectly captures the noisy teacher in binary classification: when the teacher label is different than the ground-truth  $g(x)$  then it has to be equal to the "flipped" ground-truth  $1 - g(x)$ .
We now further refine our model so that it is realistic for datasets with thousands of classes. Even though the top-1 accuracy of the teacher model may not be very high on the unlabeled data  $U$ , the true label is much more likely to belong in the top-5 or top-10 predictions of the teacher rather than being completely arbitrary. For example, training a ResNet50 network on  $10\%$  of ImageNet [49] yields an average top-1 accuracy about  $52.78\%$  on the test dataset whereas the top-10 accuracy of the same model is about  $83.55\%$ . In datasets with a large number of classes, this observation significantly reduces the number of potential correct classes of the examples where the teacher label is incorrect. Motivated by the above, we assume the following structured, semi-random noise model for the teacher, tailored to multi-class settings.
Definition 3.2 (Noisy Teacher Model). Let  $x$  be any example of the unlabeled data  $U$  and denote by  $g(x)$  its ground-truth label. Let  $y_{s}(x)$  resp.  $y(x)$  be the random variable corresponding to the soft resp. hard prediction of the teacher model for the example  $x$ . We assume that for every  $x$  there exist (unknown to the learner)  $\alpha(x) \in [0,1]$  and  $k(x) \in \{2,\dots,L\}$  such that the teacher's top-1 prediction  $y$  agrees with the ground-truth  $g(x)$  with probability  $\alpha(x)$  and, with probability  $1 - \alpha(x)$ : (i) the
ground-truth belongs in the top-  $k(x)$  predictions of the teacher; and (ii) the teacher's (hard)-prediction is a uniformly random incorrect class out of the top-  $k(x)$  predictions of the teacher soft-label  $y_{s}(x)^{1}$ .
Remark 3.3. We remark that the model of Definition 3.2 captures having a "perfect" teacher model by setting  $\alpha(x) = 1$  for all  $x$  and also generalizes the binary case described above by taking  $k(x) = 2$  for all  $x \in X$ .
Given the above noise model for the teacher, the problem of improving knowledge-distillation consists of two main tasks: (i) obtaining estimates for accuracy statistics  $\alpha(x), k(x)$  for each example  $x \in U$ ; and (ii) using those estimated values to improve the training of the student model so that it is affected less by the mistakes of the teacher on dataset  $B$ .
Training Better Students Using  $\alpha(x), k(x)$  We first assume that for every  $x$  we have oracle access to the values  $\alpha(x), k(x)$  and present our Student-Label Mixing loss function. Instead of using  $\alpha(x), k(x)$  to "denoise" the teacher's label, we use them to add noise to the student's predictions. To make notation more compact, in what follows, given a vector  $z \in \mathbf{R}^L$  we denote by  $\operatorname{top}(z; k)$  the vector that has the value 1 in the positions of the 1-st up to  $k$ -th largest elements of  $z$  and 0 in all other positions, e.g.,  $\operatorname{top}((1,2,3);1) = (0,0,1)$  and  $\operatorname{top}((-1,1,0,2);3) = (0,1,1,1)$ . Assuming that the student-label for some  $x \in U$  is  $f(x;w)$  we "mix" it (hence the name Student-Label Mixing) using  $\alpha(x), k(x)$  to obtain the mixed prediction
$$
\operatorname {m i x} (f (x; w); \alpha (x), k (x)) = \alpha (x) f (x; w) + (1 - \alpha (x)) \operatorname {t o p} \left(y _ {s} (x); k (x)\right) * \frac {1 - f (x ; w)}{k (x) - 1}, \tag {1}
$$
where  $q * p$  is the element-wise multiplication of the vectors  $p, q$ . We then train the mixed student model, on the "noisy" dataset  $B$ :
$$
\min  _ {w \in \mathcal {W}} \frac {1}{| A \cup B |} \left(\sum_ {(x, z) \in A} \ell (z, f (x; w)) + \sum_ {(x, y) \in B} \ell (y, \operatorname {m i x} (f (x; w); \alpha (x), k (x))\right) \tag {2}
$$
The main intuition behind the mixing of the student's labels is that by training the "noisy" student to match the "noisy" teacher label  $y$  on dataset  $B$ , the underlying (non-mixed) student  $f(x; w)$  will eventually learn the ground-truth. In particular, when  $\ell(\cdot, \cdot)$  is the Cross-Entropy loss we have that the expected mixed loss conditioned on any  $x$  is
$$
\mathbf {E} [ \ell (y; \operatorname {m i x} (f (x; w), a (x), k (x))) \mid x ] = \ell (\operatorname {m i x} (g (x); \alpha (x), k (x)), \operatorname {m i x} (f (x; w); \alpha (x), k (x))),
$$
where we used the fact that the cross-entropy is linear in its first argument, and that by the definition of our noise model (Definition 3.2) it holds that  $\mathbf{E}[y\mid x] = \mathrm{mix}(g(x);\alpha (x),k(x))$ . Therefore, when the student is equal to the ground-truth  $f(x;w) = g(x)$ , we obtain that the mixed student-model will satisfy  $\mathrm{mix}(g(x);\alpha (x),k(x)) = \mathrm{mix}(f(x;w);\alpha (x),k(x))$  for all  $x\in X$ , and (by Gibb's inequality), we obtain that  $g(x)$  is a minimizer of the SLaM loss. We show the following proposition, see Appendix C for the formal statement and proof.
Proposition 3.4 (SLaM Consistency (Informal)). Let  $D$  be the distribution of the teacher-labeled examples of dataset  $B$ , i.e., we first draw  $x \sim X$  and then label it using the noisy teacher of Definition 3.2. Moreover, assume that there exists some parameter  $w^{*} \in \mathcal{W}$  such that the ground-truth  $g(x) = f(x; w^{*})$ . Then  $w^{*}$  is the minimizer of the (population) SLaM objective:  $\min_{w} \mathbf{E}_{(x,y) \sim D}[\mathrm{ce}(y, \mathrm{mix}(f(x; w); \alpha(x), k(x))))$ , where  $\mathrm{ce}(\cdot, \cdot)$  is the Cross-Entropy loss.
Estimating the Teacher's Accuracy Statistics  $\alpha(x), k(x)$  via Isotonic Regression We first show how we estimate  $\alpha(x)$  for each  $x$  of dataset  $B$ , i.e., the dataset labeled by the teacher model. In [27] the authors empirically observed that  $\alpha(x)$  correlates with metrics of teacher's confidence such as the "margin", i.e., the difference between the probabilities assigned in the top-1 class and the second largest class according to the teacher's soft label  $y_s$ . In particular, the larger the margin is the more likely is that the corresponding teacher label is correct. We exploit this monotonicity by employing isotonic regression on a small validation dataset to learn the mapping from the teacher's margin at an example  $x$  to the corresponding teacher's accuracy  $\alpha(x)$ . For more details, see Appendix B.1.
To perform this regression task we use a small validation dataset  $V$  with correct labels that the teacher has not seen during training. For every example  $x \in V$  we compute the corresponding
Given that the teacher's prediction is incorrect and that the ground-truth belongs in the top- $k(x)$  predictions of the teacher, assumption (ii) describes a uniform distribution on  $k(x) - 1$  labels.
soft-teacher label  $y_{s}(x)$  and compute its margin margin(x) = max1(ys(x)) - max2(ys(x)). For every  $x\in V$  we also compute the hard-prediction of the teacher and compare it with the ground-truth, i.e., for every  $x\in V$  the covariate and response pair is (margin(x), 1 - err(g(x), y(x))). We then use isotonic regression to fit a piecewise constant, increasing function to the data. We remark that isotonic regression can be implemented very efficiently in  $O(n\log n)$  time (where  $n$  is the size of the validation dataset).