An Optimal Transport Approach to Personalized Federated Learning
This addresses the problem of data heterogeneity in federated learning for distributed clients, offering a novel method but with incremental improvements over existing personalized approaches.
The paper tackles the challenge of non-identical data distributions across clients in federated learning by proposing FedOT, a personalized scheme based on optimal transport that learns transport maps to a common distribution and a prediction model, achieving improved performance in numerical experiments under heterogeneous data.
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.