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

1. Problem Statement

The paper addresses the limitations of Knowledge Distillation (KD) when applied to scenarios with unknown covariate shift. While foundation models demonstrate strong zero-shot capabilities and distributional robustness, deploying them in resource-constrained environments often requires transferring knowledge to smaller student networks. Standard KD relies on "dark knowledge" (soft labels) from a teacher to improve student generalization. However, when the training data ($P_{train}$) does not sufficiently cover the test distribution ($P_{test}$), the student often fails to learn the true decision boundary.

Instead, the student exploits spurious features (e.g., hair color or age in gender classification) that are correlated with the label in the training set but not causally related. This leads to shortcut learning, where the model performs well on training data but fails on test groups that were absent or underrepresented during training (e.g., old blond females or young non-blond males). Existing diffusion-based data augmentation methods often fail to target these specific biases because they do not explicitly leverage the discrepancy between a robust teacher and a biased student.

2. Methodology: ConfiG

The authors propose ConfiG (Confidence-Guided), a novel data augmentation framework designed to mitigate covariate shift by generating synthetic samples that specifically challenge the student's reliance on spurious features.

Core Mechanism

ConfiG utilizes a diffusion model (based on Stable Diffusion) guided by the disagreement between a robust teacher model ($t$) and a biased student model ($f$). The goal is to generate images where:

1. The teacher maintains high confidence in the correct class (preserving causal features).
2. The student has low confidence in the correct class (indicating reliance on spurious features).

Technical Implementation

The method operates in two main stages:

1. Inter-Class Mixing (Warm-Starting):
   To avoid local minima and ensure the generation of challenging samples, ConfiG first performs inter-class mixing. Using prompt-to-prompt editing, the original image is edited to a random target class (different from the original). This provides a latent initialization that is semantically distinct from the original, facilitating the exploration of new feature combinations.

2. Latent Optimization:
   The core of ConfiG involves optimizing the latent code ($z_0$), null-text ($\emptyset$), and conditioning signal ($C$) to maximize the disagreement between the teacher and student. The optimization minimizes the following loss function:
   $$\min_{z_0, \emptyset, C} \; f(\mathcal{D}(\mathcal{M}_{\epsilon_\theta}(z_0, \emptyset, C)))_y^\gamma + \left(1 - t(\mathcal{D}(\mathcal{M}_{\epsilon_\theta}(z_0, \emptyset, C)))_y\right)^\gamma$$
   Where:

   • $f$ and $t$ are the student and teacher models.
   • $y$ is the target class confidence.
   • $\gamma$ is a non-linear transformation (empirically set to 2) to handle cases where the student is initially highly confident.
   • $\mathcal{M}_{\epsilon_\theta}$ is the noise prediction model, and $\mathcal{D}$ is the decoder.

   This process generates images where the student is "tricked" by spurious features (e.g., a male face with blond hair being classified as female) while the teacher correctly identifies the true class.

3. Training Pipeline:

   • An auxiliary student is first distilled on the limited, biased real training data. This student naturally picks up the spurious correlations.
   • ConfiG uses this auxiliary student to generate synthetic augmentations.
   • A robust student is then trained on a mix of real data and the generated ConfiG augmentations (with a replacement probability $\alpha = 0.5$).

3. Key Contributions

1. Problem Identification: The authors demonstrate that while KD from a robust teacher improves overall performance, it remains insufficient under unknown covariate shift, leading to poor performance on unseen groups in the test set.
2. ConfiG Framework: They introduce a confidence-guided diffusion augmentation method that explicitly targets covariate shift by maximizing teacher-student disagreement, without requiring explicit knowledge of the bias or group labels.
3. Empirical and Theoretical Validation: The paper provides theoretical motivation showing that augmentations challenging to the student but easy for the teacher can reduce generalization error ($\Omega$). Empirically, ConfiG outperforms state-of-the-art diffusion-based baselines across multiple datasets.

4. Experimental Results

The authors evaluated ConfiG on four datasets: CelebA-HQ, SpuCo Birds, BAR (Biased Activity Recognition), and a subset of ImageNet (Spurious ImageNet).

• Group Shift Datasets (CelebA-HQ, SpuCo Birds, BAR):

  • ConfiG consistently achieved the best or second-best Sample Mean Accuracy (SMA) and the best Group Mean Accuracy (GMA) and Worst Group Accuracy (WGA) compared to six baselines (GIF, DistDiff, DA-Fusion, Diff-Mix, Diff-II, ActGen).
  • On CelebA-HQ, ConfiG improved the Worst Group Accuracy by over 34% compared to training with real data only, narrowing the gap to the teacher significantly.
  • On SpuCo Birds and BAR, ConfiG reduced the performance gap to the teacher by 23% and 37% respectively, outperforming all baselines in worst-group performance.

• ImageNet-100 (Spurious Features):

  • ConfiG achieved the highest validation accuracy and Spurious Score (SpuScore), indicating improved robustness against class-wise spurious features.
  • Improvements were observed for both ViT-T and ViT-S student architectures.

• Ablation Studies:

  • Components: Combining inter-class mixing and latent optimization yielded the best results; using either component alone was less effective.
  • Stronger Augmentations: ConfiG combined with stronger non-synthetic augmentations (e.g., AutoAugment, MixUp) further improved performance.
  • Teacher Strength: ConfiG remained effective even when using weaker teacher models (e.g., ResNet50 or smaller CLIP variants), though absolute performance decreased.
  • Model-Centric Bias Mitigation: ConfiG augmentations significantly improved the performance of model-centric bias mitigation methods (LfF, uLA, ERM) when combined with them.

5. Significance and Claims

The paper claims that ConfiG offers a data-centric solution to the problem of limited training data coverage and unknown covariate shift. By leveraging the discrepancy between a robust teacher and a biased student, the method generates targeted synthetic data that forces the student to unlearn spurious correlations and approximate the teacher's decision boundary more accurately.

The authors emphasize that their approach does not require explicit group annotations or knowledge of the specific biases, making it applicable to real-world scenarios where such information is unavailable. While the method shows superior performance in group robustness and handling spurious features, the authors note limitations:

• The method is currently designed for global image augmentation in classification tasks and may require adaptation for location-aware tasks like object detection.
• It relies on pre-trained teacher models and diffusion models, which have computational costs, though the authors demonstrate compatibility with faster single-step diffusion models (Hyper-SD).
• The theoretical analysis assumes certain conditions (e.g., teacher accuracy) that are approximately satisfied in practice but difficult to characterize purely from discrete augmentation sets.

In conclusion, ConfiG represents a step toward more robust knowledge distillation in data-scarce and biased environments, effectively bridging the gap between the capabilities of large foundation models and smaller, deployable student networks.