Distilling Lightweight Domain Experts from Large ML Models by Identifying Relevant Subspaces

TL;DR

SubDistill addresses subtask distillation by identifying task-relevant subspaces at each layer and transferring only this information from a large teacher to a smaller student. It introduces an orthogonal subspace matching loss $\mathcal{E}_l(a_T, a_\theta) = \mathbb{E}\left[\| V(a_\theta - \mu_\theta) - U^\top (a_T - \mu_T)\|_2^2\right]$ with $V^\top V = I_K$ and an orthogonal $U$, plus a layer-wise weight $\alpha_l$ scaled by $\mathbb{E}\left[\|(U^{(l)})^\top (a_T^{(l)} - \mu_T^{(l)})\|_2^2\right]$. Subspace $U$ is chosen via an adapted PRCA objective that maximizes a task-relevant margin, formulated as $\max_U \{ \mathbb{E}[ \langle \tilde{a}_T, c_T \rangle_U] + \beta^{-1} \mathbb{E}[ \langle \tilde{a}_T, \tilde{a}_T \rangle_U] + \beta \mathbb{E}[ \langle c_T,c_T \rangle_U] \}$ under $U^\top U = I_K$. Empirically, SubDistill yields substantial accuracy gains over baselines on CIFAR-100 and ImageNet subtasks for both CNN and ViT teachers and, via Explainable AI analyses (e.g., pixel-wise heatmaps), better preserves the teacher’s decision strategy, particularly in low-data regimes. A decoupled training variant further suggests subtask-aware distillation can reduce hyperparameter tuning while remaining competitive, pointing to improved reproducibility and practical deployment in resource-constrained settings.

Abstract

Knowledge distillation involves transferring the predictive capabilities of large, high-performing AI models (teachers) to smaller models (students) that can operate in environments with limited computing power. In this paper, we address the scenario in which only a few classes and their associated intermediate concepts are relevant to distill. This scenario is common in practice, yet few existing distillation methods explicitly focus on the relevant subtask. To address this gap, we introduce 'SubDistill', a new distillation algorithm with improved numerical properties that only distills the relevant components of the teacher model at each layer. Experiments on CIFAR-100 and ImageNet with Convolutional and Transformer models demonstrate that SubDistill outperforms existing layer-wise distillation techniques on a representative set of subtasks. Our benchmark evaluations are complemented by Explainable AI analyses showing that our distilled student models more closely match the decision structure of the original teacher model.

Figures (5)

• Figure 1: Schematic depiction of our SubDistill approach tailored for subtask distillation. Our proposed method combines (1) orthogonal transformation, (2) centering, and (3) an Explainable AI analysis called PRCA chormai-tpami24 for identifying task-relevant components of the teacher's representation. Besides producing accurate student models, our approach also better preserves the teacher's decision strategy, as revealed by Explainable AI pixel-wise heatmaps.
• Figure 2: Demonstration of SubDistill on a synthetic example where the data follows a one-dimensional manifold, and where only a fraction of the (one-dimensional) input domain possesses task-relevant features (A). Distillation requires the student to reproduce the teacher's internal representation, which we subsume by its kernel (B). The teacher's kernel exhibits a perfect band structure, with only its middle part being task-relevant (C). The classical $(W,b)$-formulation fails to focus on the relevant section of the kernel and numerical issues further hinder distillation, leading to sporadic gaps in its band structure (D). Only the SubDistill approach, owing to its advantageous numerical properties and its focus on what is task-relevant, is able to reliably capture the teacher's band structure (E).
• Figure 3: Top: Pixel-wise explanations of the predictions of the teacher and the distilled students trained on the ImageNet 'wading bird' subtask using 80% and 25% of training data. Results are shown for a random selection of input images. We compare students produced by our SubDistill approach and the VKD, and 'output only' baselines. Pixel-wise explanations are averaged over three training runs. Bottom: Scatter plots comparing the teacher and the students explanations on a patch level (patches of size $8\times 8$) and for the whole data distribution. The diagonal line in the plots represents exact matching between teacher and student. Next to the scatter plots, we show the Pearson correlation coefficient (corr) between teacher and student explanations.
• Figure 4: Ablation study on the layer-wise losses included in the overall distillation objective in Eq. \ref{['eq:model-distillation-loss']}. The teacher is a ResNet18 trained on CIFAR-100 and the student is the ResNet18-S. The shaded region indicates standard error from three random initializations. From right to left, we remove the top-layer loss, the second-to-top layer loss, and so on until the input layer, at which point the distillation algorithm becomes equivalent to 'output only'.
• Figure 5: Comparison between the performance of students trained with SubDistill using decoupled training and other layer-wise distillation approaches trained using end-to-end training. Results are shown for distillation on the CIFAR100 Vehicles 2 subtask. The teacher model is ResNet18 (used in the main experiment), and the student is a small ResNet18 with layer dimensions equal to (48, 64, 16, 8). Accuracy scores are shown relative to the simple output-only baseline.