Weighted Distillation with Unlabeled Examples
Abstract
Distillation with unlabeled examples is a popular and powerful method for training deep neural networks in settings where the amount of labeled data is limited: A large “teacher” neural network is trained on the labeled data available, and then it is used to generate labels on an unlabeled dataset (typically much larger in size).
These labels are then utilized to train the smaller “student” model which will actually be deployed.
Naturally, the success of the approach depends on the quality of the teacher’s labels, since the student could be confused if trained on inaccurate data.
This paper proposes a principled approach for addressing this issue based on a “debiasing" reweighting of the student’s loss function tailored to the distillation training paradigm.
Our method is hyper-parameter free, data-agnostic, and simple to implement.
We demonstrate signiﬁcant improvements on popular academic datasets and we accompany our results with a theoretical analysis which rigorously justiﬁes the performance of our method in certain settings.
Introduction
In many modern applications of deep neural networks, where the amount of labeled examples for training is limited, distillation with unlabeled examples [8, 20] has been enormously successful.
In this two-stage training paradigm a larger and more sophisticated “teacher model” — typically non-deployable for the purposes of the application — is trained to learn from the limited amount of available training data in the ﬁrst stage.
In the second stage, the teacher model is used to generate labels on an unlabeled dataset, which is usually much larger in size than the original dataset the teacher itself was trained on.
These labels are then utilized to train the “student model”, namely the model which will actually be deployed.
Notably, distillation with unlabeled examples is the most commonly used training-paradigm in applications where one ﬁnetunes and distills from very large-scale foundational models such as BERT [13] and GPT-3 [5] and, additionally, it can be used to signiﬁcantly improve distillation on supervised examples only (see e.g. [47]).
While this has proven to be a very powerful approach in practice, its success depends on the quality of labels provided by the teacher model.
Indeed, often times the teacher model generates inaccurate labels on a non-negligible portion of the unlabeled dataset, confusing the student.
As deep neural networks are susceptible to overﬁtting to corrupted labels [50], training the student on the teacher’s noisy labels can lead to degradation in its generalization performance.
As an example, Figuredepicts an instance based on CIFAR-10 where ﬁltering out the teacher’s noisy labels is quite beneﬁcial for the student’s performance.
We address this shortcoming by introducing a natural “noise model” for the teacher which allows us to modify the student’s loss function in a principled way in order to produce an unbiased estimate of the student’s clean objective.
From a practical standpoint, this produces a fully “plug-and-play” method which adds minimal implementation overhead, and which is composable with essentially every other known distillation technique.
The idea of compressing a teacher model into a smaller student model by matching the predictions of the teacher was initially introduced by Buciluˇa, Caruana and Niculescu-Mizil [6] and, since then, variations of this method [20, 26, 31, 35, 39, 42, 49] have been applied in a wide variety of contexts [37, 48, 51].
(Notably, some of these applications go beyond compression — see e.g. [8, 47, 48, 52] for reference.)
In the simplest form of the method [26], and using classiﬁcation as a canonical example, the labels produced by the teacher are one-hot vectors that represent the class which has the maximum predicted probability — this method is often referred to as “hard”-distillation.
More generally, Hinton et.
al.
[8, 20] have shown that it is often beneﬁcial to train the student so that it minimizes the cross-entropy (or KL-divergence) with the probability distribution produced by the teacher while also potentially using a temperature higher than 1 in the softmax of both models (“soft”-distillation).
(Temperatures higher than 1 are used in order to emphasize the difference between the probabilities of the classes with lower likelihood of corresponding to the correct label according to the teacher model.)
The main contribution of this work is a principled method for improving distillation with unlabeled examples by reweighting the loss function of the student.
That is, we assign importance weights to the examples labeled by the teacher so that each weight reﬂects (i) how likely it is that the teacher has made an inaccurate prediction regarding the label of the example and (ii) how “distorted” the unweighted loss function we use to train the student is (measured with respect to using the groundtruth label for the example instead of the teacher’s label).
More concretely, our reweighting strategy is based on introducing and analyzing a certain noise model designed to capture the behavior of the teacher in distillation.
In this setting, we are able to come up with a closed-form solution for weights that “de-noise” the objective in the sense that (on expectation) they simulate having access to clean labels.
Crucially, we empirically observe that the key characteristics of our noise model for the teacher can be effectively estimated through a small validation dataset, since in practice the teacher’s noise is neither random nor adversarial, and it typically correlates well with its “conﬁdence” — see e.g. Figure 1.
In particular, we use the validation dataset to learn a map that takes as input the teacher’s and student’s conﬁdence for the label of a certain example and outputs estimates for the quantities mentioned in items (i) and (ii) above.
Finally, we plug in these estimates to our closed-form solution for the weights so that, overall, we obtain an automated way of computing the student’s reweighted objective in practice.
A detailed description of our method can be found in Section 2.3.
Our main ﬁndings and contributions can be summarized as follows:
• We propose a principled and hyperparameter-free reweighting method for knowledge distillation with unlabeled examples.
The method is efﬁcient, data-agnostic and simple to implement.
• We present extensive experimental results which show that our method provides signiﬁcant improvements when evaluated in standard benchmarks.
• Our reweighting technique comes with provable guarantees: (i) it is information-theoretically optimal; and (ii) under standard assumptions SGD optimization of the reweighted objective learns a solution with nearly optimal generalization.
1.1 Related work
Fidelity vs Generalization in knowledge distillation.
Conceptually, our work is related to the paper of Stanton et al.
[43] where the main message is that “good student accuracy does not imply good distillation ﬁdelity”, i.e., that more closely matching the teacher does not necessarily lead to better student generalization.
In particular, [43] demonstrates that when it comes to enlarging the distillation dataset beyond the teacher training data, there is a trade-off between optimization complexity and distillation data quality.
Our work can be seen as a principled way of improving this trade-off.
Advanced distillation techniques.
Since the original paper of Hinton et.
al.
[20], there have been several follow-up works [2, 7, 31, 45] which develop advanced distillation techniques which aim to enforce greater consistency between the teacher and the student (typically in the context of distillation on labeled examples).
These methods enforce consistency not only between the teacher’s predictions inaccurate harmful beneﬁcial instances the
Variational Information Distillation Knowledge Transfer how it is indeed beneﬁcial to combine it with our method in Section
We also show that our approach provides beneﬁts on top of any improvement obtained via temperature-scaling in Appendix B.2.
Learning with noisy data techniques.
As we have already discussed, the main conceptual contribution of our work is viewing the teacher model in distillation with unlabeled examples as the source of stochastic noise with certain characteristics.
Naturally, one may wonder what is the relationship between our method and works from the vast literature of learning with noisy data (e.g. [4, 14, 16, 23, 25, 27, 29, 33, 36, 38]).
The answer is that our method does not attempt to solve the generic “learning with noisy data” problem, which is a highly non-trivial task in the absence of assumptions (both in theory and in practice).
Instead, our contribution is to observe and exploit the fact that the noise introduced by the teacher in distillation has structure, as it correlates with several metrics of conﬁdence such as the margin-score, the entropy of the predictions etc.
We use and quantify this empirical observation to our advantage in order to formulate a principled method for developing a debiasing reweighting scheme (an approach inspired by the “learning with noisy data"-literature) which comes with theoretical guarantees.
An additional difference is that works from the learning with noisy data literature typically assume that the training dataset consists of corrupted hard labels, and often times their proposed method is not compatible with soft -distillation (in the sense that they degrade its performance, making it much less effective) — see e.g. [32] for a study of this phenomenon in the case of label smoothing [28, 44].
Uncertainty-based weighting schemes.
Related to our method are approaches for semi-supervised learning where examples are downweighted or ﬁltered out when the teacher is “uncertain.”
(A plethora of ways for deﬁning and measuring “uncertainty” have been proposed in the literature, e.g. entropy, margin-score, dropout variance etc.)
These methods are independent of the student model (they only depend on the teacher model), and so they cannot hope to remove the bias from the student’s loss function.
In fact, these methods can be viewed as preprocessing steps that can be combined with the student-training process we propose.
To demonstrate this we both combine and compare our method to the ﬁdelity-based weighting scheme of [12] in Section 3.3.
1.2 Organization of the paper.
In Section 2, we present our method in detail, while in Section 3 we present our key experimental results.
In Section 4, we discuss the theoretical aspects of our work.
In Section 5, we summarize our results, we discuss the beneﬁts and limitations of our method, and future work.
Finally, we present extended experimental and theoretical results in the Appendix.
Weighted distillation with unlabeled examples
In this section we present our method.
In Section 2.1, we review multiclass classiﬁcation.
In
Section 2.2, we introduce the noise model for the teacher in distillation, which motivates our approach.
And ﬁnally, in Section 2.3, we describe our method in detail.
2.1 Multiclass classiﬁcation
In multiclass classiﬁcation, we are given a training sample S =
{ ( x 1 , y 1 ) , ( x 2 , y 2 ) , . . .
( x n , y n ) } drawn from P n , where P is an unknown distribution over instances X and labels Y
=
[ L ] = { 1 , 2 , . . .
, L } .
Our goal is to learn a predictor f : X → R L , namely to minimize the risk of f .
The latter is deﬁned as the expected loss of f :
[Equation]
where ( x, y ) is drawn from P , and (cid:96) :
[ L ] × R L → R + is a loss function such that, for a label y ∈
[ L ] and prediction vector f ( x ) ∈ R L , (cid:96) ( y, f ( x )) is the loss incurred for predicting f ( x ) when the true label is y .
The most common way to approximate the risk of a predictor f is via the so-called empirical risk :
[Equation]
That is, given a hypothesis class of predictors F , our goal is typically to ﬁnd min f ∈F R S ( f ) as a way to estimate min f ∈F R ( f ) .
2.2 Debiasing weights
As we have already discussed, the main drawback of distillation with unlabeled examples is essentially that the empirical risk minimizer corresponding to the dataset labeled by the teacher cannot be trusted, since the teacher may generate inaccurate labels.
To quantify this phenomenon and guide our algorithmic design, we consider the following simple but natural noise model for the teacher.
Let X be an unknown distribution over instances X .
We assume the existence of a ground-truth classiﬁer so that each x ∈ X is associated with a ground-truth label f true ( x ) ∈
[ L ] = { 1 , 2 , . . .
, L } .
In other words, a clean labeled example is of the form ( x, f true ( x )) ∼ P (and x ∼ X ).
Additionally, we consider a stochastic adversary that given an instance x ∈ X , outputs a “corrupted” label y adv ( x ) with probability p ( x ) , and the ground-truth label f true ( x ) with probability 1 − p ( x ) .
Let D denote the induced adversarial distribution over instances and labels.
It is not hard to see that the empirical risk with respect to a predictor f and sample from D is not an unbiased estimator of the risk
[Equation]
— see Proposition 2.1.
On the other hand, the following weighted empirical risk (5) is indeed an unbiased estimator of R ( f ) .
For each x ∈ X let
[Equation]
and deﬁne
[Equation]
where S =
( x i , y i ) } ni
=1 ∼ D n .
Observe that the weight for each instance x depends on (i) how likely it is the adversary corrupts its label; and on (ii) how this corrupted label “distorts” the loss we observe at instance x .
In the following proposition we establish that the standard (unweighted) empirical risk with respect to distribution D and a predictor f is a biased estimator of the risk of f under the clean distribution P , while the weighted empirical risk (5) is an unbiased one.
Proposition 2.1 (Debiasing Weights) .
Let S = { ( x i , y i ) }
ni
=1 ∼ D n be a sample from the adversarial distribution.
Deﬁning Bias( f ) =
E x ∼ X
[ p ( x ) · (distortion f ( x ) − 1) · (cid:96) ( f true ( x ) , f ( x ))] we have: “reweighted” E [ R wS ( f )]
= R ( f ) vs “standard” E [ R S
( f )]
= R ( f ) + Bias( f ) .
Notice that, as expected, the bias of the unweighted predictor is a function of the “power” of the adversary, i.e., a function of how often they can corrupt the label of an instance, and the “distortion” this corruption causes to the loss we observe.
The proof of Proposition 2.1 follows from simple, direct calculations and it can be found in Appendix D.3.
Intuitively, given a sufﬁciently large sample S ∼ D n , optimizing an unbiased estimator for the risk should provide a better approximation for min f ∈ F R ( f ) compared to optimizing an estimator with constant bias.
We formalize this intuition in Section 4 and in Appendix D.
2.3 Our method
We consider the standard setting for distillation with unlabeled examples where we are given a dataset S (cid:96) =
{ ( x i , y i ) } mi =1 of m labeled examples from an unknown distribution P , and a dataset
S u = { x i } m + n i = m +1 of n unlabeled examples — typically, n ≥ m .
We also assume the existence of a (small) clean validation dataset S v = { ( x i , y i ) } i = m
+ n + q i = m
+ n +1 of size q .
Finally, let (cid:96) :
R L × R L → R + be a loss function that takes as input two vectors over the set of labels [ L ] .
We describe our method below and more formally in Algorithm 1 in Appendix A.
Remark 2.1.
The only assumption we need to make about the validation set S v is that it is not in the train set of the teacher model.
That is, set S v can be used in the train set of the student model if needed — we chose to present S v as a completely independent hold out set to make our presentation as conceptually clear as possible.
Training the teacher.
The teacher model is trained on dataset S (cid:96) , and then it is used to generate labels for the instances in S u .
The labels can be one-hot vectors or probability distributions on [ L ] , depending on whether we apply “hard” or “soft” distillation, respectively.
Training the student.
We start by pretraining the student model on dataset S (cid:96) .
Then, the idea is to think of the teacher model as the source of noise in the setting of Section 2.2, compute a weight w f ( x ) for each example x based on (4), and ﬁnally train the student on the union of labeled and teacher-labeled examples by minimizing the weighted empirical risk (examples from S (cid:96) are assigned unit weight).
We point out two remarks.
First, in order to apply (4) to compute the weight of an example x , we need to have estimates of p ( x ) and distortion f ( x ) .
To obtain these estimates we use the validation dataset S v as we describe in the next paragraph.
Second, observe that, according to (4), the weight of an example is a function of the predictor, namely the parameters of the model in the case of neural networks.
This means that ideally we should be updating our weights assignment every time the parameters of the student model get updated during training.
However, we empirically observe that even estimating the weight assignment only once during the whole training process (right after training the student model on S (cid:96) )
sufﬁces for our purposes, and so our method adds minimal overhead to the standard training process.
More generally, the fact that our process of computing the weights is simple and inexpensive allows us to recompute them during training (say every few epochs) to improve our approximation of the theoretically optimal weights.
We explore the effect of updating the weights during training in Section 3.2.
Estimating the weights.
We estimate the weights using the Nearest Neighbors method on S v to learn a map that takes as input the teacher’s and student’s “conﬁdence” for the label of a certain example x , and outputs estimates for p ( x ) and distortion f ( x ) so we can apply (4).
In our experiments we measure the conﬁdence of a model either via the margin-score , i.e., the difference between the largest two predicted classes for the label of a given example (see e.g. the so-called “margin” uncertainty sampling variant [40]), or via the entropy of its prediction.
However, one could use any metric (not necessarily conﬁdence) that correlates well with the accuracy of the corresponding models.
More concretely, we reduce the task of estimating the weights to a two-dimensional (i.e., two inputs and two outputs) regression task over the validation set which we solve using the Nearest Neighbors method.
In particular, our Nearest Neighbor data structure is constructed as follows.
Each example x of the validation set is assigned the two following pairs of points: (i) (teacher conﬁdence at x , student conﬁdence at x ) — this is the covariate of the regression task; (ii) ( 1 , distortion at ( x )) if the teacher correctly predicts the label of x , or ( 0 , distortion at x ), if the teacher does not correctly predict the label of x — this is the response of the regression task.
The query corresponding to an unlabeled example x (cid:48) is of the form (teacher conﬁdence at x (cid:48) , student conﬁdence at x (cid:48) ).
The Nearest Neighbors data structure returns the average response over the k closest in euclidean distance pairs (teacher conﬁdence at x , student conﬁdence at x ) in the validation set.
The value of k is speciﬁed in the next paragraph.
The pseudocode for our method can be found in Algorithm 2 in Appendix A.
We point out two remarks.
First, the number k of neighbors we use for our weights-estimation is always k = √ | S v | .
This is because choosing k = Θ( q 2 / (2+dim) ) , where q is the size of the validation dataset and dim is the dimension of the underlying metric space ( dim = 2 in our case), is asymptotically optimal, and 1 / 2 is a popular choice for the hidden constant used in practice, see e.g. [9, 18].
Second, notice that (4) implies that the weight of an example could be larger than 1 if (and only if) the corresponding distortion value (4) at that example is less than 1 .
This could happen for example if both the student and teacher have the same (or very similar) inaccurate prediction for a certain example.
In such a case, the value of the weight in (4) informs us that the loss at this example should be larger than the (low) value the unweighted loss function suggests.
However, since we do not have the ground-truth label for a point during training — but only the inaccurate prediction of the teacher — having a weight larger than 1 in this case would most likely guide our student model to ﬁt an inaccurate label.
To avoid this phenomenon, we always project our weights onto the [0 , 1] interval (Line 16 of Algorithm 2).
In Appendix D.2 we discuss an additional reason why it is beneﬁcial to project the weights of examples of low distortion onto [0 , 1] based on a MSE analysis.
Experimental results
In this section we present our experimental results.
In Section 3.1 we consider an experimental setup according to which the weights are estimated only once during the whole training process.
In
Section 3.2 we study the effect of updating the weights during training.
Finally, in Section 3.3 we demonstrate that our method can be combined with uncertainty-based weighting techniques by both combining and comparing our method to the ﬁdelity-based weighting scheme of [12].
[Figure]
3.1 Improvements via one-shot estimation of the weights
Here we show how applying our reweighting scheme provides consistent improvements on several popular benchmarks even when the weights are estimated only once during the whole training process.
We compare our method against conventional distillation with unlabeled examples, but we also show that our method can provide beneﬁts when combined with more advanced distillation techniques such as the framework of [2] (see Figure 9). [Figure]
More concretely, we evaluate our method on benchmark vision datasets.
We compare our method to conventional distillation with unlabeled examples both in terms of the best test accuracy achieved by the student, and in terms of convergence speed (see Figure 2).
We also evaluate the comparative advantage of our method as a function of the number of labeled examples available (size of dataset
S (cid:96) ).
We always choose the temperature in the softmax of the models to be 1 for simplicity and consistency, and our metric for conﬁdence is always the margin-score.
Implementation details for our experiments and additional results can be found in Appendices B and C. We implemented all algorithms in Python making use of the TensorFlow deep learning library [1].
We use 64 Cloud TPU v4s each with two cores.
3.1.1 Experimental setup
Our experiments are of the following form.
The academic dataset we use each time is ﬁrst split into two parts A and B .
Part A , which is typically smaller, is used as the labeled dataset S (cid:96) where the teacher model is trained on (recall the setting we described in Section 2.3).
Part B is randomly split again into two parts which represent the unlabeled dataset S u and validation dataset S v , respectively.
Then, (i) the teacher and student models are trained once on the labeled dataset S (cid:96) ; (ii) the teacher model is used to generate soft-labels for the unlabeled dataset S u ; (iii) we train the student model on the union of S (cid:96) and S u using our method and conventional distillation with unlabeled examples.
We repeat step (iii) a number of times: in each trial we partition part B randomly and independently, and then the student model is trained using the (student-)weights reached after completing the training on dataset S (cid:96) in step (i) as initialization, both for our method and the competing approaches.
3.1.2 CIFAR-{10, 100} and SVHN experiments
SVHN [34] is an image classiﬁcation dataset where the task is to classify street view numbers (classes).
The train set of SVHN contains 73257 labeled images and its test set contains 26032 images.
We use a MobileNet [21] with depth multiplier 2 as the teacher, and a MobileNet with depth multiplier 1 as the student 1 .
The tables in Figure 4 contain the results of our experiments (averages over 3 trials).
In each experiment we use the ﬁrst N ∈ { 7500 , 10000 , 12500 , 15000 , 17500 , 20000 } examples as the labeled dataset S (cid:96) , and then the rest 73257 − N images are randomly split to a labeled validation dataset S v of size 2000 , and an unlabeled dataset S u of size 71257 − N .
CIFAR-10 and CIFAR-100 [24] are image classiﬁcation datasets with 10 and 100 classes respectively.
Each of them consists of 60000 labeled images, which we split to a training set of 50000 images, and a test set of 10000 images.
For CIFAR-10, we use a Mobilenet with depth multiplier 2 as the teacher, and a Mobilenet with depth multiplier 1 as the student.
For CIFAR-100, we use a ResNet-110 as a teacher, and a ResNet-56 as the student.
We use a validation set of 2000 examples.
The results of our experiments (averages over 3 trials) can be found in the tables of Figures 5, 6. [Figure] [Figure] [Figure]
3.1.3 CelebA experiments: considering more advanced distillation techniques
As we have already discussed in Section 1.1, our method can be combined with more advanced distillation techniques which aim to enforce greater consistency between the teacher and the student.
We demonstrate this fact by implementing the method of Variational Information Distillation for
Knowledge Transfer (VID) [2] and showing how it is indeed beneﬁcial to combine it with our method.
We chose the gender binary classiﬁcation task of CelebA [15] as our benchmark (see details in the next paragraph), because it is known (see e.g. [31]) that the more advanced distillation techniques tend to be more effective when applied to classiﬁcation tasks with few classes.
In the table below, “Unweighted VID” corresponds to the implementation of loss described in equations (2), (4) and (6) of [2], and “Weighted VID” corresponds to the reweighting of the latter loss using our method.
CelebA [15] is a large-scale face attributes dataset with more than 200000 celebrity images, each with forty attribute annotations.
Here we consider the binary male/female classiﬁcation task.
The train set of CelebA contains 162770 images and its test set contains 19962 images.
We use a MobileNet with depth multiplier 2 as the teacher, and a ResNet-11 [19] as the student.
The tables in Figurecontain the results of our experiments (averages over 3 trials).
In each experiment we use the ﬁrst
N ∈ { 10000 , 15000 , 20000 , 25000 , 30000 , 35000 } examples as the labeled dataset S (cid:96) , and then the rest 162770 − N images are randomly split to a labeled validation dataset S v of size 2000 , and an unlabeled dataset S u of size 160770 − N .
[Figure] [Table]
3.1.4 ImageNet experiments
ImageNet [41] is a large-scale image classiﬁcation dataset with 1000 classes consisting of approximately I ≈ 1 .
3 M images.
For our experiments, we use a ResNet-50 as the teacher [Table] and a ResNet-18 as the student.
In each experiment we use the ﬁrst N ∈ { 64058 , 128116 } labeled examples ( 5% and 10% of I , respectively) as the labeled dataset S (cid:96)
[Table]
and the rest I − N examples are randomly split to a labeled validation dataset S v of size
[Table]
and an unlabeled dataset S u of size
I − N − 10000 .
The results of our experiments (averages over 10 trials) can be found in Figure 8.
[Figure]
3.2 Updating the weights during training
In this section we consider the effect of updating our estimation of the optimal weights during training.
For each dataset we consider, the experimental setup is identical to the corresponding setting in
Section 3.1, except for that we always use a validation set of size 500 and the entropy of a model’s prediction as the metric for its conﬁdence.
We note that the time required for computing the weight for each example is insigniﬁcant compared to total training time (less than 1% of the total training time in all of our experiments), which allows us to conduct experiments in which we update our estimation at the end of every epoch.
We see that doing this typically signiﬁcantly improves the resulting student’s performance (however, in CIFAR-100 we do not observe substantial beneﬁts).
[Figure]
3.3 Combining and comparing with uncertainty-based weighting schemes
As we have already discussed in Section 1.1, our method can be combined with uncertainty-based weighting schemes as these are independent of the student model and, therefore, they can be seen as a preprocessing step (modifying the loss function) before applying our method.
We demonstrate this by combining and comparing our method with the ﬁlelity-based weighting scheme of [12] on
CIFAR-10 and CIFAR-100.
Details about our experimental setup can be found in Appendix C.
[Figure]
Theoretical motivation
In this section we show that our debiasing reweighting method described in Section 2.2 comes with provable statistical and optimization guarantees.
We ﬁrst show that the method is statistically consistent in the sense that, with a sufﬁciently large dataset, the reweighted risk converges to the true or “clean” risk.
We show (see Appendix D)
the following convergence guarantee.
Theorem 4.1 (Uniform Convergence of Reweighted Risk) .
Assume that S is a dataset of i.i.d.
“noisy” samples from D .
Under standard capacity assumptions for the class of models F and regularity assumptions for the loss (cid:96) ( · ) , for every f ∈ F it holds that
[Equation]
To prove our optimization guarantees we analyze the reweighted objective in the fundamental case where the model f ( x ; Θ) is linear, i.e., f ( x ; Θ)
=
Θ x ∈ R L , and the loss (cid:96) ( y, z ) is convex in z for every y .
In this case, the composition of the loss and the model f ( x ; Θ) is convex as a function of the parameter Θ ∈ R L × d .
Recall that we denote by f true ( x ) : R d (cid:55)→ R L the ground truth classiﬁer and by P the “clean” distribution, i.e., a sample from P has the form ( x, f true ( x )) where x is drawn from a distribution X supported on (a subset of) R d .
Finally, we denote by D the “noisy” labeled distribution on R d × R L and assume that the x -marginal of D is also X .
We next give a general deﬁnition of debiasing weight functions, i.e., weighting mechanisms that make the corresponding objective function an unbiased estimator of the clean objective R (Θ) for every parameter vector Θ ∈ R d .
Recall that the weight function deﬁned in Section 2.2 is debiasing.
Deﬁnition 4.2 (Debiasing Weights) .
We say that a weight function w ( x, y adv ; Θ) is a debiasing weight function if it holds that
[Equation]
Since the loss is convex in Θ , one could try to optimize the naive objective that does not reweight and simply minimizes (cid:96) ( · ) over the noisy examples, R naive (Θ) (cid:44) E ( x,y adv ) ∼ D [ (cid:96) ( y adv , Θ x )] .
We show (unsurprisingly) that doing this is a bad idea: there are instances where optimizing the naive objective produces classiﬁers with bad generalization error over clean examples.
For the formal statement and proof we refer the reader to Appendix E.
SGD on the naive objective R naive ( · ) learns parameters with arbitrarily bad generalization error over the “clean” data.
Our main theoretical insight is that optimizing linear models with the reweighted loss leads to parameters with almost optimal generalization.
Given a debiasing weight function w ( · ) , SGD on the reweighted objective R w ( · ) learns a parameter with almost optimal generalization error over the “clean” data.
The main issue with optimizing the reweighted objective is that, in general, we have no guarantees that the weight function preserves its convexity (recall that it depends on the parameter Θ ).
However, we know that its population version corresponds to the clean objective R ( · ) which is a convex objective.
We show that we can use the convexity of the underlying clean objective to show results for stochastic gradient descent, by proving the following key structural property.
Proposition 4.3 (Stationary Points of the Reweighted Objective Sufﬁce (Informal)) .
Let S be a dataset of n = poly( dL/(cid:15) )
i.i.d.
samples from the noisy distribution D .
Let (cid:98) Θ be any stationary point of the weighted objective R wS (Θ) constrained on the unit ball (with respect to the Frobenious norm (cid:107) · (cid:107) F ).
Then, with probability at least 99% , it holds that
[Equation]
Conclusion
We propose a principled reweighting scheme for distillation with unlabeled examples.
Our method is hyper-parameter free, adds minimal implementation overhead, and comes with theoretical guarantees.
We evaluated our method on standard benchmarks and we showed that it consistently provides signiﬁcant improvements.
We note that investigating improved data-driven ways of estimating the weights (4) could be of interest, since potential inaccurate estimation of the weights is the main limitation of our work.
We leave this question open for future work.
A Formal description of our method
In this section we present pseudocode for our method.
Require: model Teacher , model Student , labeled dataset S (cid:96) =
{ ( x i , y i ) } mi =1 , unlabeled dataset
S u = { x i } m + n i = m +1 , validation dataset S v = { ( x i , y i ) } m + n + q i = m
+ n +1 , number of weights-estimating iterations r 1: Train Teacher and Student on S (cid:96) 2: Use Teacher to generate labels for S u to obtain set S (cid:96)u = { ( x, Teacher( x )) | x ∈ S u } 3: for i = m
+ 1 to n + m do 4: y
i ← Teacher( x i ) 5: S ← { ( x i , y i ) }
m
+ n i = 6: for i
= 1 to m do 7: w ( x i ) ←8:
for i = 1 to r do 9: { w ( x m +1 ) , . . .
, w ( x m + n ) } ← E STIMATE W EIGHTS ( Teacher , Student , S v , S (cid:96)u ) 10:
Train Student on S using the weighted empirical risk: m + n m + n (cid:88) i =
w ( x i ) (cid:96) ( y i , Student( x i ))
1: procedure E STIMATE W EIGHTS ( Teacher , Student , V, D ) 2: (cid:46) V is the validation dataset and D is the teacher-labeled dataset 3:
U ← ∅ , k ← (cid:100) 12 (cid:112) | V |(cid:101) 4: for every ( x, y ) ∈ V do 5:
X ← (Conﬁdence(Teacher( x )) , Conﬁdence(Student( x ))) 6: if arg max(Teacher( x ))
=
arg max( y ) then : 7: ( p, distortion) ← (0 , 1) 8: else : 9: ( p, distortion) ← (cid:16) 1 , (cid:96) (Teacher( x ) , Student( x ))
(cid:96) ( y, Student( x )) (cid:17) 10:
Y ← ( p, distortion) 11:
U ← U ∪ { ( X, Y ) } 12:
Weights = ∅ (cid:46) Initialize and empty list for the weights 13: for every ( x, y ) ∈ D do 14:
Query ← (Conﬁdence(Teacher( x )) , Conﬁdence(Student( x ))) 15: (ˆ p, ˆ d ) ← k - NN
( U, Query) (cid:46) Predict p ( x ) and distortion f ( x ) from the k nearest neighbors of Query in U 16: w ( x ) ← min (cid:110) 1 , 1 1+ˆ p (ˆ d − 1) (cid:111) 17:
Append w ( x ) to Weights 18: return Weights
B Extended experiments B.1 The student’s test-accuracy-trajectory
In this section we provide extended experimental results that show the student’s test accuracy over the training trajectory corresponding to experiments we mentioned in Section 3.1.
Notice that in the vast majority of cases our method signiﬁcantly outperforms the conventional approach almost throughout the training process.
B.2 Considering the effect of temperature
Temperature-scaling, a technique introduced in the original paper of Hinton et.
al.
[20], is one the most common ways for improving student’s performance in distillation.
Indeed, it is known (see e.g. [43]) that choosing the right value for the temperature can be quite beneﬁcial, to the point it can outperform other more advanced techniques for improving distillation.
Here we demonstrate that our approach provides beneﬁts on top of any improvement on can get via temperature-scaling by conducting an ablation study on the effect of temperature on CIFAR-100.
In our experiment, the teacher model is a Resnet-110 achieving accuracy 56 . 0% , the student model is a Resnet-56, the number of labeled examples is 12500 , the validation set consists of 500 examples, and we use the entropy of a prediction as a metric of conﬁdence.
We apply our method using one-shot estimation of the weights.
We compare training the student model using conventional distillation to using our method for different values of temperature.
The results can be found in the table below.
We see that in almost all cases the student-model trained using our method outperforms the student-model trained using conventional distillation and, in particular, the best student overall is the result of choosing 2 .for
the value of temperature and applying our method.
[Figure]
See Appendix B.2 for details
C Implementation details
In this section we describe the implementation details of our experiments.
Recall the description of our method in Section 2.3.
C.1 Experiments on CelebA, CIFAR-10, CIFAR-100, SVHN
All of our experiments are performed according to the following recipe.
In all cases, the loss function (cid:96) :
R L × R L → R + we use is the cross-entropy loss.
We train the teacher model for 200 epochs on dataset S (cid:96) .
We pretrain the student model for 200 epochs on dataset S (cid:96) and save its parameters.
Then, using the latter saved parameters for initialization each time, we train the student model for 200 epochs optimizing either the weighted or conventional (unweighted) empirical risk, and report its average performance over three trials.
We use the Adam optimizer.
The initial learning rate is lr = 0 .
We proceed according to the following learning rate schedule (see e.g., [19]):
[Equation]
Finally, we use data-augmentation.
In particular, we use random horizontal ﬂipping and random width and height translations with width and height factor, respectively, equal to 0 .
C.2 Experiments on ImageNet
For the ImageNet experiments we follow a similar although not identical recipe to the one described in Appendix C.1.
In each training stage above, we train the model (teacher or student) forepochs instead of 200 .
We also use SGD with momentum 0 .
9 instead of Adam as the optimizer.
For data-augmentation we use only random horizontal ﬂipping.
Finally, the learning rate schedule is as follows.
For the ﬁrst 5 epochs the learning rate lr is increased from 0 . 0 to 0 .
1 linearly.
After that, the learning rate changes as follows: lr =   0 .
01 , if #epochs >0 .
001 , if #epochs >0 . 0001 , if #epochs >
C.3 Details on the experimental setup of Section 3.
In the CIFAR-10 experiments of Section C.3 the teacher model is a MobileNet with depth multiplier 2, and the student model is a MobileNet with depth multiplier 1.
In the CIFAR-100 experiments of the same section, the teacher model is a ResNet-110, and the student model is a ResNet-56.
We use a validation set consisting of 500 examples (randomly chosen as always).
The student of each method has access to the same number of labeled examples, i.e., the validation set is used for training the student model as we describe in Remark 2.1.
We compare the following three methods:
• Fidelity weighting scheme [12].
For every example x we use the entropy of the teacher’s prediction as an uncertainty/conﬁdence measure, which we denote by entropy( x ) .
We then compute the exponential weights described in [12] as w ( x ) = exp( − entropy( x ) / entropy) , where entropy is the average entropy of the teacher’s predictions over all training examples.
• Our method.
We use the entropy as the metric for conﬁdence.
In the case of CIFAR-10 we re-estimate the weights at the end of every epoch.
In the case of CIFAR-100 the weights are estimated only once in the beginning of the process.
• Composition.
We reweight each example in the loss function by multiplying the weights resulting from the two methods above.
D Extended theoretical motivation: statistical aspects
In this section we study the statistical aspects of our approach.
In Section D.1 we revisit and formally state Theorem 4.1 — see Corollary D.2 and Remark D.1.
In Section D.2 we perform a Mean-SquaredError analysis that provides additional justiﬁcation of our choice to always project the weights on the [0 , 1] interval (recall Line 16 of Algorithm 2).
Finally, in Section D.3 we provide the proof of
Proposition 2.1 which was omitted from the main body of the paper.
D.1 Statistical motivation
Recall the background on multiclass classiﬁcation in Section 2.1.
In this section we study hypothesis classes F and loss functions (cid:96) :
R L × R L : → R that are “well-behaved” with respect to a certain (standard in the machine learning literature) complexity measure we describe below.
For (cid:15) > 0 , a class H of functions h : X → [0 , 1] and an integer n , the “growth function" N ∞ ( (cid:15), H , n ) is deﬁned as
[Equation]
where H ( x ) = { ( h ( x 1 ) , . . .
, h ( x n )) :
h ∈ H} ⊆ R n and for A ⊆ R n the number N ( (cid:15), A, (cid:107) · (cid:107) ∞ ) is the smallest cardinality
A 0 of a set
A 0 ⊆
A such that A is contained in the union of (cid:15) -balls centered at points in A 0 , in the metric induced by (cid:107) · (cid:107)
∞
The growth number is a complexity measure of function classes commonly used in the machine learning literature [3, 17].
The following theorem from [30] provides large deviation bounds for function classes of polynomial growth.
Theorem D.1 (Theorem 6, [30]) .
Let Z be a random variable taking values in Z distributed according to distribution µ , and let H : Z → [0 , 1] be a class of functions.
Fix δ ∈ (0 , 1) , n ≥and set
M ( n ) = 10 N ∞ (1 /n, H , 2 n ) .
Then with probability at least 1 − δ in the random vector Z =
( Z 1 , . . .
, Z n ) ∼ µ n , for every h ∈ H we have:
[Equation]
where V n ( h, Z ) is the sample variance of the sequence { h ( Z i ) } ni =1 .
A straightforward corollary of Theorem D.1 and Proposition 2.1, and the main motivation for our method, is the following corollary.
Corollary D.2.
Let (cid:96) :
R L × R L → [0 , 1] be a loss function and ﬁx δ > 0 .
Consider any hypothesis class F of predictors f : X → R L , and the two induced classes H ⊆
[0 , 1] R L × R L , H w ⊆ [0 , 1] R L × R L of functions
h f ( x, y ) :
= (cid:96) ( y, f ( x )) and h wf ( x, y ) :
=
w f ( x ) (cid:96) ( y, f ( x )) , respectively.
Fix δ > 0 , n ≥ 16 , and set M ( n ) = 10 N ∞ (1 /n, H , 2 n ) and M w ( n ) = 10 N ∞ (1 /n, H w , 2 n ) .
Then, with probability at least 1 − δ over S
=
{ x i , y i } ni
=1 ∼ D n ,
[Equation]
where V S ( f ) , V wS ( f ) are the sample variances of the loss values { h f ( x i , y i ) } ni =1 , { h wf ( x i , y i ) } ni =1 , respectively.
The following remark formally captures Theorem 4.1.
Remark D.1.
Under the assumptions of Corollary D.2, if we additionally have that M ( n ) and
M w ( n ) are polynomially bounded in n , then, for every f ∈ F it holds that
[Equation]
D.2 Studying the MSE of a ﬁxed prediction
In this section we study the Mean-Squared-Error (MSE) of a ﬁxed prediction f ( x ) for an arbitrary instance x ∈ X , predictor f , and loss function (cid:96) :
R L × R L → R + , in order to gain some understanding on when the importance weighting scheme could potentially underperform the standard unweighted approach from a bias-variance perspective (i.e., when the training sample is “small enough” so that asymptotic considerations are ill-suited).
These considerations lead us to an additional justiﬁcation for always projecting the weights to the [0 , 1] interval (recall Line 16 of Algorithm 2).
Formally, we study the behavior of the quantities:
[Equation]
Recalling the deﬁnition of distortion (4) we have the following proposition.
Proposition D.3.
Let (cid:96) :
R L × R L → R + be a bounded loss function.
Fix x ∈ X and a predictor f : X → R L .
We have MSE( x ) <
MSE w ( x ) if and only if:
distortion f ( x ) < 1 / 2 ; and 2.
p ( x ) ∈ (cid:16) 0 , 1 − 2 · distortion f ( x ) (1 − distortion f ( x )) 2 (cid:17) .
Proof sketch.
Via direct calculations we obtain:
[Equation]
and
[Equation]
Recalling the deﬁnition of weights (4), and combining it with (9) and (10) implies the claim.
In words, Proposition D.3 implies that when the adversary does not have the power to corrupt the label of an instance x with high enough probability, i.e., p ( x ) is sufﬁciently small, and the prediction of the student is “close enough" to the adversarial label (i.e.,
when distortion f ( x ) is small enough), then it potentially makes sense to use the unweighted estimator instead of the weighted one from a bias-variance trade-off perspective, as the former has smaller MSE in this case.
Notice that this observation aligns well with our method as we always project w f ( x ) to [0 , 1] (observe that w f ( x ) >
iff distortion f ( x ) < 1 and p ( x ) > 0 ).
D.3 Proof of Proposition 2.
Recall the weight, distortion and bias deﬁnitions in (4) and Proposition 2.1.
We prove the ﬁrst claim of Proposition 2.1 via direct calculations:
[Equation]
x
[Equation]
x
[Equation]
= R ( f ) + Bias( f ) .
Similarly for the second claim:
[Equation]
concluding the proof.
E Extended theoretical motivation: optimization aspects
To prove our optimization guarantees, we analyze the reweighted objective in the fundamental case where the model f ( x ; Θ) is linear, i.e., f ( x ; Θ) =
Θ x ∈ R L , and the loss (cid:96) ( y, z ) is convex in z for every y .
In this case, the composition of the loss and the model f ( x ; Θ) is convex as a function of the parameter Θ ∈ R L × d .
Recall that we denote by f true ( x ) : R d (cid:55)→ R L the ground truth classiﬁer and by P the “clean” distribution, i.e., a sample from P has the form ( x, f true ( x )) where x is drawn from a distribution X supported on (a subset of) R d .
Finally, we denote by D the “noisy” labeled distribution on R d × R L and assume that the x -marginal of D is also X .
Notation In what follows, for any elements r, q of the same dimensions we denote by r · q their inner product.
For example for two vectors r, q ∈ R d we have r · q = (cid:80) di
=1 r i q i .
Similarly, for two matrices Θ , Q ∈ R L × d we have Θ · Q = (cid:80) Li
=1 (cid:80) dj
=1 Θ ij Q ij .
We denote by (cid:107) · (cid:107) 2 the (cid:96) 2 for vectors and the spectral norm for matrices.
We use ⊗ to denote the standard tensor (Kronecker) product between two vectors or matrices.
For example, for two matrices A, B we have ( A ⊗ B ) ijkl =
A ij B kl and for two vectors v, u we have ( v ⊗ u ) ij
=
v
i
u j .
We denote by (cid:107) · (cid:107) F the Frobenious norm for matrices.
We remark that we use standard asymptotic notation O ( · ) , etc.
and (cid:101) O ( · ) to omit factors that are poly-logarithmic (in the appearing arguments).
For example, training a linear model f ( x ; Θ) =
Θ x with the Cross Entropy loss corresponds to using (cid:96) ( t, y )
=
(cid:80) Li
=1 t i log( e yi (cid:80)
Lj =1 e yj ) and minimizing the objective
[Equation]
More generally, in what follows we shall refer to the population loss over the clean distribution P as
R ( · ) , i.e.,
[Equation]
We next give a general deﬁnition of debiasing weight functions, i.e., weighting mechanisms that make the corresponding objective function an unbiased estimator of the clean objective R (Θ) for every parameter vector Θ ∈ R d .
Deﬁnition E.1 (Debiasing Weights) .
We say that a weight function w ( x, y adv ; Θ) : R d × R L (cid:55)→ R is a debiasing weight function if it holds that
[Equation]
Remark E.1.
We remark that the weight function w ( · ) depends on the current hypothesis, Θ , and also on the noise advice p ( x ) that we are given with every example.
In order to keep the notation simple, we do not explicitly track these dependencies and simply write w ( x, y adv ; Θ) .
We also remark that, in general, in order to construct the weight function w we may also use “clean” data, which may be available, e.g., as a validation dataset, as we did in Section 2.2.
Our main result is that, given a convex loss (cid:96) ( · ) and a debiasing weight function w ( · ) that satisfy standard regularity assumptions, stochastic gradient descent on the reweighted objective produces a parameter vector with good generalization error.
We ﬁrst present the assumptions on the example distributions, the loss, and the weight function.
In what follows, we view the gradient of a function q (Θ) :
R L × d (cid:55)→ R as an L × d -matrix and the hessian ∇ 2 q (Θ) as an ( L × d ) × ( L × d ) -tensor (or equivalently as a dL × dL -matrix).
Deﬁnition E.2 (Regularity Assumptions) .
The x -marginal X of D and P is supported on (a subset of) the ball of radius R > 0 , B R (cid:44) { x ∈ R
d : (cid:107) x (cid:107) 2 ≤ R } .
The training model is linear f ( x ; Θ) =
Θ x and the parameter space is the unit ball, i.e., (cid:107) Θ (cid:107) F ≤ 1 .
For every label y adv ∈ R L in the support of D , the loss z (cid:55)→ (cid:96) ( y adv , z ) is a twice differentiable, convex function in z .
Moreover (cid:96) ( y adv , z ) is M (cid:96) -bounded, L (cid:96) -Lipschitz, and B (cid:96) -smooth, i.e., | (cid:96) ( y adv , z ) | ≤
M (cid:96) , (cid:107)∇ z (cid:96) ( y adv , z ) (cid:107) 2 ≤ L (cid:96) , and (cid:107)∇ 2 z (cid:96) ( y adv , z ) (cid:107) 2 ≤ B (cid:96) , for all z with (cid:107) z (cid:107) 2 ≤ R .
For every example ( x, y adv ) ∈ R d × R L in the support of D the weight function Θ (cid:55)→ w ( x, y adv ; Θ) is twice differentiable, M w -bounded, L w -Lipschitz, and B w -smooth, i.e., | w ( x, y adv ; Θ) | ≤ M w , (cid:107)∇ Θ w ( x, y adv ; Θ) (cid:107) F ≤ L w , and (cid:107)∇ 2 Θ w ( x, y adv ; Θ) (cid:107) 2 ≤ B w for all Θ
Remark E.2.
Observe that if a property in the above deﬁnition is satisﬁed by some parameter-value
Q , then it is also satisﬁed for any other Q (cid:48) > Q .
For example, if the loss function is 0 .
5 -Lipschitz it is also 1 -Lipschitz.
Therefore, to simplify the expressions, in what follows we shall assume (without loss of generality) that all the regularity parameters, i.e., R, M (cid:96) , L (cid:96) , B (cid:96) , M w , L w , B w , are larger than 1 .
Since the loss is convex, it is straightforward to optimize the naive objective that does not reweight the loss and simply minimizes (cid:96) ( · ) over the noisy examples, R naive (Θ)
(cid:44) E ( x,y adv ) ∼ D [ (cid:96) ( y adv , Θ x )] .
We ﬁrst show that (unsurprisingly) it is not hard to construct instances (even in binary classiﬁcation) where optimizing the naive objective produces classiﬁers with large generalization error over clean examples.
For simplicity, since in the following lemma we consider binary classiﬁcation, we assume that the labels y ∈ {± 1 } and the parameter of the linear model is a vector θ ∈ R d .
Proposition E.3 (Naive Objective Fails) .
Fix any c ∈ [0 , 1] .
Let (cid:96) ( · ) be the Binary Cross Entropy loss, i.e., (cid:96) ( t ) =
log(1 + e − t ) .
There exists a “clean” distribution P and a noisy distribution D on
R d × {± 1 } so that the following hold.
The x -marginal of both P and D is uniform on a sphere.
The clean labels of P are consistent with a linear classiﬁer sign( θ ∗ · x ) .
D has (total) label noise Pr ( x,y adv ) ∼ D
[ y adv (cid:54)
= sign( θ ∗ · x )] =
c ∈ [0 , 1] .
The minimizer (cid:98) θ of the (population) naive objective R naive ( θ ) =
E ( x,y ) ∼ D [ (cid:96) ( y adv θ · x )] , constrained on the unit has generalization error
[Equation]
( (cid:98) θ ) ≥ min (cid:107) θ (cid:107)
[Equation]
where R ( θ ) is the “clean” risk, R ( θ )
=
E ( x,y ) ∼ P [ (cid:96) ( yθ · x )] .
Our positive results show that, having a debiasing weight function w ( · ) that is not very “wild” (see the regularity assumptions of Deﬁnition E.2) and optimizing the corresponding weighted objective
R w (Θ)
=
E ( x,y adv ) ∼ D [ w ( x, y adv ; Θ) (cid:96) ( y adv , Θ x )] with SGD, gives models with almost optimal generalization.
The main issue with optimizing the reweighted objective is that, in general, we have no guarantees that the weight function preserves its convexity (recall that it depends on the parameter Θ ).
However, we know that its population version corresponds to the clean objective R ( · ) which is a convex objective.
We show that we can use the convexity of the underlying clean objective to show results for both single- and multi-pass stochastic gradient descent.
We ﬁrst focus on single-pass stochastic gradient descent where at every iteration a fresh noisy sample ( x, y adv ) is drawn from D , see Algorithm 3.
Input: Number of iterations T , Step size sequence η ( t )
Output: Parameter vector Θ ( T ) .
Initialize Θ (1) ← 0 .
For t = 1 , . . .
, T :
Draw sample ( x ( t ) , y ( t ) adv ) ∼ D .
Update using the gradient of the weighted objective:
[Equation]
B
[Equation]
t
[Equation] [Equation]
t )
[Equation] [Equation]
[Equation]
( t
[Equation]
(cid:17)(cid:17)
Return Θ ( T ) .
Theorem E.4 (Generalization of Reweighted Single-Pass SGD) .
Assume that the example distributions P , D , the loss (cid:96) ( · ) , and the weight function w ( · ) satisfy the assumptions of Deﬁnition E.2.
Set κ =
L w M (cid:96) + RM w L (cid:96) .
After T = Ω( κ 2 /(cid:15) 2 ) SGD iterations (with a step size sequence that depends on the regularity parameters of Deﬁnition E.2, see Algorithm 3), with probability at least 99% , it holds
[Equation]
The main observation in the single-pass setting is that, since the weight function w ( · ) is debiasing, we can view the gradients of the reweighted objective as stochastic gradients of the true objective over the clean samples.
Therefore, as long as we draw a fresh i.i.d.
noisy sample ( x, y ) ∼ D at each round, the corresponding sequence of gradients corresponds to stochastic unbiased estimates of the gradients of the true loss R (Θ) .
We next turn our attention to multi-pass SGD (see Algorithm 4), where at each round we pick one of the N samples with replacement and update according to its gradient.
The key difference between single- and multi-pass SGD is that the expected loss over the stochastic algorithm for single-pass SGD is the population risk, while the expected loss for multi-pass SGD is the empirical risk.
In other words, in the multi-pass setting we have a stochastic gradient oracle to the empirical reweighted objective (cid:98) R w (Θ) = 1 N (cid:80)
Ni =1 w ( x ( i ) , y ( i ) adv ; Θ) (cid:96) ( y ( i ) adv , Θ x ( i ) ) , which is not necessarily convex.
Our second theorem shows that under the regularity conditions of Deﬁnition E. multi-pass SGD also achieves good generalization error.
Theorem E.5 (Generalization of Reweighted Multi-Pass SGD) .
Assume that the example distributions P , D , the loss (cid:96) ( · ) , and the weight function w ( · ) satisfy the assumptions of
Deﬁnition E.2.
Set κ = RM (cid:96) L (cid:96) B (cid:96)
M w L w B w and deﬁne the empirical reweighted objective with
N = ( dL ) 2 /(cid:15) 2 poly( κ ) i.i.d.
samples ( x (1) , y (1)adv ) , . .
, ( x ( N ) , y ( N ) adv ) from the noisy distribution D as (cid:98)
[Equation]
=
[Equation]
N
N (cid:88)
[Equation]
Then, after T = poly( κ ) /(cid:15) 4 iterations, multi-pass SGD with constant step size sequence η ( t ) =
C/ √ T 3 (see Algorithm 4) on (cid:98) R w ( · ) outputs a list Θ (1) , . .
, Θ ( T ) that, with probability at least 99% , contains a vector (cid:98) Θ that satisﬁes
[Equation]
We remark that our analysis also applies to the multi-pass SGD variant where, at every epoch we pick a random permutation of the N samples and update with their gradients sequentially.
Draw N i.i.d.
samples ( x (1) , y (1)adv ) , . .
, ( x ( N ) , y ( N ) adv ) ∼ D .
Initialize Θ (1) ← 0 .
For t = 1 , . . .
, T :
Pick I uniformly at random from { 1 , . . .
, N } and update using the gradient of the reweighted objective:
[Equation]
Return Θ (1) , . .
, Θ ( T ) .
E.1 The proof of Proposition E.
In this subsection we restate and prove Proposition E.3.
Proposition E.6 (Naive Objective Fails (Restate of E.3) ) .
Fix any c ∈ [0 , 1] .
Let (cid:96) ( · ) be the Binary
Cross Entropy loss, i.e., (cid:96) ( t ) =
log(1 + e − t ) .
There exists a “clean” distribution P and a noisy distribution D on R d × {± 1 } so that the following hold.
The x -marginal of both P and D is uniform on a sphere.
The clean labels of P are consistent with a linear classiﬁer sign( θ ∗ · x ) .
D has (total) label noise Pr ( x,y adv ) ∼ D
[ y adv (cid:54)
= sign( θ ∗ · x )] =
c ∈ [0 , 1] .
The minimizer (cid:98) θ of the (population) naive objective R naive ( θ ) =
E ( x,y ) ∼ D [ (cid:96) ( y adv θ · x )] , constrained on the unit has generalization error
[Equation]
where R ( θ ) is the “clean” risk, R ( θ )
=
E ( x,z ) ∼ P [ (cid:96) ( zθ · x )] .
Proof.
We set the X -marginal to be the uniform distribution on a sphere of radius R > 0 to be speciﬁed later in the proof.
We ﬁrst observe that the unit vector θ ∗ minimizes the (clean) Binary
Cross Entropy R ( θ ) .
We can now pick a different parameter vector (cid:101) θ with angle φ ∈ [0 , π ] with θ ∗ , and construct a noisy instance as follows: we ﬁrst draw x ∼ X (recall that we want the x -marginal of the noisy distribution to be the same as the “clean”) and then set y
adv = sign( (cid:101) θ · x ) .
By the symmetry of the uniform distribution on the sphere we have that
[Equation]
Therefore, by picking the angle φ to be equal to πc we obtain that Pr x ∼ D x
[sign( (cid:101) θ · x ) (cid:54) =
sign( θ ∗ · x )]
= c as required by Proposition E.3.
Moroever, we have that the minimizer of the “naive” BCE objective (constrained on the unit ball) is (cid:101) θ and the minimizer of the clean objective is θ ∗ .
We have that
[Equation]
Moreover, we have that
[Equation]
We next need to bound from below the “margin” of the optimal weight vector θ ∗ , i.e., provide a lower bound on | θ ∗ · x | that holds with high probability.
We will use the following anti-concentration inequality on the probability of a origin-centered slice under the uniform distribution on the sphere.
For a proof see, e.g., Lemma 4 in [10].
Lemma E.7 (Anti-Concentration of Uniform vectors, [10]) .
Let v ∈ R
d be any unit vector and let X be the uniform distribution on the sphere.
It holds that
[Equation]
Using Lemma E.7, we obtain that
[Equation]
where, at the last step, we used the elementary inequality log(1 + x ) ≤ x .
Assuming that R/ √ d is much larger than 1 , we can pick γ = ( √ d/R ) log( R/ √ d ) .
For this choice of γ we obtain that
E x ∼ X
[log(1 + e −| θ ∗ · x | )]
= O ( √ d/R log( d/R )) .
Therefore, for R = O ( √ d/c ) we obtain that
[Equation]
Therefore, combining the above bounds we obtain that R ( (cid:98) θ ) − R ( θ ∗ ) ≥ c − c/ 2 ≥ c/ 2 .
E.2 The Proof of Theorem E.
In this section we restate and prove our result on the generalization error of single-pass stochastic gradient descent on the weighted objective.
Theorem E.8 (Generalization of Reweighted Single-Pass SGD (Restate of E.4)) .
Assume that the example distributions P , D and the (cid:96) ( · ) and weight function w ( · ) satisfy the assumptions of
Deﬁnition E.2.
Set κ =
L w M (cid:96) + RM w L (cid:96) .
After T = Ω( κ 2 /(cid:15) 2 )
SGD iterations (see Algorithm 3), with probability at least 99% , it holds
[Equation]
Proof.
We observe that, since w is a debiasing weight function, given a sample ( x ( t ) , y ( t ) adv ) ∼ D it holds that ∇ ( w ( x ( t ) , y ( t )
adv ; Θ) (cid:96) ( y adv , Θ x ( t ) )) is an unbiased gradient estimate of ∇ Θ R (Θ) .
We will use the following result on the convergence of the last-iterate of SGD for convex objectives.
For simplicity, we state the following theorem for the case where the parameter θ is a vector in R d (instead of a L × d matrix).
Lemma E.
(Last Iterate Stochastic Gradient Descent [22]) .
Let W be a closed convex set of diameter
R .
Moreover, let F : R d (cid:55)→ R be a convex, L -Lipschitz function. Deﬁne the stochastic gradient descent iteration as
[Equation]
where g ( t ) ( θ ( t ) ) is an unbiased gradient estimate of ∇ θ f true ( θ ( t ) ) .
Assume that for all t ∈ [ T ] it holds (cid:107) g
( t ) ( θ ) (cid:107) 2 ≤ L for all θ ∈ W .
There exists a step size sequence η ( t ) that depends only on
T, L, R such that, with probability at least 1 − δ , it holds
[Equation]
To simplify notation we let (cid:96) ( t ) (Θ) (cid:44) (cid:96) ( y ( t ) adv , Θ x ( t ) ) and w ( t ) (Θ) (cid:44) w ( x ( t ) , y ( t ) adv ; Θ) .
For a sample ( x ( t ) , y ( t ) adv ) , the gradient g ( t ) of the weighted loss is:
[Equation]
Using the triangle inequality for the Frobenious norm and the assumptions of Deﬁnition E.2 on the functions w ( · ) and (cid:96) ( · ) , we obtain that (cid:107) g ( t ) (cid:107) F ≤ L w M (cid:96) + RM w L (cid:96) .
Using Lemma E.9 we obtain that with T =
Ω
(( L w M (cid:96) + RM w L (cid:96) ) 2 /(cid:15) 2 ) , the last iteration of Algorithm 3 satisﬁes the claimed guarantee.
E.3 The proof of Theorem E.
In this section we prove our result on multi-pass SGD.
For convenience, we ﬁrst restate it.
Theorem E.10 (Generalization of Multi-Pass SGD (Restate of E.5)) .
Set κ = RM (cid:96) L (cid:96) B (cid:96)
M w L w B w and deﬁne the empirical reweighted objective with N = d 2 /(cid:15)
2 poly( κ ) i.i.d.
samples ( x (1) , y (1)adv ) , . .
, ( x ( N ) , y ( N ) adv ) from the noisy distribution D as
[Equation]
Then, after T = poly( κ ) /(cid:15) 4 iterations, multi-pass SGD (see Algorithm 4) on (cid:98) R w ( · ) outputs a list θ (1) , . .
, θ ( T ) that, with probability at least 99% , contains a vector (cid:98) θ that satisﬁes
[Equation]
Proof.
To prove the theorem we shall ﬁrst show that all stationary points of the empirical objective (which for arbitrary weight functions w ( · ) may be non-convex) will have good generalization guarantees.
Before we proceed we formally deﬁne approximate stationary points.
To simplify notation we shall assume that the parameter is a vector θ ∈ R d .
The deﬁnition extends directly to the case where L is a function of a parameter matrix Θ by using the corresponding matrix inner product.
Deﬁnition E.11 ( (cid:15) -approximate
Stationary Points) .
Let L : R d (cid:55)→ R be a differentiable function and
C be any convex subset of R d .
A vector θ ∈ R
d is an (cid:15) -approximate
stationary point of L ( · ) if for every θ (cid:48) ∈ C it holds that
[Equation]
Proposition E.12.
Set κ = RM (cid:96) L (cid:96) B (cid:96)
M w L w B w and deﬁne the empirical reweighted objective with N = (cid:101) O (( dL/(cid:15) ) 2 ) poly( κ ) log(1 /δ )
i.i.d.
samples ( x (1) , y (1) ) , . .
, ( x ( N ) , y ( N ) ) from the noisy distribution D as
[Equation]
Let (cid:98) Θ be any (cid:15) -stationary point of (cid:98) R w (Θ) constrained on B R .
Then, with probability at least 1 − δ , it holds that
[Equation]
Proof.
We ﬁrst show that, as long as the empirical gradients are close to the population gradients, any stationary point of the weighted empirical objective will achieve good generalization error.
In what follows we shall denote by Θ ∗ the parameter that minimizes the clean objective:
[Equation]
Since the population objective is convex in Θ , we have that for any Θ it holds that
[Equation]
We have that the contstraint set (cid:107) Θ (cid:107) F ≤ 1 is convex and therefore for a stationary point (cid:98) Θ of L w (Θ) we have that |∇ θ (cid:98) R w (Θ) · (Θ − Θ ∗ ) | ≤ (cid:15) (cid:107) Θ − Θ ∗ (cid:107) F ≤ 2 (cid:15) .
Therefore, (cid:98) Θ satisﬁes
[Equation]
∗
[Equation]
≤
[Equation]
(cid:107)∇
[Equation]
(cid:98) Θ
[Equation]
(cid:107) 2 + 2 (cid:15) .
Since w ( · ) is a debiasing weighting function, we know that, as the number of samples N → ∞ , the empirical gradients of the reweighted objective will converge to the gradients of the population clean objective R ( · ) , i.e., it holds
that ∇ θ (cid:98) R w (Θ) → ∇ Θ R (Θ) .
Therefore, to ﬁnish the proof, we need to provide a uniform convergence bound for the gradient ﬁeld of the empirical objective.
We ﬁrst consider estimating the gradient of some ﬁxed parameter matrix Θ .
We will use McDiarmid’s inequality.
Lemma E.13 (McDiarmid’s Inequality) .
Let x 1 , . . .
, x n be n i.i.d.
random variables taking values in X .
Let φ : X n → R be such that | φ ( x ) − φ ( x (cid:48) )
| ≤ b i whenever x and x (cid:48) differ only on the i -th coordinate.
It holds that
[Equation]
We consider the n i.i.d.
random variables ( x ( t ) , y ( t ) adv ) .
We have that the empirical gradient of the weighted loss function is equal to
[Equation]
We have that w ( x ( t ) , y ( t ) adv ; Θ) is M w -bounded and L w -Lipschitz, (cid:96) is M (cid:96) -bounded and L (cid:96) -Lipschitz, and (cid:107) x
( t ) (cid:107) 2 ≤ R .
Therefore, the maximum value of each coordinate of each term in the sum of the empirical gradient (cid:98) g
( θ ) is bounded by L q (cid:44) L w M (cid:96) + RM w L (cid:96) .
Using this fact we obtain that each coordinate of the empirical gradient is a function of the N i.i.d.
random variables ( x ( t ) , y ( t ) adv ) that satisﬁes the bounded differences assumption with constants b 1 , . . .
, b N that satisfy b t ≤ L q /N .
From
Lemma E.13, we obtain that
[Equation]
We next need to provide a uniform convergence guarantee over the whole parameter space (cid:107) Θ (cid:107) F ≤ 1 .
We will use the following standard lemma bounding the cardinality of an (cid:15) -net of the unit ball in d -dimensions.
For a proof see, e.g., [46].
Lemma E.14 (Cover of the Unit Ball [46]) .
Let B be the d -dimensional unit ball around the origin.
There exists an (cid:15) -net of B with cardinality at most (1 + 2 /(cid:15) )
d .
Since we plan to construct a net for the gradient of w ( x, y adv ; Θ) (cid:96) ( y adv ; Θ x ) we ﬁrst need to show that the weighted loss w ( x, y adv ; Θ) (cid:96) ( y adv ; Θ x ) is a smooth function of its parameter Θ or, in other words, that its gradients do not change very fast with respect to Θ .
The following lemma follows directly from the regularity assumptions of Deﬁntion E.2 and the chain and product rules for the derivatives.
Lemma E.15.
For all ( x, y ) ∈
B R × R L , it holds that the function q (Θ) =
w ( x, y ; Θ) (cid:96) ( y ; Θ x ) is
B q -smooth for all Θ with (cid:107) Θ (cid:107) F ≤ 1 , with B q = M (cid:96) B w + 2 L (cid:96) L w R + M w B (cid:96) R 2 .
Proof.
For simplicy we shall denote ∇ z (cid:96) ( y ; z ) simply by ∇ (cid:96) ( y ; z ) and similarly ∇ 2 z (cid:96) ( y ; z ) by ∇ 2 (cid:96) ( y ; z ) .
Using the chain rule, we have that the gradient of the weighted loss q (Θ) is equal to
[Equation]
Using again the chain and product rules we ﬁnd the Hessian of q (Θ) :
[Equation]
where H is the ( L × d ) ×
( L × d ) tensor with element H ijkl = ∇ 2 (cid:96) ( y ; Θ x )
ik x j x l .
Recall that we view ∇ 2Θ q (Θ) as an Ld × Ld and to prove that it is smooth we have to ﬁnd its operator (spectral) norm.
Using the assumptions of Deﬁnition E.2 we obtain that (cid:107)∇ 2Θ w ( x, y ; Θ) (cid:107) 2 ≤ B w M (cid:96) .
For the term ( ∇ (cid:96) ( y ; Θ x ) x T ) ⊗ ∇ Θ w ( x, y ; Θ) we consider any q ∈ R Ld with (cid:107) q (cid:107) 2 = 1 .
We assume that q is indexed as q ij for i = 1 , . . .
, L and j = 1 , . . .
, d .
We have
[Equation]
Similarly, we bound the spectral norm of the term ∇ Θ w ( x, y ; Θ) ⊗ ( ∇ (cid:96) ( y ; Θ x ) x T ) .
Finally for the term H we have
[Equation]
where s ∈ R L has s i
= (cid:80)
j q ij x j .
Observe that since (cid:107) x (cid:107) 2 ≤ R and (cid:107) q (cid:107) 2 = 1 we have that (cid:107) s (cid:107) 2 ≤ R .
Therefore, from the assumption of Deﬁnition E.2, we obtain that (cid:107) H (cid:107) 2 ≤ R 2 B (cid:96) .
We conclude that the function q ( θ ) is B q -smooth on the unit ball B with B q = M (cid:96) B w + 2 L (cid:96) L w R +
M w B (cid:96) R 2 .
Let N (cid:15) be an (cid:15) -net of the unit ball B .
Using Lemma E.15 we ﬁrst observe that the vector maps θ (cid:55)→ (cid:101) g (Θ) and Θ (cid:55)→ ∇ Θ R (Θ) are both B q -Lipschitz, where B q is the constant deﬁned in Lemma E.15.
Using the triangle inequality and the fact that (cid:98) g ( · ) and ∇ Θ R ( · ) are B q -Lipschitz, we have that
[Equation]
Combining the above, and performing a union bound over the (cid:15) -net N (cid:15) , we obtain that
[Equation]
We conclude that with N = (cid:101) Ω(( dL ) 2 L 2 q B 2 q /(cid:15) 2 log(1 /δ ))
samples, it holds that (cid:107) (cid:98) g
(Θ) − ∇ Θ R (Θ) (cid:107) 2 ≤ (cid:15) , uniformly for all parameters Θ with (cid:107) Θ (cid:107) F ≤ 1 , with probability at least 1 − δ .
We now have to show that the multi-pass SGD ﬁnds an approximate stationary point of the empirical objective.
We will use the following result on non-convex projected SGD.
To simplify notation, we state the following optimization lemma assuming that the parameter is a vector θ ∈ R d .
Lemma E.16 (Non-Convex Projected Stochastic Gradient Descent [11]) .
Let W be a closed convex set of diameter R .
Moreover, let F : R d (cid:55)→ R be an L -Lipschitz and B -smooth function.
Deﬁne the stochastic gradient descent iteration as
[Equation]
where g ( t ) ( θ ( t ) ) is an unbiased gradient estimate of ∇ θ F ( θ ( t ) ) .
Fix a number of iterations T ≥and assume that for all t ∈ [ T ] it holds (cid:107) g
( t ) ( θ ) (cid:107) 2 ≤ L for all θ ∈ W .
Set the step-size η ( t ) =
Θ( (cid:112) R/ ( BL 2 T )) .
With probability at least 99% , there exists a t ∈ { 1 , . .
, T } such that θ ( t ) is an
O (cid:16) √ BLRT 1 / 4 (cid:17) -stationary point of F ( · ) constrained on W .
Theorem E.5 now follows directly by applying Lemma E.16 on the empirical objective to ﬁnd an (cid:15)
-approximate stationary point and then using Proposition E.12.