Efficient Federated Tumor Segmentation via Normalized Tensor Aggregation and Client Pruning

Abstract

Federated learning, which trains a generic model for different institutions without sharing their data, is a new trend to avoid training with centralized data, which is often impossible due to privacy issues. The Federated Tumor Segmentation (FeTS) Challenge 2021 has two tasks for participants. Task 1 aims at effective weight aggregation methods given a pre-defined segmentation algorithm for clients training. While task 2 looks for robust segmentation algorithms evaluated on unseen data from remote independent institutions. In federated learning, heterogeneity in the local clients’ datasets and training speeds results in non-negligible variations between clients in each aggregation round. The naive weighted average aggregation of such models causes objective inconsistency. As for task 1, we devise a tensor normalization approach to solve the objective inconsistency. Furthermore, we propose a client pruning strategy to alleviate the negative impact on the convergence time caused by the uneven training time among local clients. Our method achieves a projected convergence score of 74.32% during the training phase. For task 2, we dynamically adapt model weights at test time by minimizing the entropy loss to address the domain shifting problem for unseen data evaluation. Our method finally achieves dice scores of 90.67%, 86.23%, and 78.90% for the whole tumor, tumor core, and enhancing tumor, respectively, on the task’s validation data. Overall, the proposed solution ranked first for task 2 and third for task 1 in the FeTS Challenge 2021.