Personalized Federated Learning via Variance-Aware Nonparametric Empirical Bayes

This paper introduces the Variance-Aware Nonparametric Empirical Bayes (VANEB) framework for personalized federated learning (PFL). The authors, Jae Ho Chang, Arnab Auddy, and Subhadeep Paul from the Department of Statistics at The Ohio State University, develop a new approach that casts personalized federated learning as a compound decision problem governed by empirical Bayes principles.

The paper addresses a fundamental limitation in existing Nonparametric Empirical Bayes (NPEB) theory and algorithms: "A fundamental limitation of existing NPEB theory and algorithms [Jiang and Zhang, 2009, Saha and Guntuboyina, 2020, Soloff et al., 2025] is their reliance on the assumption of known and parameter-independent covariance structures."

In the context of federated learning, this premise is frequently violated because "for M-estimators and generalized linear models (GLMs), the covariance of a client's local estimator is unknown and is often a function of the client parameter that is being estimated (e.g., the inverse of the Fisher information or a 'sandwich' formula)."

The authors formulate the problem as follows: Let θ1,..., θK ∈ Rd denote latent client parameters drawn i.i.d. from an unknown measure G0 on Rd, and suppose client k reports a local estimator θ̂k. They adopt a heteroskedastic Gaussian working model:

√nk(θ̂k − θk) θk ∼ N(0d, Σk(θk)), k = 1,..., K

where Σk(θ) is a known covariance function that depends on the parameter θ. The model captures differences across clients in both sample size and covariance structure.

The framework considers two distinct sources of heterogeneity: (1) cross-client heterogeneity in the latent parameters θk, encoded in the unknown prior G0, and (2) heteroskedasticity within each client's model where the local covariance Σk(θ) depends on the parameter θ itself.

The authors derive a generalized Tweedie's formula for parameter-dependent covariances. The oracle Bayes estimator is given by:

θko = θ̂k/√nk + qk,G0(θ̂k)/(√nk · fk,G0(θ̂k))

where qk,G0 involves variance-weighted marginal densities: qk,G0(θ̂k):= ∫ Σk(θ)∇φ(k)(θ̂k; θ) G0(dθ)

The key insight is that the correction term now involves variance-weighted densities f̃k,G,i, whose coordinate-wise weights σk,i(θ)2 make the shrinkage adaptive to the local noise level of each client.

The NPMLE is defined as a solution to: Ĝ ∈ arg maxG∈P(Rd) lK(G), lK(G):= (1/K) ΣKk=1 log fk,G(θ̂k)

The authors note that when Σk(θ) depends on θ, the score of the mixture density... involves the gradients of the variance functions σk,i, and no closed-form solution exists for the critical point(s).

Algorithm 1 addresses the intractability by "alternating between (i) solving a fixed-covariance NPMLE subproblem, where substituting Sk = Σk(aj) recovers the tractable ridgeline property, and (ii) updating the covariances evaluated at the newly estimated atom locations."

Hellinger Accuracy of the NPMLE (Theorem 4.1): The authors establish non-asymptotic convergence rates showing that any approximate NPMLE satisfying the likelihood gap condition achieves:

P[h̄2(f•,Ĝ, f•,G0) ≳d t2ε2M(S, G0)] ≤ 2exp(−t2nd/2 logd K)

Oracle Denoising Inequality (Theorem 4.2): The excess risk of the VANEB estimator relative to the oracle posterior mean is bounded by:

(1/K)ΣKk=1 E∥θ̂kv − θko∥2 ≤d,s,s̄ ((log K)3∨(d/2−1)/n) · ε2M(S•, G0)

For special cases:

• Discrete support: Risk scales as (k*/Kn)(log K)d+(d/2∨4)

• Compact support: Risk scales as (Vol(S1)/Kn)(log K)d+(d/2∨4)

• Gaussian Mixture Model: Risk scales as (Vol(Sr)/Kn)(log K)d+(d/2∨4)

van Trees Lower Bound (Theorem 4.3): The oracle baseline satisfies:

d/(∥J0∥+n̄/s + dL2Σ/2) ≤ R(Θo) ≤ ds̄/n

The authors demonstrate that the gain from federation enters entirely through learning the shared prior since local M-estimators have excess risk of order 1/n regardless of K, while VANEB achieves Õ((Kn)−1).

The authors evaluate VANEB in three heterogeneous federated learning scenarios in R3 with a complex five-curve prior (trefoil knot, helix, tilted ellipse, figure-eight, and Viviani curve):

1. Quadratic variance scenario: VANEB improves steadily with an increasing number of clients and consistently outperforms fixed-variance EB once K is moderate and strongly outperforms NPEB in the small sample regime.

2. Multiclass logistic regression: VANEB substantially outperforms NPEB, AdaMix, and FedAvg in the hardest regimes with small client sample sizes though at K = 3200, VANEB has the best performance among the compared methods.

3. Poisson regression: VANEB is consistently the strongest non-oracle method and maintains a strict advantage over NPEB for every K.

The authors extend VANEB to deep neural networks through two variants:

• VANEB-Head: applies the EB update intermittently during federated training (every five communication rounds in our experiments), while the backbone is updated in every round.

• VANEB-FT: first trains a common model by FedAvg, then fine-tunes each client's head with the backbone fixed, and finally applies one EB update across all locally adapted heads.

For DNNs, clients report a diagonal covariance estimate based on an approximation to the observed Fisher information matrix using the gradient outer product.

The authors evaluate on MNIST and CIFAR-10 with joint label and covariate heterogeneity, comparing against 10 state-of-the-art methods including FedAvg, FedAvg+FT, FedRep, FedBABU, FedPer, LG-FedAvg, Per-FedAvg, SCAFFOLD, Ditto, and local-only training.

Key findings:

• VANEB-FT is particularly strong under balanced evaluation: it achieves the highest accuracy across all 13 reported MNIST and CIFAR-10 settings.

• It is the only personalized method that outperforms FedAvg across all balanced and matched MNIST and CIFAR-10 experiments.

• VANEB-Coordinate FT outperforms FedAvg in every reported comparison: all nine settings under both balanced and client-matched evaluation on both datasets (36 of 36 comparisons).

Under balanced evaluation (testing on the complete label space), the VANEB-FT variants give the strongest and most consistent balanced performance with either Direct or Coordinate FT is best in every balanced column.

The authors note that VANEB-FT Coordinate is not always the overall best-performing method, especially in client-matched tests where FedAvg with fine-tuning, LG-FedAvg, and related local-adaptation methods are often stronger.

The authors conclude: "Several directions remain open. First, our guarantees are restricted to the diagonal covariance approximation; extending them to full, non-diagonal Fisher information matrices is a natural next step. Second, bridging the gap between the theory and deep learning applications requires relaxing the asymptotic normality condition and addressing high dimensionality."

They suggest that A non-asymptotic analysis based on empirical process theory, controlling the concentration of the local empirical risks directly, could provide guarantees for VANEB in non-convex, finite-sample settings.

Based on this paper, I can improve AI systems in the following specific ways:

Improvement: Replace fixed-weight or heuristic personalization with a variance-aware empirical Bayes update that automatically adjusts shrinkage based on each client's local noise level and sample size.

Capability: The AI system can now automatically determine how much to trust each client's local model versus the global prior, dynamically balancing personalization and generalization without manual tuning. It handles clients with vastly different data volumes (e.g., 10 samples vs. 10,000 samples) by shrinking noisy estimates more aggressively while preserving signal from reliable ones.

Improvement: Incorporate parameter-dependent covariance structures (e.g., Fisher information for logistic/Poisson regression) into the federated learning objective, rather than assuming homoskedastic noise.

Capability: The system can now handle classification and count data correctly, where the variance of local estimators naturally depends on the parameter values (e.g., high uncertainty near decision boundaries). This yields more accurate personalization for tasks like medical diagnosis with rare classes or recommendation systems with skewed item popularity.

Improvement: Use the NPMLE approach to learn an arbitrary, data-driven prior distribution over client parameters, avoiding restrictive parametric assumptions (e.g., Gaussian priors).

Capability: The AI system can capture complex, multi-modal client heterogeneity—such as distinct user clusters with different behavioral patterns—without requiring the user to specify the number of clusters or prior family. It automatically discovers latent structure (e.g., trefoil knots, curves in parameter space) that parametric methods miss.

Improvement: Leverage the oracle denoising inequality and van Trees lower bound to certify that the system's personalization error decreases at rate Õ((Kn)(-1)) rather than O(n(-1)) for local-only methods.

Capability: The system can now provide formal performance guarantees: with K clients each having n samples, the personalized model achieves accuracy comparable to a centralized model with Kn samples, while preserving privacy. This enables deployment in safety-critical applications (e.g., personalized medicine) where worst-case error bounds are required.

Improvement: Implement the alternating optimization scheme that separates covariance estimation from NPMLE updates, making the computation tractable for high-dimensional parameters.

Capability: The system can now scale to models with thousands of parameters (e.g., neural network heads) while maintaining computational efficiency. The algorithm alternates between fixing covariances to solve a tractable subproblem and updating covariances at new atom locations, enabling practical deployment on resource-constrained edge devices.

Improvement: Apply the empirical Bayes update only to the final layer (head) after standard FedAvg training, with the backbone frozen.

Capability: The system achieves state-of-the-art personalization on vision tasks (MNIST, CIFAR-10) with minimal additional computation. It outperforms 10 existing methods (FedAvg, FedRep, Per-FedAvg, etc.) under balanced evaluation, and is the only method that consistently beats FedAvg across all settings—making it ideal for production systems where the backbone is expensive to retrain but head adaptation is cheap.

Improvement: Use diagonal covariance approximations with per-coordinate variance weighting, enabling feature-specific shrinkage.

Capability: The system can now handle scenarios where different features have vastly different noise levels (e.g., image pixels with varying illumination vs. stable structural features). It automatically applies stronger shrinkage to noisy coordinates while preserving informative ones, improving personalization quality in high-dimensional settings with heterogeneous feature reliability.

Improvement: Combine the federated framework with the empirical Bayes correction that only requires sharing local estimators (not raw data) and has provable excess risk bounds.

Capability: The system can now offer formal privacy-utility trade-offs: it achieves near-oracle performance (within (log K) 3/n of the optimal posterior mean) while only exchanging model parameters, not raw data. This enables deployment in regulated industries (healthcare, finance) where data sharing is restricted but personalization is critical.

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.