Robust and Resource-Efficient Data-Free Knowledge Distillation by Generative Pseudo Replay

TL;DR

This work tackles data-free knowledge distillation when no validation data is available, addressing catastrophic forgetting caused by distribution shifts in synthetic data. It introduces PRE-DFKD, which employs two specialized generators and a memory-aware Variational Autoencoder to perform generative pseudo replay without storing samples, ensuring monotonic improvement with a small memory footprint. A synthetic data-aware VAE reconstruction loss and class-balanced memory inference are proposed to preserve the teacher’s categorical information and maintain diverse memory samples during distillation. Empirical results across MNIST, CIFAR, and Tiny ImageNet show PRE-DFKD achieves higher average student accuracy and dramatically reduced memory usage compared with replay-based and replay-free baselines, demonstrating practical, privacy-preserving data-free KD.

Abstract

Data-Free Knowledge Distillation (KD) allows knowledge transfer from a trained neural network (teacher) to a more compact one (student) in the absence of original training data. Existing works use a validation set to monitor the accuracy of the student over real data and report the highest performance throughout the entire process. However, validation data may not be available at distillation time either, making it infeasible to record the student snapshot that achieved the peak accuracy. Therefore, a practical data-free KD method should be robust and ideally provide monotonically increasing student accuracy during distillation. This is challenging because the student experiences knowledge degradation due to the distribution shift of the synthetic data. A straightforward approach to overcome this issue is to store and rehearse the generated samples periodically, which increases the memory footprint and creates privacy concerns. We propose to model the distribution of the previously observed synthetic samples with a generative network. In particular, we design a Variational Autoencoder (VAE) with a training objective that is customized to learn the synthetic data representations optimally. The student is rehearsed by the generative pseudo replay technique, with samples produced by the VAE. Hence knowledge degradation can be prevented without storing any samples. Experiments on image classification benchmarks show that our method optimizes the expected value of the distilled model accuracy while eliminating the large memory overhead incurred by the sample-storing methods.

Figures (5)

• Figure 1: Example of student accuracy degradation over distillation steps due to catastrophic forgetting. Red vertical line marks the epoch with peak accuracy on the validation dataset.
• Figure 2: Method overview. (a) Proposed PRE-DFKD framework. The student and the generators are trained alternately. First the generators are fixed and the student is trained by a combination of novel and memory samples. In the next stage, student is fixed and the generators are trained by the learning signals received from the KD. (b) Training process of VAE. First, the memory generator is frozen and a batch of memory samples are inferred. Later the memory batch is combined with a novel sample batch to train the encoder-decoder pair.
• Figure 3: Example real and synthetic CIFAR10 samples are given before and after noise injection in the first row. In the second row, the projection of real and synthetic samples injected with noise on 2-dimensional planes are shown. Colors represent the class labels assigned to these samples by the teacher model. The original samples are marked with purple crosses.
• Figure 4: Visualized student accuracy curves. The first row contains MNIST results with mean and variance values over runs. The second row contains student accuracies averaged over runs for CIFAR100.
• Figure 5: Ablation experiment results on MNIST.