S CAFFOLDING A S TUDENT TO I NSTILL K NOWLEDGE
Anil Kag ∗ anilkag@bu.edu
Durmus Alp Emre Acar ∗ alpacar@bu.edu Aditya Gangrade † agangra2@andrew.cmu.edu Venkatesh Saligrama ∗ srv@bu.edu
A BSTRACT
We propose a novel knowledge distillation (KD) method to selectively instill teacher knowledge into a student model motivated by situations where the student’s capacity is significantly smaller than that of the teachers.
In vanilla KD, the teacher primarily sets a predictive target for the student to follow, and we posit that this target is overly optimistic due to the student’s lack of capacity.
We develop a novel scaffolding scheme where the teacher, in addition to setting a predictive target, also scaffolds the student’s prediction by censoring hard-to-learn examples.
The student model utilizes the same information as the teacher’s soft-max predictions as inputs, and in this sense, our proposal can be viewed as a natural variant of vanilla KD.
We show on synthetic examples that censoring hard-examples leads to smoothening the student’s loss landscape so that the student encounters fewer local minima.
As a result, it has good generalization properties.
Against vanilla KD, we achieve improved performance and are comparable to more intrusive techniques that leverage feature matching on benchmark datasets.
I NTRODUCTION
A fundamental problem in machine learning is to design efficient and compact models with near state-of-the-art (SOTA) performance.
Knowledge Distillation (KD) (Zeng & Martinez, 2000; Bucila et al., 2006; Ba & Caruana, 2014; Hinton et al., 2015) is a widely used strategy for solving this problem wherein the knowledge from a large pre-trained teacher model with SOTA performance is distilled onto a small student network.
Vanilla KD.
Hinton et al.
(2015) proposed the popular variant of KD by matching the student soft predictions, s ( x ) with that of the pre-trained teacher, t ( x ) on inputs x .
Informally, during student training, an additional loss term, D KL ( t ( x ) , s ( x )) is introduced that penalizes the difference between student and teacher predictive distributions using Kullback-Leibler (KL) divergence.
This promotes inter-class knowledge learned by the teacher.
We will henceforth refer to Vanilla KD as KD.
Capacity mismatch between student and teacher.
One of the primary issues in KD is that the loss function is somewhat blind to the student’s capacity to interpolate.
In particular, when the student’s capacity is significantly lower than the teacher’s, we expect the student to follow the teacher only on those inputs realizable by the student.
We are led to the following question: What can the teacher provide by way of predictive hints for each input, so that the student can leverage this information to learn to its full capacity?
Our Proposal: Scaffolding a Student to Distill Knowledge (DiSK).
To address this question, we propose that the teacher, during training, not only set a predictive target, t ( x ) , but also provide hints on hard to learn inputs.
Specifically, the teacher utilizes its model to output a guide function, g ( x ) , such that the student can selectively focus only on those examples that it can learn.
• if g ( x ) ≈ 1 , teacher discounts loss incurred by the student on the input x . • if g ( x ) ≈ 0 , teacher signals the input x as learnable by student.
With this in mind we modify the KL distance in the KD objective and consider,
D KL ( t ( x ) , ϕ ( s ( x ) , g ( x ))) , where ϕ ( s , g ) , which will be defined later, is such that, ϕ ( s , 0) =
s if g offers no scaffolding.
We must impose constraints on the guide function g to ensure that only hard-to-learn examples are scaffolded. In the absence of such constraints, the guide can declare all examples to be hard, and the student would no longer learn.
We propose to do so by means of a budget constraint B ( s , g ) ≤ δ to ensure that the guide can only help on a small fraction of examples.
While more details are described in Sec. 3, we note that, in summary, our proposed problem is to take the empirical linear combination of the aforementioned KL distance and a cross-entropy term as the objective, and minimize it under the empirical budget constraint.
We emphasize that g ( x ) is used only during training .
The inference logic for the student remains the same as there is no need for g ( x ) during inference.
The guide function supported student training has three principal benefits.
The benefits are explored in Sec
• Censoring Mechanism.
Our guide function censors examples that are hard to learn for the student.
In particular, when there is a large capacity gap, it is obvious that the student cannot fully follow the teacher.
For this reason, the teacher must not only set an expectation for the student to predict, but also selectively gather examples that the student has the ability to predict.
• Smoothen the Loss landscape.
We also notice in our synthetic experiments that whenever scaffolding is powerful, and can correct student’s mistakes, the loss landscape undergoes a dramatic transformation.
In particular, we notice fewer local minima in the loss viewed by the guided student.
• Good Generalization.
The solution to our constrained optimization problem in cases where the guide function is powerful can ensure good student generalization.
Specifically, we can bound the statistical error in terms of student complexity and not suffer additional complexity due to the teacher.
Contributions.
We summarize our main results.
• We develop a novel approach to KD that exploits teacher representations to adjust the predictive target of the student by scaffolding hard-to-learn points.
This novel scaffolding principle has wider applicability across other KD variants, and is of independent interest.
• We design a novel response-matching KD method (Gou et al., 2021) which is particularly relevant in the challenging regime of large student-teacher capacity mismatch.
We propose an efficient constrained optimization approach that produces powerful training scaffolds to learn guide functions.
• Using synthetic experiments, we explicitly illustrate the structural benefits of scaffolding.
In particular, we show that under our approach, guides learn to censor difficult input points, thus smoothening the student’s loss-landscape and often eliminating suboptimal local minima in it.
• Through extensive empirical evaluation, we demonstrate that the proposed DiSK method; – yields large and consistent accuracy gains over vanilla KD under large student-teacher capacity mismatch (upto 5% and 2% on CIFAR-100 and Tiny-Imagenet).
– produces student models that can get near-teacher accuracy with significantly smaller model complexity (e g. 8 × computation reduction with ∼ 2% accuracy loss on CIFAR-100).
– improves upon KD even under small student-teacher capacity mismatch, and is even competitive with modern feature matching approaches.
I LLUSTRATIVE E XAMPLES
We present two synthetic examples to illustrate the structural phenomena of the censoring mechanism and smoothening of student’s loss landscape enabled by the scaffolding approach DiSK, which lead to globally optimal test errors.
We defer exact specification of the algorithm to Sec.3.
Example 1 (1D Intervals) .
Consider a toy dataset with one dimensional features x ∈ [0 , 9] and binary class labels
y ∈
{ Red , Blue } as shown in Figure 1.
There are two Blue labelled clusters as in [2 , 3] and in [5 , 7] .
The remaining points are labelled as Red .
We sample 1000 i.i.d.
data points as the training set and 100 data points as the test set with balanced data from both classes.
We describe the details of the experiment setup such as models and learning procedure in Appx. A.1.
Teacher T belongs to the 2-interval function class, and the capacity-constrained student S belongs to the 1-interval function class.
Since teacher capacity is sufficient to separate the two classes without error, it learns the correct classifier (see Figure 1).
In contrast, the best possible student hypothesis cannot correctly separate the two classes.
Hence, the student will have to settle onto one of the many [Figure] [Table]
[Figure] [Table] local minima. We show these minima and the contour plot for the student in Figure 1.
We present the results of training student models with different initializations in Table 1.
KD suffers from bad local minima.
KD loss landscape contains many local minima (see Figure 1).
Due to a big gap between student and teacher capacity, it is unable to help the student discern between these minima.
Hence, KD fails to distinguish between the different minima (see Table 1).
DiSK censors interval [2 , 3] and in addition focuses training on learnable datapoints.
If we analyze the guide function at the end of the training, we see that it covers (censors) the first Blue cluster.
Indeed, both clusters are not simultaneously l earnable with the available student capacity.
Once we censor the interval [2 , 3] , then the problem becomes realizable for the student model.
The guide function thus captures the excess capacity of the data.
DiSK smoothens loss landscape.
The guide function and the budget constraint enable our method to have a smooth loss landscape thanks to the guide-function censoring points, which eliminate the local minima.
Hence, DiSK solution lands in the global minimum with high probability.
Example 2 (2D Gaussians).
Consider another toy dataset with two dimensional features x ∈ Rand three class labels y ∈
{ Red , Green , Blue } .
Here we wish to show that DiSK can allow for globally optimal solutions reaching 100% accuracy, which appears unachievable with cross-entropy minimization regardless of data size.
Figure 2a shows the labelled data.
There are six cluster centers, two with each class label.
Data points are drawn using Gaussian balls around the cluster centers with small radii.
We sample 1000 i.i.d.
data points as the training set and 1000 data points as the test set with equal representation from all three classes.
We provide details (hypothesis classes, learning procedure, etc.) in Appx.
A.2.
The teacher is a 3-layer neural network with 8, 16, and 3 neurons.
The student is a 2-layer neural network with 2 and 3 neurons.
We point out that the teacher being an over-parameterized network in this feature space, easily learns the correct decision boundary.
While the student being severely constrained network suffers in learning the task.
Different training runs lead to different local minima.
We show teacher solution and student local minima in Figure 2a.
The contour plots for the student models under KD loss and DiSK loss are shown in Figure 2b-2c using Li et al.
The results are similar to the 1D example - KD converges to a poor local minimum with at least 43% of the initializations, while in contrast, DiSK escapes these by focusing on the learnable part of the input space (Fig. 2c), converging to the global minimum nearly always (Table 1). [Figure]
To conceptualize our findings in these examples, let us attempt to intuitively infer the examplecensoring, landscape-smoothening, and good generalization, by utilizing the following conditions that appear to be satisfied for these synthetic examples.
Realizability.
Suppose we are in a situation where the guide function g ∈ G is sufficiently powerful that there is a student and guide function capable of interpolation,
i.e., predictions supported by the guide function, ϕ ( s , g ) , interpolates to mimic the labels.
Example: For instance, consider a binary classification problem with the labels y ∈ {− 1 , 1 } .
Let ϕ ( s, g )
=
y ( s + g ) with s ( x ) ∈ [ − 1 , 1] .
Our realizability condition is that we always satisfy y ( s ( x ) + g ( x )) > 0 .
As such, this leads to the condition that if ys ( x ) ≤ 0 , then yg ( x ) > 0 .
Therefore, E [ 1 [ ys ( x ) < 0] ] ≤ E [ 1 [ ys ( x ) < 0 ,yg ( x ) > 0] ] ≤ E [ 1 [ yg ( x ) > 0] ] .
Small Guide Function Capacity.
In addition to realizability suppose the class of guide functions g ∈ G has a small capacity (for instance, small VC dimension).
For our case this condition is satisfied because our guidance function is obtained by using an MLP on teacher’s last layer features.
Example: Continuing with the example above, say we now have m training instances, ( x i , y i ) , i ∈ [ m ] , ˆ g ( x ) is guide function output of DiSK.
We can infer by standard statistical learning results (Shalev-Shwartz & Ben-David, 2014) that, for the estimated function ˆ g ∈ G , it follows with probability greater than 1 − η that E [ 1 [ y ˆ g ( x ) > 0] ] ≤ 1 m (cid:80) mi =1
[ y i ˆ g ( x i ) > 0]
+ (cid:113) V C ( G )
+log 1 η m .
As a result, we can say that if there is a student, s ( x ) (not necessarily that output by DiSK), which complements ˆ g ( x ) and satisfies realizability, then with probability greater than 1 − η : E
[Equation]
Remarks.
The key point is that the student capacity is considerably larger since we typically train an entire DNN, and student complexity-based bound can be vacuous.
While the guidance function does bound the student generalization error in terms of guide function complexity, there are strong caveats— we require the strong assumption of realizability on the entire domain, and additionally, while the guide function can witness student error, we are not in a position to precisely estimate it without additional training data.
Furthermore, the RHS is a relaxed bound on the student training error.
This motivates having a budget constraint to ensure that student learns with small training error.
D EFINITIONS AND F ORMULATIONS
Notation .
Let X and Y = { 1 , . .
, C } be the feature and label spaces respectively, focusing on a Cclass classification task.
We assume that we have a training set of N i.i.d.
data points D
= { x i , y i } Ni =1 , where x i ∈ X
and y i ∈ Y .
We use symbols S and T to denote the student and teacher models respectively.
Let l S ( x ) ∈ R |Y| and l T ( x
) ∈ R |Y| be the score vector, logits, predicted by S and T on input x .
We use τ as the temperature used to soften the probability distribution.
We write the resulting softened student and teacher probabilities as s τ ( x ) and t τ ( x ) , i.e.,
[Equation]
The standard prediction probabilities correspond to s 1 ( x ) and t 1 ( x ) .
We will use s τy ( x ) to denote the y th coordinate in s τ ( x ) , and similarly t τy ( x ) .
The hard prediction of the student is p S ( x ) = [Equation]
max y [Equation] s 1 y ( x ) , and similarly p T ( x ) = arg max y [Equation]
t 1 y ( x ) for the teacher.
We use g ( x ) ∈ [0 , 1] to denote the helper guide function for the student and teacher pair ( S, T ) .
Guide takes input x and any other feature processed by ( S, T ) pair and decides whether or not the student needs help on the input x .
Finally, we define ReLU activation as ( · ) + = max(0 , · ) .
3. V ANILLA K NOWLEDGE D ISTILLATION
KD relaxes the 0 − 1 error between the student predictions and the true labels y using the cross-entropy loss L CE .
Then, KD denotes the distance between the student and teacher softened probability distributions using the KL divergence.
We summarize the corresponding losses as,
[Equation]
For hyperparameters α ∈ [0 , 1] , τ > 0 , KD minimizes a mixture of the above losses, as shown below
[Equation]
3. S ELECTIVE K NOWLEDGE D ISTILLATION .
KD attempts to transfer the knowledge from the teacher to the student on all training data points, which is a sub-optimal objective when there is a capacity mismatch between the student and the teacher.
Instead, we propose distilling selective knowledge (DiSK) to allow the student to selectively ignore some hard-to-learn data points during training, transferring the teacher’s knowledge only on easy-to-learn inputs, and matching the learning to student capacity.
Our objective is to minimize min s
[Equation]
g,δ
[Equation]
where ϕ interpolates student predictions based on the guide’s help.
The divergence term helps in minimizing the distributional distance between the teacher and student probabilities after the guide g is included.
While the budget term in the optimization constrains the helper g to provide help only when necessary, the amount of help given to the student should be within the budget δ ∈ [0 , 1] .
Function g Construction.
As previously stated, we use the teacher’s last layer features and soft predictions as input to the guide g .
The guide is structured as a light-weight three-layer neural network with these inputs, with a sigmoid activation at the last layer.
We re-emphasise that g is not used at inference time, and only aids training.
More details are left to Appx.
A.4.
Relaxed Losses, Lagrangian & Optimization Algorithm.
We relax Eq. 2 and construct a Lagrangian by integrating the constraint into the minimization.
Budget constraint relaxation.
We relax the indicator loss in the budget to a cross-entropy, and treat δ as a hyperparameter to get
[Equation]
We view the scaffold as a way for the student to interpolate the uncensored data.
It suggests that a good initialization for the budget is the error of cross-entropy trained model when the student does not have the teacher supervision.
Thus, we scan the budget in a small interval around this initialization. Input
Training = x
Teacher
[Equation] [Equation] [Equation]
: Return Distillation objective. Motivated from KL loss, we construct a distance loss with guide function as,
[Equation]
We point out two modifications in the distillation loss.
First, L dist explicitly adds guide value to softened student probabilities in selected class indices.
The class indices guide function adds value are picked as top K classes based on the teacher probabilities for any input
x i where K is a hyperparameter of our method.
The rest of the class indices do not get any value from g .
Second, we use different temperature parameters for teacher and student.
Temperature parameter for teacher, τ , is a hyperparameter.
The student temperature is found by minimizing the KL loss between teacher softened probabilities and the student softened probabilities over the training dataset, .i.e
[Equation]
Similar to KD, we incorporate standard cross entropy loss between student model predictions and the ground truth labels for stability.
We construct our Lagrangian by combining Eq.
3 and 4 as,
[Equation]
where α is a hyper-parameter and λ is the dual parameter of DiSK.
We optimize Obj.
5 using a primal dual update scheme as explained in Algorithm 1.
Primal Parameter Updates ( s , g ).
We learn the student s and the guide function g using alternating minimization.
We approximate arg min with running SGD for a small number of epochs on the full dataset.
In each iteration, we first learn the guide function g to select the data partition from which the knowledge needs to be distilled.
Next, given the function g , we learn the student using the help g .
We empirically found that not optimizing the student model on budget loss gives more stable results.
Hence, we minimize the student model only on the distillation and cross-entropy losses.
Dual Parameter Update ( λ ) Intuition.
Although it is tempting to optimize the above via a dual ascent and primal descent scheme (wherein the dual parameter λ is increased by residual term in the constraint until constraint satisfaction), recent work, Sun & Sun (2021), has proposed to decrease the λ in the non-convex regime.
Inspired by this, we update the dual parameter by a fixed schedule between [ λ min , λ max ] .
λ min encourages exploration and allows student model to distill knowledge from all points.
On the other hand, λ max enforces the constraint and forces the student model to learn on uncensored inputs.
We choose R ≈ 4 λ T , so that the algorithm is exposed to a few exploratory periods.
For the final period, we increase λ monotonically so that budget is more strictly enforced at termination.
Computational Efficiency.
Algorithm 1 trains both student and guide networks.
The guide network being small (three-layer MLP) relative to the student (CNN model), the additional cost in training the guide is relatively insignificant, and as such DiSK efficiency is similar to KD.
E XPERIMENTS
We evaluate DiSK in various capacity mismatch scenarios on benchmark datasets. We avail our code at https://github.com/anilkagak2/DiSK_Distilling_Scaffolded_Knowledge
Datasets.
We use publicly available CIFAR-100 (Krizhevsky, 2009), Tiny-Imagenet (Le & Yang, 2015) datasets.
CIFAR-100 contains 50 K training and 10 K test images from 100 classes with size [Table]
32 × 32 × 3 . While Tiny-Imagenet contains 100 K training and 10 K test images from 200 classes with size 64 × 64 × 3 .
We provide the dataset setup and data augmentations used in detail in Appx. A.3.
Models.
We evaluate standard convolutional models on these datasets including ResNet(He et al., 2016), Wide-ResNet(Zagoruyko & Komodakis, 2016), MobileNet(Sandler et al., 2018)
[Table]
and ShuffleNet(Ma et al., 2018).
Table 2 shows the storage and computational requirements of all the models used in this work.
We provide explicit model configurations in Appx. A.4, including the tiny models we generate from the ResNet architectures.
Methods.
We study performance against standard cross-entropy (CE) based learning and the vanilla
KD methods.
For each method, we train models for 200 epochs using SGD as the optimizer with 0 .momentum and 0 .
1 learning rate.
See Appx. A.5 for more training details.
We have recorded the mean in our results as the variance of 3 trials in our experiments is not larger than 0 .
1 in most cases.
We perform evaluations in different settings.
Below, we explain individual setups.
We cover more ablative experiments in Appx.
A due to page limit.
Large Capacity Mismatch Setting.
We distill knowledge from a teacher model into a student model where the student has much less capacity compared to the teacher model.
We use four large capacity
ResNet teachers and five tiny ResNet students and train these students using CE, KD, and DiSK methods.
Performances, and the gains of DiSK are reported in Table 3.
Small Capacity Mismatch Setting.
While DiSK has been designed for the scenario when student capacity is very low, we further evaluate it in the setting where teacher and student capacities are similar, to probe how far the power of the method extends.
The model classes used are the standard choice for this scenario (Chen et al., 2022; Tung & Mori, 2019).
Performance is reported in Table 4.
Table 4 further reports the results of the feature matching distillation methods: FitNets (Romero et al., 2015), SemCKD (Chen et al., 2021), and SimKD (Chen et al., 2022).
Such methods can often outperform response matching KD on large students, due to student representations that are more aligned with the teacher, but typically at an increased training cost.
While feature matching methods are not the main focus of our work (and in principle scaffolding idea can be extended to them), we observe that DiSK often improves upon their performance without any direct feature matching.
Experiment results .
Below, we highlight salient features of DiSK based on empirical data.
DiSK outperforms the baselines uniformly across all datasets and student sizes.
As shown by Table 3,
DiSK significantly improves the student performance in CIFAR-100 and the (more challenging)
Tiny-Imagenet dataset, respectively showing accuracy gains of up to 5% and 2% compared to KD.
These gains are consistent across a wide range of student and teacher capacities.
DiSK achieves better performance with worse teachers than KD does with even the best teachers.
In Table 3, we point out that the student performance increases for KD as the teacher complexity is increased for a given student.
But note that for the same student, DiSK achieves much better performance with even worse teacher.
For instance, for the ‘ResNet10-m‘ student, KD accuracy increases from 66 .
96% to 68 .
09% by using high capacity teachers.
But ‘ResNet10-m‘ trained with even the worst teacher ( ‘ResNet10- ℓ ‘ ) achieves 70 .
03% accuracy.
This saves a lot of resources in any application as large teacher requires more training time, and larger compute resources.
DiSK is competitive even in small capacity difference setting.
As shown by Table 4, DiSK does not loose its competitive edge over the KD even when the student is relatively similar sized as [Table] [Table] teachers, and shows gains of up to 2 .
5% relative to KD.
We conjecture that the observed gains arise from the fact that DiSK provides scaffolding for hard points to the student in initial training stages, which promotes the student to learn easy examples first.
As training progresses, DiSK removes the discounted help from hard inputs.
As a result, the student evolves from simpler hypothesis to the ones consistent with both easy and hard inputs.
This justifies our dual parameter ( λ ) update in Algorithm 1, wherein we periodically increase and decrease λ to enforce and relax the budget constraint.
DiSK students achieve near teacher accuracy while saving up to 8 × MACs & 5 × Params.
As reported in Table 3, student (‘ResNet10-m‘) trained with the teacher (‘ResNet10- ℓ ‘) achieves close to the teacher accuracy of 71 .
In this process, it saves 4 × compute and requires 4 × less parameters.
Similarly, student (‘ShuffleNetV2‘) trained with the teacher (‘ResNet32x4‘) achieves close to the teacher accuracy of 81 .
In this process, it saves 24 × compute and requires 5 × less parameters.
DiSK cleverly selects a subset of datapoints and smoothens the loss landscape.
As illustrated in
Figure 1 and 2, DiSK judiciously selects a subset of hard-to-learn data points for the students and provide discounted help to the student focus on easily learnable inputs.
As a result, it eliminates some bad local minima in the student loss-landscape, and smoothens tihs surface.
DiSK enables the student to reach saturation capacity.
In Table 3, the performance of KD often suffers as the teacher size is increased, e.g., student (‘ResNet10-s’) accuracy decreases substantially with the teacher ‘ResNet34’ versus the teacher ‘ResNet18’.
In contrast, DiSK saturates the student performance across different teachers.
For instance, student (‘ResNet10-m‘) accuracy is ≈ 70% for all the teachers.
Thus, we point out that DiSK enables the student to reach saturation.
This may be due to the fact that guide g identifies the same set of ‘easy’ points across different teachers.
R ELATED W ORK
We refer the reader to Gou et al.
(2021) for a comprehensive survey on knowledge distillation.
Response Matching.
Zeng & Martinez (2000); Bucila et al.
(2006) distill the response of an ensemble of classifiers into a single neural network by creating a pseudo-labeled dataset using the ensemble of classifiers.
Ba & Caruana (2014) extend this to the setting with the neural network as the teacher.
Hinton et al.
propose vanilla KD that distills knowledge from an ensemble of neural networks into a single network by matching their output logits.
This work provided a simple recipe for aligning the teacher and student predictive distributions using the Kullback-Leibler (KL) divergence.
Recently,
Beyer et al.
(2021) modify the KD procedure to include patient and consistent teacher resulting in substantial gains.
Knowledge consistency is enforced by using the same aggressive data augmentation and image views in the student as in the teacher.
Patience is promoted using a very long training schedule.
This results in a computationally very expensive training process.
Stanton et al.
(2021) analyze response matching KD and suggests that difficulty in optimization leads to poor knowledge distillation.
Thus, the teacher and student predictions do not always match, even on the training data.
Cho & Hariharan (2019) study vanilla KD through the lens of mismatched student and teacher capacities.
They show that small students are unable to mimic complex teachers.
They propose early stopping teacher training as to remedy to achieve a student-learnable teacher.
The above works fail miserably when the gap between student and teacher complexities is large.
Specifically, the student cannot learn the complex teacher decision boundaries primarily due to the small student capacity.
It becomes imperative to selectively choose only easy-to-learn data points and transfer the teacher knowledge from these points and ignore the hard-to-learn data points during distillation.
Thus, our proposal targets the problem of severe capacity gaps between student and teacher models.
Additionally, our experiments with standard student and teacher configurations show that DiSK is still competitive when the capacity difference is small.
Feature Matching.
In response matching, teacher supervision is limited to its logits.
We can enforce intermediate layer feature matching for refined teacher supervision.
FitNets (Romero et al., 2015) extend the KD by including the feature matching in the middle layers.
Zagoruyko & Komodakis (2017) use feature map attention as the teacher supervision.
Tung & Mori (2019) preserve pairwise similarity in feature maps amongst data points during the distillation.
Chen et al.
(2022) modify the student by projecting the student features onto the teacher feature space and by reusing the teacher classifier.
While our work focuses primarily on selective distillation in vanilla KD for simplicity.
We can easily extend the proposed framework to incorporate it into feature-matching distillation.
Curriculum Learning & Hard Instance Mining.
Curriculum Learning (CL) (Bengio et al., 2009;
Hacohen & Weinshall, 2019; Graves et al., 2017) sorts the data based on their hardness as measured by some scoring function ( ex., predictive entropy, softmax margin score, etc. ).
It presents the data points during training in the order of increasing hardness.
Similarly, Hard Instance Mining (HIM) (Zhou et al., 2020) reduces the weight of the easy example and increases the weight on hard inputs to promote hard-example learning.
We point out that our method is only conceptually related to these works via input hardness.
Our method helps the student with hard examples by providing explicitly discounted help g .
We learn the helper function g (through teacher representation) that decides whether the student needs help on a given input.
Thus, we do not prioritize learning hard examples keeping in mind the fact that student capacity is much smaller than the teacher.
C ONCLUSION
We develop a new knowledge distillation method that utilizes teacher predictions in novel ways combining predictive targets with scaffolding the student on hard-to-learn points by means of a guide function.
Our method is particularly relevant when there is large gap between student and teacher capacities.
We show that our method allows for convergence to better minima based on two key properties.
Our guide function allows for censoring hard-to-learn examples, and the predictive targets set by the teacher on remaining points allow for eliminating bad local minima and smoothening the resulting student loss landscape.
Against vanilla KD, we achieve improved performance and are comparable to more intrusive techniques that leverage feature matching on benchmark datasets. A CKNOWLEDGMENTS This research was supported by Army Research Office Grant W911NF2110246, the National Science
Foundation grants CCF-2007350 and CCF-1955981, AFRL Contract no.
FA8650-22-C-1039, and a gift from ARM Corporation.
A A PPENDIX A. D ETAILS FOR I LLUSTRATIVE E XAMPLE (1D I NTERVALS )
Dataset Overview.
We generate a synthetic toy dataset with one dimensional features x ∈ [0 , 9] and binary class labels y ∈
{ Red , Blue } .
We use σ ( x ) to denote the sigmoid function with a scaled by parameter κ > 0 , i.e., σ ( x ) = 1 1+exp( − κx ) .
Function Classes.
Let H be the 1-interval function class parametrized by two variables { a, b } , i.e., for h ∈ H
[Equation]
Similarly, let F be the 2-interval family parametrized by four variables { a, b, c, d } , i.e., for f ∈ F
[Equation]
Note that any function in H behaves as an indicator for the interval ( a, b ) .
Similarly, any function in
F behaves as an indicator for two exclusive intervals { ( a, b ) , ( c, d ) } .
Data Generation.
We assume that the data is generated using the function f ∗ ∈ F with parameters ( a ∗ , b ∗ , c ∗ , d ∗ ) .
Dataset is sampled with balanced data from both classes.
We label x as red if f ∗ ( x ) < 0 .
5 , otherwise we label the point as blue.
We sample 1000 i.i.d.
data points as the training set and 100 data points as the test set.
Figure 1 shows the train data.
We draw an independent validation set of 100 data points for hyper-parameter tuning.
Large Capacity Teacher T belongs to the 2-interval function class F and is learnt with all the training data points.
We learn the teacher with cross-entropy loss.
We use the SGD optimizer with momentum 0 .
9 , learning rate 0 .
1 , weight decay 0 .
01 , and minimize the loss for 200 epochs.
Note that the teacher recovers the underlying function f ∗ as shown by the two intervals in Figure 1.
Capacity Constrained Student S belongs to the 1-interval function class H and it has access to all the training dataset.
Note that best possible hypothesis in H cannot recover the performance of the function f ∗ and hence the student will have to settle on one of the many local minima.
We show these minima as well the contour plot for the student in the Figure 1.
We learn the student with three different loss functions ( cross-entropy L CE , vanilla KD L τ,αKD ( s ) , and DiSK Algorithm 1 ).
We use similar training setup as the teacher in terms of the optimizer and training steps.
For DiSK method, our guide function g has similar capacity as the student but utilizes the teacher features to learn the decision as to which points are hard-to-learn for the student.
For both, KD and DiSK, we scan the α hyper-parameter over the range { 0 . 0 , 0 . 1 , 0 . 5 , 0 . 9 , 1 . 0 } .
Similarly, we scan the temperature τ in the range { 1 , 2 , 4 } .
For DiSK, we scan the different hyper-parameters in the following ranges: (a) τ s ∈ { 1 , 2 , 4 } , (b)
K ∈ { 1 , 2 , 3 } , (c) λ min ∈ { 0 . 01 , 0 . 1 , 1 , 5 , 10 } , (d) λ max ∈ { 1 , 5 , 10 , 20 , 50 , 100 , 1000 } , (e) Budget δ ∈ { 0 . 1 , 0 . 05 , 0 . 0 } , and (f) λ T ∈ { 20 , 50 } .
We replace the arg min in Algorithm 1, with three gradient steps.
Note that although the hyper-parameter scan looks daunting, the default hyper-parameters: τ s = τ (teacher temperature), K = 2 , λ min = 0 .
{+1 , λ max = 50 , δ = 0 . 
+}
(approximate error of the global minima), λ T = 50 , work well in this setup as well as the 2D gaussian example described below.
Vanilla KD suffers from local minima.
The loss landscape of the Vanilla KD contains many local minima (see Figure 1(b)).
Since there is a big gap between student and teacher capacity, the teacher is unable to help the student discern between these bad minima.
Hence, Vanilla KD leads to one of the bad local minima with high probability (see Table 1).
DiSK removes bad local minima.
In contrast, DiSK deletes harder points from the landscape and as a result settles onto the global minima for S with high probability (see Figure 1(c) and Table 1 where one cluster of blue points have been removed).
Note that this also removes bad minima from the S loss landscape.
Finally, DiSK learns the student that has the best performance.
A. D ETAILS FOR I LLUSTRATIVE E XAMPLE (2D G AUSSIANS )
Dataset Overview & Data Generation.
We generate a synthetic toy dataset with two dimensional features x ∈ R 2 and three class labels y ∈
{ Red , Green , Blue } .
We generate six cluster centers.
We
• (0 , 0) , Red • (1 . 5 , 0) , Blue • (3 , 0) , Green • (0 , 1 . 5) , Blue • (1 . 5 , 1 . 5) , Green • (3 , 1 . 5) , Red
Given a cluster center c , we draw input features x using a Gaussian ball with radius r = 0 .
05 around the center using multi-variate Gaussian N ( c , r I ) , where I is the Identity matrix.
Figure 2a shows the labelled data.
We sample 1000 i.i.d.
data points as the training set and 1000 data points as the test set with equal representation from all three classes.
Function Classes.
We use two feed-forward neural networks as function classes in this example.
Let ϕ ( · ) denote the Batch-Norm followed by ReLU operation.
Let H be the two feed-forward layer neural network.
Any h ∈ H can be written as
[Equation]
where W 1 ∈ R 2 × 2 and W 2 ∈ R 3 × 2 .
Note that h has only two neurons and hence a very small network.
Let F be the three feed-forward layer neural network.
Any f ∈ F can be written as
[Equation]
where ˆ W 1 ∈ R 8 × 2 , ˆ W 2 ∈ R 16 × 8 and ˆ W 3 ∈ R 3 × 2 .
Note that f has 8 neurons in first and 16 neurons in the second layer.
The final layer in above networks is the classifier layer that transforms the features into the class probabilities.
Large Capacity Teacher T is a 3 layer neural network with 8, 16 and 3 neurons.
In between each feed-forward layer, we have batch-norm and ReLU activation non-linearity.
We point out that the teacher being an over-parameterized network in this feature space, easily learns the correct decision boundary.
We show this decision boundary in the Figure 2a.
We learn the teacher with cross-entropy loss.
We use the SGD optimizer with momentum 0 .
9 , learning rate 0 .
1 , weight decay 0 .
01 , and minimize the loss for 200 epochs.
Capacity Constrained Student S is a 2 layer neural network with 2 and 3 neurons.
Similar to the teacher, we have batch-norm and ReLU non-linearity in between the feed-forward layers.
Since the student is severely constrained as compared to the teacher, it suffers in learning the task.
Different training runs lead to some popular local minima.
We show the teacher solution as well as the student
local
minima in Figure 2a.
For DiSK method, our guide function g has similar capacity as the student but utilizes the teacher features to learn the decision as to which points are hard-to-learn for the student.
The contour plots for the student models under KD loss and DiSK loss are shown in
Figure 2b-2c
using the visualization toolkit described in Li et al.
We following similar setup for hyper-parameter tuning as in Sec.
A.1.
We see a similar result as in 1D example.
KD suffers from bad local minima and converges to the global minima with only 43% of the initializations.
Differently, DiSK escapes the local minima solutions and focus on the learnable part of the input space as shown in Figure 2c.
Our method converges to the global minima with very high probability ( see Table 1).
A. D ATASET D ETAILS
We use publicly available CIFAR-100 (Krizhevsky, 2009), Tiny-Imagenet (Le & Yang, 2015) and
ImageNet-1K (Russakovsky et al., 2015) datasets.
CIFAR-100 contains 50 K training and 10 K test images from 100 classes with size 32 × 32 × 3 .
While Tiny-Imagenet contains 100 K training and
10 K test images from 200 classes with size 64 × 64 × 3 .
Imagenet contains 1 . 2 M and 100 K test images from 1000 classes with size
×
× 3 .
For the CIFAR-100 and Tiny-Imagenet datasets, we use the standard data augmentations including ‘RandomCrop‘, ‘RandomHorizontalFlip‘, ‘AutoAugment‘(Cubuk et al., 2019), ‘Cutout‘ (DeVries & Taylor, 2017), and ‘Mean-Std-Normalization‘.
We use the same augmentation strategy in all our experiments, across baselines and different model families.
For the Imagenet-1K dataset, following the previous work Chen et al.
(2022), we use the ‘RandomCrop‘, ‘RandomHorizontalFlip‘, and ‘Mean-Std-Normalization‘.
A. M ODEL D ETAILS
In this section, we list the model characteristics as well as their accuracy obtained using standard cross-entropy (CE) loss.
Table 5 lists all the models used in large capacity mismatch setting.
While
Table 6 lists all the models in the small capacity mismatch setting.
Below, we describe individual model for completeness.
Large Student-Teacher Capacity Mismatch All models in the Table 5 belong to the same ResNet family and use the standard ‘BasicBlock‘ as the building block.
It consists of a convolutional block, followed by four residual block stages, followed by the adaptive average pooling layer and the classifier layer.
Different capacity models in this family differ only in the number of repetitions of the residual block and the number of filters in each stage.
Below, we write the different of repetitions and the number of filters for the four different residual stages.
• ResNet34 has [64 , 128 , 256 , 512] filters and repeats the ‘BasicBlock‘ [3 , 4 , 6 , 3] times.
• ResNet18 has [64 , 128 , 256 , 512] filters and repeats the ‘BasicBlock‘ [2 , 2 , 2 , 2] times.
• ResNet10 has [64 , 128 , 256 , 512] filters and repeats the ‘BasicBlock‘ [1 , 1 , 1 , 1] times.
• ResNet10- ℓ has [32 , 64 , 128 , 256] filters and repeats the ‘BasicBlock‘ [1 , 1 , 1 , 1] times.
• ResNet10-m has [16 , 32 , 64 , 128] filters and repeats the ‘BasicBlock‘ [1 , 1 , 1 , 1] times.
• ResNet10-s has [8 , 16 , 32 , 64] filters and repeats the ‘BasicBlock‘ [1 , 1 , 1 , 1] times.
• ResNet10-xs has [8 , 16 , 16 , 32] filters and repeats the ‘BasicBlock‘ [1 , 1 , 1 , 1] times.
• ResNet10-xxs has [8 , 8 , 16 , 16] filters and repeats the ‘BasicBlock‘ [1 , 1 , 1 , 1] times.
Small Student-Teacher Capacity Mismatch Definitions of all models in the Table 6 are borrowed from Chen et al.
We refer the reader to their official github repository ( https://github. com/DefangChen/SimKD.git ) for the exact definition.
We trained these models on our end using the data augmentations mentioned above and found that our cross-entropy baseline as well as the vanilla KD baselines are much better than the ones reported in their work.
Guide Function Our guide function g is a three layer feed-forward network.
It uses the last layer features and logits of the teacher as the input.
It has 64 , 128 , and 1 neurons in the three layers.
We include batch-norm followed by ReLU non-linearity between these layers.
The final layer contains a sigmoid activation to contain the scaler output in the range [0 , 1] .
Warm Start We note that we warm start each student model by first training them with cross entropy loss without teacher.
We observe that the warm start benefits both DiSK and KD.
Note that, we do not change the algorithms.
We only start from a CE pre-trained student model.
[Table]
[Table]
[Table]
A.
[Table]
H YPER - PARAMETERS
For both, KD and DiSK
[Table]
we scan the α hyper-parameter over the range { 0 . 0 , 0 . 1 , 0 . 5 , 0 . 9 , 1 . 0 } .
As per recommendations from previous works(Chen et al., 2022; Cho & Hariharan, 2019; Tung & Mori, 2019), we use τ = 4 as the temperature in Eq. 1.
For DiSK
[Table]
we scan the different hyper-parameters in the following ranges: (a) τ s ∈ { 1 , 2 , 4 } , (b)
K ∈ { 1 , 3 , 5 , 10 , 20 , 50 } , (c) λ min ∈ { 0 . 01 , 0 . 1 , 1 , 5 , 10 } , (d) λ max ∈ { 1 , 5 , 10 , 20 , 50 , 100 , 1000 } , (e) Budget δ within 0 .
2 distance from the cross-entropy trained student model’s error, and (f) λ T ∈ { 20 , 50 } .
We replace the arg min in the Algorithm 1, with three SGD steps over the entire dataset.
For all our experiments (both KD and DiSK), we use the popular cosine learning rate scheduler for the SGD optimizer 0 .
1 learning rate, 0 .
9 momentum and 5 e − 4 weight decay.
We use 200 as the batch size.
Note that the default hyper-parameters: τ s = τ (teacher temperature), K = 20 , λ min = 0 .
1 , λ max = 50 , δ = approximate error of the global minima (replaced by the cross-entropy error), λ T = 50 , work well in most of our experiments.
We point out that the denominator N in the budget constraint should be calibrated for the correct numerical implementation.
{+Instead of N we use the number of wrong student predictions as the normalizer, i.e., (cid:80) Ni =
+}
y i
̸ = arg max y s y
( x i ) as this is the term that appears in the Eq.
2 alongside the g term in the budget constraint.
We note that we use similar parameters to train with CE in ResNet based models.
Improving CE training would improve DiSK and KD as well since both are initialized with the CE trained model.
A. D IFFERENT E XPERIMENT S ETUPS
Scaling upto ImageNet setting.
We show that DiSK scales easily to the large-scale ImageNet-1K dataset.
We train two configurations with DiSK, namely (a) ResNet18 student and ResNet50 teacher, and (b) ViT-Tiny student and ViT-Large teacher.
We borrow these models from the timm(Wightman, 2019) library.
Table 7 shows the DiSK performance on these two configurations along with the baseline.
It clearly shows that DiSK scales well to ImageNet-1K setup and achieves significant improvements over the baselines.
A. R ELATED W ORK C ONTINUED
We continue a section of related work here due to the page constraint in the main part.
Privileged Information.
Vapnik & Izmailov (2015) propose the ‘learning under privileged information‘ (LUPI) framework wherein a support vector machine is trained using privileged information
[Table]
unavailable during the inference stage.
Later, Lopez-Paz et al.
(2016) unify LUPI and vanilla KD into generalized distillation, wherein the teacher is learned using the privileged information.
Next, the student is trained using the ground truth and the teacher labels.
These works rely on privileged information in the application domain and are shown to work on toy setups.
Since our guide function, g is available only during training, it can be thought of as privileged information from a teacher.
A. E XTENSIONS BEYOND V ANILLA KD.
In this section, we discuss potential extension of our method beyond Vanilla KD to feature based distillation as well as self-supervised distillation.
Feature based KD methods.
Let f s and f t denote the features for the student and teacher respectively and Ψ denote the operator such that Ψ( f s ) lies in the same feature space as f t .
Commonly used feature transfer strategy is to minimize the distance between these two representations via loss function such as mean-squared error as shown below by the loss L ft
[Table]
[Equation]
A simple extension of this feature alignment loss to the selective distillation is shown by the loss
L ft − g that weighs each data point with the helper function decision.
Table 8 shows this feature matching extension for the SimKD(Chen et al., 2022) scheme.
We leave question of finding better selective distillation losses in feature alignment to future work.
Self-Supervision.
We can substitute the teacher with a moving average of the student to help guide in the distillation process (a surrogate that helps in learning the hints).
Other Domains.
Although our work focuses mainly on models in the vision domain.
We point out that conceptually the function g is applicable on models in other domains albeit with some modifications.
For instance, we can utilize the censoring mechanism in sequential decision making with recurrent neural networks (Kag et al., 2020; Kag & Saligrama, 2021).
We already show that DiSK can be applied to the vision transformers in the Table 7.
We leave the application and modifications to DiSK for transformers in language domain for future research.