Explainable Knowledge Distillation for On-device Chest X-Ray Classification
This work addresses the need for efficient AI models in clinical settings with limited hardware, though it is incremental as it applies existing knowledge distillation methods to a specific domain.
The paper tackled the problem of deploying deep learning models for multi-label chest X-ray classification on compact devices by proposing a knowledge distillation strategy to create a smaller, efficient student model, achieving AUC scores of up to 88.7% with 4.7 million parameters and 0.3 billion FLOPS.
Automated multi-label chest X-rays (CXR) image classification has achieved substantial progress in clinical diagnosis via utilizing sophisticated deep learning approaches. However, most deep models have high computational demands, which makes them less feasible for compact devices with low computational requirements. To overcome this problem, we propose a knowledge distillation (KD) strategy to create the compact deep learning model for the real-time multi-label CXR image classification. We study different alternatives of CNNs and Transforms as the teacher to distill the knowledge to a smaller student. Then, we employed explainable artificial intelligence (XAI) to provide the visual explanation for the model decision improved by the KD. Our results on three benchmark CXR datasets show that our KD strategy provides the improved performance on the compact student model, thus being the feasible choice for many limited hardware platforms. For instance, when using DenseNet161 as the teacher network, EEEA-Net-C2 achieved an AUC of 83.7%, 87.1%, and 88.7% on the ChestX-ray14, CheXpert, and PadChest datasets, respectively, with fewer parameters of 4.7 million and computational cost of 0.3 billion FLOPS.