Personalized Federated Segmentation with Shared Feature Aggregation and Boundary-Focused Calibration

TL;DR

FedOAP addresses organ-agnostic tumor segmentation under privacy-preserving personalized federated learning by learning a shared feature space across organs while keeping client data local.It introduces Decoupled Cross-Attention to fuse globally aggregated features with local queries and a client-specific spatial adapter for personalization, plus a Perturbed Boundary Loss for boundary-focused calibration during local fine-tuning.Across brain, liver, and breast tasks, plus unseen lung data for generalization, FedOAP achieves superior Dice scores and robust cross-domain performance relative to state-of-the-art FL and PFL baselines.This approach enables effective cross-organ knowledge transfer with privacy guarantees and practical applicability in heterogeneous medical imaging scenarios.

Abstract

Personalized federated learning (PFL) possesses the unique capability of preserving data confidentiality among clients while tackling the data heterogeneity problem of non-independent and identically distributed (Non-IID) data. Its advantages have led to widespread adoption in domains such as medical image segmentation. However, the existing approaches mostly overlook the potential benefits of leveraging shared features across clients, where each client contains segmentation data of different organs. In this work, we introduce a novel personalized federated approach for organ agnostic tumor segmentation (FedOAP), that utilizes cross-attention to model long-range dependencies among the shared features of different clients and a boundary-aware loss to improve segmentation consistency. FedOAP employs a decoupled cross-attention (DCA), which enables each client to retain local queries while attending to globally shared key-value pairs aggregated from all clients, thereby capturing long-range inter-organ feature dependencies. Additionally, we introduce perturbed boundary loss (PBL) which focuses on the inconsistencies of the predicted mask's boundary for each client, forcing the model to localize the margins more precisely. We evaluate FedOAP on diverse tumor segmentation tasks spanning different organs. Extensive experiments demonstrate that FedOAP consistently outperforms existing state-of-the-art federated and personalized segmentation methods.

Figures (5)

• Figure 1: Motivation for FedOAP. FedOAP aims to learn a shared feature representation across clients by aggregating informative tokens from different organs while preserving client-specific data privacy. This enables improved tumor segmentation by leveraging inter-organ knowledge and producing personalized predictions guided by cross-client feature dependencies.
• Figure 2: (a) Overview of the proposed FedOAP framework. Each client $k \in \mathcal{K}$ trains a local model $\theta_k$ (composed of an encoder-decoder architecture with a client-specific spatial adapter $\phi_k$) on its private dataset $\mathcal{D}_k$ for $\mathcal{T}$ rounds. During this phase, DCA separates the shared parameters $\gamma_k$ for global aggregation. Finally, each client fine-tunes its model using PBL, resulting in the personalized parameters ${\theta}_k^{\ast}$. (b) The proposed Decoupled Cross-Attention (DCA) mechanism. The private query embeddings $\mathbf{q}_k$ attend to the concatenated key-value $(\mathbf{k}_k, \mathbf{v}_k)$ pairs, capturing cross-organ, long-range dependencies without exposing private client representations. (c) The Perturbed Boundary Loss (PBL) mechanism. The inconsistencies between predicted masks $\hat{\mathcal{Y}}_i$ and ground truth $\mathcal{Y}_i$ is used as a guide to inject noise, reinforcing supervision on uncertain regions to improve boundary precision.
• Figure 3: Qualitative comparison of FedOAP with existing FL and PFL methods across different organ datasets. Each row depicts a random sample from a client, illustrating segmentation performance across methods.
• Figure 4: Average inference time (in ms) per sample for each FL method.
• Figure 5: Comparison of FedOAP using $\mathcal{K}=30$ in various training scenarios. Here, $R$ denotes the number of communication rounds and $F$ denotes the number of local fine-tuning epochs.