Adaptive Personalized Federated Learning via Multi-task Averaging of Kernel Mean Embeddings

Personalized Federated Learning enables a collection of agents to collaboratively learn individual models without sharing raw data. We propose a new 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 automatically transitions 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.

We consider a personalized federated learning framework in which multiple agents collaborate to learn their own models without sharing raw data. We focus on identifying, for each agent, which other agents provide useful data, namely data that are similar to its own while still enriching its local dataset. To this end, we learn weights between agents that capture these similarities directly from the data, without requiring prior knowledge about how heterogeneous the agents are. The originality of our work lies in connecting this problem to multiple mean estimation (multi-task averaging), which allows us to design a fully adaptive collaborative learning method with strong statistical guarantees. Once the weights are learned, the personalized models are trained using standard federated learning methods such as FedAvg. We also propose a practical implementation that reduces communication costs while maintaining good statistical performance, and we validate our approach through numerical experiments.