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.
Federated Learning (FL) is a decentralized machine learning paradigm that enables collective model training across multiple data sources without the need to share raw data (Li et al., 2020a). By ensuring that data remains local to its source, FL inherently adheres to privacy regulations and data protection standards (Voigt and Von dem Bussche, 2017). This makes FL particularly well-suited for sensitive domains such as healthcare (Horst et al., 2025;Kaissis et al., 2020;Kwak and Bai, 2023). However, non-independent and identically distributed (Non-IID) data or heterogeneous data poses as a great hurdle for most FL methods (Zhao et al., 2018). A universally shared global model is usually insufficient to tackle the diverse types of data encountered across different clients. Heavily heterogeneous data may introduce client drift which severely degrades the model's performance (Karimireddy et al., 2020). Specially in the domain of medical image analysis, it is common to have inter-client incon-sistencies as individual clients often address distinct clinical tasks, patient populations, cultural contexts, and imaging equipment (Guan et al., 2024).
Personalized federated learning (PFL) tries to address this problem by maintaining distinct local models for each client that conforms to its unique distribution while keeping the data localized (Tan et al., 2022). This is achieved by a variety of methods such as regularization of local and global weight divergence (Karimireddy et al., 2020;Li et al., 2020b), meta-learning (Fallah, Mokhtari, and Ozdaglar, 2020;T Dinh, Tran, and Nguyen, 2020), transfer-learning (Afzali and Shamsinejadbabaki, 2025;Tan et al., 2023), decoupling parameters into shared and private segments (Arivazhagan et al., 2019), knowledge distillation (Li and Wang, 2019;Lin et al., 2020), clustering clients based on similar tasks or data distributions (Briggs, Fan, and Andras, 2020;Sattler, Müller, and Samek, 2020), among others. In medical image segmentation, a field often bound by strict data privacy policies, PFL is heavily relied upon. The most pervasive strategy is to share the feature extractor while training a separate mask prediction head for each client to accommodate local data distributions (Wang et al., 2023;Xie et al., 2024;Wang, Jin, and Wang, 2022;Jiang et al., 2023;Liu et al., 2025).
Despite recent strides made in tumor segmentation using PFL, existing approaches have largely overlooked the potential benefits of leveraging cross-organ information to enhance segmentation performance. In a federated setting, clients often possess segmentation data regarding anatomically different organs such as brain, liver, or breast. We hypothesize that a shared feature extractor with appropriate personalization can greatly benefit the segmentation process across all the clients. However, most existing methods are not well suited to accommodate data originating from different organs and often from different modalities (Madni, Umer, and Foresti, 2023;Babar, Qureshi, and Koubaa, 2024). This can potentially be handled by modeling long-range dependencies via cross-attention mechanisms, enabling the aggregation of global contextual information across the feature space of the clients. Furthermore, introducing personalization during feature extraction across clients not only aids in the adaption of the local models to the heterogeneous data distribution but also preserves data privacy (Wang et al., 2023). For effective client personalization local fine-tuning is typically required, where each client further trains its model for a few epochs on its private dataset. During this step, inconsistencies in the predicted masks greatly degrades the segmentation performance (Wang et al., 2023;Wang, Jin, and Wang, 2022). Moreover, the few methods that attempt to resolve this often introduce substantial computational or memory overhead. In general, our contributions are as follows:
• We propose a novel PFL framework for tumor segmentation specifically designed to handle segmentation tasks across heterogeneous organ datasets named FedOAP.
• We introduce DCA, a mechanism that decouples the query and key-value embedding layers. The query remains local to each client, while the aggregated key-value pairs from all clients are used to enable cross-organ feature fusion through crossattention.
• In addition, we propose PBL to improve proposed tumor mask accuracy. It identifies inconsistencies between predicted and ground truth boundaries while injecting stochastic noise into these regions, and formulates a targeted loss to enhance supervision on boundary pixels, thereby improving segmentation robustness and precision during client side fine-tuning.
We conduct comprehensive experiments on three tu-mor segmentation tasks from different organs: brain, liver, and breast. FedOAP achieves superior performance compared to existing state-of-the-art PFL methods, demonstrating its strong generalization capabilities and its effectivenes
This content is AI-processed based on open access ArXiv data.