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

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. * Authors contributed equally. 35th Conference on Neural Information Processing Systems (NeurIPS 2021). impeded by difficulties in data sharing stemming from the privacy issues and limited consent of patients [50, 38, 53] . To alleviate this problem, the distributed machine learning methods, devised to enable the computation on multiple clients and servers leaving data to reside on the source devices, can be effectively leveraged for healthcare research [7, 43] . Federated learning (FL) is one of these methods which enables model training on a large corpus of decentralized data [26, 32, 59 ]. However, FL still holds several limitations in that it depends on clients' computational resources for its client-side parallel computation strategy for update and is not free from privacy concerns [27, 31, 52] . In contrast to FL, another distributed machine learning method, split learning (SL) offers better privacy and requires lower computational resources of clients by splitting the network between clients and the server [22, 52] , but still possess problems that it shows significant slower convergence than FL and can not learn under non-independent and non-identically distributed (non-IID) data [20] . Especially under unprecedented pandemic of an emerging pathogen like COVID-19, under which direct multi-national collaboration is deterred for prevention of epidemics, the collaboration via these distributed machine learning approaches is becoming increasingly important, since these enable to build a model with performance tantamount to data-centralized learning without any direct sharing of raw data between institutions to offer privacy. Recently proposed Vision Transformer (ViT) architecture [14], inspired by astounding results of Transformer-based models on natural language processing (NLP), have demonstrated impeccable performance on many vision tasks by enabling to model long dependencies within images. Besides this strength, the straightforward design of the Transformer allows to easily decompose the entire network into parts: the head for extracting features from the input image, the Transformer body to model the dependency between features, and the tail used for mapping features to task-specific output. One of the important contributions in this paper is the observation that this configuration is optimal for SL where a network should be split into the parts for clients and servers. In addition, as suggested in [9] , the Transformer body with sufficient capacity can be shared between various tasks, being suitable for multi-task learning (MTL) to leverage robust representation from multiple related tasks to enhance the generalization performance of individual tasks. Accordingly, here we propose a novel Federated Split Task-Agnostic (FESTA) framework equipped with a Transformer to simultaneously process multiple chest X-ray (CXR) tasks including diagnosis of COVID-19, emulating a real collaboration between several hospitals. To validate the practicability of FESTA, we also implemented the framework using a friendly federated learning framework (Flower) protocol [2], confirming seamless integration of various components. Experimental results show t