# An Optimal Transport Approach to Personalized Federated Learning

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

## Key facts

- Cited by 12 papers in the map (data through 14 Oct 2026).
- Builds on 41 of its 57 references that are in the map.
- First posted on arXiv 6 Jun 2022.

## Builds on (most cited first, 30 of 41)

- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [Going Deeper with Convolutions](https://mlsanity.com/paper/1409.4842/) (2014)
- [Advances and Open Problems in Federated Learning](https://mlsanity.com/paper/1912.04977/) (2019)
- [Federated Optimization in Heterogeneous Networks](https://mlsanity.com/paper/1812.06127/) (2018)
- [Invariant Risk Minimization](https://mlsanity.com/paper/1907.02893/) (2019)
- [On the Convergence of FedAvg on Non-IID Data](https://mlsanity.com/paper/1907.02189/) (2019)
- [Federated Learning with Non-IID Data](https://mlsanity.com/paper/1806.00582/) (2018)
- [Spectrally-normalized margin bounds for neural networks](https://mlsanity.com/paper/1706.08498/) (2017)
- [Federated Multi-Task Learning](https://mlsanity.com/paper/1705.10467/) (2017)
- [TensorFlow: Large-Scale Machine Learning on Heterogeneous Distributed Systems](https://mlsanity.com/paper/1603.04467/) (2016)
- [FedBN: Federated Learning on Non-IID Features via Local Batch Normalization](https://mlsanity.com/paper/2102.07623/) (2021)
- [Exploiting Shared Representations for Personalized Federated Learning](https://mlsanity.com/paper/2102.07078/) (2021)
- [An Efficient Framework for Clustered Federated Learning](https://mlsanity.com/paper/2006.04088/) (2020)
- [Agnostic Federated Learning](https://mlsanity.com/paper/1902.00146/) (2019)
- [Adaptive Personalized Federated Learning](https://mlsanity.com/paper/2003.13461/) (2020)
- [Think Locally, Act Globally: Federated Learning with Local and Global Representations](https://mlsanity.com/paper/2001.01523/) (2020)
- [Three Approaches for Personalization with Applications to Federated Learning](https://mlsanity.com/paper/2002.10619/) (2020)
- [On Gradient Descent Ascent for Nonconvex-Concave Minimax Problems](https://mlsanity.com/paper/1906.00331/) (2019)
- [Improving Federated Learning Personalization via Model Agnostic Meta Learning](https://mlsanity.com/paper/1909.12488/) (2019)
- [Model Fusion via Optimal Transport](https://mlsanity.com/paper/1910.05653/) (2019)
- [Fast Computation of Wasserstein Barycenters](https://mlsanity.com/paper/1310.4375/) (2013)
- [Personalized Federated Learning using Hypernetworks](https://mlsanity.com/paper/2103.04628/) (2021)
- [Federated Learning of a Mixture of Global and Local Models](https://mlsanity.com/paper/2002.05516/) (2020)
- [Statistical Aspects of Wasserstein Distances](https://mlsanity.com/paper/1806.05500/) (2018)
- [SCAFFOLD: Stochastic Controlled Averaging for Federated Learning](https://mlsanity.com/paper/1910.06378/) (2019)
- [Wasserstein Distributionally Robust Optimization: Theory and Applications in Machine Learning](https://mlsanity.com/paper/1908.08729/) (2019)
- [Solving a Class of Non-Convex Min-Max Games Using Iterative First Order Methods](https://mlsanity.com/paper/1902.08297/) (2019)
- [Multi-Center Federated Learning: Clients Clustering for Better Personalization](https://mlsanity.com/paper/2108.08647/) (2021)
- [Distributionally Robust Federated Averaging](https://mlsanity.com/paper/2102.12660/) (2021)
- [Robust Federated Learning: The Case of Affine Distribution Shifts](https://mlsanity.com/paper/2006.08907/) (2020)

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

- [Global and Local Prompts Cooperation via Optimal Transport for Federated Learning](https://mlsanity.com/paper/2403.00041/) (2024)
- [Federated Domain Generalization: A Survey](https://mlsanity.com/paper/2306.01334/) (2023)
- [Robust Decentralized Learning with Local Updates and Gradient Tracking](https://mlsanity.com/paper/2405.00965/) (2024)
- [PersA-FL: Personalized Asynchronous Federated Learning](https://mlsanity.com/paper/2210.01176/) (2022)
- [pFedMMA: Personalized Federated Fine-Tuning with Multi-Modal Adapter for Vision-Language Models](https://mlsanity.com/paper/2507.05394/) (2025)
- [Distribution Alignment for One-Shot Federated Learning via Optimal Transport](https://mlsanity.com/paper/2606.16655/) (2026)
- [Certifiably Robust Model Evaluation in Federated Learning under Meta-Distributional Shifts](https://mlsanity.com/paper/2410.20250/) (2024)
- [Federated Learning for Discrete Optimal Transport with Large Population under Incomplete Information](https://mlsanity.com/paper/2411.07841/) (2024)
- [Optimal Transport-based Domain Alignment as a Preprocessing Step for Federated Learning](https://mlsanity.com/paper/2506.04071/) (2025)
- [Cost-Free Personalization via Information-Geometric Projection in Bayesian Federated Learning](https://mlsanity.com/paper/2509.10132/) (2025)
- [Data Valuation and Selection in a Federated Model Marketplace](https://mlsanity.com/paper/2509.18104/) (2025)
- [Collaborative and Efficient Fine-tuning: Leveraging Task Similarity](https://mlsanity.com/paper/2602.07218/) (2026)

## Abstract (arXiv)

Federated learning is a distributed machine learning paradigm, which aims to train a model using the local data of many distributed clients. A key challenge in federated learning is that the data samples across the clients may not be identically distributed. To address this challenge, personalized federated learning with the goal of tailoring the learned model to the data distribution of every individual client has been proposed. In this paper, we focus on this problem and propose a novel personalized Federated Learning scheme based on Optimal Transport (FedOT) as a learning algorithm that learns the optimal transport maps for transferring data points to a common distribution as well as the prediction model under the applied transport map. To formulate the FedOT problem, we extend the standard optimal transport task between two probability distributions to multi-marginal optimal transport problems with the goal of transporting samples from multiple distributions to a common probability domain. We then leverage the results on multi-marginal optimal transport problems to formulate FedOT as a min-max optimization problem and analyze its generalization and optimization properties. We discuss the results of several numerical experiments to evaluate the performance of FedOT under heterogeneous data distributions in federated learning problems.