Venue: ICLR 2022 Conference

Reviews:

final_score: rating: 5: marginally below the acceptance threshold
review_content: 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.  

The proposed is similar to the method "Splitfed" (Thapa et al., 2020) but is extended for different tasks (as described above).

Experimental results demonstrate: 
(i) successful training of the neural network models with the proposed method.
(ii) better/comparable performance to prior works on distributed/privacy-preserving methods.
(iii) better performance using the Vit backbone as compared to CNN backbones, and also with the proposed multi-task vs. single-task setting.
main_review: Strengths:
(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.
(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.
(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).

Weaknesses:
(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.

(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.

(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.




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.

-------------------------------

Post authors' response:

Thank 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.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: 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.
Experimental 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.
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.

However, 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:
- 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."
- 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.
- 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?

If 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.

--------------------
Post-rebuttal:
I 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.

Minor comment:
- The "Federated Learning" paragraph on pg. 3 contains an unresolved reference due to a typo.



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").
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: 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.
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.
- 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.
- 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)?
- 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.
- Prior works  [*1,*2] have also addressed the task-agnostic problem in federated learning. What's the significant difference between those works?

[*1] Federated Learning with Unbiased Gradient Aggregation and Controllable Meta Updating
[*2] Task-Agnostic Privacy-Preserving Representation Learning via Federated Learning
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. 

After reading the responses from the authors, the authors partially solved my concerns. Thus, I decided to increase the rating.
review_title: null

final_score: rating: 8: accept, good paper
review_content: 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.
main_review: # Strengths
* It gives a new and practical distributed learning framework for image restoration tasks. It is capable of handling multiple tasks with maintaining privacy.
* It gives state-of-the-art or competitive results on the evaluated restoration tasks.
* The paper is easy to follow.

# Weaknesses
* The emphasized privacy-preserving property of the given framework is not experimentally or theoretically validated.
* No explanation about why only choosing deblocking, denoising, deraining, and deblurring in clients. What about super-resolution and inpainting?
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.
review_title: null

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