FedPOD: the deployable units of training for federated learning

TL;DR

This work tackles efficient federated learning for brain tumor segmentation under heterogeneous, time-constrained FeTS conditions. It introduces FedPOD, a POD-inspired, Poisson-based node-selection and validation-entropy–driven aggregation approach that eliminates the need to retain the same participants across rounds and supports scale-out via Kubernetes-like POD units. FedPOD improves training efficiency and achieves Dice scores and convergence comparable to FedPIDAvg while mitigating stragglers and data-sparsity effects through dynamic participation and task-wise data provisioning. The proposed method demonstrates practical deployability for clinically relevant FL by balancing communication, computation, and performance in non-IID, skewed data settings.

Abstract

This paper proposes FedPOD, which ranked first in the 2024 Federated Tumor Segmentation (FeTS) Challenge, for optimizing learning efficiency and communication cost in federated learning among multiple clients. Inspired by FedPIDAvg, we define a round-wise task for FedPOD to enhance training efficiency. FedPIDAvg achieved performance improvement by incorporating the training loss reduction for prediction entropy as weights using differential terms. Furthermore, by modeling data distribution with a Poisson distribution and using a PID controller, it reduced communication costs even in skewed data distribution. However, excluding participants classified as outliers based on the Poisson distribution can limit data utilization. Additionally, PID controller requires the same participants to be maintained throughout the federated learning process as it uses previous rounds' learning information in the current round. In our approach, FedPOD addresses these issues by including participants excluded as outliers, eliminating dependency on previous rounds' learning information, and applying a method for calculating validation loss at each round. In this challenge, FedPOD presents comparable performance to FedPIDAvg in metrics of Dice score, 0.78, 0.71 and 0.72 for WT, ET and TC in average, and projected convergence score, 0.74 in average. Furthermore, the concept of FedPOD draws inspiration from Kubernetes' smallest computing unit, POD, designed to be compatible with Kubernetes auto-scaling. Extending round-wise tasks of FedPOD to POD units allows flexible design by applying scale-out similar to Kubernetes' auto-scaling. This work demonstrated the potentials of FedPOD to enhance federated learning by improving efficiency, flexibility, and performance in metrics.

Figures (11)

• Figure 1: Federated learning workflow of the FeTS challenge. Simulates N nodes performing computations in parallel and distributed manner.
• Figure 2: The standard U-Net architecture applied identically to all participants in the FeTS challenge.
• Figure 3: Two data distributions provided in FeTS. Partition 1 (left) is divided by actual institution IDs. Partition 2 (right) is divided based on actual institutions and data characteristics (tumor size). The above figure is plotted on a $log_2$ scale.
• Figure 4: In cases of (a) partition1 and (b) partition2: In each task, all nodes in $S^{primary}$ participate and supply a subset of data with a size around the Poisson mean. Nodes in $S^{secondary}$ participate randomly in tasks. Since the data size in $S^{secondary}$ is smaller than the Poisson mean, they supply all their data when participating.
• Figure 5: Diagram of PID controller