# Tackling Data Heterogeneity in Federated Learning via Loss Decomposition

arXiv: https://arxiv.org/abs/2408.12300 · HTML: https://mlsanity.com/paper/2408.12300/

## Key facts

- Cited by 9 papers in the map (data through 14 Oct 2026).
- Builds on 25 of its 35 references that are in the map.
- First posted on arXiv 22 Aug 2024.

## Builds on (most cited first, 25 of 25)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [Federated Optimization in Heterogeneous Networks](https://mlsanity.com/paper/1812.06127/) (2018)
- [Shortcut Learning in Deep Neural Networks](https://mlsanity.com/paper/2004.07780/) (2020)
- [WILDS: A Benchmark of in-the-Wild Distribution Shifts](https://mlsanity.com/paper/2012.07421/) (2020)
- [An overview of gradient descent optimization algorithms](https://mlsanity.com/paper/1609.04747/) (2016)
- [Measuring the Effects of Non-Identical Data Distribution for Federated Visual Classification](https://mlsanity.com/paper/1909.06335/) (2019)
- [FedBN: Federated Learning on Non-IID Features via Local Batch Normalization](https://mlsanity.com/paper/2102.07623/) (2021)
- [Federated Learning with Matched Averaging](https://mlsanity.com/paper/2002.06440/) (2020)
- [Representation Learning on Graphs: Methods and Applications](https://mlsanity.com/paper/1709.05584/) (2017)
- [Multi-site fMRI Analysis Using Privacy-preserving Federated Learning and Domain Adaptation: ABIDE Results](https://mlsanity.com/paper/2001.05647/) (2020)
- [Rethinking Architecture Design for Tackling Data Heterogeneity in Federated Learning](https://mlsanity.com/paper/2106.06047/) (2021)
- [Personalized Federated Learning with Feature Alignment and Classifier Collaboration](https://mlsanity.com/paper/2306.11867/) (2023)
- [Layer-wised Model Aggregation for Personalized Federated Learning](https://mlsanity.com/paper/2205.03993/) (2022)
- [Federated Learning Enables Big Data for Rare Cancer Boundary Detection](https://mlsanity.com/paper/2204.10836/) (2022)
- [FedCM: Federated Learning with Client-level Momentum](https://mlsanity.com/paper/2106.10874/) (2021)
- [On Large-Cohort Training for Federated Learning](https://mlsanity.com/paper/2106.07820/) (2021)
- [Decentralized Federated Learning through Proxy Model Sharing](https://mlsanity.com/paper/2111.11343/) (2021)
- [Label-Efficient Self-Supervised Federated Learning for Tackling Data Heterogeneity in Medical Imaging](https://mlsanity.com/paper/2205.08576/) (2022)
- [Don't blame Dataset Shift! Shortcut Learning due to Gradients and Cross Entropy](https://mlsanity.com/paper/2308.12553/) (2023)
- [SplitAVG: A heterogeneity-aware federated deep learning method for medical imaging](https://mlsanity.com/paper/2107.02375/) (2021)
- [An Experimental Study of Data Heterogeneity in Federated Learning Methods for Medical Imaging](https://mlsanity.com/paper/2107.08371/) (2021)
- [PGrad: Learning Principal Gradients For Domain Generalization](https://mlsanity.com/paper/2305.01134/) (2023)
- [Federated Learning with Manifold Regularization and Normalized Update Reaggregation](https://mlsanity.com/paper/2311.05924/) (2023)
- [Tackling the Non-IID Issue in Heterogeneous Federated Learning by Gradient Harmonization](https://mlsanity.com/paper/2309.06692/) (2023)

## Built on by (most cited first, 9 of 9)

- [Selective Aggregation for Low-Rank Adaptation in Federated Learning](https://mlsanity.com/paper/2410.01463/) (2024)
- [One-shot Federated Learning via Synthetic Distiller-Distillate Communication](https://mlsanity.com/paper/2412.05186/) (2024)
- [Exploring the Vulnerabilities of Federated Learning: A Deep Dive into Gradient Inversion Attacks](https://mlsanity.com/paper/2503.11514/) (2025)
- [FedVLMBench: Benchmarking Federated Fine-Tuning of Vision-Language Models](https://mlsanity.com/paper/2506.09638/) (2025)
- [Spatio-Temporal Representation Decoupling and Enhancement for Federated Instrument Segmentation in Surgical Videos](https://mlsanity.com/paper/2506.23759/) (2025)
- [Exploring Federated Pruning for Large Language Models](https://mlsanity.com/paper/2505.13547/) (2025)
- [Federated CLIP for Resource-Efficient Heterogeneous Medical Image Classification](https://mlsanity.com/paper/2511.07929/) (2025)
- [VGS-ATD: Robust Distributed Learning for Multi-Label Medical Image Classification Under Heterogeneous and Imbalanced Conditions](https://mlsanity.com/paper/2507.18657/) (2025)
- [SurgFed: Language-guided Multi-Task Federated Learning for Surgical Video Understanding](https://mlsanity.com/paper/2603.09496/) (2026)

## Abstract (arXiv)

Federated Learning (FL) is a rising approach towards collaborative and privacy-preserving machine learning where large-scale medical datasets remain localized to each client. However, the issue of data heterogeneity among clients often compels local models to diverge, leading to suboptimal global models. To mitigate the impact of data heterogeneity on FL performance, we start with analyzing how FL training influence FL performance by decomposing the global loss into three terms: local loss, distribution shift loss and aggregation loss. Remarkably, our loss decomposition reveals that existing local training-based FL methods attempt to reduce the distribution shift loss, while the global aggregation-based FL methods propose better aggregation strategies to reduce the aggregation loss. Nevertheless, a comprehensive joint effort to minimize all three terms is currently limited in the literature, leading to subpar performance when dealing with data heterogeneity challenges. To fill this gap, we propose a novel FL method based on global loss decomposition, called FedLD, to jointly reduce these three loss terms. Our FedLD involves a margin control regularization in local training to reduce the distribution shift loss, and a principal gradient-based server aggregation strategy to reduce the aggregation loss. Notably, under different levels of data heterogeneity, our strategies achieve better and more robust performance on retinal and chest X-ray classification compared to other FL algorithms. Our code is available at https://github.com/Zeng-Shuang/FedLD.