Adaptive Personalized Federated Learning via Multi-task Averaging of Kernel Mean Embeddings
This work addresses the challenge of data heterogeneity in federated learning for distributed agents, offering an adaptive solution that can transition between global and local learning regimes, though it is incremental in building on existing PFL methods.
The paper tackles the problem of personalized federated learning by proposing a method where agents collaboratively learn individual models using data-driven weights, derived from kernel mean embeddings and multi-task averaging, without requiring prior knowledge of data heterogeneity. The result includes finite-sample guarantees on local excess risks and a practical implementation with random Fourier features to manage communication costs.
Personalized Federated Learning (PFL) enables a collection of agents to collaboratively learn individual models without sharing raw data. We propose a new PFL approach in which each agent optimizes a weighted combination of all agents' empirical risks, with the weights learned from data rather than specified a priori. The novelty of our method lies in formulating the estimation of these collaborative weights as a kernel mean embedding estimation problem with multiple data sources, leveraging tools from multi-task averaging to capture statistical relationships between agents. This perspective yields a fully adaptive procedure that requires no prior knowledge of data heterogeneity and can automatically transition between global and local learning regimes. By recasting the objective as a high-dimensional mean estimation problem, we derive finite-sample guarantees on local excess risks for a broad class of distributions, explicitly quantifying the statistical gains of collaboration. To address communication constraints inherent to federated settings, we also propose a practical implementation based on random Fourier features, which allows one to trade communication cost for statistical efficiency. Numerical experiments validate our theoretical results.