Distill on a Diet: Efficient Knowledge Distillation via Learnable Data Pruning

Knowledge Distillation (KD) is a standard technique for compressing large teacher models into efficient student models. However, the distillation process itself is computationally expensive, often requiring the student to train on the entire dataset. This paper investigates whether data pruning—selecting a smaller, high-quality subset of training data—can maintain or improve student performance while significantly reducing the computational cost of the distillation process.

The authors propose IF-Beta, a framework that addresses two primary challenges in KD-oriented data pruning: the high cost of estimating sample importance and the rigidity of heuristic selection rules.

1. Efficient Scoring: Instead of relying on expensive training dynamics (like retraining models to track forgetting or loss trajectories), the authors use Influence Functions (IF) combined with Sharpness-Aware Minimization (SAM). This allows them to estimate the impact of each training sample on the teacher's performance using only a pretrained teacher and a small validation set.
2. Learnable Sampling: Rather than using fixed rules (like top-k or sliding windows), the authors parameterize the selection process using a Beta distribution. This allows the model to learn an adaptive sampling policy that can adjust to different data distributions and pruning ratios.
3. Bilevel Optimization: The framework is trained via a bilevel objective. The inner loop performs fast proxy training of a linear classifier on frozen teacher features, while the outer loop optimizes the Beta distribution parameters to maximize the student's performance on a validation set.

Empirical results on CIFAR-10, CIFAR-100, and ImageNet demonstrate that IF-Beta consistently outperforms existing pruning baselines. By effectively identifying the most informative samples, the method allows student models to achieve higher accuracy than those trained on the full dataset, even when using significantly fewer training samples and less total compute. The authors show that their IF-based scoring correlates strongly with established difficulty metrics, validating its effectiveness as a post-hoc estimator.