Stein Latent Optimization for Generative Adversarial Networks

1 INTRODUCTION . GANs have shown remarkable results in the synthesis of realistic data conditioned on a specific class ( Odena et al. , 2017 ; Miyato & Koyama , 2018 ; Kang & Park , 2020 ) . Training conditional GANs requires a massive amount of labeled data ; however , data are often unlabeled or possess only a few labels . For unsupervised conditional generation , the salient attributes of the data are first identified by unsupervised learning and used for conditional generation of data . Recently , several unsupervised conditional GANs have been proposed ( Chen et al. , 2016 ; Mukherjee et al. , 2019 ; Pan et al. , 2021 ; Armandpour et al. , 2021 ) . By maximizing a lower bound of mutual information between latent codes and generated data , they cluster the attributes of the underlying data distribution in their latent spaces . These GANs achieve satisfactory performance when the salient attributes of data are balanced . However , the attributes of real-world data can be imbalanced . For example , in the CelebA dataset ( Liu et al. , 2015 ) , examples with one attribute ( not wearing eyeglasses ) outnumber the other attribute ( wearing eyeglasses ) . Similarly , the number of examples with disease-related attributes in a biomedical dataset might be miniscule ( Hwang et al. , 2019 ) . Thus , the imbalanced nature of real-world attributes must be considered for unsupervised conditional generation . Most of existing unsupervised conditional GANs are not suitable for real-world attributes , because they assume balanced attributes if the imbalance ratio is unknown ( Chen et al. , 2016 ; Mukherjee et al. , 2019 ; Pan et al. , 2021 ) . Examples where existing methods fail to learn imbalanced attributes are shown in Figure 1 ( a ) , ( b ) and ( c ) . In this paper , we propose unsupervised conditional GANs , referred to as Stein Latent Optimization for GANs ( SLOGAN ) . We define the latent distribution of the GAN models as Gaussian mixtures to enable the imbalanced attributes to be naturally clustered in a continuous latent space . We derive reparameterizable gradient identities for the mean vectors , full covariance matrices , and mixing coefficients of the latent distribution using Stein ’ s lemma . This enables stable learning and makes ∗Corresponding author 9/21 latent distribution parameters , including the mixing coefficient , learnable . We then devise a GAN framework with an encoder network and an unsupervised conditional contrastive loss ( U2C loss ) , which can interact well with the learnable Gaussian mixture prior ( Figure 2 ) . This framework facilitates the association of data generated from a Gaussian component with a single attribute . For the synthetic dataset , our method ( Figure 1 ( d ) ) shows superior performance on unsupervised conditional generation , with the accurately learned mixing coefficients . We performed experiments on various real-world datasets including MNIST ( LeCun et al. , 1998 ) , Fashion-MNIST ( Xiao et al. , 2017 ) , CIFAR-10 ( Krizhevsky et al. , 2009 ) , CelebA ( Liu et al. , 2015 ) , CelebA-HQ ( Karras et al. , 2017 ) , and AFHQ ( Choi et al. , 2020 ) using architectures such as DCGAN ( Radford et al. , 2016 ) , ResGAN ( Gulrajani et al. , 2017 ) , and StyleGAN2 ( Karras et al. , 2020 ) . Through experiments , we verified that the proposed method outperforms existing unsupervised conditional GANs in unsupervised conditional generation on datasets with balanced or imbalanced attributes . Furthermore , we confirmed that we could control the attributes to be learned when a small set of probe data is provided . The contributions of this work are summarized as follows : • We propose novel Stein Latent Optimization for GANs ( SLOGAN ) . To the best of our knowledge , this is one of the first methods that can perform unsupervised conditional generation by considering the imbalanced attributes of real-world data . • To enable this , we derive the implicit reparameterization for Gaussian mixture prior using Stein ’ s lemma . Then , we devise a GAN framework with an encoder and an unsupervised conditional contrastive loss ( U2C loss ) suitable for implicit reparameterization . • SLOGAN significantly outperforms the existing methods on unsupervised learning tasks , such as cluster assignment , unconditional data generation , and unsupervised conditional generation , on datasets that include balanced or imbalanced attributes . 2 BACKGROUND . 2.1 GENERATIVE ADVERSARIAL NETWORKS . Unsupervised conditional generation Several models including InfoGAN ( Chen et al. , 2016 ) , ClusterGAN ( Mukherjee et al. , 2019 ) , Self-conditioned GAN ( Liu et al. , 2020 ) , CD-GAN ( Pan et al. , 2021 ) , and PGMGAN ( Armandpour et al. , 2021 ) have been proposed to perform conditional generation in a completely unsupervised manner . However , these models primarily have two drawbacks : ( 1 ) Most of these methods embed the attributes in discrete variables , which induces discontinuity among the embedded attributes . ( 2 ) Most of them assume uniform distributions of the attributes , and thus fail to learn the imbalance in attributes when the imbalance ratio is not provided . In Appendix B.2 , we discuss the above models in detail . By contrast , our work addresses the aforementioned limitations by combining GANs with the gradient estimation of the Gaussian mixture prior via Stein ’ s lemma and representation learning on the latent space . 5 Method GANs with Gaussian mixture prior DeLiGAN ( Gurumurthy et al. , 2017 ) is analogous to the proposed method , as it assumes a Gaussian mixture prior and learns the mean vectors and covariance matrices via the reparameterization trick . However , DeLiGAN assumes uniform mixing coefficients without updating them . As a result , it fails to perform unsupervised conditional generation on datasets with imbalanced attributes . In addition , it uses the explicit reparameterization trick , which inevitably suffers from high variance in the estimated gradients . This will be discussed further in Section 2.3 . 2.2 CONTRASTIVE LEARNING . Contrastive learning aims to learn representations by contrasting neighboring with non-neighboring instances ( Hadsell et al. , 2006 ) . In general , contrastive loss is defined as a critic function that approximates the log density ratio log p ( y|x ) /p ( y ) of two random variables X and Y . By minimizing the loss , the lower bound of the mutual information I ( X ; Y ) is approximately maximized ( Poole et al. , 2019 ) . Several studies have shown that contrastive losses are advantageous for the representation learning of imbalanced data ( Kang et al. , 2021 ; 2020 ; Wanyan et al. , 2021 ) . Motivated by these observations , we propose a contrastive loss that cooperates with a learnable latent distribution . 2.3 GRADIENT ESTIMATION FOR GAUSSIAN MIXTURE . Stein ’ s lemma Stein ’ s lemma provides a first-order gradient identity for a multivariate Gaussian distribution . The univariate case of Stein ’ s lemma can be described as follows : Lemma 1 . Let function h ( · ) : R 7→ R be continuously differentiable . q ( z ) is a univariate Gaussian distribution parameterized by the mean µ and variance σ . Then , the following identity holds : Eq ( z ) [ σ−1 ( z − µ ) h ( z ) ] = Eq ( z ) [ ∇zh ( z ) ] ( 1 ) Lin et al . ( 2019b ) generalized Stein ’ s lemma to exponential family mixtures and linked it to the implicit reparameterization trick . Stein ’ s lemma has been applied to various fields of deep learning , including Bayesian deep learning ( Lin et al. , 2019a ) and adversarial robustness ( Wang et al. , 2020 ) . To the best of our knowledge , our work is the first to apply Stein ’ s lemma to GANs . Reparameterization trick A simple method to estimate gradients of the parameters of Gaussian mixtures is explicit reparameterization , used in DeLiGAN . When the c-th component is selected according to the mixing coefficient p ( c ) , the latent variable is calculated as follows : z = µc+ϵ ·Σ1/2c , where ϵ ∼ N ( 0 , I ) . Derivatives of a loss function ∂L ( z ) ∂µ and ∂L ( z ) ∂Σ only update the mean and covariance matrices of the selected ( c-th ) component µc and Σc , respectively . Gradient estimation using explicit reparameterization is unbiased ; however , it has a distinctly high variance . For a single latent vector z , the implicit reparameterization trick ( Figurnov et al. , 2018 ) updates the parameters of all the latent components . Gradient estimation using implicit reparameterization is unbiased and has a lower variance , which enables a more stable and faster convergence of the model . The gradients for the parameters of the Gaussian mixture prior in our method are implicitly reparameterizable . 3 PROPOSED METHOD . In the following sections , we propose Stein Latent Optimization for GANs ( SLOGAN ) . We assume a Gaussian mixture prior ( Section 3.1 ) , derive implicit reparameterization of the parameters of the mixture prior ( Section 3.2 ) , and construct a GAN framework with U2C loss ( Section 3.3 ) . Additionally , we devise a method to manipulate attributes to be learned if necessary ( Section 3.4 ) . An overview of SLOGAN is shown in Figure 2 . 3.1 GAUSSIAN MIXTURE PRIOR . We consider a GAN with a generator G : Rdz 7→ Rdx and a discriminator D : Rdx 7→ R , where dz and dx are the dimensions of latent and data spaces , respectively . In the latent space Z ∈ Rdz , we consider a conditional latent distribution q ( z|c ) = N ( z ; µc , Σc ) , c = 1 , ... , K , where K is the number of components we initially set and µc , Σc are the mean vector and covariance matrix of the cth component , respectively . Subsequently , we consider a Gaussian mixture q ( z ) = ∑K c=1 p ( c ) q ( z|c ) parameterized by µ = { µc } Kc=1 , Σ = { Σc } Kc=1 and π = { πc } Kc=1 = { p ( c ) } Kc=1 as the prior . We hypothesize that a mixture prior in a continuous space could model some continuous attributes of real-world data ( e.g. , hair color ) more naturally than categorical priors which could introduce discontinuity ( Mukherjee et al. , 2019 ) . Because we use implicit reparameterization of a mixture of Gaussian priors ( derived in Section 3.2 ) , SLOGAN can fully benefit from implicit reparameterization and U2C loss . By contrast , the implicit reparameterization of prior distributions that do not belong to the exponential family ( e.g. , categorical priors ) remains an open question . In the experiments , the elements of µc were sampled from N ( 0 , 0.1 ) , and we selected Σc = I and πc = 1/K as the initial values . For the convenience of notation , we define the latent distribution q = q ( z ) , the mixing coefficient πc = p ( c ) , and δ ( z ) = { δ ( z ) c } Kc=1 , where δ ( z ) c = q ( z|c ) /q ( z ) . q ( c|z ) , the responsibility of component c for a latent vector z , can be expressed as follows : q ( c|z ) = q ( c , z ) q ( z ) = q ( z|c ) p ( c ) q ( z ) = δ ( z ) cπc ( 2 )

Summary: This work proposes a method for performing unsupervised conditional generation from imbalanced attributes. The key components of this work are the following:  * A (architecture agnostic) GAN with a GMM latent space, where we have one component per attribute in the data.   * The parameters of the components $(\mu, \Sigma)$ are learned via implicit reparameterization, and leveraging stein's lemma to derive gradients for $\mu$ and $\Sigma$.   * The components are encouraged to be distinct by introducing a contrastive objective. An encoder is used to obtain a low dimensional vector representation of a generated sample, and this sample is encouraged to be close to the mean vector of the corresponding component in the latent space.