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

TL;DR

This paper presents FedOAP, a novel personalized federated learning method for tumor segmentation that leverages shared features and boundary-focused calibration to improve accuracy across diverse organs.

Contribution

The work introduces a decoupled cross-attention mechanism and a perturbed boundary loss to enhance federated segmentation performance on heterogeneous data.

Findings

FedOAP outperforms existing methods on multiple tumor segmentation tasks.

The decoupled cross-attention captures long-range inter-organ dependencies.

Boundary-focused calibration improves segmentation boundary accuracy.

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.