PaperlayerView exact revision on arXiv

Source-linked AI summary

Personalized Federated Learning with Gaussian Processes

Idan Achituve, Aviv Shamsian, Aviv Navon, Gal Chechik, Ethan Fetaya

arXiv:2106.15482v2cs.LGstat.ML

TL;DR

Personalized federated learning must handle heterogeneous clients with limited local data while retaining accurate uncertainty estimates. pFedGP shares a deep kernel across clients but fits personal Gaussian-process classifiers, adding inducing-point variants for sample efficiency and computation. It reports well-calibrated predictions, non-vacuous novel-client guarantees, and up to 21% accuracy gains over competing methods.

  • Problem

    Personalized federated learning must support heterogeneous clients with limited data while providing well-calibrated predictions.

  • Method

    pFedGP learns a shared neural-network kernel with client-specific Gaussian-process classifiers and adds inducing-point variants for low-data generalization and computational efficiency.

  • Results

    pFedGP consistently outperforms competing methods, provides well-calibrated predictions, and achieves up to 21% accuracy improvement.

  • Takeaways & Limitations

    The method is particularly useful for clients with little data, strongly different distributions, or out-of-distribution novel clients.

  • Takeaways & Limitations

    The federated communication protocol does not explicitly guarantee that private information cannot be inferred.

Abstract

from arXiv · show

Federated learning aims to learn a global model that performs well on client devices with limited cross-client communication. Personalized federated learning (PFL) further extends this setup to handle data heterogeneity between clients by learning personalized models. A key challenge in this setting is to learn effectively across clients even though each client has unique data that is often limited in size. Here we present pFedGP, a solution to PFL that is based on Gaussian processes (GPs) with deep kernel learning. GPs are highly expressive models that work well in the low data regime due to their Bayesian nature. However, applying GPs to PFL raises multiple challenges. Mainly, GPs performance depends heavily on access to a good kernel function, and learning a kernel requires a large training set. Therefore, we propose learning a shared kernel function across all clients, parameterized by a neural network, with a personal GP classifier for each client. We further extend pFedGP to include inducing points using two novel methods, the first helps to improve generalization in the low data regime and the second reduces the computational cost. We derive a PAC-Bayes generalization bound on novel clients and empirically show that it gives non-vacuous guarantees. Extensive experiments on standard PFL benchmarks with CIFAR-10, CIFAR-100, and CINIC-10, and on a new setup of learning under input noise show that pFedGP achieves well-calibrated predictions while significantly outperforming baseline methods, reaching up to 21% in accuracy gain.

1 Introduction

Personalized federated learning addresses heterogeneous clients, limited local data, input shifts, and the need for calibrated predictions. pFedGP shares a neural-network kernel while retaining client-specific Gaussian-process classifiers, with inducing-point variants and PAC-Bayes guarantees.

  • Motivation: Personalized federated learning targets settings where a single shared model performs poorly because client data distributions vary significantly.The motivation includes limited data per client, input-distribution shifts, and well-calibrated predictions for safety-critical applications.
  • pFedGP: pFedGP jointly learns a shared neural-network parameterized kernel and a personal Gaussian-process classifier for each client.This design addresses the dependence of GP performance on kernel quality while sharing representation learning across clients.
  • Inducing-point variants: pFedGP extends the method with two inducing-point variants: one improves low-data generalization, while the other reduces computational cost.The variants are designed for complementary constraints in personalized federated learning.
  • Guarantees and evaluation: pFedGP derives a PAC-Bayes generalization bound for novel clients and reports meaningful empirical guarantees.The contribution also includes extensive evaluation across personalized federated-learning settings.
  • Results: Up to 21% accuracy improvement is reported across the paper’s experiments.The contribution statement describes state-of-the-art results across a wide array of experiments.

2 Related work

Related work spans federated optimization, personalization, Bayesian federated learning, and Gaussian-process classification. These approaches address communication, heterogeneity, uncertainty, or scalability, while pFedGP combines personalized GP classification with shared deep kernels.

  • Federated learning: Federated-learning methods commonly optimize a global model through local updates, but this limits their ability to handle heterogeneous datasets.FedAvg averages locally updated models, while other methods target convergence, privacy, robustness, or communication cost.
  • Personalized federated learning: Personalized federated learning introduces client-specific models through multitask learning, meta-learning, model mixing, regularization, and clustering approaches.These methods address client heterogeneity by varying how parameters or models are shared.
  • Bayesian federated learning: Bayesian federated-learning studies include Bayesian neural networks, matching-based local models, Thompson sampling, and global Gaussian processes.The cited GP framework uses random Fourier features and addresses regression without deep kernel learning, unlike this study’s personal GP classifiers.
  • Gaussian-process classification: Gaussian-process classification requires approximations because its non-Gaussian likelihood prevents a closed-form posterior.Prior approaches include Laplace, expectation-propagation, least-squares, and Pólya-Gamma-based methods.

3 Gaussian processes background

The paper’s background covers Gaussian-process priors, approximate classification inference, deep kernel learning, inducing points, and GP-Tree classification. These components motivate a scalable GP classifier with learned representations and federated personalization.

  • Gaussian processes: A Gaussian process places a distribution over latent functions through a mean function and positive-semidefinite kernel, yielding Gaussian-distributed function evaluations.The kernel determines covariance among input points and therefore strongly influences the model.
  • Gaussian-process classification: Gaussian-process classification lacks closed-form posterior inference because its likelihood is not Gaussian.Pólya-Gamma augmentation restores Gaussian conditional densities and enables Gibbs sampling for inference and prediction.
  • Deep kernel learning: Deep kernel learning applies a standard kernel to neural-network features because common kernels may not measure semantic similarity well for images.The neural network therefore parameterizes a learned representation used by the kernel.
  • Inducing points: Inducing-point methods approximate the full kernel using M ≪ N pseudo-inputs, making the main bottleneck inversion of M × M matrices.This reduces computation relative to storing and inverting a kernel matrix over the entire training set.
  • GP-Tree: GP-Tree converts multiclass classification into binary GP decisions along tree nodes, with leaves representing dataset classes.Each node uses Pólya-Gamma augmentation and the data associated with that node.

4 pFedGP: federated learning with Gaussian processes

pFedGP extends GP-Tree to personalized federated learning by sharing neural-network kernel parameters while fitting client-specific GP models. Two global inducing-point variants address low-data performance and computational scalability.

  • Core pFedGP: pFedGP maps each client’s data through shared neural-network features, then builds a client-specific hierarchical GP classifier using those representations.The tree is rebuilt each round from current features, and posterior-sample gradients update the shared network.
  • Federated optimization: Federated training samples clients, performs local kernel learning, and averages updated shared parameters with the FedAvg rule.The server holds shared neural-network parameters, while each client retains its local dataset.
  • Sample efficiency: Global inducing points are shared across clients to increase information sharing when each client has very few data points.They are defined in the shared network’s final embedding space and sent with the model from the server to clients.
  • Sample efficiency: The inducing-point classifier corrects predictions for distorted class frequencies caused by evenly allocating inducing inputs across classes.The correction maps predictions back toward the original class ratios at each tree node.
  • Computational efficiency: The computational variant uses FITC and globally defined inducing inputs so that inference requires inversion of M × M or diagonal matrices.This procedure is intended to reduce complexity in low-resource environments and scale to larger datasets.

5 Generalization bound

The paper derives PAC-Bayes generalization bounds for pFedGP on novel clients and evaluates whether they provide meaningful estimates of test error.

  • 5 Generalization bound: The bound applies to a GP classifier on a novel client while treating the learned neural-network kernel as fixed.This targets deployment scenarios where new clients arrive after kernel training, without retraining the kernel function.
  • 5 Generalization bound: The PAC-Bayes analysis uses the KL divergence between posterior and prior Gaussian processes, equivalently evaluated on their training-sample values.With Pólya-Gamma augmentation, the resulting non-Gaussian posterior requires estimating the KL term rather than using a closed-form expression directly.
  • 5 Generalization bound: The bound was tested on CIFAR-10 using 90 training clients and 10 novel clients across four training-set sizes.The experiment evaluates both Gibbs-classifier generalization and test error under δ = 0.01.
  • 5 Generalization bound: The estimated bound exceeded actual test error for every evaluated point, remained non-vacuous, and strongly correlated with the observed error.Unlike worst-case bounds based only on sample counts, it depends on the actual data.

6 Experiments

Experiments evaluate pFedGP and its inducing-point variants across standard personalized federated-learning benchmarks, calibration, and client-specific input noise. pFedGP improves over competing methods especially when local data are limited, while retaining strong calibration and robustness across heterogeneous settings.

  • 6 Experiments: The experiments compare full pFedGP, pFedGP-IP-data, and pFedGP-IP-compute against local, FedAvg, FOLA, LG-FedAvg, pFedMe, FedU, and pFedHN baselines.Results use the federated-accuracy protocol and report averages with standard errors over three random seeds unless otherwise stated.
  • 6.1 Standard PFL setting: The standard evaluation varies client counts from 50 to 500 across CIFAR-10, CIFAR-100, and CINIC-10, with the number of samples per client decreasing as client count increases.Each client sees the same classes in its training and test sets, while the total number of system samples remains fixed.
  • 6.1 Standard PFL setting: Up to 21% accuracy improvement over competing methods is reported when local data are limited, including 9% and 21% gains on CIFAR-100 with 100 and 500 clients.pFedGP is competitive or comparable to prior state-of-the-art methods when local data are sufficient, and the pFedGP-IP-compute variant is often comparable to or better than pFedGP-IP-data.
  • 6.1 Standard PFL setting: pFedGP classifiers are best calibrated across nearly all ECE, MCE, and Brier-score comparisons on CIFAR-100 with 50 clients.Baseline calibration can improve with temperature scaling, but that requires a separate validation set whereas pFedGP does not.
  • 6.2 PFL with input noise: The pFedGP variants show a significant accuracy advantage over baseline methods under client-specific input noise, although globally shared inducing points are less beneficial when noise types differ by client.The noisy-CIFAR evaluation assigns distinct corruption-noise distributions to clients.

7 Conclusion

The paper concludes that pFedGP combines a shared neural-network kernel with personal GP classifiers, while supporting calibrated predictions, novel-client generalization, and strong performance. It also relies on Pólya-Gamma augmentation and GP-Tree inference, with a PAC-Bayes bound providing non-vacuous guarantees.

  • pFedGP learns a neural-network-parameterized kernel shared across clients while retaining a personal GP classifier for each client.
  • The method provides well-calibrated predictions and generalizes well to out-of-distribution novel clients.
  • pFedGP consistently outperforms competing methods, with the paper reporting accuracy improvements of up to 21%.
  • The GP classifier uses Pólya-Gamma augmentation so the augmented likelihood becomes Gaussian and posterior sampling can use block Gibbs sampling.
  • GP-Tree extends binary Pólya-Gamma classification to multiple classes through a binary tree whose nodes receive GP models.

B pFedGP-IP-compute detailed derivation

The pFedGP-IP-compute derivation factorizes inference through shared inducing points, allowing client- and tree-node-specific quantities to be computed separately and aggregated afterward. Training uses block Gibbs sampling and Fisher’s identity, while cached quantities accelerate test-time inference.

  • pFedGP-IP-compute assumes conditional independence through shared inducing points, factorizing inference over clients and tree nodes before aggregation.
  • The inducing-point model represents the joint GP prior over data and inducing latent values, then factors the likelihood using Pólya-Gamma variables.
  • Posterior inference obtains the inducing-variable distribution by Bayes’ rule and the data-function posterior by marginalizing over inducing values.
  • Predictions for a novel input use the posterior over inducing values, while the full data-function values remain necessary for sampling Pólya-Gamma variables.
  • Model parameters are learned with block Gibbs sampling and Fisher’s identity, and test-time inference caches input-independent quantities such as ω, Λ, and a Cholesky decomposition.

C Generalization bound - derivation

The bound derivation expresses the PAC-Bayes KL term between Gaussian GP posterior and prior distributions in closed form. Figure 5 compares estimated upper bounds with test error across clients and training-set sizes.

  • The KL-divergence term compares the Gaussian prior P(f) = N(µ, K) with the Gaussian posterior Q(f|ω) = N(Σ(K−1µ + κ), Σ).
  • The derivation yields a closed-form KL expression involving determinant ratios, a trace term, and the posterior-mean displacement under K−1.
  • Figure 5 plots test error against an estimated upper bound for 10 CIFAR-10 clients across varying training-set sizes, with dots representing client–size combinations.
  • The figure reports the average difference between empirical and test error in parentheses for each plotted setting.

D Experimental details

The experiments cover CIFAR-10, CIFAR-100, and CINIC-10 under heterogeneous client partitions, with validation-based hyperparameter selection. Additional analyses vary client data size, model architecture, and input corruption.

  • Data-size analysis: Figure 6 measures model performance over 50 CIFAR-10 clients as the average number of training samples per client changes on a logarithmic x-axis.
  • Datasets: pFedGP is evaluated on CIFAR-10, CIFAR-100, and the more diverse CINIC-10 dataset, which extends CIFAR-10 with down-sampled ImageNet images.
  • Data assignment: Client partitions assign each client a varying number of samples and a unique set of classes, typically using 2, 10, and 4 classes for CIFAR-10, CIFAR-100, and CINIC-10.
  • Hyperparameter tuning: Hyperparameters and early stopping are selected with validation sets, using configurations tuned on the 50-client learning setup.
  • Architecture analysis: Figure 7 compares test error and estimated upper bounds for ResNet-18 and MobileNetV2 across clients and training-set sizes.
  • Noisy datasets: The noise experiments apply 19 corruption types at severities 3, 4, and 5, producing 57 unique noise models for CIFAR-10 and CIFAR-100.

E.1 Generalization bound - additional experiments

Additional experiments evaluate the PAC-Bayes generalization bound with deterministic Bayes classifiers and larger neural networks. The bound follows similar behavior across classifiers and backbones, while Bayes classifiers generally outperform Gibbs classifiers.

  • The Bayes-risk analysis extends the bound evaluation from Gibbs risk to deterministic predictors closer to the estimated label.The Bayes risk is used to assess error for a classifier closer to the Gauss-Hermite estimate of y∗.
  • Bayes classifiers generally outperform Gibbs classifiers when estimating generalization error on novel clients.Figure 5 compares the estimated bound with actual error and shows patterns similar to the Gibbs-risk analysis.
  • The generalization-bound behavior remains similar when evaluated with larger ResNet18 and MobileNetV2 backbones.These networks contain approximately 11.4M and 2.8M parameters, respectively.

E.2 Varying the training set size

Across varying data availability, client heterogeneity, homogeneous federated learning, computational settings, and calibration tests, pFedGP variants retain strong accuracy and calibration while addressing low-data and efficiency constraints.

  • E.2 Varying the training set size: 77.7% accuracy with only 1000 training examples shows that pFedGP-IP-data remains effective in the extremely low-data regime.As training data decreases, all methods lose accuracy, but pFedGP methods degrade less severely than baselines; LG-FedAvg was excluded for low accuracy.
  • E.2 Varying the training set size: Sharing inducing inputs improves low-data performance by effectively increasing the training data available per client.The pFedGP-IP-data variant shares inducing inputs across clients, and its test-time use of personal data is especially effective when each client has limited data.
  • E.3 PFL with input noise under a homogeneous class distribution: pFedGP and its variants surpass all baseline methods by a large margin under heterogeneous input noise with balanced class distributions.This result holds on noisy CIFAR-10 with 100 clients, while FedAvg outperforms competing non-pFedGP methods in the balanced-class setting.
  • E.4 Homogeneous federated learning with CIFAR-10: pFedGP outperforms all PFL baselines by a large margin in homogeneous CIFAR-10 federated learning, where FedAvg also performs well.The setup distributes all classes equally across clients; pFedGP can model absent classes using an inducing-point variant.
  • Computational requirements: pFedGP has reasonable computational requirements for current federated-learning systems, using approximately 1.4/1.6 GB memory and 1/2 hours on CIFAR-10/100.The computational comparison also evaluates reduced-cost inducing-point variants and predictive-posterior inference runtime.
  • E.9 Reliability diagrams: pFedGP is calibrated by design and does not benefit from temperature scaling, unlike baseline methods that require validation-based temperature selection.Reliability diagrams cover CIFAR-100 with 50 and 100 clients under default and optimized temperatures.
Loading 2106.15482v2…