# FedAVOT: Exact Distribution Alignment in Federated Learning via Masked Optimal Transport

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

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 18 of its 28 references that are in the map.
- First posted on arXiv 17 Sep 2025.

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

- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [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)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [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)
- [Tackling the Objective Inconsistency Problem in Heterogeneous Federated Optimization](https://mlsanity.com/paper/2007.07481/) (2020)
- [Towards Federated Learning at Scale: System Design](https://mlsanity.com/paper/1902.01046/) (2019)
- [Computational Optimal Transport](https://mlsanity.com/paper/1803.00567/) (2018)
- [Agnostic Federated Learning](https://mlsanity.com/paper/1902.00146/) (2019)
- [Fair Resource Allocation in Federated Learning](https://mlsanity.com/paper/1905.10497/) (2019)
- [Optimal Client Sampling for Federated Learning](https://mlsanity.com/paper/2010.13723/) (2020)
- [Federated Learning Under Intermittent Client Availability and Time-Varying Communication Constraints](https://mlsanity.com/paper/2205.06730/) (2022)
- [Optimal Importance Sampling for Federated Learning](https://mlsanity.com/paper/2010.13600/) (2020)
- [FedGS: Federated Graph-based Sampling with Arbitrary Client Availability](https://mlsanity.com/paper/2211.13975/) (2022)
- [Federated Learning Under Restricted User Availability](https://mlsanity.com/paper/2309.14176/) (2023)
- [The Stochastic Multi-Proximal Method for Nonsmooth Optimization](https://mlsanity.com/paper/2505.12409/) (2025)
- [Convergence of Agnostic Federated Averaging](https://mlsanity.com/paper/2507.10325/) (2025)

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

- [FedeRage: Provably Convergent Agnostic Federated Learning under General Client Drift](https://mlsanity.com/paper/2609.21057/) (2026)

## Abstract (arXiv)

Federated Learning (FL) allows distributed model training without sharing raw data, but suffers when client participation is partial. In practice, the distribution of available users (availability distribution q) rarely aligns with the distribution defining the optimization objective (importance distribution p), leading to biased and unstable updates under classical FedAvg. We propose textbfFereated AVerage with Optimal Transport (FedAVOT), which formulates aggregation as a masked optimal transport problem aligning q and p. Using Sinkhorn scaling, FedAVOT computes transport-based aggregation weights with provable convergence guarantees. FedAVOT achieves a standard O(1/√T) rate under a nonsmooth convex FL setting, independent of the number of participating users per round. Our experiments confirm drastically improved performance compared to FedAvg across heterogeneous, fairness-sensitive, and low-availability regimes, even when only two clients participate per round.