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: The paper proposes a deterministic autoencoder that imposes a GMM on the latent space to learn a complex multimodal latent space. The model uses two regularizers including the KS distance and a constant covariance term to minimize the discrepancy between the marginal CDF of the empirical posterior and the given GMM prior. The authors provide multiple experimental results to show the reconstruction and sampling quality, deep clustering, and the latent model capability to generate discrete and complex structures such as chemical molecules.