Personalized Federated Learning with Gaussian Processes
Year: 2021

Abstract:
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.

heading: 1 Introduction
text: 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 cannot 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.

heading: 2 Related work
text: 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].

heading: 3 Gaussian processes background
text: 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.

heading: 4 pFedGP: federated learning with Gaussian processes
text: Now we describe our approach for applying personalized federated learning (PFL) with Gaussian processes. First, we extend GP-Tree to the FL setup and show how to use Gibbs sampling to learn the NN parameters. Then, we present two alternatives for this method that use inducing points. The first is for extremely limited-size datasets, while the second allows controlling the computational resources. We name our method pFedGP. An illustration of our method is given in Figure 1.

heading: 4.1 A full GP model
text: The training procedure follows the standard protocol in this field [4, 44, 51]. We assume the existence of a server that holds the shared parameters θ (a NN). Let C denote the set of clients. For each client c ∈ C we denote by Dc its local dataset of size Nc. At each training iteration (round) the model is sent to S clients to perform kernel learning (|S| ≤ |C|). Each client c ∈ S updates its copy of the global model and then sends the updated model to the server. The server then averages over the updates to obtain a new global model.
At each client c, we perform kernel learning in the following manner. We first compute the feature representation of the data samples associated with the client using the shared network. Then, we build the hierarchical classification tree as discussed in Section 3 & Appendix A.2. In [2] the tree was built only once after a pre-training stage and the model parameters were learned using a variational inference approach. Here, we re-build the tree at each round using the most recent features and we use a Gibbs sampling procedure, as it allows this flexibility in building the tree and performs better when not prohibitive by computational limitations. Learning the network parameters θ with the Gibbs sampling approach can be done with two common objectives, the marginal likelihood, and the predictive distribution.
We denote by Xv the data associated with the tree node v, i.e., the data points which have v on the path from the root node to their class leaf node. We denote by yv the binary label of these points, i.e., does their path go left or right at this node. And we denote by ωv the Pólya-Gamma random variables associated with node v. The marginal likelihood term for the full hierarchical classification tree is the sum of the separate marginal likelihood terms of all the nodes v in the tree:
LMLc (θ;Dc) = ∑ v log pθ(yv|Xv) = ∑ v log ∫ pθ(yv|ωv,Xv)p(ωv)dωv. (1)
Similar to [69] we use a gradient estimator based on Fisher’s identity [19]: ∇θLMLc (θ;Dc) = ∑ v ∫ pθ(ωv|yv,Xv)∇θlog pθ(yv|ωv,Xv)dωv ≈ ∑ v 1 L L∑ l=1 ∇θlog pθ(yv|ω(l)v ,Xv). (2) Here, ω(1)v , ...,ω (L) v are samples from the posterior at node v. Due to the Pólya-Gamma augmentation pθ(yv|ω(l)v ,Xv) is proportional to a Gaussian density. The exact expression is give in Appendix A.2.
To use the predictive distribution as an objective, in each training iteration, after building the tree model, at each node we randomly draw a portion from the (node) training data and use it to predict the class label for the remaining part. We denote with Xv and yv the training portion, x∗v and y ∗ v the input and the label of the point we are predicting, and P y ∗
the path from the root node to the y∗ leaf node (i.e., the original class). Here we also take advantage of the independence between nodes to maximize the predictive distribution per node individually. The predictive distribution for a single data point:
LPDc (θ;x∗, y∗) = ∑ v∈Py∗ log pθ(y ∗ v |x∗v,yv,Xv) = ∑ v∈Py∗ log ∫ pθ(y ∗ v |ωv,x∗v,yv,Xv)p(ωv|yv,Xv)dωv. (3) We use an approximate-gradient estimator based on posterior samples of ω:
∇θLPDc (θ; x∗, y∗) ≈ ∑ v∈Py∗ 1 L L∑ l=1 ∇θlog pθ(y∗v |ω(l)v ,x∗v,yv,Xv). (4)
Where pθ(y∗v |ω (l) v ,x∗v,yv,Xv) = ∫ p(y∗v |f∗)pθ(f∗|ω (l) v ,x∗v,yv,Xv)df
∗ does not have an analytical expression, but pθ(f∗|ω(l)v ,x∗v,yv,Xv) = ∫ pθ(f ∗|f ,x∗v,Xv)pθ(f |ω (l) v ,yv,Xv)df
∗ is Gaussian with known parameters. We then compute the predictive distribution by performing Gauss-Hermite integration over f∗. See exact expression in Appendix A.2.

heading: 4.2 Augmenting the model with inducing points: sample efficiency
text: The GP model described in Section 4.1 works well in most situations. However, when the number of data points per client is small, performance naturally degrades. To increase information sharing between clients and improve the per-client performance, we suggest augmenting the model with global inducing points shared across clients. When sending the model from the server to a client, we also send the inducing inputs and their labels. To streamline optimization and reduce the communication burden, we define the inducing inputs in the feature space of the last embedding layer of the shared NN. Therefore, usually, their size will be negligible compared to the network size.
We denote by X̄ the learned inducing inputs and by ȳ their fixed class labels. They are set evenly across classes. During training, we regard only the set of inducing inputs-labels (X̄, ȳ) as the available (training) data and use them for posterior inference. More formally, we first compute pθ(f |ω̄, ȳ, X̄,X) = ∫ pθ(f |f̄ , X̄,X)pθ(f̄ |ω̄, ȳ, X̄)df̄ using its analytical expression for the actual training data and then compute the probability of y using Gauss-Hermite integration. Then we use Eq. 3 & 4 for learning the network parameters and the inducing locations. At test time, to make full use of the training data, we combine the inducing inputs with the training data and use both to obtain the GP formulas and to make predictions. We note that with just using the inducing inputs at test time the model performs remarkably well, despite having almost no personalization component. See Appendix E.8 for a further discussion.
One potential issue with using IPs in this manner is that it distorts the true class distribution. As a result, the classifier may be more likely to predict a low-probability class during test time. We address this issue by adjusting the output distribution. In general, let p(y,x) and q(y,x) be two distributions that differ only in the class probabilities, i.e. p(x|y) = q(x|y), the predictive distribution follows:
q(y∗|x∗) p(y∗|x∗) ∝ q(x ∗|y∗)q(y∗) p(x∗|y∗)p(y∗) =⇒ q(y∗|x∗) ∝ q(y ∗) p(y∗) p(y∗|x∗). (5)
We use this to correct the GP predictions to the original class ratios at each tree node. We found in our experiments that this correction generally improves the classifier performance for class imbalanced data. As an example for this phenomena, consider a binary classification problem having 90 examples from the first class and 10 examples from the second class (therefore, q(y = 0) = 0.9, and q(y = 1) = 0.1). Assume we defined 50 inducing inputs per class, so now during test time the model sees 140 samples from the first class and 60 samples from the second class which corresponds to probabilities p(y = 0) = 0.7 and p(y = 1) = 0.3.

heading: 4.3 Augmenting the model with inducing points: computational efficiency
text: Learning the full GP model described in Section 4.1 requires inverting a matrix of size Nc in the worst case (at the root node), which hasO(N3c ) run-time complexity andO(N2c ) memory complexity.
Algorithm 1 pFedGP. C clients indexed by c; E - number of local epochs; |S| - number of sampled clients; M - number of inducing inputs per class
Server executes: Initialize shared network θ ← θ0 Initialize M inducing inputs per class for all classes in the system # in pFedGP-IP variants only for each round t← 1, 2, ... do:
Sample S clients uniformly at random for each client c ∈ S in parallel:
θct+1,M c t+1 ← ClientUpdate(θt,Mt) # obtain updates from client c
Update θt+1,Mt+1 using FedAvg [51] update rule.
ClientUpdate(θ,M ): for each local epoch e← 1, ..., E do:
if e = 1: Build GP tree classifier using the personal dataset Dc
Update θ,M using gradient-based optimization methods on Dc with LML or LPD return θ,M
Therefore, we propose an additional procedure based on inducing points to allow reduced complexity in low resource environments and scalability to larger dataset sizes.
This variant is based on the fully independent training conditional (FITC) method [70]. The key idea is to cast all the dependence on the inducing points and assume independence between the latent function values given the inducing points. Here for brevity, we omit the subscripts denoting the client and the tree node. However, all quantities and data points are those that belong to a specific client and tree node. Let X̄ ∈ RM×d denote the pseudo-inputs (defined in the embedding space of the last layer of the NN), and f̄ ∈ RM the corresponding latent function values. Here as well, the inducing inputs are defined globally at the server level and they are set evenly across classes.
We assume the following GP prior p(f , f̄) = N ( 0, [ KNN KNM KMN KMM ]) , where KMM is the kernel between the inducing inputs, KNN is a diagonal matrix between the actual training data, KNM is the kernel between the data and the inducing inputs, and we placed a zero mean prior. The likelihood of the dataset when factoring the inducing variables and the Pólya-Gamma variables (one per training sample), and the posterior over f , both have known analytical expressions. We can then obtain the posterior and marginal distributions by marginalizing over f̄ . Here we will present the posterior and marginal distributions:
p(f |X,y,ω, X̄) = N (f |KNMB−1KMNΛ−1Ω−1κ, KNN −KNM (K−1MM −B −1)KMN ),
(6)
p(y|X,ω, X̄) ∝ N (Ω−1κ|0, Λ + KNMK−1MMKMN ). (7)
Where Ω = diag(ω), κj = yj − 1/2, Λ = Ω−1 + diag(KNN −KNMK−1MMKMN ), and B = KMM + KMNΛ
−1KNM . Importantly, we only need to invert M ×M or diagonal matrices. See full derivation in Appendix B. During test time, we use f̄ to get the posterior of f∗ to compute the predictive distribution.
Now we can use either the marginal or the predictive distribution to learn the shared NN parameters and the inducing locations. The complexity of applying this procedure is reduced toO(M2Nc+M3) in run-time, and O(MNc + M2) in memory. While the (conditional) independence assumption between the latent function values may be restrictive, we found this method to be comparable with the full GP alternative in our experiments. Potentially, this can be attributed to the effect of sharing the inducing inputs among clients and the information that ω stores on f .

heading: 5 Generalization bound
text: It is reasonable to expect that after we learned the system new clients will arrive. In such cases, we would like to use pFedGP without re-training the kernel function. Under this scenario, we can derive generalization bounds concerning only the GP classifier without taking into account the fixed
neural network using PAC-Bayes bound [50]. Having meaningful guarantees can be very important in safety-critical applications. The PAC-Bayes bound for GPC [64] (with the Gibbs risk):
Theorem 1. Given i.i.d. samples Dc = {(xi, yi)}Nci=1 of size Nc drawn from any data distribution over X × {−1, 1}, a GP posterior Q, and a GP prior P , the following bound holds, where the probability is over random data samples:
PrDc{R(Q) > RDc(Q) +KL−1ber(RDc(Q), (δ, n, P,Q))} ≤ δ. (8)
Here, we have,
R(Q) = E(x∗,y∗)[Prf∗∼Q(f∗|x∗,Dc){sign f ∗ 6= y∗}], RDc(Q) =
1
Nc Nc∑ i=1 Prfi∼Q(fi|Dc){sign fi 6= yi}
(δ,Nc, P,Q) = 1
Nc
( KL[Q || P ] + log Nc + 1
δ
) , KL−1ber(q, ) = maxp∈[0,1]KLber[q || p] ≤
(9)
An important observation in [64] is that the KL-divergence between the posterior and prior Gaussian processes is equivalent to the KL-divergence between the posterior and prior distribution of their values on the Nc training samples. While [64] assumed Q to be Gaussian, this observation still holds even without this assumption. However, when Q is no longer Gaussian, as is the case here, KL[Q(f) || P (f)] no longer has a closed-form expression. We can show that for the Pólya-Gamma augmentation:
KL[Q(f) || P (f)] = EQ(ω){KL[Q(f |ω)||P (f)]} −MI[f ;ω] = EQ(ω){KL[Q(f |ω)||P (f)]}+ EQ(f ,ω) [ log Q(ω)
Q(ω|f)
] (10)
where MI denotes the mutual information. Since Q(f |ω) and P (f) are Gaussian, the KL[Q(f |ω)||P (f)] term has a close form expression so we only need to perform MonteCarlo approximation on the expectation on ω on the first element. In the second expectation, Q(ω) does not have a known expression. To estimate it, given {(ωi, fi)}Ni=1 samples, we use Q(ωi) ≈ 1N−1 ∑ j 6=iQ(ωi|fj). Note that if the summation for j includes fi, it might result in a biased estimator. Further details on estimating KL[Q(f) || P (f)] are in Appendix C.
To assess the quality of the bound, we partitioned the CIFAR-10 dataset to 100 clients. We trained a shared network using our full-GP variant on 90 clients and then recorded the generalization and test error on the remaining 10 clients four times, each with a different training set size. Figure 2 shows the estimation of the generalization error bound (δ = 0.01) vs the actual error on the novel clients with the Gibbs classifier. First, we observe that indeed
the bound is greater than the actual test error for all points and that it is not vacuous. There is a strong correlation between the actual error and the bound. Secondly, unlike worst-case bounds (e.g., VC-dimension), this bound depends on the actual data and not only the number of data points.

heading: 6 Experiments
text: We evaluated pFedGP against baseline methods in various learning setups. We present the result for the following model variants: (i) pFedGP, the full GP model (Section 4.1); (ii) pFedGP-IP-data, the model with IPs described in Section 4.2; and (iii) pFedGP-IP-compute, the model with IPs described in Section 4.3. For pFedGP and pFedGP-IP-compute, the results obtained by maximizing the predictive and marginal likelihood were similar, with a slight advantage to the former. Therefore, we present here the results only for the predictive alternative and defer the results of the marginal alternative to the Appendix. Additional experiments, ablation study, and further analyses are provided
in Appendix E. Unless stated otherwise, we report the average and the standard error of the mean (SEM) over three random seeds of the federated accuracy, defined as the average accuracy across all clients and samples.
Datasets. All methods were evaluated on CIFAR-10, CIFAR-100 [38], and CINIC-10 [15] datasets. CINIC-10 is more diverse since it combines images from CIFAR-10 and ImageNet [16].
Compared methods. We compared our method against the following baselines: (1) Local, pFedGP full model on each client with a private network and no collaboration with other clients; (2) FedAvg [51], a standard FL model with no personalization component; (3) FOLA [47], a Bayesian method that used a multivariate Gaussian product mechanism to aggregate local models; (4) FedPer [4], a PFL approach that learns a personal classifier for each client on top of a shared feature extractor; (5) LG-FedAvg [44], a PFL method that uses local feature extractor per client and global output layers; (6) pFedMe [72], a PFL method which adds a Moreau-envelopes loss term; (7) FedU [18], a recent multi-task learning approach for PFL that learns a model per client; (8) pFedHN [65], a recent PFL approach that uses a hypernetwork to generate client-specific networks.
Training protocol. We follow the training strategy proposed in [65]. We limit the training process to 1000 communication rounds, in each we sample five clients uniformly at random for model updates. The training procedure is different in the FOLA and pFedHN baselines, so we used an equivalent communication cost. In LG-FedAvg, we made an extra 200 communication rounds after a pre-training stage with the FedAvg model for 1000 communication rounds. In the local model, we performed 100 epochs of training for each client. In all experiments, we used a LeNet-based network [40] having two convolution layers followed by two fully connected layers and an additional linear layer. We tuned the hyperparameters of all methods using a pre-allocated held-out validation set. Full experimental details are given in Appendix D.

heading: 6.1 Standard PFL setting
text: We first evaluated all methods in a standard PFL setting [65, 72]. We varied the total number of clients in the system from 50 to 500 and we set the number of classes per client to two/ten/four for CIFAR-10/CIFAR-100/CINIC-10 respectively. Since the total number of samples in the system is fixed, the number of samples per client changed accordingly. For each client, the same classes appeared in the training and test set.
The results are presented in Table 1. They show that: (1) The performance of the local baseline is significantly impaired when the number of samples per client decreases, emphasizing the importance of federated learning in the presence of limited local data. (2) FedAvg and FOLA, which do not use personalized FL, perform poorly in this heterogeneous setup. (3) pFedGP outperforms or is on par with previous state-of-the-art approaches when local data is sufficient (e.g., 50 clients on all datasets). When the data per client becomes limited, pFedGP achieves significant improvements over competing methods; note the 9% and 21% difference in CIFAR-100 over 100 and 500 clients, respectively. (4) pFedGP-IP-compute often achieves comparable results to pFedGP and is often superior to pFedGP-IP-data. We believe that it can be attributed to the fact that in pFedGP-IP-compute
the training data take an active part in the GP inference formulas (Eq. 6), while in pFedGP-IP-data the data impact in a weak manner only through the loss function. (5) pFedGP-IP-data is especially helpful when few samples per class are available, e.g., CIFAR-100 with 500 clients. That last point is further illustrated by decoupling the effect of the number of clients from that of the training set size. To illustrate that, in Appendix E.2 we fixed the number of clients and varied the number of training samples per class. From this experiment, we deduced that both factors (individually) contribute to pFedGP success.
A desired property from PFL classifiers is the ability to provide uncertainty estimation. For example, in decision support systems, such as in healthcare applications, the decision-maker should have an accurate estimation of the classifier confidence in the prediction. Here, we quantify the uncertainty through calibration. Figure 3 compares all methods both visually and using common metrics [7, 25, 55] on the CIFAR-100 dataset with 50 clients. Expected calibration error (ECE) measures the weighted average between the classifier confidence and accuracy. Maximum calibration error (MCE) takes the maximum instead of the average. And, Brier score (BRI) [7] measures the average squared error between the labels and the prediction prob-
abilities. The figure shows that pFedGP classifiers are best calibrated across all metrics in almost all cases. We note that with temperature scaling, the calibration of the baseline methods can be improved [25]; however, choosing the right temperature requires optimization over a separate validation set, which our model does not need. Additional calibration results, including temperature scaling, are presented in Appendix E.9.

heading: 6.2 PFL with input noise
text: In real-world federated systems, the clients may employ different measurement devices for data collection (cameras, sensors, etc.), resulting in different input noise characteristics per client. Here, we investigate pFedGP performance in this type of personalization. To simulate that, we partitioned CIFAR-10/100 to 100 clients similar to the protocol described in Section 6.1, we defined 57 unique distributions of image corruption noise [29], and we assigned a noise model to each client. Then for each example in each client, we sampled a corruption noise according to the noise model allocated to that client. Here we show the results for the noisy CIFAR-100 dataset in Table 2. Further details
on the perturbations performed and result for the noisy CIFAR-10 are given in the Appendix2. We observe a significant gap in favor of the pFedGP variants compared to baseline methods. Note that using global inducing points is slightly less beneficial in this case since they are defined globally and therefore are not tied to a specific noise type as the real client data is.

heading: 6.3 Generalization to out-of-distribution (OOD) novel clients
text: FL are dynamic systems. For example, novel clients may enter the system after the model was trained, possibly with a data distribution shift. Adapting to a new OOD client is both challenging and important for real-world FL systems. To evaluate pFedGP in this scenario, we followed the learning protocol proposed in [65]. We partitioned the CIFAR-10 dataset into two groups. The data in the first group was distributed between 90 clients for model training. The remaining data from the second group was distributed between an additional 10 clients that were excluded during training. Within each group, we set the class probabilities in each client by sam-
pling from a Dirichlet distribution with the same α parameter. For the training group, we set α = 0.1, trained the shared model using these clients, and froze it. Then, we evaluated the models on the second group by varying α ∈ {.1, .25, .5, .75, 1}, on the remaining 10 clients. As α moves away from 0.1 the distribution shift between the two groups increases, resulting in more challenging OOD clients. Figure 4 reports the generalization gap as a function of the Dirichlet parameter α. The generalization gap is computed by taking the difference between the average test accuracy of the (ten) novel clients and the average test accuracy of the (ninety) clients used for training. From the figure, here as well, pFedGP achieves the best generalization performance for all values of α. Moreover, unlike baseline methods, pFedGP does not require any parameter tuning. Several baselines were excluded from the figure since they had a large generalization gap.

heading: 7 Conclusion
text: In this study, we proposed pFedGP, a novel method for PFL. pFedGP learns a kernel function, parameterized by a NN, that is shared between all clients using a personal GP classifier on each client. We proposed three variants for pFedGP, a full model approach that generally shows the best performance and two extensions to it. The first is most beneficial when the number of examples per class are small while the second allows controlling the computational requirements of the model. We also derived PAC-Bayes generalization bound on novel clients and empirically showed that it gives non-vacuous guarantees. pFedGP provides well-calibrated predictions, generalizes well to OOD novel clients, and consistently outperforms competing methods.
Broader impact: Our method shares the standard communication procedure of FL approaches, where no private data is directly communicated across the different nodes in the system. This protocol does not explicitly guarantee that no private information can be inferred at this time. As we show, pFedGP is particularly useful for clients with little data, and for clients that have strongly different distribution. This has great potential to improve client personalization in real-world systems, and do better at handling less common data. The latter is of great interest for decision support systems in sensitive domains such as health care or legal.

heading: Acknowledgements
text: This study was funded by a grant to GC from the Israel Science Foundation (ISF 737/2018), and by an equipment grant to GC and Bar-Ilan University from the Israel Science Foundation (ISF 2332/18). IA was funded by a grant from the Israeli innovation authority, through the AVATAR consortium.
2The noisy CIFAR-10/100 datasets are available at: https://idanachituve.github.io/projects/pFedGP

Review summaries:
This work considers the task of learning user personalised predictive models in the federated learning setting. It is realised via a local Gaussian Process (GP) on each client which operates on top of features extracted from a neural network that is shared across the entire federation, named pFedGP. The authors motivate the usage of GPs via their Bayesian nature which allows for good performance in the low data regime (common in federated learning) as well as the extraction of prediction probabilities that are calibrated (crucial for safety-critical applications of FL).   Since the authors operate on federated classification tasks,  they employ the GP classification model proposed at [1], which decomposes a multi-class prediction problem into a series of binary predictions. The benefit of this is that one can use the Polya-Gamma augmentation scheme, which can allow for efficient sampling of the posterior distributions, necessary for learning the hyper parameters of the model. The authors then  introduce two additional variants of their overall model which employ inducing points, a common technique in the GP literature. The first one is motivated for improving the performance of the model in the case of limited data, whereas the second one is about making the model more scalable in the case of large local datasets. Finally, the authors also provide a PAC-Bayesian generalization bound for the model performance one new clients that enter the federation.   The paper concludes with a series of experiments that compare the performance of pFedGP against several baselines for federated learning and personalised federated learning.   [1] GP-Tree: A Gaussian Process Classifier for Few-Shot Incremental Learning, Achituve et al., 2021

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.