Improving Knowledge Distillation Under Unknown Covariate Shift Through Confidence-Guided Data Augmentation

Large foundation models trained on extensive datasets demon-strate strong zero-shot capabilities in various domains. Knowledge dis-tillation has become an established tool for transferring knowledge fromfoundation models to small student networks when data and model sizeare constrained. However, the efficacy of distillation is often hampered bylimited training data coverage. This can result in a covariate shift betweentraining and test data, which in turn can lead the student to exploitspurious features or even shortcut learning. We address this problemby introducing a novel diffusion-based data augmentation strategy thatgenerates images by maximizing the disagreement between the teacherand the student, effectively creating challenging samples that the studentstruggles with, thus mitigating the problem of covariate shift. Experi-ments demonstrate that, compared to state-of-the-art diffusion-based dataaugmentation baselines, our approach is best or second-best in samplemean accuracy and improves the worst group and mean group accuracyon CelebA-HQ, SpuCo Birds and BAR as well as the spurious score onSpurious ImageNet under covariate shift.