Anatomical 3D Style Transfer Enabling Efficient Federated Learning with Extremely Low Communication Costs

TL;DR

A novel federated learning approach that utilizes 3D style transfer for the multi-organ segmentation task and can maintain its accuracy even in cases where the communication cost is highly limited, demonstrating that the method has high practicality in real-world scenarios where low communication costs and a simple pipeline are required.

Abstract

In this study, we propose a novel federated learning (FL) approach that utilizes 3D style transfer for the multi-organ segmentation task. The multi-organ dataset, obtained by integrating multiple datasets, has high scalability and can improve generalization performance as the data volume increases. However, the heterogeneity of data owing to different clients with diverse imaging conditions and target organs can lead to severe overfitting of local models. To align models that overfit to different local datasets, existing methods require frequent communication with the central server, resulting in higher communication costs and risk of privacy leakage. To achieve an efficient and safe FL, we propose an Anatomical 3D Frequency Domain Generalization (A3DFDG) method for FL. A3DFDG utilizes structural information of human organs and clusters the 3D styles based on the location of organs. By mixing styles based on these clusters, it preserves the anatomical information and leads models to learn intra-organ diversity, while aligning the optimization of each local model. Experiments indicate that our method can maintain its accuracy even in cases where the communication cost is highly limited (=1.25% of the original cost) while achieving a significant difference compared to baselines, with a higher global dice similarity coefficient score of 4.3%. Despite its simplicity and minimal computational overhead, these results demonstrate that our method has high practicality in real-world scenarios where low communication costs and a simple pipeline are required. The code used in this project will be publicly available.

Figures (5)

• Figure 1: (a) Our proposed setting with limitted communication rounds and shared style information. (b)(c) Model accuracy evaluation. The hatched bars indicate the accuracy when the number of communications is reduced from 400 to 5.
• Figure 2: Overview of our A3DFDG. (Stage1) First, we calculate the 3D visual style of each client in the frequency space and store them in clusters based on the predicted slice scores (slice position). (Stage2) During training, we retrieve 3D styles from the same cluster as the samples in the minibatch and perform style transfer without losing organ information.
• Figure 3: (a)(b) Distribution of the predicted slice scores (= predicted slice position at which images are captured) across six datasets. (c) The distribution of extracted style. Here, color indicates its slice score and PCA is implemented for visualization.
• Figure 4: Data augmentation in the frequency domain and subsequent post-processing
• Figure 5: Qualitative results in the out-of-federation setting.