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.
Knowledge distillation under covariate shift with ConfiG data augmentations.
Our goal is to maximize student performance even on groups that are fully absent from the training data.
Left: without having seen similar samples, the student is incapable of correctly
classifying the unseen groups in the test set.
Right: leveraging the discrepancy between a robust teacher and the biased student, ConfiG
generates and adds challenging samples to the training set. Further training improves the student's
performance on previously missing groups by removing shortcuts.
Video
Problem
While knowledge distillation improves overall student performance, under unknown covariate shift the student still performs poorly on test samples from unseen groups
Covariate Shift
Change of input features between training and test datasets while the conditional distribution of the output given the input remains the same, e.g. group shift or class-wise spurious features
Goal
Use a diffusion-based data augmentation method that leverages the disagreement between a teacher that is robust to covariate shift and a student to identify and mitigate student biases without requiring explicit knowledge of what these biases are
ConfiG: Confidence-Guided Data Augmentation
Motivation
For a student distilled only on the biased real training data, the agreement between teacher and student is
lower on groups which are not present in the training data.
For ConfiG, we make use of the observation that teacher and student predictions tend to diverge on samples from groups that are absent from the training data.
In this experiment we distill students on subsets of the CelebA-HQ and SpuCo Birds dataset.
On Celeb A-HQ the task is again gender classification where age and hair color are spurious features.
Female faces in the training data are young and have blond hair while male faces are old and do not have blond hair.
The test data contains all combinations of age and hair color for both classes.
SpuCo Birds includes waterbirds or landbirds with the background as bias.
The training data only contains waterbirds with water background while the test data features waterbirds in front of land backgrounds and vice versa.
We observe that the agreement between teacher and student is high for groups that are present in both the training and test data but lower for groups absent from the training data.
Latent optimization
ConfiG uses an auxiliary student model that was trained only on real images and maximizes the confidence of the teacher model while minimizing the confidence of the student with respect to a target class.
For this purpose we use the DiG-IN [1] framework with the loss:
where $\mathbf{f}$ is the student, $\mathbf{t}$ the teacher and $y$ the target class.
$\mathbf{z}_0$ denotes the initial latent, $\mathcal{D}$ the VAE decoder, $\emptyset$ the null-text, $C$ the conditioning and $\mathcal{M}_{\epsilon_\theta}$ the iterative denoising map of the diffusion model for the input image $\mathbf{x}$.
The purpose of the nonlinear transformation $(\cdot)^{\gamma}$ with $\gamma = 2$ is to speed up the optimization since we observe that the teacher and student frequently agree on the initial image with either high or low confidences.
Following DiG-IN, we use Adam to perform the optimization and use a foreground-aware distance regularization.
Warm-Starting the Optimization through Inter-Class Mixing
Inter-class mixing of training samples has been shown to greatly benefit diffusion-based data augmentation [3].
For every augmentation we therefore sample a random target class, constrained to be different from the
original class. The optimization is warm-started by first performing null-text inversion with the original
image label, followed by prompt-to-prompt editing [2] towards the target class label. The latent of the
resulting image initializes the optimization problem above, which is then solved with respect to the target
class.
Illustration of ConfiG compared to other diffusion-based augmentation methods.Left: the task is gender prediction (male/female), where all men in the training data are
old and not blond and all women are young and blond.
Right: existing augmentations yield high confidences for “male” or
“female” for both the teacher and the student. In contrast, ConfiG maximizes the difference
between teacher and student and thus generates a challenging example for the student — here an old and
non-blond female face carrying the attributes the student picked up as spurious features.
Summary of Quantitative Results
CelebA-HQ
SpuCo Birds
BAR
IN-100
Method
SMA
GMA
WGA
SMA
WGA
SMA
SMA
SpuScore
Teacher
99.0
97.7
91.9
96.9
93.8
95.7
96.1
85.5
Real Data
94.0
86.9
53.4
56.5
13.0
37.7
71.5
54.1
DA-Fusion
95.4
90.1
68.0
71.5
30.2
53.0
74.9
57.0
Diff-Mix
97.4
93.1
70.6
69.5
32.1
59.3
76.9
59.5
ActGen
97.0
92.6
75.7
64.5
26.0
55.7
78.4
60.5
ConfiG (ours)
97.8
95.7
88.2
73.4
39.5
58.7
80.0
62.4
All values in %. Bold marks the highest and underline the second highest score per column.
The table shows the top three baselines; the results for the remaining five baselines can be found in the
full paper.
Evaluation MetricsSMA sample mean accuracy · GMA group mean accuracy ·
WGA worst group accuracy · SpuScore class-wise separation of images from the
Spurious ImageNet test set (spurious feature but no class object) versus the validation set
Diffusion-Based Data Augmentation
Real samples are replaced by synthetic augmentations with probability 0.5.
We generate one ConfiG augmentation per image. The number of augmentations per baseline is adjusted to approximately
match the generation FLOPs. On the group shift datasets (CelebA-HQ, SpuCo Birds, BAR) our method uses
5 optimization steps on average and we use the multipliers 8 (ActGen), 12 (DA-Fusion) and
28 (Diff-Mix). On ImageNet the average number of steps is higher at approximately 7 and we use the
multipliers 10 (ActGen), 16 (DA-Fusion) and 36 (Diff-Mix). More details on the computational
cost of ConfiG and the baseline multipliers can be found in Appendix C of the paper.
Main insight
For improving knowledge distillation under covariate shift, fewer high-quality data augmentations from
ConfiG are more beneficial than more augmentations from existing diffusion-based methods that are
computationally cheaper per sample.
Qualitative Examples
Qualitative examples of ConfiG on CelebA-HQ, SpuCo Birds and BAR. The target class is chosen
randomly to be different from the original class. The generated images have low student confidence and reveal
the spurious features the student relies on: a young male with blond hair is recognized as female
(CelebA-HQ), or a landbird standing in water as a waterbird (SpuCo Birds).
Main References
Augustin, M. et al.: DiG-IN: Diffusion guidance for investigating networks – uncovering classifier
differences, neuron visualisations, and visual counterfactual explanations. CVPR 2024 (Link)
Hertz, A. et al.: Prompt-to-prompt image editing with cross attention control. arXiv preprint 2022
(2208.01626)
Wang, Z. et al.: Enhance image classification via inter-class image mixup with diffusion model. CVPR 2024 (Link)
