Personalized Federated Learning via Variance-Aware Nonparametric Empirical Bayes

2026-08-09

AI Summary

This work addresses the degradation of personalized performance in federated learning caused by client heterogeneity, particularly when local estimation variances are unknown and parameter-dependent. To tackle this challenge, the authors propose the VANEB framework, which treats local parameters as noisy observations and leverages a nonparametric empirical Bayes approach to estimate a shared prior for personalization. A key innovation is the introduction of a variance-aware mechanism that explicitly handles heteroscedastic noise whose variance depends on the underlying parameters. The method integrates the generalized Tweedie formula with nonparametric maximum likelihood estimation and establishes non-asymptotic error bounds for density estimation along with an oracle denoising inequality. In deep neural network settings, the authors design two practical strategies—VANEB-head and VANEB-FT—which demonstrate significant performance gains over existing baselines on MNIST and CIFAR-10 benchmarks.

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.

Problem

• Personalized Federated Learning
• Nonparametric Empirical Bayes
• Heterogeneous Clients
• Unknown Variances
• M-estimation

Innovation

• Personalized Federated Learning
• Nonparametric Empirical Bayes
• Variance-Aware Estimation
• Heteroskedastic Tweedie's Formula
• M-estimation