Knowledge distillation: A good teacher is patient and consistent

Introduction

Large-scale vision models currently dominate many areas of computer vision. Recent state-of-the-art models for image classification , object detection or semantic segmentation push model size to the limits allowed by modern hardware. Despite their impressive performance, these models are rarely used in practice due to high computational costs. Instead, practitioners typically use much smaller models, such as ResNet-50 or MobileNet , which are order(s) of magnitude cheaper to run. According to the download counts of five BiT models from Tensorflow Hub, the smallest ResNet-50 model has been downloaded for significantly more times than the larger ones. As a result, many recent improvements in vision do not translate to real-world applications.

To address this problem, we concentrate on the following task: given a specific application and a large model that performs very well on it, we aim to compress the model to a smaller and more efficient architecture without compromising performance. There are two widely used paradigms that target this task: model pruning and knowledge distillation . Model pruning reduces the large model’s size by stripping away its parts. This procedure can be restrictive in practice: first, it does not allow changing the model family, say from a ResNet to a MobileNet. Second, there may be architecture-dependent challenges, e.g. if the model uses group normalization , pruning channels may result in the need to dynamically re-balance channel groups.

Instead, we concentrate on the knowledge distillation approach which does not suffer from these drawbacks. The idea behind knowledge distillation is to “distill” a teacher model, in our case a large and cumbersome model or ensemble of models, into a small and efficient student model. This works by forcing the student’s predictions (or internal activations) to match those of the teacher, thus naturally allowing a change in the model family as part of compression. We closely follow the original distillation setup from and find it surprisingly effective when done right: We interpret distillation as a task of matching the functions implemented by the teacher and student, as illustrated in Figure 2. With this interpretation, we discover two principles of knowledge distillation for model compression. First, teacher and student should process the exact same input image views or, more specifically, same crop and augmentations. Second, we want the functions to match on a large number of support points to generalize well. Using an aggressive variant of mixup , we can generate support points outside the original image manifold. With this in mind, we experimentally demonstrate that consistent image views, aggressive augmentations and very long training schedules are the key to make model compression via knowledge distillation work well in practice.

Despite the apparent simplicity of our findings, there are multiple reasons that may commonly prevent researchers (and practitioners) from making the design choices that we suggest. First, it is tempting to precompute the teacher’s activations for an image offline once to save compute, especially for very large teachers. As we will show, this fixed teacher approach does not work well. Second, knowledge distillation is also commonly used in different contexts (other than model compression), where authors recommend different or even opposite design choices , see Figure 2. Third, knowledge distillation requires an atypically large number of epochs to reach best performance, much more than commonly used for supervised training. Finally, choices which may look suboptimal in training of regular length often end up being best for long runs, and vice-versa.

In our empirical study, we mostly concentrate on compressing the large BiT-ResNet-152x2 from that was pretrained on the ImageNet-21k dataset and fine-tuned to the relevant datasets of interest. We distill it to a standard ResNet-50 architecture (but replace batch normalization with group normalization) on a range of small and mid-sized datasets without compromising accuracy. We also achieve very strong results on the ImageNet dataset: with a total number of 9600 epochs for distillation, we set the new ResNet-50 SOTA 82.8% on ImageNet. This is 4.4% better than the ResNet-50 model from , and 2.2% better than the best ResNet-50 model in the literature, which uses a more complex setup . Finally, we demonstrate that our distillation recipe also works when simultaneously compressing and changing the model family, e.g. BiT-ResNet architecture to the MobileNet architecture.

Experimental setup

In this section, we introduce the experimental setup and benchmarks we use throughout the paper. Given a large-scale vision model (the teacher, or T) with high accuracy on a particular task, we aim to compress this model to a much smaller model (the student, or S) without compromising its performance. Our compression recipe relies on knowledge distillation, as introduced in , and a careful investigation of several key ingredients in the training setup.

Datasets, metrics and evaluation protocol. We conduct experiments on five popular image classification datasets: flowers102 , pets , food101 , sun397 and ILSVRC-2012 (“ImageNet”) . These datasets span diverse image classification scenarios; In particular, they vary in the number of classes, from $37$ to $1000$ classes, and total number of training images, from $1020$ to $1281167$ training images. This allows us to verify our distillation recipe for a broad range of practical settings and ensure its robustness.

As a metric, we always report classification accuracy. For all datasets, we perform design choices and hyperparameters selection using a validation split, and report final results on the test set. These splits are defined in the appendix E.

Teacher and student models. Throughout the paper, we opt for using pre-trained teacher models from BiT , which provides a large collection of ResNet models pretrained on ILSVRC-2012 and ImageNet-21k datasets, with state-of-the-art accuracy. The only significant differences between BiT-ResNets and standard ResNets is their use of group normalization layer and weight standardization , which are used instead of batch normalization .

In particular, we concentrate on the BiT-M-R152x2 architecture: a BiT-ResNet-152x2 (152 layers, ‘x2’ indicates the width multiplier) pretrained on ImageNet-21k. This model demonstrates excellent performance on a variety of vision benchmarks and it is still manageable to run extensive ablation studies with it. It is highly expensive to deploy (requires roughly 10x more compute than the standard ResNet-50), and thus effective compression of this model is of practical importance. For the student’s architecture, we use a BiT-ResNet-50 variant, referred to as ResNet-50 for brevity.

Distillation loss. We use the KL-divergence between the teacher’s $p_{t}$, and the student’s $p_{s}$ predicted class probability vectors as a distillation loss, as was originally introduced in . We do not use any additional loss term with respect to the original dataset’s hard labels:

where $\mathcal{C}$ is a set of classes. Also, as in , we introduce a temperature parameter $T$, which is used to adjust the entropy of the predicted softmax-probability distributions before they are used in the loss computation: $p_{s}\propto\exp(\frac{\log p_{s}}{T})$ and $p_{t}\propto\exp(\frac{\log p_{t}}{T})$.

Training setup. For optimization, we train our models with the Adam optimizer with default parameters, except for the initial learning rate that is part of our hyperparameter exploration. We use a cosine learning rate schedule without warm restarts. We also sweep over the weight decay loss coefficient for all our experiments (for which we use a “decoupled” weight decay convention ). To stabilize training we enable gradient clipping with a threshold of 1.0 on the global L2-norm of a gradient. Finally, we use batch size 512 for all our experiments, except for models trained on ImageNet, where we train with batch size 4096. For the remaining hyperparameters, we discuss their sweeping range together with corresponding experiments in the next section.

One additional important component of our recipe is the mixup data augmentation strategy . In particular, we introduce a mixup variant in our “function matching” strategy (see Section • ‣ 3.1.1), in which we use “agressive” mixing coefficients sampled uniformly from $$, which can be seen as an extreme case of the originally proposed sampling from$\beta$-distribution.

Unless explicitly specified otherwise, for prepossessing we use an “inception-style” crop and then resize images to a fixed square size. Furthermore, in order to make our extensive analysis computationally feasible (overall we trained dozens of thousands of models), we use relatively low input resolution and resize input images to $128\times 128$ size, except for our ImageNet experiments, that use the standard input $224\times 224$ resolution.

For all our experiments we use Google Cloud TPU accelerators . We also report our batch sizes, epochs or total number of update steps, which allow to estimate resource requirements for any particular experiment of interest. Code and checkpoints are made publicly available .

Distillation for model compression

In this section, we perform an experimental verification of our hypothesis formulated in the introduction and visualised in Figure 2, that distillation works best when seen as function matching, i.e. when the student and teacher see consistent views of the input images, synthetically ”filled” via mixup, and when student is trained using long training schedule (i.e. the “teacher” is patient).

To make sure that our findings are robust, we perform a very thorough analysis on four small and medium scale datasets, namely Flowers102 (1020 training images), Pets (3312 training images), Food101 (about 68k training images), and SUN397 (76k training images).

In an effort to remove any confounding factors, for each individual distillation setting we sweep over all combinations of learning rates $\{0.0003,0.001,0.003,0.01\}$, weight decays $\{1\cdot 10^{-5},3\cdot 10^{-5},1\cdot 10^{-4},3\cdot 10^{-4},1\cdot 10^{-3}\}$, and distillation temperatures $\{1,2,5,10\}$. In all reported figures, we show every single run as a low opacity curve, and highlight the one with the best final validation accuracy. We provide corresponding test accuracies in Appendix A.

First, we demonstrate that the consistency criterion, i.e. student and teacher seeing the same views, is the only way of performing distillation which reaches peak student performance across all datasets consistently. For this study, we define multiple distillation configurations which correspond to instantiations of all four options sketched in Figure 2, with the same color coding:

Fixed teacher. We explore several options where the teacher’s predictions are constant for a given image (precomputed target). The simplest (and worst) method is fix/rs, where the image is just resized to $224^{2}$px for both student and teacher. fix/cc follows a more common approach of using a fixed central crop for the teacher and a mild random crop for the student. fix/ic_ens is a heavy data augmentation approach where the teacher’s prediction is the average of 1k inception crops, which we verified to improve the teacher’s performance. The student also uses random inception crops. The two latter settings are similar to the input noise strategy from the “noisy student” paper .

Independent noise. We instantiate this common strategy in two ways: ind/rc computes two independent mild random crops for the teacher and student respectively, while ind/ic uses the heavier inception crop instead. A similar setup was used in .

Consistent teaching. In this approach, we randomly crop the image only once, either with mild random cropping (same/rc) or heavy inception crop (same/ic), and use this same crop for the input to both the student and the teacher.

Function matching. This approach extends consistent teaching, by expanding an input manifold of images through mixup (mix), and, again, providing consistent inputs to the student and the teacher. For brevity, we sometimes refer to this approach as “FunMatch”.

Figure 3 shows 10 000 epoch training curves on Flowers102 dataset in all of these configurations. These results clearly show that “consistency” is the key: all “inconsistent” distillation settings plateau at a lower score, while consistent settings increase student performance significantly, with the function matching approach working the best. Furthermore, the training losses show that, for such small datasets, using a fixed teacher leads to strong overfitting. In contrast, function matching never reaches such loss on the training set while generalizing much better to the validation set. Due to space constraints, we show analogous results for other datasets and training durations in Appendix C.

1.2 Importance of “patient” teaching

One can interpret distillation as a variant of supervised learning, where labels (potentially soft) are provided by a strong teacher model. This is especially true when the teacher predictions are (pre)computed for a single image view. This approach inherits all problems of the standard supervised learning, e.g. aggressive data augmentations may distort actual image label, while less aggressive augmentations may cause overfitting.

However, things change if we interpret distillation as function matching, and, crucially, make sure to provide consistent inputs to the student and teacher. In this case we can be very aggressive with image augmentations: even if an image view is too distorted, we still will make a progress towards matching the relevant functions on this input. Thus, we can be more opportunistic with augmentations and avoid overfitting by doing aggressive image augmentations and, if true, optimize for very long time until the student’s function comes close to the teacher’s.

We empirically confirm our intuition in Figure 4, where for each dataset we show the evolution of test accuracy during training of the best function matching student (according to validation), for different amounts of training epochs. The teacher is shown as a red line and is always reached eventually, after a much larger number of epochs than one would ever use in a supervised training setup. Crucially, there is no overfitting even when we optimize for a 1M epochs.

We also trained and tuned two more baselines for reference: training a ResNet-50 from scratch using the dataset original hard labels, as well as transferring a ResNet-50 that was pre-trained on ImageNet-21k. For both of these baselines, we heavily tune learning rate and weight decay as described in Section 3.1. The model trained from scratch using the original labels is substantially outperformed by our student. The transfer model fares much better, but is eventually also outperformed.

Notably, training for a relatively short but common duration of 100 epochs leads to much worse performance than the transfer baseline. Overall, the ResNet-50 student patiently and consistently matches the very strong but much more expensive ResNet-152x2 teacher across the board.

2 Scaling up to ImageNet

Based on our insights from the previous sections, we now investigate how the proposed distillation recipe scales to the widely used and more challenging ImageNet dataset . Following the same protocol as before, in Figure 5 (left), we report student accuracy curves throughout training for three distillation settings: (1) fixed teacher, (2) consistent teaching and (3) function matching. For reference, our base teacher model reaches a top-1 accuracy of 83.0%. Fixed teacher again suffers from long training schedules, and starts overfitting after 600 epochs. In contrast, the consistent teaching approaches continuously improves performance as the training duration increases. From this we can conclude that consistency is a key to make distillation work on ImageNet, similar to the behaviors on the previously discussed small and mid-sized datasets.

Compared to simple consistent teaching, function matching performs slightly worse with short schedules, which likely happens due to underfitting. But when we increase the length of training schedule, the improvement of function matching becomes apparent: for instance with only 1200 epochs, it is able to match the performance of consistent teaching at 4800 epochs, thus saving 75% compute resource. Finally, for the longest run of function matching we experimented on, the vanilla ResNet-50 student architecture achieves 82.31% top-1 accuracy on ImageNet.

3 Distilling across different input resolutions

So far, we have assumed that both the student and teacher receive the same standard input resolution of 224px. However, it is possible to pass images of different resolution to the student and the teacher, while still being consistent: one simply has to perform the crop on the original high-resolution image, and subsequently resize it differently for the student and the teacher: their views will be consistent, albeit at different resolutions. This insight can be leveraged for learning from a better, higher resolution, teacher , but also for training a smaller, faster student . We investigate both directions: first, following , we train a ResNet-50 student with an input resolution of 160px while leaving the teacher resolution unchanged (224px). This results in a twice faster model, which still achieves remarkable 80.49% top-1 accuracy (see Table 1), compared to the best published 78.8% at this resolution using an array of modifications .

Second, following , we distill a teacher that was fine-tuned at a resolution of 384px (and attains 83.7% top-1 accuracy), this time leaving the student resolution unchanged, i.e. consuming a 224px input image. Compared to the baseline teacher, this provides a modest but consistent improvement across the board, as shown in Table 1.

4 Optimization: A second order preconditioner () improves training efficiency

We observe that optimization efficacy creates a computational bottleneck for our distillation recipe with “function matching” perspective due to long training schedules. Intuitively, we believe that optimization difficulties stem from the fact that it is much harder to fit a general function with multivariate outputs, rather than fixed image-level labels. Thus, we conduct an initial exploration, whether more powerful optimizers can do a much better job at our task.

To this end, we change the underlying optimizer from Adam to Shampoo , with the second order preconditioner. In Figure 5 (middle) we observe that Shampoo achieves the same test accuracy reached by Adam at 4800 epochs in just 1200 epochs, and with minimal step time overhead. And, in general, we observe consistent improvement over Adam in all our experimental settings. Experimental details on the Shampoo optimizer are provided in the Appendix D.

5 Optimization: A good initialization improves short runs but eventually falls behind

Motivated by transfer learning literature and , where a good initialization is able to significantly shorten the training cost and achieve a better solution, we try to initialize the student model with a pre-trained BiT-M-ResNet50 weights and show the results in Figure 5 (right).

The BiT-M initialization improves more than 2% when the distillation duration is short. However, the gap closes when the training schedule is long enough. Our observation is similar to the conclusion of . Starting from 1200 epochs, distilling from scratch matches the BiT-M initializated student, and slightly overtakes it for 4800 epochs.

6 Distilling across different model families

Going beyond using different input resolutions for student and teacher, nothing in principle prevents us from using architectures of different families entirely, as our consistent patient teacher approach still applies in this setting. This allows us to efficiently transfer knowledge from stronger and more complex teachers, e.g. ensembles, while keeping the simple architecture of a ResNet50 student, but also transfer the state-of-the-art performance of large ResNet models to more efficient architectures e.g. MobileNet. We demonstrate this via two experiments. First, we use an ensemble of two models as teacher and show that this further improves performance. Second, we train a MobileNet v3 student and obtain the best reported MobileNet v3 model to date.

MobileNet student. We use MobileNet v3 (Large) as a student, for most experiments we opt for the variant which uses GroupNorm (with the default of 8 groups) instead of BatchNorm. We do not use any of the training tricks used in the original paper,we simply perform function matching. Our student reaches 74.60% after 300 epochs, and 76.31% after 1200 epochs, resulting in the best published MobileNet v3 model. More results are in the Appendix A.

Ensemble teacher. We now try a better teacher: we create a model which consists of averaging the logits from our default teacher at 224px resolution, and our teacher at 384px resolution from the previous section. This is a different, though closely related, type of teacher which is significantly more powerful but also slower. This teacher’s student is better than our default teacher’s student at every duration we tried (Appendix A) and, after 9600 epochs, reaches a new state-of-the-art top-1 ImageNet accuracy of 82.82%.

7 Comparison to the results from literature.

Now, when we introduced our key experiments, we compare our best ResNet-50 models to the best ResNet-50 models available in the literature, see Table 2. In particular, for $224\times 224$ input resolution we compare against the original ResNet-50 model from , BiT-M-ResNet-50 pretrained on ImageNet-21k dataset and previous state-of-the-art model from . For $160\times 160$ input resolution we compare against very recent and competitive model from . We observe that our distillation recipe leads the state-of-the-art performance in both cases and by a significant margin.

8 Distilling on the ”out-of-domain” data

By looking at knowledge distillation as “function matching”, one can draw a reasonable hypothesis that distillation can be done on arbitrary image inputs. In this section we investigate this hypothesis.

We conduct experiments on pets and sun397 datasets. We use our distillation recipe to distill pets and sun397 models using out-of-domain images from the food101 and ImageNet datasets and, for the reference results, also run distillation with “in-domain” images from pets and sun397 datasets.

Figure 6 summarizes our results. First we observe that distilling using in-domain data works the best. Somewhat surprisingly, even if the images are completely unrelated, distillation still works to some extent, though results get worse. This, for example, means that the student model can learn to classify pets with roughly 30% accuracy by only seeing food images (softly) labeled as breeds of pets. Finally, if distillation images are somewhat related or overlapping with the actual “in-domain” images (e.g. Pets and ImageNet, or sun397 and ImageNet), then results can be as good (or almost as good) as using “in-domain” data, but extra long optimization schedule may be required.

9 Finetuning ResNet-50 with augmentations

To make sure that our observed state-of-the-art distillation results are not an artifact of our well-tuned training setup, namely very long schedule and aggressive mixup augmentations, we train corresponding baseline ResNet-50 models. More specifically, we reuse the distillation training setup for supervised training on ImageNet dataset without distillation loss. To further strengthen our baseline, we additionally try SGD optimizer with momentum, which is known to often work better for ImageNet than Adam optimizer.

Results are shown in Figure 7. We observe that training with labels and without distillation loss leads to significantly worse results and starts to overfit for long training schedules. Thus, we conclude that distillation is necessary to make our training recipe work well.

Related work

There are many paradigms for compressing neural networks. One of them is pruning, where the general idea is to discard parts of the trained model while making it much more efficient and incurring little or no sacrifise in performance. Model pruning comes in many different flavours: it can be unstructured (i.e. focus on pruning individual connections) or structured (i.e. focus on pruning larger building blocks, e.g. whole channels). It can also come with or without an additional finetuning step, or be iterative or not. Balanced and fair discussion of this topic goes beyond the scope of this paper, so we refer interested reader to recent overviews as a starting point .

Knowledge distillation is a technique for transferring knowledge from one model (teacher) to another (student), by optimizing a student model to match certain outputs (or intermediate activations) of a teacher model. This technique is used in numerous distinct contexts, such as semi-supervised learning or even self-supervised learning . In this paper we only consider knowledge distillation as a tool for model compression. The efficiency of distillation has been showcased in numerous works, e.g. , under different depth/width patterns of the student and teacher architectures, and even combined with other compression techniques . Notably, MEAL proposes to distill an ensemble of large ResNet teachers to a smaller ResNet student with an adversarial loss and achieves strong results. The main difference of our work to similar works on knowledge distillation for compression, is that our method is simultaneously the simplest and best performing: we do not introduce any new components, but rather discover that correct training setup is sufficient to attain state-of-the art results.

Weights quantization and decomposition aim to accelerate and reduce the memory footprint of CNNs by replacing large matrices operations with their lightweight approximations. This line of research is largely orthogonal to this work and can generally be combined with the method from this paper, especially during for the final model deployment stage. We leave exploration of this topic for future research.

Finally, there is a line of work, which approaches our goal (compact and high performing models) from a different angle, by focusing on altering the architecture and getting good compact models trained from scratch, so there is no need to compress large models. Some notable examples include ResNeXt , Squeeze-and-Excitation Networks and Selective Kernel , which propose modifications that improve model accuracy for a fixed compute budget. These improvements are complementary to the research question tackled in this paper and can be compounded.

Conclusion

Instead of proposing a new method for model compression, we closely look at the existing common knowledge distillation process and identify how to make it work really well in the context of model compression. Our key findings stem from a specific interpretation of knowledge distillation: we propose to see it as a function matching task. This is not the typical view of knowledge distillation, as normally it is seen as “a strong teacher generates better (soft) labels that are useful for training a better and smaller student model”.

Based on our interpretation we simultaneously incorporate three ingredients: (i) make sure that teacher and student always get identical inputs, including noise, (ii) introduce aggressive data augmentations to enrich the input image manifold (through mixup) and (iii) use very long training schedules. Even though each component of our recipe may seem trivial, our experiments show that one has to apply all of them jointly to get top results.

We attain very strong empirical results for compressing very large models to the more practical ResNet-50 architecture. We believe that they are very useful from a practical point of view and are a very strong baseline for future research on compressing large-scale models.

Acknowledgements. We thank Daniel Keysers and Frances Hubis for their valuable feedback on this paper; Ilya Tolstikhin and the Google Brain team at large for providing a supportive research environment.

References

Appendix A Full results tables

We provide a full summary of our experiments on ImageNet in Table 3, with a dash “-” marking settings we did not deem necessary to run, as the cost outweights the potential insights.

Furthermore, Table 4 gives numerical results for the models shown in Figure 4, our best models (according to validation) on the four smaller datasets at $128$px resolution, together with baselines and the teacher.

Appendix B BiT models download statistics

In Figure 8, we show the download statistics for models with different sizes: ResNet50, ResNet50x3, ResNet101, ResNet101x3 and ResNet152x4. It’s clear that the smallest ResNet50 model is the most used, with a significant gap compared to the other models. The practitioners’ behavior motivates our work of getting the best possible ResNet50 model.

Appendix C More consistency plots

In Figures 9 to 12, we show the “consistency” plots (cf Figure 3 in the main paper) for all datasets and across all training durations. It is noteworthy that (relatively) short runs may provide deceptive signal on the best method, and only with the addition of “patience”, e.g. when distilling for a long time, does it become clear that the full function-matching approach is the best choice.

Appendix D Shampoo optimization details

For all experiments the learning rate schedule was a linear warm-up up to 1800 steps followed by a quadratic decay towards zero. Overhead of Shampoo is quite minimal due blocking trick (each preconditioner is atmost 128x128) and inverse is run in a distributed manner across the TPU cores every step, with nesterov momentum. These settings are identical to the the training recipe in for training a ResNet-50 architecture on ImageNet from scratch efficiently at large batch sizes. All experiments uses weight decay of 0.000375.

Appendix E Training, validation and test splits

Throughout our experiments we rely on the tensorflow datasets libraryhttps://www.tensorflow.org/datasets to access all datasets. A huge advantage of this library is that it enables a unified and reproducible way to access diverse datasets. To this end, we report our train, validation and test splits (following the library’s notation) in Table 5.

Appendix F Configuration file for ImageNet distillation

We present the configuration for performing distillation on ImageNet following the big_vision conventions.