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

1 Introduction . After its earlier success in many fields , deep neural networks have found a pervasive suite of applications in healthcare research including medical imaging , becoming a new de facto standard [ 54 , 19 , 13 , 49 , 17 , 63 , 6 , 58 , 24 ] . Training these networks requires a vast amount of data to achieve robust performance [ 8 , 11 , 47 ] . Despite the fact that multi-center collaboration is mandatory due to the shortage of labeled data in a single institution , collaboration in healthcare research is heavily ∗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 that our framework can show stable performance even under non-IID settings which is a frequently faced situation for collaboration between hospitals while offering privacy , by amalgamating FL and SL to maximally exploit their main advantages . In addition , we show that the proposed FESTA Transformer along with MTL improves the performances of individual tasks . In summary , our contributions are two folds : • We proposed a novel FESTA learning framework equipped with the ViT by utilizing its decomposable design to amalgamate the merit of FL and SL . • We showed that the model trained with the FESTA framework can leverage the robust representations from multiple related tasks to improve the performance of the individual task . 2 Related Work . Distributed machine learning . FL is a distributed machine learning approach originally proposed by Google [ 26 ] to enable the training of a model via distributed devices and data . It eliminates the need to aggregate the raw data in a centralized way by enabling the model to be updated on the edge devices ( e.g . mobile phones , computers in hospitals ) . Specifically , during the training , the server initializes a global model and sends it to each client . The clients then train the model with their local data in parallel and return the updated model to the server . Then , the server aggregates and distributes those models by a method such as federated averaging ( FedAvg ) [ 32 ] to update the global model . This process ( called round ) continues repeatedly until the model converges . Though FL enables decentralized training in a privacy-preserving manner , it still holds limitations as it largely depends on the computation resources of clients and is vulnerable to model inversion attacks [ 56 ] . Different from FL , SL divides the neural network into several sub-networks , and these separated sub-networks are trained under distributed setting [ 16 ] . In detail , the first sub-network is trained on the client-side with local data and then passes the feature to the second sub-network located in the server . The server can access the only feature from the first sub-network , and train the second sub-network to send the subsequent feature to the third sub-network on the client . Finally , the third sub-network is able to provide the output of the overall network . By inserting black-box sub-networks in both client and server sides , it is possible to offer better privacy than FL . Besides the privacy benefit , SL uses less computational resources than FL at the client-side . Nevertheless , since the one cycle of forward and backward passes is finished after data and gradients move back and forth across the sub-networks distributed on multiple sides , the convergence of SL is considerably slower than FL . In addition , it was reported that the convergence is not reached at all under the non-IID setting [ 20 ] . Multi-task learning . MTL is a learning strategy to improve the generalization performance of a specific task with the help of information from other related tasks . In MTL , models for multiple related tasks are trained simultaneously [ 5 , 16 ] . In its early era , the motivation of MTL is to mitigate the data insufficiency problem where the number of labeled data is limited for each task to train an accurate learner . MTL helps to reuse knowledge and thereby reduce the requirement for labeled data for each task , exhibiting that the MTL model can achieve better performance than the single-task counterpart in many fields ranging from computer vision [ 15 , 34 , 36 , 28 ] to NLP [ 30 , 21 , 29 , 42 ] . With increased data , the MTL model can learn more robust representations via knowledge sharing among multiple tasks , resulting in improved performance and less overfitting for the individual task . Vision Transformer . Transformer [ 51 ] , which was originally developed for NLP , is a deep neural network based on an attention mechanism that utilizes an appreciably large receptive field . After achieving state-of-the-art ( SOTA ) performance in NLP , it has also inspired the vision community to explore its application on vision tasks to utilize its ability to model long-range dependency within an image [ 25 ] . The ViT was one of the successful attempts to apply Transformer directly to images , achieving excellent results compared to the SOTA convolutional neural networks in image classification tasks [ 14 ] . Furthermore , in addition to the superb performance , the straightforward modular design of ViT facilitates broad applications in many tasks only with minimal change . Chen et al . [ 9 ] proposed image processing transformer , which is one of the successful multi-task model for various computer vision tasks , by splitting ViT into shared body and task-specific heads and tails , suggesting that Transformer body with sufficient capacity can be shared across relevant tasks . However , they leveraged encoder-decoder design and the usefulness of the multi-task ViT model was not evaluated along with the distributed learning methods . Recently , ViT was successfully used for diagnosis and severity prediction of COVID-19 , showing the SOTA performance [ 35 ] . Specifically , to alleviate the overfitting problem with limited data available , the overall framework is decomposed into two steps : the pre-trained backbone network to classify common low-level CXR features , which was leveraged in the second step by Transformer for high-level diagnosis and severity prediction of COVID-19 . By maximally utilizing the merit of the large-scale database containing more than 220,000 CXR images , the model has attained stable generalization performance as well as SOTA performance in a variety of external test data from different institutions , even with the limited number of labeled cases for COVID-19 . 3 Split Task-Agnostic Transformer for CXR COVID-19 Diagnosis . Inspired by these works , here we are interested in utilizing ViT for distributed learning in COVID19 CXR diagnosis , where the collaboration via these distributed machine learning approaches is becoming increasingly important by offering privacy and still allowing similar performance to the data-centralized learning . The reason we are interested in ViT architecture for this purpose is that the natural configuration of ViT may be optimal for MTL as well as SL where the easily decomposable modular design of the network is preferred , suggesting a possibility to maximally reconcile the merits of MTL and SL through ViT architecture . Specifically , the clients just train the head and tail parts of the network , whereas the Transformer body is shared across multiple clients . Then , the embedded features from the head network from multiple clients can be leveraged in the second step by Transformer to process individual tasks including diagnosis of COVID-19 . Furthermore , by maximally utilizing the merit of the large-scale database from CXR for various tasks , our goal is to demonstrate stable generalization performance as well as SOTA performance in the external test dataset . 3.1 FESTA : Federated Split Task-Agnostic Learning . The concept of the proposed Federated Split Task-Agnostic ( FESTA ) learning process is illustrated in Fig 1 ( a ) and ( b ) . Let C = ⋃K k=1 Ck be a group of client sets with different CXR tasks , where K denotes the number of tasks and Ck has one or more clients with different datasets each other for the k-th task , i.e . Ck = { ck1 , ck2 , . . . , ckNk : Nk ≥ 1 } . Each client c ∈ Ck has its own task-specific network architecture for a headHc and a tail Tc , which are connected to Transformer B in the server . In our FeSTA framework , the server first initializes the weights of the Transformer body and taskspecific heads , tails for each task k. Then , it distributes the initialized weights of heads and tails to each client c ∈ Ck . For round i = 1 , 2 , . . . R , each client ( e.g . hospital ) perform the forward propagation on their task-specific head and pass their intermediate feature to the server . Specifically , using the local training data { ( x ( i ) c , y ( i ) c ) } Nci=1 , the head network Hc encodes the smashed feature maps h ( i ) c and sent it to the server : h ( i ) c = Hc ( x ( i ) c ) . Then , the server-side Transformer body B receives the feature from the all clients and generate the features b ( i ) c for c in parallel : b ( i ) c = B ( h ( i ) c ) . Resulting smashed features from the Transformer are allocated to task-specific tail to produce the final prediction ŷ ( i ) c = Tc ( b ( i ) c ) according to each task , and the forward path finishes . Subsequently , the loss for each client can be calculated as ℓc ( y ( i ) c , Tc ( B ( Hc ( x ( i ) c ) ) ) ) where ℓc ( y , ŷ ) refers to the c client- specific loss between the target y and the estimate ŷ . By minimizing the loss with respect to the tail weight , the gradients of the local tails are passed reversely to the server . After receiving the gradients , the server performs the back-propagation on the server-side body model , sends the gradients back to the clients . Specifically , for the task-agnostic body update , the following optimization problem is solved : min B ∑ c∈C Nc∑ i=1 ℓc ( y ( i ) c , Tc ( B ( Hc ( x ( i ) c ) ) ) ) , ( 1 ) For the task-specific fine-tuning , the following optimization problem is solved : min Hc , Tc Nc∑ i=1 ℓc ( y ( i ) c , Tc ( B ( Hc ( x ( i ) c ) ) ) ) , ( 2 ) Finally , the server aggregates and averages the weights of local heads and tails from the clients to update the global heads and tails among the clients for the same tasks via FedAvg , and distributes back the updated global weights of heads and tails for each task k to the clients . The algorithm is formally presented in Algorithm 1 . Algorithm 1 : FESTA : Federated Split Task-Agnostic learning 1 Function ServerMain : 2 Initialize the body weight w ( 1 ) B and client head/tail weights ( w̄H , k , w̄T , k ) for each task k ∈ { 1 , ... , K } in server 3 for rounds i = 1 , 2 , . . . R do 4 for tasks k ∈ { 1 , 2 , . . .K } do in parallel 5 for clients c ∈ Ck do in parallel 6 if i = 1 or ( i− 1 ) ∈ UnifyingRounds then 7 Set client ( w ( i ) Hc , w ( i ) Tc ) ← ( w̄H , k , w̄T , k ) 8 h ( i ) c ← ClientHead ( c ) 9 b ( i ) c ← B ( h ( i ) c ) 10 ∂L ( i ) c ∂b ( i ) c ← ClientTail ( c , b ( i ) c ) & Backprop . 11 ( w ( i+1 ) Hc , w ( i+1 ) Tc ) ← ClientUpdate ( c , ∂L ( i ) c ∂h ( i ) c ) 12 Update body w ( i+1 ) B ← w ( i ) B − η K K∑ k=1 ∑ c∈Ck ∂L ( i ) c Nk∂w ( i ) B 13 if i ∈ UnifyingRounds then 14 for tasks k ∈ { 1 , 2 , . . .K } do 15 Update ( w̄H , k , w̄T , k ) ← ( 1Nk ∑ c∈Ck w ( i+1 ) Hc , 1 Nk ∑ c∈Ck w ( i+1 ) Tc ) 16 Function ClientHead ( c ) : 17 xc ← Current batch of input from client c 18 returnHc ( xc ) 19 Function ClientTail ( c , bc ) : 20 yc ← Current batch of label from client c 21 Lc ← ℓc ( yc , Tc ( bc ) ) & Backprop . 22 return ∂Lc∂bc 23 Function ClientUpdate ( c , ∂Lc∂hc ) : 24 Backprop . & ( wHc , wTc ) ← ( wHc − η ∂Lc∂wHc , wTc − η ∂Lc ∂wTc ) 25 return ( wHc , wTc )

Summary: This paper combines: split learning, federated learning, multi-task learning, vision transformers, and COVID-19 chest x-ray (CXR) classification. Specifically, the authors propose the Federated Split Task-Agnostic (FESTA) framework, motivated by privacy concerns of training models with patient level healthcare data (CXRs). This framework is for training large deep learning models based on private datasets distributed over clients (i.e. hospitals). Here, each client has a local copy of a "head" and "tail" of a vision transformer model, while the server has a shared body of a vision transformer. The training process is as follows: 1.) clients perform forward passes using their local data through their head networks and send the embedded representations to the server 2.) the server performs a forward pass with the embedded representations and sends the resulting representations back to their respective clients 3.) the clients finish the forward pass through their tail networks, compute the loss, and send gradients back through the server 4.) server continues backprop, client finishes backprop, everyone updates their weights (the server's network weights will be updated considering data from _all_ clients) 5.) occasionally, all the weights from the clients head and tail networks will be averaged together (FedAvg)  Additionally, the authors consider the case where some clients are training on different tasks (e.g. image classification, segmentation, detection) using the same modality (CXRs). Here, step 5 is applied over clients with the same task.