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 authors targeted federated generative modelling in an unsupervised setting. Specifically, the work is built on top of Bayesian GANs. In order to aggregate the information from different clients, the authors proposed to use expectation propagation (EP). It makes sense. Despite being a well established Bayesian inference algorithm, EP operating on the neural network parameters can suffer from intractability. The authors presented a low complexity solution. The experiment results showed improved FID and IS over multiple baseline methods. However, the overall performance is quite poor on the rather simple dataset CIFAR10, owing to the scalability issue of Bayesian models.