Source-linked AI summary
Generalizable Knowledge Distillation from Vision Foundation Models for Semantic Segmentation
Chonghua Lv, Dong Zhao, Shuang Wang, Dou Quan, Ning Huyan, Nicu Sebe, Zhun Zhong
TL;DR
Conventional KD compresses segmentation models but can sacrifice the out-of-domain robustness of vision foundation models. GKD decouples domain-agnostic representation learning from task adaptation and uses query-based soft distillation to transfer spatial knowledge, outperforming existing methods across five benchmarks with average gains of +1.9% in F2F and +10.6% in F2L.
Problem
Conventional KD prioritizes in-domain accuracy, while distilling VFMs often fails to transfer their robustness to unseen domains.
Method
GKD separates domain-agnostic representation learning from task adaptation and uses query-based soft distillation to selectively retrieve transferable spatial knowledge.
Results
GKD consistently outperforms existing KD methods across five domain-generalization benchmarks, with average gains of +1.9% in F2F and +10.6% in F2L.
Takeaways & Limitations
GKD transfers VFM generalization to compact models under F2F and F2L settings, including strong performance with limited annotations.
Abstract
from arXiv · showhide
Knowledge distillation (KD) has been widely applied in semantic segmentation to compress large models, but conventional approaches primarily preserve in-domain accuracy while neglecting out-of-domain generalization, which is essential under distribution shifts. This limitation becomes more severe with the emergence of vision foundation models (VFMs): although VFMs exhibit strong robustness on unseen data, distilling them with conventional KD often compromises this ability. We propose Generalizable Knowledge Distillation (GKD), a multi-stage framework that explicitly enhances generalization. GKD decouples representation learning from task learning. In the first stage, the student acquires domain-agnostic representations through selective feature distillation, and in the second stage, these representations are frozen for task adaptation, thereby mitigating overfitting to visible domains. To further support transfer, we introduce a query-based soft distillation mechanism, where student features act as queries to teacher representations to selectively retrieve transferable spatial knowledge from VFMs. Extensive experiments on five domain generalization benchmarks demonstrate that GKD consistently outperforms existing KD methods, achieving average gains of +1.9% in foundation-to-foundation (F2F) and +10.6% in foundation-to-local (F2L) distillation. The code will be available at https://github.com/Younger-hua/GKD.
1. Introduction
Conventional knowledge distillation compresses segmentation models but often fails to preserve out-of-domain generalization, especially when distilling vision foundation models. GKD addresses this by separating representation learning from task adaptation and selectively transferring spatial knowledge.
- Domain generalization is critical for segmentation because deployment spans changing weather, lighting, devices, and clinical sites.
- VFM distillation can compromise the strong unseen-domain robustness of foundation models, creating a need for compact models that retain this ability.
- F2F distillation uses foundation-model teachers and students, whereas F2L distillation transfers from a large foundation model to a small locally trained model.
- Conventional KD and enhanced variants can produce students that generalize worse than teachers, particularly in the F2L setting.
- GKD first learns domain-agnostic representations through selective feature distillation, then freezes them during supervised task adaptation to mitigate domain overfitting.
- Query-based soft distillation uses student features as queries to retrieve transferable spatial knowledge from teacher representations through attention.
2. Related work
Prior segmentation distillation methods transfer semantic structure, while earlier multi-stage designs remain coupled to task supervision and can bias students toward source domains. GKD instead isolates domain-general representation learning and adds query-based retrieval of transferable spatial knowledge.
- Semantic-segmentation KD methods extend beyond classification by transferring structural semantic correlations and inter-class relations.
- Earlier multi-stage methods jointly optimize task supervision with feature or logit distillation, improving in-domain performance but biasing students toward source domains.
- GKD isolates domain-agnostic representation learning from task optimization to prevent domain overfitting during adaptation.
- Query-based soft distillation lets students selectively retrieve transferable spatial knowledge from teachers rather than relying only on coupled task-oriented distillation.
3. Methodology
GKD separates domain-general representation learning from task learning and uses query-based soft distillation to transfer spatial knowledge from vision foundation models. The student first learns through staged distillation, then adapts to segmentation while preserving its encoder representations.
- Task learning: During task learning, the student encoder is frozen and only the decoder is optimized with labeled source-domain supervision.Freezing the encoder preserves the distilled domain-general representations while grounding them in the segmentation task.
- Overall framework: GKD separates domain-general distillation for representation learning from task learning for downstream segmentation.The domain-general stage includes task-agnostic and domain-agnostic distillation, while task learning uses supervised segmentation training.
- Domain-general distillation: The student first distills task-agnostic knowledge on a diverse proxy dataset, then distills task-relevant domain-agnostic features on source images.The proxy dataset is used to provide generic visual representations, while source-domain distillation adds task-relevant features without domain-specific supervision bias.
- Query-based Soft Distillation: Query-based soft distillation uses student features as queries to retrieve teacher spatial semantics through attention rather than enforcing point-wise feature matching.The mechanism reweights spatial responses so the student internalizes relational structure and global dependencies from the teacher.
- Query-based Soft Distillation: The attention-based reconstruction redistributes student features with teacher-derived context, then aligns reconstructed features with teacher features using mean squared error.A linear projection adapts feature dimensions, softmax produces attention weights, and the reconstructed representation is constrained by feature-level MSE.
- Additional distillation objectives: Masked patch-level distillation and CLS-token distillation extend QSD to hidden patch knowledge and global semantics.Masked images provide masked student features, while the CLS token transfers global semantic information using the same reconstruction procedure.
4. Experiments
Experiments evaluate GKD across multiple domain-generalization benchmarks, distillation settings, source-domain counts, label fractions, and ablations. Results consistently support multi-stage optimization and query-based spatial distillation for transferring domain-general representations.
- Experimental Setup: GKD is evaluated on driving-scene and remote-sensing benchmarks under foundation-to-local and foundation-to-foundation settings.The setup includes five driving-scene datasets and two cross-urban remote-sensing datasets, with comparisons against conventional KD methods.
- Comparison with various KD Methods: 59.8% average mIoU is achieved by DINOv2-B with GKD on GTAV → Citys + BDD + Map, outperforming Vanilla KD by 1.6%.GKD also improves average mIoU from 66.2% to 68.0% on ACDC and from 64.1% to 65.9% on remote-sensing datasets.
- Scaling Up: Student performance steadily improves as more source domains are incorporated, even when additional domains are used only for distillation.This supports transfer of richer visual representations and domain-agnostic knowledge.
- Scaling Up: 51.4% mIoU is maintained with only 1/16 labels in F2L, exceeding Af-DCD by 5.4% and the vanilla student by 15.7%.Improvements also appear across other label fractions and target domains, including F2F settings.
- Ablation Study and Analysis: QSD produces more compact student-teacher feature distances and diagonal attention with off-diagonal aggregation of related teacher features.The visualization links selective spatial aggregation with internalization of domain-invariant structure.
- Ablation Study and Analysis: Single-stage MSE† and QSD† perform significantly worse, whereas multi-stage optimization and QSD improve cross-domain generalization.The comparison identifies both decoupled optimization and relational spatial knowledge as important for transferring domain-general representations.
- Ablation Study and Analysis: Domain-agnostic distillation contributes most of the ablation gain, while task-agnostic distillation and freezing the encoder provide further improvements.Freezing the encoder also reduces training cost and limits source-domain bias during task learning.
5. Conclusion
The paper concludes that GKD addresses conventional KD’s limited unseen-domain generalization by decoupling representation learning from task adaptation and selectively transferring VFM spatial knowledge. Across F2L and F2F benchmarks, it also performs strongly with limited annotations and additional source domains.
- GKD transfers robust VFM generalization to compact models while mitigating domain overfitting through decoupled representation learning and task-specific adaptation.