A Federated Multi-stage Light-Weight Vision Transformer for Respiratory Disease Detection

Abstract

Artificial Intelligence-based computer-aided diagnosis (CAD) has been widely applied to assist medical professionals in several medical applications. Although there are many studies on respiratory disease detection using Deep Learning (DL) approaches from radiographic images, the limited availability of public datasets limits their interpretation and generalization capacity. However, radiography images are available through different organizations in various countries. This condition is suited for Federated Learning (FL) training, which can collaborate with different institutes to use private data and train a global model. In FL, the local model on the client’s end is critical because there must be a balance between the model’s accuracy, communication cost, and client-side memory usage. The current DL or Vision Transformer (ViT)-based models have large parameters, making the client-side memory and communication costs a significant bottleneck when applied to FL training. The existing state-of-the-art (SOTA) FL techniques on respiratory disease detection either use small CNNs with insufficient accuracy or assume clients have sufficient processing capacity to train large models, which remains a significant challenge in practical applications. In this study, we tried to find one question: Is it possible to maintain higher accuracy while lowering the model parameters, leading to lower memory requirements and communication costs? To address this problem, we propose a federated multi-stage, light-weight ViT framework that combines the strengths of CNNs and ViTs to build an efficient FL framework. We conduct extensive experiments and show that the proposed framework outperforms a set of current SOTA models in FL training with higher accuracy while lowering communication costs and memory requirements. We adapted Grad-CAM for the infection localization and compared it with an experienced radiologist’s findings.