Federated Learning via Meta-Variational Dropout
Authors: Insu Jeon
Minui Hong
Junhyeog Yun
Gunhee Kim
Venue: NeurIPS 2023 poster (2023)
Keywords: Personalized Federated Learning
Variational Dropout
Meta-Learning
Bayesian Neural Network
TL;DR: We propose a new efficient Bayesian personalized federated learning algorithm based on the hierarchical variational dropout and sparse prior.

Abstract: Federated Learning (FL) aims to train a global inference model from remotely distributed clients, gaining popularity due to its benefit of improving data privacy. However, traditional FL often faces challenges in practical applications, including model overfitting and divergent local models due to limited and non-IID data among clients. To address these issues, we introduce a novel Bayesian meta-learning approach called meta-variational dropout (MetaVD). MetaVD learns to predict client-dependent dropout rates via a shared hypernetwork, enabling effective model personalization of FL algorithms in limited non-IID data settings. We also emphasize the posterior adaptation view of meta-learning and the posterior aggregation view of Bayesian FL via the conditional dropout posterior. We conducted extensive experiments on various sparse and non-IID FL datasets. MetaVD demonstrated excellent classification accuracy and uncertainty calibration performance, especially for out-of-distribution (OOD) clients. MetaVD compresses the local model parameters needed for each client, mitigating model overfitting and reducing communication costs. Code is available at https://github.com/insujeon/MetaVD.

Federated Learning via Meta- Meta-Variational Dropout
Insu Jeon'
Minui Hong' Junhyeog Yun" Gunhee Kim
Seoul National University, Seoul, South Korea {insuj3on, antemrdm) antemrdn//ggail... mommail.com com
[alsdml 123, gunhee) @snu. ac. kr
Abstract
Federated Learning (FL) aims to train global inference model from remotely distributed clients, gaining popularity due to its benefit of improving data privacy. However, traditional FL often faces challenges in practical applications, including model overfitting and divergent local models due to limited and non-IID data among clients. To address these issues, we introduce novel Bayesian meta- learning approach called meta- variatiaaaai dropout (MetaVD). Meta VD learns to predict client-dependent dropout rates via shared hypernetwork, enabling effec- tive model personalization of FL algorithms in limited non-IID data settings. We also emphasize the posterior adaptation view of meta-learning and the posterior ag- gregation view of Bayesian FL via the conditional dropout posterior. We conducted extensive experiments on various sparse and non-IID FL datasets. MetaVD demon- strated excellent classification accuracy and uncertainty calibration performance, especially for out-of-disee of-distrinttion (OOD) clients. MetaVD compresses the local model parameters needed for each client, mitigating model overfitting and reducing communication costs. Code is available at htttps//gituub.co///njjjo///////
Introduction
Federated learning (FL) aims training a global model from distributed clients without sharing or
collecting their sensitive raw data. Thanks to its
aspect of FL M. it is increasingly
popular be applied to various applications such as image classification 2 3 object detection
4. 5) keyboard suggestion 7 recommendation 8 [9. and healthcare [10 1 The conventional
algorithm could achieve convergence when the data from different clients independently and identically distributed (IID) 14]. However, due to the differences in preferences, locations, and usage habits of clients, the private data in FL. are usually non-lllD. When the data distributions of the clients vary, the local model learned from each client can diverge, and thus learning optimal global model could fail [5 [6]. Furthermore, the scale of client data may not be sufficient to train local model with large parameters, causing model overfitting and poor generalization
To overcome the challenge caused by non-lID data, personalized federated learning (PFL) has emerged [15 [6]. In the PFL, each client is allowed to have its own personalized model trained on each client's 's local data while still participating in the global model training. There are many branches of PFI., such as those based on multitask learning [19 20), meta-learning (21426), and transfer learning (27-3) Although these approaches improve training convergence in non-lll) data settings, they may still experience model overfitting with limited client data. Recently, the Bayesian learning paradigm has also been introduced in FL to address the overfitting by incorporating uncertainty the model parameters 37] However, they may also struggle with divergent local models when the data from different clients exhibit significant statistical variability. Motivated by these challenges, we aim to address the issues of FL with the non-lID and limited client data simultancously.
In this paper, we present Meta- Variational Dropout (MetaVD), a novel Bayesian meta- leeannea approach 38 42 developed for PFL. MetaVD learns to predict client-dependemt dropout rates
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

This mechanism facilitates a data-efficient estimation of client- specific posterior simply by modulating a shared global NN parameter across all clients. The adaptation of MetaVD to conventional Fl. algorithms (e.g., FedAvg [.. Reptile [22). MAMI. [21], or PerFedAvg (2 allows flexible model personalization across variety of non-IlID client distributions. MetaVD) also incorporates unique model aggregation strategy based on client-specific dropout uncertainty, providing a principled Bayesian way to consolidate local models into a global model. This strategy significantly enhances the convergence of the FL algorithm
via a hypernetwork [25 43 45
non-IlD data. In addition, MetaVD inherits the merit of Variational Dropout (VD) [46 47), which facilitates model compression. This not only improves model generalization for oft-of-data (OOD) clients but also reduces communication costs for parameter exchange. We performed an extensive analysis on Me covering a wide range of FL scenarios [18), including different scales, non-IID degrees, partitions, and multi-domain settings [48]. In all of these experiments, MetaVD achieves excellent results compared other state-of-the-are baselines.
2 Background
A standard approach to FL (e.g.. FedAvg (1) iterates between the local training on the client devices and global optimization at the server. Given M clients, each of which has a data set the FL problem can be formulated as follows:
D = [[(²", (("""""""))
M
Server: min I(w) = Σ m=1
gth gm"JT(((() (w)) Client: (m(()= = 1Dm Σ E(aT", :
(I)
w denotes the model parameter, and the global learning objective (w)) at the server is a weighted average of the local objectives 7 (w) over M clients. The weight 9rd is proportional to the size of the local dataset (c.g., |DT I/DD). The local loss in a client device is usually defined as the empirical negative log- ligelekddd on the m-th client's dataset DM (i.c., ((D""; w) = log p/y"">" w)).
The local training is carried out in parallel fully (or partly) in each client device, with multiple Stochastic Gradient Descent (SGD) [49] epochs to get the fine-tuned local parameter wt" Then, the aggregation step computes the global parameter in the server by taking the weighted average of the local parameters (i.e., T + Em= """ The global model parameter w is then used as the initial parameter for each client in the next round of local training. FL aims to train models on large distributed datasets by only exchanging model parameters (e.g., w and w") between server and local devices, thereby minimizing privacy leakage of client data.
Challenges in FL. There are many challenges to real-world FL applications. (1) Heterogeneity of client data. The original FL algorithm converges well when the client data is IID 12.14. However, the client data often has different characteristics (e.g., classes or tasks follow non-IID). The w would drift away from each other, causing the w be suboptimal [15] [6]. (2) Sparse participation. In practice, the total number of clients M can be extremely large, while communication between the server and the clients can be intermittent or unreliable. This creates the challenge of inconsistent training due to small subset of participating clients in each round of communication [18). (3) Poor generalization due to limited data. When the training data available on each local device is limited, the local model can easily overfit, resulting in poor generalization to unseen clients [17 18] (4)
Communication cost. FL optimization requires frequent communication between local devices and the central server to exchange model parameters. This process is slow and could introduce additional privacy concerns. Therefore, reducing model size is also an important area of research (50)
Bayesian FL. While conventional FL methods use point estimate of the parameter as in Eq recent work 33 has incorporated probability distributions over the model parameters. In these Bayesian FL approaches, the client device first estimates the local posterior from its data, and then the server aggregates the partially updated local posteriors into a global posterior. Based on FedAvg, Gaussian distribution is used to represent each parameter in FedAG 32). and posterior aggregation strategy using the MCMC technique proposed in FedPA [33] FedBE [3 uses a Bayesian ensemble global model with Gaussian or Dirichlet distributions. These methods improve prediction confidence and model convergence. However, FedAG [32] and FedBE (31 assume global Gaussian posterior, but only consider point estimates of local parameters. FedPA 33 uses global and local Gaussian posteriors, but only maintains the global posterior mean on the server, making difficult to track local uncertainties. In addition, the performance of previous methods can still be compromised by the

---

statistical discrepancies in the client data, as they are not tailored to the heterogeneous FL scenario. Recently, Bayesian PFL methods have also been introduced for non-IlI data [34 51 52]. For example, pFedGP R4 learns common Gaussian kernel for all clients and infers personalized classifiers for each client. pFedBayes 35 assumes an independent posterior model for each device while learning shared global prior. These approaches achieve substantial predictive performance improvement non-lID data senerios compared to the previous Bayesian FL approaches. However, their approach still has limitations, such as the high computational cost of inverting large kernel matrix in pFedGP [E and imposing strong constraints in predBayes [35]. Their model aggregation rule does not account for parameter uncertainty directly. In addition, probabilistic modeling typically requires additional parameters to approximate the density, which can increase communication costs.
3 Approach: Meta- Variatiooa Dropout
We propose a new Bayesian PFL approach, called Meta- Variational Dropout (MetaVD), which can simultancously handle model personalization, regularization, and compression. MetaVD also exploits the uncertainties in model aggregation, thereby improving the training convergence non-IID data. In addition, MetaVD is a universal approach that is compatible with various existing FL algorithms.
3.1 Variational Inference for FL
Instead of approximating the local posterior using MCMC in Bayesian FL [31333] we can utilize the (amortized) Variational Inference (VI) framework of Bayesian meta-learning [3 42) that is originally developed for few-shot multi-task learning [53) Considering each m-th client in FL as an individual task, evidence lower-bound (ELBO) LELR over all the M datasets is defined
M
max (o)
= Σ 9 (uu( [lg
w² -
(2)
m=1
Here, ply"'> """h ) represents likelihood model constructed using neural network (NN) on each m-th client's data [56]59) and w is a client-specific NN parameter. qww"; Φ) is a variational posterior (or probabilistic) distribution over the w", which is also characterized by o. 9thh is the local weight as defined in [ p(w") is a prior distribution that acts as a regularizer for the q(u"): Φ). The local ELBO defined each m-th client in Eq.(2) makes tradeoff's between the expected log-likelihood on its local dataset Dmm and the KL divergence with a prior. In fact, maximizing the (F)BO(#) with respect the variational parameter
equivalent to minimizing Ee 9g KL(q(w"; x) " [D")). In theory, q(w";; o) is trained to approximate the true local posterior distribution over the w (i.e., q("";; ; σ) R3 ppp""''" (ww") (DT)) 59).59)
straightforward Bayesian PFL approach might be utilizing a separate local posterior model q(w": :)) for each client device (e.g.. ø (o',... .... (M)). However, only sparse subset of clients can participate in each FL round due to the communication availability of edge devices. Moreover, the number of clients M can be extremely large in practice, and some clients might only have small datasets. Thus, learning the variational parameter ntTh independently for each device is difficult. As an alternative, we introduce new hypernetwork-basee [25] 43 43 conditional dropout posterior modeling approach that can be data-efficiently trained across multiple clients FL.
3.2 Meta- Vetta-Varittonall Dropout
Posterior model. To promote efficient model personalization and reduce overfitting in Bayesian FL, we define the posterior model in Eq based on a Variational Dropout (VD) technique that multiplies continuous Gaussian noise to the NN parameters during training to prevent overfitting 446 47 62). Me extends the VD posterior by employing global hypernetwork that leams to predict client-specific dropout rates (or personal model structure). In MetaVD, the variational posterior distribution for each m-th client's parameter, q("";: "ou"d o) in Eq can be constructed as
K
qww"; (0, = II
W1o 10,, of o where at = =e)
(3)
R=1
3

---

Server
(b)
Server
Model Update (6, 0.
Client
Adaptation
XG
Suz
(KAMI/Septii)
hypernety
Hypernat
X
Clarts
Noc
Clent
12
Client
BC Pi
9
0g
Global NiNg
parameters 1p
specific
Local Data
Posterior Aggregation
0
50
GLobal N/N
Figure Overview of the Meta- Variational Dropout algorithm. (a) The server's hypernetwork predicts -lient-specific dropout rates from client embedding, e". The global parameters, 6, and dropout variables, a"", are sent to each m-th client. (b) few-shot local adaptation is performed on each client's s device in parallel; then, the updated parameters are sent back to the server. (c) The server then aggregates those parameters to update its variational parameters t/, and e.
The variational parameter è is characterized by three distinct parameters (6, v..e). The is global NN model parameter kept at the server (with no client index m)) where K is the size of the model parameter. The posterior distribution over the m-th client's 's weight w is described as a product of independent Gaussian distributions. In the product term in Eq. (3. each k-th factor describes conditional Gaussian noise multiplication to the global NN parameter (e.g., w = 0 x e where eM (ii, a" and 1 is an K K-dimensional all-one vector). The ox represents the client- specific dropout variable on each k-th index of NN parameter 8g. The hy indicates hypernetwork parameterized by / [43) 45 predicts the client-specific dropout rate a.M The e is a learnable client embedding used as input to the hig. In essence, MetaVD is technique that modulates global NN parameter with Gaussian noises. By predicting the fffent-specifer dropout variable via a hypernetwork, single NN can be reconfigured for various different clients. Since learning a hypernetwork across multiple local clients is more efficient than learning all the local posteriors independently, this approach can mitigate the sparse client participation and limited data issues in FL.
Prior model. To optimize the posterior model q qw""; : )). we need to specify the KL divergence term and the prior model p( ("" in (q. Prior modeling in MetaVD requires two criteria: (i) the KL divergence terms must be independent of the NN parameter to ensure the lower-bound assumption in 61 62], and (ii) all clients must share the same prior model to support the multiplicative discussed in 47 because of its straightforward analytic KI. term derivation and proven effectiveness in network regularization and sparsification. Under the hierarchical prior assumption, the KI. divergence term in [)][(("")) ())
VD 46 47
posterior aggregation rule in Bayesian FL.. We adopt the hierarchical prior 66
Eq.21 2 is simplified to KL(g(w""
= Σk= 0.5 log(1 + (ag))"); see Appendix for a more detailed derivation. This KI. term is independent of the global NN parameter and efficiently regularizes the dropout variable. The same hierarchical prior is applied to all M clients.
Client-side objective. Initially, the hypernetwork in the server approximates the dropout variable at for each m-th client. The global parameter and a are transmitted to each client device. Then, the posterior model in m-th client is trained on local data using the following ELBO objective:
m ELOO (0, 8m
max
1
Σ
loggly? af (e;00,aa)
Σ R=1
0.5 log(1 + (ax)
(4)
1Dm
which maximizes Cm FL.BO on the client dataset D = ((]]" y." )!""" with respect to the variational
parameters (i.e., 0, )) ). The optimization can be done by the stochastic gradient variational Bayes
(SGVB). [5).58). which reparameterizes the random weight variable wtt using a differentiable transformation as wh f(c;; 0k, a. = Ax + /al' ff with a random IID noise 6 AO, 1). The intermediate weight wh is now differentiable with respect to 0g hence a. and can be optimized via the SGD 49 The analytical KL term acts as a regularization for o.
Server-side objective. Once local posterior adaptation is done in each client device, the updated parameters 0th and a, are returned the server, then the server updates the variational parameters
'The dropout rate p a/(1 a) 00
63 46 We refer the dropout variable for simplicity.

---

Algorithm 1 MetaVD algorithm with MAML and Reptile variant for FL
of communication round R. of client N. server learning rate n, client learning rate y, inner learning rate I. local steps E, inner steps I, KL divergence parameter B. Initialize a glohal parameter 6, hypemetwork h (-) and e. for 1.... R do
Input:
procedure LOCALADAPTATION MAML(A. for Each local step i from to do Set 0 8; and < = a; Sample dataset D and Dm from D for Each inner step from 1 to do
> FL Rounds
0f <- 0g 1V CELBO (0, D())
Sample Mf clients from 1,... N clients
for m = 1,.., .M do
0g 0, Ve. L EL,
BO 18, o Dm)) a; 04 > u, Lm ELBO (%, a Dm)
Set 0T = and at = hy (c"))
Send 0th and am to the me-th client
Set = 0g av oq
0 "r LOCALADAPIATIONOO" a
Send 8, and a. to the server
Compute global param 6a8 using 0Th a and Eq. A6 -TTH
procedure LOCALADAPTATION REPTILE( for Each local step i from to do
Aam <- at at
0 0 1100
8; + 0 - Vo. L (0i, og; D=")) Set 0. 0 and oq
V < + 17 ir Em (J Aa
> FE.BO (0,, agi Dm)
for m 1,--. M do
em et vl Vm
(V
Send and a, the server
(e.g., 0, v, e" using the updated and local parameters. To update the global NN parameter 0, we first
assume the Bayesian posterior aggregation rule: pDw)D) x II p(w" [D") [31333 66). Since each local posterior in Eq. is a Gaussian (dropout) distribution, their product is also Gaussian: N( .)) ~ I N (wnT (om"l0"" 10Th a=" ((m))) This gives us an exact aggregation rule to compute
the maximum a posterior (MAP) solution of the 6aee as follows:
y T where yT 9 (a (8TH )2)
(5)
M
Em (A (0Th (2)-
m
Note that Eq. has the intuitive interpretation that the aggregation weight Tr inversely propor- tional to its corresponding dropout variable (or noise variance) a." Thus, parameters with high uncertainty have correspondingly less influence on the global mean prediction. In this way, we can fully exploit the uncertainty of the model parameters in the aggregation. To fully update each global NN parameter Os, we follow the similar parameter update rule of FedAvg, except that the Bayesian aggregation Eq. 5) is utilized as in A For the parameter of hypernetwork th, more general update rule is used following 225.. We compute the changes the updated dropout for each client Aa described in Algo [ then approximated gradient of the hypemetwork parameter is computed using the chain rule as V C (a = (Voaa") "'"""""" Aa"" where V y.a" is the gradient of the output of hypernetwork with respect to v. Aa is approximation of the vector-jacobiat (a"))
product, which we are inspired by the work of (25] and [69]. The gradient for V.m Lo (a.") can be derived using the similar chain rule. The detailed update rules for each parameter are in Algo T
Combination with other meta-learning algorithms. Since the KL regularization in Eq is independent of the global (or initial) parameter 0, MD can be combined with several existing meta-learning based PFL algorithms (e.g., Reptile 22 70 MAML 25), PerFedAAg 23 71)).
For example, MAMI [21 [55] requires some internal update steps using subsampled dataset to compute the uecond-order gradient for 8. Reptile [2 70 uses only a first-order gradient computation as described in Algo D which is similar to FedAvg except for the addition of the learning rate r). Each local adaptation step is performed using the CH in Eq.44 and the local data set D Unlike conventional meta-learning based PFL algorithms, which maintain only one global initialization parameter, MetaVD allows to change the mode of the initialization parameters for each client.
4 Experiments
To validate the Meta VD approach, we conducted extensive experiments in various scenarios following the FL benchmark research [18) including different degrees of non-IID and client participation rates.
We adapted this rule from simplities to es MM r n²th where T' = (()=))))-1 /////> m=1 ( (()- 677. Interestingly, Bq. is
3R In general, the mode of the product of M Gaussians, A N( r²)).
equivalent to maximizing the logarithm of the product of the weighted posteriors (68] In the heterogeneous
data, the product rule can achieve smaller aggregation error than the mixture of Gaussian posteriors.
5

---

CIPAR-IO0 dataset
CIFAR- dataset
Hetrogenity
à 5.0
à 0.5
à 0.1
Method
Test (%) OOD (%)
A
Test (%)
OOD (%)
A
Test (%)
OOD (%)
A
FedAvg 1
42.35
43.08
+00.33
41.92
41.96
+0.04
71.65
71.57
-0.88
FedAvg+FT 18
41.49
42.45
+0.96
40.99
39.83
-1.16
69.62
68.38
1.24
FedProx 73
42.23
44.11
+1.88
42.03
40.51
-1,52
72.27
73.75
+1,48
FedBF [31
45.17
45.43
+0.26
44.29
44.23
-0.06
70.23
69.19
-1.04
pFedGP 34
42.69
43.07
+0.38
42.44
42.53
+0.09
71.94
76.83
+4.89
Reptile 22
47.87
47.73
-0.14
46.13
45.94
0.19
73.93
76.36
+2.43
MAML 21
48.30
49.14
+0.84
46.33
46.65
+0.32
76.06
74.89
-1.17
PerFedAvg (HI--MML) 23
48.19
47.35
0.84
46.22
46.36
+0.14
75.42
79.56
+4114
FedAvg+MetaVD (ours)
47.82
50.26
+2.44
47.54
47.55
+0.01
76.87
76.25
0.62
ReptieeMettaVD (ours)
53.71
54.50
+0.79
52.06
51.50
-055
76.51
82.07
+5.56
MAML+MetaVD (ours)
52.40
51.78
00.62
50.21
49.75
0.46
77.27
79.05
+1.78
PedFedAvg+MeetaDD (ours)
51.67
51.70
+0.03
50.02
48.70
- 1.32
76.06
81.77
+5.71
Table 1: Classification accuracies with different (non-III)) heterogeneity degrees of & CIFAR-100 and á 0.1 in CIFAR-10. The higher score, the better.
[5.0, 0.5] in
We used multiple FL. datasets 722.. including CIFAR- 10, CIFAR- 100, FEMINIST, and CelebA. To evaluate the effectiveness of the hypernetwork, also performed an ablation study comparing MetaVD to regular VD. Additionally, we assessed the uncertainty calibration and model compression ability of MetaVD, Finally, we MetaVD with multi-domain datasets.
Baselines. We compare our method with standard FL methods such as FedAvg [ and FedProx 73 meta-learning PFL algorithms like Reptile [22], MAML [21]. and PerFedAvg [23, and Bayesian FL methods including FedBE [3 and PFddGP 34 To ensure consistency, we employ the widely- used CNN model for FL 3 74 25 across all baselines. Additionally, "fine-tuning" (FT) refers to
performing few-shot adaptation steps with FedAvg before evaluation. An overview of baselines can be found in Appendix D.
For reproducibility, we run experiments in a containerized environment that simulates FL communication with clients only server. We test T 1000 of total FL rounds, following the conventions [8]. One baseline model can be run on a single GPU. All experiments cluster of 32 NVIDIA GTX 1080 GPUs. MetaVD's hypernetwor consists of an embedding layer of dimension (1 +1+ M/4). followed by three fully connected NNs with Reaky ReLU activation and exponential activation for the dropout logit output. The predicted dropout variable is then applied to the global weight of other baselines. In our study, we apply Meta VD to only one fully connected layer before the output layer (75.78). which leads to significant performance improvements in all experiments. See Appendix E for implementation details. Our code is available at httt:s///gguubb..o//nssjj///////
Implementation.
4.1 Generalization on Non-IID Settings
Datasets and training. To evaluate the gen- eralization capabilities of the models under non-IID data conditions, we perform tests on both CIFAR- 10 and CIFAR-100 datasets with varying degrees of heterogeneity. We follow the similar evaluation protocols of the pFI.- Bench [18. using Dirichlet allocation to par- tition each dataset into 130 clients with dif- ferent Dirichlet parameters, denoted à
0.5
500
700
750
100-
620
100
500
200
200
100
150
[5. 0.5, 0.1]. As shown in Figure the class labels and data size per client are heteroge- neous across clients. A smaller à represents a higher degree heterogeneity. the test accuracy and generalization performance of the baselines new clients, we randomly select 30 out of 130 clients as out-of-distribution (OOD) clients, which are not involved in the training phase. See Appendix G for more details.
9 220
49 2
$$
Figure 2: Visualization of client's data distribution in different non-llI degrees (a = [0.5, 0.1)).
evaluate

---

Results. Table I shows the weighted average classification accuracy for participating (Test) and particciiipi (OOD) clients on CIFAR-I0 and CIFAR-10 100 datasets with varying non-IID degrees. The generalization gap, denoted by A, represents the difference between OOD and Test accuracy. As shown in [ PFL. methods such as Reptile, MAMI., and PerFedAvg generally outperform non-PFL. methods such as FedAvg and FedProx. While the Bayesian ensemble approach, FedBE, improves on FedAvg. it still lags behind PFL methods in OOD accuracy. Additionally, pFedGP exhibits suboptimal performance in the presence of heterogeneous distribution of data among clients. As á changes from 5.0 to 0.5, the test and OOD accuracy of all models decreases as the degree of non-llI increases. When combined with Meta VD, all baselines show significant performance improvements, regardless of whether they are FL or PFL algorithms (e.g.. Reptile enhances from 47.87 to 53. 71 adapting MetaVD).. These results demonstrate the adaptability and effectiveness of Me in mitigating model overfitting and handling non-IID client data in FL contexts.
4.2 Ablation Study
Settings. To evaluate the capability of the hypernetwork MetaVD, we perform an ablation study by comparing MetaVD with naive VD [46] and EnsembleVD [79 approaches on the CIFAR-100 dataset. The naive (global) VD model maintains a global dropout parameter shared by all clients; the dropout parameter is treated independent dropout parameters for all clients. The client-specific dropout parameter can be stored in each client analogous to the partial FL 80 8 [ In contrast, MetaVD utilizes hypernetwork to learn the personal dropout rate across all clients. Bayesian posterior aggregation rules is applied in all model based on dropout rates to update the global model parameter
global model parameter, as in FedAvg, Ensemble VD maintains M
83
Results. Table 2 outlines the results of the abla- tion study; VD's hypernetwork-basee pos- terior modeling outperforms all other baselines.
CIFAR-100 dataset
Method
Test (%) OOD (%)
A
The dropout rates in baselines such VD or En- SD could not fully learn the independent dropout variables well due to restricted client participation. On the other hand, both the hyper- network and the global parameter converge well in MetaVD. This observation demonstrates that Meta VD D' 's hypemeework provides more data- eefta-effiientt approach learning client-specific model uncertainty compared to other baselines. Further ablation studies performed on the FEMNIST dataset shown in Appendix H.
Reptile
22
47.87
47.73
-0.14
Reptile+ VD 46
50.20
49.28
-0.92
Reptile+Ensemblee 79
52.49
52.36
-0.13
Reptile+MetaVD (ours)
53.71
54.50
+ 0.779
Table 2: MetaVD ablation study in CIFAR-100.
4.3
Uncertainty Calibration
(a) Repale
(b) MAML
CIFAR-100 dataset
ECE
0.77% MCE
ECE = 0.70%
19%
MOE 68.10%
Method
ECE (%)
MCE (%)
FedAvg 1
0.60
36.79
FedAvg+FT
0.69
15.04
FedProx 76
0.67
39.69
0.2
FedBE [31
Reptile 222
0.77
0.50
34.66
50.52
Accuracy
060
(4) Reptile + MctaVD
0.2
0.4
34
0.8
680
(4) MAML McsVD
MAMI.
0.75
46.57
1.0
DCE 0.32%
3,0
ECE 0.30%
PerFedAvg (HF-MAMI)
2
0.69
45.27
0.8
MCEE 86%
0.5
MCE 22.55%
FedAvg+MetaVD (ours)
0.39
25.27
0.6
0.5
Reptile+MetaVD (ours)
0.57
42.40
0.4
MAML-MetaVD (ours)
0.52
37.26
02
PorFedAvg+MettaD (ours)
0.43
30.20
00
0.2
04
0.5
0.8
100
0g00
02
0.4
0.5
0A
1.0
Confidence
Table 3: Uncertainty calibration scores (ECE and MCE) measured on the OOD client in CIFAR-100 (à- 0.1). The lower, the better.
Figure 3: Reliability diagrams for (a) Rep- tile, (b) MAML, (c) Reptile+MetaVD, and (d) MAML-MeetaV in CIFAR-100.
Settings. Identifying any potential bias in the model's prediction is important to avoid serious consequences, especially when the model is used to make important decisions [84186). In the FL
7

---

environment, where customers have limited non-llI data, it is even more critical to properly calibrate the prediction model. Therefore, we investigate whether the proposed MetaVD approach can also improve the calibration measures for FL. baselines. The Expected Calibration Error (ECE) measures the expected deviation between a model's predicted probability and the actual positive class frequency, while the Maximum Calibration Error (MCE) measures the maximum difference. These calibration metrics are commonly used to evaluate the reliability of probabilistic predictions.
Results. Table 3 summarizes the ECE and MCE values tested with the OOD clients in the CIFAR- 100 dataset and shows that the meta-learning based PFL algorithms (e.g., Reptile, MAML, and PerFedAvg) tend to have higher ECE and MCE values than the conventional FL algorithms (e.g.. FedAvg, FedProx, and FedBE). This means that PFL baselines achieve high classification accuracy, but their probability predictions are more likely to be biased. This may be a byproduct of the additional optimization-
adaptation steps of meta-learning with limited local client data, On the other hand, our MetaVD approach significantly reduces both ECE and MCE values for all meta- learning-basse -aaeed methods, indicating that MetaVD effectively mitigates overfitting and reduces bias on the OOD clients. Figure 3 shows the reliability diagrams that visualize the model calibration (84 in CIFAR-I00 (á+ = 0.5). They plot the expected sample accuracy as a function of confidence. If the model is perfectly calibrated, then the diagram becomes the identity function. Any deviation from perfect diagonal represents miscalibration. While Reptile and MAML tend to be overconfident in their predictions, Reptile+Meta VD comes fairly close to the desired diagonal function. Remarkably, in most instances, adapting MetaVD leads to improved model calibration.
4.4 Client Participation
FEMNIST dataset
Sparsity
8
0.2
0.1
8
0.05
Method
Test (%) OOD (%)
A
Test (%) OOD (%)
A
Test (%) OOD (%)
A
FedAvg 1
88.08
85.29
-2.79
88.13
84.70
-3.43
88.00
86.22
-1.84
FedAvg+FT 18
88.33
86.37
-1.9
87.85
86.95
-0.90
87.66
87.11
-0.55
Reptile 22
88.55
86.52
-2.03
88.39
87.20
-1.19
87.86
88.22
+0.36
FedAvg+MetaVD (ours)
88.81
85.28
-3,53
88.69
85.71
-2.98
88.66
86.21
-2.4
Reptile+MettVD (ours)
89.90
89.04
-0.86
89.86
88.63
-1.23
89.43
88.71
-0.72
Table 4: Results of classification accuracies with different participant client rates of 0.2, 8 0.1,
and s = 0.05 in the FEMNIST dataset. We report the results of participating clients (Test) and non-patiicipating clients (OOD). The higher the better.
Settings. In real-world FL scenarios, such as intermittent connections between clients and servers or limited client device performance, numerous clients may be unable to participate in each FL round. This is important for cross-device FL with a large number of clients or resource-limited clients. In this experiment, we evaluated the performance of the methods under different levels of client participation in each FL round. We experimented with 200 clients in the FEMNIST dataset. For each FL round, we randomly selected 40, 20, and 10 clients to participate during training, with participating client rates s of 0.2, 0. and 0.05, respectively. To measure OOD accuracy, we excluded 40 preselected clients from the selection of 200 clients that they do not participate in the entire training.
Results. The overall classification results with different participant client rates & are summarized in table 4 Reptile performs better than FedAvg. The effect of data heterogeneity on performance degradation becomes more severe as more clients participate training. The decrease in test accuracy as the participating client rate becomes smaller is due to having less training data as fewer clients participate in each round. Mcanwhile, Reptile+Meta' the performance drop for Reptile+ is not as significant as for Reptile, showing that Meta VD can adapt well to FL scenarios with smaller participant sizes.
outperforms the other basclines. Interestingly,
4.5 Muti-domain Datasets
Settings. Existing federated learning algorithms usually assume a single-domain approach, where only one dataset used in the experiment. Multi-domain learning [87-89) aims to utilize all available training data across different domains to improve the performance of the model. In this section, we
8

---

2 Domains (a)
2 Domains (h)
2 Domains (c)
3 Domains (d)
Method
Test (%) OOD (%)
Test (%) OOD (%)
Test (%) OOD (%)
Test (%) OOD (%)
FedAvg
43.65
43.45
64.02
52.88
86.81
81.01
63.73
35.55
Reptile
12
18.92
18.93
66.13
56.22
87.50
83.05
66.98
57.12
MAML
21
47.39
48.51
66.56
55.
88.57
84.52
67.08
58.15
PerFedAvg (HF-MAML) 26
49.21
50.91
66.57
55.24
87.94
83.75
67.20
57.03
FedAvg+MetaVD (ours)
48.23
48.62
65.58
56.85
87.38
82.67
65.93
58.58
Reptile+-MetaVD (ours)
52.26
54.75
68.35
59.07
88.63
85.03
68.78
61.59
MAML+MetaVD (ours)
51.34
52.82
68.24
61.21
88.59
85.02
68.
61.60
PerFedAvg+MettVD (ours)
51.18
53.06
67.93
61.32
88.27
85.32
68.05
61.17
Table 5: Classification accuracies with multi-domain datasets. (a) CelebA + CIFAR-100, (b) CIFAR- 100 FEMNIST. (c) CelebA FEMNIST, (d) CelebA CIFAR-100 + FEMNIST.
further evaluate the performance of PFL algorithms on multi-domain FL dataset, where we assume that each client have data from different domains. We use three different FL datasets to construct the multi-domain task distributions: FEMNIST, CIFAR-100, and CelebA, To sample each client's local data, we use the Dirichlet sampling technique (& 0.5) used in the $4.1
Results. Table 5 shows the classification accuracies for the Test and OOD clients when using a combination of two or three datasets. The meta-learning based PFL algorithm consistently outper- forms FedAvg in terms of classification accuracy. In addition, when MetaVD is applied to either the FL or PFL algorithms, it significantly improves prediction accuracy in all multi-domain settings. Notably, MetaVD shows greater improvements in OOD accuracy compared to the improvement in test accuracy. Overall, these results demonstrate the versatility of Ve in improving robustness and generalization even in multi-domain FL datasets.
4.6 Model Compression
Settings. Federated learning optimiza- tion requires frequent communication of model parameters between devices and the central server, which can be slow and may raise privacy concerns. There- fore, minimizing communication costs by reducing the size of model parameters an important issue in FL. To explore the compression capabilities of MetaVD, we performed an additional experiment on the CIFAR-I0 10 dataset. The sign DP indicates that each model parameter is dropped during the FL communication.
CIFAR-1O dataset (å = 0.5)
Method
Test (%) OOD (%) Sparsity(%)
Reptile+MetaVD
83. 20
83
0
MAML+MetaVD PerFedAvg+MeuaVD
81.32
81.81
0
81.06
81.47
Reptile+MetaVD+DD MAML+MetaVD+DP
81.40
80.98
80.06
81.48
81.73
79.49
PadfedAng-MeaaV--D
82.43
82.19
78.20
Table 6: Results of model compression. MetaVD+DP does not communicate the model parameters whose dropout rates are larger than 0.8. We used the thresholding technique to drop the model parameter; Any parameter with dropout rate greater than 0.8 was dropped during the FL rounds.
the test and OOD accuracy results, well the sparsity (%) in the CIFAR- 100 dataset. The sparsity represents the proportion of zero-value model parameters in the personalized layer. higher sparsity percentage indicates more parameter pruning or elimination performed on the weight. In our experiment, Meta VD was able to prune about 80% of the weights in the personalized layer. In addition, when we dropped the communication of the parameters between the client and the server using the dropout thresholding technique, MetaVD still showed relatively good performance. In the case of Reptile+Meta VD+DP, the performance decreased by about 2% while using only 20% of the weights. On the other hand, MAML+MetaVDDe an improvement in performance of about 1%. This shows that Meta can compress the model parameters required in the personalized layer, reducing the communication cost in FL without sacrificing much performance. Appendix shows more experiments on model compression results on the CIFAR-100 dataset with different non-IID settings.
Results.
and PerFedAvg+MMtta MettMMettDDDDD show

---

5 Conclusion
In this study, we presented a new novel Bayesian personalized federated learning (PFL) approach called meta-variational dropout (MetaVD). MetaVD utilizes a hypernetwork that predicts the dropout rates for each independent NN parameter, which enables effective model personalization and adapta- tion in federated learning (FL) with non-IlI and limited data scenarios. In addition, MetaVD is the first approach to exploit variational dropout uncertainty in posterior aggregation in PFL. MetaVD's dropout posterior modeling provides a principled Bayesian aggregation strategy to consolidate local models into a global model, thereby improving training convergence. Meta is also a generic approach that is compatible with any other existing saaa-leerning-baseel PFI. algorithms to avoid model overfitting. In addition, Meta ability to compress model parameters can be used to reduce communication costs. A potential limitation of our approach is that it may increase the complexity of the model by introducing an additional hypernetwork. However, the hypernetwork is kept on the server, and we have verified that applying MetaVD to just one last layer before the output layer yields significant performance improvements in all experiments. Experimentally. MetaVD's performance has been validated several PFL benchmarks, including CIFAR-10 10, CIFAR- 100, FEMINIST, and CelebA, as well as multi-domar datasets. It demonstrates superior classification accuracy and uncertainty calibration, especially for out-of-distribution (OOD) clients. Overall, the experimental results show Meta to be a highly versatile approach capable of addressing many challenges in FI..
Acknowledgment
This work was supported by Center for Applied Research in Artificial IntellgeenceCARRA) grant funded by Defense Acquisition Program (Adinistraaionn(AAA)) and Ageney for Defense Devel- opment(ADD) (UD190031RD)
Institute of Information & communications Technology Planning & Evaluation (ITT)) grant (No. (No.2019-0-01082, SW Starlab and N011-01-0- Artificial Intel- ligence Graduate School Program (Seoul National University)), Basic Science Research Program through National Research Foundation of Korea (NRF) funded by the Korea government (MSIT) Gunhee Kim is the corresponding author.
(NRT-200ORR
References
[1] Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication--fficieee learning of deep networks from decentralized data. Artificial intelligence and statistics, pages 1273-1282. PMLR, 2017.
[2] Yue Zhao, Meng Li, Liangzhen Lai, Naveen Suda, Damon Civin, and Vikas Chandra. Federated learning with Non-IID data. arXiv preprint arXiv: 00582, 2018.
[3] Jachoon Oh, Sangmook Kim, and Se- Young Yun. FedBABU: Towards enhanced representation for federated image classification. arXiv preprint arXiv:2106.06042 2021.
[4] Yang Liu, Anbu Huang, Yun Luo, He Huang, Youzhi Liu, Yuanyuan Chen, Lican Feng. Tianjian Chen, Han Yu, and Qiang Yang. FedVision: An online visual object detection platform powered by federated learning. Proceedings of the AAAI conference on artificial intelligence, volume 34, pages 13172-13179, 2020.
[5] Benyuan Sun, Hongxing Huo, Yi Yang, and Bo Bai. PartialFed: Cross-domain personalized federated learning via partial initialization, Advances in Neural Information Processing Systems, 34:23309-23322 2021.
[6] Andrew Hard, Kanishka Rao, Rajiv Mathews, Swaroop Ramaswamy, Françoise Beaufays, Sean Au- genstein, Hubert Eichner, Chloé Kiddon, and Daniel Ramage. Federated learning for mobile keyboard prediction. arXiv preprint arXiv: 1.03604, 2018.
[7] Timothy Yang, Galen Andrew, Hubert Eichner, Haicheng Sun, Wei Li, Nicholas Kong. Daniel Ramage, and Françoise Beaufays. Applied federated learning: Improving google keyboard query suggestions. arXiv preprint arXiv: (8/../2220000 2018
[8] Khalil Muhammad, Qinqin Wang, Diarmuid O' James Geraci, and Aonghus Lawlor. FedFast: Going beyond average for faster training of federated recommender systems. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pages 1234 1242, 2020.
Elias Tragos, Barry Smyth, Neil Hurley,
10

---

[9] Chuhan Wu, Fangzhao Wu, Yang Cao. Yongfeng Huang. and Xing Xie. FedGNN: Federated graph neural network for Srivac-Prreervvngg recommendation. arXiv preprint arXiv:2102.04925, 2021.
[10] Qian Yang, Jianyi Zhang, Weituo Hao, Gregory P Spell, and I Lawrence Carin. FL.OP: Federated learning on medical datasets using partial networks. In Proceedings of the 27th ACM SIGKDD Conference on Knowledge Discovery & Data Mining, pages 3845-3853, 2021.
[11] Ittai Dayan, Holger R Roth, Aoxiao Zhong, Ahmed Harouni, Amilcare Gentili, Anas Z. Abidin, Andrew Liu, Anthony Beardsworth Costa, Bradford J Wood, Chien-Sung Tsai, et al. Federated learning for predicting clinical outcomes in patients with COVID-19. Nature medicine, 27(10)11355-174, 2021.
[12] Xiang Li, Kaixuan Huang, Wenhao Yang, Shusen Wang, and Zhihua Zhang. On the convergence of FedAvg on Non-IID data. arXiv preprint a90i 081. 2019.
[13] Peter Kairouz, HI Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Kallista Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, et al. Advances and open problems in federated learning, Foundations and Trends in Machine Learning, 14(1-2):1-21 2021.
[14] Mohammed Aledhari, Rehma Razzak, Reza M Parizi, and Fahad Saeed. Federated learning: A survey on enabling technologies, protocols, and applications. IEEE Access, 8:140699- 140725, 2020.
[15] Viraj Kulkarni, Milind Kulkarni, and Aniruddha Pant. Survey of personalization techniques for federated learning. In 2020 Fourth World Conference on Smart Trends in Systems, Security and Sustainabiliit (WorldS4), pages 794- 797. IEEE, 2020.
[16] Alysa Ziying Tan, Han Ya, Lizhen Cui, and Qiang Yang. Towards personalized federated learning. JEEE Transactions on Neural Networks and Learning Systems, 2022.
[17] Honglin Yuan, Warren Morningstrr, Lin Ning, and Karan Singhal. What do we mean by generalization in federated learning? arXiv preprint arXiv:2110..1m 14216, 2021.
[18] Daoyuan Chen, Dawei Gao. Weirui Kuang. Yaliang Li, and Bolin Ding, PFL-Bench: A comprehensive benchmark for personalized federated learning, Advances in Neural Information Processing Systems. 35:9344-9360 0064444930,,2-- 2022.
[19] Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai. Exploiting shared representations for personalized federated learning. International conference on machine learning, pages 2089-2099. PMLR, 2021.
[20] Tian Li, Shengyuan Hu, Ahmad Beirami, and Virginia Smith. Ditto: Fair and robust federated learning through personalization. In International Conference on Machine Learning, pages 6357-6368. PMLR, 2021.
[21] Fei Chen, Mi Luo, Zhenhua Dong, Zhenguo Li, and Xiuqiang He. Federated Meta-Learning with fast convergence and efficient communication. arXiv preprint 1..2:v:1802077880 2018.
[22] Yihan Jiang, Jakub Konečný, Keith Rush, and Sreeram Kannan. Improving federated learning personal- ization via model agnostic meta learning. arXiv preprint arXiv: 0ttp//999./244. /2488, 2019.
[23] Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. Personalized federated learning with theoretical guarantees: A model-agnostic meta-leaming approach. Advances in Neural Information Processing Systems, 33:3557-3568, 2020.
[24] Canh T Dinh, Nguyen Tran, and Josh Nguyen. Personalized federated learning with moreau envelopes. Advances in Neural Information Processing Systems, 33:21394 21405, 2020.
[25] Aviv Shamsian, Aviy Navon, Ethan Fetaya, and Gal Chechik. Personalized federated learning using hypemetworks. In International Conference on Machine Learning, pages 9489-9502. PMLR, 2021. PMLR,
[26] Jianqing Zhang, Yang Hua, Hao Wang, Tao Song, Zhengui Xue, Ruhui Ma, and Haibing Guan. FedALA: Adaptive local aggregation for personalized federated learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pages 1237-11244, 2023.
[27] Qiang Yang, Yang Liu, Tianjian Chen, and Yongxin Tong. Federated machine learning: Concept and applications. ACM Transactions on Intelligent Systems and Technology (TIST), 10(2):1-19, 2019.
[28] Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. SCAFFOLD: Stochastic controlled averaging for federated learning. In Hal Daumé li and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages
PMLR, 2020.
11

---

[29] Zhuangdi Zhu, Junyuan Hong, and Jiayu Zhou. Data-free knowledge distillation for heterogeneous federated learning. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 12878-12889 12889, PMLR, 18 24 Jul 2021.
[30] Tiansheng Huang, Li Shen, Yan Sun, Weiwei Lin, and Dacheng Tao. Fusion of glohal and local knowledge for personalized federated learning, arXiv preprint arXiv:2302. 2023.
[31] Hong-You You Chen and Wei-Lun Chao. FedBE: Making bayesian model ensemble applicable to federated learning, arXiv preprint arXiv:2009.01974 2020.
[32] Adam Thor Thorgeirsson and Frank Gauterin. Probabilistic predictions with federated learning. Entropy. 23(1):41, 2020.
[33] Maruan Al-Shedivat, Jennifer Gillenwater, Erie Xing, and Afshin Rostamizadeh. Federated learning via posterior averaging: A new perspective and practical algorithms. arXiv preprint arXiv:2010.05273, 2020.
[34] Idan Achituve, Aviv Shamsian, Aviv Navon, Gal Chechik, and Ethan Fetaya. Personalized federated learning with gaussian processes. Advances Neural Information Processing Systems, 34:8392-8406, 2021.
[35] Xu Zhang, Yinchuan Li, Wenpeng Li, Kaiyang Guo, and Yunfeng Shao. Personalized federated learning via variational bayesian inference. In International Conference on Machine Learning, pages PMLR, 2022.
[36] Haolin Yu, Kaiyang Guo, Mahdi Karami, Xi Chen, Guojun Zhang, and Pascal Poupart. Federated bayesian neural regression: scalable global federated gaussian process. arXiv preprint arXiv:2206. 06357, 2022.
[37] Shrey Bhatt, Aishwarya Gupta, and Piyush Rai. Bayesian federated learning via predictive distribution distillation. arXiv preprint arXiv:2206.07562 2022.
[38] Sachin Ravi and Alex Beatson. Amortized bayesian Metannnarning. International Conference on Learning Representations, 2018.
[39] Marta Garnelo, Jonathan Schwarz, Dan Rosenbaum, Fabio Viola, Danilo J Rezende, SM Eslami, and Yee Whye Teh. Neural processes. arXiv preprint 1807.0022, 2018.
[40] Jonathan Gordon, John Bronskill, Matthias Bauer, Schastian Nowozin, and Richard E Turrnr. Mcta- Learning probabilistic inference for prediction. arXiv preprint arXiv: 180500800,,,921, 2018.
[41] Insu Jeon, Youngjin Park, and Gunhee Kim. Neural variational dropout processes. In International Conference Learning Representations, 2022.
[42] Ekaterina lakovleva, Jakob Verbeek, and Karteek Alahari. Meta-Learning with shared amortized vari- ational inference. In Hal Daumé li and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 4572-4582, PMLR, 2020.
[43] David Ha, Andrew Dai, and Quoe V. Le. HyperNetworks. arXiv preprint 16Xiv: 2016.
[44] Johannes Von Oswald, Christian Henning, Benjamin F Grewe, and João Sacramento. Continual learning with hypernetworks. arXiv preprint arXiv:1906.000e 00695, 2019.
[45] Dominic Zhao, Seijin Kobayashi, João Sacramento, and Johannes von Oswald. Meta-Learning via hypernetworks. In 4th Workshop on
NeurlPS 2020 (MetaLearn 2020). NeurlPS, 2020.
[46] Durk P Kingma, Tim Salimans, and Max Welling, Variational dropout and the local reparameterizzatoo trick. Advances Neural Information Processing Systems, 28, 2015.
[47] Yuhang Liu, Wenyong Dong, Lei Zhang, Dong Gong, and Qinfeng Shi. Variational bayesian dropout with a hierarchical prior. 2019 IEFF/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 7124-7133. IEEF, June 2019.
[48] Risto Vuorio, Shao-Hua Sun, Hexiang Hu, and Joseph J Lim. Multimodal model-agnostic meta-learning via task-aware modulation. Advances in Neural Information Processing Systems, 32, 2019.
[49] Yingjie Tian, Yugi Zhang, and Haibin Zhang. Recent advances in stochastic gradient descent in deep learning, Mathematics, 11(3):682, 2023.
12

---

[50] Farzin Haddadpour, Mohammad Mahdi Kamani. Aryan Mokhtari, and Mehrdad Mahdavi. Federated learning with compression: Unified analysis and sharp guarantees. In International Conference on Artificial Intelligence and Statisticx, pages 2358- PMLR, 2021.
[51] Nikita Kotelevskii, Maxime Vono, Alain Durmus, and Eric Moulines. FedPop: A bayesian approach for personalised federated learning. Advances in Neural Information Processing Systems, 85:8677-8701, 2022.
[52] Shiyu Liu, Shaogao Lv, Dun Zeng, Zenglin Xu. Hui Wang, and Yue Yu. Personalized federated learning via amortized bayesian Meta-Leamning arXiv preprint 0tt:::23070022222 2023.
[53] Brenden M Lake, Ruslan Salakhutdinov, and Joshua B Tenenbaum. Human-level concept learning through probabilistic program induction, Science, 350(6266):1332-3333 2015.
[54] Jake Snell, Kevin Swersky, and Richard Zemel. Prototypical networks for few-shot learning. Advances in Neural Information Processing Systems, 30, 2017.
[55] Chelsea Finn, Picter Abbeel. and Sergey Levine. Model- Aonel-Agnostic Meta-Learning for fast adaptation of deep networks. In Doina Precup and Yee Whye Teh, editors, Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pages 1135.113.110 PMLR, 2017.
[56] Radford M Neal, Bayesian Learning for Neural Nenworks, volume 118, Springer Science & Business Media, 2012.
[57] Alex Kendall and Yarin Gal. What uncertainties do we need in bayesian deep learning for computer vision? Advances in Neural Information Processing Systems, 30, 2017.
[58] Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra. Weight uncertainty in neural network. In Francis Bach and David Blei, editors, Proceedings of the 32nd International Conference on Machine Learning, volume 37 of Proceedings of Machine Learning Research, pages 1613-1622, Lille, France, 2015. PMLR.
[59] Kirill Neklyudov, Dmitry Molchanov, Arsenii Ashukha, and Dmitry Vetrov. Variance networks: When expectation does not meet your expectations. arXiv preprint 1rXiv:1803.08: 2018.
[60] Yarin Gal and Zoubin Ghahramani. Dropout as bayesian approximation: Representing model uncertainty deep learning. Maria Florina Balcan and Kilian Q Weinberger, editors, Proceedings of The 33rd International Conference on Machine Learning, volume 48 of Proceedings of Machine Learning Research, pages 1050-1059, New York, New York, USA, 2016. PMLR
[61] Dmitry Molchanov, Arseni Ashukha, and Dmitry Vetrov. Variational dropout sparsifies deep neural networks. In Doina Precup and Yee Whye Teh, editors, Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pages 2498-2507. PMLR, 2017.
[62] Jiri Hron, Alex Matthews, and Zoubin Ghahramani. Variational bayesian dropout: pitfalls and fixes. In Jennifer Dy and Andreas Krause, editors, Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 2019-2028 PMLR, 2018.
[63] Sida Wang and Christopher Manning. Fast dropout training. In Sanjoy Dasgupta and David McAllester, editors, Proceedings of the 30th International Conference on Machine Learning, volume 28 of Proceedings of Machine I Learning Research, pages 11 118-126, Atlanta, Georgia, USA, 2013. PMI.R
[64] Michael E Tipping. Sparse baycsian learning and the relevance vector machine. Journal of machine learning research, 1(Jun):211 1-244, 2001.
[65] S Derin Babacan, Rafael Molina, Minh N Do. and Aggelos K Katsaggelos. Bayesian blind deconvolution with general sparse image priors. In Computer Vision- 2012: 12th European Conference on Computer Vision, Florence, Italy, October 7-13, 2012, Proceedings, Part VI 12, pages 341-355. Springer, 2012.
[66] Willie Neiswanger, Chong Wang, and Eric Xing, Asymptotically exact, embarrassingly parallel MCMC. arXiv preprint arXiv:1311.4780, 2013.
[67] Paul Bromiley. Products and convolutions of gaussian probability density functions. Ting- Vision Memo, 3(4):1, 2003.
13

---

[68] Liangxi Liu,, Xi Jiang, Feng Zheng, Hong Chen, Guo-Jun Qi. Heng Huang, and Ling Shao. A bayesian federated learning framework with online laplace approximation. IEEE Transactions Pattern Analysis and Machine Intelligence, 2023.
[69] Michael Zhang, James Lucas, Jimmy Ba, and Geoffrey E. Hinton. Lookahead optimizer: k steps forward, step back. Advances in neural information processing systems, 32, 2019.
[70] Alex Nichol, Joshua Achiam, and John Schulman. On First-Order Meta- Leta-Learning algorithms. arXiv preprint 1rXiv:1803 003.02999, 2018.
[71] Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. the convergence theory of Gradient-Band Based Model-Agnostic Meta- Learning algorithms.. In Silvia Chiappa and Roberto Calandra, editors, Proceedings of the Twenty Third International Conference on Artificial Intelligence and Statistics, volume 108 of Proceedings of Machine Learning Research, pages 1082-1092. PMLR, 2020.
[72] Sebastian Caldas, Sai Meher Karthik Duddu, Peter Wu, Tian Li, Jakub Konečný, H Brendan McMahan, Virginia Smith, and Ameet Talwalkar. LEAF; A benchmark for federated settings. arXiv preprint arXiv:1812.01097 2018.
[73] Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogencous networks. Proceedings of Machine learning and systems, 2:429-450, 2020.
[74] Paul Pu Liang. Terrance Liu, Liu Ziyin, Nicholas Allen, Randy P Auerbach, David Brent, Ruslan Salakhutdinov, and I Louis-Philippe Morency. Think locally, act globally: Federated learning with local and global representations, arXiv preprint arXiv:2001 015523, 2020.
[75] Agustinus Kristiadi, Matthias Hein, and Philipp Hennig. Being bayesian, even just bit, fixes overconfi- dence in ReLU networks. In Hal Daumé li and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 5436 5446. PMLR, 2020.
[76] Noah Weber, Janez Stare, Arpit Mittal, Roi Blanco, and Lluís Mårquez. Optimizing over a bayesian last layer, 2018.
[77] Joe Watson, Jihao Andreas Lin, Pascal Klink, Joni Pajarinen, and Jan Peters. Latent derivative bayesian last layer networks. In Arindam Banerjee and Kenji Fukumizu, editors, Proceedings of The 24th International Conference on Artificial Intelligence and Statistics, volume 130 of Proceedings of Machine Learning Research, pages 1198-1200 1206. PMLR, 2021.
[78] Alexander A Alemi, Ian Fischer, Joshua V Dillon, and Kevin Murphy. Deep variational information bottleneck. arXiv preprint 1rXiv:1612.00410 2016.
[79] Wei Du, Xiao Zeng. Ming Yan, and Mi Zhang. Efficient federated learning via variational dropout. Accessed: 2023-10-21, September 2018
[80] Krishna Pillutla, Kshitiz Malik, Abdel-Rahman Mohamed, Mike Rabbat, Maziar Sanjabi, and Lin Xiao. Federated learning with partial model personalization. In International Conference on Machine Learning, pages 17716-17758. PMLR, 2022.
[81] Karan Singhal, Hakim Sidahmed, Zachary Garrett, Shanshan Wu, John Rush, and Sushant Prakash. Federated reconstruction: Partially local federated learning. Advances in Neural Information Processing Systems, 34:1 12201122111 11232, 2021.
[82] Dingzhu Wen, Ki-Jun Jeon, and Kaibin Huang. Federated dropout a simple approach for enabling federated learning on resource constrained devices. IEEE wireless communications letters, 11(5):923 927. 2022.
[83] Gary Cheng, Zachary Charles, Zachary Garrett, and Keith Rush. Does federated dropout actually work? Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) Workshops, pages 3387 0395,,3355000 2022.
[84] Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q Weinberger. On calibration of modern neural networks. In Doina Precup and Yee Whye Tch, editors, Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pages 1321-1330. PMLR, 2017.
[85] Alexandru Niculescu- Mizil and Rich Caruana. Predicting good probabilities with supervised learning. In Proceedings of the 22nd international conference on Machine learning, pages 625 625-632, 2005.
14

---

[86] Max-Heinrich Laves, Sontje Ihler. Karl-Philipp Kortmann, and Tobias Ortmaier. Calibration of model uncertainty for dropout variational inference. arXiv preprint arXiv:2006. 1584, 2020.
[87] Shoushan Li and Chengqing Zong. Multi-domain sentiment classification. In Proceedings of ACL-08: HLT, short papers, pages
2008.
[88] Yongxin Yang and Timothy M Hospedales. A unified perspective on Multi-Dommain and Multi-Tase learning. arXiv preprint arXiv: 1412.7489, 2014.
[89] Hakan Bilen and Andrea Vedaldi. Universal representations:the missing link between faces, text, plank- tons, and cat breeds. arXiv preprint arXiv:1701.072755 2017.
[90] Hongyi Wang. Mikhail Yurochkin, Yuekai Sun, Dimitris Papailiopoulos, and Yasaman Khazaeni. Feder- ated learning with matched averaging. arXiv preprint arXiv:2002.06440. 2020.
[91] In S Jeon, Deokyoung Kang, and Suk I Yoo. Blind image deconvolution using student's s-t prior with overlapping group sparsity. In 2017 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pages 1817-1821, IEEE, March 2017.
[92] Christopher M Bishop and Nasser M Nasrabadi. Pattern Recognition and Machine Learning, volume 4. Springer, 2006.
[93] David IC MacKay, A practical bayesian framework for hackpropagation networks. Neural computation, 4(3):448 472. 1992.
[94] David Barber and Christopher Bishop. Ensemble learning in bayesian neural networks. Nato ASI Series F Computer and Systems Sciences, 168:215-233 238, 1998.
[95] Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: A simple way to prevent neural networks from overfiting, 2014.
[96] Geoffrey E Hinton, Nitish Srivastava, Alex Krizhevsky, Ilya Sutskever, and Ruslan R Salakhutdinov. Im- proving neural networks by preventing co-adaptation of feature detectors. arXiv preprint w:Xiv:1277.055 2012.
[97] Li Wan, Matthew Zeiler, Sixin Zhang, Yann Le Cun, and Rob Fergus. Regularization of neural networks using DropConnect. In Sanjoy Dasgupta and David McAllester, editors, Proceedings of the 30th International Conference on Machine Learning, volume 28 of Proceedings of Machine Learning Research, pages 1058-1066, Atlanta, Georgia, USA, 2013. PMLR.
[98] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images, 2009.
[99] Gregory Cohen, Saced Afshar, Jonathan Tapson, and Andre Van Schaik. EMNIST: Extending MNIST to handwritten letters. In 2017 international joint conference on neural networks (IJCN)), pages 2921 -926.. IEEE, 2017.
[100] Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes the wild. In 2015 IEEE International Conference on Computer Vision (ICCV), pages 3730-3738. IEEE, December 2015.
[101] Namhoon Lec, Thalaiyasingam Ajanthan, and Philip HS Tor SNIP: Single-shot network pruning based on connection sensitivity, arXiv preprint arxiv:1810.02340. 2018
[102] Yuang Jiang, Shiqiang Wang, Victor Valls, Bong Jun Ko, Wei-Han Lee, Kin K Leung, and Leandros Tassiulas, Model pruning enables efficient federated learning on edge devices, 2022.
[103] Takuya Akiba, Shotaro Sano, Toshihiko Yanase, Takeru Ohta, and Masanori Koyama. Optuna: A next-gencration hyperparameter optimization framework. In Proceedings of the 25th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, KDD '19, pages 2623-2631, New York, NY, USA, July 2019. Association for Computing Machinery.
[104] Zachary Lipton, Yu-Xiang Wang. and Alexander Smola. Detecting and correcting for label shift with black box predictors. In International conference machine learning, pages 3122-3130. PMLR, 2018.
[105] Ali Ramezani- Ramezani-Kebrya, Fanghui Liu, Thomas Pethick, Grigorios Chrysos, and Volkan Cevher. Federated lcarning under covariate shifts with generalization guarantees. arXiv preprint 20xx:33060053255 2023.
15

---

A The Global Posterior Decomposition in Bayesian FL
In Bayesian estimation, an alternative to Maximum Likelihood Estimation (MLE) is the inference or estimation of the posterior distribution of the parameters given all the data, denoted as p(w[D = ..U....D This posterior distribution is proportional to the product of the likelihoods and a prior, p(w(D) x p(w) I (1" (w). In the case of a uniform prior, the modes of the global posterior coincide with the MLE solutions. This establishes an equivalence between the inference of the posterior mode and optimization. Under the uniform prior, any global posterior distribution that exists decomposes into product of local posteriors:
P() x II ((O)
(6)
This proposition is well discussed in the recent Bayesian FL work [33] In the FL context, the global model can be computed in the server by multiplicatively aggregating the local models adapted to each client. However, in practice, aggregating the posterior with the overfitted local models into the global one is difficult due to the heterogeneity among clients' data and the permutation invariance property of the NN architecture [0] Thus, frequent communications hetween the server and the client are still demanded the FL.. The challenge of inferencing the local and global model and improving communication efficiency remains an active research area for real federated applications.
B The Motivation of Hierarchical Prior in Meta VD
proposed in + for several reasons. The first reason is simply that it is a well-posed Bayesian prior that can avoid a degenerate posterior problem of the conventional VD prior 46 61 62]. The second reason is that we want a sparse prior to reduce communication costs by compressing the model; the hierarchical prior has been proven effective for parameter pruning. Although we briefly mentioned the KL divergence term as KL(g(W"; d)] p(w"")) = Ex= 0.5 log(1 + manuscript for readability, here we provide more detailed descriptions of applying the hierarchical prior.
Section 3.2, we adopt the hierarchical prior 66
") in the
Under the hierarchical prior assumption, we consider the joint prior and joint posterior distributions (2). The joint prior, p(w" 7m) = p(w"' p>(>"("""))) is defined a combination of zero-mean Gaussian distribution. p(w" I>"" - (0,"""),, and uniform hyper-prior, p(y"))
14(y")) |aa, b), over the variance. Then, we define a (conditional) joint variational posterior, q(wu", ; (0)) = q(w" (<)g(>>") comprising the (conditional) dropout posterior q(w"";; é = (8, é, e" e"") defined in Section 3.2 of the manuscript and an additional Dirac delta distribution, 4(>"",, approximate the true (joint) posterior p(w"" ,7 (D"" (given the client's 's dataset D"
ELBO. With the hierarchical prior p(w""," ") = p(w"" ((""""""") ("(""") and the (conditional) joint posterior q(w" (<)) q(w" ("){("")"" of the MetaVD, can derive the local objective for each client as follows:
q(w",, "(w"","""""" ,"""("b) Ow th 0yy"
KL(g(W"", " (1) [pp["p , ""D""")
"(w"", """ (bb) log p(w" ,0 / ym)
f alw" " (o) log g(w0 nnth "0))p(""") (x"" "wma
p(y amm nh (m mm)
f alw"., ,7 ( log g(w", (0)
logp(y" (I"" - log 1x" w
p( ,m)
KL(g(W", " (b) lp(w", y "(")) + log (x") - Betum - ply "h |x"", w0))).
(8)
Eq. (2) is derived from Bayes' rule: p(w", nm (Dm) = p(u" (a """"""""""""""" "o(""" "" By reordering (3), we get
p(u" [1=")
log p(y" "0""""/00" [x="" 2 Bo(""". " (0) [o> p(y" w w"""] - (L(g(w" >"""""") le(w". ,"""))
(9)
= -E (6) Byc [log p(y" "a""" w w"]] (L(((W", y "")))ll(("" """")
-yg 16) [log p(y"' [z." )="")] KL(g(w", >>"")] [[w"","""" ")"
(10)
The lower-bound in (4) is due to the positivity of the (L(g(w"" """"""""""""""""" ("l)))" (D))) Here, Eq. (5) corresponds to the EL objective of Eq. (3) in the manuscript.
KL term. If we further decompose the KI. divergence term in Fq. (5), KL(g(w", , [#)]lp(("" ))) ))- KL(g(w" (o)l""""))"""""""""""""")""""""" "> (L( ())(?""))
log (1((+))))))) - Σ (log(b }}}
(1I)
The log(b- - a) is independent of the unknown variables am 0, and ymh Thus, we do not need to specify the value of hyperparameters a and b and can neglect them in practic Eq. (6) provides the rationale behind the
For more detailed proof of this, please the appendix section of 47
16

---

KL divergence term that we mentioned in Section 3.2 of the manuscript. In fact, the independence between the KL term and the parameter 0 ensures compatibility with other optimization meta-learning algorithms. Also, the two-level structure in a hierarchical system can generate much more complex distribution, expanding the potential solution spaces for selecting feasible prior. Thus, the hierarchical prior is a suitable prior for interpreting the variety of different clients' models in the FL environment. The same hicrarchical prior is uniformly applied across all 1...M clients to ensure the global posterior decomposition assumption in Eq.(1).
C Additional Related Works
Bayesian Neural Networks and Variational Dropout.
To address the issue of overfitting caused by limited data, Bayesian neural networks (BNN) [93 56 94 [7] were proposed, which impose prior distribution
on each parameter. Various inference approaches have been proposed to model the posterior distribution
of BNNs (58 46 55 Variational dropout (VD) (0) 46 62 47 encompasses a set of techniques that
model the variational posterior distribution based on the dropout regularization method. Dropout regularization
955 96 randomly deactivates some of the model parameters during training by multiplying them with discrete Bernuuli random noise. This method was initially popularized to prevent overfitting in neural network models. Subsequently, fast dropout (63 (7) demonstrated that multiplying the parameters with continuous noise sampled from Gaussian distributions yields similar results to conventional Bernoulli dropout
[95]. VD approaches often employ sparse priors for regularization (46 [61 (2 (7). which facilitate the learning of independent dropout rates on the neural network parameters as variational parameters. This property distinguishes VD approaches from conventional dropout, which uses a single fixed rate for all parameters.
D Dataset and Methods
We follow the datasets and the evaluation protocol of pFL.-Bench for federated learning.
18, which is a recently proposed benchmark
Datasets. Here, we present descriptions of the dataset used in our experiment.
* The CIFAR-10 10 and CIFAR-100 datasets [98] are popular for 10-class and 100-class image classifica- tion respectively. Each dataset contains 50,000 training and 10,000 test images with a resolution of 32x32 pixels. Following the heterogeneous partition manners used in [8] we use Dirichlet allocation to split this dataset into 130 elients with different Dirichlet parameters as = [5, 0.5, 0.1] (a smaller indicates a higher heterogeneous degree).
The Federated Extended MNIST (FEMNIST) is widely used FL dataset for 62-class handwritten character recognition (72). The original FEMNIST dataset contains 3,550 clients and each client corresponds to a character writer from EMNIST [99]. Following (18. we adopt the sub-sampled version, which contains 400 clients and a total of 85350 training and 21536 test images with a
resolution of 28x28 pixels.
The CelebA is a FL dataset based on 100 for 2-class image classification; Smiling or Not. Following 18). we adopt the sub-sampled version, which contains 500 elients and a total of 8752 training and 2347 test images with a resolution of 84x84 pixels. Each client is assigned images of single celebrity.
We randomly select 30 clients for CIFAR- 10 and CIFAR- 100 and 40 clients for FEMNIST as OOD clients who do not participate in the FL processes. For all of the datasets, follow the same heterogeneous patterns exhibited in the pFL-Bench, which covers a wide range of scales, partition manners, and non-III degrees. This enables comprehensive comparisons and analysis among different methods in the non-IID data environment.
Baselines. We present an overview the baseline models used in our experiment, covering various popular and ofate-of-the-art approaches across three categories: Non-PFL, weeta-leanning- based PFL, and Bayesian FL methods.
The following Non-PFL methods are considered in our experiments:
FedAvg I is standard FL algorithm that averages gradients weighted by the data size of clients in each FL round.
[3] employs proximal term to encourage updated local models for clients not to deviate too much from the global model.
FedProx
The following (meta-lcamingg based) PFL methods are considered in our experiment:
inserts a meta- learning fine-tunn tuning phase of [70, 55 after the federated averaging algorithm stage to provide reliable personalized model.
Reptile [2
17

---

MAML 21 enhances federated learning by integrating MAML-based meta-learner [55] dividing the dataset into support and query sets for more robust local training.
PerFedAvg [
method focuses the convergence analysis of the HF-MAML algorithm in the
FL scenario, offering provably convergent method based on MAML to tackle non-convex functions.
The following Bayesian FL algorithm are also compared in our experiment:
FedBE [1] enhances robust aggregation by adopting Bayesian inference approach, sampling high- quality global models, and combining them through Bayesian model ensemble using Gaussian or Dirichlet distributions fitted to local models.
employs Gaussian processes for personalized federated learning, leveraging shared kernel functions parameterized by neural network and novel inducing point approach for the generalization in the low data regime.
PFedGP
The following Pruning algorithm is also compared in our compression experiment in the appendix:
algorithm is a deep learning pruning technique that identifies and removes less important connections in neural networks before training begins using a small subset the dataset. It is compared with our algorithm in the model compression experiment in the FL environment.
SNIP [Oo 102
Additionally, Fine-tuning (FT) in the baseline's name indicates fine-tuning the local models with a few steps before evaluation within the processes, which is similar to the adaptation step the optimization bassd meta-lcaming.
E Implementation Details.
Models. To maintain consistency with previous research, employ the widely adopted CNN model for all algorithms and baselines (25 74 B). Specifically, the global model comprises three convolutional layers with 64 filters and 3x3 kernels, followed by three fully-connected layers of 256, 128, and 64 hidden units.
The hypernetwork architecture of MetaVD consists an embedding layer, followed by two consecutive blocks containing linear layer and a LeakyRel.U activation function, and one block containing a linear layer with exponential activation output the dropout logit parameter The dimension of client embedding e" is proportional to the number of clients M and is calculated (1 + M/4). The hidden units' size in the hypemetwork was set 200. The predicted dropout logit parameter then applied to each weight the Meta VD layer within the global model. In our study, we selectively apply MetaVD to just one fully-connected layer right before the output layer of the global model. This simple adaptation of Mcta VD only in one fully- colnecconneeted layer yielded significant performance improvements across all experiments.
Hyperparameters. For all datasets, we set to 1000 ensure sufficient convergence following conventions
[18]. The batch size was set to 64, and local steps was set to 5. Personalization was executed with a batch size of 64 and 1-step update. In order to ensure fair comparison between the algorithms, the results presented in all of our experiments an extensive parameter optimization using an optimization tool called Optuna [103]. We employed both
obtained with the optimal hyperparameters for each model. To do so, conducted
Tree-structured Parzen Estimator algorithm and Random Sampler as hyperparameter samplers in Optuna.
For all methods, we investigated the server learning rate and local SGD learning rate within identical ranges. The server learning rate 7) was explored within the range of [0.6, 0.7, 0.8, 0.9, 1.0). The local SGD learning rate was investigated within the range of [0.005, 0.01, 0.015, 0.02, 0.025, 0.03]. In MAML and PerFedAvg, an additional client learning [0.01, 0.02, 0.03, 0.04, 0.05, 0.06, 0.07, 0.08, 0.09, 0.1]. For Meta VD, an additional KL divergence parameter is needed, and we sought its optimal value within the range of [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15). We follow the hyperparameter setting outlined in pFedGP, except for adjusting the batch size to 64 or 320 and investigating the learning rate within the range of [0.03 to 0.1]. To ensure the reproducibility of the experiments, will release all code, including baselines, on our GitHub repository.
y is required, for which
searched within the range of
F Experiment Outline
To evaluate the effectiveness and robustness of the proposed PFL methods, we tested the algorithms under various FL scenarios, including:
'tttps//optuna..or/
18

---

Degree of Data Heterogeneity: assessed the performance of each algorithm in scenarios where data from different clients heterogeneous, considering factors such variations in data distributions, label imbalance, and situations while some clients have limited training data available.
Ablation Study: we conducted an ablation study to verify the advantages of employing hypernetwork in MetaVD compared to naive Variational Dropout (VD) Ensemble VD approaches in Federated Learning.
Client Participation: we evaluated the performance of each algorithm under different levels of client participation rates in each FL round.
Uncertainty Calibration: used Expected Calibration Error (ECE) measures evaluate the accuracy probabilistic predictions made by the predictive models. These measures help determine whether a model is overconfident or underconfident its predictions and identify any biases in the model's s predictions.
Communication Efficiency: measured the communication overhead of each algorithm in terms of the rates of model compression.
Multi-domain Datasets:
Multi-domain learning aims to leverage all available training data across different domains enhance the performance of the model, but it typically results in suboptimal global model. We tested our approach in with multi multi-domain learning.
Model Compression: we compared MetaVD's compression capabilities with the existing pruning algorithm and evaluated the performance of Meta with and without compression.
G Additional Results with Non-IID Settings
CIFAR- datasct
Hetrogenity
= 5.0
0.5
à = 0.1
Method
Test (%) OOD (%)
A
Test (%)
OOD (%)
A
Test (%) OOD (%)
A
FedAvg [
42.35
43.08
+-0.33
11.92
41.96
+0.01
37.57
37.90
+0.33
FT 18
41.49
42.45
+0.96
40.99
39.83
-1.16
36.79
37.00
+0.21
FedProx 73
42.23
44.11
+1.88
42.03
40.51
-1.52
38.07
39.36
+1.29
FedBF [31
45.17
45.43
+0.26
44.29
44.23
-0.06
40.89
41.56
+0.67
pFedGP 34
42.69
43.07
+0.38
42.44
42.53
+0.09
37.65
38.09
+0.44
Reptile 122
47.87
47.73
-0.14
46.13
45.94
0.19
40.71
43.59
+2.88
MAML 21
48.30
49.14
+0.84
46.33
46.65
+0.32
41.56
45.19
+3.63
PerFedAvg (HF-MAML) 23
48.19
47.35
0.84
46.22
46.36
+0.14
42.46
42.92
+0.46
FedAvg+MetaVD (ours)
47.82
50.26
+2.44
47.54
47.55
+0.01
42.34
42.65
+0.31
Reptile+MetaVD (ours)
53.71
54.50
+0.79
52.06
51.50
-0.56
44.56
44.62
+0.06
MAML--MetVD (ours)
52.40
51.78
-0.62
50.21
49.75
0.46
43.84
48.17
+4.33
PedFedAvg+MMetVDD (ours)
51.67
51.70
+0.03
50.02
48.70
1.32
43.31
46.19
+2.88
Table 7: Results of classification accuracies with different (non-IID) heterogeneity strengths of á = 5.0, á = 0.5, and à 0.1 in the CIFAR-100 dataset. We report the results of participating clients during the training (Test) dataset and non-participating clients (OOD). The higher, the better.
To evaluate the generalization capabilities of the models under non-IID data conditions, we conducted tests on both CTFAR-10 and and
CIFAR- 100 datasets with varying degrees of heterogeneouss partitions. We randomly selected 30 out of 130 clients training phase. TT and Table T2 with different Dirichlet parameters a. A smaller à corresponds to a higher degree of heterogeneity. Test(%) denotes the weighted average accuracy of participating clients' test samples, 0OD(%) signifies the weighted average accuracy of the non non-participants, and A represents the participation generalization gap calculated as the difference between OOD accuracy and test accuracy.
partiiiiciiininngg oui-of-distrobution (OOD) clients, who were involved during the
display the weighted average accuracy while adjusting non-III) degrees
both datascts, PFL methods such as Reptile, MAML, and PerFedAvg generally outperform non-PFL methods like FedAvg and FedProx. The Bayesian ensemble approach on FedAvg (FedBE) offers a slight improvement in overall performance but still lags behind PFL methods in terms of OOD accuracy. Notably, when combined with MetaVD, all baselines experience significant performance enhancements, regardless of whether they employ FL or PFL approaches. Importantly, non-IID data, illustrating its versatility in effectively augmenting conventional FL algorithms without being limited by specific optimization techniques, addressing non- IID data and model overfitting issues. Additionally, our experimental results indicate that pFedGP performs relatively poorly when the data distribution among clients is heterogeneous.
method consistently improves OOD accuracy across all degrees of
19

---

CIFAR-IO dataset
Hetrogenity
à
5.0
à 0.5
à
0.1
Method
Test (%) OOD (%)
A
Test (%)
OOD (%)
A
Test (%) OOD (%)
A
FedAvg 1
80.73
79.98
0.75
78.28
78.90
+0.62
71.65
71.57
-0.88
FedAvg+FT 18
80.20
79.63
-0.57
76.96
75.83
-1.13
69.62
68.38
1.24
FedProx 73
80.79
80.30
-0.49
78.55
77.60
-0.95
72.27
73.75
+1,48
FedBF [31
81.55
81.07
-0.48
79.44
79.46
+0.02
70.23
69.19
-1.04
pFedGP 34
83.73
83.27
-0.46
79.56
79.37
-0.19
71.94
76.83
+4.99
Reptile 22
83.48
82.88
-0.60
79.35
79.41
+0.06
73.93
76.36
+2.43
MAML 21
82.61
81.69
00.9
80.31
80.08
0.28
76.06
74.89
-1.17
PerFedAvg (HI--MML) 23
82.57
81.99
0.58
80.87
80.85
0.02
75.42
79.56
+4114
FedAvg+MetaVD (ours)
82.60
82.74
+0.14
79.08
80.33
+1.25
76.87
76.25
0.62
Reptie+MettaD (ours)
84.70
84.28
-0.42
83 20
83. 40
+0.20
76.51
82.07
+5.56
MAML+MetaVD (ours)
83.93
84.95
+1.02
81.32
81.81
+0.49
77.27
79.05
+1.78
PefFedAvg+MetaVD (ours)
83.88
83.96
+0.08
81.06
81.47
+0,41
76.06
81.77
+5.71
Table 8: Results of classification accuracies with different (non-III) heterogeneity strengths of = 5.0, 0.5, and à = 0.1 in the CIFAR-10 dataset. report the results participating clients during the training (Test) dataset and con-patticipating clients (OOD). The higher the better.
In the original pFedGP paper [34], the data samples were partitioned across clients according to [25] 24 resulting each client having 2 and 10 classes for the CIFAR-10 and CIFAR-100 datasets, respectively. In contrast, our non-IlD experiment, we do not predefine the number of classes each client has; instead, determined by the Dirichlet parameter &. In addition, we found some limitations of pFedGP, which could not classify data points when encountering unseen labels within a client (especially when the label shift [I04] [OS] occurs between the training and testing distributions of the clients). To ensure that pFedGP would run the labels, circumvented this limitation allowing the client to have each class while evaluating model's 's performance.
40 50 data points for
Dirichlet parameter á
Dataset
á 5.0
0.5
0.1
CIFAR-100 CIFAR-10
99.78 = 0.5
73.18 5.2
33.78 + 7.33
10 = 0
9.10 = 0.94
4.65 + 1.49
Table 9: Statistics of the number of classes assigned to each client.
H Additional A Ablation Results in FEMNIST Dataset
FEMNIST dataset
Method
Test (%) OOD (%)
A
Reptile
87.86
88.22
+0.36
Reptile+VD
87.93
85.88
-2.05
Reptile=FnsembleeDD
87.99
87.97
-0.02
Reptile=MetaVD (ours)
89.43
88. 71
-0.72
Table 10: MetaVD ablation study results in the FEMNIST dataset.
this experiment, further performed the ablation study using the FEMNIST dataset to evaluate the benefits of MetaVD's hypernetwork in FL. We compared MetaVD to naive VD (46) and Ensemble VD 79 approaches. The naive (global) VD model keeps one global dropout parameter shared with all clients; the dropout parameter is treated as global model parameter in FedAvg. In EnsembleVD, we updated client-specific dropout rates manner analogous the local adaptation step MetaVD but maintained independent variational dropout rates different clients on the server. contrast, Me utilized a hypernetwork to learn dropout rates. All models employed Bayesian posterior aggregation rules based on dropout rates to update the global model parameter.
The ablation study's results on the FEMNIST dataset are outlined in Table 10 in the manuscript. It is evident from the table that MetaVD's conditional variational dropout-bass based hypernetwork surpasses all other baselines
20

---

in classification accuracy. Here too, Reptile-MetaVD consistently outperforms other methods, We noticed that baselines like VD EnsembleVD, client-specific dropout rates not fully optimized due restricted client participation. Consequently, the imitial dropout rates and divergence loss VD EnsembleVD remained fairly static. Conversely, MetaVD, both dropout rates and KL divergence loss converged all tests. This observation underscores that MctaVD's hypernetwork presents more data-efficiem approach to learning the client-specific model uncertainty compared to other basclines.
I Additional Results on Uncertainty Calibration
Addressing uncertainty calibration issues crucial for enhancing accuracy and reliability model's predictions, particularly when model is employed making significant decisions. Uncertainty calibration metrics evaluate the accuracy of probabilistic predictions made by a predictive model. By identifying and addressing any calihration issues, it hecomes possible improve the model's prediction accuracy and reliability, which is of essential importance in decision-making processes. clients have limited non-IID data, overfitting can easily occur, making the calibration of prediction uncertainty even more critical. We conducted experiments to evaluate the effectiveness our proposed algorithm in improving uncertainty calibration, demonstrating its potential to strengthen model performance in various real applications.
Federated Learning environment where
CIFAR-100 dataset
Hetrogenity
5.0
á
Method
FCE (%) MCE (%)
ECE (%) MCE (%) ECE (%)
MCE (%)
FedAve
0.29
22.05
0.53
32.35
0.60
36.79
FcdAvg+FT 18
0.41
28.36
0.46
31.91
0.69
45.04
FedProx 73
0.43
21.92
0.40
26.74
0.67
39.69
FedBE [1
0.43
29.75
0.38
27.54
0.50
34.66
pFedGP
0.15
22.99
0.14
16.57
0.12
17.40
Reptile
0.73
16.53
0.77
16.19
0.77
50.52
MAML 21
0.40
27.74
0.70
48.16
0.75
16.47
PerFedAvg (HF-MAML)
23
0.65
45.44
0.40
23.74
0.69
45.27
FedAvg+MetaaDD (ours)
0.26
19.99
0.32
26.60
0.39
25.27
Reptile+MetaaD (ours)
0.37
25.24
0.32
23.86
0.57
42.40
MAML+MetaVD (ours)
0.31
23.18
0.38
32.56
0.52
37.26
PorFedAvg+MettaVD (ours)
0.26
20.46
0.39
27.59
0.43
30.20
Table 11: Results of uncertainty calibration scores (ECE and MCE) in the CIFAR-100 dataset. The lower the better.
CIFAR-10 dataset
Hetrogenity
à = 5.0
â 0.5
à = 0.1
Method
ECE (%) MCE (%)
ECE (%)
MCE (%) ECE (%) MCE (%)
FedAvg 11
1.93
28.87
2.39
26.45
3.16
27.37
Fc
18
2.02
20.80
2.71
27.86
3.20
33.22
FedProx 73
2.01
25.96
2.52
25.24
3.06
28.62
FedBE [3
1.78
21.33
2.18
21.89
3.35
29.72
PFcdGP
33
1.23
8.34
0.91
6.23
0.84
6.08
Reptile 22
2.40
30.53
3.06
38.89
2.75
27.76
MAML 21
2.61
36.86
2.49
30.67
2.66
32.10
PerFedAvg (HF-MAML)
2.19
26.78
2.39
27.13
2.58
24.47
FedAvg+MetaVD (ours)
1.52
18.02
1.76
21.17
1.40
13.92
Reptile+MetaVD (ours)
1.60
25.46
2.16
29.81
1.76
17.81
MAML+MetaVD (ours)
1.39
16.46
1.97
21.51
0.35
5,80
PorFedAAg++Metav (ours)
1.14
17.55
1.51
20.38
0.29
3.63
Table 12: Results uncertainty calibration scores (ECE and MCE) in the CIFAR-10 dataset. The lower is the better.
Expected Calibration Error (ECE) Maximum Calibration Error (MCE) are widely used metrics measuring uncertainty calibration. ECE represents the average discrepancy between the model's confidence and its accuracy. To compute ECE, we group samples based on their confidence levels and calculate the average confidence and the percentage of correct samples for each group. ECE is then derived as the weighted average of the differences
21

---

between the average confidence and the percentage of correct samples across all groups. ECE values range from 0 to 1. where 0 signifies perfect calibration and I indicates complete miscalibration. MCE, the other hand, is akin to ECE but focuses on the largest gap for any group rather than the weighted average.
Table IT and Table 12 summarize the ECE and MCE results for varying non- degrees CIFAR-ION and CIFAR-IO datasets, with Dirichlet parameters represented as & = [5, 0.5, 0.1]. Typically, ECE and MCE values tend to increase as clients possess more non-IID data, which is a result of overfitting each client's i local data. In the CIFAR-10 dataset, all methods that incorporate MetaVD exhibit decline in ECE and MCE values à decreases, in contrast standard FL PFL methods that show an increase in these values. MetaVD effectively mitigates overfitting to local data and enhances calibration by leveraging the heterogeneity present the training data. Consequently. methods employing MetaVD consistently achieve the lowest ECE and MCE values in almost all scenarios.
Examining reliability diagram can provide visual comparison of uncertainty calibration results, as it illustrates the relationship between prediction probabilities and true labels. Ideally. if a model predicts specific class with a certain probability, the actual label should correspond to that confidence level, placing the point the reliability diagram's 's diagonal line, However, if the prediction probability and the true label do not align, the point would be below or above the diagonal line. In this context, ECE represents the average distance of cach point from the diagonal line, while MCE indicates the maximum distance of any point from the diagonal line, offering comprehensive understanding of the model's calibration performance.
Figure 4 and Figure present the reliability diagrams for CIFAR- 100 and CIFAR-10 with & 0.5 respectively. These ligures allow us to compare the calibration performance of FedAvg. Reptile, MAML, and PerFedAvg. both with and without the integration of MetaVD. Remarkably, in most instances, employing MetaVD leads to improved calibration, drawing the reliability diagrams closer to the identity function. This supports our findings that Meta significantly improves the uncertainty calibration of the models well.
J Additional Results on Muti-domain Datasets
CelebA CIFAR- (/FAR-100 dataset
Hetrogenity
5.0
0.5
0.1
Method
Test (%) OOD (%)
A
Test (%)
OOD (%)
A
Test (%) OOD (%)
A
FedAvg
43.42
44.45
+1.03
43.65
43.45
-0,20
40.46
41.00
+0.54
Reptile
22
50.40
49.07
-1.33
48.92
48.93
+0.01
43.41
42.16
-1.25
MAMI
21
49.97
48.40
-1.57
47.39
48.51
+1.12
44.80
44.79
-0.01
PerFedAvg (HF-MAML) 23
50.61
49.61
-1.00
49.21
50.91
+1.70
44.20
44.13
-0.07
FedAvg+MetaVD (ours)
48.00
48.88
+0.80
48.23
48.62
+0.39
42.98
45.19
+2.21
Reptile+MetaVD (ours)
53.04
53.97
-0.93
52.26
54.75
+2.49
45.83
47.05
+1.22
MAML+MetaVDD (ours)
52.82
53.47
-0.65
51.31
52.82
+1.48
46.83
48.33
+1.50
PertiedAvg+MetaaiD (ours)
53.14
54.26
+1.12
51.18
53.06
+1.88
46.65
47.11
+0.46
Table 13: Classification accuracies with different (non-IID) heterogeneity degrees of à
[5.0, 0.5, 0.1] in multi-domain datasets (a); CelebA CCFAAR-100
CIFAR-100 + FEMNIST dataset
Hetrogenity
à = 5.0
à 0.5
0.1
Method
Test (%) OOD (%)
A
Test (%)
OOD (%)
A
Test (%) OOD (%)
A
FedAvg Reptile
63.90
53.57
100.3
64.02
52.88
-11.14
62.13
49.65
-12.48
66.47
57.60
-8.87
66.13
56.22
-9.91
64.10
55.81
-8.29
MAML
21
66.86
53.25
13.61
66.56
55.72
-10.84
64.56
56.21
-8.35
PerFedAvg MAML) 23
67.47
57.15
10.33
66.57
55.24
-11.3
65.16
56.70
-8.46
Fedavyg+MetaVD (ours)
65.91
57.59
-8.32
65.58
56.85
-8.73
63.26
52.78
10.48
Reptile+MetaVD (ours)
68.61
60.95
-7.66
68.35
59.07
-9288
65.29
58.11
7.18
MAML-MetaVD (ours)
68.41
58.86
-955
68.24
61.21
-7.03
65.46
58.07
-7.9
PerFeddAng-MetaaD (ours)
68.27
58.72
-955
67.93
61.32
6661
66.81
59.13
-7.68
Table 14: Classification accuracies with different (non-lID) heterogeneity degrees of à [5.0 0.5, 0.1] in multi-domain datasets (b); CIFAR-I00 + FEMNIST.
Unlike existing federated learning algorithms that usually assume a dingle-domain approach where only one dataset is used in the experiment. In this section, we further evaluate the performance of our method on real-world non-IID FL experiments by introducing
datasets in which assume each client can have data
22

---

CelebA CIFAR-100 FEMNIST dataset
Hetrogenity
d 5.0
á= 0.5
à == 0.1
Method
Test (%) OOD (%)
Test (%) OOD (%)
Test (%) OOD (%)
A
FedAvg 1
63.57
53.10
10.47
63.73
55.55
-8.18
61.86
52.98
-8.88
Reptile [22
65.93
57.98
-7.95
66.98
57.12
-9.86
63.89
52.33
- -11.56
MAML
21
67.70
59.06
-8.64
67.08
58.15
-8.93
65.67
54.41
- 11.26
PerFedAvg (HF-MAML) 23
67.39
59.17
-8.22
67.20
57.03
10.17
65.08
54.54
10.5-
FedAvg+MetaVD (ours)
63.62
55.87
-7.75
65.93
58.38
- 7.35
64.59
55.61
88.98
Reptile+MetaVI) (ours)
67.63
60.95
- 6..8
68.78
61.39
-7.19
65.18
56.46
-8.72
MAML=MetaVD (ours)
69.42
61.52
-7.90
68.81
61.60
-7.21
66.56
56.54
- -10.02
PerFedAAg4MMeuVDD (ours)
67.83
60.74
- -7.99
68.05
61.17
6.88
65.11
55.64
-9.47
Table 15: Classification accuracies with different (non-III) heterogeneity degrees of à =
[5.0, 0.5, 0.1] in multi-domain datasets (d); CelebA + CIFAR-100 FEMNIST.
from different domains. Multi-domain learning 87 87-89) aims to leverage all available training data across different domains to enhance the performance of the model. However, directly utilizing data from different domains typically results in suboptimal global model. In the context of federated learning. it becomes even critical apply multi-domain learning effectively. We use three different FL datasets construct the multi-domain task distributions: CelebA, FEMNIST, and CIFAR-100. The non-IID heterogeneous environment has been also assumed in the field of multi-domain learning. We utilize the Dirichlet sampling technique (á = [5.0, 0.5, 0.1) to sample cach client's 's local CIFAR-100 data.
Table 13 14 and 15 llsstrrte the classification accuracies for Test and OOD clients with varying degrees of heterogeneity. Employing MD leads to significantly improved prediction accuracy, with larger improvements in OOD accuracy compared to the improvement in Test accuracy, across all multi-domain settings and degrees of heterogeneity. The results indicate that Me can improve the versatility and performance of algorithms when applied to a broad range of multi-domain datasets.
K
Additional Results on Model Compression
CIFAR-10 R0 dataset
Hetrogenity
á 5.0
á 0.5
à 0.1
Method
Test (%) OOD (%)
Sparsity (%)
Test (%) OOD (%) Sparsity (%) Test (%)
OOD (%) Sparsity (%)
SNIP
81.03
80.46
70.0
78.53
79.30
70.0
71.67
73.13
70.0
Reptile-SNIP
83.21
82.10
70.0
81.43
81.72
70.0
74.56
69.57
70.0
MAMI+SNIP
83.35
81.59
70.0
81.47
81.84
70.0
72.51
69.27
70.0
PerFelAvg+SNIP
83.33
83.49
70.0
80.97
82.47
70.0
76.27
75.90
70.0
Reptile-MetaVD
84.70
84.28
0
83. 20
83.40
0
76.51
82.07
MAML+MettVD
83.93
84.95
0
81.32
81.81
0
77.27
79.05
PefFedAng=MeeaVD
83.88
83.96
81.06
81.47
76.06
81.77
Reptlle-MeaVDDDDD
84.19
84.46
77.12
82.67
82.79
79.25
75.85
80.52
76.18
MAML-MMeaVD-DDP
84.37
84.95
76.76
82.65
83.64
78.85
77.77
80.74
72.47
PofFddAw++MetaV--DDP
84.04
83.69
78.55
82.64
83.52
79.32
77.10
83.24
72.87
Table 16: Results of study on model compression with different (non-IID) heterogeneity strengths of = 5.0, à = 0.5, and = 0.1 in the CIFAR-IO dataset. Meta VD+DP does not communicate the model parameters when the dropout rate is larger than 0.9.
Federated Learning optimization involves frequent communication of model parameters between devices and the central server, which can be slow and may raise privacy concerns. Therefore, is erucial to minimize communication costs by reducing both the size of exchanged model parameters and the number of communication rounds. MD also has the benefit of compressing the model parameters needed for each client device. To explore the compression capabilities of MetaVD, we conducted experiments on CIFAR-I 10 and CIFAR- 100 datasets, comparing
method to the baseline Pruning algorithm SNIP [I0I]. For the implementation of the SNIP algorithm in the context of Federated Learning, referred the work of Jiang et al. (2022) 102 we let
SNIP prune the original model to the target sparsity right after the first round.
Table 16 and Table demonstrate the results of Test(%), OOD(%), and Sparsity(%) percentage of the models with or without
Sparsity represents the ratio of zero- valued model parameters in the personalized layer. The DP refers to the process of dropping communication of weights in the FL algorithm. where weights with dropout rate greater than 0.9 dropped after 150 rounds. We have set these 150 rounds to allow MetaVD
23

---

CIFAR-I00 datasct
Hetrogenity
=
à 0.5
0.1
Method
Test (%) OOD (%) Sparsity (%) Test (%)
OOD (%) Sparsity (%) Test (%)
OOD (%) Sparsity (%)
FodAvg+SNIP
43.76
44.11
70.0
42.53
41.54
70.0
37.80
40.34
70.0
Reptile+SNIP
49.54
49.61
70.0
48.38
48.62
70.0
42.03
43.68
70.0
MAML+SNIP
51.06
49.56
70.0
48.16
46.99
70.0
42.03
42.49
70.0
PorFedAvg+SNIP
50.66
48.85
70.0
48.71
48.46
70.0
43.23
43.09
70.0
Reptile-MetaVD
53.71
54.50
0
52.06
51.50
44.56
44.62
MAML+MctaVD
52.40
51.78
50.21
49.75
43.84
48.17
PerFedAvg+MetaaVD
51.67
51.70
0
50.02
48.70
43.31
46.19
Reptioo-MssaDDDDP
52.56
55.03
65.83
50.94
50.85
57.26
45.13
45.19
75.72
MAML-MMctVDD+DP
52.31
52.29
74.61
50.30
50.88
77.88
45.75
45.78
72.26
RefealAvg--MtaaDDDDD
51.68
52.56
70.08
51.64
51.47
68.88
45.16
46.45
55.18
Table 17: Results of study on model compression with different (non-IID) heterogeneity strengths of 5.0, d: 0.5, and à 0.1 in the CIFAR-100 dataset. MetaVDDDP does not communicate the model parameters when the dropout rate is larger than 0.9.
to adequately leam the spent--pe dropout rates. The target sparsity the SNIP algorithm was set 0.7. Note that, in our Meta VDDP setting, even though we dropped weights with dropout rate greater than 0.9, the actual sparsity values observed be lower and tend to be around 70%, indicated in the tables.
Overall, most models removed 70% of their parameters in the personalized layer without losing much performance. In CIFAR-IO, heterogeneity increased, the performance gap between SNIP and Mcta tended increase. For instance, CIFAR-10 1-10 with MAMM.+MetaVD-DHH
5.0, OOD performance of SNIP model and model were .55 and 84.95, respectively, whereas in CIFAR-10 with á = 0.1, they were 69.27 and 80.74 respectively, demonstrating an increased difference in performance. This was similar in the of Reptile+MetaVDDDD gap, according could lead to performance improvements in some cases. For example, CIFAR-100 with 0.5, the PerFc- dAvg+Meta VD+DP model's 's OOD performance increased from 48. 70 to 51.47. This experiment demonstrates that the MetaVD method can decrease the communication cost in Federated Learning (FL) by compressing model parameters. This emphasizes the efficiency of MetaVD, this approach without significantly increasing communication costs. It also boosts generalization performance for new, unseen clients. These aspects make our approach significant addition to the field Federated Learning.
and PegFedAvg+ Meta-+MMMaDDDDDD In CIFAR-100 dataset, the performance increase heterogeneity, was not noticeable. Surprisingly. dropping model parameters
be applied numerous FL algorithms
24

---

(a) FedAvg
(b) FedAvg + MetaVD
1.0
ECE 0.53%
ECE = 0.32%
MCE 32.35%
MCE 26.60%
0.8
0.8
0.6
0.4
0.4
0.2
0.2
0.0.0
0.0.0
0.2
0.4
0.6
0.8
1.0
0.2
0.4
0.6
0.8
1.0
(c) Reptile
(d) Reptile MetaVD
1.0
1.0
ECE 0.77%
ECE 0.32%
MCE 46.19%
MCE 23.86%
0.6
0.6
0.4
0.4
0.2
0.2
Accuracy
0.0.0
0.0.0
0.2
0.6
0.8
1.0
0.2
0.4
0.6
0.8
(e) MAML
(1) MAML + MetaVD
1.0
1.0
ECE
0.70%
ECE = 0.38%
MCE 48.16%
MCE 32.56%
0.8
0.8
0.6
0.4
0.4
0.2
0.2
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.0
0.2
0.4
0.6
0.8
1.0
(g) PerFedAvg
(h) PerFedAvg MetaVD
1.0
1.0
ECE 0.40%
ECE 0.39%
MCE 23.74%
MCE 27.59%
0.8
0.6
0.6
0.4
0.4
0.2
0.2
0.0.0
0.8.0
0.2
0.6
0.8
1.0
0.2
0.4
0.6
0.8
1.0
Confidence
Figure 4: Reliability diagrams for (a) FedAvg, (b) MetAvg+ (c) Reptile, (d) Mettile+MetaVD. (e) MAML, (f) MAML+MetaVD, (g) PerFedAvg and (h) PerfeedAv--MetaDD in CIFAR-100
(à = 0.5).
25

---

(a) FedAvg
(b) FedAvg + MetaVD
1.0
ECE 2.39%
ECE = 1.76%
MCE 26.45%
MCE 21.17%
0.8
0.8
0.6
0.4
0.4
0.2
0.2
0.0.5
0.0
0.2
0.4
0.6
0.8
1.0
0.0.5
0.0
0.2
0.4
0.6
0.8
1.0
(c) Reptile
(d) Reptile MetaVD
1.0
1.0
ECE 3.06%
ECE 2.16%
MCE 38.89%
MCE 29.84%
0.6
0.6
0.4
0.4
0.2
0.2
Accuracy
0.0.0
0.0.0
0.2
0.6
0.8
1.0
0.2
0.4
0.6
0.8
1.0
(e) MAML
(f) MAML MetaVD
1.0
ECE 2.49%
ECE = 1.97%
MCE 30.67%
MCE 21.51%
0.8
0.8
0.6
0.4
0.4
0.2
0.2
0.0
0.0.0
0.2
0.4
0.6
0.8
1.0
0.2
0.4
0.6
0.8
1.0
(g) PerFedAvg
(h) PerFedAvg MetaVD
1.0
1.0
ECE 2.39%
ECE 1.51%
MCE 27.13%
MCE 20.38%
0.6
0.6
0.4
0.4
0.2
0.2
0.0.0
0.8.
0.2
0.6
0.8
1.0
0.2
0.4
0.6
0.8
1.0
Confidence
Figure 5: Reliability diagrams for (a) FedAvg, (b) MedAvg+MetaVD, (c) Reptile, (d) Mettile+MetaVD. (e) MAML, (f) MAML+MetaVD, (g) PerFedAvg and (h) PerFedAvg+MettWD in CIFAR-10 0.5).
26


Reviews:
[{"review_id": "ytBJxDElrj", "reviewer": "Reviewer_6ZgF", "summary": "The paper introduces a novel approach called meta-variational dropout (MetaVD) for addressing challenges in federated learning (FL). Traditional FL faces issues such as model overfitting and divergence of local models due to non-i.i.d. data across clients. MetaVD leverages Bayesian meta-learning to predict client-specific dropout rates using a hyper network. Extensive experiments conducted on various FL datasets demonstrate that MetaVD achieves good performance.", "questions": "(see \"weaknesses\" above).", "limitations": "The authors haven’t discussed the limitations of their work.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "1. The proposed approach is novel.\n2. Authors performed experiments on a variety of tasks and datasets and showed promising results.\n3. The proposed method showed consistent improvements over the considered baselines with a good margin.", "weaknesses": "A primary weakness is that the authors have not compared their results with existing state-of-the-art personalized FL algorithms, e.g., [1,2] and the related work/baselines referenced within them.\n\nOther notes and requests for clarification:\n1. In L66, where the authors discuss the convergence guarantees of the original FL algorithm, it is important to provide a citation.\n2. The authors have written about the “Challenges of FL” in Section 2. They should briefly describe how the proposed approach addresses these challenges.\n3. Figure 1 requires additional details to aid understanding. The authors should clarify the meaning of the \"x\" in the modules of the figure and explain how it illustrates the difference in aggregation between MetaVD and FedAvg.\n4. In Table 3, the signs of \"+\" and \"-\" should be reversed to ensure accurate representation.\n5. Table 4 should explicitly mention that the results presented are for out-of-distribution (OOD) clients to provide a clear context for the findings.\n6. Figure 1 is not referenced or discussed in the paper. It should be either referred to in the main text or removed if it does not contribute significantly to the paper's content.\n7. Additional information regarding the computation of gradients for the hypernetwork parameters should be included.\n8. The authors should provide an explanation as to why the PerFedAvg+MetaVD+DP combination in Table 7 leads to improved performance despite dropping almost 80% of the parameters.\n9. Regarding equation 5, the authors should elaborate on the implications of the inverse dependence of the aggregation weights on the square of the model weight. This information would enhance the understanding of the aggregation process and its impact on the final results.\n\n\n[1] “Fusion of Global and Local Knowledge for Personalized Federated Learning”, TMLR 2023.\n\n[2] “FedALA: Adaptive Local Aggregation for Personalized Federated Learning”.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the encouraging feedback on our work!\n\nWe appreciate introducing the important PFL works (FedSLR, FedALA). We will cite them in the paper. \n\n**PFL Baseline**\nWe summarize the baselines and related works **presented** in the submitted paper\n|**Method**|category|status|\n|-----------|--------|------|\n|FedAvg|FL|baseline|\n|FedProx|FL|baseline|\n|FedAvg+FT [1]|Personalized FL|baseline|\n|(pFed) Reptile [2]|Personalized FL|baseline|\n|(pFed) MAML [3]|Personalized FL|baseline|\n|PerFedAvg [4]|Personalized FL|baseline|\n|pFedHN [5]|Personalized FL|cite|\n|FedAG |Bayesian FL|cite|\n|FedBE |Bayesian FL|baseline|\n|FedPA [11]|Bayesian FL|cite|\n|FedAVG + SNIP [9]|FL + Pruning|baseline (in appendix)|\n|MetaVD (ours)|Bayesian PFL|ours|\n\n* Reptile [2], MAML [3], and PerFedAvg [4] are the SOTA PFL baselines extended from few-shot/meta-learning approaches. We chose these meta-learning-based PFL approaches as our key baseline since they perform well on out-of-distribution (OOD) clients. A summary of baselines is presented in section D of the appendix. \n\n* We extensively analyzed all baselines to find the best parameter settings using a hyperparameter search tool, Optuna. We conducted experiments on various FL algorithm evaluation scenarios following the recent FL benchmark paper, pFL-Bench, which makes our approach comparable to 20 existing SOTA pFL methods experimented in the pFL-Bench paper. For example, our results in Table 4 on FEMNIST are comparable to Table 2 of the pFL-Bench paper. In the CIFAR-10 and CIFAR-100 experiments, the sample sizes for participating and non-participating clients are slightly different; otherwise, they are almost the same. \n\nHere are answers to your questions:\n\n**1. L66.**\n\nIn L22, we cited two references; one discusses the convergence rate of FedAvg with IID data, and the other is a survey paper.\n\n**2. “Challenges of FL” in Section 2.**\n\n* Heterogeneity of client data: MetaVD is a technique of modulating a global NN parameter (of FedAvg, or (pFed) Reptile/MAML) with a conditional dropout technique. Predicting the client-specific dropout variable via a hypernetwork enables reconfiguring a single NN for various different tasks. We also utilized the client-specific dropout rate as a weighting factor in the global model aggregation.\n\n* Sparse connectivity (or client’s participation in training): Table 5 in our experiment illustrates that reducing the participating client size does not degrade the prediction scores when we applied MetaVD. We hypothesize this is partially due to the conditional dropout preventing the overfitting into a small subset of clients.\n\n* limited data: The initial parameters learned in meta-learning help predict a local model with only a few data adaptations; thus, meta-learning-based PFL is advantageous to the client with limited data. The conditional dropout also prevents the local model overfitting by additionally regularizing (or conditioning) the initial parameters. \n\n* Communication cost: We demonstrated the possibility of pruning the model parameter in FL communication.\n\n**3. Figure 1 requires additional details.**\n\nWe will update Figure 1 as follow:\n\n* We will add the meaning of “X” in Figure 1: it represents the pruned parameters in each communication round.\n* When aggregating local parameters, MetaVD utilizes the model’s weights as well as the weights’ uncertainty to find its global optimum. This corresponds to the global posterior mean of the product of two local Gaussian posteriors. In contrast, FedAVG only considers the mean of each local model. We will add this description with more details to the figure to explain the difference in the global aggregation between them. \n\n**4. The reversed sign in Table 3.**\n\nThank you! We will reverse the sign.\n\n**5. Specify OOD clients in Table 4.**\n\nWe will update the phrase (L261) in the caption.\n\n**6. Reference Figure 1.**\n\n(see “3.” above)\n\n**7. The computation of gradients for the hypernetwork parameters.**\n\n(Please see ``the update rule for hypernetwork'' in the general comment)\n\n**8. Why the PerFedAvg+MetaVD+DP combination in Table 7 leads to improved performance.**\n\nThe performance enhancement after pruning our model aligns with Occam's Razor, which favors simpler explanations and helps prevent overfitting. This principle explains why pruning or dropout training can increase performance, a phenomenon also observed in existing works [9, 13]. A comparison with the SNIP pruning algorithm [9] is included in the Appendix on page 9.\n\n**9. The implications of the inverse of $\\theta^2$ in Eq 5.**\n\nIn our work, the $m$-th client's posterior distribution is defined as Gaussian dropout posterior [10,12] of the form $q(w^m) = \\mathcal{N}(w^{m} \\vert \\theta, \\alpha^{m} \\theta^2)$ where a mean equal to the global NN's weight $\\theta$, and variance is the square of the weight (i.e., $\\theta^2$) times the client-specific dropout variable $\\alpha^m$. Thus, the $\\theta^2$ emphasizes the effect of larger weights on the estimated variance $\\sigma^2 = (\\alpha^m)\\theta^2$. \n\nThe $\\alpha^m(\\theta^m)^2$ in Eq 5 describes the estimated uncertainty in the model parameter for the $m$-th client. The inverse of this term is used to compute the aggregation weights, $r^m$. This means that parameters with more uncertainty (larger values or higher dropout rates) will have less influence on the aggregated global model parameter $\\theta^{\\text{agg}}$. Based on this technique, the federated learning system can adaptively adjust each client's local model's contribution to the global model. This method can help reduce the impact of outliers or anomalous clients that may have deviating model parameters due to unique local data distributions or noise. Also, in statistical learning theory, the Fisher information is directly proportional to the precision (or inverse variance) of Gaussian distribution.\n\n**Reference**\n(Please use the reference in the general comment)"}, {"review_id": "k6Cxb91SfC", "reviewer": "Reviewer_1ob4", "summary": "This paper proposes a novel Bayesian personalized federated learning approach using meta-variational dropout. The proposed approach employs a shared hypernetwork to predict the client-dependent dropout rates for each model parameter, enabling effective model personalization and adaptation in the limited non-i.i.d. data environment. The effectiveness of this approach is demonstrated empirically on various FL datasets, including CIFAR-10, CIFAR-100, FEMINIST, and CelebA and multi-domain FL datasets.", "questions": "1. Page 3, line 115: i.e., $\\phi=\\phi^1,..,\\phi^M$ should be $\\phi=\\{\\phi^1,..,\\phi^M\\}$.\n\n2. Page 3, line 132: add a definition of $\\mathbf{1}$ for  $\\epsilon^m \\sim \\mathcal{N}(\\mathbf{1}, \\alpha^m)$ .\n3. In equation (5), how was the aggregation rule derived? It is recommended to provide more details (referring to [4] is possible).\n\n4. It would be clearer to specify the specific sections of the Appendix. For example, on page 6, \"More details are in Appendix.\"\n5. The sparsity in Table 7 seems to only consider the personalized layer (the last fully connected layer). I am curious about the proportion of discarded parameters to the total parameters of the model.\n6. I hope the authors can provide a discussion on the communication complexity of the additional hypernetwork.\n7. Algorithm 1 seems confusing. It would be clearer to separate Algorithm 1 into two distinct algorithms. \n\n\n\n[4] A Bayesian Federated Learning Framework with Online Laplace Approximation.", "limitations": "NA.", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is clearly written and well organized.\n\n2. The proposed method in the paper is well-motivated and technically correct.\n\n3. The Bayesian approach to FL is interesting and seems suitable for personalization. The introduced MetaVD seems to work well in practice.\n\n4. The method is extensively tested on a variety of FL datasets.\n\n", "weaknesses": "1. The discussion with related work needs to include other Bayesian treatments for personalized FL, e.g. [1], [2]\n\n2. The experiments lack a comparison with relevant methods mentioned above [1-2], as well as the baseline in [3].\n\n[1] Personalized Federated Learning via Variational Bayesian Inference. ICML 2022.\n\n[2] Fedpop: A Bayesian approach for Personalised Federated Learning. NeurIPS 2022.\n\n[3] Personalized Federated Learning using Hypernetworks. ICML 2021.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**We thank the reviewer for the positive and encouraging feedback!**\n\nAs the review suggested, the recommended recent (Bayesian) PFL works will be thoroughly compared in the Background section. \nWe will update the difference between the recent PFL approaches [1,2,3,4] and ours. \nWe will also add the experimental results of pFedGP [17] as a Bayesian PFL baseline, as we could not find officially released codes for the other methods. \n\nHere are answers to your questions:\n\n**1. The correction for L115.**\n\nThank you for the correction.\n\n**2. A definition of for  $\\epsilon \\sim N(\\mathbf{1}, \\alpha)$ in L132.**\n\n$\\mathbf{1}$ is a K dimensional all-ones vector. K is the dimension of the NN's weight. ($\\vec{\\mathbf{1}}$ might be appropreate.)\n\n**3. The derivation of the aggregation rule.**\n\nThank you for the suggestion! Please see \"the derivation of the aggregation rule\" in the general comment.\n\n**4. Specify the specific sections of the Appendix.**\n\nWe will add the section number for all the references to Appendix.\n\n**5. The sparsity in Table 7.**\n\nYes, it is applied to the personalized layer before the last layer. Here is the computation of the proportion: 12900 (params of a personal FC layer) * 0.8 (dropout rate) / 416612 (total params) * 100 = 2.48 (\\%) approximately. We are showing the possibility of not increasing the model size while incorporating the additional model uncertainty parameter in Bayesian PFL.\n\n**6. Communication complexity of the additional hypernetwork.**\n\nThe server only sends the global parameter and the approximated dropout variable to the client while keeping the hyperparameter in the server. Hence, if we apply dropout to the full NN layers, the communication complexity of our work is $\\mathcal{O}(2K)$ where K is the dimension of full NN parameters. This doubles the cost compared to FedAVG. This is due to the modeling of uncertainty in the Bayesian framework similar to predates. However, our MetaVD is applied to just a single layer before the last output and also can prune some of the parameters. The increase in communication costs could be negligible. \n\n**7. Algorithm 1 seems confusing. It would be clearer to separate Algorithm 1 into two distinct algorithms.**\n\nThank you for the suggestion. We will try to optimize the readability of Algorithm 1. \n\n**Reference**\n\n(Please use the reference in the general comment)"}, {"review_id": "ePv9tQH3Rk", "reviewer": "Reviewer_yoXe", "summary": "This paper proposes a new federated learning (FL) framework to address the various issues of FL. Specifically, this framework (a) leverages Bayesian FL to address the issue of non i.i.d. data among clients, (b) instantiates its BFL model with a variational dropout posterior to efficiently handle large amount of clients, and (c) applies meta learning adaptation strategies (e.g., MAML and Reptile) at client-sides to personalize the local models. The paper demonstrates the performance of this approach on several vision datasets, with ablation studies to confirm the effectiveness of proposed components.", "questions": "- What is the benefit of the hyper-network, as opposed to directly learning the client-specific dropout rates ?\n- What is the overhead computation cost of the proposed framework? I believe both local meta-learning and training a server-side hypernetwork is quite costly, and it is quite unfair to compare 1000 comm rounds of MetaVD to 1000 comm rounds of FedAvg.", "limitations": "The authors have discussed some limitations of the work in the conclusion. I think the method is purely theoretical and there is no potential negative societal impact.", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 2, "strengths": "The proposed method makes sense and addresses various relevant problems in the field of FL. I think the paper did an excellent job combining various ideas, and presenting them as one coherent framework. In addition, the empirical results positively improve over existing baselines, which suggests strong practical merits. I also like that the paper is really rigorous with conducting ablation studies.", "weaknesses": "The most apparent weakness of this paper is the lack of technical novelty. Bayesian FL, Variational Dropout, and Meta-learning are all very well-known strategies. While I am not opposed to a creative combination of ideas, such idea should have a scientific merit that outweighs the sum of its individual components. Here, each component is behaving exactly as expected and there seems to be no technical challenge in terms of integrating them. I think the novel adjustment here is the hypernetwork (although I am not sure if this strategy has been previously adopted in the context of variational dropout --- e.g., hierarchical prior is sort of similar, but not exactly the same).\n\nSome other weaknesses:\n- Lack of error bar/standard deviation in any of the reported results.\n- I would prefer some sort of performance vs. communication rounds to demonstrate convergence. Result tables are acceptable, but very uninformative.\n- It was previously claimed that the hypernetwork helps with large number of clients, especially when some of them have limited data. I think this should be supported by an ablation study showing the effect of increasing no. clients/ decreasing client data (the current ablation study only compares MetaVD to NaiveVD/EnsembleVD on the default client setting).\n\nSome other minor issues:\n- Various typos (e.g., FEMINIST, line 187; Hetrogenity, table 2 header ...)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Thank you for the helpful review!**\n\nIn this work, we extended the hypernetwork to approximate the dropout variable of NN's weight to enable the full utilization of the weight uncertainty in the aggregation, regularization, and personalization of FL. Our efficient composition is also implemented on various meta-learning based baselines and PFL benchmark datasets.\nHere we briefly summarize the contribution of our works.\n\n**Contribution**\n* In FL, we show a new hyper-network-based Bayesian network approach is feasible. \n* Our method is the first approach of employing the variational dropout uncertainty in the posterior aggregation for the PFL.\n* We show that the conditional Gaussian (dropout) noise injection to NN weight method is a versatile approach that can be combined with any optimization-based meta-learning approach. \n* Experimentally, we have verified the performance of our work in various PFL benches. We release the experiment code for all baselines for reproducibility. We put much effort into implementing the stable gradient propagation to hypernetwork and the global parameter of meta-learning algorithm using the PyTorch functorch library.\n\n**Lack of error bar/standard deviation.**\nWe presented the results of baselines with the best parameter settings optimized via a hyperparameter search tool, Optuna. We performed an extensive analysis of various different FL scenarios.\n\n**An ablation study showing the effect of increasing no. clients/ decreasing client data.**\nSome portions of clients have small datasets in the non-IID setting. They increase with a higher alpha (Dirichlet parameter). The reason we reported the simple ablation study is that the EnsembleVD cannot actually learn the independent dropout parameters well. So, technically, EnsembleVD was reduced to FedAVG with Dropout(0.1).\n\n**1. The benefit of the hypernetwork.**\n\nIn our experiment, the dropout rates in EnsembleVD were not well optimized due to restricted client participation. On the other hand, joint optimization of MetaVD with hypernetwork converges well. This can be observed via the changes in the KL divergence loss term.  This demonstrates that MetaVD's hypernetwork offers a more data-efficient approach to learning the client-specific model uncertainty than Ensemble VD.\n\n**2. What is the overhead computation cost of the proposed framework?**\n\nWe found that all models converged in 1000 rounds in our experiment. The only difference between MetaVD and FedAVG is one layer before the last layer. Since we only applied the conditional dropout to only the last layer, the increase in computation cost is quite small. \nThe actual (approximate) training time in Cifar-10 dataset is  4 hours for  FedAVG, 5 hours for MetaVD, 6 hours for PerFedAvg (pFedGP in our experiment seems to take almost 12 hours)."}, {"review_id": "TyVI4LXY3u", "reviewer": "Reviewer_mtpr", "summary": "The submission proposed a novel Bayesian meta-learning approach metaVD for federated learning. metaVD learns to predict client-dependent dropout rates via a hypernetwork, helping address the model personalization and limited non-i.i.d. data problems. At the same time, metaVD compressed the model parameter, alleviated the overfitting and reduced the communication costs.", "questions": "1. “Recently, the Bayesian learning paradigm was introduced to the FL to tackle overfitting by considering the uncertainty of the model parameters [28–30]. However, they could also struggle with diverging local models if the data from different clients exhibit significant statistical variability.” This statement is incorrect. In recent years, Bayesian federated learning has made significant advancements. Many Bayesian methods have been employed not only to address overfitting and introduce uncertainty but also to tackle personalized scenarios. I list some works as reference below. \n\n2. \"Moreover, previous Bayesian FL methods are not developed for model personalization, and the performance can be degraded in non-i.i.d. client data.\" This statement is incorrect for the above reasons. \n\n3. Why was the variational distribution in Equation 3 designed in such a form? I understand that this is an engineering approach, but can this parameterized form effectively approximate the posterior distribution of the weights? Is there any theoretical analysis available? Furthermore, why is the variance parameter $\\alpha_m$ used to model the dropout rate? Is there any inherent connection between the two? While this may be an engineering approach, is there any theoretical or intuitive analysis supporting it? It is important to note that while engineering approaches are often driven by empirical performance, they can also be guided by theoretical insights or intuitions. \n\n4. In eq(5), what is the definition of $g_m$? How does the expression of $r_k^m$ come out? Any derivation?\n\n5. In the server-side optimization, when you update the parameters $\\psi,e^m$, you are trying to optimize the $L_{ELBO}^m$ w.r.t. $\\psi$ and $e^m$, but $L_{ELBO}^m$ requires the data on each client, doesn't this imply data leakage? \n\n6. In the experiment section, as I mentioned earlier, this article lacks thorough investigation into Bayesian federated learning (FL), as it fails to mention and compare many Bayesian FL approaches specifically designed to address personalized scenarios. To name a few examples but not exaustive: \n[1] Personalized federated learning with gaussian processes, NIPS 2021.\n[2] Federated Bayesian Neural Regression: A Scalable Global Federated Gaussian Process. H Yu et al. 2022. \n[3] Personalized federated learning via variational bayesian inference, ICML, 2022.", "limitations": "The authors discussed the limitations of their work, but they did not address the potential negative societal impact it may have.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "MetaVD encompasses a new posterior aggregation strategy to consolidate local models into a global one. In addition, MetaVD predicts the dropout rates of parameters via a hypernetwork, enabling parameter compression. This not only allievates the overfitting problem but also reduces the communication costs of exchanging model parameters.", "weaknesses": "In general, this article is written in a fluent, simple, and easily understandable manner. However, there are two main issues that need to be addressed:\n\n1. The proposed method in the article is straightforward and intuitive, but many aspects lack theoretical guarantees and analysis. \n\n2. In the past years, Bayesian federated learning has made significant progress, but many relevant works have not been discussed or compared in the article. \n\n3. if my understanding is right, the propsoed method has the risk of data leakage.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**We thank the reviewer for insightful comments.**\n\nHere are answers to your questions:\n\n**1. Bayesian PFL work is not discussed.**\n\nWhen we started our research last year, we were aware of only a few papers [17,15] dealing with the personalization of models using Bayesian methods. Meanwhile, we realized that measuring the algorithm’s performance on unseen clients during training is an important task [5]. In real FL applications, most clients may not participate in training. Then pFL benchmark paper [1] also extensively analyzed the gap between participating and non-participating clients across various pFL scenarios. Thus, we adopted their pFL evaluation protocols for our experiment.\n\nOn the other hand, pFedBayes [15] did not show results for OOD clients. They assume an independent local model for each device while learning a shared global model. In most of their experiments, the performance of local models showed better results than the global model, but there are no local models for unseen clients. pFedGP [17] learns a shared kernel for all clients and infers personalized classifiers for each client. However, the generalization gap in Figure 4 seems large (compared to our results in Table 2). Also, the inducing point (IP) approximation to speed up the computation of inverting a large kernel matrix did not show better results than the original pFedGP [17] in their OOD experiment. \n\nFor this reason, we focused on the meta-learning-based PFL methodology for most of our baseline because this approach is a researched methodology to apply to out-of-distribution (OOD) tasks. Also, we adopted the product of the posterior strategy for the global model aggregation. The posterior aggregation approach considers the uncertainty of weight in the model aggregation. However, pFedBayes and pFedGP update the global parameter (or the kernel) similar to FedAVG. That is why we cited Bayesian FL works related to posterior aggregation at most. \n\nWe had no intention of disregarding the existing Bayesian PFL works. We will add a thoughtful discussion paragraph about the Bayesian PFL works and report the results of pFedGP [17] as a baseline (we could not find the official codes for the other papers). Only one article [19] in the arXiv addresses PFL with the Bayesian meta-learning (updated in July 2023 and not yet officially published). Please note that our work is still novel in this direction as our work is the first approach to utilize the variational dropout uncertainty in the model aggregation in the Bayesian PFL. \n\n**2. The introduction about Bayesian FL.**\n\nThe “previous Bayeain FL” means those approaches mentioned in the same paragraph. We will change the sentence or add an explanation accordingly. (e.g., “Recently, Bayesian PFL methods [17,15,18] are introduced in this field and have been shown effective in dealing with non-IID. However, their approach still has limitations (e.g., Strong constraints in pFedBayes and high computational cost in pFedGP). Their model aggregation rule does not specifically considers the model uncertainty.” )\n\n**3. The variational distribution in equation 3.**\n\nAs in L123, we extended the variational dropout (VD) [10, 21, 12, 22] approach for the local posterior. The VD is a regularization technique in Neural Networks (NN). This involves randomly turning off NN parameters during training by applying discrete Bernoulli random noises; a method initially used to prevent over-fitting [C]. \n\nLater, it was found that continuous noises sampled from Gaussian distributions work similarly to the Bernoulli dropout due to the central limit theorem [A, B]. The $m$-th client's weight is interpreted as a form of Gaussian posterior distribution $q(w^m) = \\mathcal{N}(w^{m} \\vert \\theta, \\alpha^{m} \\theta^2)$ where a mean equal to the global NN's weight $\\theta$, and variance is the square of the weight (i.e., $\\theta^2$) times the client-specific dropout variable $\\alpha^m$. \n\n[A] Fast dropout training. ICML, 2013.\n\n[B] Regularization of neural networks using dropconnect. ICML, 2013.\n\n[C] Dropout: a simple way to prevent neural networks from overfitting. The journal of machine learning research, 2014.\n\n**4. The derivation of the aggregation rule in equation 5.**\n\nThe definition of $g^m$ is a weighting factor of each client proportional to the local data size (mentioned in L56 and L108).  \nWe presented the derivation of the aggregation rule in the general comment.\n\n**5. The risk of data leakage?**\n\nIn the FL system, the server and client only exchange the parameters (summarized statistics) instead of raw data, thus reducing the risk of data leakage. In our study, MetaVD only exchanges the NN parameter $\\theta$ and dropout parameter $\\alpha^m$. In addition, the dropout is applied to the weight of just one layer before the output layer (L202). Therefore, the risk of data leakage is relatively small or negligible compared to other existing FL baselines.\n\nComputing the gradient w.r.t $\\psi$ and $e^m$ happens in the server after the server receives $\\theta^m_*$ and $\\alpha^m_*$. The expression in L171 $\\nabla_\\psi \\mathcal{L}^m_{\\text{ELBO}}(\\alpha^m) \\approx (\\nabla_\\psi \\alpha^m)^{T} \\Delta \\alpha^m$ is an approximation of computing gradient w.r.t $\\psi$. $\\nabla_\\psi \\alpha^m$ is the gradient of a hyper network's output w.r.t $\\psi$.  $\\Delta \\alpha^m = \\alpha_*^m - \\alpha^m$ is the changes in the local dropout variable, approximating the vector-Jacobian product. For this updating rule, we followed the recently proposed PFL algorithm, pFedHN [5] (please see Section 3.2 of them). In this work, we extended the hyper-network to approximate the dropout variable (or conditional variance) to fully utilize the weight uncertainty in the (aggregation, regularization, and personalization of) FL.  \n\n**6. The investigation into Bayesian PFL.**\n\n(Please see \"1\" above and the Bayesian PFL in the general comment)\n\n**Reference**\n\n(Please use the reference list in the general comment)\n\n\n\n\n\n\n"}]

Author rebuttal:
# General comment to all reviewers

We sincerely thank the reviewer for our work's positive and encouraging feedback. 

Here, we summarize some common answers to the reviewers.

## Bayesian PFL
Since we hear the reviewers' suggestion of discussing some recent Bayesian PFL works, we will add a thoughtful discussion paragraph about the recent development in Bayesian PFL works. 
We will remove Table 1 in the Background section and add a paragraph to distinguish our work.
We will report the results of pFedGP[17] as a baseline in the experiment (we could not find the official codes for the others). 
We can find only one paper [19] in the arXiv addressing PFL with the Bayesian meta-learning (updated in July 2023 and not officially published). Our work is still rare in this direction. Also, our work is the first approach to utilize the conditional variational dropout uncertainty in the model aggregation in the Bayesian PFL. 

|**Method**|Category|Status|
|-----------|-----------|-------|
|pFedGP [17]|Bayesian PFL|→ baseline|
|pFedBayes [15]|Bayesian PFL|→ cite|
|FedBNR [18]|Bayesian PFL|→ cite|
|Fedpop [16]|Bayesian PFL|→ cite|
|FedPPD |Bayesian PFL|→ cite|
|FedABML [19]|Bayesian Meta PFL|→ cite|
|MetaVD (ours)|Bayesian Meta PFL|ours|

## The aggregation rule in Eq 5

Eq 5 is derived from the product of Gaussian local posteriors. 

 In general, the mode of the product of $M$ Gaussians, $\prod_{m=1}^{M} \mathcal{N}(\mu_m, \sigma_m)$, simplifies to $\mu_{\text{agg}} = \sum_{m=1}^M r_m \mu_m$  where $r_m = (\sigma^2_m)^{-1} / \sum_{m=1}^M (\sigma^2_m)^{-1}$. 

It is a popular statistical theory [23] (mentioned in L160).

The difference in our approach is that we utilize the Gaussian dropout posterior,  $q(w^m) = \mathcal{N}(w^{m} \vert \theta^m, \alpha^{m} (\theta^m)^2)$, for each m-th client, where $\theta^m$ and $\alpha^m$ are locally updated NN's weights and dropout parameters, respectively. If we consider the mean $\mu_m = \theta^m$ and variance $\sigma_m^2 = \alpha^m (\theta^m)^2$, we can get Eq 5 as an approximation of the mean of aggregated posterior. 

In Eq 5, the definition of  $g^m$ is a weighting factor of each client proportional to the local data size (mentioned in L56 and L108). This formulation is directly adapted from Eq 3 of the paper FedPA [11] and [6,7]. In fact, minimizing the local posteriors using the KL divergence in L112 is equivalent to maximizing the logarithm of the product of weighted posteriors as described in Eq 8 of [7]. In the heterogeneous data, the product rule can achieve a smaller aggregation than the mixture of Gaussian posteriors. We can also find a similar derivation in the continual learning paper [8].

## The update rule for hypernetwork
The expression in L171 is an approximation of computing gradient w.r.t $\psi$.: $\nabla_\psi \mathcal{L}^m_{\text{ELBO}}(\alpha^m) \approx (\nabla_\psi \alpha^m)^{T} \Delta \alpha^m $ where $\nabla_\psi \alpha^m$ is the gradient of a hyper network's output w.r.t $\psi$ and $\Delta \alpha^m = \alpha_*^m - \alpha^m$ is the changes in the dropout variable, where  $\alpha_*^m$ is the locally updated dropout rate using the local data and  $\alpha^m$ is the initially predicted dropout variable from hyper-network. $\Delta \alpha^m$ is an approximation of the vector-jacobian product, which we are inspired by the work of pFedHN[5] and LOOKAHEAD[14]. Given the difference $\Delta \alpha^m$, the update rules for $\psi$ and $e^m$ are:

For $\psi$:

$\Delta \psi = (\nabla_\psi \alpha^m)^{T} \Delta \alpha^m$

$\psi \leftarrow \psi + \eta \frac{1}{M} \sum_m g_m (\nabla_\psi \alpha^m)^{T} \Delta \alpha^m$, using the weights $g_m$ (defined in L56)

For $e^m$:

$\Delta e^m = (\nabla_{e^m} \psi)^{T} (\nabla_\psi \alpha^m)^{T} \Delta \alpha^m$, using the chain rule given $\alpha^m = h_\psi(e^m)$

$e^m \leftarrow e^m + \eta (\nabla_{e^m} \psi)^{T} (\nabla_\psi \alpha^m)^{T} \Delta \alpha^m$

## Reference
[1] pFL-bench: A comprehensive benchmark for personalized federated learning. NeurIPS, 2022.

[2] Federated meta-learning with fast convergence and efficient communication. arXiv preprint, 2018.

[3] Improving federated learning personalization via model agnostic meta learning. arXiv preprint, 2019.

[4] Personalized federated learning with theoretical guarantees: A model-agnostic meta-learning approach. NeurIPS, 2020.

[5] Personalized federated learning using hypernetworks. ICML, 2021.

[6] Asymptotically exact, embarrassingly parallel MCMC. arXiv preprint, 2013.

[7] A Bayesian Federated Learning Framework with Online Laplace Approximation, arXiv, 2021.

[8] Overcoming catastrophic forgetting by incremental moment matching. NeurIPS, 2017.

[9] Snip: Single-shot network pruning based on connection sensitivity. ICLR, 2018.

[10] Variational dropout and the local reparameterization trick. NeurIPS, 2015.

[11] Federated learning via posterior averaging: A new perspective and practical algorithms. ICLR, 2021.

[12] Variational bayesian dropout with a hierarchical prior. CVPR, 2019.

[13] Learning both weights and connections for efficient neural network. NeurIPS, 2015.

[14] Lookahead Optimizer: k steps forward, 1 step back. 2019.

[15] Personalized federated learning via variational bayesian inference. ICML, 2022.

[16] Fedpop: A bayesian approach for personalised federated learning. NeurIPS, 2022.

[17] Personalized federated learning with gaussian processes. NeurIPS, 2021.

[18] Federated Bayesian Neural Regression: A Scalable Global Federated Gaussian Process. arXiv preprint, 2022.

[19] Personalized Federated Learning via Amortized Bayesian Meta-Learning. arXiv preprint, 2023.

[20] What do we mean by generalization in federated learning?. ICLR, 2022.

[21] Variational Bayesian dropout: pitfalls and fixes. ICML, 2018.

[22] Dropout as a bayesian approximation: Representing model uncertainty in deep learning. ICML, 2016.

[23] Products and convolutions of Gaussian probability density functions. Tina-Vision Memo, 2003.

Meta-review: The submission was generally well-received by the reviewers. The combination of hypernetworks and variational dropout to target Bayesian federated learning is interesting and non-trivial. The results are also promising.

Key requested changes are: (i) discussion of and comparison to existing Bayesian federated learning techniques, (ii) more careful ablation studies + error bars. The authors have promised to have some of these changes in the next iteration. The AC thus recommended acceptance.
Decision: Accept (poster)