Variational Inference in Similarity Spaces: A Bayesian Approach to Personalized Federated Learning

Venue: NeurIPS_cc_2024_Workshop_BDU
Published: 2024-09-05 14:39:29

Abstract: Similarity space (or S-space) employs an encoder function, fed by labeled original pairwise data, to find a latent pairwise space with markers (prototypical) vector. It divides the space into regions where pairs of objects are either similar or dissimilar. This paper enhances S-space, equipping variational inference from personalized federated learning. The S-space representation aligns local representation spaces across clients, while variational inference improves generalization and reduces overfitting caused by data scarcity and client heterogeneity. Our theoretical analysis improved upper bounds on KL divergence between optimal local and optimal global variational models compared to traditional distributed Bayesian neural networks.

Main paper: 
 introduction: 
 Federated Learning (FL) has emerged as a methodology that enables learning from decentralized datasets without compromising data privacy [1] . Real-world FL applications often involve non-IID data, which decrease individual client performance due to heterogeneous local data distributions [1, 2] .Personalized Federated Learning (PFL) addresses this by tailoring models to better align with local data characteristics [3] . Various methods have been proposed to enhance PFL, including global model personalization [4, 5] , federated meta-learning [6, 7] , and parameter decoupling [3, 8] . Traditional Neural Networks (NN) are typically used in PFL, although they usually show poor calibration and overconfidence in predictions, particularly when faced with varying data distributions [9] [10] [11] .In contrast, Bayesian Neural Networks (BNNs) is a probabilistic approach that has been used in other contexts for modeling uncertainty and enabling models to learn continually by capturing past information [12] , suggesting its potential use for PFL. This paper proposes a novel framework for PFL with BNNs to address challenges in model overfitting due to limited local data (FL privacy constraints). Our approach leverages variational inference (VI) within an auxiliary representation space to enhance PFL model performance by quantifying weight uncertainty in NNs at client and server models. To achieve personalization, each client updates its local VI parameters by reusing the global distribution from the server and balancing the KL divergence between the local posterior distribution and the server variational parameters. This strategy improves the upper bounds on this KL divergence compared to traditional distributed BNNs [13, 14] . Finally, our experiments show promising results in five datasets in the literature. 
methodology: 
 Federated Learning:Consider a client u i wishing to train a machine learning model with their respective datasets X i . Traditional (centralized) machine learning training methods group all the data in the set X = X 1 ∪ • • • ∪ X N . In contrast, a FL system enables clients to collaboratively train a Submitted to Workshop on Bayesian Decision-making and Uncertainty, 38th Conference on Neural Information Processing Systems (BDU at NeurIPS 2024). Do not distribute.model with the same architecture across all clients without sharing their local datasets X i with one another. However, FL introduces challenges due to the heterogeneity of local datasets.(Neural Network-induced) Similarity Space: Let X be the set of x i ∈ X elements in an m-dimensional feature space, each associated with a label y i ∈ Y and the function ℓ : X → Y, which maps the data into their respective labels. The original feature space X is transformed into a latent feature space Z by a representation function f Θ : X -→ Z, where z i = f Θ (x i ) ∈ R d and Θ is the set of weights of a given NN. We estimate an auxiliary space, referred to as S-Space [15] , to refine the latent feature space with a mapping function f S : X × X → S, where if ℓ(x i ) = ℓ(x j ), then (x i ,x j ) is a similar pair. On the other hand, we consider (x i , x j ) to be a dissimilar pair if ℓ(x i ) ̸ = ℓ(x j ).Given a pair of elements (x i , x j ) from the original feature space, f S computes a similarity vector s ij using the absolute element-wise difference between their latent space representationsEQUATIONwhere s ij ∈ R d has the same dimension as z i , z n,k is the k-th feature of the n-th sample in the latent space representation Z, and the absolute operation ensures symmetry s ij = s ji .We define two disjoint sets in the S-Space to quantify the similarity between input pairs: similarity markers (M + ) and dissimilarity markers (M -). Their union M = M + ∪ M -is the set of all markers (or prototypes) and M + ∩ M -= ∅. The (dis)similarity between input pairs is determined by the distance of the vector s ij to the markers m e ∈ M. We use a Cauchy kernel [16] to measure the similarity q e ij between the point s ij and the marker m e as q e ij = (1+∥sij-me∥ 2 2 )-1 m e ′ ∈M (1+∥sij-m e ′ ∥ 2 2 ) -1 .Consider the pair (x i , x j ), we define the probability of s ij vector being "similar" as q + ij = p q p ij for all mp ∈ M + . Analogously, the probability of s ij vector being "dissimilar" is q - ij = n q n ij for all mn ∈ M -. Notice that q + ij + q - ij = 1 because the sets M + and M -are disjoint. The binary pair cross-entropy loss function J(•) isEQUATIONwhereu ij = 1[ℓ(x i ) = ℓ(x j )] is the indicator function, i.e., 1[•] = 1 if [•]is true and 0 otherwise;We note that similar definitions exist for this auxiliary space [15] . However, as demonstrated in the following section, this is the first instance of applying VI within FL settings.We exploit the S-Space by introducing VI techniques to estimate the markers m. Our challenge is to make statistical inferences from the posterior distribution p Θ (M | X 1 , . . . , X N ) based on the NN parameters Θ without compromising data privacy. To tackle this, each client minimizes the Kullback-Leibler (KL) divergence between each client's variational distribution q ϕ k and the true posterior, formulated as arg minEQUATIONwhere Q is a variational family of distributions.We assume that each client's marker distribution (variational parameters) in the S-Space follows a Gaussian distribution (N (µ, σ 2 ) ∈ Q), representing the variational distribution q ϕ k associated with the client k as a product of normal distributions (mean-field approximation) [13, 17, 18] . The likelihood function p θ (X k | M) ∝ exp(-J(X k )/α) is defined using an exponential loss function (Boltzmann distribution) [19] , where J(•) is the S-Space loss function (or energy function -Eq. 2)and α > 0 is a (temperature) scaling parameter [20, 21] .Denote as X \k = {X 1 , . . . , X k-1 , X k+1 , . . . , X N } the local datasets excluding the data from client u k . Note that client u k does not have access to X \k due to FL privacy constraints. We approximate the posterior distribution using Bayes' theorem and a server variational model as p Θ (M | X \k ) ≈ s(M).The KL divergence (Eq. 3) can be approximated as1 KL q ϕ k (M) s(M) p θ (X k | M) Z k = KL q ϕ k (M) ∥ s(M) + log Z k + 1 α Eq ϕ k [J(X k )] ,where Z k is a normalization constant.Following [23] , we have adjusted the scale α to enhance numerical stability. We omit the normalization constant log Z k from the optimization problem (Evidence Lower Bound) [13] . Our approach captures model performance on specific tasks and ensures regularization by minimizing divergence from the server model; cf. Appendix B for details. Consequently, our VI approach in FL is a dual optimization framework that enhances client-level personalization:Client: arg minEQUATIONServer: arg minEQUATIONTheoretical Analysis:In this section, we present a theoretical discussion about the S-space.Moreover, we state the necessary Assumptions 1, 2, 3 and analyze equal-width BNNs as in [24, 25] .Definition 2.1. (Optimal Variational Latent Space -OVLS) Consider all x i , x j ∈ X and a latent representation function f Θ : X → Z. The transformation f Θ generates an OVLS Z if (i) ∥s ij ∥ 2 ∼ δ(0) ⇒ ℓ(x i ) = ℓ(x j ), and if (ii) ℓ(x i ) ̸ = ℓ(x j ) ⇒ E[∥s ij ∥ 2 ] > 0,where δ(.) is the Dirac delta function, and ℓ : X → Y maps an unlabeled example x i into its true label y i .Definition 2.1 emphasizes the relationship between geometric proximity in the latent space Z and label identity. Unlike traditional deep metric learning methods [26, 27] , our approach estimates a more flexible representation and captures more sophisticated data patterns (see Appendix A). Furthermore, a 1-nearest neighbor classifier in OVLS is optimal, emphasizing its capacity for perfect classification in theoretical applications. Our algorithm induces a representation space as Definition 2. Based on Corollary 1 (see Appendix C), we adapt the aggregation function proposed in [13] by ordering the set M using the norm ∥•∥ 2 before transmitting the weights to the server. This permutation optimizes the aggregation function's performance, ensuring alignment with the criteria for OVLS.Theorem 1. If Assumptions 1, 2, 3 are true, then the following inequality holdsKL q ⋆ ϕ k (M) ∥ s ⋆ (M) ≤ D -1 D (C ′ nr n ) < C ′ (n -1)r n ,where D is the number of markers in M, r n and C ′ are constants.Theorem 1 provides an upper limit for the KL divergence between the local (client) optimal variational solution q ⋆ ϕ k and the global (server) optimal variational solution s ⋆ . Our approach, thus, improves this FL theoretical results by using the variational S-space to estimate an optimal global variational distribution, with a tighter upper bound on divergence compared to traditional BNNapproaches (KL(q ⋆ ϕ k (M) ∥ s ⋆ (M)) ≤ C ′ nr n ),as documented in [13, 14] . We present the proof in Appendix C.2. 
experimentation: 
 FL settings: We performed our experiments using the Flower framework [28] in a FL setting characterized by non-IID clients (quantity-based label imbalance) [29, 30] . For each dataset, we sorted the data by labels and divided it into N clients. Each client was assigned #S random nonoverlapping subsets (shard), each containing an equal number of samples [29, 30] . Thus, we used a server and 100/200 clients in our experiment to evaluate our model, and we trained our method with two NVIDIA RTX 6000 Ada Generation (48 GB) for 1000 FL epochs. For each training round, the server selects 5 % of clients to train for five local epochs of the user model. We use the F1-Score commonly used in classification tasks, which can be directly computed from the confusion matrix.We use five datasets for our evaluation: The MNIST/FMNIST datasets consist of 28x28 pixel grayscale images, with 60,000 training and 10,000 testing examples. The FedRep [31] 0.9598 0.9627 0.8351 0.8616 0.8439 0.8318 0.8874 0.9250 0.6891 0.7007 FedPer [3] 0.9553 0.9664 0.8296 0.8458 0.8988 0.9015 0.8974 0.9112 0.6599 0.6726 FedPop [17] 0.9180 0.9302 0.7734 0.8013 0.9124 0.9243 0.9033 0.9191 0.6211 0.6657 pFedSim [32] 0.8973 0.9406 0.7270 0.7670 0.9093 0.9135 0.8899 0.9017 0.6643 0.6975 pFedBayes [13] 0.8864 0.9143 0.8327 0.8452 0.8771 0.9025 0.9194 0.9236 0.6927 0.7098 DITTO [5] 0 Each malware binary code is visualized as a 64x64 grayscale image.We compared our proposal to six PFL techniques from the literature [3, 5, 13, 17, 31, 32] ,as well as the standard FedAvg (Global Model -GM) and Local model without client communication (baseline). The results, summarized in Table 1 , show that our approach outperformed FedAvg across all five datasets. For example, on the CIFAR10 dataset, our method achieved an F1-score improvement of 27.43 % with four shards per client.Our method also shows improvements across other datasets. For instance, on CIFAR10, our approach achieved the highest F1 scores of 0.7015 and 0.7202 with four and five shards per client, respectively.On the Malimg dataset, it scored 0.9324 with five shards per client, surpassing FedRep, which achieved the second-best result of 0.9250. However, on the MNIST dataset, our proposal achieved the third-best results for both shard values. Figure 1 compares the F1-scores between our proposal and the best methods from the literature across five datasets with four shards per client over 1000 FL epochs. In summary, our approach achieved the best performance in six out of ten FL settings and the second-best in two additional settings. These findings are consistent with results obtained using 200 clients, as detailed in Table 2 (Appendix D.3). 
conclusion: 
 In this work, we introduced the variational auxiliary similarity space within an FL environment, designed to enhance latent feature representation. The variational S-Space adapts to the unique data distributions of individual clients, mitigating the effects of data heterogeneity. Our evaluations, conducted across five datasets and ten different FL settings, demonstrate that our approach outperforms existing PFL methods. In addition, we have provided theoretical results showing that our approach achieves improved upper bounds compared to traditional distributed BNNs by applying an aggregation function projected onto BNNs within an FL framework. A particularly promising direction is designing a specialized aggregation function optimized for the variational S-space. 
 
 Appendix: As introduced by [26] , contrastive loss is a precursor approach for estimating latent representation based on pairs of items (z i , z j ), facilitating the discernment of their class relationships. Specifically, it is formulated asL ij c = y ij ∥z i -z j ∥ 2 2 + (1 -y ij ) max(0, ξ -∥z i -z j ∥ 2 ) 2 ,wherey ij = 1[ℓ ′ (z i ) = ℓ ′ (z j )] is the indicator function, i.e., 1[•] = 1 if [•]is true and 0 otherwise and the function ℓ ′ : Z → Y, which maps the latent data z i = f Θ (x i ) into their respective labels.This method minimizes the loss function L ij c (also known as energy) by clustering similar points (same class) and separates points from different classes by at least a predefined margin ξ.In contrast to this approach, we introduce the OVLS (Definition 2.1), a latent space representation projected for 1-nearest neighbor (1-NN) classification, achieving a classification accuracy of 100 %.In an OVLS, similar pairs are collapsed, i.e., the distance between any two points within the same cluster is effectively zero (∥z i -z j ∥ 2 = 0), while ensuring that clusters of different classes maintain a minimum separation distance of ξ > 0. (z i , z j ) from different classes (y ij = 0) is E[∥z i -z j ∥ 2 ] = ξ, leading to a contrastive loss of L ij c = 0 because max(0, ξ -∥z i -z j ∥ 2 ) 2 = 0.However, the expected distance between elements from the same class isE[∥z i -z j ∥ 2 ] = ξ √ 2/2because we have two different clustering for each class, resulting in a non-zero contrastive loss that could potentially disrupt the latent space. Our method advances the conventional contrastive loss framework by preserving the structural integrity of the OVLS. For the settings in Figure 2 , if we introduce four 2D markers (in S-space) at positions M + = {(0, 0), (ξ, ξ)} and M -= {(ξ, 0), (0, ξ)}, we achieve a loss function equal to zero. This approach maintains the structured separation between classes within the OVLS, offering an alternative to the traditional contrastive loss model.The NN architecture used for all FL approaches, including our proposed method, is identical.Following [13, 33] , the network dimensions are m-100-n for the MNIST and FMNIST datasets, where m represents the number of input features and d denotes the latent space representation dimension. Additionally, for other datasets, we utilized a LeNet-5 architecture with the same latent space dimension (d = 64). We employed SGD with a learning rate of 0.01 for all experiments. All baseline models were configured using the hyperparameters recommended in their respective original publications.This section describes the practical implementation of our proposed FL framework, emphasizing the optimization of similarity markers. Each client estimates a market set M, where M + represents similarity markers and M -represents dissimilarity markers.We hypothesize that the marker values for each client in the S-Space follow a Gaussian distribution. Therefore, the joint probability density function q ϕ (M) is modeled as a product of normal densities:q ϕ (M) = m k ∈M N µ k , Σ k = Diag(σ 2 k ) ,where µ k , σ k ∈ R d , Diag(.) denotes the diagonal matrix function, and mean µ and covariance matrix Σ = LL T are the variational parameters. In addition, L is a diagonal matrix representing a mean-field parameterization. The variational parameters µ k and σ k are initialized by sampling from the uniform distribution µ prior ∼ U(-d -1/2 , d -1/2 ), where d is the latent feature dimension, and the constant σ prior = 0.05, as referenced in [34] .Furthermore, the variance parameters σ k are reparameterized as σ k = exp(p k ) to enable the application of gradient-based optimization techniques directly, resolving issues with the non-negativity constraint on standard deviations [35] .For sampling marker instances m i , we follow the methodology in [36] , introducing a noise component ϵ ∈ R d sampled from a standard normal distribution N (0, 1): we sample m i ∼ q ϕ k (M) asm i = µ k + exp(p k ) ⊙ ϵ,where ⊙ denotes element-wise multiplication. This formulation makes m i differentiable, enabling the gradient backpropagation through the randomness introduced by ϵ.We employ Monte Carlo sampling [37, 36] to approximate the objective function for client k (Eq. 4):D B k = - n k n b 1 K b j=1 K i=1 J(B j ) + αKL(q ϕ k (M) ∥ s(M)) ,where B ⊂ X k represents a minibatch of size n b , n k denotes the total number of data points in dataset X k , and K = 10 is the number of samples used in the Monte Carlo estimation [17, 38] .Finally, via the backpropagation algorithm, we update the variational model parameters using minibatch gradient descent, denoted by ∆D B k .C.1 External results Assumption 1. (Ref. [13, 24] ) The activation function is 1-Lipschitz continuous.Assumption 2. (Ref. [13, 24] ) Consider a NN with T parameters, I hidden layers, n samples in the FL environment, an input dimension d, and M neurons per hidden layer. The parameters d, n, M, I are large enough such thatEQUATIONwhere H = BM andEQUATIONAs discussed in [13, 24] , the parameter σ n is constructed to facilitate the proof of Theorem 1, particularly in inequality 13. Given that the neural network parameters are bounded by B, their variance should be upper bounded by B 2 .Lemma C.1. (Ref. [13] ) Let s ⋆ (M) be the optimal server variational distribution based on the following FL optimization problems ⋆ (M) = arg min s(M) 1 N N i=1 KL q ⋆ ϕi (M) ∥ s(M) ,where q ⋆ ϕi (M) is the local optimal variational model for user u i , the server variational distribution parameters of s ⋆ (M) for marker m j are (µ ⋆,s j , σ ⋆,s j ). We denote µ ⋆,s j | n ∈ R and σ ⋆,s j | n ∈ R as the n-th components of the vectors µ ⋆,s j ∈ R d and σ ⋆,s j ∈ R d , respectively.Therefore, we haveEQUATIONEQUATION)where the variational distribution parameters of q ⋆ ϕi (M) for the j-th marker are (µ ⋆ i,j , σ ⋆ i,j ). Futhermore, µ i,j | a ∈ R and σ i,j | a ∈ R as the a-th components of the vectors µ ⋆ i,j ∈ R d and σ ⋆ i,j ∈ R d , respectively.Assumption 3. Let the number of positive markers be #M + , and the number of negative markers be #M -in the S-space. We assume #M + = #M -and #M + < n/2, where n represents the number of samples in the FL environment.Corollary 1. Let f i,Θ be a latent representation function that generates an OVLS, and q ⋆ ϕi (M) denote the optimal variational distribution for the i-th user, estimated by a NN with weights Θ. If the markers are permuted based on the norm, i.e., the variational parameters (µ ⋆ i,j , σ ⋆ i,j ) ∼ q ⋆ ϕi (M) are organized into an ordered set according to the norm ∥µ ⋆ i,j ∥ 2 , then the optimal server variational distribution s ⋆ (M) admits the existence of an i such that ∥µ ⋆,s i ∥ 2 ∼ δ(0).Proof. Consider a FL system comprising N clients, each utilizing an OVLS solution (see Definition 2.1). These parameters are organized into an ordered set based on the norm ∥µ ⋆ i,j ∥ 2 . Consequently, for expected positive markers with indices 1≤ j ≤ #M + -1, we have ∥µ ⋆ i,j ∥ 2 ≤ ∥µ ⋆ i,j+1 ∥ 2for all clients in the FL environment.From Definition 2.1 and the ordering, we conclude that for all client i, (µ ⋆ i,1 , σ ⋆ i,1 ) ∼ δ(0). According to Lemma C.1, the server aggregated model's mean µ ⋆,s 1 and variance σ ⋆,s 1 are given byµ ⋆,s 1 | a = 1 N N i=1 µ ⋆ i,1 | a , and (σ ⋆,s 1 | a ) 2 = 1 N N i=1 (σ i,1 | a ) 2 + (µ i,1 | a ) 2 -(µ ⋆,s 1 | a ) 2 .Therefore, we conclude than (µ ⋆ s,1 , σ ⋆ s,1 ) ∼ δ(0).Theorem 1. Suppose that the assumptions are true, then the following inequality holdsKL q ⋆ ϕi (M) ∥ s ⋆ (M) ≤ D -1 D (C ′ nr n ) < C ′ (n -1)r n ,where D = (#M + ) + (#M -) is the number of marker in M and C ′ and r n are constants.Proof. Considering the mean-field decomposition for µ ⋆ i,m , µ ⋆,s m ∈ R d , we haveq ⋆ ϕi (M) = D m=1 N (µ ⋆ i,m , (σ ⋆ n ) 2 ), and s ⋆ (M) = D m=1 N (µ ⋆,s m , (σ ⋆,s m ) 2 )•Thus, the KL divergence between q ⋆ ϕi (M) and s ⋆ (M) can be decomposed asKL(q ⋆ ϕi (M) ∥ s ⋆ (M)) = KL D m=1 N (µ ⋆ i,m , (σ ⋆ n ) 2 ) D m=1 N (µ ⋆,s m , (σ ⋆,s m ) 2 ) = D m=1 KL N (µ ⋆ i,m , (σ ⋆ n ) 2 ) ∥ N (µ ⋆,s m , (σ ⋆,s m ) 2 .For each marker m m , the KL divergence between two Gaussian distributions isEQUATIONBy Corollary 1, we know that for variational parameters(µ ⋆ i,1 , σ ⋆ i,1 ) ∼ δ(0) and (µ ⋆,s 1 , σ ⋆,s 1 ) ∼ δ(0), i.e., for any client i with optimal variational latent space, the variation parameters for the first marker m 1 (marker with lowest ∥ • ∥ 2 norm) is equal to the optimal server defined by optimization problem in Lemma C.1. Therefore, the KL divergence between those lowest norm markers resulted in zero (equal distribution). For the remaining markers m ≥ 2, we haveEQUATIONEQUATIONwhere we applied in Eq. ( 10) the bellow equality (based Lemma C.1 -Eq.A.14 in [13] ) ensuring that(σ ⋆ n ) 2 + (µ ⋆ i,m -µ ⋆,s m ) 2 (σ ⋆,s m ) 2 = 1,and the inequality applies Assumption 2 and Eq. ( 8) (as can see in) that(σ ⋆,s m ) 2 = (σ * n ) 2 -(µ ⋆ s,m ) 2 + 1 N N i=1 (µ ⋆ i,m ) 2 ≤ (σ * n ) 2 + B 2 .By bounding the variance term using Assumption 2, we obtainEQUATION)By Assumption 2, incorporating into Eq. ( 11) and following similar steps in Eq. A.19 from [13] , we which, combined with Eqs. ( 12) and ( 13), results inEQUATIONKL(q ⋆ ϕi (M) ∥ s ⋆ (M)) ≤ d(D -1) 2 log 2B 2 (σ ⋆ n ) 2 ≤ D -1 D (C ′ nr n ).Therefore, by Assumption 3, we getD -1 D < n -1 n ,and KL q ⋆ ϕi (M) ∥ s ⋆ (M) < C ′ (n -1)r n .In this experiment, we evaluate the impact of various hyperparameter settings on the performance of our approach. The results of these experiments are summarized in Figure 3 with the CIFAR10 dataset.Figure 3a shows the impact of the number of pairs used in the local training. As we increase the number of pairs, we gain more confidence about the pair distribution, improving model performance.However, this also increases the training time. Although the performance suggests that using a larger number of pairs can improve model performance, the F1-Score exhibits marginal improvements beyond 2048 pairs. For this short paper, we adopted 2048 pairs per epoch in the training step to balance performance and training efficiency.Figure 3b presents the effect of the α parameter, which controls the influence of KL divergence regularization between the local and global variational parameters (Eq. ( 4)). We also evaluate our approach with a different number of similar/dissimilar markers. As expected, our experiment indicates that the F1-Score is sensitive to α. For lower values of alpha (0.1 ≤ α ≤ 10), our model is stable. Furthermore, the number of markers affects our model's results, with more markers capturing the similarity structure more effectively. Therefore, our model achieves optimal performance with α = 100 and three similar/dissimilar markers (six markers in total).This section analyzes the norms of the positive and negative expected markers, denoted as ∥µ + ∥ 2 and ∥µ -∥ 2 , respectively, across five datasets. Figure 4 shows the positive expected markers normThe prior distribution is replaced with the global (server) distribution because the prior (for each client) is difficult to characterize in practice[13]. This approach avoids making assumptions about the prior distribution, leading to a better fit with the collected data. The global (server) distribution is also updated/recycled for each FL epoch[22]. 



Reviews: review: title: title: Paper Review: Variational Inference in Similarity Spaces: A Bayesian Approach to Personalized Federated Learning
main_review: main_review: ## Evaluation Overview
This paper presents a novel framework that incorporates Variational Inference (VI) within a similarity space (S-Space) to enhance Personalized Federated Learning (PFL). The approach builds on Bayesian Neural Networks (BNNs) and introduces a probabilistic perspective to handle client heterogeneity and data scarcity in federated learning environments. The authors claim that their method improves generalization and reduces overfitting while aligning local representation spaces across clients. Their contribution is evaluated across five datasets with comprehensive experiments.

## Quality
The paper demonstrates strong theoretical foundations by proposing a novel combination of S-space and VI to improve PFL. The authors show a clear understanding of the core challenges in federated learning and present a robust method to address these. They provide an in-depth theoretical analysis, which adds credibility to the proposed approach.
Moreover, the experiments are conducted on well-known datasets, and the results strengthen the quality of the work.

## Clarity
The paper is generally well-written, with detailed explanations of the methodology.

## Originality
The paper is original. The introduction of S-Space in a PFL setting is a novel contribution that brings together different research directions (similarity learning and Bayesian inference) in a cohesive way.

## Significance
The proposed method addresses critical issues in personalized federated learning. By improving the performance of federated models on such data, the paper has significant practical implications. The enhancement of PFL through probabilistic methods can greatly influence future research in federated systems.
The experimental results suggest that the method consistently outperforms existing solutions, positioning it as a competitive approach for personalized model training in decentralized settings.

## Pros
Novelty: The use of variational inference within similarity spaces for federated learning is an original and well-conceived idea.
Theoretical contribution: Provides an enhanced theoretical analysis with tighter KL divergence bounds, improving over existing distributed Bayesian methods.
Performance: Outperforms state-of-the-art PFL approaches across multiple benchmarks and datasets.
Robust methodology: Clear experimental setup with evaluations over multiple datasets.
title: title: New method for personalized federated learning that combines the concepts of S-Similarity and variational inference
main_review: main_review: This paper proposes a new method for personalized federated learning. It is based on the combination of the concept of similiarity space (S-space), which allows local distributions of client data to be taken into account, and variational inference to counterbalance data heterogeneity and the scarcity of available data. The paper combines solid theoretical evidence with benchmarks that test the proposed solution against methods in the literature. \
Overall, the paper is clear. However, some passages lack clarity and a focus on the literature. I think it would be a good idea to flesh out the introductory section and the benchmark section. In particular, I'd like to see a brief description of what other methods offer. In addition, a summary flowchart of the method would be a plus to better understand the new method and get an overview. \
Finally, I don't know if it would be possible to evaluate the privacy of learning using metrics. \
All in all, this new method is solid and original, and a definite plus-value for the field of personalized federated learning.