Privacy-preserving Task-Agnostic Vision Transformer for Image Processing
Authors: Boah Kim
Jeongsol Kim
Jong Chul Ye
Venue: ICLR 2022 Submitted (2022)
Keywords: Federated learning
Split learning
Transformer
Image processing

Abstract: Distributed collaborative learning approaches such as federated and split learning have attracted significant attention lately due to their ability to train neural networks using data from multiple sources without sharing data. However, they are not usually suitable in applications where each client carries out different tasks with its own data. Inspired by the recent success of Vision Transformer (ViT), here we present a new distributed learning framework for image processing applications, allowing clients to learn multiple tasks with their private data. The key idea arises from a novel task-agnostic Vision Transformer that is introduced to learn the global attention independent of specific tasks. Specifically, by connecting task-specific heads and tails at client sides to a task-agnostic Transformer body at a server side, each client learns a translation from its own task to a common representation, while the Transformer body learns global attention between the features embedded in the common representation. To enable decomposition between the task-specific and common representation, we propose an alternating training strategy in which task-specific learning for the heads and tails is run on the clients by fixing the Transformer, which alternates with task-agnostic learning for the Transformer on the server by freezing the heads and tails. Experimental results on multi-task learning for various image processing show that our method synergistically improves the performance of the task-specific network of each client while maintaining privacy.

Under review as a conference paper at ICLR 2022
PRIVACY-PRESERVING TASK-AGNOSTIC VISION TRANSFORMER FOR IMAGE PROCESSING
Anonymous authors Paper under double-blind review
ABSTRACT
Distributed collaborative learning approaches such as federated and split learning have attracted significant attention lately due to their ability to train neural net- works using data from multiple sources without sharing data. However, they are not usually suitable in applications where each client carries out different tasks with its own data. Recently. Vision Transformer (ViT) has been widely explored in computer vision applications due to its capability learn the common repre- sentation through global attention of the embedded input sequence. By leveraging the advantages ViT, here we present a new distributed learning framework for image processing tasks, allowing clients to learn multiple tasks with their private data. The key idea arises from a disentangled representation of local and non-local features using task-agnostic Vision Transformer and task-specific head/tail. By connecting task- specific heads and tails at client sides to a agnossicc Trans- former body at server side, each client learns translation from its own task to a common representation, while the Transformer body learns global attention between the features embedded in the representation. To enable decomposition between the task-specific and common representation, we propose an alternating training strategy in which task-specific learning for the heads and tails is run on the clients by fixing the Transformer, which alternates with task-agnostic learn- ing for the Transformer on the server by freezing the heads and tails. Once the Transformer body is fully trained with a sufficient number of tasks and clients, additional training of the Transformer body is no longer required when new client is added with new task, and all that is required is the training of customer-specific head and tail. Experimental results on multi-task learning for various low-level and high-level computer vision including medical image data show that our method synergistically improves the performance of the specispciiiic network of each client while maintaining privacy.
INTRODUCTION
Deep learning approaches have demonstrated the state-of-the-au performance and fast inference time in computer vision tasks (Ronneberger et al.) 20155) Zhang et al,, (2177a) 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. 20176)
Chang al. 2020), deraining (Wei al. 2019 Ren et al. 2019). deblurring (Nah al. 2017 Kupyn et al. 2019), deblocking (Li et al. 2020b]] Maleki et al. 2018), etc. However, the performance of CNN typically depends on large number of training data (Chervenak et al., 20000) 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 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 central entity or to each device, have been investigated (Konečný al.) 2016 McMahan et al. 2017aa Gupta & Raskar 2018).
For example, federated learning (FL) (McMahan et al. 201 Li et al. 2020cc 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 al. 2018) is developed as enhanced pewacy-preserving model that

---

Under review as a conference paper at ICLR 2022
splits network into clients and server 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 difficult to determine the proper layer of the network 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 (Zaagg 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. (0206). 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 clients. In contrast to the existing SI. framework where the network split is arbitrary. TAViT provides systematic way of splitting neural networks between clients and servers for 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.
server to Iearn common representation across all tasks of
One may think that the proposed method similar to the image processing transformer (IPT) (Chen iet al.] (020)), 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, 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 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.
training step is composed of task-specific
Recall that one of the most unique advantages of Transformer body is to convert "unattended imputt features into "attended output features by learning global attention and non-local interactions between the imput 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.
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 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
teassfoommrr-bbsse approach
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 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:
2

---

Under review as a conference paper at ICLR 2022
(a)
Body subscription
(b)
Task speciffi learning
Task agghostii learning
(c)
Transformer body
Clients
Server
Clients
Server
Server
xD
Head
Head
Transformer body
Positional arribcciing
C1
C1
Image
Teek
Feddig
Transformer body
Fed
Transformer body
Sound
Hoad
Cycle
Head
Client
And Norm
Text
Add Norm
Head
Heed
VXExO
Task
Tall
Overall framework of TAVIT. (a) By subscribing a task-agnostic Transformer body, (b) clients train their heads and tails with the fixed body in parallel by task-specific learning. while server trains the body with the fixed head and tail of randomly chosen client for each iteration by task-agnostic learning, (c) and that the body consists of encoder-only Transformer architecture.
Figure I:
We propose novel distributed learning framework, TAVIT, that carries out multiple image process- ing tasks using distributed data.
The proposed method consists of task-specific heads and tails on clients and a task agno-sgi Transformer body 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 experi- mental results on multiple tasks.
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.) (017a). For the efficient implementation of FL., practical challenges unstable networks, hardware capacity
difference, and statistical heterogeneity of data distributions (Li et al. 2020c Smith et al. 2017: Li
ct al. 2018)) have been actively studied. CCrinzia et al. (2019) performs FL with multiple classification server and small CNNs to clients and trains them by knowledge distillation. Yao et al. (209)) presents an unbiased gradient aggregation for FL and meta updating of the model. In contrast, our method 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 simultancously for synergistic performance improvement.
tasks, and He al. (2020) loads huge network
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-par part networks sequentially (Gupta Raskar 2018). By extending this idea, Vepakomma et al. (018)) presents several ways to use SL, and Abuadbbb al. (2020) applies SL to ID 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.
particular, Thapa et al. (2020) proposes a combination of FL and SI. 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 single task. In contrast, 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-basee image processing methods have been extensively explored (He et al. al.] 20211; et al. 202T). In particular, Dosovitskiy et al. (2020) proposes Vision Transformer (VIT) with 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

---

Under review as a conference paper at ICLR 2022
(-th cycle
(-th iteration
t. iteration
Fodlug
Update
CR
Update
P2
Update
Update
Random sample
sample
Task-specific learning
Task-agnostic leaming
Head forword
Dody forvard di
Tal forward
Feature maps
- Network parameters
leadl bewcknserrd
Budly tackwward
teckweerd
Gradients
Figure 2: Training scheme of the proposed TAVIT, where denotes common server, and denotes clients. The FedAvg with dashed line is applied at the weight aggregation steps.
3 PRIVACY-PRESERVING TASK-AGNOSTIC VISION TRANSFORMER
3.1 SUBSCRIPTION BASED SERVICE MODEL
As illustrated in Figure aa). TAVIT is designed for subscription based services. Specifically, 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 subseribe to the same Transformer body. Accordingly, each client can train its 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 shown in Figure [b), 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 that 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 [c), we reshape the features f with Yx X x D size into sequence of patches f = {f1, f2,.... .... J.). where X, D denote width, height, and channel dimension of image features, respectively, S is the number of patches, i.e. S = Y./ /p² for patch size p, and fa denotes the s-th patch of the features with p² x D size. Then, these reshaped features 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 input for the tail CNN.
Here, for the Transformer body, we employ the encoder-only architecture as aano-tgi network, compared to IPT (Chen et al. 2020) that uses both encoder and decoder of Transformer. The encoder- only Transformer learns the global relationship between features in the input corpus, and that global attention may be all we need for better performance vision tasks as demonstrated 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 learn the task-specife 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.
3.2 TRAINING SCHEME
For distributed datasets of different tasks, apply the alternating training strategy between clients and the server by considering them as two players. Specifically, as shown in Figure TAVIT trains the client-side task task-specific head and tail networks and the server-side task-agnostic body network in an alternating manner. In the task-specific learning, clients train their own heads and tails with the fixed body weights in parallel using locally stored datasets. In contrast, in the task-agnostic learning,
4

---

Under review as a conference paper at ICLR 2022
Algorithm 1 TAVIT: C = [C,, C2. C2....,Crr Cx) is group of client sets with different tasks each other. Is and la denote the number of optimization iterations for each task-specific and task-agnostic step in one cycle. He and T. are the head and the tail of client e, and B is the Transformer body on the server.
Initialization H, all clients, B a body for i in [1, num_mmmmmpmmmmppp _cccesccccs]_]__________________ do for is in [1, I.]
// task-specific learning (heads & tails)
for each client Cj C in parallel do update He, Te with fixed B
end
i aggregation step then for each client subset Cs C.C. s.t. |Ck] > 1 do
// for case of multi-elients with one task
unify He and Te of client E Ck (e.g. FedAvg)
end
end end
for in in [1, [,] do
// task-agnostic learning (body)
k + randomly selected task update B with fixed He, Te. s.t. E C₁
end end
Output: H,T, B
the server trains the Transformer body with the fixed head and tail of randomly chosen client for each iteration. More details are as follows.
3.2.1 TASK- -PECCICC LEARNING
Let C = U C* be group of client sets participating different image processing tasks, where K denotes the number of tasks, and Cg has one or more clients with different datasets for the k-th task, i.e. C {er cy ........ N2 with N₁ 2 1. Each client E Cs has task-specific own network architecture for head H. and tail Te, which are connected to the Transformer body B in the server.
In the task- aascispecifii learning, for the given freezed Transformer B at the server and the local training data (x!.) ( J the e-th client then trains the neural networks He and T. by solving the following optimization problem:
min H..
where Lo(y, 9) refers to the e-th client specific loss between the target and the estimate ÿ. The parameters of He and Te are iteratively updated using Bl./OT, and 0l./OHt. These gradients are calculated by back-propagation through the entire model which can be expressed by the chain rule:
de.
dl.
ag
Ole Hhc
de.
ÖfB Ö5н
(2)
ØT.
OT. OH.
0Sh ОH.
0Sn
0ʃn OH.
where ==== --( (8) fg B(fa) and i = Tc(f)) This implies that to update the head H. and
the tail Te, the gradient 0 transmitted to the server after back- poop-propggttion through the tail, and also the 0f./0fn computed from back propppropagttion through the body is transported to each
client.
Federated learning In the task-specific learning, when there multiple clients for the same task (i.e. N₂ > 1), their heads and tails can be trained in parallel. Suppose that oi has training dataset size of [D,] and the total size of dataset in C is Σ [P
[D). In this case, the back- propagation and optimization process are the same with the single client case, but additionally applies al. 2017a) the parameters He and T. of E C₁ for every assigned period,
which is written as:
No
D
D
(He, &
Σ
H
Te,
where 1 I $ N
(3)
D
D
i=1
5

---

Under review as a conference paper at ICLR 2022
The period of the weight aggregation is adjustable (50 epochs in our experiments). From this federated learning, clients corresponding to the l--th task share the same parameters at the end of task- specific learning as shown in Figure 2
3.2.2 TASK-AGNOSTIC LEARNING
Once the heads and tails of multiple clients are trained, the Transformer body is trained by fixing the heads and tails at the clients. To train the Transformer body that learns the common representation in a task-agnostic manner, we construct a subset of the client CB by selecting one client from each task:
CB [Cres Gr
}
A "i
E CE-
Then, the training data fae ze
(1) J,N. corresponding to the task of the client are also selected, and the Transformer body in the server is updated by solving the optimization problem as follows:
min Σ>
lolub T.(B(H.(a?
CECB
Similar to the task-specific learning, the parameters of B are updated using dl./BB, where the client is randomly chosen from CB at each optimization step. The required gradients also come from back-propage nockpproaagation as following:
Oe
Ol.-
0fg 0B
acc
Oec 00
where
(6)
0B
0ʃB
0ʃB 0g
dʃB
where fs = B(fH) and 9 Tfff)) Here, the gradient 0l./ is only transported the server after proppaatnnnn through the tail. Through this task-agnostic learning, the Transformer body in the server learns global embedding representation and provides task-agnostic self-attended features for various image processing. The pseudocode of the overall TAVIT is described in Algorithm [ with more details in Appendix
3.3 COMMUNICATION COST AND PRIVACY PRESERVATION BY TAVIT
Given that gradients have to be transmitted two-way or one-way for training head/tail and body parts of the architecture, one many wonder whether additional communication overhead is significant. However, since the Transformer body shared model on the server that does not perform any weight aggregation, our model has much smaller cost for one communication between the client and the server in the task-agnostic learning. This comes from the small size of transported features and gradients for the heads and tails. If we sample clients in the network training, the communication cost can be further controlled. Therefore, up to a certain epoch size, our model is more communication bandwidth efficient compared to the classical FL, and the advantage increases if bigger Transformer body is used for better representation global attention. For detailed analysis, see Appendix D.4]
The proposed TAVIT is designed to distributed local data for distinct tasks without sharing the data to the other clients or any central device. Although the privacy attack on the transported features between the server and clients can be occurred, yet another powerful and unique mechanism for maintaining privacy in TAViT arises when the client-side network of the proposed method has a skip connection between the head and the tail. In this case, the transported feature characteristics can contain very lossy information from the original data, and one cannot reconstruct data only using the transmitted hidden features of the proposed method as detailed in Appendix D.1
EXPERIMENTAL RESULTS
We examine the performance of TAVIT with the following image processing tasks: deblocking (JPEG artifact removal), denoising, deraining, and deblurring. Additional experiments for image inpainting and medical data are also performed to investigate its performance for high-level computer vision tasks and different domain data, respectively, which can be found in Section D. Appendix. With single server, we set two clients to carry out FL on the deblocking task and set one client for each of the other tasks, the total number of clients is five in our experiments. We evaluate results using two metrics of PSNR and SSIM.
The public datasets we used are as follows. For deblocking and denoising, we use 10,582 images from PASCAL VOC 2012 (Everingham et al. 2010) and Segmentation Boundaries Dataset
Datasets

---

Under review as a conference paper at ICLR 2022
(SBD) (Hariharan et al. 2011 Particularly, for FL on the deblocking, we split the data into two sets with 5,291 images and distribute them to each client. Deblocking results are evaluated on Berkely Segmentation Database (BSD500) (Martin et al.] 2001b) that provides 200 test images. For the demoising, we apply random Gaussian noise with the level of = 30 to images. The Denoising model evaluated on CBSD68 that contains 68 test images extracted from the BSD500. For deraining, following the experiment setting of Jiang et al. (2020), we use data from Rain14000 (Fu et al. 2017b), Rain1800 (Yang et al. 2017). Rain800 (Zhang et al. 2019a), and Rain12 (Li et al., 2016) that provide 13,711 pairs of clean and synthetic rain images. Deraining results are evaluated on Rain and Rain (Yang et al., 2017), each of which has 100 synthetic rainy images. Deblurring is performed using a GoPro dataset (Nah et al. (017) that contains 2, 2,103 and 1,111 pairs of sharp and blurry images for training and test sets, respectively.
Implementation details To implement our TAViT, we use the encoder and decoder of DDPM (Ho) al., 2020) with three stages our backbone of each head and tail at the client. For the Transformer body in the server, we use 8 encoder layers of the vanilla Transformer (Vaswani et al. 2017) with 256 words and 512 embedding dimensions. The total number of parameters for networks each client and the server is about 28M and 17M, respectively. Using 4 Nvidia Quadro RTX 6000 cards and 2 Nvidia Geforce GTX 1080Ti cards, we train the networks using Adam optimizer with learning rate of 3 x 10-5, We initialize parameters of the networks with those of the pre-trained model by an autoencoder scheme. For the data augmentation, we apply random horizontal and vertical flipping, rotating with 90 degrees, and cropping by patch size of 64 x 64 x 3. For three cycles, by setting the batch size as 20, we perform the task-specific learning for 200, 400, 400, and 2000 epochs on deblocking, denoising, deraining, and deblurring, respectively. Also, we perform the task-agnostic learning for 1000 epochs with 1/4 data for each task. We implement our TAViT using PyTorch library under stD-style license using Flower federated learning protocol (Beutel et al. 2020) under Apache 2.0 License. The details datasets and implementation are described in Appendix
4.1 RESULTS
Convergence of TAVIT for multi-task distributed learning We evaluated the results of the pro- posed TAViT of multi-task distributed learning with all participated clients and one common Trans- former body in the server. Figure shows the gradual progression of quality metrics through the alternating training scheme. The performance of all tasks from our method increased and outper- formed as the task-specific learning and the task-agnostic learning continued. This demonstrates the synergistic improvement of the task-specific heads/tails and the task-agnostic body: the heads and tails learn more accurate feature embedding for given tasks, and the common body learns the global attention general to multiple image processing tasks by looking at various datasets. Although there were some tasks in which the score of each step was slightly lower than that of the previous step by the interaction of different task datasets, the overall performance of TAVIT was improved as the cycle progressed. Detailed quantitative results for each cycle are described in Appendix C
Deblocking
Denolsing
Deraining
Debluning
POWR (68)
32.9
36.7
31.5
30.6
29.5
Tahs
27.4
30.4
27.5
26.5
9.02
00771
6.03
0.57
8.597
0.55
0.78
6.55
0.50
X: Teekkepeeiiic learning
Teelk-astoonicc Issening
Figure 3: Results of TAViT for multi-task distributed learning. Each column shows the PSNR and SSIM according to i-th cycle for the deblocking, denoising, deraining, and deblurring, respectively.
Comparison of TAVIT to other strategies We compared our TAViT with other distributed learning strategies: SL and FL. We conducted both SL and FL with the two clients assigned for the deblocking task. For SL. designed in /Veakkomaa et al. (2018), clients and the body on the server, and trained those split networks without the weight aggregation for the head and tail. For FI., we placed the entire model composed of the head, body, and tail on each
placed the head and tail networks on
7

---

Under review as a conference paper at ICLR 2022
client, and trained the network in parallel by carrying out the aggregation with FedAvg (McMahan jet al. (0177a) using a common server. Figure[sshowe these scenarios, where CI and C2 are clients for the deblocking. C3, C4, and C5 are clients for the denoising, deraining, deblurring, and S is the server. As reported in Table | the proposed method achieves better performance compared to the other strategies even though ours learns multiple tasks.
Trainn
50%
C1
(s)
50%
C1
(b)
Table 1: Comparison result of our TAVIT to other dis- tributed learning strategies. Q# denotes the quality of JPEG images for deblocking task.
50%
Q10
Q50
Method
Device
PSNR
SSIM
PSNR
SSIM
C3
Split learning (SI.)
C1,8
27.55
0.781
32,97
0.921
Split learning (SL)
C2.
27.51
0.781
32.88
0.921
Federated learning (FL)
C1.
27.46
0,780
32.79
0.919
Figure 4: Scenarios of TAVIT (pink arrow) and other strategies of (a-b) SL and (c) FL. where C# is clients and S is the server.
TAVIT (Ours)
CI. C2.
27.69
0.786
33,20
0.924
Comparison of TAVIT to learning each separate task To verify the capability of the task-agnostic Transformer body learning from multiple tasks, we compared TAViT with the models independently trained on each individual task. Under the setting of centralized data for each task, we implemented this study two manners: end-to-end learning (EI.) and single-task learning (STL.). For EL., we trained the whole network in one device through the end-to-end optimization scheme. For STL, distributed the decomposable head, body, and tail to a client and a server as the proposed method, and trained the networks by the alternating training strategy for one cycle. Table 2 the results on Benchmark datasets for each task. This shows that our TAVIT trained on multiple tasks simultaneously outperforms both EL and STL, which suggests that the task task-agnostic body of our framework does not degrade the model by task heterogeneity but enhances the performance for various tasks.
Comparison of TAVIT to CNN- models To compare the performance of TAVIT with CNN- based deep learning models, we tested several existing methods on benchmark datasets for each task. Table 2 and Figure 5 show the qualitative and visual comparison results, respectively. For the deblocking, when comparing with DnCNN [Zhang et al. (2017a), AR-CNN
al. (2015). and QCN (Li et al., 2020b). the proposed method outperforms for both the 10 and 50 levels of quantization quality. Visual comparisons also show that the proposed method removes block artifacts clearly compared to the others. For the denoising, we compared our method with CBM3D (Dabov al. 2007). DnCNN (Zhang et al. 70), FFDNet (Zhang et al. (0188)) IRCNN (Zhang et al. 2017b). DHDN (Park et al.) 2019 and SADNet (Chang et al.) 2020). The results show that our
Table 2: Comparison results on Benchmark datasets. For Transformer-basee methods, EL is the end-to-end learning, STL is the single-task learning, and ours is the multi-task learning using TAViT.
Task
Method
Datasel
Metric
CNNed
Transformer-
Dehlocking
Input
DaCNN
AR-CNN
QCN
EI.
STI.
Ours
BSDS500 PSNR
25.67
26.70
26.42
27.66
27.67
27.59
27.69
(Q10)
SSIM
0.755
0.811
BSDS500 PSNR
31.51
32.70
33.00
33.01
32.93
33.20
(Q50) SSIM
0.902
0.918
0.934
0.923
0.924
0.924
Denoising
Input
CBM3D
DECNN
FFDNet
IRCNN
DHDN
SADNet
EL
STL
Our
CBSD68
PSNR
19.03
29.71
30.31
30.22
30.64
30.43
30.65
30.69
(aw30)
SSIM
0.336
0.843
0.861
0.860
0.861
0.864
0.864
0.870
0.871
Deraining
Input
DerainNet
SEMI
UMRL
PreNet
MSPFNN
EL
STL
Ours
Rain10011
PSNR
12.13
14.92
16.56
26.01
26.77
28.66
28.88
28.95
29.35
0.349
0.592
0.486
0.858
0.860
0.863
Rain100L
PSNR
25.52
27.03
25:03
29.18
32,44
32.40
32.93
32,50
34.30
SSIM
0.825
0.884
0.842
0.923
0.950
0.933
0.937
0.935
0.949
Deblurring
Input
DeblurGAN
Nah et al. 2017
Zhang al. 20188a
DeblurGANv2
EL
STL
Ours
GoPro
PSNR
25.64
28. 28.70
29.00
29. 19
29.55
28.62
29.28
30.06
SSIM
0.790
0.858
0.914
0.931
0.934
0.864
0.877
0.894
8

---

Under review as a conference paper at ICLR 2022
(a)
Reference
Input
DrCNN
AR-CNN
QUN
TAVIT
Deblocking
36.8 IR
77.0Y
77.58 VB
37.79
Reference
Input
CBMBD
DOCNN
FFDNet
IRONN
DHD'
TAVIT
PAB dR
BBEBAR
31A2 IB
31.09
3E,99
31.3
Denoising
2013
31.30
(c)
Reference
Input
DerainNet
SEMI
UMRL
PreNet
TAVIT
Deraining
25.24 6tb
34 48
30 66 40
34.55 dB
41.85 4B
41.5-10
(d)
Reference
Input
DeblurGAN
Nah of
DehlurGAN22
TAYIT
Deblurring
36.19
31.23
31.43
Figure 5: Visual comparisons for multiple tasks: (a) deblocking, (b) denoising, (c) deraining, and (d) deblurring. Yellow values PSNR, and inset box is magnified view of a red rectangle.
TAViT achieves better PSNR/SSIM scores, and also provides more clearly denoised images while preserving structure and texture details than the comparisons. For the deraining, we tested our model addition to DerainNet (Fu et al. 2017a). SEMI (ee al. 2019). UMRL (Yasarla Patel 2019) PreNet (Ren et al., 22019), and MSPFN Uiang et al. 2020). We used Y channel in YCbCr color space followed by Diang et al. (2020) for the evaluation. As result, our model outperforms the comparative methods on both Rain 100H and Rain 1000. Also, the images restored by ours are closer to the references by removing rain streaks than the others. For the deblurring, we employed DeblurGAN (Kupyn et al. 2018), Nah et al. (2017), Zhang et al. (2018a). DeblurGANv2 (Kupyn et al. (2019) for comparisons. The results show that the proposed method achieves comparable performance to the existing approaches. Visual results show that our TAVIT restores blurry images with sharp edges, while the others still contain blurry artifacts or position shifting objects compared to the references.
CONCLUSION
In this work, we present a multi-task distributed learning framework called TAVIT. In TAViT, the task-specific head CNN and the tail CNN are distributed to clients that have their own data, which are connected to a common Transformer body placed in the server. With an alternating training scheme, the heads and tails on client sides are trained by task-specific learning, while the body trained by task-agnostic learning. We conduct experiments on four different image processing tasks, which shows the success of task-agnostic learning of the Transformer body and its synergistic improvement with the task- specific heads and tails. Through our model, the participating clients can design and train their own networks depending the task using local data in parallel. We expect that the proposed TAVIT can be efficiently used in the case that sharing data with other institutions is sensitive such medical fields.

---

Under review as a conference paper at ICLR 2022
Ethics statement As our work utilizes distributed learning models, similar to the existing FL and SL, our method may be vulnerable to privacy attacks against the server such as inversion attacks (Yin et al. 202T). Although the proposed framework is designed by encoding the feature maps and gradients under Flower protocol which makes it difficult for attackers to restore the original data, the hidden feature may leak the data some degree. Thus, privacy-related techniques such as differential privacy (McMahan et al. 2017b) and authenticated encryption of data (Rogaway) 2002) should be analyzed for practical applications.
Reproducibility statement The source code and our trained models to reproduce the proposed method are available at httpps//gil hub. com/ docode, refer to Appendix A Also, the data processing steps for datasets used the experiments are provided in Appendix
/AVVi] For the detailed pseu-
REFERENCES
Sharif Abuadbba, Kyuyeon Kim, Minki Kim, Chandra Thapa, Seyit A Camtepe, Yansong Gao, Hyoungshick Kim, and Surya Nepal. Can we use split learning on Id enn models for privacy preserving training? In Proceedings of the 15th ACM Asia Conference on Computer and Commu- nications Security, pp. 305-318, 2020.
Pablo Arbeláez, Michael Maire, Charless Fowlkes, and Jitendra Malik, Contour detection and hierarchical image segmentation. IEEE Trunsactions on Pultern Analysis and Machine Intelligence, 33(5):898-916 2011. doi: 10.1 109/TPAM12210.11..
Daniel J Beutel, Taner Topal, Akhil Mathur, Xinchi Qiu, Titouan Parcollet, and Nicholas D Lane. Flower: A friendly federated learning research framework. arXiv preprint arXlv:2007.14390 2020.
Meng Chang, Qi Li, Huajun Feng, and Zhihai Xu, Spatial-adaptive network for single image denoising. In European Conference on Computer Vision, pp. 171-187. Springer, 2020.
Hanting Chen, Yunhe Wang, Tianyu Guo, Chang Xu, Yiping Deng, Zhenhua Liu, Siwei Ma. Chun- jing Xu, Chao Xu, and Wen Gao. Pre-trained image processing transformer. arXiv preprint arXiv:2012.00360 2020.
Ann Chervenak, Ian Foster, Carl Kesselman, Charles Salisbury, and Steven Tuecke. The data grid: Towards an architecture for the distributed management and analysis of large scientific datasets. Journal of network and computer applications, 23(3):187-200, 2000.
Luca Corinzia, Ami Beuret, and Joachim M Buhmann. Variational federated multi-task learning. arXiv preprint arXiv: 0900696266660 2019.
Kostadin Dabov, Alessandro Foi, Vladimir Katkovnik, and Karen Egiazarian. Image denoising by sparse 3-d corssfoom--mmmnn collaborative filtering. IEEE Transactions on image processing, 16 10:228002095, 2007.
Jacob Devlin, Ming- Wing-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint 11Xiv:1810.04800 2018.
Chao Dong, Yubin Deng, Chen Change Loy, and Xiaoou Tang. Compression artifacts reduction by deep convolutional network. In Proceedings of the IEEE International Conference on Computer Vision, pp. 576-584, 2015.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11922 2020.
Mark Everingham, Lue Van Gool, Christopher KI Williams, John Winn, and Andrew Zisserman. The pascal visual object classes (voc) challenge. International jounnattonaJouunaa of computer vision, 88(2): 303-338, 2010.
10

---

Under review as a conference paper at ICLR 2022
Xueyang Fu, Jiabin Huang, Xinghao Ding, Yinghao Liao, and John Paisley. Clearing the skies: A deep network architecture for single-image rain removal. IEEE Transactions on Image Processing, 26(6):2944 2956, 2017a.
Xueyang Fu, Jiabin Huang, Delu Zeng, Yue Huang. Xinghao Ding, and John Paisley. Removing rain from single images via a deep detail network. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 3855-3863, 2017b.
Otkrist Gupta and Ramesh Raskar. Distributed learning of deep neural network over multiple agents.
Journal of Network and Computer Applications, 116:1-8, 2018.
Kai Han, An Xiao, Enhua Wu, Jianyuan Guo, Chunjing Xu, and Yunhe Wang. Transformer in transformer. arXiv preprint arXiv:2 /2, 2021.
Bharath Hariharan, Pablo Arbeláez, Lubomir Bourdev, Subhransu Maji, and Jitendra Malik. Semantic contours from inverse detectors. In 2011 International Conference on Computer Vision, pp. 991-998. IEBE, 2011.
Chaoyang He, Murali Annavaram, and Salman Avestimehr. Group knowledge transfer: Federated learning of large cnns at the edge. arXiv preprint 20XXi::2007.155133 2020.
Shuting He, Hao Luo, Pichao Wang, Fan Wang. Hao Li, and Wei Jiang. Transreid: Transformer-bar object re-identification. arXiv preprint 102.0437, 2021.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. arXiv preprint arXiv:2006. 11239, 2020.
Kui Jiang, Zhongyuan Wang, Peng Yi, Chen Chen, Baojin Huang, Yimin Luo, Jiayi Ma, and Junjun Jiang. Multi-scale progressive fusion network for single image deraining. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp.
2020.
Jakub Konečný, H Brendan McMahan, Daniel Ramage, and Peter Richtárik. Federated optimization: Distributed machine learning for on-device intelligence. arXiv preprint arXiv: 2016.
Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convolu- tional neural networks. Communications of the ACM, 60(6):84-90, 2017.
Orest Kupyn, Volodymyr Budzan, Mykola Mykhailych, Dmytro Mishkin, and Jifí Matas. Deblurgan: Blind motion deblurring using conditional adversarial networks. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. $183-8192, 2018.
Orest Kupyn, Tetiana Martyniuk, Junru Wu, and Zhangyang Wang. Deblurgan-v2: Deblurring (frders-of-magnituuee faster and better. In Proceedings the IEEE/CVF International Conference on Computer Vision, pp. 3878-8887, 2019.
Ang Li, Huanrui Yang, and Yiran Chen. aask-agn privacy-preserving represcntation learning via federated learning. In Federated Learning, pp. 51-65. Springer, 2020a.
Jianwei Li. Yongtao Wang, Haihua Xie, and Kai- Kuang Ma. Learning a single model with wide range of quality factors for jpeg image artifacts removal. IEEE Transactions on Image Processing, 00.8842-8854, 2020b.
Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. arXiv preprint arXiv: 1812.06777 2018.
Tian Li, Anit Kumar Sahu, Ameet Talwalkar, and Virginia Smith. Federated learning: Challenges, methods, and future directions. IEEE Signal Processing Magazine, 37(3):50-6) 60, 2020c.
Yu Li, Robby T Tan, Xiaojie Guo, Jiangbo Lu, and Michael S Brown. Rain streak removal using layer priors. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 2736-2744, 2016.
Danial Maleki, Soheila Nadalian, Mohammad Mahdi Derakhshani, and Mohammad Amin Sadeghi. Blockenn: deep network for artifact removal and image compression. In Proceedings of the IEFE Conference Computer Vision and Pattern Recognition Workshops, pp. 2555-2558 2018.
11

---

Under review as a conference paper at ICLR 2022
D. Martin, C. Fowlkes, D. Tal, and J. Malik. A database of human segmented natural images and its application to evaluating segmentation algorithms and measuring ecological statistics. In Proceedings Eighth IEEE International Conference Computer Vision. ICCV 2001, volume pp. 416-423 vol.2, 2001a. doi: 10.1
D. Martin, C. Fowlkes, D. Tal, and Malik. A database human segmented natural images and its application to evaluating segmentation algorithms and measuring ecological statistics. In Proc. 8th Int'l Conf. Computer Vision, volume 2, pp. 416-423, July 2001
Cynthia McCollough, Baiyu. Chen, III Holmes, David, Xinhu. Duan, Zhicong. Yu, Lifeng. Yu, Shuai. Leng, and Joel. Fletcher. Data from low dose ct image and projection data, 2020.
Brendan McMahan, Fider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera Arcas. Communication-effice learning of deep networks from decentralized data. Artificial Intelli- gence and Statistics, pp. 1273-1282. PMLR, 2017a.
H Brendan McMahan, Daniel Ramage, Kunal Talwar, and Li Zhang. Learning differentially private 16.0663,,, 2017b.
recurrent language models. arXiv preprint
Seungjun Nah, Tae Hyun Kim, and Kyoung Mu Lee. Deep multi-scale convolutional neural network for dynamic scene deblurring. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 3883-389 2017.
Bumjun Park. Songhyun Yu, and Jechang Jeong. Densely connected hierarchical network for image denoising. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops. pp. 0-0, 2019.
W Nicholson Price and Glenn Cohen. Privacy in the age of medical big data. Nature medicine, 25 (1):37-43, 2019.
Dongwei Ren, Wangmeng Zuo, Qinghua Hu, Pengfei Zhu, and Deyu Meng. Progressive image deraining networks: A better and simpler baseline. In Proceedings of the JEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 3937-3946, 2019.
Phillip Rogaway. Authenticated-encrypiinn with associated-data. In Proceedings of the 9th ACM Conference Computer and Communications Security, pp. 98-107, 2002.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In International Conference on Medical image computing and computer- assisted intervention, pp. 234-241. Springer. 2015.
Virginia Smith, Chao-Kai Chiang, Maziar Sanjabi, and Ameet Talwalkar. Federated multi-task learning. arXiv preprint arXiv:1705 10467, 2017.
Chandra Thapa, Mahawaga Arachchige Pathum Chamikara, and Seyit Camtepe. Splitfed: When federated learning meets split learning. arXiv preprint arXiv:2004. 202/2004..20000 /2088, 2020.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones. Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. arXiv preprint 3xiv:1706.0376: 2017.
Praneeth Vepakomma, Otkrist Gupta, Tristan Swedish, and Ramesh Raskar. Split learning for health: Distributed deep learning without sharing raw patient data. arXiv preprint 1XXi:1812.00564 2018.
Xiaolong Wang. Abhinav Shrivastava, and Abhinav Gupta. A-fast-renn: Hard positive generation via adversary for object detection. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 2606-2615, 2017.
Xiaolong Wang, Ross Girshick, Abhinav Gupta, and Kaiming He. Non-local neural networks. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp- 7794-7803, 2018.
Wei Wei, Deyu Meng, Qian Zhao, Zongben Xu, and Ying Wu. Semi-supervised transfer learning for image rain removal. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 3877-3886, 2019.
12

---

Under review as a conference paper at ICLR 2022
Wenhan Yang, Robby T Tan, Jiashi Feng, Jiaying Liu, Zongming Guo, and Shuicheng Yan. Deep joint rain detection and removal from a single image. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 1357-1366, 2017.
Xin Yao, Tianchi Huang, Rui-Xiao Zhang, Ruiyu Li, and Lifeng Sun. Federated learning with unbiased gradient aggregation and controllable meta updating. arXiv preprint arXiv: 191.:0v990002244 2019.
Rajeev Yasarla and Vishal M Patel. Uncertainty guided multi-scale residual learning-using a cycle spinning enn for single image de-raining. In Proceedings of the IEEE/CVF Conference Computer Vision and Puttern Recognition, pp. 8405-8414, 2019.
Jong Chul Ye, Yoseob Han, and Eunju Cha. Deep convolutional framelets: A general deep learning framework for inverse problems. SIAM Journal Imaging Sciences, 11(2):991- 10448, 2018.
Hongxu Yin, Arun Mallya, Arash Vahdat, Jose M Alvarez, Jan Kautz, and Pavlo Molchanov. See through gradients: Image batch recovery via gradinversion. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 16337-16346, 2021.
Jiahui Yu, Zhe Lin, Jimei Yang, Xiaohui Shen, Xin Lu, and Thomas Huang. Generative image inpainting with contextual attention. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 5505-5514, 2018.
He Zhang, Vishwanath Sindagi, and Vishal M Patel. Image de-raining using a conditional generative adversarial network. IEEE transactions circuits and systems for video technology, 30(1 1): 3943-3956, 2019a.
Jiawei Zhang, Jinshan Pan, Jimmy Ren, Yibing Song, Linchao Bao, Rynson WH Lau, and Ming- Hsuan Yang. Dynamic scene deblurring using spatially variant recurrent neural networks. In Proceedings of the IEFE Conference on Computer Vision and Pattern Recognition, pp. 2521-2529 2018a.
Kai Zhang, Wangmeng Zuo, Yunjin Chen, Deyu Meng, and Lei Zhang. Beyond a gaussian denoiser: Residual learning of deep cnn for image denoising. IEEE transactions on image processing, 26(7):
3142-3155, 2017a.
Kai Zhang, Wangmeng Zuo, Shuhang Gu, and Lei Zhang. Learning deep emn denoiser prior for image restoration. In Proceedings of the IEEE conference on computer vision and pattern recognition. pp. 3929-3938, 2017b.
Kai Zhang, Wangmeng Zuo, and Lei Zhang. Ffdnet: Toward a fast and flexible solution for enn-based image denoising. IEEE Transactions on Image Processing. 27(9)44088 4622, 2018b.
Songyang Zhang, Xuming He, and Shipeng Yan. LatentGNN: Learning efficient non-local relations for visual recognition. In International Conference on Machine Learning, pp. 7374-7383, PMLR,
2019b.
13

---

// head output Under review as a conference paper at ICLR 2022
A DETAILS OF TAVIT WITH PSEUDOCODE
As described in the main paper, the task-specific heads and tails in clients and the Transformer body in the server are trained in an alternating manner between the proposed task-specific learning and the task-agnostic learning. In the following, we describe each step more detail in terms of its implementation.
Pseudocode for the task-specific learning A shows the pseudocode for the task-specific learning. Given K image processing tasks, the task-specific learning updates the heads H and the tails T in each client with the fixed body B. The server first initializes global weights of the heads
Algorithm 2
Task- ssecc--cecffic learning TAVIT: denotes the number of iterations for the speciipe learning in one cycle. (He, Tc) denote the head and the tail in client e Cj. for the k-th task, and B denotes the Transformer body in a common server. (Hc,: To,, are global weights of the heads and the tails for the task k. (fH ff) output of the tail. & is the task-specific loss the client c. 1D,1 is the size of training data at ej. and IDI is the total size of training data at Ck. i.e. IDI=E, [Dyl-
output feature maps from the head and the body, and ÿ is the
/* run on the server (with fixed B) k/ Initialize Hc. Te,
Send Hc, Tc, to clients e E C. for ig in [1, [ do
for each client c Ck. where k (1,2..... .) in parallel do
fH + ClientPhasel ()
fB +- B(fn)
// body output
at. + ClientPhase2 (fn)
Jls
at + 8t afn
// backpropagation through body
Th Bfg 8fR
ClientUpdate 85H
8k_
end
i, is weight aggregation step for |C&l > 1 then
Get (He, Te,) from client cy. where j E (1.2....,
Hc. < Nx j=
Heg
// RedAvg for head
D
Tc, + Σ N₂ D
Tep
// FedAvg for tail
D
Send ,Ho, Tox) to all clients C
end
end
* run on client c x/ Function ClientPhasel
E. f- set current data and label
fH + H.(x)
return JH
* run on client e x/
Function ClientPhase2 (fg)
ÿ + T((n)
// tail output // computation tail gradients
la f- Loss(u, 9)
He. G6. dy
ST. St. 8f.
St. +
85B ST. 85B
a
return Gfg
run on client c +/
Function ClientUpdate ( a
8he
Oin
// computation of head gradients by optimizer e.g. Adam
0H.
3f8 dH.
update Ho, T. using DH.
86 and क7.
8h.
14

---

Under review as a conference paper at ICLR 2022
Algorithm 3 Task-agnostic learning of TAVIT: Lo denotes the number of optimization iterations for task-agnostic learning in one cycle. (He, To) denote the head and tail in a client in a client e C for the k-th task, and B denotes the Transformer body in common server. (fH, f8) are output feature maps from the head and the body, and ÿ is the output the tail. & is the task-specific loss of the client c.
/*
the server */
Initialize Cp {ch 2 (
c where oney Ck
for in [1, L,]do do
c -q E CB
// Random selection of client with task &
JH t C-ClienPPhasel ()
Jn - B(Jn)
// body output
DC. ChiientPasse2 (fB)
Ujn oc. ac. 8fg
// computation of body gradients
8fa at
update B using dB
by optimizer e.g. Adam
end
run on client (with fixed Her Te) +/ Function ClientPhasel ()
V, y.y* f- set current data and label
fH + H²(x)
// head output
return fH
run on client e (with fixed Her Te) +/ Function ClientPhase2 (fz)
① + Tc(fB)
// tail output
le t Loss(y,
8l. 0r.
// backpropagation through tail
oy dfn
a
return OT
and the tails and sends them to all clients in Ck, where Cg is a set of clients with different datasets for the k-th task. When each client E Cj takes local training data x and provides feature map fu by the head H. to the server (line 5 with C1 client Phasel ), the server-side Transformer body takes the feature map fH as an input embedding and estimates the self- features fB that is independent of specific tasks. Once the fa in the server is sent to the client, the client computes the task task-specific loss E. between the label y and the tail output ğ (line 24 in ClientP Paathh The gradient of the tail 0l./OT, is also computed in the client at C1 MeentPaase (line 25), which is used compute 0l./OfB that is transported to the server. Then, in the server, 0k.00 10/0fz fH calculated by backpropagation through the body, which is sent to the client so as to compute the head gradient H and finally update H. and T. (lines 28-30 of C1 iennttdtttte using a single optimizer.
Here, when there are multiple clients for the k-th task, i.e. |C&| > 1, we apply the federated learning for the heads and the tails those clients as described in lines 11-6 in Algorithm 2 The heads and the tails are trained in parallel, and their weights are aggregated by FedAvg (McMahan et al.) 2017a) on the server side for every weight aggregation period. Then, these updated global weights of the heads Hc. and the tails Tc, and are transmitted to all clients in C₁ that the clients train their own head and tail using the new global weights from the next step.
Pseudocode for the task- aggnssic learning In the task-agnostic learning, the Transformer body in the server is updated with the fixed heads and tails of clients. 3 the pseudocode for the task-agnostic learning of TAVIT. Given a subset of clients, CB. by selecting one client among C₁ for each task, the client E CB is randomly chosen for every iteration. Then, compared to the task-specific learning, the implementation of the task-agnostic learning is similarly conducted but does not need C1 iienttddatt process in Algorithm 2 In other words, after the gradient dl,/ 0l./O.u is computed on the client side at the server updates the Transformer body by computing the body gradients OBB (lines 7-8), which the final step of each iteration.
(line 14- 18) and transmitted to the server (line 6),
15

---

Under review as a conference paper at ICLR 2022
B DETAILS OF DATASETS AND IMPLEMENTATION
B.1 LICANSE/SOURCE FOR EACH DATASET
In experiments, use the public datasets for image deblocking, denoising, deraining, and deblurring tasks. Here, we describe the specific information of each data set such as license and source link.
PASCAL VOC 2012 The PASCAL VOC data set (Everingham et al., 2010) is publicly available, which includes images obtained from the "flickr" website under SuugMug or its third-part -arty licensors. The data are protected by the United States and international intellectual property laws. The data source is from the URL: http: /ttsttttttt////// robots........o..........
w.t..uk/paaa. pk/pasca //VO//
BSDS500 and CBSD68 The Berkeley Segmentation Data Set and Benchmarks 500 (BSDS500) (Arbelēzz et al.) 2011) data set is extended version of BSDS300 (Martin et al. 2001b). which is a public data set originally provided for image segmentation and boundary detection by Berkeley Computer Vision Group. This data set is widely used for measuring image restoration perfor- mance. The color BSD68 data set (CBSD68) is extracted from the BSDS500. The BSDS500 can be downloaded wt/ttps://www2 eecs berkeley http://.w....eeeeeeccccoooeoeeee vi
/g/ hww.gg/eesosssceeooo html
Synthetic rainy images The synthetic rainy data set for training is collected from Rain 1000 synthesized bu et al. (2017b), Rain authored by Yang et al. (2017), Rain800 created by Zhang al.) (2019a), and Rain made by [i al. (2016). We test our method on the synthetic rainy data sets of Rain T00H and Rain 100I., both of which are authored by Yang et al. (2017). All these data sets are publicly available and can be downloaded at the following links:
Rain 1400:
html
Rain/800:
removal html
Rain800:
Raiin1: httpp:://uuli..gieeetetet
RRain000L & Rain100H:
joint rain removal.
GoPro The GoPro dataset (Nah et al. 2017) provides training and test sets for deblurring, The data
available
html
B.2
DATA PROCESSING
All datasets we used in experiments provide natural images that have three RGB channels and pixel values with range of [0, 255]. Upon these datasets, performed the following data processing according to the image processing tasks.
For the image deblocking task, we quantized the images following JPEG compression. We first transformed RGB image into YUV color space using the following equations.
Y 0.257R + 0.504G + 0.098B + 16
(7)
U 10--0.1 -0.148R - 0.291G 0.439B + 128
(8)
V 0.439 0.368G 0.071 + 128
Then, split the image into 8x8 blocks without overlapping and apply Discrete Cosine Transform (DCT) to each block. According to the level of quantization quality, we divided each element of the DCT coefficients by proper predefined matrices. After that, we apply inverse DCT and aggregate all blocks into an image, and then, we transformed the image from YUV to RGB color space.
R 1.164(Y 16) 1.596(V 128)
(10)
= 1.164(Y - 16) 0.392(U 128) - 0.813(V - 128)
(11)
B 1.164(Y 16) 2.017(U 128)
(12)
16

---

Under review as a conference paper at ICLR 2022
For the denoising task, added Gaussian noise to the clean images. Specifically, we applied random Gaussian noise with the level of (H4, σ) = (0, 30) to images, and then clipped the values into [0, 255].
For the other tasks, the datasets named Rain# and GoPro provide synthetic rainy images and blurry images, respectively. Since we used these datasets for the deraining and deblurring tasks, we did not perform any preprocessing such as the synthesis of rain artifacts blurry effect.
After the above data processing for all tasks, we randomly cropped the images by a patch size of 64 64 3. Also, we applied data augmentation using random flipping and rotating with 90 degrees. Then, we normalized the images with the scale of pixel values from [0,255] to [-1, 1J. which are final inputs the model.
B.3 NETWORK ARCHITECTURES
For the task-specific head and tail for each task, we use the network architecture of DDPM (Ho et al. 2020) that is composed of residual blocks and attention modules. We set the number of two-times downsanplinguupsamping stages as 3. For each stage, the channel size is set as 128, 256, and 512, respectively. Accordingly, given an input image x E [264x44x3 the head provides a feature map fu E R 16x10x that passes through the body, and the tail generates an output of the same size as the input.
On the other hand, for the Transformer body, we use the encoder part of the vanilla Transformer (Vaswani et al., 2017). As described in the main paper, the Transformer body takes a sequence of patches f by reshaping the feature map In as an embedding of the words. In the experiments, the length of the input sequence is 256 by setting the patch size as I, and the sequence dimension 512. Then, once the input sequence is added to learnable positional encodings, the encoded features / pass through encoder layers (L = 8 in our experiments). Table 3 shows the structure of each encoder layer of the body.
Table 3: The Transformer body architecture and its parameters in our experiments. For each encoder layer 1, M HHHAA is the multi-head attention modules. LN is the layer-normalization, DropOut is the dropout layer, Linear is the fully tully-connected layer, and Lu is the ReLU activation function.
Encoder layer
Parameters
a' AHH'))) E Baxd
Notation
Value
Mcaning
u LN + DropOut(d')) e Raxd
The number of encoder layers
of ReLU (Linear Raxd,
256
The length of sequence
5/ DropOu((Lineard) e gaxd d
512
The sequence dimension
5 LN vb)) E Ruxd
dp.
1024
The hidden dimension
4sooe shows the model sizes of the task-specific head and tail, and the Transformer body we used. When comparing the number of parameters and the size of networks, we can observe that the client-side networks composed of the head and the tail is larger than the task-agnostic Transformer body. Considering the experimental results in the main paper, this model size suggests that the body estimates task-agnostic self-attended features that provide the synergy effect in the task-specific and task agnk-agnottic learning even if the body size is smaller than the sum of head and tail.
Model sizes Table 4
Table 4: Model sizes of the head, body, and tail in our experiment.
Network
Head
Body
Tail
# Parameters 22,341,891 103.54
16,822,272 64.98
5,610.629 39.00
Size (MB)
17

---

Under review as a conference paper at ICLR 2022
C EXPERIMENTAL RESULTS
C.I TAVIT ON MULTIPLE IMAGE PROCESSING TASKS
Evaluation results of TAVIT Table 5 reports the quantitative evaluation results of TAVIT multiple image processing tasks, which visualized with graphs of scores for the cycles in the main paper. Figure shows shows the qualitative results of TAViT. This shows that the performance of each task improved according to the cycles between the task-specific and task-agnostic learning.
Table Quantitative results of TAVIT according to the cycles, which are visualized with graphs in the main paper. The best results are highlighted in bold.
Task
Deblocking
Denoising
Deraining
Deblurring
BSDS500 (Q10)
BSDS500 (Q50)
CBSD68
30)
Rain
Rain100L
GoPro
Cycle
PNSR
SSIM
PNSR
SSIM
PNSR
SSIM
PNSR
SSIM
PNSR
SSIM
PNSR
SSIM
0.5
27.53
0.781
32.92
0,921
30.57
0.868
28.24
0.860
33.17
0.939
1.0
27.57
0.782
33.01
0,922
30.62
0.869
28.75
0.862
32.69
0.936
29.09
0.873
1.5
27.61
0.784
33.05
0.923
30.57
0.870
28.57
0.869
33.58
0.945
29.63
0.885
2.0
27.65
0.785
33.14
0,924
30.66
0,870
28.79
0.867
33.50
0.944
29.72
0.887
2.5
27.64
0.785
33.14
0.924
30.62
0.870
29.25
0.875
34,30
0.949
29.
0.893
3.0
27.69
0.786
33.21
0.924
30.69
0.871
29.35
0.875
33.88
0.947
30.06
0.894
Deblocking
Denoising
Deraining
Deblurring
Cycle Input Reference
32.21 dls
SK
MLUS dP
32.29 db
50
LB
DELDS
22.45 dB
SIL-4
Figure 6: Qualitative results of TAVIT according to the cycles. From the left to the right columns, the deblocking results on images with the quantization quality 10 and 50, the denoising results, the deraining results on Rain 100H and Rain 100L, and the deblurring results. The yellow value is PSNR, and an inset box a magnified view red rectangle.
18

---

Under review as a conference paper at ICLR 2022
Qualitative comparisons Besides the results presented in the main paper, here, show more visual comparisons of our TAVIT to the existing methods. Figure 7 g and [O0 display the deblocking, denoising, deraining, and deblurring results, respectively. All these results verify that TAViT as a distributed learning for multiple image processing tasks outperforms the comparisons.
Reference
Input
DrCNN
AR-CNN
QCN
TAVIT
10.0
27.47
29.74
Figure 7: Visual comparisons of image deblocking task using BSDS500 (Arbeláez al., 201 Tp. The yellow value of each result is PSNR. We compare our TAVIT with DnČNN (Zhang et al. (017a). AR-CNN (Dong et al. 2015), and QCN al. 2020b
19

---

Under review as a conference paper at ICLR 2022
Reference
Input
CBM3D
DrCNN
FFDNet
IRCNN
DIIDN
TAVIT
"TV. dil
24 AIR
25 AII
SJ JM dIl
31 25 dI
'11 49 AER
33..1 dIl
28.58 BB
$5.46
29.55
25.42
28.41 JB
20:46 IE
4B
25.000 JB
SGB 4B
2509 4B
25 dB
JUAR AB
$9,99
19.80
27.22 dB
31.26 AIR
31.25 4B
31.13 4B
31.52 1B
31.77 CB
Figure 8: Visual comparisons of image denoising task using CBSD68 Martin al. 2001a). The yellow value of each result is PSNR. We compare the proposed TAVIT with CBM3D Dabov et al. 2007), DrCNN (Zhang al. 2017a), FFDNet (Zhang al. 2018b), IRCNN (Zhang et al. 2017b).
and DHDN (Park et al. 2019
20

---

Under review as a conference paper at ICLR 2022
Reference
Input
DerainNet
SEMI
UMRL
PreNet
TAVIT
(a)
38.11
$6.00
53 3M
76
23.43
24.25
31.32
32.78 8B
35.33 XB
(b)
32.36100
36.35 dll
39.32 AB
35.50
34.38
32.61
40.34 6F
41.9
Figure 9: Visual comparisons of image deraining task using Rain 100H and Rain 100L data sets (Yang
al. 2017). The yellow value of each result is PSNR. DerainNet (Fu et al.) (2017a). SEMI (Wei et al.
2019).
UMRL (Yasarla & Patel 2019). and PreNet (Ren et al. (011)) are used to compare our TAVIT.
(a) Results 000 data (b) Results on Rain 10LL data
21

---

Under review as a conference paper at ICLR 2022
Reference
Input
DeblurGAN
Nah el.
DeblurGANv2
TAVIT
23.41 411
28.20 11
31:06:18
N0.66
in an
21 997
4 77
4 77
< 97
872-3300
-444
87
-72-3300
372-3300-4
972-3300
$5u
F8-0000 UR 5511, 55
5011
10 50
6 9011. 56
750
58 752
20
0212
0212
0212
0212
0212
0212
23.91 2
23.37 AB
27.54 1B
25.34
28.60 .T
26,68
26.05 AB:
30.39 IBB
29,56
31.07 AB
22.7 J
22.4 AB
27.11 AB
5617 4B
77.07 4B
24.50 WIS
31.51
Figure 10: Visual comparisons of image deblurring task using GoPro (aah al. 2017) dataset. For visualization, we enlarge a region of red inset box (left column) for all images and denote PSNR of each result with yellow. We compare our TAVIT with DeblurGAN (Kupyn et al., 2018), Nah et al. (2017), and DeblurGANv2 (Kupyn et al. 2019).
22

---

Under review as a conference paper at ICLR 2022
C.2 ABLATION STUDY OF TAVIT
Study on the amount of data for each task in the task-agnostic learning In the main paper, we implemented our method using 1/4 the dataset for each task in the task-agnostic learning. To verify that this amount of data is enough for the task-agnostic learning, study using different amounts of data with a 1/2 ratio for each deblocking, denoising, deraining, and deblurring task. Table Figure 11 show the quantitative results of TAVIT on the multiple tasks using 1/2 data in the task-agnostic learning. Similar to the results with a 1/4 data ratio, the scores of PSNR and SSIM tend increase the cycle continues. When comparing the best results from the 1/4 and 1/2 data ratio, we can observe that performance for each task using even 1/4 amount of data comparable or better than using 1/2 data. This suggests that using 1/4 data for each task in the task-agnostic learning is sufficient to train the Transformer body and obtain high performance.
performed the ablation
Table 6: Quantitative results of TAVIT on multiple image processing tasks with 1/2 data for each task in the task-agnostic learning. The best results are highlighted in bold.
Task
Deblocking
Denoising
Deraining
Deblurring
BSDS500 (Q10)
BSDS500 (Q50)
CBSD68
30)
100H
Rain
GoPro
Cycle
PNSR
SSIM
PNSR
SSIM
PNSR
SSIM
PNSR
SSIM
PNSR
SSIM
PNSR
SSIM
0.5
27.53
0.781
32.92
0.921
30.57
0.868
28.24
0.860
33.17
0.939
28.94
0.871
1.0
27.57
0.782
33.02
0,922
30.62
0.869
28.78
0.862
32.47
0.933
29.18
0.875
1.5
27.63
0.784
33.10
0.923
30,58
0.870
28.92
0.867
33.28
0.942
29.63
0.887
2.0
27.65
0.785
33.15
0.924
30.67
0.870
29.18
0.870
33.71
0.946
29.83
0.890
2.5
27.66
0.785
33.15
0.923
30.56
0,870
29.35
0.875
33.83
0.945
29.94
0.893
3.0
27.69
0.786
33,20
0.924
30.69
0.870
29.53
0.877
33.68
0.945
30.18
0.897
1/4 data for each task
Best
27.69
0.786
33.21
0.924
30.69
0.871
29.35
0.875
34.30
0.949
30.06
0.894
Deblocking
Demoising
Deraliing
Debluriig
Rate
30.5
PSNW (00)
30.7
20
27.8
27.4
30.4
28.5
3.8
2.0
3.8
1.1
2.0
3.1
3.0
2.0
350
3.8
2.0
3.0
4.006
0.513
0.05
0.02
8.00
0.500
0.00
3.775
0.607
9:83
0.57
3.0
X: Tesksspecific learning
Tosk-agnooric learning
Figure 11: Results of TAVIT with 1/2 data for each task in the aask-agnostic learning. Each col- umn shows PSNR and SSIM according to i-th cycle for the deblocking, denoising, deraining, and deblurring, respectively.
Study the weight aggregation period In the main paper, we conducted the experiment of TAViT by applying FL to the deblocking task that has two clients with their own data. In FL,, the weights of the network in each client are averaged in the server for every weight aggregation period that is given as a hyperparameter. Here, since this period can influence learning performance in that the clients and the common server communicate to aggregate network weights, we performed the ablation study on the weight aggregation period for training the client-side networks. As reported in Table for the deblocking task, we trained the model with the aggregation period 20, 50, and 100 epochs. When we evaluated the deblocking results, the weight aggregation per 50 epochs provides better performance with 27. 27.53dB/0.77 781 and 32. 92..2dB/0.99 of PSNR/SSIM for the quality 10 and 50, respectively, compared to the other methods. This verifies that the weight aggregation period of 50 epochs presented in the main paper is proper to train and evaluate the proposed TAVIT in our experiments.
23

---

Under review as a conference paper at ICLR 2022
Table 7: Results of study the weight aggregation period of FL for image deblocking task. The best results are highlighted in bold. Q# denotes the quantization quality of JPEG images.
Q1O
Q50
Aggregation period
PSNR
SSIM
PSNR
SSIM
per 20 epochs
26.53
0.769
30.91
0.910
50 epochs
27.53
0.781
32.92
0,921
per 100 epochs
26.73
0.772
31.16
0.910
D DISCUSSION
D.1 IKIP-CONNECTION OF HEAD AND TAIL FOR PRIVACY PRESERVATION
When configuring the task-specific heads and tails with skip-connections. our model can avoid the privacy attack in some degree while maintaining the encoding information for the tail to generate outputs. This is because the skip-connected features isolated on each client and not transported to the server. Accordingly, the transported features hetween the clients and the server may contain far less information about the original data. Figure T2 shows examples of the outputs with and without skip-connections. This shows that the network output without skip-connections barely has the property of original data, which indicates that one may not be able to reconstruct the original data using the transmitted hidden features of the proposed method.
Input
w/ skipp-cconncttion
Deraining Denoising
Figure 12: The network output with and without skip-connections between task-specific heads and tail for denoising and deraining task.
D.2 EFFECT OF TASK-AGNOSTIC TRANSFORMER BODY
As we described in the main paper, the reason for developing our model with CNN-based heads/tails and the Transformer-basee body is take advantage of each network. In particular, Transformer learns the global attention of the input sequence through self-attention modules and has recently been extensively studied for various computer vision tasks. One of the most unique advantages of Transformer is to convert "unattended input feature vectors into "attended" " output feature vectors by learning global attention and non-local interactions between the input features. Accordingly, the task- sesk-specific head / tail can be only trained to learn task task-specific local features, whereas the global feature can be learned through the Transformer body. This disentangled representation of local and non-local features has been pursued throughout the development of deep networks. Thus, the proposed Transformer-ban baasd 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 split-learning architecture.
In order to show that this design is proper to the multi-task distributed learning, we additionally conducted the experiment by replacing the Transformer body to the CNN model. Specifically, we configured the CNN body with CBR blocks, where C is convolutional layer with the consistent channel 512, B is a batch normalization layer, and R is an activation by ReLU layer. For a fair comparison, set the CBR blocks to have almost the number of learnable parameters
24

---

Under review as a conference paper at ICLR 2022
with the Transformer body we used (16.522,240 of CNN body vs. 16,953,344 of Transformer body). Then, using this CNN body, we implemented our proposed task-specific and task-agnostic learning for one cycle on the multiple image processing tasks as the main paper. As result, Table shows that our model with the Transformer body achieves higher performance in both task-specific and task-agnostic learning. This indicates that the Transformer can be used as a general task task-agnostic body for multi-task learning.
Table 8: Results of study on the effect of the Transformer body of TAVIT versus CNN body.
Deblocking
Denoising
Deraining
Dehlurring
Body
Cycle
BSDS500 (Q10)
BSDS500 (Q50)
CBSD68 (o 30)
Rain 100H
Rain
GoPro
PNSR
SSIM
PNSR
SSIM
PNSR
SSIM
PNSR
SSIM
PNSR
SSIM
PNSR
SSIM
CNN
0,5
26.42
0.752
32.92
0.921
30.22
0.866
28.23
0.851
31.58
0.930
28.78
0.865
1.0
26.49
0.755
32.95
0.921
30.50
0.866
28.26
0.852
31.77
0,930
28.94
0.868
0.5
27.53
0.781
32.92
0.921
30.57
0.868
28.24
0.860
33.17
0,939
28.94
0.871
Ours
1,0
27.57
0.782
33.01
0.922
30.62
0.869
28.75
0.862
32.69
0,936
29.09
0,873
D.3 SAMPLING STRATEGY OF CLIENTS
When there are multiple clients for one task in task-specific learning, the task-specific networks of clients are aggregated through the sampling strategy of FedAvg. On the other hand, in task-agnostic learning of the proposed TAViT, one client is sampled for each iteration. Since the networks of clients for the same task are aggregated before the task- agnostic learning, we can readily sample one client for each task. Then, choosing one client for the subset of Eq. (4) can be viewed as sampling one task, which naturally reduces the communication cost.
In fact, the performance of TAViT not affected by the number of sampled clients in the aano-agnottic learning, since the task-agnostic body is updated for sufficient iterations. To demonstrate this, we performed the task-agnostic learning for the four tasks in our experiments by varying the sampling strategy. Table 9 the results after training our model for one cycle according to the number of sampled clients in the task-agnostic learning. The results show that sampling one client achieves comparable or higher performance for all tasks, compared to the results of sampling more than one clients. This supports that our sampling strategy is an efficient way to train the Transformer body with less computational cost even when the number of clients increases.
Table 9: Results of the study on the sampling strategy of clients in task- learning.
Task
Deblocking
Denoising
Deraining
Deblurring
BSDS500 (Q10)
BSDS500 (Q50)
CBSD68 (a = 30)
R0in
Rain 100L.
GoPro
Sumpling
PNSR
SSIM
PNSR
SSIM
PNSR
SSIM
PNSR
SSIM
PNSR
SSIM
PNSR
SSIM
27.57
0.782
33.01
0.922
30.62
0.869
28.75
0.862
32.69
0.936
29.09
0.873
27.57
0.782
33.01
0,922
30.62
0.869
28.68
0.861
32.59
0.935
29.10
0.873
4
27.57
0.782
33.01
0.922
30.62
0.869
28.75
0.862
32.35
0.933
29.11
0.873
D.4 COMMUNICATION COST BETWEEN CLIENTS AND SERVER
In the proposed TAViT, the features and gradients of the networks are transported between clients and server, so one may wonder how much additional communication cost occurs. To compute the communication cost of our method, we assume that the cost is proportional to the number of transported elements. Also, since the size of features and gradients from clients to the server are the same with those from the server to client, we only consider one direction from clients to the server. Then, we computed the maximum cost for communication update our model for each task-specific and aano-agn learning, and compared our cost to the method of FL. (McMahan et al. 2017a).
Specifically, when there N₂ clients for the k-th task, let PH. Pp and Pr be the number of parameters of the head, body, and tail, respectively. In the case of FL that aggregates the whole model composed of the head, body, and tail, the communication cost per communication can be represented as:
CostEL N&(PH + PB + Pr).
(13)
25

---

Under review as a conference paper at ICLR 2022
On the other hand, our model does not require the transportation of learnable parameters except for the aggregation step in the task-specific learning. Thus, the communication cost can be computed as follows:
Na + Pr), (F + G)), F FF
if an aggregation step else if a non-aggregation step (task-specific learning) (14) otherwise,
(task-specific learning)
AOST =
(task-agnostic learning)
where F and G are the number of elements of transported features and gradients, respectively.
see that the communication cost at the network aggregation step in the task-specific learning of the proposed method is smaller than FL that needs to transport the parameters of the whole model including head, body, and tail. Specifically, instead of aggregating parameters of the Transformer body. the TAVIT transports features and gradients that have much smaller size than the body parameters, which can reduce the cost per communication significantly. For example, the proposed model for the deblocking task contains PH + PB + Pr = 44, 774, 792 parameters whose memory size is about 207.5MB. Suppose that 10 clients participate in FL to train this model. Then, 447.7M elements are transported from the clients to the server, and the network of the server should handle more than 2GB load per communication. In contrast, our model transports PH + PT = 27, 952, 520 parameters whose memory size is approximately 142.5MB. Thus, even with 10 clients, 279.5M elements are transported, and the network of the server is supposed to handle about 60% load of FL. In addition, since the number of features and gradients is F = G 20 x 16 x 16 x 512- = 2, 621, 440 which is 10MB of memory, the number of transported elements per communication for 10 clients is 52.4M, and the server is pressed by only 200MB load per communication.
From (14).
On the other hand, in task-agnostic learning, the server updates the body with the sampled client without any weight aggregation. Accordingly, only the features and gradients are transported from the client the server. In particular, in the case of the communication from the server to the client, the server does not need to transport the gradients to the client, but only transmit the features. Thus, the cost per communication in the task agsk-agnostic learning phase is significantly reduced.
Therefore, up to a certain epoch size, our model is more communication bandwidth efficient compared to the classical FL., and the advantage increases more if a bigger Transformer body is used for better representation of global attention.
Scalability Suppose that there are K tasks, and let the total number of clients connected server be Nau- For simple analysis of the scalability, we assume that each communication between clients and the server takes a constant time. The scalability computed on the time complexity for one communication between clients and the server to update the models. For our task-specific learning, the one communication has time complexity (Nau) if we update the heads and tails of all clients. This means that the communication cost would increase according to the number of clients, which can limit the scalability of the proposed method. However, if we apply the client sampling strategy of the FedAvg, we can control the number of communications and the one communication will have time complexity Q(K) This sampling strategy can be readily adapted to our model without significant modification. On the other hand, for the task-agnostic learning phase, the one communication has time complexity O ), since the network parameters of clients for the same task are aggregated before the task-agnostic optimization. Also, according to the sampling strategy of one task in the proposed method, one communication has time complexity ((1), which is studied in Appendix D.3
D.5 APPLICATION TO HIGH-LEVEL VISION TASKS AND MEDICAL DATA
In the main paper, the proposed TAViT was demonstrated on multiple low-level computer vision tasks. However, the TAVIT framework can be also used for a wide range of high-level computer vision tasks, and even with different data domains such as medical images. To demonstrate this, we additionally conduct on inpainting task for natural images and denoising task for X-ray CT images. Here, the image inpainting is a higher-level computer vision task which requires more semantic information, and the denoising of X-ray CT requires domain-specifer knowledge about the data. In particular, to show that our agnostic Transformer body provides a positive effect on the training of new performed the task-specific learning only for the client-side heads and
task-specific networks, tails by subscribing to the pre-trained Transformer body, which was trained on the four natural image processing tasks without additional fine-tuning. The details of training and results
follows.
26

---

Under review as a conference paper at ICLR 2022
Dataset For the image inpainting task, used PASCAL VOC2012 dataset which contains 10,582
natural images. The information about license and source of this dataset can be found in Appendix scaled the image from [0, 255] to [-1. 1] and randomly cropped by 128 x 128 patches. Then we multiplied the image with the zero-box mask that has a random size of width and height from 48 64 according Yu al. (2018). For the X-ray CT denoising task, used the 2016 AAPM Low-dose CT Grand Challenge dataset (McCollough et al.) (2020) that provides noisy CT images with quarter dose and clean CT images with routine dose X-ray. Since the X-ray CT data are measured in the Hounsfield unit, we divided the intensity by 4,000 and randomly cropped by 64 64 patches.
B For the preprocessing,
Implementation details For the image inpainting task, we employed the network architecture of Yu et al. (2018) and decomposed it into two parts for the task-specific head and tail. We performed task-specific learning by minimizing the adversarial generative loss for 400 epochs using Adam optimizer with learning rate 1 x 10 For the X-ray CT denoising task, we used the same network architecture of head and tail implemented in this paper. We trained the task-specific networks with the fixed task-agnostic body for 400 epochs using Adam optimization algorithm with learning rate 5 x 10-3
Results To evaluate the performance of image inpainting and medical image denoising, we com- pared our method the CNN model that has the same network architectures of head and tail with ours but does not have the Transformer body. The quantitative evaluation results are shown in Table 10 and the visual comparisons are shown in Figure 13 We can see that the performance inpainting is improved when training the client-side networks with pre-trained Transformer body, even though we use the Transformer body pre-trained using low-level computer vision tasks. This implies that the proposed method can be extended to various high-level tasks. In addition, model on the medical image denoising task achieves higher performance than the comparative CNN model and provides clean images, although the Transformer body was trained the natural image domain. From these results, we can confirm that our task-agnostic Transformer body has capability learn the domain gap even in different data sources. Also, this suggests that clients do not need to train the server-side body from the scratch when they subscribe the body for the other tasks.
observe that
Table 10: Results of applications to diverse image processing tasks using different image domains.
CNN
TAViT
Task
Domain
PSNR
SSIM
PSNR
SSIM
Inpainting Denoising
Natural Image
24.50 41.28
0.861 0.958
24.85 42.05
0.865
Medical Image (CT)
0.961
27

---

Under review as a conference paper at ICLR 2022
Reference
Input
CNN
TAViT
(a)
48 0.863
07 0.881
20.63 0.900
21.61 / 0.905
0,840
21 0.849
(b)
40.13/0.944
40.70/ 0.949
40.06 0.956
40.49/ / 0.958
Figure 13: Visual comparisons of the application results: (a) inpainting of natural images, and (b)
denoising of medical CT images. The average values of PSNR SSIM are displayed on each result.
28


Reviews:
[{"review_id": "eBxBfQ3UzHC", "reviewer": "Reviewer_kCVH", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The submission is overall clearly written and presents an embodiment of a split architecture between client(s) and server that facilitates multi-task learning, which could be adopted by further work in the future. Code and models are available, which greatly eases adoption.\n\nHowever, while most of the pages are spent on architecture description and experimental results, there are several key omissions of discussion points which I deem essential for a paper in the distributed learning space:\n- There is no mention of any sort of privacy guarantees given the proposed architecture, despite \"privacy-preserving\" being part of the title. Privacy preservation is a strong claim that needs to be backed up rigorously. Moreover, in their ethics statement, the authors correctly lay out that transmitted hidden features \"may leak the raw data to some degree.\"\n- There is no discussion of communication cost in a federated setting. The number of clients used in experiments is exceedingly small, which might serve as a proof of concept. But given that gradients have to be transmitted two-way or one-way for training head/tail and body parts of the architecture, a discussion on communication cost and scalability is necessary.\n- Several key choices are not sufficiently motivated, including the choice of both CNN and transformer architectures, or the sampling of exactly one client from each task in Eq. 5. What makes CNN architectures less suitable for the body part in a more general framework beyond the particular architecture embodiment presented here? How would the presented sampling strategy scale in the face of several orders of magnitudes of more clients?\n\nIf the focus of the paper is to position TAviT as a general distributed multi-task learning framework, then the diversity of presented experiments for validation could have been expanded. On the other hand, if the focus is to present this particular architecture as a viable means to do distributed multi-task learning for the presented tasks, then the results of Table 2 remain unconvincing, in the sense that differences in results might come from essentially uncomparable architectures, as opposed to contributions to multi-task learning or distributed learning.\n\n--------------------\nPost-rebuttal:\nI thank the authors for their valuable comments on my review and their revision. The revision has improved the submission substantially. My comments were addressed mainly due to the addition of Section 3.3 and the there referenced Appendix D. I believe that most of the other reviewers' comments were also addressed and can now recommend the submission for acceptance. This is reflected by my adjusted score.\n\nMinor comment:\n- The \"Federated Learning\" paragraph on pg. 3 contains an unresolved reference due to a typo.\n\n\n", "summary_of_the_paper": "The paper presents an architecture for image processing tasks that splits up a network into three subsequent parts: head, body, and tail. Head and tail parts are CNN-based and can be trained on multiple client devices using federated learning (FedAvg), while the body part of the architecture is transformer-based and is trained on a central server. Head and tail parts are trained for specific tasks, while the body part is trained in a task-agnostic manner by selecting clients from each task for loss optimization.\nExperimental results show benchmark and convergence results that are comparable or favorable to non-distributed models, as well as comparison results to purely FL and SL approaches with a very small nr of clients.", "summary_of_the_review": "The submission presents a specific system that combines split and federated learning for multi-task learning of various image processing applications. It offers a good proof of concept of the proposed architecture decomposition, but lacks a robust discussion on communication cost/overhead as well as privacy guarantees. In particular, the ethics statement relativizes what the title claims (\"privacy-preserving\").", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "CtnOb-viAzp", "reviewer": "Reviewer_fi2d", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n(i) the key idea of the paper is communicated effectively. The method is an extension of \"SplitFed\", Thapa et al. 2020 to multi-task learning with a shared backbone on the server side.\n(ii) the experiments validate that the proposed method is viable, and achieves comparable performance to/marginal improvements over existing privacy-preserving training methods on the tasks considered in this work.\n(iii) some synergy can be seen in the performance on various tasks due to joint multi-task training, whereby the proposed method outperforms end-to-end and single-task distributed training (table 2).\n\nWeaknesses:\n(i) the key weakness is that the \"core contribution\" of using a transformer (Vit) based task-agnostic backbone on the server has incremental novelty over the more standard/prior works (specifically, Thapa et al.). The key novelty is joint-training of a shared backbone on *different* tasks in the proposed work --- this is a natural extension of the prior method.\n\n(ii) it is unclear why the specific decision of using CNNs as heads/tails (on the clients) and a ViT as backbone (on the server) is so crucial; could the roles be swapped --- ViT based heads/tails, and a CNN backbone? Overall this particular design choice has not been shown to be critical for operation or the performance of privacy-preserving distributed learning. Any differentiable neural network can be placed at the server/client sides as long as the two can be \"plumbed\" together.\n\n(iii)  the number of clients in the \"distributed\" experiments are small (perhaps this is prevalent in this research community). For example, only 5 clients, are used for 4 tasks (section 4). Hence, it is difficult to gauge the practical deployment of this method, as common distributed learning issues like clients dropping out, asynchronous communication is not tested.\n\n\n\n", "summary_of_the_paper": "This work is aimed at distributed \"privacy preserving\" training of neural networks for image processing tasks like deblocking, denoising, deraining, and deblurring. \"One of the most important contribution\" [pg 2] is breaking down the neural network model into task-specific convolutional head and tails (trained on \"clients\"), and a common shared (across tasks) Transformer based feature backbone, which is trained on the server. The heads/tails and the transformer backbone are trained in an alternate manner by assuming the other model to be fixed.  \n\nThe proposed is similar to the method \"Splitfed\" (Thapa et al., 2020) but is extended for different tasks (as described above).\n\nExperimental results demonstrate: \n(i) successful training of the neural network models with the proposed method.\n(ii) better/comparable performance to prior works on distributed/privacy-preserving methods.\n(iii) better performance using the Vit backbone as compared to CNN backbones, and also with the proposed multi-task vs. single-task setting.", "summary_of_the_review": "This work proposes an extension of the method \"SplitFed\" of Thapa et al. 2020, where a common transformer based backbone is trained across *different* tasks; this is a natural extension of the prior work. This method has been shown to achieve similar/marginally better performance across all the five tasks considered in this work. However, the scale of experimentation is small (only 5 clients), and idea itself carries incremental novelty. In view of the above, I can recommend this paper further, but perhaps with some reservation.\n\n-------------------------------\n\nPost authors' response:\n\nThank the authors for taking the time to address the concerns raised in the review. However, as detailed in individual comments, their response is not convincing. The key points are: (1) no fundamental reason to favor Transformers over other neural modules, (2) no guarantees of privacy preservation, despite claims to this effect, (3) no large-scale distributed study to validate their design. Hence, I urge the authors to address these, and revise the paper. I am unable to recommend this work further in the current form. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "4W6zOcK04GE", "reviewer": "Reviewer_QjZj", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "- In line-5 of the abstract, what did the authors inspire from Vision Transformers? Why the core motivation of this paper comes from \"the success of ViT\". The authors should revise the statement.\n- The proposed method has been verified on four low-level tasks. Can this strategy/framework be applied to higher-level tasks, such as image inpainting, classification, and object detection? Note that those tasks need more semantic understanding during training/inference.\n- In the traditional federated learning framework, each client commonly conducts the local train process, then transfers the model update consisting of the intermediate gradient to sever. However, in this paper, the author transfers the dataset's features to sever, which may bring about unpredictable challenges. For example, the computation cost of homomorphic encryption for features may increase rapidly. Does it have any advantages (research or application value)?\n- As for different tasks, the authors use a unified task-agnostic Transformer body, how the authors bridge the domain gap caused by the other tasks' knowledge. Furthermore, do different data sources (e.g., nature, satellite, and medical images) share the same Transformer body in Figure-1? Does it still work well? The authors should provide more experimental results and in-depth analysis to verify this point.\n- Prior works  [*1,*2] have also addressed the task-agnostic problem in federated learning. What's the significant difference between those works?\n\n[*1] Federated Learning with Unbiased Gradient Aggregation and Controllable Meta Updating\n[*2] Task-Agnostic Privacy-Preserving Representation Learning via Federated Learning", "summary_of_the_paper": "In this work, the authors present a multi-task distributed learning framework called TAViT. The task-specific head CNN and the tail CNN are distributed to clients with their data connected to a standard Transformer body placed in the server. With an alternating training scheme, the heads and tails on client sides are trained by task-specific learning, while the body is trained by task-agnostic learning. Experiments on four different image processing tasks show the success of task-agnostic learning of the Transformer body and its synergistic improvement with the task-specific heads and tails.", "summary_of_the_review": "This paper is well-written and presented. However, some key experiments and designs confuse me a lot. Overview, it approaches the borderline of the ICLR community. The authors should address the above concerns (# Main Review) in the rebuttal period. \n\nAfter reading the responses from the authors, the authors partially solved my concerns. Thus, I decided to increase the rating.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "-5SuxV4JhMC", "reviewer": "Reviewer_iupq", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "# Strengths\n* It gives a new and practical distributed learning framework for image restoration tasks. It is capable of handling multiple tasks with maintaining privacy.\n* It gives state-of-the-art or competitive results on the evaluated restoration tasks.\n* The paper is easy to follow.\n\n# Weaknesses\n* The emphasized privacy-preserving property of the given framework is not experimentally or theoretically validated.\n* No explanation about why only choosing deblocking, denoising, deraining, and deblurring in clients. What about super-resolution and inpainting?", "summary_of_the_paper": "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.", "summary_of_the_review": "I vote for accepting this paper. Its technical novelties and contributions are sufficient, and the given system seems practical and effective. It applies federal learning to image restoration tasks. It leverages ViT for universal representation learning, and task-specific heads and tails for training different tasks. The results are convincing. I think it is worth giving a brief study about how this framework reacts to privacy attacks.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper aims to devise a distributed multi-task privacy preserving framework for image processing. In this regard, author propose partitioning neural network models into task specific heads/tails and a common task-agnostic feature backbone (body). A training procedure is designed which is claimed to be privacy preserving wherein the head and tail is trained locally on the client or using federated learning when multiple clients share a task, while the main backbone/body is trained in a centralized manner by collecting appropriate gradients from the clients. Making easy to follow code is also highly appreciated. We thank the reviewers and authors for engaging in an active discussion and also updating the paper. While the new version is definitely resolves some of the concerns of the reviewers, some still remain. Privacy preserving in title and in main body of the paper seems misleading. Proposed method doesn't provide any guarantees for privacy (also pointed out by many reviewers). The author response doesn't seem to be convincing and other federated learning papers do not claim privacy unless having some specific mechanism like adding noise, secure aggregate, etc. Also, the reviewers are in consensus that novelty as well as large scale empirical evaluation is limited.
Decision: Reject