Uncertainty-Aware Multi-Expert Knowledge Distillation for Imbalanced Disease Grading (2505.00592v1)

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.

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

Disease image grading, an essential application of AI in healthcare, faces notable challenges due to domain shifts and data imbalance which hinder model accuracy and deployment in clinical settings. The paper "Uncertainty-Aware Multi-Expert Knowledge Distillation for Imbalanced Disease Grading" presents a novel framework, UMKD, designed to tackle these issues by effectively distilling knowledge from multiple expert models into a singular student model. This framework aims to enhance the reliability and performance of disease grading systems, specifically focusing on histological prostate grading (SICAPv2) and fundus image grading (APTOS).

Technical Contributions

UMKD introduces significant methodological innovations:

1. Feature Alignment Mechanisms: The framework encompasses both shallow feature alignment (SFA) and compact feature alignment (CFA). These mechanisms decouple structural (task-agnostic) features from semantic (task-specific) features, ensuring that the student model can capture intricate pathological details effectively. SFA applies multi-scale low-pass filtering, preserving essential image structures, while CFA projects features into a common spherical space, facilitating consistent knowledge transfer.
2. Uncertainty-aware Decoupled Distillation (UDD): The paper addresses expert model bias engendered by class imbalance using an uncertainty-aware mechanism. UDD dynamically adjusts the weight of knowledge transferred based on the uncertainty in expert model predictions, thus mitigating bias propagation and fostering robust distillation. This approach ensures that the student model remains reliable even when trained with heterogeneous architectures.

Empirical Validation

Extensive experiments substantiate UMKD's efficacy across distinct scenarios:

• Source-Imbalanced Distillation: In which expert models are trained on imbalanced datasets but tested on balanced target datasets. UMKD outperforms existing state-of-the-art (SOTA) methods, achieving higher mean accuracy and overall performance metrics, thereby demonstrating its ability to correct expert bias.
• Target-Imbalanced Distillation: Where expert models are trained on balanced datasets and tested on imbalanced target datasets. Here, UMKD consistently surpasses SOTA methods, including the prominent feature-based Relational Knowledge Distillation (RKD) and logits-based methods, emphasizing its capability to handle real-world data imbalances effectively.

Implications and Future Work

The implementation of UMKD has profound implications for clinical practice. By mitigating bias and enhancing grading accuracy, it promises improved patient assessments, chiefly in conditions like diabetic retinopathy and prostate cancer, where early detection and intervention are crucial. The paper sets a precedent for future research focused on uncertainty-driven knowledge distillation techniques in medical imaging. Future efforts might aim to expand UMKD's applicability to a broader range of medical conditions and explore its integration into existing clinical workflows to further improve decision-making processes.

In summary, UMKD represents a robust and practical advancement in medical image analysis, leveraging uncertainty and expert diversity to enhance the precision and reliability of disease grading systems.