Personalized Federated Learning With Gaussian Processes

1 Introduction . In recent years , there is a growing interest in applying learning in decentralized systems under the setup of federated learning ( FL ) [ 37 , 51 , 66 ] . In FL , a server node stores a global model and connects to multiple end-devices ( “ clients '' ) , which have private data that can not be shared . The goal is to learn the global model in a communication-efficient manner . However , learning a single shared model across all clients may perform poorly when the data distribution varies significantly across clients . Personalized Federated Learning ( PFL ) [ 67 ] addresses this challenge by jointly learning a personalized model for each client . While significant progress had been made in recent years , leading approaches still struggle in realistic scenarios . First , when the amount of data per client is limited , even though this is one of the original motivations behind federated learning [ 4 , 51 , 72 ] . Second , when the input distribution shifts between clients , which is often the case , as clients use different devices and sensors . Last , when we require well-calibrated predictions , which is an important demand from medical and other safety-critical applications . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Here , we show how Gaussian Processes ( GPs ) with deep kernel learning ( DKL ) [ 80 ] is an effective alternative for handling these challenges . GPs have good predictive performance in a wide range of dataset sizes [ 2 , 81 ] , they are robust to input noise [ 75 ] , can adapt to shifts in the data distribution [ 48 ] , and provide well-calibrated predictions [ 69 ] . While regression tasks are more natural for GPs , here we focus on classification tasks for consistency with common benchmarks and learning procedures in the field ; however , our approach is also applicable to regression tasks . Consider a naive approach that fits a separate GP classifier to each client based on its personal data . Its performance heavily depends on the quality of the kernel , and standard kernels tend to work poorly in domains such as images . A popular solution to this problem is to use deep kernel learning ( DKL ) [ 80 ] , where a kernel is applied to features outputted by a neural network ( NN ) . Unfortunately , GPs with DKL can strongly overfit , often even worse than standard NNs [ 56 ] , and thus negate the main benefit of using a GP . We solve this issue by jointly learning a shared kernel function across clients . As the kernel captures similarities between inputs , a single kernel should work well across clients , while using a separate GP per client will give the required flexibility for personalization . We adapt a GP classifier recently proposed in [ 2 ] which uses the Pólya-Gamma augmentation [ 57 ] in a tree-structure model to the federated setting . We term our method pFedGP . We extend pFedGP by tailoring two inducing points ( IPs ) methods [ 58 , 70 ] . The first helps generalization in the low data regime and , unlike common inducing point methods , does not reduce the computational costs . The second does focus on reducing the computational cost to make our approach scalable and work in low-resource clients . We also adjust previous PAC-Bayes generalization bounds for GPs [ 60 , 64 ] to include the Pólya-Gamma augmentation scheme . These bounds are suitable for cases where the kernel is not learned , such as when new clients arrive after the shared NN was already learned . Therefore , this paper makes the following contributions : ( i ) introduce pFedGP as a natural solution to PFL ; ( ii ) develop two IP methods to enhance GP classifiers that use the Pólya-Gamma augmentation scheme and integrate them with pFedGP ; ( iii ) derive a PAC-Bayes generalization bound on novel clients and show empirically that it gives meaningful guarantees ; ( iv ) achieve state-of-the-art results in a wide array of experiments , improving accuracy by up to 21 % 1 . 2 Related work . Federated learning . In FL , clients collaboratively solve a learning task while preserving data privacy and maintaining communication efficiency [ 1 , 34 , 42 , 51 , 54 , 82 , 85 ] . FedAvg [ 51 ] is an early but effective FL approach that updates models locally and averages them into a global model . Several optimization methods have been proposed for improving convergence in FL [ 41 , 45 , 71 , 77 ] . Other approaches focus on preserving client privacy [ 3 , 20 , 52 , 90 ] , improving robustness to statistical diversity [ 26 , 27 , 31 , 35 , 87 , 88 ] , and reducing communication cost [ 13 , 61 ] . These methods aim to learn a global model across clients , limiting their ability to deal with heterogeneous datasets . 1Our code is publicly available at https : //github.com/IdanAchituve/pFedGP Personalized federated learning . To overcome client heterogeneity , PFL aims to introduce some personalization for each client in the federation [ 39 , 73 ] . Recent methods include adapting multitask learning [ 18 , 67 ] , meta-learning approaches [ 6 , 21 , 22 , 33 , 43 , 89 ] , and model mixing , where clients learn a mixture of the global and local models [ 4 , 17 , 27 , 44 ] . Other approaches utilize different regularization schemes to enforce soft parameter sharing [ 32 , 72 ] . Personalization in FL has also been explored through clustering approaches in which similar clients within the federation have a greater effect on one another [ 49 , 86 ] . Recently , [ 65 ] proposed learning a central hypernetworks that acts on client representation vectors for generating personalized models . Bayesian FL . Some studies put forward a Bayesian treatment to the FL setup . [ 8 , 12 ] used variational inference with Bayesian NNs . [ 76 , 84 ] proposed a matching algorithm between local models based on the Beta-Bernoulli process to construct a global model . [ 14 ] extended Bayesian optimization to FL setting via Thompson sampling . To scale the GP model they used random Fourier features . We use inducing points instead . [ 83 ] proposed a federated learning framework that uses a global GP model for regression tasks and without DKL . Unlike this study , we focus on classification tasks with a personal GP classifier per client and advocate sharing information between clients through the kernel . [ 74 ] used GPs in a client selection strategy . In [ 36 ] an approach based on stein variational gradient descent was suggested . This method does not scale beyond small-sized networks . [ 47 ] proposed a multivariate Gaussian product mechanism to aggregate local models . As we will show , this method is less suited when the data heterogeneity between clients is large . Gaussian process classification . Unlike regression , in classification approximations must be used since the likelihood is not a Gaussian [ 59 ] . Classic approaches include the Laplace approximation [ 79 ] , expectation-propagation [ 53 ] , and least squares [ 62 ] . Recently , several methods were proposed based on the Pólya-Gamma augmentation [ 57 ] for modeling multinomial distributions [ 46 ] , GP classification [ 23 , 24 , 78 ] , few-shot learning [ 69 ] , and incremental learning [ 2 ] . Here we build on the last approach . Classification with GPs is commonly done with variational inference techniques [ 30 ] , here we wish to exploit the conjugacy of the model to take Gibbs samples from the posterior . This approach yields well calibrated [ 69 ] and more accurate models [ 2 ] . 3 Gaussian processes background . We first provide a brief introduction to the main components of our model . Detailed explanations are deferred to the Appendix . Scalars are denoted with lower-case letters ( e.g. , x ) , vectors with bold lowercase letters ( e.g. , x ) , and matrices with bold capital letters ( e.g. , X ) . In general , y = [ y1 , ... , yN ] T is the vector of labels , and X ∈ RN×d is the design matrix with N data points whose ith row is xi . Gaussian processes . GPs map input points to target output values via a random latent function f . f is assumed to follow a Gaussian process prior f ∼ GP ( m ( x ) , k ( x , x′ ) ) , where the evaluation vector of f on X , f = [ f ( x1 ) , ... , f ( xN ) ] T , has a Gaussian distribution f ∼ N ( µ , K ) with means µi = m ( xi ) and covariance Kij = k ( xi , xj ) . The mean m ( x ) is often set to be the constant zero function , and the kernel k ( x , x′ ) is a positive semi-definite function . The target values are assumed to be independent when conditioned on f . For Gaussian process regression the likelihood is Gaussian , p ( y|f ) = N ( f , σ2 ) . Therefore , the posterior p ( f |y , X ) is also Gaussian , and both the marginal and the predictive distributions have known analytic expressions . This is one of the main motivations behind using GPs , as most other Bayesian models have intractable inference . Unfortunately , for Gaussian process classification ( GPC ) the likelihood , p ( y|f ) , is not a Gaussian and the posterior does not admit a closed-form expression . One approach for applying GPs to binary classification tasks is the Pólya-Gamma augmentation [ 57 ] . Using this approach , we can augment the GP model with random variables ω from a Pólya-Gamma distribution , one for each example . As a result , p ( f |y , X , ω ) is a Gaussian density and p ( ω|y , X , f ) is a Pólya-Gamma density . This allows to use Gibbs sampling to efficiently sample from the posterior p ( f , ω|y , X ) for inference and prediction . A key advantage of the Pólya-Gamma augmentation is that it benefits from fast mixing and has the ability of even a single value of ω to capture much of the volume of the marginal distribution over function values [ 46 ] . Full equations and further details on the Pólya-Gamma augmentation scheme are given in Appendix A.1 . Deep kernel learning ( DKL ) . The quality of the GP model heavily depends on the kernel function k ( xi , xj ) . For many data modalities , such as images , common kernels are not a good measure of semantic similarity . Therefore , in [ 10 , 80 ] standard kernels are used over features outputted by a neural network gθ . For example , the RBF kernel kθ ( xi , xj ) = exp ( − ||gθ ( xi ) −gθ ( xj ) || 2 2 ` 2 ) . In regression , it is possible to directly backpropagate through the GP inference as it is given in closed-form . In our case , we use Fisher ’ s identity [ 19 ] to obtain stochastic gradients [ 69 ] . Inducing points . GPs require storing and inverting a kernel matrix on the entire training set which often limits its usage . A common solution to this problem is to use inducing point methods [ 58 , 70 ] . The key idea is to replace the exact kernel with an approximation for fast computation . Usually , M N pseudo-inputs are learned such that the main computational bottleneck is in inverting M ×M matrices . GP-Tree . We build on GP-Tree [ 2 ] , a recent GP classifier that was shown to scale well with dataset size and the number of classes . GP-Tree turns the multi-class classification problem into a sequence of binary decisions along the tree nodes . Each node in the tree fits a binary GP classifier based on the Pólya-Gamma augmentation scheme and the data associated with that node . The leaf nodes correspond to the classes in the dataset . The tree is constructed by first computing a prototype for each class and then recursively performing divisive hierarchical clustering on these prototypes to two clusters at each node . Further details are given in Appendix A.2 .

Summary: This paper presents a new framework for personalized federated learning via Gaussian process classification (GPC) model. The key idea is to have each client node maintain its own GPC model but share the NN feature map f(x). Here, the GPC model of client c is equipped with kernel k_c(x, x') = k_c( f(x), f(x'); w_c)  with personalized parameters w_c.   Thus, per iteration, a subset of clients is selected. Each client c downloads the shared f(x) from server and re-build its GPC model (i.e., re-fitting for w_c) using local data D_c. A local updated feature map f(x) (via gradient propagation) is then computed for each client and communicated to the server, which averages the updated model to form a new feature map f(x).  The proposed framework is evaluated on 3 image datasets: CIFAR-10, CIFAR-100 and CINIC-10. The reported performance shows that it outperforms a selected set of baselines substantially when the amount of data per client is small & heterogeneous.