Dirichlet Variational Autoencoder

1 INTRODUCTION . A Variational Autoencoder ( VAE ) ( Kingma & Welling , 2014c ) brought success in deep generative models ( DGMs ) with a Gaussian distribution as a prior distribution ( Jiang et al. , 2017 ; Miao et al. , 2016 ; 2017 ; Srivastava & Sutton , 2017 ) . If we focus on the VAE , the VAE assumes the prior distribution to be N ( 0 , I ) with the learning on the approximated µ̂ and Σ̂ . Also , Stick-Breaking VAE ( SBVAE ) ( Nalisnick & Smyth , 2017 ) is a nonparametric version of the VAE , which modeled the latent dimension to be infinite using a stick-breaking process ( Ishwaran & James , 2001 ) . While these VAEs assume that the prior distribution of the latent variables to be continuous random variables , recent studies introduce the approximations on discrete priors with continuous random variables ( Jang et al. , 2017 ; Maddison et al. , 2017 ; Rolfe , 2017 ) . The key of these approximations is enabling the backpropagation with the reparametrization technique , or the stochastic gradient variational Bayes ( SGVB ) estimator , while the modeled prior follows a discrete distribution . The applications of these approximations on discrete priors include the prior modeling of a multinomial distribution which is frequently used in the probabilistic graphical models ( PGMs ) . Inherently , the multinomial distributions can take a Dirichlet distribution as a conjugate prior , and the demands on such prior have motivated the works like Jang et al . ( 2017 ) ; Maddison et al . ( 2017 ) ; Rolfe ( 2017 ) that support the multinomial distribution posterior without explicit modeling on a Dirichlet prior . When we survey the work with a explicit modeling on the Dirichlet prior , we found a frequent approach such as utilizing a softmax Laplace approximation ( Srivastava & Sutton , 2017 ) . We argue that this approach has a limitation from the multi-modality perspective . The Dirichlet distribution can exhibit a multi-modal distribution with parameter settings , see Figure 1 , which is infeasible to generate with the Gaussian distribution with a softmax function . Therefore , the previous continuous domain VAEs can not be a perfect substitute for the direct approximation on the Dirichlet distribution . Utilizing a Dirichlet distribution as a conjugate prior to a multinomial distribution has an advantage compared to the usage of a softmax function on a Gaussian distribution . For instance , Figure 1 illustrates the potential difficulties in utilizing the softmax function with the Gaussian distribution . Given the three-dimensional probability simplex , the Gaussian-Softmax distribution can not generate the illustrated case of the Dirichlet distribution with a high probability measure at the vertices of the simplex , i.e . the multi-modality where the necessity was emphasized in Hoffman & Johnson ( 2016 ) . Additionally , the Griffiths-Engen-McCloskey ( GEM ) distribution ( Pitman , 2002 ) , which is the prior distribution of the SBVAE , is difficult to model the multi-modality because the sampling procedure of the GEM distribution is affected by the rich-get-richer phenomenon , so a few components tend to dominate the weight of the samples . This is different from the Dirichlet distribution that does not exhibit such phenomenon , and the Dirichlet distribution can fairly distribute the weights to the components , and the Dirichlet distribution is more likely to capture the multi-modality by controlling the prior hyper-parameter ( Blei et al. , 2003 ) . Then , we conjecture that an enhanced modeling on Dirichlet prior is still needed 1 ) because there are cases that the Gaussian-Softmax approaches , or the softmax Laplace approximation , can not imitate the Dirichlet distribution ; and 2 ) because the nonparametric approaches could be influenced by the biases that the Dirichlet distribution does not suffer from . Given these motivations for modeling the Dirichlet distribution with the SGVB estimator , this paper introduces the Dirichlet Variational Autoencoder ( DirVAE ) that shows the same characteristics of the Dirichlet distribution . The DirVAE is able to model the multi-modal distribution that was not possible with the Gaussian-Softmax and the GEM approaches . These characteristics allow the DirVAE to be the prior of the discrete latent distribution , as the original Dirichlet distribution is . Introducing the DirVAE requires the configuration of the SGVB estimator on the Dirichlet distribution . Specifically , the Dirichlet distribution is a composition of the Gamma random variables , so we approximate the inverse Gamma cumulative distribution function ( CDF ) with the asymptotic approximation . This approximation on the inverse Gamma CDF becomes the component of approximating the Dirichlet distribution . We compared this approach to the previously suggested approximations , i.e . approaches with the Weibull distribution and with the softmax Gaussian distribution , and our approximation shows the best log-likelihood among the compared approximations . Moreover , we report that we had to investigate the component collapsing along with the research on DirVAE . It has been known that the component collapsing issue is resolved by the SBVAE because of the meaningful decoder weights from the latent layer to the next layer . However , we found that SBVAE has latent value collapsing issue resulting in many near-zero values on the latent dimensions that leads to the incomplete utilization of the latent dimension . Hence , we argue that Gaussian VAE ( GVAE ) suffers from the decoder weight collapsing , previously limitedly defined as component collapsing ; and SBVAE has a problem of the latent value collapsing . Finally , we suggest that the definition of component collapsing should be expanded to represent both cases of decoder weight and latent value collapsings . The proposed DirVAE shows neither the near-zero decoder weights nor the near-zero latent values , so the reconstruction uses the full latent dimension information in most cases . We investigated this issue because our performance gain comes from resolving the expanded version of the component collapsing . Due to the component collapsing issues , the existing VAEs have less meaningful latent values or could not effectively use its latent representation . Meanwhile , DirVAE does not have component collapsing due to the multi-modal prior which possibly leads to superior qualitative and quantitative performances . We experimentally showed that the DirVAE has more meaningful or disentangled latent representation by image generation and latent value visualizations . Technically , the new approximation provides the closed-form loss function derived from the evidence lower bound ( ELBO ) of the DirVAE . The optimization on the ELBO enables the representation learning with the DirVAE , and we test the learned representation from the DirVAE in two folds . Firstly , we test the representation learning quality by performing the supervised and the semi-supervised classification tasks on MNIST , OMNIGLOT , and SVHN . These classification tasks conclude that DirVAE has the best classification performances with its learned representation . Secondly , we test the applicability of DirVAE to the existing models , such as topic models with DirVAE priors on 20Newsgroup and RCV1-v2 . This experiment shows that the augmentation of DirVAE to the existing neural variational topic models improves the perplexity and the topic coherence , and most of best performers were DirVAE augmented . 2 PRELIMINARIES . 2.1 VARIATIONAL AUTOENCODERS . A VAE is composed of two parts : a generative sub-model and an inference sub-model . In the generative part , a probabilistic decoder reproduces x̂ close to an observation x from a latent variable z ∼ p ( z ) , i.e . x ∼ pθ ( x|z ) = pθ ( x|ζ ) where ζ = MLP ( z ) is obtained from a latent variable z by a multilayer perceptron ( MLP ) . In the inference part , a probabilistic encoder outputs a latent variable z ∼ qφ ( z|x ) = qφ ( z|η ) where η = MLP ( x ) is computed from the observation x by a MLP . Model parameters , θ and φ , are jointly learned by optimizing the below ELBO with the stochastic gradient method through the backpropagations as the ordinary neural networks by using the SGVB estimators on the random nodes . log p ( x ) ≥ L ( x ) = Eqφ ( z|x ) [ log pθ ( x|z ) ] − KL ( qφ ( z|x ) ||pθ ( z ) ) ( 1 ) In GVAE ( Kingma & Welling , 2014c ) , the prior distribution of p ( z ) is assumed to be a standard Gaussian distribution . In SBVAE ( Nalisnick & Smyth , 2017 ) , the prior distribution becomes a GEM distribution that produces samples with a Beta distribution and a stick-breaking algorithm . 2.2 DIRICHLET DISTRIBUTION AS A COMPOSITION OF GAMMA RANDOM VARIABLES . The Dirichlet distribution is a composition of multiple Gamma random variables . Note that the probability density functions ( PDFs ) of Dirichlet and Gamma distributions are as follows : Dirichlet ( x ; α ) = Γ ( ∑ αk ) ∏ Γ ( αk ) ∏ xαk−1k , Gamma ( x ; α , β ) = βα Γ ( α ) xα−1e−βx ( 2 ) where αk , α , β > 0 . In detail , if there are K independent random variables following the Gamma distributions Xk ∼ Gamma ( αk , β ) or X ∼ MultiGamma ( α , β · 1K ) where αk , β > 0 for k = 1 , · · · , K , then we have Y ∼ Dirichlet ( α ) where Yk = Xk/ ∑ Xi . It should be noted that the rate parameter , β , should be the same for every Gamma distribution in the composition . Then , the KL divergence can be derived as the following : KL ( Q||P ) = ∑ log Γ ( αk ) − ∑ log Γ ( α̂k ) + ∑ ( α̂k − αk ) ψ ( α̂k ) ( 3 ) for P = MultiGamma ( α , β ·1K ) andQ = MultiGamma ( α̂ , β ·1K ) where ψ is a digamma function . The detailed derivation is provided in Appendix B . 2.3 SGVB FOR GAMMA RANDOM VARIABLE AND APPROXIMATION ON DIRICHLET DISTRIBUTION . This section discusses several ways of approximating the Dirichlet random variable ; or the SGVB estimators for the Gamma random variables which compose a Dirichlet distribution . Utilizing SGVB requires a differentiable non-centered parametrization ( DNCP ) for the distribution ( Kingma & Welling , 2014d ) . The main SGVB for Gamma random variables , used in DirVAE , is using the inverse Gamma CDF approximation explained in the next section . Prior works include two approaches : the use of the Weibull distribution and the softmax Gaussian distribution , and the two approaches are explained in this section . Approximation with Weibull distribution . Because of the similar PDFs between the Weibull distribution and the Gamma distribution , some prior works used the Weibull distribution as a posterior distribution of the prior Gamma distribution ( Zhang et al. , 2018 ) : Weibull ( x ; k , λ ) = k λ ( x λ ) k−1 e− ( x/λ ) k where k , λ > 0 . ( 4 ) The paper Zhang et al . ( 2018 ) pointed out that there are two useful characteristics when approximating the Gamma distribution with the Weibull distribution . One useful property is that the KL divergence expressed in a closed form , and the other is the simple reparametrization trick with a closed form of the inverse CDF from the Weibull distribution . However , we noticed that the Weibull distribution has a component of e− ( x/λ ) k , and the Gamma distribution does not have the additional power term of k in the component . Since k is placed in the exponential component , small changes on k can cause a significant difference that limits the optimization . Approximation with softmax Gaussian distribution . As in MacKay ( 1998 ) ; Srivastava & Sutton ( 2017 ) , a Dirichlet distribution can be approximated by a softmax Gaussian distribution by using a softmax Laplace approximation . The relation between the Dirichlet parameter α and the Gaussian parameters µ , Σ is explained as the following : µk = logαk − 1 K ∑ i logαi , Σk = 1 αk ( 1− 2 K ) + 1 K2 ∑ i 1 αi , ( 5 ) where Σ is assumed to be a diagonal matrix , and we use the reparametrization trick in the usual GVAE for the SGVB estimator .

Summary: This paper proposes to change the typical Gaussian posterior distribution (and prior) for the latent features z associated to an image x that is used in Variational Autoencoders by a Dirichlet distribution. The work improves over previous attempts based on a soft-max + Gaussian distribution and the soft-max + Weibull distribution. The trick proposed to make feasible training the model includes approximating the inverse CDF of the gamma distribution and using the fact that the Dirichlet distribution can also be obtained as a normalized sum of gamma random variables. The method is compared in several problems. Some analysis of the reasons why it performs better is also carried out.