FedAVOT: Exact Distribution Alignment in Federated Learning via Masked Optimal Transport
This addresses stability and fairness issues in federated learning for distributed systems, offering a novel solution with provable guarantees.
The paper tackled the problem of biased and unstable updates in Federated Learning due to misalignment between client availability and importance distributions, proposing FedAVOT which uses masked optimal transport for aggregation and achieves an O(1/√T) convergence rate independent of user participation.
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 (\emph{availability distribution} $q$) rarely aligns with the distribution defining the optimization objective (\emph{importance distribution} $p$), leading to biased and unstable updates under classical FedAvg. We propose \textbf{Fereated AVerage with Optimal Transport (\textbf{FedAVOT})}, which formulates aggregation as a masked optimal transport problem aligning $q$ and $p$. Using Sinkhorn scaling, \textbf{FedAVOT} computes transport-based aggregation weights with provable convergence guarantees. \textbf{FedAVOT} achieves a standard $\mathcal{O}(1/\sqrt{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.