Verbetering van Knowledge Distillation onder Onbekende Covariate Shift via Vertrouwensgeleide Data Augmentatie

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

June 2, 2025

Auteurs: Niclas Popp, Kevin Alexander Laube, Matthias Hein, Lukas Schott

cs.AI

Samenvatting

Grote foundationmodellen die getraind zijn op uitgebreide datasets tonen sterke zero-shot-capaciteiten in verschillende domeinen. Om hun succes te repliceren wanneer data en modelgrootte beperkt zijn, is kennisdistillatie een gevestigd hulpmiddel geworden voor het overdragen van kennis van foundationmodellen naar kleine studentnetwerken. De effectiviteit van distillatie wordt echter kritisch beperkt door de beschikbare trainingsdata. Dit werk behandelt het veelvoorkomende praktische probleem van covariate shift in kennisdistillatie, waarbij spurious features tijdens de training verschijnen maar niet tijdens de testfase. We stellen de vraag: wanneer deze spurious features onbekend zijn, maar een robuuste leraar beschikbaar is, is het dan mogelijk voor een student om ook robuust te worden tegen deze features? We pakken dit probleem aan door een nieuwe, op diffusie gebaseerde data-augmentatiestrategie te introduceren die afbeeldingen genereert door de onenigheid tussen de leraar en de student te maximaliseren, waardoor uitdagende samples worden gecreëerd waar de student moeite mee heeft. Experimenten tonen aan dat onze aanpak de nauwkeurigheid van de slechtste groep en de gemiddelde groep significant verbetert op CelebA en SpuCo Birds, evenals de spurious mAUC op spurious ImageNet onder covariate shift, en daarbij state-of-the-art, op diffusie gebaseerde data-augmentatiebaselines overtreft.

English

Large foundation models trained on extensive datasets demonstrate strong zero-shot capabilities in various domains. To replicate their success when data and model size are constrained, knowledge distillation has become an established tool for transferring knowledge from foundation models to small student networks. However, the effectiveness of distillation is critically limited by the available training data. This work addresses the common practical issue of covariate shift in knowledge distillation, where spurious features appear during training but not at test time. We ask the question: when these spurious features are unknown, yet a robust teacher is available, is it possible for a student to also become robust to them? 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. Experiments demonstrate that our approach significantly improves worst group and mean group accuracy on CelebA and SpuCo Birds as well as the spurious mAUC on spurious ImageNet under covariate shift, outperforming state-of-the-art diffusion-based data augmentation baselines