Uncertainty-Aware Multi-Expert Knowledge Distillation for Imbalanced Disease Grading

TL;DR

The paper tackles imbalance-driven domain shifts in automatic disease image grading by proposing UMKD, an uncertainty-aware multi-expert knowledge distillation framework. It decouples image representations into task-agnostic and task-specific components through frequency-domain shallow feature alignment (SFA) and compact feature alignment (CFA) in a common spherical space, while an uncertainty-aware distillation (UDD) module dynamically weights knowledge transfer based on expert uncertainties. The approach yields state-of-the-art performance on SICAPv2 (prostate grading) and strong results on APTOS (fundus grading) under both source-imbalanced and target-imbalanced conditions, with ablations confirming the contribution of each component. By addressing model heterogeneity and distribution discrepancies, UMKD offers a robust, practical solution for reliable disease grading in real-world clinical settings.

Abstract

Automatic disease image grading is a significant application of artificial intelligence for healthcare, enabling faster and more accurate patient assessments. However, domain shifts, which are exacerbated by data imbalance, introduce bias into the model, posing deployment difficulties in clinical applications. To address the problem, we propose a novel \textbf{U}ncertainty-aware \textbf{M}ulti-experts \textbf{K}nowledge \textbf{D}istillation (UMKD) framework to transfer knowledge from multiple expert models to a single student model. Specifically, to extract discriminative features, UMKD decouples task-agnostic and task-specific features with shallow and compact feature alignment in the feature space. At the output space, an uncertainty-aware decoupled distillation (UDD) mechanism dynamically adjusts knowledge transfer weights based on expert model uncertainties, ensuring robust and reliable distillation. Additionally, UMKD also tackles the problems of model architecture heterogeneity and distribution discrepancies between source and target domains, which are inadequately tackled by previous KD approaches. Extensive experiments on histology prostate grading (\textit{SICAPv2}) and fundus image grading (\textit{APTOS}) demonstrate that UMKD achieves a new state-of-the-art in both source-imbalanced and target-imbalanced scenarios, offering a robust and practical solution for real-world disease image grading.

Figures (2)

• Figure 1: Domain shifts between source and target data (left) and the performance of methods (right) for sources-imbalanced and target-imbalanced KD tasks.
• Figure 2: Model of uncertainty-aware multi-expert knowledge distillation.