Privacy-preserving Task-Agnostic Vision Transformer for Image Processing

1 INTRODUCTION . Deep learning approaches have demonstrated the state-of-the-art performance and fast inference time in computer vision tasks ( Ronneberger et al. , 2015 ; Zhang et al. , 2017a ; Wang et al. , 2017 ) . In particular , convolutional neural networks ( CNN ) can learn the hierarchy of complex image features , so that a variety of CNN-based methods have been developed for denoising ( Zhang et al. , 2017b ; Chang et al. , 2020 ) , deraining ( Wei et al. , 2019 ; Ren et al. , 2019 ) , deblurring ( Nah et al. , 2017 ; Kupyn et al. , 2019 ) , deblocking ( Li et al. , 2020b ; Maleki et al. , 2018 ) , etc . However , the performance of CNN typically depends on a large number of training data ( Chervenak et al. , 2000 ; Krizhevsky et al. , 2017 ) , and it is often difficult to collect data from various entities due to privacy and regulation issues ( Price & Cohen , 2019 ) . Since the amount of data from a single source may not be enough , a deep learning framework that can leverage many datasets without violating privacy is required in real-world applications . To address this , distributed collaborative learning ( DCL ) approaches , which jointly train a single network on multiple systems or devices without revealing distributed data to a central entity or to each device , have been investigated ( Konečnỳ et al. , 2016 ; McMahan et al. , 2017a ; Gupta & Raskar , 2018 ) . For example , federated learning ( FL ) ( McMahan et al. , 2017a ; Li et al. , 2020c ) is studied to aggregate all data to the center under privacy constraints . Thanks to the parallel communication between each client , FL enables fast training of the network across multiple clients . Also , split learning ( SL ) ( Gupta & Raskar , 2018 ; Vepakomma et al. , 2018 ) is developed as an enhanced privacy-preserving model that splits a network into clients and server so that each client does not share all network parameters but only train a part of networks . From the advantages of FL and SL , a combination of split and federated learning , named SplitFed learning ( SFL ) ( Thapa et al. , 2020 ) , has been recently proposed to provide efficient training and a high level of privacy with a less computational burden . However , the existing CNN-based methods are difficult to determine the proper layer of the network to split . Also , although training data are distributed across each client , all clients usually consider a common learning task . Meanwhile , in many practical image processing applications , it is unlikely that all the clients are interested in the same applications . For example , some of the clients may be interested in image denoising ( Zhang et al. , 2017b ) , whereas the other clients are focused on image deblurring ( Nah et al. , 2017 ) , deraining ( Wei et al. , 2019 ) , deblocking ( Li et al. , 2020b ) , etc . As each task is different from the others , the existing distributed learning framework may not work . That said , these image processing tasks still require understanding of common image representation , so one may wonder whether there is any systematic way of synergistically learning multiple image processing tasks in a privacy-preserving manner . One of the most important contributions of this work is to show that Task-agnostic Vision Transformer ( TAViT ) , composed of the CNN-based head and tail and ViT-based body , is nicely fit to this purpose . Specifically , the head and tail are placed on each client to learn specific image processing tasks , while the body is stored and trained on a server to learn common representation across all tasks of clients . In contrast to the existing SL framework where the network split is arbitrary , TAViT provides a systematic way of splitting neural networks between clients and servers for privacy-preserving training without losing any performance . Furthermore , TAViT allows clients to use a common Transformer body model to learn multiple image processing tasks and synergistically improve the performance of their task-specific networks . One may think that the proposed method is similar to the image processing transformer ( IPT ) ( Chen et al. , 2020 ) , which consists of CNN-based heads and tails and a Transformer body . However , IPT requires centralized data and large computation resources for both pretraining and task-specific fine-tuning the whole model . Also , the Transformer in IPT has an encoder-decoder architecture which needs an explicit conditioning vector to convert the Transformer for a specific task . Thus , to our best knowledge , IPT is not suitable for distributed learning . In contrast , the body of TAViT is made of an encoder-only Transformer architecture to learn global embedding features of multiple tasks without any condition . Besides , by imposing computation of this Transformer body on the server rather than clients , our framework enables clients to reduce the computational burden while maintaining the overall performance for specific image processing tasks . In addition , TAViT views the heads and tail at the clients and the body at the server as two-part players and updates them alternately . Specifically , our training step is composed of task-specific learning and task-agnostic learning : the former is to train the client-side heads and tails to learn each task of the client , while the latter is to train the server-side Transformer body to learn general feature interpretation over multiple tasks . When there are more than two clients for any single task , parameters of their heads and tails can be aggregated through FL . Accordingly , TAViT offers seamless integration between SL and FL approaches to protect privacy . Recall that one of the most unique advantages of Transformer body is to convert “ unattended ” input features into “ attended ” output features by learning global attention and non-local interactions between the input features . Accordingly , with the help of aforementioned alternating training scheme , the task-specific head/tail can be only trained to learn task-specific local features , whereas the global feature can be learned through the Transformer . In fact , this disentangled representation of local and non-local features has been pursued throughout the development of deep networks ( Ye et al. , 2018 ; Zhang et al. , 2019b ; Wang et al. , 2018 ) . Thus , the proposed Transformer-based approach is considered to be one of the most advanced architectures for achieving this goal , as it improves synergistically overall performance , and at the same time leads the privacy-preserving split learning . We validate the performance of TAViT on multiple image processing tasks . Experimental results show that our multi-task distributed learning framework using the alternating training strategy outperforms the end-to-end learning of each individual task thanks to the decomposition of the task-agnostic Transformer body and task-specific networks . This suggests that our framework is a promising approach for learning multiple tasks with distributed privacy-sensitive data . In sum , our contributions are summarized as follows : • We propose a novel distributed learning framework , TAViT , that carries out multiple image processing tasks using distributed data . • The proposed method consists of task-specific heads and tails on clients and a task-agnostic Transformer body on a server , which reduces the computational cost of clients and does not require centralized data for multi-task learning . • An alternating training strategy between the task-specific and task-agnostic learning for the split networks shows the synergy effect of performance improvement , which is demonstrated by experimental results on multiple tasks . 2 RELATED WORKS . Federated learning In the FL setting , multiple clients learn locally stored data while one server aggregates information of clients by various methods including FedAvg ( McMahan et al. , 2017a ) . For the efficient implementation of FL , practical challenges of unstable networks , hardware capacity difference , and statistical heterogeneity of data distributions ( Li et al. , 2020c ; Smith et al. , 2017 ; Li et al. , 2018 ) have been actively studied . Corinzia et al . ( 2019 ) performs FL with multiple classification tasks , and He et al . ( 2020 ) loads a huge network to a server and small CNNs to clients and trains them by knowledge distillation . Yao et al . ( 2019 ) presents an unbiased gradient aggregation for FL and meta updating of the model . In contrast , our method is presented for effectively learning on task heterogeneity using distributed data . Although Li et al . ( 2020a ) presents task-agnostic FL method based on the feature extractor , each client trains the task-specific network independently , while our model can learn multiple tasks simultaneously for synergistic performance improvement . Split learning Split learning ( SL ) is designed to train networks over distributed data by splitting networks into two parts , which updates client-part and server-part networks sequentially ( Gupta & Raskar , 2018 ) . By extending this idea , Vepakomma et al . ( 2018 ) presents several ways to use SL , and Abuadbba et al . ( 2020 ) applies SL to 1D CNN models . However , the existing SL methods are designed using CNN , and to our best knowledge , there is no principle way of splitting the network for the best performance . In particular , Thapa et al . ( 2020 ) proposes a combination of FL and SL , but the server requires labels from clients to update the split networks , which may lose data privacy . Also , since outputs are generated from a shared network on the server when there are multiple clients , these methods can be narrowly applied to a single task . In contrast , our model presents Transformer-based shared body that enables multi-task learning of clients without sharing data . Vision Transformer for image processing Recently , inspired by the success of Transformer in natural language processing ( Vaswani et al. , 2017 ; Devlin et al. , 2018 ) , Transformer-based image processing methods have been extensively explored ( He et al. , 2021 ; Han et al. , 2021 ) . In particular , Dosovitskiy et al . ( 2020 ) proposes a Vision Transformer ( ViT ) with an encoder-only architecture to learn image recognition tasks . Also , Chen et al . ( 2020 ) presents an image processing Transformer ( IPT ) that learns low-level vision tasks by pretraining and task-specific fine-tuning . However , to the best of our knowledge , there are no existing works that exploit ViT architecture for distributed learning applications . 3 PRIVACY-PRESERVING TASK-AGNOSTIC VISION TRANSFORMER . 3.1 SUBSCRIPTION-BASED SERVICE MODEL . As illustrated in Figure 1 ( a ) , TAViT is designed for subscription-based services . Specifically , a client subscribes to a task-agnostic Transformer model at the server side that has learned global attention over the image features from other datasets . Then , the client can build the head and tail proper to its own image processing task , and connect them to the Transformer body at the server . At the subscription time , there may be already multiple clients that subscribe to the same Transformer body . Accordingly , each client can train its own head and tail using its local data whereas the common Transformer body is regularly updated using embedding features from all subscribers through alternating training strategy as shown in Figure 1 ( b ) , or even fixed if training has been performed with sufficient number of tasks and clients . As a result , the latest version of the Transformer body trained using more training data can be maintained on the server side so that it can be offered to new clients at the next subscription . Since the local data are not centralized to one device and are not shared with other clients , our framework can preserve data privacy . In the proposed framework , we consider the features from the head as a sequence of tokens similar to natural language processing . Specifically , as shown in Figure 1 ( c ) , we reshape the features f with Y ×X×D size into a sequence of patches f = { f1 , f2 , . . . , fS } , whereX , Y , D denote width , height , and channel dimension of image features , respectively , S is the number of patches , i.e . S = Y X/p2 for patch size p , and fs denotes the s-th patch of the features with p2 ×D size . Then , these reshaped features f are taken into the Transformer body as an input sequence , which is added to learnable positional embeddings to keep the position information of each feature patch . The Transformer body consists of several encoder layers proposed in Vaswani et al . ( 2017 ) so that the encoded features pass through several multi-head self-attention modules and feed-forward modules for each layer . And then , the body output of transformed features is reshaped into the original shape of features f to be used as input for the tail CNN . Here , for the Transformer body , we employ the encoder-only architecture as a task-agnostic network , compared to IPT ( Chen et al. , 2020 ) that uses both encoder and decoder of Transformer . The encoderonly Transformer learns the global relationship between features in the input corpus , and that global attention may be all we need for better performance in vision tasks as demonstrated in ViT . Therefore , the body of our framework can be trained to translate the input embedding features into globally self-attended features independent of specific tasks . Moreover , the heads are guided to learn the task-specific embedding from the input images to the common feature representation , and the tails are trained to learn the attended features for the specific image processing tasks . This architectural modification enables the framework to be suitable for multi-task distributed learning .

Summary: This paper presents a new distributed learning framework exploiting the vision transformer for various image processing applications. It gives impressive quantitative and qualitative results on multiple image restoration tasks meanwhile keeping privacy. Specifically, it employs a task-agnostic vision transformer to learn universal representation at the server, and several CNN-based task-specific heads and tails to handle different image restoration tasks at the client side. It also gives a training strategy to learn this model.