On the Surprising Efficacy of Distillation as an Alternative to Pre-Training Small Models

TL;DR

The paper addresses the high cost of pre-training small models and proposes distilling knowledge from a large, pre-trained teacher using a contrastive-based embedding loss. It integrates dataset augmentation with synthetic samples from open diffusion models to simulate pre-training effects without touching large pre-training corpora. The authors show that this approach can achieve comparable or superior accuracy to pre-training for several vision datasets, with training time reductions up to 94%. The results suggest a practical pathway for deploying small models on resource-constrained devices while leveraging foundation-model progress.

Abstract

In this paper, we propose that small models may not need to absorb the cost of pre-training to reap its benefits. Instead, they can capitalize on the astonishing results achieved by modern, enormous models to a surprising degree. We observe that, when distilled on a task from a pre-trained teacher model, a small model can achieve or surpass the performance it would achieve if it was pre-trained then finetuned on that task. To allow this phenomenon to be easily leveraged, we establish a connection reducing knowledge distillation to modern contrastive learning, opening two doors: (1) vastly different model architecture pairings can work for the distillation, and (2) most contrastive learning algorithms rooted in the theory of Noise Contrastive Estimation can be easily applied and used. We demonstrate this paradigm using pre-trained teacher models from open-source model hubs, Transformer and convolution based model combinations, and a novel distillation algorithm that massages the Alignment/Uniformity perspective of contrastive learning by Wang & Isola (2020) into a distillation objective. We choose this flavor of contrastive learning due to its low computational cost, an overarching theme of this work. We also observe that this phenomenon tends not to occur if the task is data-limited. However, this can be alleviated by leveraging yet another scale-inspired development: large, pre-trained generative models for dataset augmentation. Again, we use an open-source model, and our rudimentary prompts are sufficient to boost the small model`s performance. Thus, we highlight a training method for small models that is up to 94% faster than the standard pre-training paradigm without sacrificing performance. For practitioners discouraged from fully utilizing modern foundation datasets for their small models due to the prohibitive scale, we believe our work keeps that door open.

Figures (7)

• Figure 1: By only mimicking the relevant slice of a pre-trained manifold, a small model can achieve the same or better performance than if it had been fully pre-trained and finetuned. Adding synthetic samples leads to better generalization.
• Figure 2: (Left) Our proposed alternative to pre-training and finetuning. (Right) The benefits of our method compared to standard approaches.
• Figure 3: Cost/accuracy comparison of our method (A/U) to supervised training (Fresh) and pre-training then linear probing. The teacher was a ResNet50 and the student was a MobileNetV2. All timing experiments were done on one NVIDIA P100 GPU.
• Figure 4: Our Alignment/Uniformity (A/U) based contrastive loss.
• Figure 5: Examples of synthetic samples generated for each dataset, (a) MIT-67, (b) DTD, (c) Caltech-101, with the respective prompt.