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

TL;DR

This paper introduces an adaptive personalized federated learning method that learns collaborative weights from data using kernel mean embeddings, enabling automatic adaptation to data heterogeneity without prior knowledge.

Contribution

It formulates weight estimation as a kernel mean embedding problem, providing a fully adaptive, data-driven approach with theoretical guarantees and a communication-efficient implementation.

Findings

Finite-sample guarantees on local excess risks.

Automatic transition between global and local learning regimes.

Effective communication-efficient implementation with random Fourier features.

Abstract

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.