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

2025-11-24

AI Summary

In medical image segmentation, client data exhibit non-independent and identically distributed (Non-IID) characteristics with heterogeneous organ categories, yet existing personalized federated learning methods neglect the collaborative potential of shared features across clients. Method: We propose an organ-agnostic personalized federated segmentation framework comprising three key components: (1) a decoupled cross-attention mechanism to model long-range inter-organ dependencies across clients; (2) a global key-value aggregation module for robust fusion of shared features; and (3) a perturbation boundary loss to enhance boundary robustness and generalization. Contribution/Results: Evaluated on multi-organ tumor segmentation, our method significantly outperforms state-of-the-art federated and personalized segmentation approaches. It demonstrates superior capability in effectively integrating shared knowledge with client-specific modeling under severe data heterogeneity, achieving improved accuracy, boundary fidelity, and cross-client generalizability.

Technology Category

• Computer Vision: Segmentation
• Machine Learning: Distributed Machine Learning & Federated Learning
• Search and Optimization: Distributed Search

Application Category

• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization
• Systems and Infrastructure for Web, Mobile and WoT: Federated Web and WoT systems, including distributed, federated and edge-based data processing
• Search and Retrieval-Augmented AI: Personalized, context-aware and across-device search

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.

Problem

• Addresses data heterogeneity in federated medical image segmentation
• Improves segmentation accuracy through boundary-focused calibration
• Leverages shared feature aggregation across different organ segmentation tasks

Innovation

• Utilizes cross-attention for inter-client feature dependencies
• Employs boundary-aware loss to improve segmentation precision
• Aggregates shared features while preserving local client data