EP-GAN: Unsupervised Federated Learning with Expectation-Propagation Prior GAN

1 INTRODUCTION . Generative Adversarial Networks ( GANs ) ( Goodfellow et al. , 2014 ) are overwhelming in unsupervised learning tasks due to their expressive power in modeling fine-grained data distributions . GANs usually consist of a generator , which takes white noises as input and generates data points , and a discriminator , which encourages the generator to generate high-quality data points via discriminating whether a data point follows a real distribution . Successful training of GANs requires large-scale public data in a centralized setting . Recently there has also been an increasing need for training GANs on decentralized private data . Unfortunately , due to concern about data privacy , conventional ways of aggregating cross-silo data have become controversial . On the other hand , federated learning is an increasingly important machine learning paradigm where many clients jointly train a powerful global model with cross-silo training data . The major target of federated learning is to utilize the massive data created and collected by different clients while obeying privacy protection regulations such as the European Union ’ s General Data Protection Regulation ( GDPR ) ( Voss , 2016 ) . The most representative federated learning algorithm is FederatedAveraging ( FedAvg ) ( McMahan et al. , 2017 ) , which successfully trains a powerful global model while keeping training data on each client ( i.e. , each mobile phone ) on mobile applications . Since then , FedAvg has been quickly adopted in many other applications , such as medical image classification ( Choudhury et al. , 2019 ) . Recently , more and more federated learning researchers turn their attention to a more practical setting where the data distributions are non-i.i.d . A federation usually consists of different clients from diverse sources whose data distributions are intrinsically distinct . Vanilla federated learning algorithms fail under such non-i.i.d settings , as the simple averaging tends to over-smooth the distributions of each client . Such an over-smoothing effect is exacerbated in GANs , which notoriously often suffer from mode collapse . To gain initial insight into this issue , we conduct a toy experiment of learning a mixture of Gaussians ( see Section 5.1 for more details ) . Fig . 1 shows the learned density map of the data centralized setting ( the first column ) and the extremely non-i.i.d federated learning setting ( the last column ) , where ‘ red ’ indicates higher scores , and ‘ blue ’ indicates lower scores . Ideally , GANs should be able to rebuild four red circles around the data points , as shown in the centralized settings . We can observe that although both WGAN ( Goodfellow et al. , 2014 ) and ProbGAN ( He et al. , 2019 ) work well in the centralized setting , but in federated learning , their locally learned data distributions are over-smoothed after FedAvg , resulting in a collapsed mode . This phenomenon can be interpreted from the Bayesian perspective . WGAN can be treated as a special case of Bayesian GAN by assuming a uniform prior instead of a Gaussian prior for GAN ’ s parameters and fitting the point estimate of GAN ’ s parameters during training . ProbGAN is a special case of Bayesian GAN , which simply adopts the generator distribution at the previous turn as the prior . Therefore , both WGAN and ProbGAN have the similar issue of enforcing an over-general prior , limiting their capability in capturing the latent multi-modal structure of non-i.i.d data partitions under federated learning setting and thus leading to mode collapse . Similar problems are also observed in the VAE literature , where researchers refer to such problems as ” posterior collapse ” ( i.e. , the learned posterior collapses to the over-general prior ) . To tackle this challenge , we present a novel GAN , dubbed expectation propagation prior GAN ( EP-GAN ) , which is equipped with a complex prior that is learned automatically from data partitions . EP-GAN treats the prior of parameter distributions in each client as a latent variable and models the joint distribution of all clients . In doing so , EP-GAN introduces several approximate factors , where each approximates the posterior contribution from each data partition . As these approximate factors are updated iteratively and parallelly in a similar spirit of EP , we name the prior of our EP-GAN as EP-prior . The major contributions of our work are three-fold : • We identify the mode collapse problem of Bayesian GANs under non-i.i.d cross-silo unsupervised learning scenarios . Inspired by EP , we propose an effective solution , dubbed EP-GAN , that imposes a partition-invariant prior distribution on a Bayesian GAN . • In the solution , we derive closed-form expressions of the moments of the hybrid/tilted distribution approximations for neural networks . Unlike the previous work on EP with deep learning that requires numerical quadratures , our EP-GAN obtains a relatively more efficient closed-form solution without requiring numerical quadratures . • Our proposed method is evaluated on real-world image and speech datasets . The experimental results demonstrate that our framework effectively alleviates the performance deterioration caused by non-i.i.d federated data compared to other state-of-the-art Bayesian GANs . 2 RELATED WORK . In this work , we develop a new Bayesian GAN equipped with an expectation propagation ( EP ) prior for unsupervised federated learning under non-i.i.d settings ; therefore , recent studies on federated learning , Bayesian GANs and EP with neural networks are reviewed in this section . Federated Learning with Non-i.i.d Data . Federated learning , as a new collaborative learning paradigm , has been gaining more attention in recent years ( McMahan et al. , 2017 ; Yang et al. , 2019 ) , and the challenges of federated learning from non-i.i.d data have been noticed ( Li et al. , 2021 ; Sahu et al. , 2018 ; Zhao et al. , 2018 ; Xie et al. , 2019 ; Yu et al. , 2020 ) . However , they only focus on the supervised federated learning settings , where the label information can help reduce the heterogeneity ; in contrast , we focus on the unsupervised federated learning setting , where the heterogeneity embeds in the distributions across clients and these methods can not handle well . On the other hand , some works have been proposed to adapt GANs to federated settings to deal with unsupervised federated learning tasks , but they are either proposed to improve the privacy protection of multiple GANs ( Augenstein et al. , 2019 ; Zhang et al. , 2018 ; Jordon et al. , 2018 ; Xie et al. , 2018 ) or improve the communication efficiency when training in a federated manner ( Rasouli et al. , 2020 ) . More importantly , they ignore another critical challenge : mode collapse of federated GANs under the non i.i.d federated learning setting . Addressing this problem is one of the main focuses of our model . The detailed discussion of existing federated GANs with non-i.i.d data is reported in Appendix C. Bayesian GANs . Recent advances in Bayesian formulation of GANs , such as BayesGAN ( Saatchi & Wilson , 2017 ) and ProbGAN ( He et al. , 2019 ) , are proposed to address the notorious mode-collapse issue . The BayesGAN and ProbGAN assume a Gaussian prior and a carefully constructed prior for GANs , respectively . However , from the Bayesian perspective , both BayesianGAN and ProbGAN have the similar issue of enforcing an over-general prior , limiting their capability in capturing the latent multi-modal structure of non-iid data partitions under federated learning setting and thus leading to mode collapse . In contrast , our proposed EP-GAN successfully avoids such an issue by imposing a partition-invariant prior distribution inferred by EP . EP with Neural Networks . To efficiently model the joint distribution from partitioned data , we follow the spirit of Expectation Propagation ( EP ) . EP is one of the most popular Bayesian inference methods in Bayesian neural networks ( Minka , 2013 ) , which approximates the posterior with exponential-family factors that are iteratively updated via moment matching . However , the moment matching can be intractable if the likelihood term involves many latent variables or has a complex form , which results in the intractability of the moments of the hybrid/tilted distribution . An intuitive solution is to approximate the likelihood term by numerical quadrature ( Jylänki et al. , 2014 ; Soudry et al. , 2014 ) , but it fails to scale to large datasets or complex neural networks . The closest work to ours is proposed by Bui et al . ( 2018 ) , which adopts variational inference ( VI ) to simulate EP , allowing to identify new ways of performance VI on federated settings with synthetic data , but their EP update is inefficient as it is simulated with gradient descent with a deliberately designed KL loss . In contrast , we propose a closed-form solution for moment matching without requiring numerical approximation and is more efficient than estimating through gradient descent . The detailed comparisons with related works are described in Appendix C , and some of the representative methods are compared with empirical experiments , shown in Section 5 . 3 PRELIMINARIES . Federated Learning . In this paper , we consider the problem of training Bayesian GANs under the cross-silo FL setting1 ( Kairouz et al. , 2019 ) , where the datasets are non-identically and isolatedly distributed on different clients . To avoid losing the focus , we leave the formulation of federated learning in Appendix D. Expectation Propagation ( EP ) . Before introducing our EP-GAN , we begin by reviewing the traditional EP algorithm ( Minka , 2013 ) . EP is a deterministic approximation algorithm that is often used for Bayesian inference of posterior distributions of model parameters , which is believed to be able to provide significantly more accurate approximations than variational inference methods ( Jordan et al. , 1999 ) . Consider a dataset X = { xi } Ni=1 sampled from a distribution p ( X|θ ) whose parameter θ follows the prior p0 ( θ ) . As the posterior distribution p ( θ | X ) is computationally intractable , EP attempts to approximate it with a tractable approximating distribution q ( θ ) , which can be further factorized into multiple approximate factors : q ( θ ) = p0 ( θ ) ∏N i=1 qi ( θ ) , where the approximate factor qi ( θ ) is 1In this work , we focus on presenting a novel solution for modeling non-i.i.d federated data , without diving into the privacy and security protection techniques . As our proposed method is a mild modification to the existing federated learning algorithm in implementation , the current advanced privacy and security protection techniques such as secure aggregation ( Bonawitz et al. , 2017 ) can be directly applied to our method . iteratively refined so that they capture the contribution of each datapoint xi to the posterior p ( θ | X ) : p ( θ | X ) ∝ p0 ( θ ) N∏ i=1 li ( xi ) , ( 1 ) where the likelihood li ( xi ) = p ( xi | θ ) and X represents the data distribution . More specifically , EP iterates over the following four steps : 1 . Construct the cavity distribution by removing one of the approximate factor , i.e. , the i-th factor . It can be written as : q−i ( θ ) = p0 ( θ ) ∏ j 6=i qj ( θ ) 2 . Integrate the likelihood contribution li ( xi ) to the cavity to produce the hybrid distribution : hi ( θ ) ∝ q−i ( θ ) li ( xi ) . 3 . Update the parameters of the i-th approximate factor qi ( θ ) through minimizing the KL divergence between the hybrid distribution hi ( θ ) and the approximating distribution q ( θ ) , namely , KL [ hi ( θ ) ‖q−i ( θ ) qi ( θ ) ] . 4 . Update the approximating distribution q ( θ ) by including the updated approximate factor qi ( θ ) : q ( θ ) ← q−i ( θ ) qi ( θ ) . It is worth noting that when qi ( θ ) is assumed to follow an exponential family distribution ( e.g . the Gaussian ) , the minimizing of KL divergence in Step 3 can be reduced to the moment matching ( ichi Amari & Nagaoka , 2000 ) . However , when applying EP to deep neural networks , this moment matching step is computationally intractable , requiring numerical approximation or sampling and therefore have to compromise between accuracy and efficiency . Addressing this problem is one of the main focuses of our model .

Summary: The goal of this paper is to train a Bayesian GAN on non-i.i.d. federated data. Specifically, the authors propose to adopt the newly-introduced expectation propagation (EP) prior, being partition-invariant, to address the non-i.i.d. federated challenge. Experiments on synthetic and real datasets are conducted.