Investigating the Robustness of Subtask Distillation under Spurious Correlation

TL;DR

The paper investigates how spurious correlations in distillation data affect subtask distillation, comparing traditional and state-of-the-art methods including SubDistill across CNN and transformer teacher–student pairs. It demonstrates that SubDistill maintains high performance and closely mirrors the teacher's representations and decision strategies even when training data are heavily tainted, whereas simpler baselines degrade toward random predictions. By combining quantitative results with t-SNE and XAI analyses, the study connects robustness to both representation alignment and the preservation of task-focused decision strategies. The findings underscore the need for alignment-centric distillation and careful data curation to ensure reliable deployment of compact models in real-world environments with confounding patterns.

Abstract

Subtask distillation is an emerging paradigm in which compact, specialized models are extracted from large, general-purpose 'foundation models' for deployment in environments with limited resources or in standalone computer systems. Although distillation uses a teacher model, it still relies on a dataset that is often limited in size and may lack representativeness or exhibit spurious correlations. In this paper, we evaluate established distillation methods, as well as the recent SubDistill method, when using data with spurious correlations for distillation. As the strength of the correlations increases, we observe a widening gap between advanced methods, such as SubDistill, which remain fairly robust, and some baseline methods, which degrade to near-random performance. Overall, our study underscores the challenges of knowledge distillation when applied to imperfect, real-world datasets, particularly those with spurious correlations.

Figures (3)

• Figure 1: Spurious correlation model considered in this work, where MNIST digits are inserted in the top-right corner of ImageNet images. On the training data, the digit's class is spuriously correlated to the true image class (bird type), whereas on the test data the digits are assigned randomly.
• Figure 2: Data representations produced by the ResNet18 teacher and the ResNet18-S students trained with different distillation methods under the presence of spurious correlation (100% level), and visualized using t-SNE. The color of the marker corresponds to the true class label (i.e. the type of wading bird present in the image), while the annotated digit corresponds to the type of artifact (superposed MNIST digits 0 to 4). The representations are taken from the best run of each method from the output of the global average pooling layer.
• Figure 3: Top: Explainable AI analysis of the teacher and students trained with the Output Only, VID, and SubDistill approaches using the 0%-MNIST-contaminated training data (no spurious correlation) and the 100%-MNIST-contaminated training data (spurious correlation). Red highlighting indicates visual features that contribute positively to the prediction, and blue highlighting indicates features that contribute negatively. Bottom: Scatter plots between each patch's relevance obtained from the teacher and the student at the level of 56$\,\times\,$56 patches. The relevance scores are computed from all MNIST-contaminated test samples, averaged over the three training runs of each student. The depicted 'corr' is the Pearson correlation coefficient.