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

Conference: ECCV2026
Paper: ECCV Paper
Area: Model Compression
Keywords: knowledge distillation, covariate shift, confidence guidance, synthetic data augmentation, group robustness

TL;DR

ConfiG combines inter-class diffusion editing with teacher-student confidence disagreement to generate examples the student has not mastered, improving distillation without knowing the spurious attributes and raising CelebA-HQ worst-group accuracy from 53.44% with real data alone to 88.15%.

Background & Motivation

Knowledge distillation lets a small model learn the class probabilities of a large teacher instead of only the hard labels in a training set. However, a teacher knowing a useful rule does not mean a limited distillation dataset will force the student to learn it. If training waterbirds always appear on water and landbirds always appear on land, the student can imitate the teacher on training examples by recognizing backgrounds alone. When a waterbird appears on land at test time, a low training distillation loss no longer indicates that the student understands bird identity. The covariate shift studied here changes the input distribution while preserving the conditional label distribution, rather than introducing new classes.

The harder setting is that some attribute combinations are entirely absent during training, not merely underrepresented. Reweighting minority groups then has no examples of those groups to reweight, while ordinary diffusion augmentation can reproduce the existing correlations. Generating difficult images using only the student is also insufficient: the image might have changed class, so low student confidence need not reveal a shortcut. The procedure must establish both that the student struggles and that the example still supports the specified class, without relying on training annotations for missing groups.

The authors assume access to a teacher that is reasonably reliable on both training and test distributions, and use its disagreement with a biased student to locate useful generation regions. In source Table 1, 2256 of 2351 teacher-student disagreements on CelebA-HQ occur in missing training groups; the corresponding count on SpuCo Birds is 779 of 828. These test-set statistics motivate the approach empirically; they are neither test samples accessed during generation nor a guarantee for arbitrary distributions. Core Idea: constrain semantics with high teacher confidence in the target class, locate knowledge gaps with low student confidence in that class, and add the resulting images to distillation.

Method

Overall Architecture

The inputs are a small labeled real-image dataset, a pretrained classification teacher, and a pretrained diffusion model. An auxiliary student is first distilled on real data to expose the shortcuts supported by the available examples, then used as a fixed evaluator during generation. Each image undergoes inter-class editing warm-start, confidence-guided latent optimization, and finally mixed-data distillation using real and synthetic images. The output is a small classifier that can be deployed independently, not a system that must run diffusion at test time.

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    Input["Real images and labels"] --> Warm["Inter-Class Editing<br/>Warm-Start"]
    Input --> Aux["Real-only distillation<br/>of auxiliary student"]
    Warm --> Opt["Confidence-Guided<br/>Latent Optimization"]
    Aux -->|Low target-class confidence| Opt
    Teacher["Fixed classification teacher"] -->|High target-class confidence| Opt
    Opt --> Mix["Mixed-Data Distillation"]
    Input -->|Real samples| Mix
    Teacher -->|Soft-label supervision| Mix
    Mix --> Output["Deployed student<br/>No generator or teacher"]

The two teacher roles should not be conflated: the classification teacher supplies task judgments, while the diffusion model supplies an editable image prior. The auxiliary student helps find counterexamples; the final student learns from the expanded dataset rather than undergoing adversarial optimization at test time. Group labels are used to construct and evaluate the benchmarks, not as required inputs to the ConfiG generator. An unknown shift does not mean the method is prior-free: it still depends on useful teacher and generator knowledge in the missing regions.

Key Designs

1. Inter-Class Editing Warm-Start: move beyond the original class neighborhood

Each real image is encoded into the diffusion model's latent representation, and null-text inversion finds an initial latent and null-text conditioning that reconstruct the image. This connects a real example to a differentiable diffusion trajectory instead of searching entirely from random noise. The classifiers do not read the latent directly; they read the image decoded from the final latent representation. Subsequent confidence optimization must therefore propagate gradients through the generation process.

Full ConfiG randomly selects a target class different from the original class. It performs inversion with the original class, then applies prompt-to-prompt editing adapted to DiG-IN to move the image toward the target class. For example, starting from a waterbird on water and targeting a landbird can retain scene content while creating an attribute combination missing from training. This is not pixel-wise averaging of two images or the soft-label interpolation used in MixUp. It is inter-class semantic editing inside the generative model, providing a better initialization for subsequent optimization. Searching only near the original class can be hindered by small gradients in high-confidence regions and can fail to escape the correlations already present in training.

2. Confidence-Guided Latent Optimization: keep teacher approval while challenging the student

Confidence here means the probability assigned to the specified target class, not the maximum class probability or prediction entropy. If the target is landbird, the procedure examines both models' landbird probabilities even when their current predicted classes differ. Generation aims to lower the student's target-class confidence while maintaining high teacher confidence in that class. Difficulty should thus arise from teacher-student disagreement rather than simply making an image blurry, unrecognizable, or semantically incorrect. The teacher constraint acts as a semantic check, although its correspondence to correct semantics still depends on teacher reliability.

The optimized quantities include the inverted initial latent, null-text conditioning, and CLIP text conditioning. After each generation and decoding pass, the fixed teacher and auxiliary student provide target-class probabilities, and gradients update those inputs through the differentiable diffusion trajectory. The paper uses DiG-IN and Adam; this stage does not retrain the diffusion model's weights. The authors observe that the auxiliary student is often highly confident initially with small gradients, and apply a power transformation to student confidence with exponent \(\gamma=2\). They also use foreground-aware distance regularization to constrain deviation from the original content. The cache does not supply its full details or coefficient, so it should not be expanded into an additional segmentation-supervision module.

The text extraction of source Equation (7) is corrupted, preventing reliable recovery of its complete operations. This note therefore describes the optimization directions and variables confirmed by the surrounding prose, without presenting a guessed probability difference, ratio, or product as the authors' exact objective. This does not prevent a mechanistic explanation, but it means the generation loss cannot be reproduced term by term from this note alone. In particular, maximizing disagreement should not be reinterpreted as maximizing KL divergence between the complete class distributions; the prose does not define it that way.

3. Mixed-Data Distillation: turn discovered knowledge gaps into ordinary training examples

The auxiliary student initially sees only real training data and therefore exposes biases that this dataset can induce. It is used to generate synthetic augmentations offline, after which the final student learns teacher predictions on a mixture of real and synthetic data. The main experiments generate 1 ConfiG augmentation per real image and replace real images with synthetic ones during training with probability \(\alpha=0.5\). This is a sampling replacement probability, not a half-and-half pixel blend or a weight between teacher supervision and hard-label loss. The new images provide counterexamples where the student could previously exploit shortcuts, encouraging a closer approximation of the teacher's decision boundary.

Final training uses the teacher's full class probabilities as soft labels, rather than treating the inter-class editing target as the sole supervision. This separates the target-class constraint used during generation from the predictive-distribution matching used during training. Even when a generated image is not assigned to a single class with absolute certainty, the teacher probabilities preserve relative class information. The described main pipeline is not an adaptive online loop that repeatedly regenerates data after updating the final student; augmentation generation precedes student training. Only the final student is retained at inference time, so the main added cost is offline dataset construction rather than the deployed computation graph.

A Worked Example

In the CelebA-HQ example in source Figure 3, the starting image is an older, non-blond male, matching the training correlation. ConfiG changes the target class to female and first warm-starts optimization with inter-class editing. After changing the class prompt alone, both models can still agree through familiar cues, providing little useful disagreement. After further optimization, the final image retains older and non-blond attributes, while the teacher assigns 99.6% probability to female and the auxiliary student only 1.8%. In the authors' visual interpretation, the teacher recognizes the target semantics while the student remains misled by spurious attributes; this is not an automatic correctness proof for every generated image. Once such an image enters distillation, continuing to classify by hair color or age conflicts with the teacher's supervision. The example concerns the dataset's binary attribute classification and does not imply that these appearance cues should determine real people's identities.

Loss & Training

The final student uses response-based knowledge distillation, minimizing cross-entropy between teacher and student probabilities. The following standard equivalent expression restates the prose definition of distillation training; it does not reconstruct corrupted Equation (7):

\[ \mathcal{L}_{\mathrm{KD}}=-\frac{1}{N}\sum_{i=1}^{N}\sum_{c=1}^{C}t_c(x_i)\log f_c(x_i). \]

Here \(t_c(x_i)\) and \(f_c(x_i)\) are the teacher and final student's probabilities for class \(c\), with training samples drawn through the mixture described above. The main setup does not add a real hard-label cross-entropy term, and this note does not assume a distillation temperature absent from the cache. Group-shift experiments use an ImageNet-21k AugReg-pretrained ViT-T and a DataComp-XL-pretrained CLIP ViT-L/14 teacher. CelebA-HQ and BAR use class names for zero-shot teacher predictions; SpuCo Birds predicts fine-grained bird categories before aggregating them into waterbirds and landbirds. ImageNet-100 uses a BEiTv2-B teacher trained on full ImageNet, ViT-T or ViT-S students, and A1 standard augmentations. Diffusion augmentation baselines all use Stable Diffusion 1.4, with synthetic sample counts adjusted to approximately match generation FLOPs rather than equal image counts. The main-text Diff-Mix comparison uses 28 augmentations per image, showing that each ConfiG image is more expensive; generating only 1 image is not itself a speed advantage.

Key Experimental Results

Main Results

CelebA-HQ and SpuCo Birds retain 250 training images per class while fully removing some groups; BAR uses 100 training images per class. SMA is accuracy averaged over individual samples, GMA is the equally weighted mean of group accuracies, and WGA is the minimum group accuracy. CelebA-HQ forms 8 test groups from binary class, age, and hair-color attributes; SpuCo Birds has balanced test groups, so SMA equals GMA. BAR test locations do not overlap training locations, and the paper reports only SMA; a BAR WGA should not be invented. The table below selects results from source Table 2, page 10; values are percentages, with student means ± standard deviations over 6 training runs and random resized crops as standard augmentation.

Method	CelebA-HQ SMA	CelebA-HQ GMA	CelebA-HQ WGA	SpuCo Birds SMA	SpuCo Birds WGA	BAR SMA
Real data only	93.98 ±1.68	86.89 ±2.48	53.44 ±8.12	56.53 ±2.40	12.97 ±6.65	37.69 ±3.58
Diff-Mix	97.36 ±0.16	93.07 ±0.73	70.63 ±3.76	69.49 ±2.33	32.13 ±2.22	59.28 ±2.32
ActGen	97.04 ±0.48	92.61 ±0.67	75.68 ±4.22	64.54 ±2.07	25.97 ±1.80	55.71 ±2.70
ConfiG	97.76 ±0.25	95.67 ±0.55	88.15 ±2.10	73.38 ±1.28	39.50 ±4.53	58.66 ±3.76

CelebA-HQ WGA improves by 34.71 percentage points over real data alone and by 12.47 percentage points over ActGen. However, BAR SMA remains below Diff-Mix, so the method is not first on every dataset and metric. The SpuCo Birds teacher has 93.80% WGA, leaving a substantial gap to the student's 39.50%; removing some shortcuts does not fully recover teacher capabilities.

ImageNet-100 uses 100 training images per class across 100 classes; the following selection is from source Table 3, page 12, with values in percent. SpuScore measures separation in class probability between genuine class images and images containing only the spurious feature without the class object; higher is better. It is not classification accuracy on the latter images; the main text does not expand its full formula, so this note does not assume a particular ranking statistic.

Method	ViT-T SMA	ViT-T SpuScore	ViT-S SMA	ViT-S SpuScore
Real data only	71.46	54.05	86.46	68.62
ActGen	78.36	60.48	91.04	73.17
ConfiG	80.02	62.36	91.68	75.49

Ablation Study

The following selection from source Table 4, page 12, reports percentages under the same teacher, student, and training setup as the main group-shift experiment. Latent optimization only omits inter-class editing and retains the original target class; p2p only performs inter-class editing without subsequent confidence-guided optimization.

Config	CelebA-HQ WGA	SpuCo Birds SMA	SpuCo Birds WGA	BAR SMA
Real data only	53.44 ±8.12	56.53 ±2.40	12.97 ±6.65	37.69 ±3.58
Latent optimization only	71.11 ±3.74	61.53 ±2.03	21.63 ±3.61	39.40 ±3.53
p2p only	86.99 ±1.61	66.64 ±1.83	33.20 ±3.41	44.09 ±3.77
ConfiG	88.15 ±2.10	73.38 ±1.28	39.50 ±4.53	58.66 ±3.76

Key Findings

• Inter-class editing alone explains most of the CelebA-HQ improvement; full ConfiG adds only 1.16 percentage points of WGA, so confidence optimization should not receive credit for the entire gain.
• On BAR, full ConfiG exceeds p2p only by 14.57 percentage points, showing that the complementarity of warm-starting and subsequent optimization is task-dependent.
• Source Table 5, page 13, reports 53.10 ±3.42% SpuCo Birds WGA when ConfiG is combined with stronger crops, flips, and CutMix, demonstrating complementarity with standard augmentation.
• In source Table 6, page 14, increasing ConfiG from 1 to 2 augmentations per image raises SpuCo Birds WGA from 39.50% to 46.90%, but changes CelebA-HQ WGA from 88.15% to 87.83%; more samples do not improve every setting.

Highlights & Insights

• Disagreement determines where training data are missing, rather than serving only as a distillation loss. It converts teacher knowledge not yet transferred to the student into a trainable input distribution.
• Inter-class editing is not incidental. It can break the association between classes and backgrounds or attributes, providing more useful starting points for confidence optimization.
• Reporting SMA, GMA, and WGA together matters. A high overall accuracy can coexist with severe worst-group failure, obscuring the robustness deficit when only averages are considered.

Limitations & Future Work

• The authors explicitly limit the method to global image augmentation for classification; localization tasks such as object detection require constrained editing and are not validated here.
• The method assumes a reliable teacher on missing groups. Gains with weaker teachers do not establish that correct counterexamples can be found when teacher and student share the same shortcut.
• Reader analysis: high teacher confidence is not sufficient for semantic correctness, and generation could exploit model vulnerabilities; independent quality review or multiple-teacher agreement would help investigate this risk.
• Approximate generation-FLOPs matching does not imply matched end-to-end time, memory, or data-preparation costs, so the main results do not directly establish a predeployment cost-benefit ratio.
• The available full-text cache contains the main paper and references, but not the cited appendices; Equation (7) is corrupted, and theoretical derivations, regularization details, full hyperparameters, and Hyper-SD appendix results were not independently verified.

Related Work & Insights

• vs ordinary response distillation: ordinary distillation imitates a teacher on a given dataset; ConfiG changes the coverage of the data available for distillation, making knowledge distillation and model compression its primary contribution area.
• vs Diff-Mix: both use inter-class generation; ConfiG additionally uses teacher and student target-class confidence to locate difficult regions, at the cost of extra optimization per augmentation.
• vs ActGen: ActGen also uses an auxiliary student for augmentation but lacks the same explicit classification-teacher semantic constraint; ConfiG guides both difficulty and target semantics.
• vs LfF / uLA: these methods mitigate bias on fixed data, whereas ConfiG expands coverage; source Table 8 shows gains from combining them, but uses teacher labels for synthetic images and should not be conflated with pure soft-label distillation.
• Extension suggested by the reader: investigate regions where teacher and student fail together and audit generated-sample diversity to avoid covering only shortcuts exposed by a single auxiliary student.

Rating

• Novelty: 4/5. Inter-class generation and teacher-student confidence jointly address unknown missing groups with a clear problem formulation, while individual components largely build on existing methods.
• Experimental Thoroughness: 4/5. Multiple datasets, component ablations, augmentation strengths, and weaker teachers are covered, but the evidence remains classification-focused.
• Writing Quality: 4/5. The main motivation and ablations connect clearly; corrupted formulas in the available cache limit reproducibility checks.
• Value: 4/5. Useful for compressing robust teachers under limited data, provided generation costs and teacher reliability are assessed alongside accuracy.