Source-linked AI summary
One-for-All: Bridge the Gap Between Heterogeneous Architectures in Knowledge Distillation
Zhiwei Hao, Jianyuan Guo, Kai Han, Yehui Tang, Han Hu, Yunhe Wang, Chang Xu
TL;DR
Existing hint-based knowledge distillation largely assumes homogeneous architectures, while heterogeneous models learn distinct features that complicate direct alignment. OFA-KD projects intermediate features into logits space and adaptively enhances target information, consistently outperforming baselines across heterogeneous settings, with gains up to 8.00% on CIFAR-100 and 0.71% on ImageNet-1K. The approach nevertheless has architecture-dependent limitations and requires tuning a modulating parameter.
Problem
Existing hint-based distillation assumes homogeneous architectures, but heterogeneous models learn distinct features that impede direct feature matching.
Method
OFA-KD uses student exit branches to project intermediate features into logits space and adds a modulating parameter for adaptive target-information enhancement.
Results
OFA-KD consistently outperforms baselines in heterogeneous distillation, with gains of 0.28% to 8.00% on CIFAR-100 and up to 0.71% on ImageNet-1K.
Takeaways & Limitations
Projecting away architecture-specific feature information makes hint-based distillation feasible across CNN, Transformer, and MLP architectures.
Takeaways & Limitations
For some architectures such as ResNet18, heterogeneous-teacher distillation underperforms homogeneous-teacher distillation, and the modulating parameter requires optimal tuning.
Abstract
from arXiv · showhide
Knowledge distillation~(KD) has proven to be a highly effective approach for enhancing model performance through a teacher-student training scheme. However, most existing distillation methods are designed under the assumption that the teacher and student models belong to the same model family, particularly the hint-based approaches. By using centered kernel alignment (CKA) to compare the learned features between heterogeneous teacher and student models, we observe significant feature divergence. This divergence illustrates the ineffectiveness of previous hint-based methods in cross-architecture distillation. To tackle the challenge in distilling heterogeneous models, we propose a simple yet effective one-for-all KD framework called OFA-KD, which significantly improves the distillation performance between heterogeneous architectures. Specifically, we project intermediate features into an aligned latent space such as the logits space, where architecture-specific information is discarded. Additionally, we introduce an adaptive target enhancement scheme to prevent the student from being disturbed by irrelevant information. Extensive experiments with various architectures, including CNN, Transformer, and MLP, demonstrate the superiority of our OFA-KD framework in enabling distillation between heterogeneous architectures. Specifically, when equipped with our OFA-KD, the student models achieve notable performance improvements, with a maximum gain of 8.0% on the CIFAR-100 dataset and 0.7% on the ImageNet-1K dataset. PyTorch code and checkpoints can be found at https://github.com/Hao840/OFAKD.
1 Introduction
Knowledge distillation trains compact students to learn from pretrained teachers, but heterogeneous architectures create feature-space mismatches that undermine direct hint matching. OFA-KD addresses this by transferring intermediate representations through logits space and adaptively emphasizing useful target information.
- 1 Introduction: Knowledge distillation trains a compact student to mimic outputs or soft labels from a pretrained teacher, with intermediate features later added as hint knowledge.Hint-based approaches aim to make student representations resemble teacher representations.
- 1 Introduction: Existing hint-based methods mainly target homogeneous teacher-student architectures, leaving cross-architecture distillation relatively unexplored despite practical needs.The introduction identifies CNN, Transformer, and MLP combinations as relevant heterogeneous settings.
- 1 Introduction: Heterogeneous models learn features in different latent spaces, so directly matching intermediate representations has no guaranteed successful alignment.The mismatch reflects architecture-specific representation differences.
- 1 Introduction: OFA-KD projects intermediate features into aligned logits space through student exit branches, discarding architecture-specific information before matching teacher outputs.The framework targets distillation across CNNs, Transformers, and MLPs.
- 1 Introduction: OFA-KD adaptively enhances target information using teacher predictive confidence to reduce irrelevant information in heterogeneous logits.The modified loss adds a modulating parameter to the reformulated vanilla KD loss.
2 Related works
Related work covers Transformer and MLP vision architectures alongside logits-based and hint-based knowledge distillation. Existing hint-based methods generally assume homogeneous architectures, leaving heterogeneous feature alignment unresolved.
- Related architectures: Transformers use attention to capture long-term dependencies, while Vision Transformers process non-overlapping image patches as embedding tokens.Vision Transformer designs adapt the Transformer architecture from NLP to computer vision.
- Related architectures: MLP-Mixer alternately mixes channel and spatial information from image-patch tokens and performs comparably to leading CNN and ViT models.ResMLP is another MLP architecture introduced for vision tasks.
- Related architectures: CNN, Transformer, and MLP architectures achieve similar performance while retaining distinct inductive biases and representation-learning preferences.These differences motivate examining feature alignment across architecture families.
- Knowledge distillation: Knowledge distillation trains lightweight students to mimic pretrained teacher logits, with later methods using structural information, ensembles, contrastive learning, or intermediate features.Hint-based distillation projects student features toward teacher feature space and can also match attention maps or solution procedures.
- Knowledge distillation: Existing hint-based methods assume homogeneous architectures, whereas heterogeneous student and teacher features may be distinct enough to make direct distillation fail.The text identifies heterogeneous hint-based distillation as an open problem.
3 Method
OFA-KD addresses the feature-space mismatch that limits hint-based distillation across heterogeneous architectures by aligning intermediate representations in logits space and adaptively enhancing target information.
- Discrepancy among architectures: CNNs, Transformers, and MLPs develop distinct representations because their inductive biases differ, making direct feature matching unreliable across architectures.CNNs emphasize locality and spatial invariance, whereas Transformers and MLPs more readily learn global representations.
- Centered kernel alignment analysis: CKA compares cross-architecture features despite differing dimensions and shows that heterogeneous models have markedly different layerwise representations.Homogeneous models tend to align at similarly positioned layers, while MobileNetV2 resembles only shallower layers of ViT-Small or Mixer-B/16.
- Challenges in heterogeneous feature distillation: Direct hint-based distillation is therefore unsuitable for heterogeneous models, while logits-only distillation lacks intermediate supervision needed for potentially better results.The analysis motivates new intermediate-layer designs for heterogeneous knowledge distillation.
- Learning in the logits space: OFA-KD adds student exit branches that project intermediate features into logits space and matches them to the teacher’s classifier outputs.Logits space contains less architecture-specific information than intermediate feature spaces, enabling cross-architecture intermediate distillation.
- Generic heterogeneous knowledge distillation: OFA-KD targets distillation across CNN, Transformer, and MLP architectures through aligned logits and multibranch learning.The framework is presented as a one-for-all method intended to perform cross-architecture KD between these model families.
- Adaptive target information enhancement: The OFA loss adaptively enhances target-class information according to teacher confidence, reducing learning from less confident teacher predictions.Its high-order target term decays more slowly for confident target predictions and faster otherwise; γ = 1 recovers logits-based KD.
4 Experiment
Experiments evaluate OFA-KD across heterogeneous CNN, Transformer, and MLP models on ImageNet-1K and CIFAR-100, with ablations and homogeneous-architecture comparisons. OFA-KD consistently outperforms baselines in heterogeneous settings, with gains of up to 8.00% on CIFAR-100 and 0.77% for CNN students on ImageNet-1K.
- Experimental setup: Experiments cover CNN, Transformer, and MLP models on CIFAR-100 and ImageNet-1K, using logits-based and hint-based KD baselines.ImageNet-1K includes fifteen heterogeneous teacher-student combinations; CIFAR-100 includes twelve.
- ImageNet-1K results: 0.20% to 0.77%: OFA-KD improves over second-best baselines when the ImageNet-1K student uses a CNN architecture.For ViT or MLP students, the maximum reported accuracy gain is 0.71%.
- ImageNet-1K results: OFA-KD consistently outperforms all baselines across the evaluated heterogeneous ImageNet-1K combinations.The second-best method varies among FitNet, vanilla KD, and DIST, while OFA-KD remains best in every case.
- CIFAR-100 results: 0.28% to 8.00%: OFA-KD improves over second-best baselines on CIFAR-100 heterogeneous distillation.Hint-based approaches are particularly weak with ViT or MLP students; FitNet reaches 24.06% for ConvNeXt-T to Swin-P.
- CIFAR-100 results: OFA-KD adaptively enhances target information, while DIST is more frequent second-best on ImageNet-1K and DKD on CIFAR-100.The paper attributes this pattern to DIST favoring stronger ImageNet-1K teachers and DKD benefiting from smaller CIFAR-100 teachers.
- Ablation study: Learning at the end of all four student stages is the best exit-branch configuration for OFA-KD.ResNet18 shows no preference among branch positions, whereas the overall ablation identifies all four stage ends as best.
- Ablation study: The optimal modulation parameter γ depends on the teacher-student pair: 1.4 for DeiT-T to ResNet18 and 1.1 for ResNet50 to DeiT-T.The paper relates the larger γ to compensating for additional interference from the weaker DeiT-T teacher.
- Homogeneous architectures: OFA matches the best distillation baseline for homogeneous ResNet34-to-ResNet18 distillation on ImageNet-1K.This comparison indicates that the framework also remains competitive when teacher and student architectures are homogeneous.
5 Conclusion
The paper proposes OFA-KD to make hint-based distillation feasible across heterogeneous architectures by projecting features into logits space and enhancing target information adaptively. Experiments demonstrate effectiveness, while the method remains sensitive to architecture choice and modulation-parameter tuning.
- OFA-KD projects intermediate student features into logits space through exit branches, discarding architecture-specific information for cross-architecture distillation.
- The method introduces a modulating parameter into a reformulated KD loss for adaptive target information enhancement.
- Experiments on CIFAR-100 and ImageNet-1K demonstrate the effectiveness of the proposed cross-architecture KD method.
- For some architectures, including ResNet18, heterogeneous-teacher distillation performs worse than homogeneous-teacher distillation.
- Optimal modulation-parameter selection adds complexity, and improper settings can produce inferior performance.
- OFA-KD suggests that heterogeneous architectures contain useful transferable information and motivates better alignment approaches.