FedLPPA: Learning Personalized Prompt and Aggregation for Federated Weakly-supervised Medical Image Segmentation

TL;DR

FedLPPA tackles federated learning under heterogeneous weak supervision for medical image segmentation by introducing a Tri-prompt Dual-attention Fusion (TDF) module and a Prompt similarity Dual-decoder with Learnable Aggregation (PDLA). It maintains three learnable prompts—universal knowledge prompt $p_U$, data distribution prompt $p_{D,i}$, and an annotation sparsity prompt $p_S$—and fuses them with sample features through a dual-attention mechanism, enabling personalized adaptation of each local decoder. The server-side PDLA mechanism uses an affinity-based prompt selection strategy and learnable aggregation to generate high-quality pseudo-labels, while the local LA adjusts decoder parameters on a per-client basis. Across four medical-imaging tasks, FedLPPA outperforms standard FL and state-of-the-art personalized FL baselines, closely matching fully supervised centralized performance and demonstrating effective privacy-preserving, annotation-efficient learning for heterogeneous clinical data. The approach offers practical impact for scalable, cross-institutional segmentation with diverse weak supervision formats.

Abstract

Federated learning (FL) effectively mitigates the data silo challenge brought about by policies and privacy concerns, implicitly harnessing more data for deep model training. However, traditional centralized FL models grapple with diverse multi-center data, especially in the face of significant data heterogeneity, notably in medical contexts. In the realm of medical image segmentation, the growing imperative to curtail annotation costs has amplified the importance of weakly-supervised techniques which utilize sparse annotations such as points, scribbles, etc. A pragmatic FL paradigm shall accommodate diverse annotation formats across different sites, which research topic remains under-investigated. In such context, we propose a novel personalized FL framework with learnable prompt and aggregation (FedLPPA) to uniformly leverage heterogeneous weak supervision for medical image segmentation. In FedLPPA, a learnable universal knowledge prompt is maintained, complemented by multiple learnable personalized data distribution prompts and prompts representing the supervision sparsity. Integrated with sample features through a dual-attention mechanism, those prompts empower each local task decoder to adeptly adjust to both the local distribution and the supervision form. Concurrently, a dual-decoder strategy, predicated on prompt similarity, is introduced for enhancing the generation of pseudo-labels in weakly-supervised learning, alleviating overfitting and noise accumulation inherent to local data, while an adaptable aggregation method is employed to customize the task decoder on a parameter-wise basis. Extensive experiments on four distinct medical image segmentation tasks involving different modalities underscore the superiority of FedLPPA, with its efficacy closely parallels that of fully supervised centralized training. Our code and data will be available.

Figures (10)

• Figure 1: A: Data samples from different centers showcase the domain gaps in their distributions; B: Examples of typical weak labels and their corresponding masks, with annotation granularity from sparse to dense.
• Figure 2: Schematic illustration of the proposed FedLPPA framework. Note that the segmentation heads and the MLP block are incorporated within the decoders and are not depicted independently, except in the left panel of the scheme. The auxiliary decoder aggregation on the right panel builds its basis on the selection of the Prompt Similarity Aggregated strategy.
• Figure 3: Detailed illustration of the Tri-prompt Dual-attention Fusion (TDF) module and the Learnable Aggregation (LA) for dual-decoder mechanism. The term 'context-injected' denotes that the training conditions or contextual information have been integrated into the original image features, thereby resulting in enhanced features.
• Figure 4: Preprocessing operations for box annotations.
• Figure 5: Visualization results from FedLPPA and other SOTA methods (CT denotes centralized training, with Weak and Full respectively indicating the use of sparse annotations and full masks for model training).