Source-linked AI summary

Knowledge distillation: A good teacher is patient and consistent

Lucas Beyer, Xiaohua Zhai, Amélie Royer, Larisa Markeeva, Rohan Anil, Alexander Kolesnikov

arXiv:2106.05237v2cs.CVcs.AIcs.LG

TL;DR

Large vision models are difficult to deploy because of their computational cost, motivating compression into practical architectures. The paper empirically identifies a robust knowledge-distillation recipe centered on consistent teacher-student views and patient training. It reports strong results across datasets, including an 82.8% top-1 ImageNet ResNet-50.

  • Problem

    Large-scale vision models achieve state-of-the-art performance but are often too computationally expensive for practical use.

  • Method

    The paper empirically investigates knowledge distillation as a compression recipe, emphasizing consistent image views, function matching, and long training schedules.

  • Results

    82.8% top-1 accuracy was achieved by a distilled ResNet-50 on ImageNet, with strong results across a range of vision datasets.

  • Takeaways & Limitations

    Carefully designed knowledge distillation can make large-scale vision models more practical while preserving performance and allowing changes in model family.

  • Takeaways & Limitations

    Distillation requires atypically long training schedules, and choices that work for regular-length training may differ from those best for long runs.

Abstract

from arXiv · show

There is a growing discrepancy in computer vision between large-scale models that achieve state-of-the-art performance and models that are affordable in practical applications. In this paper we address this issue and significantly bridge the gap between these two types of models. Throughout our empirical investigation we do not aim to necessarily propose a new method, but strive to identify a robust and effective recipe for making state-of-the-art large scale models affordable in practice. We demonstrate that, when performed correctly, knowledge distillation can be a powerful tool for reducing the size of large models without compromising their performance. In particular, we uncover that there are certain implicit design choices, which may drastically affect the effectiveness of distillation. Our key contribution is the explicit identification of these design choices, which were not previously articulated in the literature. We back up our findings by a comprehensive empirical study, demonstrate compelling results on a wide range of vision datasets and, in particular, obtain a state-of-the-art ResNet-50 model for ImageNet, which achieves 82.8% top-1 accuracy.

1. Introduction

Large vision models deliver strong performance but are often too expensive for practical deployment. The paper studies knowledge distillation as a flexible way to compress them, identifying training and data-view choices associated with strong results.

  • Motivation: Large-scale vision models achieve state-of-the-art results but are rarely used in practice because of high computational costs.Smaller models such as ResNet-50 and MobileNet are orders of magnitude cheaper to run.
  • Problem and approach: The compression goal is to convert a high-performing large model into a smaller, more efficient architecture without compromising performance.The paper focuses on knowledge distillation rather than pruning because distillation can change the model family.
  • Problem and approach: Knowledge distillation forces a student model’s predictions or internal activations to match those of a large teacher or ensemble.This formulation naturally permits changing the student architecture during compression.
  • Key findings: The study identifies patient training, consistent teacher-student image views, and other implicit design choices as important for effective distillation.Precomputing teacher activations can perform poorly, and choices that seem suboptimal on short runs may work best over long schedules.
  • Results: 82.8% top-1 accuracy was achieved by a distilled ResNet-50 on ImageNet, establishing a new state-of-the-art result.This result used 9600 distillation epochs; the model was also 2.2% better than the best ResNet-50 model in the literature cited by the paper.
  • Results: The distillation recipe also works when compressing a BiT-ResNet model while changing the model family to MobileNet.The empirical study covers small and mid-sized datasets as well as ImageNet.

2. Experimental setup

The experiments evaluate a knowledge-distillation recipe for compressing large vision models into smaller architectures across diverse image-classification settings. The setup examines teacher-student consistency, mixup, optimization, and distillation-loss choices.

  • Models and objective: The experiments compress a high-accuracy large teacher into a much smaller student without compromising performance.The study primarily uses a BiT-ResNet-152x2 teacher and a standard ResNet-50 student with group normalization.
  • Datasets and evaluation: The evaluation covers flowers102, pets, food101, sun397, and ImageNet, spanning 37 to 1000 classes and 1020 to 1281167 training images.This dataset range is used to assess robustness across practical image-classification settings.
  • Datasets and evaluation: The study uses validation splits for design and hyperparameter selection and reports classification accuracy on test sets.Input resolution is generally 128 × 128, except for ImageNet experiments at 224 × 224.
  • Models and objective: The teacher models are pretrained BiT-ResNets using group normalization and weight standardization instead of batch normalization.The main teacher is a 152-layer, x2-width BiT-ResNet pretrained on ImageNet-21k.
  • Models and objective: The distillation loss is the KL-divergence between teacher and student class-probability vectors, with a temperature parameter adjusting softmax-distribution entropy.The setup does not add a loss term for the original dataset’s hard labels.
  • Distillation design: The recipe investigates consistent teacher-student views, long training schedules, and mixup-based function matching.Its mixup variant samples aggressive mixing coefficients uniformly from [0, 1].

3. Distillation for model compression

The study tests knowledge distillation as a practical compression recipe, finding that consistent teaching and sufficiently long training are central to matching strong teachers. This recipe yields strong ImageNet and cross-dataset results while supporting changes in input resolution and model family.

  • 3.1. Distillation for model compression: Consistency—student and teacher seeing the same views—is the only tested criterion that consistently reaches peak student performance across datasets.The study evaluates multiple distillation configurations across Flowers102, Pets, Food101, and SUN397.
  • 3.1. Distillation for model compression: Fixed-teacher settings plateau lower and overfit, whereas function matching improves student performance and generalizes better.On Flowers102, function matching avoids the strong training overfitting observed with fixed teachers.
  • 3.1. Distillation for model compression: Long schedules eventually let the student match the teacher across datasets without overfitting, while 100 epochs performs much worse than the transfer baseline.The teacher is reached only after substantially more training than typical supervised schedules.
  • 3.2. Scaling up to ImageNet: 1200 epochs of function matching match consistent teaching at 4800 epochs, saving 75% compute, while the longest run reaches 82.31% top-1 ImageNet accuracy.Function matching can underfit on short schedules but improves with longer training.
  • 3.3. Distilling across different input resolutions: The recipe produces 80.49% top-1 accuracy for a twice-faster 224px model and supports stronger 384px teachers with unchanged 224px student inputs.The 384px teacher reaches 83.7% top-1 accuracy and provides a modest, consistent improvement across experiments.
  • 3.6. Distilling across different model families: Distillation transfers across model families and ensemble teachers, including strong results for MobileNet v3 and a new 82.82% top-1 ImageNet result from an ensemble teacher.The ensemble averages logits from 224px and 384px teachers and outperforms the default teacher’s student at every tested duration.
  • 3.7. Comparison to the results from literature: The distillation recipe leads state-of-the-art performance against published ResNet models on both evaluated input resolutions.The comparison uses ImageNet test accuracy for the best models and literature baselines.
  • 3.8. Distilling on the “out-of-domain” data: In-domain distillation data works best, while related or overlapping data can approach it with extra-long optimization and unrelated data performs worse.Even unrelated images retain some transfer ability, such as learning pet classification from food images labeled as pet breeds.

4. Related work

The paper situates its compression approach among pruning, knowledge distillation, lightweight approximations, and compact architectures. It emphasizes knowledge distillation as a simple, complementary route to high-performing smaller models.

  • Knowledge distillation transfers knowledge from a teacher to a student by matching outputs or intermediate activations.
  • Unlike pruning, distillation allows the student to change model family while matching the teacher’s predictions or activations.
  • The paper claims its compression method is simultaneously simpler and better performing because it introduces no new components.
  • Quantization and decomposition reduce computation or memory through lightweight matrix approximations and are generally complementary to this work.
  • Compact architectures such as ResNeXt, Squeeze-and-Excitation, and Selective Kernel improve accuracy at fixed compute budgets and can be combined with compression.

5. Conclusion

The conclusion presents knowledge distillation as function matching and identifies consistency, aggressive augmentation, and long schedules as jointly necessary ingredients. The resulting recipe compresses large models to practical architectures while retaining strong performance.

  • The paper reframes knowledge distillation for compression as a function matching task rather than primarily soft-label generation.
  • The recipe uses identical teacher and student inputs, aggressive mixup augmentation, and very long training schedules.
  • All three recipe components must be applied jointly to obtain top results in the experiments.
  • The approach achieves strong empirical results when compressing very large models to the practical ResNet-50 architecture.

A. Full results tables

The appendix summarizes ImageNet experiments and reports smaller-dataset results alongside baselines and the teacher. These tables provide the numerical results underlying the paper’s model comparisons.

  • Table 3 provides a full summary of the ImageNet experiments, using dashes for settings judged unnecessary to run.
  • The appendix omits some experiments when their computational cost outweighs the potential insights.
  • Table 4 reports numerical results for Figure 4 models, including best validation-selected models on four smaller datasets, baselines, and the teacher.

B. BiT models download statistics

Download statistics compare BiT models of different sizes and show that ResNet-50 is used most often. This usage gap motivates optimizing ResNet-50 performance.

  • ResNet-50, ResNet-50x3, ResNet-101, ResNet-101x3, and ResNet-152x4 are compared using model download statistics.
  • ResNet-50 is the most downloaded model, with a significant gap over the larger alternatives.
  • Practitioner usage motivates the paper’s goal of obtaining the best possible ResNet-50 model.

C. More consistency plots

Consistency plots across datasets and training durations show that short runs can misidentify the best distillation method, while longer training reveals full function matching as best.

  • Consistency plots cover all datasets and multiple training durations.
  • Short training runs may provide a deceptive signal about the best method.
  • Longer distillation makes full function matching emerge as the best choice.The paper characterizes this longer training as adding “patience.”

D. Shampoo optimization details

The experiments use a warm-up followed by quadratic learning-rate decay, while Shampoo overhead is kept minimal through blocking and distributed inversion.

  • The learning rate warms up linearly to 1800 steps, then decays quadratically toward zero.
  • Shampoo overhead is minimal because each preconditioner is at most 128x128 and inversion runs distributed across TPU cores every step.
  • All experiments use Nesterov momentum and settings identical to an efficient large-batch ResNet50 ImageNet training recipe.

E. Training, validation and test splits

The study uses TensorFlow Datasets to access diverse datasets reproducibly and reports their train, validation, and test splits according to the library’s notation.

  • All experiments rely on the TensorFlow Datasets library for dataset access.
  • The library provides a unified and reproducible way to access diverse datasets.
  • Train, validation, and test splits are reported in Table 5 using the library’s notation.

F. Configuration file for ImageNet distillation

The ImageNet distillation configuration follows Big Vision conventions and includes preprocessing, KL-based distance, and 1000 output classes; the accompanying tables and figures document the setup and consistency analyses.

  • The ImageNet distillation configuration follows Big Vision conventions.
  • The preprocessing pipeline decodes JPEG images, applies an Inception crop, and flips images horizontally.
  • Listing 1 presents the full ImageNet distillation configuration, with evaluation details shortened and deferred to the code release.
  • Table 3 summarizes ImageNet validation top-1 accuracy for a ResNet50 student and BiT-M-R152x2 teacher.
  • Figures 9–12 and Tables 4–5 provide consistency-plot results and tabulated result or split representations.
Loading 2106.05237v2…