# Optimal Transport-based Domain Alignment as a Preprocessing Step for Federated Learning

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 16 of its 22 references that are in the map.
- First posted on arXiv 4 Jun 2025.

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

- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Federated Learning: Challenges, Methods, and Future Directions](https://mlsanity.com/paper/1908.07873/) (2019)
- [Generalizing to Unseen Domains: A Survey on Domain Generalization](https://mlsanity.com/paper/2103.03097/) (2021)
- [Model-Contrastive Federated Learning](https://mlsanity.com/paper/2103.16257/) (2021)
- [Federated Learning with Matched Averaging](https://mlsanity.com/paper/2002.06440/) (2020)
- [Clustered Federated Learning: Model-Agnostic Distributed Multi-Task Optimization under Privacy Constraints](https://mlsanity.com/paper/1910.01991/) (2019)
- [Towards Personalized Federated Learning](https://mlsanity.com/paper/2103.00710/) (2021)
- [Learning to Generate Novel Domains for Domain Generalization](https://mlsanity.com/paper/2007.03304/) (2020)
- [No Fear of Heterogeneity: Classifier Calibration for Federated Learning with Non-IID Data](https://mlsanity.com/paper/2106.05001/) (2021)
- [Fast Computation of Wasserstein Barycenters](https://mlsanity.com/paper/1310.4375/) (2013)
- [FedHealth: A Federated Transfer Learning Framework for Wearable Healthcare](https://mlsanity.com/paper/1907.09173/) (2019)
- [Federated Visual Classification with Real-World Data Distribution](https://mlsanity.com/paper/2003.08082/) (2020)
- [Regularized Discrete Optimal Transport](https://mlsanity.com/paper/1307.5551/) (2013)
- [Federated Learning with Domain Generalization](https://mlsanity.com/paper/2111.10487/) (2021)
- [An Optimal Transport Approach to Personalized Federated Learning](https://mlsanity.com/paper/2206.02468/) (2022)

## Abstract (arXiv)

Federated learning (FL) is a subfield of machine learning that avoids sharing local data with a central server, which can enhance privacy and scalability. The inability to consolidate data leads to a unique problem called dataset imbalance, where agents in a network do not have equal representation of the labels one is trying to learn to predict. In FL, fusing locally-trained models with unbalanced datasets may deteriorate the performance of global model aggregation, and reduce the quality of updated local models and the accuracy of the distributed agents' decisions. In this work, we introduce an Optimal Transport-based preprocessing algorithm that aligns the datasets by minimizing the distributional discrepancy of data along the edge devices. We accomplish this by leveraging Wasserstein barycenters when computing channel-wise averages. These barycenters are collected in a trusted central server where they collectively generate a target RGB space. By projecting our dataset towards this target space, we minimize the distributional discrepancy on a global level, which facilitates the learning process due to a minimization of variance across the samples. We demonstrate the capabilities of the proposed approach over the CIFAR-10 dataset, where we show its capability of reaching higher degrees of generalization in fewer communication rounds.