Reduce, Reuse, Recycle: Improving Training Efficiency with Distillation

TL;DR

This paper demonstrates that knowledge distillation can significantly accelerate training of deep neural networks like ResNet-50 and BERT, providing up to nearly 2x speed-up, with simple optimizations further enhancing efficiency.

Contribution

The study systematically investigates how distillation improves training efficiency and introduces practical methods like early-stage distillation and single-teacher sampling to maximize benefits.

Findings

Distillation speeds up ResNet-50 training by up to 1.96x.

Distillation speeds up BERT training by up to 1.42x on GLUE.

Early-stage distillation (20-50%) yields optimal results for BERT.

Abstract

Methods for improving the efficiency of deep network training (i.e. the resources required to achieve a given level of model quality) are of immediate benefit to deep learning practitioners. Distillation is typically used to compress models or improve model quality, but it's unclear if distillation actually improves training efficiency. Can the quality improvements of distillation be converted into training speed-ups, or do they simply increase final model quality with no resource savings? We conducted a series of experiments to investigate whether and how distillation can be used to accelerate training using ResNet-50 trained on ImageNet and BERT trained on C4 with a masked language modeling objective and evaluated on GLUE, using common enterprise hardware (8x NVIDIA A100). We found that distillation can speed up training by up to 1.96x in ResNet-50 trained on ImageNet and up to 1.42x on BERT when evaluated on GLUE. Furthermore, distillation for BERT yields optimal results when it is only performed for the first 20-50% of training. We also observed that training with distillation is almost always more efficient than training without distillation, even when using the poorest-quality model as a teacher, in both ResNet-50 and BERT. Finally, we found that it's possible to gain the benefit of distilling from an ensemble of teacher models, which has O(n) runtime cost, by randomly sampling a single teacher from the pool of teacher models on each step, which only has a O(1) runtime cost. Taken together, these results show that distillation can substantially improve training efficiency in both image classification and language modeling, and that a few simple optimizations to distillation protocols can further enhance these efficiency improvements.