Simplifying Knowledge Transfer in Pretrained Models

AI summary

Overview

Research area: Machine learning — knowledge distillation and transfer learning between pretrained neural networks.

Technical level: Intermediate. Familiarity with knowledge distillation, softmax outputs, KL divergence, and standard vision backbones (ViT, ResNet, Mask2Former) is assumed.

Scope: The paper proposes Bi-KD, a single-stage method that lets two or more pretrained models teach each other on a per-sample basis, improving every participant simultaneously across image classification, semantic segmentation, and video saliency prediction.

What This Paper Is About

Standard knowledge distillation transfers knowledge one-way — from a frozen, pretrained "teacher" to a "student." The authors argue this ignores a key fact: pretrained models trained with different designs, objectives, or datasets make different mistakes, so each holds complementary knowledge the other lacks. They introduce a method where pretrained models dynamically swap roles as teacher or student on a per-sample basis, letting both improve together in one training stage.

Key Contributions

1. Bidirectional knowledge transfer between pretrained models. A per-sample data-partitioning scheme where the model with the highest confidence (classification) or lowest loss (dense tasks) on a given instance temporarily acts as teacher, with a stop-gradient applied so only the student updates from the distillation signal.

2. Dynamic, evolving partition rather than a fixed one. Unlike prior work (Roth et al., 2024) that fixes teacher/student roles for the whole training run, Bi-KD reassigns the teacher role every batch, so the partition adapts as models improve.

3. Consistent gains across three distinct tasks. Improvements on ImageNet classification, ADE20K semantic segmentation (within and across backbone architectures), and video saliency prediction, where the method sets new state-of-the-art results on DHF1K and Hollywood-2.

4. Extension to multidirectional transfer among K models. A general formulation where, for each sample, the most confident model teaches all others, yielding gains that grow as more models join the collaboration.

Main Findings

• Both models improve in one training stage. On ImageNet, Bi-KD beat fixed-partition KL+DP in 19 of 22 comparisons and improved every participating model, including pairs where one model was weaker, smaller, or differently pretrained.

• ViT-B gains ~1.4% Top-1 accuracy when paired with the tiny ViT-T. Gains held even for same-architecture pairs (ViT-B/ViT-S/ViT-T), and the smaller ViT-T often transferred more useful knowledge to larger ViTs than the other way around.

• Diversity of training method matters more than raw accuracy. Pairing self-supervised DINOv2 with any supervised model produced gains above 1% on both sides, supporting the idea that models with different training methodologies make uncorrelated errors.

• Classical feature-space convergence. Canonical Correlation Analysis on ViT-S and ViT-T showed mean correlation of features rising from 0.6768 to 0.7487 after Bi-KD, indicating the two models' internal representations become more aligned.

• Individual models recover much of the ensemble gain. Using a "Recovered" metric, larger models (ViT-B, SeNet154, SWSL-ResNext101, MAE, PiT-B) captured around 60% of the ensemble's improvement on average; smaller models (DINOv2, ViT-S, ViT-T) around 40%. ViT-B approached 90% in one case and exceeded the ensemble in another.

• Case-analysis of distillation labels. ViT-S + ViT-T analysis showed Case 1 (both correct) rose from 70.9% to 73.1% and Case 2 (teacher correct, student wrong) dropped from 12.5% to 10.1% after Bi-KD, while Case 3 (both wrong) stayed near 16.6%, showing limited harm from imperfect teacher signals.

• Semantic segmentation gains across and within architectures. Mask2Former with Swin-T backbone improved mIoU from 47.7 to 48.32 when paired with R50, and to 48.45 when paired with the stronger Swin-S. All four metrics (mIoU, fwIoU, mACC, pACC) improved.

• New state-of-the-art in video saliency prediction. On Hollywood-2, ViNet-A (Bi-KD) raised CC from 0.756 to 0.762 and NSS from 4.119 to 4.198. On DHF1K, TMFI-Net (Bi-KD) set a new SOTA at CC 0.558 and NSS 3.216.

• Multidirectional scaling holds. ImageNet gains grew as team size increased: CoaT-lite-mini went from +0.38 (2 models) to +0.46 (3) to +0.572 (4), and PiT-XS from +0.45 to +0.476 to +0.63.

Methodology in Plain English

Imagine two students who have each studied different textbooks and developed different strengths. Instead of one always lecturing the other, on every practice problem the pair compares notes: whoever is more confident about the right answer explains their reasoning, and the other adjusts. On the next problem the roles may flip.

That is Bi-KD. For each mini-batch, both pretrained models produce predictions. The model with the higher softmax probability on the true class (or lower task loss for dense tasks such as segmentation and saliency, where per-pixel labels make confidence less meaningful) is assigned the teacher role for that instance. A stop-gradient blocks updates through the teacher's path, so only the student is pulled toward the teacher's softened output via a KL-divergence loss. Both models still receive their own task-specific loss (cross-entropy, mask loss, etc.) to keep them grounded on the ground truth and to stabilize training.

The overall loss is simply the sum of the two task losses plus the masked distillation loss. No data augmentation, learning-rate schedule, or layer-wise tricks are used except for standard timm defaults on the ViT family. Training runs for 20 epochs at batch size 128 on a single NVIDIA RTX A6000.

For multiple models, the rule generalizes straightforwardly: for each sample, the single most confident model becomes the teacher and every other model distills from it, with the same stop-gradient treatment. This lets the framework scale to K models without architectural coupling or extra training stages.

Why This Matters

Impact on research. The paper reframes knowledge distillation as a mutual, simultaneous process rather than a one-way compression pipeline. It shows that pretrained models are valuable exchange partners for each other even without any new data, architectural surgery, or multi-stage pipelines, and it directly extends the data-partitioning line of work by Roth et al. with a dynamic, self-adapting teacher selection rule. It also provides mechanistic evidence (CCA alignment, case-level prediction analysis) that the improvement is real representation convergence, not just noise.

Real-world applications.

• Model zoo curation: Public repositories of pretrained checkpoints (HuggingFace, timm) could be mined for complementary pairs and improved in-place with a single training stage.
• Medical and scientific imaging: Where labeled data is scarce but many domain-pretrained models exist, mutual distillation can sharpen each specialist without leaking private data.
• Edge deployment: Pairs of small models can each be improved without introducing a large teacher at inference time.
• Video understanding pipelines: Saliency and segmentation models used in content editing, autofocus, and compression can be upgraded in tandem, as demonstrated by the SOTA gains on DHF1K and Hollywood-2.

Industry relevance. The method is inexpensive: one GPU, 20 epochs, no extra unlabeled data, no architecture change, and no inference-time overhead. It slots into existing distillation codebases and works across ResNets, ViTs, MLP-Mixers, and hybrid models. For teams that already maintain a catalog of pretrained checkpoints, Bi-KD offers a cheap path to squeeze out additional accuracy — and the "Recovered" analysis shows individual models can absorb roughly half of what a full ensemble would deliver, at no inference cost.

Future Directions

• Scaling to large model populations. The paper demonstrates up to four models; open questions remain about how the teacher-selection rule, computational cost, and convergence behave with dozens or hundreds of checkpoints of varying quality.

• Beyond the confidence-based selection rule. Using only ground-truth-class probability is simple but crude, especially in Case 3 (both models wrong). Entropy-based, calibration-aware, or learned teacher-selection criteria could improve the partition quality.

• Knowledge transfer across modalities and tasks. The current experiments stay within vision (classification, segmentation, saliency). Extending Bi-KD to cross-modal pretrained pairs — vision-language, audio-visual, or code-text — is an obvious next step.

• Theoretical grounding. There is no formal account yet of when bidirectional transfer is guaranteed to help versus when two models can reinforce a shared error. A theory of complementary error structure could predict which pairs are worth distilling.

Target Audience

Researchers and engineers working on knowledge distillation, transfer learning, or model compression will find this paper most useful, especially those who already keep a pool of pretrained checkpoints and want a low-cost way to improve them. Practitioners in computer vision — particularly in segmentation and video saliency — will find directly applicable results. Readers new to distillation can follow the algorithm, but should first understand softmax temperature, KL divergence, and standard teacher-student framing to appreciate why the dynamic teacher selection is the key novelty.

Authors’ abstract

Pretrained models are ubiquitous in the current deep learning landscape, offering strong results on a broad range of tasks. Recent works have shown that models differing in various design choices exhibit categorically diverse generalization behavior, resulting in one model grasping distinct data-specific insights unavailable to the other. In this paper, we propose to leverage large publicly available model repositories as an auxiliary source of model improvements. We introduce a data partitioning strategy where pretrained models autonomously adopt either the role of a student, seeking knowledge, or that of a teacher, imparting knowledge. Experiments across various tasks demonstrate the effectiveness of our proposed approach. In image classification, we improved the performance of ViT-B by approximately 1.4% through bidirectional knowledge transfer with ViT-T. For semantic segmentation, our method boosted all evaluation metrics by enabling knowledge transfer both within and across backbone architectures. In video saliency prediction, our approach achieved a new state-of-the-art. We further extend our approach to knowledge transfer between multiple models, leading to considerable performance improvements for all model participants.