A Novel Approach to Breast Cancer Segmentation using U-Net Model with Attention Mechanisms and FedProx

AI summary

Overview

• Research area: Medical image analysis / computer vision — breast cancer (BC) tumour segmentation in ultrasound images using federated learning with an attention-augmented U-Net.
• Technical level: Intermediate. Readers need some familiarity with segmentation metrics (Dice loss, IoU, sensitivity, specificity, F1, accuracy), U-Net architectures, and the basic federated learning client/server training loop.
• Scope: The paper proposes and evaluates a federated learning system that combines an Attention U-Net segmentation model with the FedProx aggregation algorithm across three clients holding non-IID breast ultrasound datasets, evaluated over 6 communication rounds.

What This Paper Is About

Breast ultrasound imaging is a low-cost, reliable screening tool, but building accurate AI segmentation models from it is difficult because medical data is sensitive and cannot easily be pooled in one place, and because images collected at different sites differ in quality, resolution, and contrast (the non-IID problem). The authors set out to segment breast tumours accurately without centralising patient data, by training an Attention U-Net locally at three clients and aggregating the weights on a server using FedProx. Their goal is to show that privacy-preserving distributed training can reach segmentation accuracy comparable to, or better than, previously reported approaches on non-IID ultrasonic breast cancer imaging (USBCI) data.

Key Contributions

1. A federated learning architecture for BC segmentation in which three clients train identical Attention U-Net models on distinct, non-IID breast ultrasound data sources and a server aggregates their weights with FedProx, so raw images never leave the client.
2. Use of a modified U-Net with attention mechanisms in place of a standard U-Net, motivated by the claim that standard U-Net and similar architectures (ResNet, V-Net) struggle to capture fine details and complex structures in USBCI.
3. A deliberately non-IID three-client data partition built from two public datasets (BUS A and BUS B), each client holding a different mix of benign, malignant, and normal cases, with augmentation (flip, rotation, translation, scaling, contrast and brightness adjustments) applied to client data.
4. A comparative evaluation of the resulting global model against five related published studies on six metrics, reporting a global model accuracy of 96%.

Main Findings

• Global model performance at convergence (Round 6): Dice Loss 0.2924, IoU 0.5494, Sensitivity 0.6066, Specificity 0.9919, F1 Score 0.7076, Accuracy 0.9607.
• Steady improvement across rounds: In Round 1 the global model had Dice Loss 0.895, IoU 0.0555, Sensitivity 0.1575, Specificity 0.8452, F1 0.105, and Accuracy 0.9204. By Round 6 these had moved to the values above, with loss decreasing and IoU, F1, and accuracy increasing.
• High specificity throughout: Specificity rose from 0.8452 in Round 1 to 0.9919 in Round 6, which the authors interpret as strong ability to identify true negatives.
• Low sensitivity relative to other metrics: Sensitivity reached 0.6066 at Round 6, the weakest of the reported metrics, though it improved substantially from 0.1575 in Round 1.
• Comparison with related studies (Table 2): This study's FedProx model reported Dice Loss 0.29, IoU 0.55, Sensitivity 0.64, Specificity 0.99, F1 0.71, Accuracy 0.96. Against Roth et al. (0.29 Dice Loss, 0.21 IoU, 0.59 sensitivity, 0.99 specificity, 0.33 F1, 0.87 accuracy) the proposed model matched Dice Loss and specificity while improving IoU and F1. Against Yang et al. (0.84, 0.77, 0.87, 0.98, 0.82, 0.96) and Wicaksana et al. (0.81, 0.73, 0.82, 0.98, 0.81, 0.96) it reached the same accuracy of 0.96 but lower values on the other metrics. Against Camajori et al. (0.78, 0.78, 0.65, 0.88, 0.98, 0.78) and Jiménez-Sánchez et al. (0.59, 0.44, 0.72, 0.94, 0.64, 0.90) it reported higher accuracy, specificity and F1 but lower IoU and sensitivity.
• Qualitative results: Visualisations of segmented test scans show predicted lesion boundaries aligning closely with ground-truth masks, with the authors reporting minimal noise in the predicted images.
• Data split used: Client 1 held 400 benign and 50 normal cases (BUS A); Client 2 held 200 malignant and 50 normal cases (BUS A); Client 3 held 110 benign and 53 malignant cases (BUS B). The server test set contained 97 benign, 23 malignant, and 34 normal cases drawn from both datasets.
• Datasets: BUS A ("Dataset of breast ultrasound images," Al-Dhabyani et al.) was collected in 2018 and contains 600 female patients and 780 scans with annotated masks (487 benign, 210 malignant, 133 normal). BUS B (Yap et al.) was collected in 2012 at the UDIAT Diagnostic Centre of the Parc Tauli Company in Sabadell, Spain, using a Siemens ACUSON Sequoia C512 system with a 17L5 HD linear array transducer, and contains 163 scans of breast lesions with masks (110 benign, 53 malignant).
• Training configuration reported: 6 federated rounds, each consisting of 10 local epochs per client; Attention U-Net trained with Dice loss and the Adam optimizer at a learning rate of 0.0001 and batch size of 16; per-client evaluation on 20% of local data. FedProx used a learning rate of 0.01, a regularization parameter of 0.001, a proximal parameter of 0.01, and a hyperparameter of μ = 0.1.
• Not reported: No statistical significance testing, no ablation isolating the effect of the attention mechanism versus a plain U-Net, no comparison against centralised (non-federated) training, and no communication cost or convergence-speed analysis are given.

Methodology in Plain English

The authors split breast ultrasound data across three simulated hospitals (clients) so that no single site holds the same mix of cases — this is what makes the setup non-IID. Each client runs the same Attention U-Net, a U-Net whose skip connections are gated by an attention mechanism that learns which features to amplify and which to suppress before passing them across the network. Inside the attention block, a lower-resolution feature vector guides a higher-resolution one: both are convolved, summed, passed through ReLU and a 1x1 convolution that collapses the dimensions to 1x32x32, then squeezed through a sigmoid into attention coefficients between 0 and 1, upsampled back to the original resolution by trilinear interpolation, and multiplied element-wise with the original feature vector.

Each client trains locally with Dice loss and the Adam optimizer, then sends only its updated weights to the server. The server aggregates those weights using FedProx, whose objective adds a proximal term (weighted by μ) that penalises local models for drifting too far from the previous global weights — a device intended to stabilise training when client data distributions differ. The aggregated model is sent back to the clients and the cycle repeats for 6 rounds. After each round, the server evaluates the global model on a held-out test set spanning both datasets, using dice loss, IoU, sensitivity, specificity, F1 score, and accuracy, and the authors compare the final results to five prior studies.

Why This Matters

The work speaks to a practical bind in medical AI: the best segmentation models need large, diverse datasets, but patient data cannot easily be shared across institutions. Showing that FedProx plus an attention-based U-Net can reach 96% global accuracy on non-IID ultrasound data — while keeping specificity at 0.9919 — supports the argument that privacy-preserving training need not come at a large accuracy cost.

Real-world applications:

• Multi-hospital breast screening networks, where several clinics could jointly improve a tumour-delineation model without transferring patient scans.
• Radiologist decision support, using boundaries the model draws around lesions to direct attention to regions of interest.
• Low-resource or remote care settings, where ultrasound is more available than mammography and where the paper notes US can detect lesions that may not be visible on mammography.
• Regulatory-sensitive deployments, where keeping raw data on-premises simplifies data-handling obligations compared with centralised training.

Industry relevance: the paper sits at the intersection of medical imaging vendors, cloud health platforms, and privacy-focused ML infrastructure. It offers a concrete demonstration that off-the-shelf federated aggregation (FedProx) can be paired with a modest architectural change (attention gates) to serve a clinically meaningful segmentation task, which is relevant to anyone building distributed clinical AI products.

Future Directions

• Raise sensitivity. At 0.6066, sensitivity is the model's weakest reported metric, meaning a meaningful share of true lesion regions may still be missed. Improving recall without sacrificing the 0.9919 specificity is the most obvious clinical priority.
• Run ablation studies. The paper attributes gains to the attention mechanism and FedProx but does not isolate each component, nor compare against plain U-Net, centralised training, or other aggregation strategies such as FedAvg.
• Scale the federation. The setup uses only 3 clients and 2 datasets over 6 rounds; testing with more clients, more heterogeneous scanners, and different non-IID partitions would indicate whether the approach holds up.
• Address communication and clinical validation. Communication overhead, straggler behaviour, and convergence speed are not analysed, and evaluation is limited to public datasets rather than prospective clinical data.

Target Audience

This paper is most useful to researchers and graduate students working on federated learning for medical imaging, to practitioners building privacy-preserving segmentation pipelines, and to clinical AI teams evaluating whether distributed training can match centralised performance on ultrasound breast cancer data. Readers with a background in segmentation metrics and deep learning will get the most out of the comparative tables, while those new to the area can use the paper as an accessible example of how a federated segmentation study is structured.

Authors’ abstract

Breast cancer is a leading cause of death among women worldwide, emphasizing the need for early detection and accurate diagnosis. As such Ultrasound Imaging, a reliable and cost-effective tool, is used for this purpose, however the sensitive nature of medical data makes it challenging to develop accurate and private artificial intelligence models. A solution is Federated Learning as it is a promising technique for distributed machine learning on sensitive medical data while preserving patient privacy. However, training on non-Independent and non-Identically Distributed (non-IID) local datasets can impact the accuracy and generalization of the trained model, which is crucial for accurate tumour boundary delineation in BC segmentation. This study aims to tackle this challenge by applying the Federated Proximal (FedProx) method to non-IID Ultrasonic Breast Cancer Imaging datasets. Moreover, we focus on enhancing tumour segmentation accuracy by incorporating a modified U-Net model with attention mechanisms. Our approach resulted in a global model with 96% accuracy, demonstrating the effectiveness of our method in enhancing tumour segmentation accuracy while preserving patient privacy. Our findings suggest that FedProx has the potential to be a promising approach for training precise machine learning models on non-IID local medical datasets.