KDAS: Knowledge Distillation via Attention Supervision Framework for Polyp Segmentation

TL;DR

This work tackles the challenge of achieving high-accuracy polyp segmentation with lightweight models. It proposes KDAS, a Knowledge Distillation framework that uses Attention-Supervision (Attention-KD) and a Symmetrical Guiding Module (SGM) to transfer knowledge from a frozen teacher to a trainable student, while addressing feature-inconsistency between the two via covariance-driven guidance. The method combines attention-based distillation losses $\\mathcal{L}_{AT}$ and $\\mathcal{L}_{SGM}$ with a segmentation loss $\\mathcal{L}_{seg}$ into a final objective $\\mathcal{L}_{KDAS}$, yielding a compact model (~3.7M parameters) that achieves competitive Dice and IoU scores and outperforms several real-time methods on multiple datasets. Experiments on merged and unseen datasets (Kvasir-SEG, ClinicDB, ColonDB, CVC-300, ETIS) demonstrate KDAS’s effectiveness and generalization, highlighting its potential for practical deployment in clinical settings. The work suggests a viable path to deploy accurate, efficient polyp segmentation tools in real-world medical imaging workflows.

Abstract

Polyp segmentation, a contentious issue in medical imaging, has seen numerous proposed methods aimed at improving the quality of segmented masks. While current state-of-the-art techniques yield impressive results, the size and computational cost of these models create challenges for practical industry applications. To address this challenge, we present KDAS, a Knowledge Distillation framework that incorporates attention supervision, and our proposed Symmetrical Guiding Module. This framework is designed to facilitate a compact student model with fewer parameters, allowing it to learn the strengths of the teacher model and mitigate the inconsistency between teacher features and student features, a common challenge in Knowledge Distillation, via the Symmetrical Guiding Module. Through extensive experiments, our compact models demonstrate their strength by achieving competitive results with state-of-the-art methods, offering a promising approach to creating compact models with high accuracy for polyp segmentation and in the medical imaging field. The implementation is available on https://github.com/huyquoctrinh/KDAS.

Figures (2)

• Figure 1: General KDAS framework
• Figure 2: Qualitative results for the improvement by integrating KDAS