Source-linked AI summary
HarmoFL: Harmonizing Local and Global Drifts in Federated Learning on Heterogeneous Medical Images
Meirui Jiang, Zirui Wang, Qi Dou
TL;DR
Medical federated learning is challenged by non-iid feature shifts that create coupled local and global optimization drifts, while prior methods generally address only one side. HarmoFL combines frequency-domain amplitude normalization with gradient-generated weight perturbation to harmonize features and produce flatter client optima without extra communication. Across three heterogeneous medical-image tasks, the authors report consistent effectiveness and state-of-the-art performance, supported by bounded-drift analysis.
Problem
Non-iid feature shifts from scanners, protocols, and other acquisition differences create local and global optimization drifts, while jointly addressing these coupled drifts remains unclear.
Method
HarmoFL normalizes frequency-space image amplitudes for local feature harmonization and applies locally gradient-generated weight perturbations to produce flat client optima without extra communication.
Results
HarmoFL consistently demonstrates effectiveness across one classification and two segmentation tasks and significantly outperforms recent state-of-the-art federated-learning methods.
Takeaways & Limitations
The framework supports simultaneously handling coupled local and global drifts, with theoretical analysis showing the overall non-iid drift is bounded.
Takeaways & Limitations
The analysis focuses on non-iid feature shift, including cases where client feature distributions vary while selected label-related conditionals remain unchanged.
Abstract
from arXiv · showhide
Multiple medical institutions collaboratively training a model using federated learning (FL) has become a promising solution for maximizing the potential of data-driven models, yet the non-independent and identically distributed (non-iid) data in medical images is still an outstanding challenge in real-world practice. The feature heterogeneity caused by diverse scanners or protocols introduces a drift in the learning process, in both local (client) and global (server) optimizations, which harms the convergence as well as model performance. Many previous works have attempted to address the non-iid issue by tackling the drift locally or globally, but how to jointly solve the two essentially coupled drifts is still unclear. In this work, we concentrate on handling both local and global drifts and introduce a new harmonizing framework called HarmoFL. First, we propose to mitigate the local update drift by normalizing amplitudes of images transformed into the frequency domain to mimic a unified imaging setting, in order to generate a harmonized feature space across local clients. Second, based on harmonized features, we design a client weight perturbation guiding each local model to reach a flat optimum, where a neighborhood area of the local optimal solution has a uniformly low loss. Without any extra communication cost, the perturbation assists the global model to optimize towards a converged optimal solution by aggregating several local flat optima. We have theoretically analyzed the proposed method and empirically conducted extensive experiments on three medical image classification and segmentation tasks, showing that HarmoFL outperforms a set of recent state-of-the-art methods with promising convergence behavior. Code is available at https://github.com/med-air/HarmoFL.
1 Introduction
Medical federated learning faces feature heterogeneity that creates coupled local and global optimization drifts. HarmoFL addresses both through frequency-space amplitude normalization and weight perturbation for flatter, more aggregatable client solutions.
- Motivation: Feature shifts from staining, scanners, and imaging protocols make non-iid data a persistent challenge in medical federated learning.These shifts occur across institutions even when collaborative training avoids direct data sharing.
- Motivation: Data heterogeneity drives local models toward different client optima and causes server aggregation to drift among mismatching solutions.The coupled drifts can make convergence slow and unstable.
- Related problem: Prior methods mainly address heterogeneity either through local normalization or through global aggregation and optimization strategies.The introduction identifies joint treatment of the coupled drifts as unresolved.
- HarmoFL: HarmoFL normalizes frequency-space image amplitudes to harmonize client features, then applies locally generated weight perturbations to reduce server update drift.The perturbation targets a neighborhood around each local optimum with uniformly low loss and requires no extra communication cost.
- Contributions: The framework is theoretically analyzed and evaluated on three medical image tasks, where it reportedly outperforms recent state-of-the-art federated learning methods.The tasks include breast cancer histology classification, histology nuclei segmentation, and prostate MRI segmentation.
2 Related Work
Related work addresses federated-learning heterogeneity through client-side training or server-side aggregation. HarmoFL instead jointly mitigates the coupled local and global drifts.
- Local client training: Client-side approaches reduce local drift using preprocessing, domain losses, simulated data, proximal regularization, or locally retained normalization layers.These methods target heterogeneous feature distributions during local training.
- Global server aggregation: Server-side approaches reduce global drift through shared data, adaptive aggregation weights, normalized gradients, or adaptive optimization.Examples include inverse-distance reweighting, FedNova, and FedAdam.
- Research gap: Existing methods partially address drift from either the local-client or global-server perspective.The paper positions its contribution as simultaneous mitigation of both coupled drifts.
3 Methodology
HarmoFL jointly addresses local and global drift in heterogeneous medical-image federated learning. It harmonizes client features through frequency-domain amplitude normalization, then uses gradient-based weight perturbation to promote flat local optima for aggregation.
- Framework overview: HarmoFL targets both client-side and server-side update drift caused by heterogeneous feature distributions across medical-image clients.The method is analyzed theoretically for overall drift under convex and non-convex settings.
- Local harmonization: Amplitude normalization decomposes images in the frequency domain and unifies amplitude components while preserving phase information locally.Amplitude captures low-level visual features such as color and contrast, whereas phase is retained locally to preserve structure-related information.
- Local harmonization: The client amplitude statistic is updated with a moving average, and a global amplitude is formed from shared average-amplitude information for later rounds.Only average amplitude information is communicated; the original image content and local phase components are not shared.
- Global harmonization: Gradient-normalized weight perturbation is generated from harmonized features and applied directly to client parameters without extra communication cost.The perturbation is controlled by α and is used while minimizing loss on the perturbed model.
- Global harmonization: The perturbation drives clients toward flat optima with uniformly low loss across a neighborhood, making local models easier to aggregate globally.Amplitude normalization reduces feature-related gradient dissimilarity, while flat optima constrain loss variation and gradient changes.
- Optimization procedure: The training procedure repeatedly normalizes local batches, applies weight perturbation during client updates, and aggregates the resulting client models over communication rounds.Algorithm 1 uses client and global learning rates η_l and η_g, mini-batch steps K, and perturbation hyper-parameter α.
4 Experiments
Experiments evaluate HarmoFL on heterogeneous medical-image classification and segmentation tasks, comparing performance, convergence, and robustness against recent federated-learning methods. HarmoFL achieves stronger results and more stable behavior across the reported settings.
- Experimental settings: HarmoFL is evaluated on breast cancer histology classification, histology nuclei segmentation, and prostate MRI segmentation using heterogeneous client data.The experiments report averages across three repeated runs with standard deviations across random seeds.
- Comparison with state-of-the-art methods: 95.48% average accuracy on breast histology classification, 8% higher than FedBN, with smaller accuracy variation across clients.The reported result indicates that heterogeneous clients benefit more evenly under HarmoFL.
- Comparison with state-of-the-art methods: At least 4% higher mean accuracy on histology nuclei segmentation than the other methods, including FedBN.For prostate MRI segmentation, HarmoFL reaches the highest Dice of 94.28% with a smaller standard deviation across clients.
- Ablation study: HarmoFL converges smoothly over communication rounds, whereas FedBN shows unstable convergence and lower accuracy.The convergence plot averages testing accuracy across five clients over 100 communication rounds with one local update epoch.
- Comparison with state-of-the-art methods: Qualitative segmentation comparisons cover prostate MRI in the top two rows and histology nuclei segmentation in the bottom two rows.The comparison examines how heterogeneous features cause competing methods to over- or under-segment relative to ground truth.
- Ablation study: With more local training epochs, FedAvg and FedBN suffer severe performance drops, while HarmoFL remains robust to the resulting larger drift.The ablation attributes HarmoFL’s robustness to weight perturbation trained with harmonized features.
- Ablation study: HarmoFL reaches its highest accuracy at perturbation degree α = 5e−2, while exceeding 90% accuracy even at α = 1.The experiment searches α ∈{1, 5e−1, 5e−2, 5e−3, 5e−4}.
5 Conclusion
The conclusion presents HarmoFL as a harmonizing strategy that addresses local client and global server drifts together in federated learning for heterogeneous medical images. Experiments and theoretical analysis support its effectiveness and bounded overall non-iid drift.
- 5 Conclusion: HarmoFL combines amplitude normalization and weight perturbation to address local client and global server drifts.The framework is presented as jointly handling the coupled drifts rather than treating them separately.
- 5 Conclusion: Experiments across one classification and two segmentation tasks consistently demonstrate the effectiveness of the proposed approach.
- 5 Conclusion: Theoretical analysis shows that HarmoFL bounds the overall non-iid drift caused by data heterogeneity.
- 5 Conclusion: The authors position HarmoFL as beneficial for wider real-world medical applications of federated learning.
A Notation Table
The notation table defines the objectives, data distributions, normalization and perturbation variables, optimization parameters, and drift quantities used throughout HarmoFL.
- Objectives and data: F denotes the global objective function, while F_i denotes the i-th client’s local objective function.
- Objectives and data: (X, Y) is the joint image-label space, and (x, y) is a data sample pair.
- Optimization dimensions: T, N, K, and M denote communication rounds, clients, mini-batch steps, and sampled images per batch, respectively.
- Harmonization: Ψ(·) is amplitude normalization, D_i is client data distribution, and D_i is the distribution harmonized by amplitude normalization.
- Harmonization: A_i,x_m, P_i,x_m, and A_i,k denote image amplitude, image phase, and average batch amplitude, respectively.
- Perturbation and optimization: δ and α represent parameter perturbation and its degree-control hyper-parameter, while η_l and η_g are client and global learning rates.
- Drift analysis: Γ denotes the overall non-iid drift term, and ϵ is a non-negative constant bounding gradient differences.
B.1 Assumptions
The assumptions and lemmas formalize bounded gradient dissimilarity, bounded stochastic-gradient variance, vector inequalities, and separation of means from variance for the analysis.
- Assumptions: Assumption B.2 imposes bounded gradient dissimilarity through a Lipschitz-type bound involving β.
- Assumptions: For convex local objectives, the bounded-gradient-dissimilarity assumption can be relaxed.
- Assumptions: Assumption B.3 treats g_i(θ) as an unbiased stochastic gradient with bounded variance.
- Supporting lemmas: Lemma B.4 provides a relaxed triangle inequality for vectors in R^d.
- Supporting lemmas: Lemma B.5 separates mean and variance for potentially dependent random variables with bounded variance.
B.2 Theorem of bounded drift and proof
Theorem B.6 analyzes HarmoFL under stated function assumptions and derives a bounded-drift result using recursive inequalities for local updates. The proof covers both non-convex and convex objectives.
- Assumptions: Theorem B.6 assumes the functions {Fi} satisfy assumptions B.1, B.2, and B.3.
- Update analysis: HarmoFL uses effective step-size ˜η = Kηgηl in its update analysis.
- Proof: The proof handles the K = 1 case directly and then analyzes K ≥ 2 through recursive drift inequalities.
- Results: For non-convex functions {Fi}, averaging over clients and local steps yields a bounded-drift relationship.
- Results: For convex functions {Fi}, the analysis likewise provides a bounded-drift relationship under the theorem’s assumptions.
C Complete experiment details and results
The experiments use a fixed PyTorch and hardware setup, with controlled randomness across three specified seeds.
- Implementation: All experiments use PyTorch with Python 3.6.10 and run training and testing on a TITAN RTX GPU.
- Implementation: Randomness is controlled with three random seeds: 0, 1, and 2.
C.1 Experimental details
The experiments cover breast cancer histology classification and nuclei segmentation using institutionally partitioned medical-image datasets with specified preprocessing and training protocols.
- Breast cancer histology image classification: Camelyon17 provides 450,000 breast-cancer patches from five hospitals for tumor-region classification.Images are preprocessed to 96 × 96 × 3 for all clients.
- Nuclei segmentation: MoNuSAC2020 is divided into four clients using the official multi-organ split, with no overlap between hospital groups.Images are reshaped to 256 × 256 × 3, and 20% of training data is reserved for validation.
- Nuclei segmentation: The MoNuSAC2020 segmentation models use U-Net and train for 500 epochs with segmentation Dice loss.
C.2 Ablation studies
Ablation studies across classification and two segmentation tasks show successive gains from amplitude normalization and weight perturbation over FedAvg.
- Experimental design: The ablation studies evaluate the proposed components across classification, nuclei segmentation, and prostate segmentation tasks.The corresponding results are reported in Tables 4, 5, and 6.
- Amplitude normalization: Amplitude normalization consistently improves performance over FedAvg across all clients and all three datasets.This result is reported as evidence for harmonizing local drifts.
- Weight perturbation: Adding weight perturbation on top of amplitude normalization further boosts performance with a clear margin.The ablation attributes this component to mitigating global drifts.
C.3 Additional results
Additional experiments show that HarmoFL produces flatter, more stable solutions than FedAvg and preserves semantic structure during amplitude normalization. Qualitative segmentation comparisons also show more complete and accurate boundaries across histology and prostate MRI images.
- Loss landscape visualization: HarmoFL shows smaller loss variation and flatter local and global solutions than FedAvg under parameter changes.The method remains robust even with changing scales eight times larger.
- Amplitude normalization: Amplitude normalization changes low-level appearance while preserving high-level semantic structures.Large receptive fields can compensate for tiny changes near ambiguous boundaries using contextual information.
- Segmentation visualization: HarmoFL produces more complete histology nuclei segmentations, separating small cells and recovering nuclei that competing methods miss.Compared methods may fragment large cells, merge small nuclei, or fail against similar backgrounds.
- Segmentation visualization: The additional visualizations compare HarmoFL with other state-of-the-art methods on prostate MRI and histology nuclei segmentation tasks.These comparisons are presented in Figures 9 and 10.
- Segmentation visualization: On prostate MRI images, HarmoFL delineates accurate boundaries despite feature-distribution shifts across hospitals.Other methods particularly struggle to identify structures in difficult examples.