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

arXiv: https://arxiv.org/abs/2506.02294 · HTML: https://mlsanity.com/paper/2506.02294/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Placed in General knowledge distillation methods, one of 12 directions in Knowledge distillation.
- Builds on 70 of its 93 references that are in the map.
- First posted on arXiv 2 Jun 2025.

## Placed in

- Knowledge distillation / [General knowledge distillation methods](https://mlsanity.com/distill/general-knowledge-distillation-methods/)

## Builds on (most cited first, 30 of 70)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Distilling the Knowledge in a Neural Network](https://mlsanity.com/paper/1503.02531/) (2015)
- [Qwen3 Technical Report](https://mlsanity.com/paper/2505.09388/) (2025)
- [ImageNet Large Scale Visual Recognition Challenge](https://mlsanity.com/paper/1409.0575/) (2014)
- [mixup: Beyond Empirical Risk Minimization](https://mlsanity.com/paper/1710.09412/) (2017)
- [LAION-5B: An open large-scale dataset for training next generation image-text models](https://mlsanity.com/paper/2210.08402/) (2022)
- [Deep Learning Face Attributes in the Wild](https://mlsanity.com/paper/1411.7766/) (2014)
- [A ConvNet for the 2020s](https://mlsanity.com/paper/2201.03545/) (2022)
- [Progressive Growing of GANs for Improved Quality, Stability, and Variation](https://mlsanity.com/paper/1710.10196/) (2017)
- [CutMix: Regularization Strategy to Train Strong Classifiers with Localizable Features](https://mlsanity.com/paper/1905.04899/) (2019)
- [Benchmarking Neural Network Robustness to Common Corruptions and Perturbations](https://mlsanity.com/paper/1903.12261/) (2019)
- [Prompt-to-Prompt Image Editing with Cross Attention Control](https://mlsanity.com/paper/2208.01626/) (2022)
- [DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps](https://mlsanity.com/paper/2206.00927/) (2022)
- [Shortcut Learning in Deep Neural Networks](https://mlsanity.com/paper/2004.07780/) (2020)
- [Knowledge Distillation: A Survey](https://mlsanity.com/paper/2006.05525/) (2020)
- [IP-Adapter: Text Compatible Image Prompt Adapter for Text-to-Image Diffusion Models](https://mlsanity.com/paper/2308.06721/) (2023)
- [Null-text Inversion for Editing Real Images using Guided Diffusion Models](https://mlsanity.com/paper/2211.09794/) (2022)
- [Latent Consistency Models: Synthesizing High-Resolution Images with Few-Step Inference](https://mlsanity.com/paper/2310.04378/) (2023)
- [ConvNeXt V2: Co-designing and Scaling ConvNets with Masked Autoencoders](https://mlsanity.com/paper/2301.00808/) (2023)
- [DataComp: In search of the next generation of multimodal datasets](https://mlsanity.com/paper/2304.14108/) (2023)
- [How to train your ViT? Data, Augmentation, and Regularization in Vision Transformers](https://mlsanity.com/paper/2106.10270/) (2021)
- [Just Train Twice: Improving Group Robustness without Training Group Information](https://mlsanity.com/paper/2107.09044/) (2021)
- [Consistency Trajectory Models: Learning Probability Flow ODE Trajectory of Diffusion](https://mlsanity.com/paper/2310.02279/) (2023)
- [Last Layer Re-Training is Sufficient for Robustness to Spurious Correlations](https://mlsanity.com/paper/2204.02937/) (2022)
- [Segment Anything in High Quality](https://mlsanity.com/paper/2306.01567/) (2023)
- [ResNet strikes back: An improved training procedure in timm](https://mlsanity.com/paper/2110.00476/) (2021)

## Built on by (most cited first, 2 of 2)

- [Anomaly Detection by Effectively Leveraging Synthetic Images](https://mlsanity.com/paper/2512.23227/) (2025)
- [Do Student LLMs Inherit OOD Robustness? Invariance-Weighted Distillation for Reliable Knowledge Transfer](https://mlsanity.com/paper/2609.22566/) (2026)

## Abstract (arXiv)

Large foundation models trained on extensive datasets demonstrate strong zero-shot capabilities in various domains. Knowledge distillation has become an established tool for transferring knowledge from foundation models to small student networks when data and model size are constrained. However, the efficacy of distillation is often hampered by limited training data coverage. This can result in a covariate shift between training and test data which in turn can lead the student to exploit spurious features or even shortcut learning. We address this problem by introducing a novel diffusion-based data augmentation strategy that generates images by maximizing the disagreement between the teacher and the student, effectively creating challenging samples that the student struggles with, thus mitigating the problem of covariate shift. Experiments demonstrate that, compared to state-of-the-art diffusion-based data augmentation baselines, our approach is best or second-best in sample mean accuracy and improves the worst group and mean group accuracy on CelebA-HQ, SpuCo Birds and BAR as well as the spurious score on Spurious ImageNet under covariate shift.