Shape your Space: A Gaussian Mixture Regularization Approach to Deterministic Autoencoders

1 Introduction . Variational autoencoders ( VAEs ) constitute one of the popular generative learning frameworks widely used for applications such as image understanding and generation , sentence modeling , and optimizing discrete data and graph-based structures [ 7 , 23 , 34 , 40 , 48 ] . The VAE framework elegantly combines autoencoders with variational inference [ 24 ] . The encoder of the model maps the input data into a lower-dimensional latent space according to a given inference model . The decoder provides a mapping from the latent space back to the original input space . Both are jointly optimized by maximizing a lower bound on the model evidence , regularizing the latent space towards a fixed prior distribution , usually a uni-modal Gaussian . By sampling from the latent space prior , we can utilize the decoder network to efficiently generate new samples from the training distribution . Due to the variational formulation , optimizing the VAE training objective poses significant practical challenges . Further , the over simplistic prior assumption often leads to an unsatisfying trade-off between the quality of reconstructed samples and the prior regularization [ 2 ] . Recent work has shown that choosing more flexible priors helps to improve the generative performance of VAEs [ 44 ] . Since the initial introduction of VAEs , various novel training objectives have been proposed . One line of work focuses on different regularization techniques derived from alternative probabilistic metrics to shape the latent space of the model during training , e.g . using the Wasserstein distance [ 43 ] . In contrast to the KL-divergence , the Wasserstein distance measure induces a metric on probability 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . distributions . Practically , this facilitates smoother convergence even for initially non-overlapping distributions . Further , it overcomes the over-regularization effect in VAEs . To be precise , it prevents the undesired behaviour of multiple data points being mapped to the same latent representation by the encoder . Since closed-form solutions for metrics like the Wasserstein distance can only be derived for very few prior distributions , these approaches rely on numerical approximations during training . Recent work by Ghosh et al . [ 12 ] reinterprets deterministic autoencoders as variational models , even when trained with a deterministic loss . During training , this approach maximizes the negative log-marginal likelihood of the latent samples under a Gaussian normal distribution as a regularization in addition to minimizing the reconstruction loss . Experimental results show that this regularization alone does not suffice to generate high quality samples using the Gaussian prior . To overcome this , Gosh et al . propose to use a multi-modal Gaussian mixture model ( GMM ) to fit arbitrary , learned latent spaces . While this approach leads to good sampling efficiency and generalization if the post-hoc fit is reasonable , sampling quality can suffer significantly if the learned latent space can not be modeled well by a GMM . In this work , we propose a deterministic training scheme for autoencoders that is applicable to expressive priors and overcomes the necessity of a post-hoc density estimation step for deterministic training . To be precise , we derive a deterministic regularization loss from the distance metric used in the non-parametric Kolmogorov-Smirnov ( KS ) test for equality of probability distributions . The resulting training objective can be derived in closed form for a class of expressive multi-modal prior distributions and provides a strong signal to efficiently shape the latent space of the model during training . We chose our experiments to evaluate the proposed approach in terms of sampling quality and expressiveness . In the first line of experiments , we compare the quality of newly generated and reconstructed samples from our model with those from a variety of other VAE variants . In the second line , we investigate our method ’ s capability to model discrete and complex structured inputs such as arithmetic expressions and molecules . In these domains , VAEs have recently been proposed as a tool for dimensionality reduction in optimization . Applying our regularization scheme effectively utilizes multi-modal prior distributions in this context and significantly improves optimization performance . 2 Related Work . Since the introduction of VAEs , many follow up works tried to overcome the practical and theoretical limitations of the framework , e.g . [ 2 , 43 , 44 ] , and make them applicable to specific applications such as clustering [ 6 , 39 ] or anomaly detection [ 49 ] . We first review some seminal examples of VAE models with different priors and probability metrics for latent regularization . Since our proposed regularization term structures the latent space to a Gaussian mixture model , we also compare it to prior work on deep clustering . Lastly , we discuss VAEs in the context of black-box optimization approaches such as Bayesian Optimization ( BO ) . VAEs In the standard VAE framework , the prior distribution is commonly assumed to be a Gaussian normal distribution . This might lead to simplified representations learned by the model which is unable to represent the rich semantics in the data distribution . Several methods were proposed to introduce more flexible and expressive priors to the VAE formulation . Casale et al . [ 3 ] employ Gaussian process priors to account for correlations between the data samples . In [ 15 ] , a Bayesian non-parametric prior is used with a hierarchical non-parametric variational autoencoder for video representation learning . Chen et al . [ 4 ] use an auto-regressive prior to achieve improved generative performance on image datasets . Berger et al . [ 2 ] propose to replace the standard spherical Gaussian prior with a more general version with an arbitrary covariance matrix and learn the correlations by optimizing the evidence lower bound of the model . Although the proposed methods offer competitive performance , they often employ complex architectures [ 4 ] to achieve desired performance . In another line of work , multi-modal priors were utilized in VAE models . Zong et al . [ 49 ] propose to use a GMM prior in autoendocers for unsupervised anomaly detection by training an additional network estimating the parameters of the GMM . Lee et al . [ 27 ] address unsupervised meta-learning using a GMM prior in VAEs to shape the latent space by employing an extension of the evidence lower bound to complex variational inference schemes . Tomczak et al . [ 44 ] propose to replace the GMM prior by a coupling of the posterior and prior of the model . Adversarial autoencoders [ 31 ] improve the generative performance of VAEs by incorporating adversarial learning into the VAE framework and offer competitive performance in image generation at an increased computational complexity and decreased training stability . To account for the over regularization effect of the KL divergence term in the standard VAE framework , [ 43 ] minimize the Wasserstein distance between the representations learned by the model and the target prior . Recently introduced regularized autoencoders [ 12 ] ( RAEs ) , question the variational framework adopted by the VAEs and propose a deterministic approach to achieve comparable or better image generation performance . The authors use the negative log likelihood for regularization , but require a post-hoc step to derive a strong sampling procedure from the model . The state of the art VAE model for high fidelity image generation , VQ-VAE [ 37 , 41 ] , can be also considered as a deterministic autoencoder . Similar to RAEs , training VQ-VAE involves two stages of training relying on complex discrete autoregressive density estimators . Moreover , the training loss of VQ-VAE is non-differentiable due to the quantization of the latent vector . Our approach elegantly combines the idea of new training objectives with the extension to multimodal priors without increasing training complexity or compromising sampling quality . We derive a strong training signal which can be derived in closed form for multimodal priors . This ensures stable training and reliable regularization of the latent space , improving sampling quality . Deep Clustering Deep Clustering approaches benefit from well structured latent spaces . Thus , several methods employ Gaussian mixture VAEs for data encoding [ 6 , 39 ] or establish a GMM-like latent space structure through k-means models in the latent space . For example , Xie et al . [ 46 ] train an autoencoder and apply a KL-divergence loss for better k-means clustering while Ghasedi et al . [ 8 ] combine the autoencoder reconstruction loss with the relative cluster entropy . Similar approaches have been proposed in the literature [ 8 , 16 , 18 , 22 , 42 , 47 ] . Caron et al . [ 33 ] iteratively group points using k-means during optimization . In the context of clustering , generative adversarial networks have been considered in [ 11 , 35 ] . While we are not considering the clustering task in this paper , we hypothesise that the proposed regularization can be beneficial in this context since it implicitly optimizes for mode assignments . Structural VAEs and optimization High-dimensional optimization problems in structured discrete input domains are ubiquitous . VAEs have been used in this context to learn low dimensional , continuous representations of high dimensional , structured data like molecules or arithmetic expressions . Recent work proposes to use such representations to perform efficient optimization by running BO in the latent space of VAEs [ 25 , 30 ] . In this setting , prior knowledge of the structure of the latent space is crucial to allow for an efficient exploration and generation of valid samples . Yet , as discussed above , VAEs can suffer from simplistic prior assumptions . Thus , sampling from the latent space of such models can result in invalid samples , reducing the sampling efficiency of BO [ 17 ] . Kusner et al . [ 25 ] overcome this issue if data follows a specific grammar . Lu et al . [ 30 ] propose a VAE that directly works on parse trees from context-free grammars to represent discrete data . Yet , those only work with unimodal priors which limits the generalization capabilities . Our approach can be readily used to extend these models to better encode structural data and improve BO performance . 3 Method . We introduce a novel loss function to regularize the latent representation learned by deterministic autoencoders towards a given prior distribution . The definition of our loss builds on the non-parametric statistical Kolmogorov-Smirnov ( KS ) test for equality of one-dimensional probability distributions . We propose a multivariate variant of the distance measure used in the KS test , that allows for gradient based optimization and can easily be applied to expressive multi-modal prior distributions . For ease of exposition , we start with introducing our regularization loss for unimodal Gaussian priors in section 3.1 and extend the formulation to expressive multi-modal Gaussian mixture models in Section 3.2 . Finally , in Section 3.3 , we provide an explicit way to estimate the weighting parameters of our loss . 3.1 Uni-Modal latent regularization . The KS test can be used to determine whether a collection of N , one-dimensional samples follow a given reference distribution . It compares the cumulative distribution function ( CDF ) of the reference distribution with the empirical CDF F̄ ( N ) of the samples . It is often applied to the class of one-dimensional Gaussian distributions , which has important analytical properties . For spherical Gaussians , the one-dimensional KS test quantifies a distance between the empirical distribution function of the data and the cumulative distribution function Φ ( z ) = 1 σ √ 2π ∫ x −∞ exp − ( t− µ ) 2 2σ2 dt ( 1 ) of the univariate Gaussian Z ∼ N ( µ , σ ) as supz∈R |F̄ ( N ) ( z ) − Φ ( z ) | . Extending this KS distance to higher dimension is particularly challenging , since it requires matching joint CDFs [ 10 , 14 , 38 ] . Especially in higher dimensions this becomes infeasible [ 29 ] . The continuous ranked probability score [ 13 ] shares the same theoretical basis as the KS distance . However it tests whether two sets of samples are consistent with each other , i.e. , they could originate from the same distribution , and is thus not suitable to regularize a collection of latent samples towards a given prior distribution . Alternative multi-variate normality tests , like the Mardia test [ 32 ] and the BEHP test [ 1 ] suffer from slow convergence rates . To derive a regularization loss from the KS distance , we propose to overcome this issue by taking into consideration the marginal CDFs and correlations in the prior distribution separately . Given d-dimensional latent samples z1 , . . . , zN , the empirical marginal CDF in dimension j is given by F̄ ( N ) j ( z ) = 1 n N∑ n=1 1 [ zn ] j≤z . ( 2 ) We aim to regularize the latent space of our models by comparing the empirical marginal CDFs with the one-dimensional CDFs of the marginal distributions of the prior . To strengthen the training signal of our regularization scheme and make it suitable for gradient-based optimization , we replace the supremum in the original KS distance by a smoother MSE loss , that compares the distances between those functions at the latent representations . For a uni-modal Gaussian prior with mean µ and covariance matrix Σ , this results in LKS ( z1 , ... , N ) = 1 d d∑ j=1 MSE ( F̄ ( N ) j ( zj ) , Φ ( z̄j ) ) , z̄j = zj − µj [ Σ ] j , j . ( 3 ) Here , F̄ ( N ) j ( zj ) denotes the vector with entries F̄ ( N ) j ( [ zi ] j ) and Φ ( z̄j ) is defined accordingly . This loss is minimized , if the empirical marginal CDFs of the latent samples match those of the uni-modal Gaussian prior . Using the above loss alone will not account for correlations between different latent dimensions . In the case of a spherical Gaussian prior with identity covariance matrix for example , samples with perfectly correlated Gaussian components [ zi ] j = [ zi′ ] j , will also minimize this objective , see Figure 1 . To overcome this problem , we equip our loss with an additional term , that matches covariances between different latent distributions explicitly . Following a similar reasoning to the MSE above , we define an additional loss term , LCV ( z1 , ... , N ) = 1 d2 d∑ l , j=1 ( [ Σ̄ ] l , j − [ Σ ] l , j ) 2 , ( 4 ) where Σ̄ is the empirical covariance matrix of the latent representations and Σ stands for the prior covariance . Compared to the negative log marginal regularization proposed in [ 12 ] , our loss will actually enforce the latent representations to be spread across the entire support of the Gaussian prior , instead of being minimal when all latent collapse to the origin .

Summary: This paper proposes to learn GMM priors in autoencoders, to be able to construct a generative model. They compare the generation quality of the proposed approach with VAEs using the FID metric. One interesting contribution seems to be the fact that they are using a loss function other than the GMM likelihood to fit the GMM.