Title: Federated Learning via Meta-Variational Dropout
Venue: NeurIPS 2023
Accepted: yes
Keywords: Personalized Federated Learning

Variational Dropout

Meta-Learning

Bayesian Neural Network
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-Variational Dropout

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 metalearning 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.

1 Introduction

Federated learning (FL) aims at training a global model from distributed clients without sharing or collecting their sensitive raw data. Thanks to its privacy-preserving aspect of FL [1], it is increasingly popular to be applied to various applications such as image classification [2, 3], object detection [4, 5], keyboard suggestion [6, 7], recommendation [8, 9], and healthcare [10, 11]. The conventional FL algorithm could achieve convergence when the data from different clients is independently and identically distributed (IID) [12-14]. However, due to the differences in preferences, locations, and usage habits of clients, the private data in FL are usually non-IID. When the data distributions of the clients vary, the local model learned from each client can diverge, and thus learning an optimal global model could fail [15, 16]. Furthermore, the scale of client data may not be sufficient to train a local model with large parameters, causing model overfitting and poor generalization [17, 18].

To overcome the challenge caused by non-IID data, personalized federated learning (PFL) has emerged [15, 16]. In the PFL, each client is allowed to have its own personalized model trained on each client's local data while still participating in the global model training. There are many branches of PFL, such as those based on multitask learning [19, 20], meta-learning [21-26], and transfer learning [27-30]. Although these approaches improve training convergence in non-IID 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 in the model parameters [31-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-IID and limited client data simultaneously.

In this paper, we present Meta-Variational Dropout (MetaVD), a novel Bayesian meta-learning approach [38-42] developed for PFL. MetaVD learns to predict client-dependent dropout rates

via a hypernetwork [25, 43-45]. This mechanism facilitates a data-efficient estimation of clientspecific posterior simply by modulating a shared global NN parameter across all clients. The adaptation of MetaVD to conventional FL algorithms (e.g., FedAvg [1], Reptile [22], MAML [21], or PerFedAvg [23]) allows flexible model personalization across a variety of non-IID client distributions. MetaVD also incorporates a 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 on non-IID data. In addition, MetaVD inherits the merit of Variational Dropout (VD) [46, 47], which facilitates model compression. This not only improves model generalization for out-of-data (OOD) clients but also reduces communication costs for parameter exchange. We performed an extensive analysis on MetaVD 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 to other state-of-the-art 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 D m = { ( x i m , y i m ) } |D m | i =1 , the FL problem can be formulated as follows:

w denotes the model parameter, and the global learning objective J ( w ) at the server is a weighted average of the local objectives J m ( w ) over M clients. The weight g m is proportional to the size of the local dataset (e.g., |D m | / |D| ). The local loss in a client device is usually defined as the empirical negative log-likelihood on the m -th client's dataset D m (i.e., ℓ ( D m ; w ) = -log p ( y m | x m , 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 w m . Then, the aggregation step computes the global parameter in the server by taking the weighted average of the local parameters (i.e., ¯ w ← ∑ M m =1 g m w m ). 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 m ) 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 m would drift away from each other, causing the ¯ w to be suboptimal [15, 16]. (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 a 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 a point estimate of the parameter as in Eq.1, recent work [31-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, a Gaussian distribution is used to represent each parameter in FedAG [32], and a posterior aggregation strategy using the MCMC technique is proposed in FedPA [33]. FedBE [31] 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 a 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 it 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-IID data [34-37, 51, 52]. For example, pFedGP [34] learns a 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 a shared global prior. These approaches achieve a substantial predictive performance improvement in non-IID data senerios compared to the previous Bayesian FL approaches. However, their approach still has limitations, such as the high computational cost of inverting a large kernel matrix in pFedGP [34] and imposing strong constraints in pFedBayes [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-Variational Dropout

We propose a new Bayesian PFL approach, called Meta-Variational Dropout (MetaVD), which can simultaneously handle model personalization, regularization, and compression. MetaVD also exploits the uncertainties in model aggregation, thereby improving the training convergence on 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 [31-33], we can utilize the (amortized) Variational Inference (VI) framework of Bayesian meta-learning [38-42] that is originally developed for few-shot multi-task learning [53-55]. Considering each m -th client in FL as an individual task, an evidence lower-bound (ELBO) L ELBO over all the M datasets is defined as

Here, p ( y m | x m , w m ) represents a likelihood model constructed using a neural network (NN) on each m -th client's data [56-59], and w m is a client-specific NN parameter. q ( w m ; ϕ ) is a variational posterior (or probabilistic) distribution over the w m , which is also characterized by ϕ . g m is the local weight as defined in Eq.(1). p ( w m ) is a prior distribution that acts as a regularizer for the q ( w m ; ϕ ) . The local ELBO defined on each m -th client in Eq.(2) makes tradeoffs between the expected log-likelihood on its local dataset D m and the KL divergence with a prior. In fact, maximizing the L ELBO ( ϕ ) with respect to the variational parameter ϕ is equivalent to minimizing ∑ M t =1 g m KL ( q ( w m ; ϕ ) || p ( w m |D m )) . In theory, q ( w m ; ϕ ) is trained to approximate the true local posterior distribution over the w m (i.e., q ( w m ; ϕ ) ≈ p ( w m |D m )) [56-59].

Astraightforward Bayesian PFL approach might be utilizing a separate local posterior model q ( w m ; ϕ ) for each client device (e.g., , ϕ = ( ϕ 1 , · · · , ϕ M ) ). However, only a 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 ϕ m independently for each device is difficult. As an alternative, we introduce a new hypernetwork-based [25, 43-45] conditional dropout posterior modeling approach that can be data-efficiently trained across multiple clients in FL.

3.2 Meta-Variational Dropout

Posterior model. To promote efficient model personalization and reduce overfitting in Bayesian FL, we define the posterior model in Eq.2 based on a Variational Dropout (VD) technique that multiplies continuous Gaussian noise to the NN parameters during training to prevent overfitting [46, 47, 60-62]. MetaVD extends the VD posterior by employing a global hypernetwork that learns to predict client-specific dropout rates (or personal model structure). In MetaVD, the variational posterior distribution for each m -th client's parameter, q ( w m ; ϕ ) in Eq.2, can be constructed as

ҧ

×

×

×

𝐻𝑦𝑝𝑒𝑟𝑛𝑒𝑡𝜓

Server

Client embedding

Client

Server

𝜃1

𝛼1

𝑚

𝑒 𝑚

𝐺𝑙𝑜𝑏𝑎𝑙 𝑁𝑁𝜃

𝜃1

𝜃2

𝜃3

𝜃2

𝛼2

𝑚

𝜃3

𝛼3

𝑚

×

×

×

×

×

×

×

×

×

×

×

×

×

×

×

×

×

×

Local Data

Few-shot Adaptation

(MAML/Reptile)

Train

Val

∆𝛼 𝑀

𝑒 1

𝑒 𝑀

Model Update (

𝜃

,

𝜓

,

𝑒

)

⋮

Posterior Aggregation

𝐺𝑙𝑜𝑏𝑎𝑙 𝑁𝑁𝜃

Client specific

parameters

𝑤 𝑚

∆𝛼 1

𝜃 𝑎𝑔𝑔

MetaVD

FedAvg

𝛼 𝑚

𝜃1

𝜃2

𝜃3

Server

to

m

-th

Client

Clients

to Server

∆𝜃

𝜃 𝑎𝑔𝑔

𝜃

×𝑝𝑟𝑢𝑛𝑒𝑑 𝑝𝑎𝑟𝑎𝑚𝑒𝑡𝑒𝑟

𝜃 𝐴

𝜃 𝐵

Mode

Mean

⋮

(a)

(b)

(c)

×

×

×

×

×

×

Client

Client

𝜃1,1

𝜃1,2

𝜃2,1

𝜃2,2

𝜃2,3

𝜃3,1

𝜃3,2

𝜃3,1

𝜃3,2

𝜃2,1

𝜃2,2

𝜃2,3

𝜃1,1

𝜃1,2

𝛼1,1

𝑚

𝛼1,2

𝑚

𝛼1,3

𝑚

𝛼2,1

𝑚

𝑚

𝛼2,3

𝑚

𝛼3,1

𝑚

𝛼3,2

𝑚

𝛼3,3

𝑚

The variational parameter ϕ is characterized by three distinct parameters ( θ, ψ, e ) . The θ is a 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 weight w m is described as a product of independent Gaussian distributions. In the product term in Eq.(3), each k -th factor describes a conditional Gaussian noise multiplication to the global NN parameter (e.g., w m = θ ∗ ϵ m where ϵ m ∼ N ( ⃗ 1 , α m ) and ⃗ 1 is an K -dimensional all-one vector). The α m k represents the clientspecific dropout variable 1 on each k -th index of NN parameter θ k . The h ψ indicates a hypernetwork parameterized by ψ [43-45, 45] predicts the client-specific dropout rate α m . The e m is a learnable client embedding used as an input to the h ψ . In essence, MetaVD is a technique that modulates a global NN parameter with Gaussian noises. By predicting the client-specific dropout variable via a hypernetwork, a 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 ( w m ; ϕ ) , we need to specify the KL divergence term and the prior model p ( w m ) in Eq.(2). 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 VD [46, 47, 61, 62], and (ii) all clients must share the same prior model to support the multiplicative posterior aggregation rule in Bayesian FL. We adopt the hierarchical prior [64, 65, 47] discussed in [47] because of its straightforward analytic KL term derivation and proven effectiveness in network regularization and sparsification. Under the hierarchical prior assumption, the KL divergence term in Eq.(2) is simplified to KL ( q ( w m ; ϕ ) || p ( w m )) = ∑ K k =1 0 . 5 log(1 + ( α m k ) -1 ) ; see Appendix B for a more detailed derivation. This KL 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 α m for each m -th client. The global parameter θ and α m are transmitted to each client device. Then, the posterior model in a m -th client is trained on local data using the following ELBO objective:

which maximizes L m ELBO on the client dataset D m = { ( x i m , y i m ) } |D m | i =1 with respect to the variational parameters (i.e., θ , α m ). The optimization can be done by the stochastic gradient variational Bayes (SGVB). [56-58], which reparameterizes the random weight variable w m using a differentiable transformation as w m k = f ( ϵ k ; θ k , α m k ) = θ k + √ α m k θ k ϵ k . with a random IID noise ϵ k ∼ N (0 , 1) . The intermediate weight w m k is now differentiable with respect to θ k , hence α m k , and can be optimized via the SGD [49]. The analytical KL term acts as a regularization for α m k .

Server-side objective. Once a local posterior adaptation is done in each client device, the updated parameters θ m ∗ and α m ∗ are returned to the server, then the server updates the variational parameters

Algorithm 1 MetaVD algorithm with MAML and Reptile variant for FL

Input: # of communication round R , # of client N , server learning rate η , client learning rate γ , inner learning rate l , local steps E , inner steps I , KL divergence parameter β .

procedure

LOCALADAPTATION_REPTILE(

for

θ

,

α

)

Each local step

i

from

1

to

E

do

Set

θ

α

∗

=

α

i

Send

∗

θ

=

∗

θ

i

and and

α

∗

to the server

(e.g., θ, ψ, e m ) using the updated and local parameters. To update the global NN parameter θ , we first assume the Bayesian posterior aggregation rule: p ( w |D ) ∝ ∏ M m =1 p ( w m |D m ) [31-33, 66]. Since each local posterior in Eq.(3) is a Gaussian (dropout) distribution, their product is also Gaussian: N ( w | θ agg ∗ , · ) ≈ ∏ M m =1 N ( w m | θ m ∗ , α m ∗ ( θ m ∗ ) 2 ) . This gives us an exact aggregation rule 2 to compute the maximum a posterior (MAP) solution of the θ agg ∗ as follows:

Note that Eq.(5) has the intuitive interpretation that the aggregation weight r m is inversely proportional to its corresponding dropout variable (or noise variance) α m ∗ . 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 θ k , we follow the similar parameter update rule of FedAvg, except that the Bayesian aggregation Eq.(5) is utilized as in Algo.1. For the parameter of hypernetwork ψ , a more general update rule is used following [25]. We compute the changes in the updated dropout for each client ∆ α m as described in Algo.1, then an approximated gradient of the hypernetwork parameter is computed using the chain rule as ∇ ψ L m ELBO ( α m ) = ( ∇ ψ α m ) T ∆ α m , where ∇ ψ α m is the gradient of the output of hypernetwork with respect to ψ . ∆ α m is an approximation of the vector-jacobian product, which we are inspired by the work of [25] and [69]. The gradient for ∇ e m L m ELBO ( α m ) can be derived using the similar chain rule. The detailed update rules for each parameter are in Algo.1.

Combination with other meta-learning algorithms. Since the KL regularization in Eq.(4) is independent of the global (or initial) parameter θ , MetaVD can be combined with several existing meta-learning based PFL algorithms (e.g., Reptile [22, 70], MAML [21, 55], PerFedAvg [23, 71]). For example, MAML [21, 55] requires some internal update steps using a subsampled dataset to compute the second-order gradient for θ . Reptile [22, 70] uses only a first-order gradient computation as described in Algo.1, which is similar to FedAvg except for the addition of the learning rate η . Each local adaptation step is performed using the L m ELBO in Eq.(4) and the local data set D m . 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 MetaVD 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 used multiple FL datasets [72], including CIFAR-10, CIFAR-100, FEMINIST, and CelebA. To evaluate the effectiveness of the hypernetwork, we also performed an ablation study comparing MetaVD to regular VD. Additionally, we assessed the uncertainty calibration and model compression ability of MetaVD. Finally, we test MetaVD with multi-domain datasets.

Baselines. We compare our method with standard FL methods such as FedAvg [1] and FedProx [73], meta-learning PFL algorithms like Reptile [22], MAML [21], and PerFedAvg [23], and Bayesian FL methods including FedBE [31] and pFedGP [34]. To ensure consistency, we employ the widelyused 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.

Implementation. For reproducibility, we run experiments in a containerized environment that simulates FL communication with clients only on a server. We test T = 1000 of total FL rounds, following the conventions in [18]. One baseline model can be run on a single GPU. All experiments are run on a cluster of 32 NVIDIA GTX 1080 GPUs. MetaVD's hypernetwork consists of an embedding layer of dimension (1 + M/ 4) , followed by three fully connected NNs with a Reaky ReLU activation and an 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 MetaVD 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 https://github.com/insujeon/MetaVD.

4.1 Generalization on Non-IID Settings

Datasets and training. To evaluate the generalization 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 pFLBench [18], using Dirichlet allocation to partition each dataset into 130 clients with different Dirichlet parameters, denoted as ˙ α = [5 , 0 . 5 , 0 . 1] . As shown in Figure 2, the class labels and data size per client are heterogeneous across clients. A smaller ˙ α represents a higher degree of heterogeneity. To evaluate

0

10

20

30

40

50

60

70

80

90

100

110

120

Client

0

100

200

300

400

500

600

700

800

# of data samples

=0.5

0

10

20

30

40

50

60

70

80

90

100

110

120

Client

0

100

200

300

400

500

600

700

800

# of data samples

=0.1

Dirichlet distribution (CIFAR-100)

the test accuracy and generalization performance of the baselines on 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.

Results. Table 1 shows the weighted average classification accuracy for participating (Test) and non-participating (OOD) clients on CIFAR-10 and CIFAR-100 datasets with varying non-IID degrees. The generalization gap, denoted by ∆ , represents the difference between OOD and Test accuracy. As shown in Table 1, PFL methods such as Reptile, MAML, 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-IID increases. When combined with MetaVD, 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 MetaVD 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 in 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 as a global model parameter, as in FedAvg. EnsembleVD maintains M independent dropout parameters for all clients. The client-specific dropout parameter can be stored in each client analogous to the partial FL [80, 81, 5]. In contrast, MetaVD utilizes a 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 [82, 83].

Results. Table 2 outlines the results of the ablation study; MetaVD's hypernetwork-based posterior modeling outperforms all other baselines. The dropout rates in baselines such as VD or EnsembleVD could not fully learn the independent dropout variables well due to restricted client participation. On the other hand, both the hypernetwork and the global parameter converge well in MetaVD. This observation demonstrates that

MetaVD's hypernetwork provides a more data-efficient approach to learning client-specific model uncertainty compared to other baselines. Further ablation studies performed on the FEMNIST dataset are shown in Appendix H.

4.3 Uncertainty Calibration

(a)Reptile

(b) MAML

1.0

1.0

ECE = 0.77%

ECE = 0.70%

0.8

MCE = 46.19%

0.8

MCE = 48.16%

0.6

0.6

0.4

0.4

0.2

0.2

0.8.0

0.8.0

0.20.40.60.8

1.0

0.2

0.4

0.6

0.8

1.0

(c)Reptile+MetaVD

(d) MAML + MetaVD

1.0

ECE = 0.32%

1.0

ECE = 0.38%

0.8

MCE = 23.86%

0.8

MCE = 32.56%

0.6

0.6

0.4

0.4

0.2

0.2

08:0

0.8.0

0.20.40.6

0.8

1.0

0.20.40.6

0.8

1.0

Confidence

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 [84-86]. In the FL

environment, where customers have limited non-IID 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 CIFAR100 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-based 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 metalearning-based 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-86] in CIFAR-100 ( ˙ α = 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 a perfect diagonal represents miscalibration. While Reptile and MAML tend to be overconfident in their predictions, Reptile+MetaVD comes fairly close to the desired diagonal function. Remarkably, in most instances, adapting MetaVD leads to improved model calibration.

4.4 Client Participation

Table 4: Results of classification accuracies with different participant client rates of s = 0 . 2 , s = 0 . 1 , and s = 0 . 05 in the FEMNIST dataset. We report the results of participating clients (Test) and non-participating 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.1, and 0.05, respectively. To measure OOD accuracy, we excluded 40 preselected clients from the selection of 200 clients so that they do not participate in the entire training.

Results. The overall classification results with different participant client rates s 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 in training. The decrease in test accuracy as the participating client rate s becomes smaller is due to having less training data as fewer clients participate in each round. Meanwhile, Reptile+MetaVD outperforms the other baselines. Interestingly, the performance drop for Reptile+MetaVD is not as significant as for Reptile, showing that MetaVD can adapt well to FL scenarios with smaller participant sizes.

4.5 Muti-domain Datasets

Settings. Existing federated learning algorithms usually assume a single-domain approach, where only one dataset is 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

further evaluate the performance of PFL algorithms on a multi-domain FL dataset, where we assume that each client can 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 outperforms 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 MetaVD in improving robustness and generalization even in multi-domain FL datasets.

4.6 Model Compression

Settings. Federated learning optimization requires frequent communication of model parameters between devices and the central server, which can be slow and may raise privacy concerns. Therefore, minimizing communication costs by reducing the size of model parameters is an important issue in FL. To explore the compression capabilities of MetaVD, we performed an additional experiment on the CIFAR-10 dataset. The sign DP indicates that each model parameter is dropped during the FL communication.

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 a dropout rate greater than 0 . 8 was dropped during the FL rounds.

Results. Table 6 shows the test and OOD accuracy results, as well as the sparsity (%) in the CIFAR100 dataset. The sparsity represents the proportion of zero-value model parameters in the personalized layer. A higher sparsity percentage indicates more parameter pruning or elimination performed on the weight. In our experiment, MetaVD 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+MetaVD+DP, the performance decreased by about 2% while using only 20% of the weights. On the other hand, MAML+MetaVD+DP and PerFedAvg+MetaVD+DP show an improvement in performance of about 1% . This shows that MetaVD can compress the model parameters required in the personalized layer, reducing the communication cost in FL without sacrificing much performance. Appendix K shows more experiments on model compression results on the CIFAR-100 dataset with different non-IID settings.

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 adaptation in federated learning (FL) with non-IID 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. MetaVD is also a generic approach that is compatible with any other existing meta-learning-based PFL algorithms to avoid model overfitting. In addition, MetaVD's 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 on several PFL benchmarks, including CIFAR-10, CIFAR-100, FEMINIST, and CelebA, as well as multi-domain datasets. It demonstrates superior classification accuracy and uncertainty calibration, especially for out-of-distribution (OOD) clients. Overall, the experimental results show MetaVD to be a highly versatile approach capable of addressing many challenges in FL.

Acknowledgment

This work was supported by Center for Applied Research in Artificial Intelligence(CARAI) grant funded by Defense Acquisition Program Administration(DAPA) and Agency for Defense Development(ADD) (UD190031RD), Institute of Information & communications Technology Planning &Evaluation (IITP) grant (No.2019-0-01082, SW StarLab and No.2021-0-01343, Artificial Intelligence Graduate School Program (Seoul National University)), Basic Science Research Program through National Research Foundation of Korea (NRF) funded by the Korea government (MSIT) (NRF-2020R1A2B5B03095585). Gunhee Kim is the corresponding author.

References

Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication-efficient learning of deep networks from decentralized data. In Artificial intelligence and statistics , pages 1273-1282. PMLR, 2017.

Yue Zhao, Meng Li, Liangzhen Lai, Naveen Suda, Damon Civin, and Vikas Chandra. Federated learning with Non-IID data. arXiv preprint arXiv:1806.00582 , 2018.

Jaehoon Oh, Sangmook Kim, and Se-Young Yun. FedBABU: Towards enhanced representation for federated image classification. arXiv preprint arXiv:2106.06042 , 2021.

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. In Proceedings of the AAAI conference on artificial intelligence , volume 34, pages 13172-13179, 2020.

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-23320, 2021.

Andrew Hard, Kanishka Rao, Rajiv Mathews, Swaroop Ramaswamy, Françoise Beaufays, Sean Augenstein, Hubert Eichner, Chloé Kiddon, and Daniel Ramage. Federated learning for mobile keyboard prediction. arXiv preprint arXiv:1811.03604 , 2018.

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:1812.02903 , 2018.

Khalil Muhammad, Qinqin Wang, Diarmuid O'Reilly-Morgan, Elias Tragos, Barry Smyth, Neil Hurley, 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.

Chuhan Wu, Fangzhao Wu, Yang Cao, Yongfeng Huang, and Xing Xie. FedGNN: Federated graph neural network for Privacy-Preserving recommendation. arXiv preprint arXiv:2102.04925 , 2021.

Qian Yang, Jianyi Zhang, Weituo Hao, Gregory P Spell, and Lawrence Carin. FLOP: 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.

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):1735-1743, 2021.

Xiang Li, Kaixuan Huang, Wenhao Yang, Shusen Wang, and Zhihua Zhang. On the convergence of FedAvg on Non-IID data. arXiv preprint arXiv:1907.02189 , 2019.

Peter Kairouz, H 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-210, 2021.

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.

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 Sustainability (WorldS4) , pages 794-797. IEEE, 2020.

Alysa Ziying Tan, Han Yu, Lizhen Cui, and Qiang Yang. Towards personalized federated learning. IEEE Transactions on Neural Networks and Learning Systems , 2022.

Honglin Yuan, Warren Morningstar, Lin Ning, and Karan Singhal. What do we mean by generalization in federated learning? arXiv preprint arXiv:2110.14216 , 2021.

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, 2022.

Liam Collins, Hamed Hassani, Aryan Mokhtari, and Sanjay Shakkottai. Exploiting shared representations for personalized federated learning. In International conference on machine learning , pages 2089-2099. PMLR, 2021.

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.

Fei Chen, Mi Luo, Zhenhua Dong, Zhenguo Li, and Xiuqiang He. Federated Meta-Learning with fast convergence and efficient communication. arXiv preprint arXiv:1802.07876 , 2018.

Yihan Jiang, Jakub Koneˇ cn` y, Keith Rush, and Sreeram Kannan. Improving federated learning personalization via model agnostic meta learning. arXiv preprint arXiv:1909.12488 , 2019.

Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. Personalized federated learning with theoretical guarantees: A model-agnostic meta-learning approach. Advances in Neural Information Processing Systems , 33:3557-3568, 2020.

Canh T Dinh, Nguyen Tran, and Josh Nguyen. Personalized federated learning with moreau envelopes. Advances in Neural Information Processing Systems , 33:21394-21405, 2020.

Aviv Shamsian, Aviv Navon, Ethan Fetaya, and Gal Chechik. Personalized federated learning using hypernetworks. In International Conference on Machine Learning , pages 9489-9502. PMLR, 2021.

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 11237-11244, 2023.

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.

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é Iii and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning , volume 119 of Proceedings of Machine Learning Research , pages 5132-5143. PMLR, 2020.

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. PMLR, 18-24 Jul 2021.

Tiansheng Huang, Li Shen, Yan Sun, Weiwei Lin, and Dacheng Tao. Fusion of global and local knowledge for personalized federated learning. arXiv preprint arXiv:2302.11051 , 2023.

Hong-You Chen and Wei-Lun Chao. FedBE: Making bayesian model ensemble applicable to federated learning. arXiv preprint arXiv:2009.01974 , 2020.

Adam Thor Thorgeirsson and Frank Gauterin. Probabilistic predictions with federated learning. Entropy , 23(1):41, 2020.

Maruan Al-Shedivat, Jennifer Gillenwater, Eric Xing, and Afshin Rostamizadeh. Federated learning via posterior averaging: A new perspective and practical algorithms. arXiv preprint arXiv:2010.05273 , 2020.

Idan Achituve, Aviv Shamsian, Aviv Navon, Gal Chechik, and Ethan Fetaya. Personalized federated learning with gaussian processes. Advances in Neural Information Processing Systems , 34:8392-8406, 2021.

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 26293-26310. PMLR, 2022.

Haolin Yu, Kaiyang Guo, Mahdi Karami, Xi Chen, Guojun Zhang, and Pascal Poupart. Federated bayesian neural regression: A scalable global federated gaussian process. arXiv preprint arXiv:2206.06357 , 2022.

Shrey Bhatt, Aishwarya Gupta, and Piyush Rai. Bayesian federated learning via predictive distribution distillation. arXiv preprint arXiv:2206.07562 , 2022.

Sachin Ravi and Alex Beatson. Amortized bayesian Meta-Learning. In International Conference on Learning Representations , 2018.

Marta Garnelo, Jonathan Schwarz, Dan Rosenbaum, Fabio Viola, Danilo J Rezende, SM Eslami, and Yee Whye Teh. Neural processes. arXiv preprint arXiv:1807.01622 , 2018.

Jonathan Gordon, John Bronskill, Matthias Bauer, Sebastian Nowozin, and Richard E Turner. MetaLearning probabilistic inference for prediction. arXiv preprint arXiv:1805.09921 , 2018.

Insu Jeon, Youngjin Park, and Gunhee Kim. Neural variational dropout processes. In International Conference on Learning Representations , 2022.

Ekaterina Iakovleva, Jakob Verbeek, and Karteek Alahari. Meta-Learning with shared amortized variational inference. In Hal Daumé Iii 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.

David Ha, Andrew Dai, and Quoc V. Le. HyperNetworks. arXiv preprint arXiv:1609.09106 , 2016.

Johannes Von Oswald, Christian Henning, Benjamin F Grewe, and João Sacramento. Continual learning with hypernetworks. arXiv preprint arXiv:1906.00695 , 2019.

Dominic Zhao, Seijin Kobayashi, João Sacramento, and Johannes von Oswald. Meta-Learning via hypernetworks. In 4th Workshop on Meta-Learning at NeurIPS 2020 (MetaLearn 2020) . NeurIPS, 2020.

Durk P Kingma, Tim Salimans, and Max Welling. Variational dropout and the local reparameterization trick. Advances in Neural Information Processing Systems , 28, 2015.

Yuhang Liu, Wenyong Dong, Lei Zhang, Dong Gong, and Qinfeng Shi. Variational bayesian dropout with a hierarchical prior. In 2019 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) , pages 7124-7133. IEEE, June 2019.

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.

Yingjie Tian, Yuqi Zhang, and Haibin Zhang. Recent advances in stochastic gradient descent in deep learning. Mathematics , 11(3):682, 2023.

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 Statistics , pages 2350-2358. PMLR, 2021.

Nikita Kotelevskii, Maxime Vono, Alain Durmus, and Eric Moulines. FedPop: A bayesian approach for personalised federated learning. Advances in Neural Information Processing Systems , 35:8687-8701, 2022.

Shiyu Liu, Shaogao Lv, Dun Zeng, Zenglin Xu, Hui Wang, and Yue Yu. Personalized federated learning via amortized bayesian Meta-Learning. arXiv preprint arXiv:2307.02222 , 2023.

Brenden M Lake, Ruslan Salakhutdinov, and Joshua B Tenenbaum. Human-level concept learning through probabilistic program induction. Science , 350(6266):1332-1338, 2015.

Jake Snell, Kevin Swersky, and Richard Zemel. Prototypical networks for few-shot learning. Advances in Neural Information Processing Systems , 30, 2017.

Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-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 1126-1135. PMLR, 2017.

Radford M Neal. Bayesian Learning for Neural Networks , volume 118. Springer Science & Business Media, 2012.

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.

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.

Kirill Neklyudov, Dmitry Molchanov, Arsenii Ashukha, and Dmitry Vetrov. Variance networks: When expectation does not meet your expectations. arXiv preprint arXiv:1803.03764 , 2018.

Yarin Gal and Zoubin Ghahramani. Dropout as a bayesian approximation: Representing model uncertainty in deep learning. In 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.

Dmitry Molchanov, Arsenii 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.

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.

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 Learning Research , pages 118-126, Atlanta, Georgia, USA, 2013. PMLR.

Michael E Tipping. Sparse bayesian learning and the relevance vector machine. Journal of machine learning research , 1(Jun):211-244, 2001.

S Derin Babacan, Rafael Molina, Minh N Do, and Aggelos K Katsaggelos. Bayesian blind deconvolution with general sparse image priors. In Computer Vision-ECCV 2012: 12th European Conference on Computer Vision, Florence, Italy, October 7-13, 2012, Proceedings, Part VI 12 , pages 341-355. Springer, 2012.

Willie Neiswanger, Chong Wang, and Eric Xing. Asymptotically exact, embarrassingly parallel MCMC. arXiv preprint arXiv:1311.4780 , 2013.

Paul Bromiley. Products and convolutions of gaussian probability density functions. Tina-Vision Memo , 3(4):1, 2003.

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 on Pattern Analysis and Machine Intelligence , 2023.

Michael Zhang, James Lucas, Jimmy Ba, and Geoffrey E Hinton. Lookahead optimizer: k steps forward, 1 step back. Advances in neural information processing systems , 32, 2019.

Alex Nichol, Joshua Achiam, and John Schulman. On First-Order Meta-Learning algorithms. arXiv preprint arXiv:1803.02999 , 2018.

Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. On the convergence theory of Gradient-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.

Sebastian Caldas, Sai Meher Karthik Duddu, Peter Wu, Tian Li, Jakub Koneˇ cn` y, H Brendan McMahan, Virginia Smith, and Ameet Talwalkar. LEAF: A benchmark for federated settings. arXiv preprint arXiv:1812.01097 , 2018.

Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. Proceedings of Machine learning and systems , 2:429-450, 2020.

Paul Pu Liang, Terrance Liu, Liu Ziyin, Nicholas B Allen, Randy P Auerbach, David Brent, Ruslan Salakhutdinov, and Louis-Philippe Morency. Think locally, act globally: Federated learning with local and global representations. arXiv preprint arXiv:2001.01523 , 2020.

Agustinus Kristiadi, Matthias Hein, and Philipp Hennig. Being bayesian, even just a bit, fixes overconfidence in ReLU networks. In Hal Daumé Iii 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.

Noah Weber, Janez Starc, Arpit Mittal, Roi Blanco, and Lluís Màrquez. Optimizing over a bayesian last layer, 2018.

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-1206. PMLR, 2021.

Alexander A Alemi, Ian Fischer, Joshua V Dillon, and Kevin Murphy. Deep variational information bottleneck. arXiv preprint arXiv:1612.00410 , 2016.

Wei Du, Xiao Zeng, Ming Yan, and Mi Zhang. Efficient federated learning via variational dropout. Accessed: 2023-10-21, September 2018.

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.

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:11220-11232, 2021.

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.

Gary Cheng, Zachary Charles, Zachary Garrett, and Keith Rush. Does federated dropout actually work? In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) Workshops , pages 3387-3395, 2022.

Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q Weinberger. On calibration of modern 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 1321-1330. PMLR, 2017.

Alexandru Niculescu-Mizil and Rich Caruana. Predicting good probabilities with supervised learning. In Proceedings of the 22nd international conference on Machine learning , pages 625-632, 2005.

Max-Heinrich Laves, Sontje Ihler, Karl-Philipp Kortmann, and Tobias Ortmaier. Calibration of model uncertainty for dropout variational inference. arXiv preprint arXiv:2006.11584 , 2020.

Shoushan Li and Chengqing Zong. Multi-domain sentiment classification. In Proceedings of ACL-08: HLT, short papers , pages 257-260, 2008.

Yongxin Yang and Timothy M Hospedales. A unified perspective on Multi-Domain and Multi-Task learning. arXiv preprint arXiv:1412.7489 , 2014.

Hakan Bilen and Andrea Vedaldi. Universal representations:the missing link between faces, text, planktons, and cat breeds. arXiv preprint arXiv:1701.07275 , 2017.

Hongyi Wang, Mikhail Yurochkin, Yuekai Sun, Dimitris Papailiopoulos, and Yasaman Khazaeni. Federated learning with matched averaging. arXiv preprint arXiv:2002.06440 , 2020.

In S Jeon, Deokyoung Kang, and Suk I Yoo. Blind image deconvolution using student'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.

Christopher M Bishop and Nasser M Nasrabadi. Pattern Recognition and Machine Learning , volume 4. Springer, 2006.

David JC MacKay. A practical bayesian framework for backpropagation networks. Neural computation , 4(3):448-472, 1992.

David Barber and Christopher M Bishop. Ensemble learning in bayesian neural networks. Nato ASI Series F Computer and Systems Sciences , 168:215-238, 1998.

Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: A simple way to prevent neural networks from overfitting, 2014.

Geoffrey E Hinton, Nitish Srivastava, Alex Krizhevsky, Ilya Sutskever, and Ruslan R Salakhutdinov. Improving neural networks by preventing co-adaptation of feature detectors. arXiv preprint arXiv:1207.0580 , 2012.

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.

Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images, 2009.

Gregory Cohen, Saeed Afshar, Jonathan Tapson, and Andre Van Schaik. EMNIST: Extending MNIST to handwritten letters. In 2017 international joint conference on neural networks (IJCNN) , pages 2921-2926. IEEE, 2017.

Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In 2015 IEEE International Conference on Computer Vision (ICCV) , pages 3730-3738. IEEE, December 2015.

Namhoon Lee, Thalaiyasingam Ajanthan, and Philip HS Torr. SNIP: Single-shot network pruning based on connection sensitivity. arXiv preprint arXiv:1810.02340 , 2018.

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.

Takuya Akiba, Shotaro Sano, Toshihiko Yanase, Takeru Ohta, and Masanori Koyama. Optuna: A next-generation 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.

Zachary Lipton, Yu-Xiang Wang, and Alexander Smola. Detecting and correcting for label shift with black box predictors. In International conference on machine learning , pages 3122-3130. PMLR, 2018.

Ali Ramezani-Kebrya, Fanghui Liu, Thomas Pethick, Grigorios Chrysos, and Volkan Cevher. Federated learning under covariate shifts with generalization guarantees. arXiv preprint arXiv:2306.05325 , 2023.

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 ≡ D 1 ∪· · ·∪D M ) . This posterior distribution is proportional to the product of the likelihoods and a prior, p ( w |D ) ∝ p ( w ) ∏ M m =1 p ( D m | 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 a product of local posteriors:

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 [90]. Thus, frequent communications between the server and the client are still demanded in 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 MetaVD

In Section 3.2, we adopt the hierarchical prior [64, 65, 91] proposed in [47] 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 ( q ( w m ; ϕ ) || p ( w m )) = ∑ K k =1 0 . 5 log(1 + ( α m k ) -1 ) in the manuscript for readability, here we provide more detailed descriptions of applying the hierarchical prior.

Under the hierarchical prior assumption, we consider the joint prior and joint posterior distributions [92]. The joint prior, p ( w m , γ m ) = p ( w m | γ m ) p ( γ m ) , is defined as a combination of a zero-mean Gaussian distribution, p ( w m | γ m ) = N ( w m | 0 , γ m ) , and a uniform hyper-prior, p ( γ m ) = U ( γ m | a, b ) , over the variance. Then, we define a (conditional) joint variational posterior, q ( w m , γ m | ϕ ) = q ( w m | ϕ ) q ( γ m ) , comprising the (conditional) dropout posterior q ( w m ; ϕ = ( θ, ϕ, e m )) defined in Section 3.2 of the manuscript and an additional Dirac delta distribution, q ( γ m ) , to approximate the true (joint) posterior p ( w m , γ m |D m ) (given the client's dataset D m ).

ELBO. With the hierarchical prior p ( w m , γ m ) = p ( w m | γ m ) p ( γ m ) and the (conditional) joint posterior q ( w m , γ m | ϕ ) = q ( w m | ϕ ) q ( γ m ) of the MetaVD, we can derive the local objective for each client as follows:

Eq. (2) is derived from Bayes' rule: p ( w m , γ m |D m ) = p ( y m | x m ,w m ) p ( w m ,γ m ) p ( y m | x m ) . By reordering Eq.(3), we get

The lower-bound in Eq.(4) is due to the positivity of the KL ( q ( w m , γ m | ϕ ) || p ( w m , γ m |D m )) . Here, Eq. (5) corresponds to the ELBO objective of Eq. (3) in the manuscript.

KL term. If we further decompose the KL divergence term in Eq. (5),

The log( b -a ) is independent of the unknown variables α m , θ , and γ m . Thus, we do not need to specify the value of hyperparameters a and b and can neglect them in practice 3 . Eq. (6) provides the rationale behind the

KL divergence term that we mentioned in Section 3.2 of the manuscript. In fact, the independence between the KL term and the parameter θ ensures compatibility with other optimization meta-learning algorithms. Also, the two-level structure in a hierarchical system can generate a 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 hierarchical 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, 57] were proposed, which impose a prior distribution on each parameter. Various inference approaches have been proposed to model the posterior distribution of BNNs [58, 46, 59]. Variational dropout (VD) [60, 46, 61, 62, 47] encompasses a set of techniques that model the variational posterior distribution based on the dropout regularization method. Dropout regularization [95, 96] randomly deactivates some of the model parameters during training by multiplying them with discrete Bernoulli random noise. This method was initially popularized to prevent overfitting in neural network models. Subsequently, fast dropout [63, 97] 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, 62, 47], 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 [18], which is a recently proposed benchmark for federated learning.

Datasets. Here, we present descriptions of the dataset used in our experiment.

The CIFAR-10 and CIFAR-100 datasets [98] are popular for 10-class and 100-class image classification 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 [18], we use Dirichlet allocation to split this dataset into 130 clients with different Dirichlet parameters as ˙ α = [5, 0.5, 0.1] (a smaller ˙ α indicates a higher heterogeneous degree).

The Federated Extended MNIST ( FEMNIST ) is a 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 clients and a total of 8752 training and 2347 test images with a resolution of 84x84 pixels. Each client is assigned images of a 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, we follow the same heterogeneous patterns exhibited in the pFL-Bench, which covers a wide range of scales, partition manners, and non-IID degrees. This enables comprehensive comparisons and analysis among different methods in the non-IID data environment.

Baselines. We present an overview of the baseline models used in our experiment, covering various popular and state-of-the-art approaches across three categories: Non-PFL, meta-learning-based PFL, and Bayesian FL methods.

The following Non-PFL methods are considered in our experiments:

FedAvg [1] is a standard FL algorithm that averages gradients weighted by the data size of clients in each FL round.

FedProx [73] employs a proximal term to encourage updated local models for clients not to deviate too much from the global model.

The following (meta-learning-based) PFL methods are considered in our experiment:

Reptile [22] inserts a meta-learning fine-tuning phase of [70, 55] after the federated averaging algorithm stage to provide a reliable personalized model.

MAML[21] enhances federated learning by integrating a MAML-based meta-learner [55], dividing the dataset into support and query sets for more robust local training.

PerFedAvg [23] method focuses on the convergence analysis of the HF-MAML algorithm [71] in the FL scenario, offering a provably convergent method based on MAML to tackle non-convex functions.

The following Bayesian FL algorithm are also compared in our experiment:

FedBE [31] enhances robust aggregation by adopting a Bayesian inference approach, sampling highquality global models, and combining them through Bayesian model ensemble using Gaussian or Dirichlet distributions fitted to local models.

pFedGP [34] employs Gaussian processes for personalized federated learning, leveraging shared kernel functions parameterized by a neural network and a novel inducing point approach for the generalization in the low data regime.

The following Pruning algorithm is also compared in our compression experiment in the appendix:

SNIP [101, 102] algorithm is a deep learning pruning technique that identifies and removes less important connections in neural networks before the training begins using a small subset of the dataset. It is compared with our algorithm in the model compression experiment in the FL environment.

Additionally, Fine-tuning (FT) in the baseline's name indicates fine-tuning the local models with a few steps before evaluation within the FL processes, which is similar to the adaptation step of the optimization-based meta-learning.

E Implementation Details.

Models. To maintain consistency with previous research, we employ the widely adopted CNN model for all algorithms and baselines [25, 74, 3]. 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 of an embedding layer, followed by two consecutive blocks containing a linear layer and a LeakyReLU activation function, and one block containing a linear layer with exponential activation to output the dropout logit parameter α . The dimension of client embedding e m is proportional to the number of clients M and is calculated as (1 + M/ 4) . The hidden units' size in the hypernetwork was set to 200 . The predicted dropout logit parameter is then applied to each weight of the MetaVD 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 MetaVD only in one fully-connected layer yielded significant performance improvements across all experiments.

Hyperparameters. For all datasets, we set T to 1000 to 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 a 1-step update. In order to ensure a fair comparison between the algorithms, the results presented in all of our experiments are obtained with the optimal hyperparameters for each model. To do so, we conducted an extensive parameter optimization using an optimization tool called Optuna 4 [103]. We employed both the 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 η 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 rate γ is required, for which we searched within the range of [0 . 01 , 0 . 02 , 0 . 03 , 0 . 04 , 0 . 05 , 0 . 06 , 0 . 07 , 0 . 08 , 0 . 09 , 0 . 1] . For MetaVD, 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, we will release all code, including baselines, on our GitHub repository.

F Experiment Outline

To evaluate the effectiveness and robustness of the proposed PFL methods, we tested the algorithms under various FL scenarios, including:

Degree of Data Heterogeneity: we assessed the performance of each algorithm in scenarios where data from different clients are heterogeneous, considering factors such as 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 a hypernetwork in MetaVD compared to naive Variational Dropout (VD) or 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: we used Expected Calibration Error (ECE) measures to evaluate the accuracy of probabilistic predictions made by the predictive models. These measures help determine whether a model is overconfident or underconfident in its predictions and identify any biases in the model's predictions.

Communication Efficiency: we 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 to enhance the performance of the model, but it typically results in a suboptimal global model. We tested our approach in FL with multi-domain learning.

Model Compression: we compared MetaVD's compression capabilities with the existing pruning algorithm and evaluated the performance of MetaVD with and without compression.

G Additional Results with Non-IID Settings

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 CIFAR-10 and CIFAR-100 datasets with varying degrees of heterogeneous partitions. We randomly selected 30 out of 130 clients as non-participating out-of-distribution (OOD) clients, who were not involved during the training phase. Table 11 and Table 12 display the weighted average accuracy while adjusting the non-IID degrees with different Dirichlet parameters ˙ α . A smaller ˙ α corresponds to a higher degree of heterogeneity. Test(%) denotes the weighted average accuracy of participating clients' test samples, OOD(%) signifies the weighted average accuracy of the non-participants, and ∆ represents the participation generalization gap calculated as the difference between OOD accuracy and test accuracy.

In both datasets, 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, our method consistently improves OOD accuracy across all degrees of 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.

Table 8: Results of classification accuracies with different (non-IID) heterogeneity strengths of ˙ α = 5 . 0 , ˙ α = 0 . 5 , and ˙ α = 0 . 1 in the CIFAR-10 dataset. We report the results of participating clients during the training (Test) dataset and non-participating clients (OOD). The higher the better.

In the original pFedGP paper [34], the data samples were partitioned across clients according to [25, 24], resulting in each client having 2 and 10 classes for the CIFAR-10 and CIFAR-100 datasets, respectively. In contrast, in our non-IID experiment, we do not predefine the number of classes each client has; instead, it is 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 [104, 105] occurs between the training and testing distributions of the clients). To ensure that pFedGP would run on the unseen labels, we circumvented this limitation by allowing the client to have access to 40 or 50 data points for each class while evaluating the model's performance.

H Additional Ablation Results in FEMNIST Dataset

In this experiment, we 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 a global model parameter as in FedAvg. In EnsembleVD, we updated client-specific dropout rates in a manner analogous to the local adaptation step of MetaVD but maintained independent variational dropout rates for different clients on the server. In contrast, MetaVD utilized a hypernetwork to learn the 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-based hypernetwork surpasses all other baselines

in classification accuracy. Here too, Reptile+MetaVD consistently outperforms other methods. We noticed that in baselines like VD or EnsembleVD, client-specific dropout rates were not fully optimized due to restricted client participation. Consequently, the initial dropout rates and KL divergence loss in VD and EnsembleVD remained fairly static. Conversely, in MetaVD, both dropout rates and KL divergence loss converged in all tests. This observation underscores that MetaVD's hypernetwork presents a more data-efficient approach to learning the client-specific model uncertainty compared to other baselines.

I Additional Results on Uncertainty Calibration

Addressing uncertainty calibration issues is crucial for enhancing the accuracy and reliability of a model's predictions, particularly when the model is employed for making significant decisions. Uncertainty calibration metrics evaluate the accuracy of probabilistic predictions made by a predictive model. By identifying and addressing any calibration issues, it becomes possible to improve the model's prediction accuracy and reliability, which is of essential importance in decision-making processes. In a Federated Learning environment where 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 of our proposed algorithm in improving uncertainty calibration, demonstrating its potential to strengthen model performance in various real applications.

Expected Calibration Error (ECE) and Maximum Calibration Error (MCE) are widely used metrics for 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

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 1 indicates complete miscalibration. MCE, on the other hand, is akin to ECE but focuses on the largest gap for any group rather than the weighted average.

Table 11 and Table 12 summarize the ECE and MCE results for varying non-IID degrees in CIFAR-100 and CIFAR-10 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 local data. In the CIFAR-10 dataset, all methods that incorporate MetaVD exhibit a decline in ECE and MCE values as ˙ α decreases, in contrast to standard FL or PFL methods that show an increase in these values. MetaVD effectively mitigates overfitting to local data and enhances calibration by leveraging the heterogeneity present in the training data. Consequently, methods employing MetaVD consistently achieve the lowest ECE and MCE values in almost all scenarios.

Examining a reliability diagram can provide a visual comparison of uncertainty calibration results, as it illustrates the relationship between prediction probabilities and true labels. Ideally, if a model predicts a specific class with a certain probability, the actual label should correspond to that confidence level, placing the point on the reliability diagram'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 each point from the diagonal line, while MCE indicates the maximum distance of any point from the diagonal line, offering a comprehensive understanding of the model's calibration performance.

Figure 4 and Figure 5 present the reliability diagrams for CIFAR-100 and CIFAR-10 with ˙ α = 0 . 5 respectively. These figures 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 MetaVD significantly improves the uncertainty calibration of the models as well.

J Additional Results on Muti-domain Datasets

Unlike existing federated learning algorithms that usually assume a single-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 multi-domain datasets in which we assume each client can have data

from different domains. Multi-domain learning [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 a suboptimal global model. In the context of federated learning, it becomes even more critical to apply multi-domain learning effectively. We use three different FL datasets to 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 each client's local CIFAR-100 data.

Table 13, 14, and 15 illustrate the classification accuracies for Test and OOD clients with varying degrees of heterogeneity. Employing MetaVD 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 MetaVD can improve the versatility and performance of FL algorithms when applied to a broad range of multi-domain datasets.

K Additional Results on Model Compression

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-10 dataset. MetaVD+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, it is crucial to minimize communication costs by reducing both the size of exchanged model parameters and the number of communication rounds. MetaVD 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-10 and CIFAR-100 datasets, comparing our method to the baseline Pruning algorithm SNIP [101]. For the implementation of the SNIP algorithm in the context of Federated Learning, we referred to 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 17 demonstrate the results of Test( % ), OOD( % ), and Sparsity( % ) percentage of the models with or without MetaVD+DP. 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 a dropout rate greater than 0 . 9 are dropped after 150 rounds. We have set these 150 rounds to allow MetaVD

Table 17: Results of study on model compression with different (non-IID) heterogeneity strengths of ˙ α = 5 . 0 , ˙ α = 0 . 5 , and ˙ α = 0 . 1 in the CIFAR-100 dataset. MetaVD+DP does not communicate the model parameters when the dropout rate is larger than 0.9.

to adequately learn the client-specific dropout rates. The target sparsity of the SNIP algorithm was set to 0 . 7 . Note that, in our MetaVD+DP setting, even though we dropped weights with a dropout rate greater than 0 . 9 , the actual sparsity values are observed to be lower and tend to be around 70 % , as indicated in the tables.

Overall, most models removed over 70% of their parameters in the personalized layer without losing much performance. In CIFAR-10, as heterogeneity increased, the performance gap between SNIP and MetaVD+DP tended to increase. For instance, in CIFAR-10 with ˙ α = 5 . 0 , the OOD performance of SNIP model and MAML+MetaVD+DP model were 81 . 59 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 case of Reptile+MetaVD+DP and PerFedAvg+MetaVD+DP. In CIFAR-100 dataset, the performance gap, according to the increase in heterogeneity, was not noticeable. Surprisingly, dropping model parameters could lead to performance improvements in some cases. For example, in CIFAR-100 with ˙ α = 0 . 5 , the PerFedAvg+MetaVD+DP model'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 can be applied to numerous FL algorithms without significantly increasing communication costs. It also boosts the generalization performance for new, unseen clients. These aspects make our approach a significant addition to the field of Federated Learning.

(a) FedAvg

(b) FedAvg + MetaVD

1.0

1.0

ECE = 0.53%

ECE = 0.32%

MCE =32.35%

MCE = 26.60%

0.8

0.8

0.6

0.6

0.4

0.4

0.2

0.2

0.0

0.0

8.0

0.2（

0.4

0.6

0.81.0

8.0

0.2

0.4

0.6

0.8

1.0

(c) Reptile

1.0

1.0

ECE = 0.77%

ECE = 0.32%

MCE = 46.19%

MCE =23.86%

0.8

0.8

0.6

0.6

0.4

0.4

0.2

0.2

0.0

8.0

0.2

0.4

0.6

0.8

1.0

0.0

8.0

0.2

0.4

0.6

0.8

1.0

(e) MAML

(f) 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.6

0.4

0.4

0.2

0.2

0.8.0

0.0

0.2

0.4

0.6

0.8

1.0

8.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.8

0.6

0.6

0.4

0.4

0.2

0.2

0.80

0.20.40.6

0.8

1.0

0.20.4

0.6

¥0.8

¥1.0

Confidence

(a) FedAvg

(b) FedAvg + MetaVD

1.0

1.0

ECE = 2.39%

ECE = 1.76%

MCE = 26.45%

MCE = 21.17%

0.8

0.8

0.6

0.6

0.4

0.4

0.2

0.2

0.0

0.2

0.40.6

0.8

1.0

8.0

0.2

0.4

0.6

0.8

1.0

(c) Reptile

1.0

1.0

ECE = 3.06%

ECE = 2.16%

MCE =38.89%

MCE = 29.84%

0.8

0.8

0.6

0.6

0.4

0.4

0.2

0.2

0.8.0

0.0

0.2

0.4

0.6

0.8

1.0

8.0

0.2

0.4

0.6

0.8

1.0

(e) MAML

(f) MAML + MetaVD

1.0

1.0

ECE = 2.49%

ECE = 1.97%

MCE =30.67%

MCE = 21.51%

0.8

0.8

0.6

0.6

0.4

0.4

0.2

0.2

0.0

0.2

0.40.6

0.81.0

8.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.8

0.8

0.6

0.6

0.4

0.4

0.2

0.2

0.80

0.20.40.6

0.8

1.0

0.20.40.6

¥0.8

¥1.0

Confidence

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nMetaVD 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.\n\nWEAKNESSES:\nIn 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.\n\nQUESTIONS:\n1. “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.\n\nLIMITATIONS:\nThe authors discussed the limitations of their work, but they did not address the potential negative societal impact it may have.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nThe 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 ...)\n\nQUESTIONS:\n- 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.\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nNA.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\nA 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”.\n\nQUESTIONS:\n(see \"weaknesses\" above).\n\nLIMITATIONS:\nThe authors haven’t discussed the limitations of their work.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}