InDistill: Information flow-preserving knowledge distillation for model compression

TL;DR

InDistill is a novel knowledge distillation method that enhances model compression by preserving critical information flow paths through curriculum learning and adaptive pruning, leading to improved student model performance.

Contribution

It introduces a warmup stage for KD that focuses on critical information flow transfer and employs pruning to match teacher-student layer widths, broadening applicability.

Findings

Consistently improves baseline KD performance across datasets.

Effective in classification and retrieval tasks.

Applicable to various teacher-student architectures.

Abstract

In this paper, we introduce InDistill, a method that serves as a warmup stage for enhancing Knowledge Distillation (KD) effectiveness. InDistill focuses on transferring critical information flow paths from a heavyweight teacher to a lightweight student. This is achieved via a training scheme based on curriculum learning that considers the distillation difficulty of each layer and the critical learning periods when the information flow paths are established. This procedure can lead to a student model that is better prepared to learn from the teacher. To ensure the applicability of InDistill across a wide range of teacher-student pairs, we also incorporate a pruning operation when there is a discrepancy in the width of the teacher and student layers. This pruning operation reduces the width of the teacher's intermediate layers to match those of the student, allowing direct distillation without the need for an encoding stage. The proposed method is extensively evaluated using various pairs of teacher-student architectures on CIFAR-10, CIFAR-100, and ImageNet datasets demonstrating that preserving the information flow paths consistently increases the performance of the baseline KD approaches on both classification and retrieval settings. The code is available at https://github.com/gsarridis/InDistill.