Teacher-Guided One-Shot Pruning via Context-Aware Knowledge Distillation

AI summary

Overview

Research area: Model compression for computer vision — specifically unstructured network pruning combined with knowledge distillation (KD) for image classification.

Technical level: Advanced. The paper assumes familiarity with gradient-based pruning, knowledge distillation (forward/reverse KL divergence), exponential moving averages, and sparsity budgets.

Scope: The paper proposes and empirically evaluates a one-shot global pruning framework in which a pre-trained teacher network actively guides which student weights to remove, tested on CIFAR-10, CIFAR-100, and TinyImageNet with ResNet18/ResNet34 students and a ResNet50 teacher.

What This Paper Is About

Unstructured pruning usually demands repeated train–prune–retrain cycles, and importance scores are typically computed from heuristics such as weight magnitude rather than dynamic learning signals. The authors propose folding knowledge distillation into the pruning decision itself — using gradients from a combined cross-entropy plus context-aware distillation loss to score weight importance — so that a single aggressive pruning step can identify the weights critical both for the task and for teacher–student knowledge transfer. After pruning, the sparse model is retrained under a fixed mask, with or without KD, to recover accuracy.

Key Contributions

1. A teacher-guided gradient importance metric that scores parameters using gradients derived jointly from the task loss and a Context-Aware Kullback–Leibler Divergence (CA-KLD) distillation loss, rather than from magnitude alone.
2. KD as an active pruning signal rather than post-hoc recovery: prior work applies distillation after pruning; here the teacher's soft targets shape the importance calculation used to decide what to prune.
3. A one-shot global pruning pipeline that is computationally cheaper than iterative schemes such as the Lottery Ticket Hypothesis (LTH) and COLT, while claiming to outperform EPG and EPSD baselines at high sparsity levels.
4. Empirical evaluation across CIFAR-10, CIFAR-100, and TinyImageNet using ResNet18 and ResNet34 students distilled from a ResNet50 teacher, including comparisons against six distillation-based baselines and against COLT and one-shot LTH.

Main Findings

• Pre-pruning accuracy (Table I): The ResNet50 teacher reaches 95.40% on CIFAR-10, 80.89% on CIFAR-100, and 78.46% on TinyImageNet. ResNet18 students reach 95.92%, 81.12%, and 62.19% respectively; ResNet34 students reach 96.50% and 82.77% on CIFAR-10 and CIFAR-100 (the main text states 82.71% for the latter), with TinyImageNet not reported for ResNet34.

• High-sparsity retention (Table II, ResNet18 with 11.69M parameters): On CIFAR-10 the method retains 90.79% accuracy at 98.41% sparsity (0.18M parameters), 94.97% at 90.00% sparsity, and peaks at 96.08% at 50.46% sparsity. On CIFAR-100 it reports 67.06% at 98.01% sparsity (the main text states 66.32%), rising to 80.99% at 50.75% sparsity (text states 81.01%). On TinyImageNet it reports 50.64% at 97.56% sparsity and 59.29% at 50.02% sparsity.

• Outperforms six distillation baselines (Figures 4–6, five sparsity levels: 36%, 59%, 79%, 90%, 95%): On CIFAR-100 the method reaches 81.74% at 36% sparsity (vs. 81.32% for DLB EPSD, +0.42 pp) and 81.32% at 59% sparsity (+1.78 pp over DLB EPSD), but at 95% sparsity it trails PS-KD EPSD by 0.86 pp (74.68% vs. 75.54%) while far exceeding DLB EPSD (46.86%). On TinyImageNet it leads at every reported level, including 61.35% at 36% sparsity (+3.14 pp over DLB EPSD) and 53.42% at 95% sparsity (+3.25 pp over DLB EPSD). On CIFAR-10 it reports 96.05%, 95.84%, 95.68%, 94.83%, and 95.88% across the five sparsity levels. The baselines are CS-KD Simple, CS-KD EPSD, PS-KD Simple, PS-KD EPSD, DLB Simple, and DLB EPSD. EGP is named as a baseline in the contributions but no EGP numbers appear in the provided content.

• The main advantage over COLT is speed, not accuracy (Table III): At matched sparsity (97.7% for CIFAR-10, 97.4% for TinyImageNet and CIFAR-100), COLT-2 achieves higher accuracy on CIFAR-10 (92.40% vs. 91.87%) and TinyImageNet (53.90% vs. 51.14%), while the proposed method is slightly better on CIFAR-100 (68.66% vs. 68.40%). Latency is the decisive difference: 27.82 vs. 276 minutes on CIFAR-10, 42.43 vs. 1756 minutes on TinyImageNet, and 19.76 vs. 355 minutes on CIFAR-100 (the text elsewhere writes 335 minutes), described as roughly 10×, 41×, and 18× reductions. All latency measurements were taken on an RTX 3090 GPU.

• Comparison with one-shot LTH (Table IV): Gains over one-shot LTH appear at high sparsity — on CIFAR-10 at 98.41% sparsity, 90.79% vs. 89.47% (a 1.32% gain) — and are larger on harder datasets, e.g. on TinyImageNet at 50.02% sparsity, 60.93% vs. 56.91%. The table and the Table II results disagree for some matched settings (for instance TinyImageNet at 74.35% sparsity is listed as 58.58% in Table IV and 57.42% in Table II), and ResNet34 pruning results are not reported in the main tables.

Methodology in Plain English

The pipeline has four stages.

1. Teacher-supervised student training. A pre-trained ResNet50 teacher supervises a ResNet18 or ResNet34 student using a combined loss: cross-entropy on hard labels plus a CA-KLD distillation term. Logits are normalized (subtract the mean, divide by the standard deviation plus a small constant) and divided by a temperature T > 1 to soften the distributions, and the distillation loss blends forward KL (teacher-to-student) and reverse KL (student-to-teacher) with a mixing factor β, then is scaled by T² to keep gradient magnitudes stable.

2. Teacher-guided importance scoring. For each batch, the authors backpropagate the combined loss and compute a raw importance score for each weight as the absolute element-wise product of the weight and its gradient. Because single-batch gradients are noisy, scores are smoothed with an exponential moving average (decay γ = 0.9) and then bias-corrected by dividing by (1 − γ^t). Averaging in this way yields a stable per-weight importance estimate that reflects both the classification task and alignment with the teacher.

3. One-shot global pruning. All importance scores across the network are flattened and a single global threshold is computed so that exactly the target percentage of lowest-scoring weights is removed. A binary mask retains the rest, and the pruned model is the original weights multiplied element-wise by that mask. This is done once, not iteratively.

4. Sparsity-aware retraining. The sparse model is retrained with gradients masked so that pruned weights never receive updates, and momentum is corrected so residual optimizer velocity cannot reactivate them. Retraining is done in two ways: plain fine-tuning with early stopping (patience = 5 epochs), and KD-aware fine-tuning using the CA-KLD loss. Key settings: α = 0.7 for the loss balance, β = 0.5 for the forward/reverse KL mix, EMA decay γ = 0.9, importance accumulated over 3 epochs, and KD retraining run with (T = 3.0, α = 0.7) and (T = 5.0, α = 0.7). Experiments ran on a single NVIDIA RTX 3090 GPU (24 GB VRAM) with 16 vCPUs and 125 GB RAM.

Why This Matters

Research impact: The paper argues that pruning importance should be measured against the signal a model is actually trained on — including teacher supervision — rather than static weight magnitude. It positions distillation as a pruning-time decision aid, connecting the pruning and distillation literatures that are usually treated as separate stages.

Real-world applications (implications, as the paper frames them):

• Deploying image classifiers on mobile phones and edge devices where memory and compute are tightly constrained.
• Time-sensitive production pipelines where multi-cycle pruning schedules are impractical, given the reported minutes-scale pruning cost versus hours for iterative methods.
• Rapid prototyping and experimentation, where a one-shot pipeline shortens the loop between model design and a deployable compressed model.
• Compression of high-capacity classifiers into sparse students for settings where accuracy loss must be minimal even at extreme sparsity.

Industry relevance: The central practical claim is the cost profile — comparable accuracy to much slower iterative methods at roughly an order of magnitude less training time, which matters directly for engineering budgets and iteration speed in applied model-compression work.

Future Directions

• Closing the accuracy gap with iterative methods: COLT-2 still wins on CIFAR-10 and TinyImageNet at matched sparsity; whether better importance estimation or longer retraining can recover those 0.53% and 2.76% gaps remains open.
• Extending beyond convolution layers: the importance calculation algorithm is described for conv layers specifically, so applicability to fully connected, normalization, attention, and transformer layers is not established here.
• Investigating the extreme-sparsity regime: the method trails PS-KD EPSD at 95% sparsity on CIFAR-100, suggesting teacher-guided scoring alone may not suffice at the most aggressive budgets.
• Expanding the evaluation: results for ResNet34 under pruning are not reported, EGP is named as a baseline without corresponding numbers, and the work covers only image classification on three benchmarks — detection, segmentation, and larger-scale datasets remain untested.

Target Audience

Graduate students and researchers working on model compression, efficient deep learning, or knowledge distillation; practitioners who need to compress vision models for edge deployment and are weighing one-shot pipelines against iterative lottery-ticket-style methods; and readers interested in how auxiliary supervision signals can be folded into pruning criteria.

Authors’ abstract

Unstructured pruning remains a powerful strategy for compressing deep neural networks, yet it often demands iterative train-prune-retrain cycles, resulting in significant computational overhead. To address this challenge, we introduce a novel teacher-guided pruning framework that tightly integrates Knowledge Distillation (KD) with importance score estimation. Unlike prior approaches that apply KD as a post-pruning recovery step, our method leverages gradient signals informed by the teacher during importance score calculation to identify and retain parameters most critical for both task performance and knowledge transfer. Our method facilitates a one-shot global pruning strategy that efficiently eliminates redundant weights while preserving essential representations. After pruning, we employ sparsity-aware retraining with and without KD to recover accuracy without reactivating pruned connections. Comprehensive experiments across multiple image classification benchmarks, including CIFAR-10, CIFAR-100, and TinyImageNet, demonstrate that our method consistently achieves high sparsity levels with minimal performance degradation. Notably, our approach outperforms state-of-the-art baselines such as EPG and EPSD at high sparsity levels, while offering a more computationally efficient alternative to iterative pruning schemes like COLT. The proposed framework offers a computation-efficient, performance-preserving solution well suited for deployment in resource-constrained environments.