Personalized Federated Learning via Variance-Aware Nonparametric Empirical Bayes
2026-08-10 • Machine Learning
Machine Learning
AI summaryⓘ
The authors present a new method for personalized federated learning, where multiple devices train models that can be different but share some common traits. Their approach uses a statistical technique called Nonparametric Empirical Bayes (NPEB) to estimate a shared prior distribution from noisy local parameter estimates. To handle the fact that variances in these estimates vary across devices, they develop a new framework called Variance-Aware Nonparametric Empirical Bayes (VANEB) with theoretical error guarantees. They also adapt their method to modern deep neural networks and demonstrate good results on MNIST and CIFAR-10 image datasets.
Personalized Federated LearningNonparametric Empirical BayesEmpirical Risk MinimizationM-estimationHeteroskedasticityTweedie's FormulaDeep Neural NetworksMNISTCIFAR-10
Authors
Jae Ho Chang, Arnab Auddy, Subhadeep Paul
Abstract
We develop a new approach to Personalized Federated Learning across heterogeneous clients using Nonparametric Empirical Bayes (NPEB). Leveraging the asymptotic normality of local parameter estimates obtained from Empirical Risk Minimization or M-estimation, our method formulates these estimates as noisy observations to estimate an unknown shared prior via Nonparametric Maximum Likelihood. A key challenge in applying NPEB in this setting is that existing approaches assume known fixed variances, which is not true in practice. To address this, we introduce a Variance-Aware Nonparametric Empirical Bayes (VANEB) framework that leverages the parameter-dependent asymptotic variance of local M-estimators. A key technical contribution is a generalized Tweedie's formula for this heteroskedastic setting. We then establish non-asymptotic error rates for density estimation in the average squared Hellinger distance and derive an oracle denoising inequality that provides error bounds for our estimator. While our theoretical guarantees are rooted in the asymptotic regime of M-estimators, we empirically explore heuristic extensions of VANEB to modern federated learning settings involving Deep Neural Networks (DNNs). For DNNs, we propose VANEB-head and VANEB-FT, which personalize the last fully connected layer via an NPEB step using an approximate diagonal variance estimator. We show that our method has strong performance on popular vision datasets MNIST and CIFAR-10, using a convolutional neural network architecture.