Personalized Federated Learning under Model Dissimilarity Constraints
Samuel Erickson, Mikael Johansson
TL;DR
This work tackles statistical heterogeneity in federated learning by proposing KARULA, a personalized FL framework that constrains pairwise client model dissimilarities using a Wasserstein-distance surrogate computed via data embeddings. It introduces an efficient inexact, variance-reduced projected SGD algorithm that supports partial client participation and provides a convergence guarantee to an $\epsilon$-stationary neighborhood with rate $O(1/K)$ for smooth, potentially non-convex losses. Theoretical results connect ideal model distances to distributional Wasserstein distances, justifying the constraint-based personalization, and empirical studies on synthetic data and Federated MNIST demonstrate superior performance and interpretable learned dissimilarities. KARULA thus offers a scalable, data-driven approach to cross-client personalization that respects local data nuances while leveraging distributional similarity to improve generalization in heterogeneous FL environments.
Abstract
One of the defining challenges in federated learning is that of statistical heterogeneity among clients. We address this problem with KARULA, a regularized strategy for personalized federated learning, which constrains the pairwise model dissimilarities between clients based on the difference in their distributions, as measured by a surrogate for the 1-Wasserstein distance adapted for the federated setting. This allows the strategy to adapt to highly complex interrelations between clients, that e.g., clustered approaches fail to capture. We propose an inexact projected stochastic gradient algorithm to solve the constrained problem that the strategy defines, and show theoretically that it converges with smooth, possibly non-convex losses to a neighborhood of a stationary point with rate O(1/K). We demonstrate the effectiveness of KARULA on synthetic and real federated data sets.