Federated Split Vision Transformer for COVID-19 CXR Diagnosis using Task-Agnostic Training

Abstract

Federated learning, which shares the weights of the neural network across clients, is gaining attention in the healthcare sector as it enables training on a large corpus of decentralized data while maintaining data privacy. For example, this enables neural network training for COVID-19 diagnosis on chest X-ray (CXR) images without collecting patient CXR data across multiple hospitals. Unfortunately, the exchange of the weights quickly consumes the network bandwidth if highly expressive network architecture is employed. So-called split learning partially solves this problem by dividing a neural network into a client and a server part, so that the client part of the network takes up less extensive computation resources and bandwidth. However, it is not clear how to find the optimal split without sacrificing the overall network performance. To amalgamate these methods and thereby maximize their distinct strengths, here we show that the Vision Transformer, a recently developed deep learning architecture with straightforward decomposable configuration, is ideally suitable for split learning without sacrificing performance. Even under the non-independent and identically distributed data distribution which emulates a real collaboration between hospitals using CXR datasets from multiple sources, the proposed framework was able to attain performance comparable to data-centralized training. In addition, the proposed framework along with heterogeneous multi-task clients also improves individual task performances including the diagnosis of COVID-19, eliminating the need for sharing large weights with innumerable parameters. Our results affirm the suitability of Transformer for collaborative learning in medical imaging and pave the way forward for future real-world implementations.

arXiv commentAccepted for NeurIPS 2021

References in corpus (22)

• An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale
• Communication-Efficient Learning of Deep Networks from Decentralized Data
• TransUNet: Transformers Make Strong Encoders for Medical Image Segmentation
• Transformers in Vision: A Survey
• Federated Learning: Strategies for Improving Communication Efficiency
• Federated Learning for Mobile Keyboard Prediction
• Learning Differentially Private Recurrent Language Models
• Deep Learning in Medical Image Registration: A Review
• iDLG: Improved Deep Leakage from Gradients
• Automatic Instrument Segmentation in Robot-Assisted Surgery Using Deep Learning
• Embracing Imperfect Datasets: A Review of Deep Learning Solutions for Medical Image Segmentation
• LEAF: A Benchmark for Federated Settings
• Split learning for health: Distributed deep learning without sharing raw patient data
• Federated Learning: Opportunities and Challenges
• BS-Net: learning COVID-19 pneumonia severity on a large Chest X-Ray dataset
• Deep Learning for Automatic Pneumonia Detection
• Pre-Trained Image Processing Transformer
• Variational Federated Multi-Task Learning
• Flower: A Friendly Federated Learning Research Framework
• Weakly Supervised Lesion Localization With Probabilistic-CAM Pooling
• End-to-End Evaluation of Federated Learning and Split Learning for Internet of Things
• Vision Transformer using Low-level Chest X-ray Feature Corpus for COVID-19 Diagnosis and Severity Quantification

Cited by in corpus (2)

• Recent Progress in Transformer-based Medical Image Analysis
• Multi-Modal Federated Learning for Cancer Staging over Non-IID Datasets with Unbalanced Modalities