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

Abstract:
Variational Autoencoders (VAEs) are powerful probabilistic models to learn representations of complex data distributions. One important limitation of VAEs is the strong prior assumption that latent representations learned by the model follow a simple uni-modal Gaussian distribution. Further, the variational training procedure poses considerable practical challenges. Recently proposed regularized autoencoders offer a deterministic autoencoding framework, that simplifies the original VAE objective and is significantly easier to train. Since these models only provide weak control over the learned latent distribution, they require an ex-post density estimation step to generate samples comparable to those of VAEs. In this paper, we propose a simple and end-to-end trainable deterministic autoencoding framework, that efficiently shapes the latent space of the model during training and utilizes the capacity of expressive multi-modal latent distributions. The proposed training procedure provides direct evidence if the latent distribution adequately captures complex aspects of the encoded data. We show in experiments the expressiveness and sample quality of our model in various challenging continuous and discrete domains. An implementation is available at https://github.com/boschresearch/GMM_DAE.

heading: 1 Introduction
text: 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.

heading: 2 Related Work
text: 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.

heading: 3 Method
text: 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.

heading: 3.1 Uni-Modal latent regularization
text: 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.

heading: 3.2 Multi-Modal latent regularization
text: One advantage of our approach is the applicability to more expressive, multi-modal prior distributions. While the Gaussian distribution has important analytical properties, it suffers from significant limitations when modelling real data sets. In contrast, a linear combination of Gaussians can give rise to very complex densities while still allowing for closed form computations of important quantities, like CDFs and covariances. A d-dimensional K-modal Gaussian mixture model is a weighted superposition of K Gaussian distributions in Rd, that are often referred to as the modes of the model. For k ≤ K, let µk and Σk be the mean and covariance matrix of the k-th mode in the model. Further, let pk > 0 be the weight of the k-th mode. Then, the marginal CDFs of a GMM model can be computed
w ei
gh te
d la
te nt
re gu
la ri
za tio n K S di st an ce
from the CDFs of univariate Gaussians as follows
FGMM,j(z) = K∑ k=1 pkΦ ( z − [µk]j [Σ]j,j ) , (5)
i.e. the marginal CDFs in the GMM are weighted sums of CDFs of one-dimensional Gaussians. The covariance matrix of the GMM can be computed as
ΣGMM = K∑ k=1 pkΣk + K∑ k=1 pk (µk − µ̄) (µk − µ̄) T , µ̄ = 1 k K∑ k=1 µk. (6)
Extending our proposed regularization scheme to multimodal GMMs is straight forward. Our first loss term is defined as
LKS,K(z1,...,N ) = 1
d d∑ j=1 MSE ( F̄ (N) j (zj), FGMM,j(zj) ) . (7)
Similarly, the second loss term is defined to be
LCV,K(z1,...,N ) = 1
d2 d∑ l,j=1 ( [Σ̄]l,j − [ΣGMM]l,j )2 . (8)
The total loss of the model is a combination of the reconstruction loss and a regularization loss, that enforces the latent representations of the encoded data to match a predefined multi-modal prior distribution. The reconstruction loss LREC(x′1,...,N ) equals the mean squared error between inputs xi and their reconstructions x′i. Given positive weights λKS and λCV, our final loss is given by
L(x1,...,N ) = λRECLREC(x′1,...,N ) + λKSLKS,K(z1,...,N ) + λCVLCV,K(z1,...,N ). (9)
Formally, the weights λKS, λCV and λREC are hyperparameters of the model. Nevertheless, we propose an explicit way to set λKS and λCV and a simple heuristic to estimate λREC to avoid an extensive optimization of these weights.

heading: 3.3 Loss weight estimation
text: Balancing the two regularization losses appropriately poses a key challenge as they potentially vary on very different scales. For example, if modes of the GMM prior are far spread, the covariance LCV,K loss will dominate the marginal CDF LKS,K loss by far. Nevertheless, given a target GMM prior, the
dimension of the latent space and the batch size n used during training, there is a concise way to fix those hyperparameters beforehand. To be precise, for m = 1, . . . ,M samples z(m)1 , . . . , z (m) N from the prior GMM, we propose to set
λ−1KS = 1
M LKS
( z (m) 1,...,N ) , λ−1CV = 1
M LCV
( z (m) 1,...,N ) . (10)
Formally, we can not overcome the necessity of tuning the weight of the reconstruction loss, which has significant impact on performance of the model. Nevertheless, a reasonable approximation to it can be obtained by training an autoencoder model and using the inverse of the best obtained loss for λREC. Using this scaling, all loss terms in our regularization loss will ultimately converge to one if the target prior is matched successfully.

heading: 4 Experiments
text: With our experiments we strive to investigate the potential of the proposed model when compared to other VAE variants in generating new samples, analyse the effect of the defined prior to effectively cluster the latent space and to shape the latent space efficiently in highly structured domains such as discrete spaces. We provide all the experimental settings and hyperparameters used in the Appendix. All experiments were run on a GPU cluster, with single GPU per individual experiments. Since the cluster is part of a carbon-neutral framework, these experiments did not contribute to climate change.

heading: 4.1 Image generation
text: We consider four dataset, MNIST [26], FASHIONMNIST [45], SVHN [36] and CELEBA [28] to evaluate the proposed method in image generation experiments. The qualitative analysis of the generated samples for MNIST, SVHN amd CELEBA images are shown in Figure 2 along with the reconstructed samples and interpolated samples in the latent space of the trained model. In order to assess the quality of the generated images, we evaluate the Fréchet Inception Distance (FID) [19] for each dataset, see Table 1. For baseline comparison, we evaluate the following models: vanilla variational autoencoder (VAE [24]), Gaussian mixture variational autoencoder (GMVAE) [6], Wasserstein autoencoder (WAE) [43] with MMD loss, 2stage VAEs (2s-VAE) [5], constant variance-VAE (CV-VAE) [12] and regularized autoencoders (RAEs) [12]. We consider the following evaluation metrics: 1. Sampling FID (Samp.) - FID score of the generated random samples (evaluated by generating random samples from the prior distribution of the respective models and by fitting a Gaussian distribution to models trained without any prior assumptions), 2. reconstruction FID (Rec.) - measured by computing the FID between the test samples and their corresponding reconstructions by the model and 3. interpolation FID (Inter.) - measured by computing the FID between the interpolated samples in the latent space and test samples. As pointed out by [12], fitting an ex-post density estimator on the learned embedding after training of VAEs further improves the generation quality. Hence, we also report the FID values by fitting a GMM in the learned latent space of the trained model (GMM column in Table 1, not evaluated for 2s-VAE as they perform ex-post density estimation using another VAE).
As shown in Table 1, our method achieves better FIDs (Samp.) on all datasets considered, when compared to all considered baselines sampled by fitting a single Gaussian in the latent space. We also improved the generation quality as argued above by fitting a mixture of Gaussians in the latent space and achieve better FIDs in MNIST, FASHION MNIST and CELEBA images, whereas for SVHN, WAEs achieved the overall best score. It is also important to note that the proposed method performs comparably or even better without employing the ex-post density estimation. The proposed method also achieves better reconstruction quality than the other VAEs except for SVHN images where RAEs performs better. The interpolation FID indicates the overall structure of the learned latent space and the obtained FID values show that the proposed method shapes the latent space better than the other approaches except for the CELEBA images where RAEs performs slightly better than our method. For a fair comparison, we use the same architecture and experimental settings in all the considered baseline evaluations. Please refer to the Appendix for more details on the experimental settings.

heading: 4.2 Unsupervised image clustering
text: We evaluate the potential of our method to naturally cluster the data points in the learned latent space in two dataset, MNIST and FASHION-MNIST. The Gaussian mixture model prior with k
components in our method could be considered as k different classes/clusters to which the data points are mapped by the encoder. We train the model with latent space dimension 10 for both dataset and visualize the random samples generated from each Gaussian component of our prior as shown in Figure 3. The figure shows that visually similar images fall into the same cluster. For a quantitative analysis of the clustering performance, we evaluated the unsupervised classification accuracy (similar to [21]) and compare the performance with JointVAE [9] and CascadeVAE [21]. The observed values are reported in the table in Figure 4. We observed a comparable performance to both baselines. We also observed that the distance between the modes in the GMM prior is a deciding factor in better clustering performance. Figure 4 (right) shows the performance comparison of both image generation
and clustering performance with increasing distance between different modes in the GMM prior. The result shows that with increasing distance, the clustering performance is improved wheres the quality of the generated images gets reduced. Our experimental analysis indicates that natural clustering happens with the multi-modal GMM prior in our method.

heading: 4.3 Modelling discrete data structures
text: In this section, we investigate the ability of our model to generate complex discrete data structures such as arithmetic expressions and molecules. The objective of this experiments is to analyze the model performance on shaping the latent space of such structured discrete spaces effectively. The learned latent space of the model is traversed to generate new samples with the desired properties by performing Bayesian Optimization (BO). We perform experiments in two sequence optimization problems similar to [25].
Arithmetic Expression Given a dataset of 50, 000 univariate (functions of x) arithmetic expressions following a formal grammar [25], the task is to find the expression that best fits a target dataset. This is done by minimizing log(1 + MSE), where the MSE is computed between values of the generated expression and the target points. For our evaluation, we choose similar target data points as in [25].
Chemical Design Given the ZINC250k dataset of drug molechules [20], the objective is to generate new drug like molecules. The drug likeliness of a molecule is quantified by the water-octanol partition coefficient, which is maximized in our line of experiments.
Results We extend the architecture and experimental settings of [25] to include our proposed losses during training. For baseline comparison, we consider Grammar VAE (GVAE) [25], Character VAE (CVAE) [17], Grammar constant variance VAE (GCVVAE) [12] and Grammar based RAE (GRAE) [12] frameworks. The three best scores found by our method for arithmetic expressions and the molecule experiments are reported in Table 2. Our model performs comparatively better than the considered baselines and achieves the best first score for both tasks. In addition to the optimization performance, it is also important to consider the validity of the new samples generated by the models. A well-structured latent space should yield valid samples following the defined grammar/rules of the used dataset. Our model achieves better validation and average scores as shown in Table 3 except for GCVVAE which achieves a better average score in the arithmetic expression task. All reported values are evaluated by averaging across 5 BO trials.

heading: 4.4 Ablation study and hyperparameter sensitivity analysis
text: We perform an ablation study on the regularization loss terms in the proposed model. When the model is trained without the KS distance loss for MNIST images, we observed an FID of 49.82 and when trained without the covariance matching loss, we observed an FID of 38.45. These values are significantly worse than the FID that we achieve when training with the weighted combination of both regularization losses i.e 13.11. These empirical evaluations show that the combination of the two regularization terms facilitates a better prior-posterior match and hence better image generation. Please refer to the Appendix for qualitative evaluation of the ablation study on MNIST images.
From a conceptual point of view, the most important hyperparameters of our model are (a) the weights of the different terms in the training objective and (b) the number of components in the prior. In section 3.4, we propose an explicit way how to fix the weights in the loss function. We investigate the sensitivity of our model performance to the number of components in the GMM prior. We trained our model on the MNIST dataset using a GMM prior with 1, 5, 10, 15, 20 and 25 modes respectively. The observed FID scores for the respective number of components are shown in Figure 5. The result show that with increasing number of components in the chosen prior, the performance of our model improves significantly. As a consequence, choosing a large number of components can be beneficial for practical considerations.

heading: 5 Limitations and Future Work
text: One limitation of our work is the necessity to chose the prior distribution in advance. We showed that fixing a suitable number of modes for the GMM is important to provide better sampling quality. Also, by considering marginal CDFs, we simplified the original distance metric from the KS test. While reducing computational complexity during training, this comes at the cost of an additional loss term. Further, our proposed addition of the KS distance is not suitable for matching higher order moments of the latent representations to the target prior, which at least from a conceptual point of view can lead to a mismatch to the prior. Further, our loss only facilitates matching empirical marginal CDFs of latent representations to the marginal CDFs of the prior evaluated at latent vectors. As a consequence, our regularization loss might be a less stable training signal for small batch sizes in high dimensions.
Using the Frobenius norm in the covariance matching loss reflects our assumption that the latent dimensions should all be independent from each other and should all simultaneously match the prior’s values. The choice for the MSE for covariance matching loss is purely based on its prevalence in the literature. Additionally, this makes all three loss terms (reconstruction loss, KS distance loss and the covariance matching loss) behave similarly, as they are all squares. While we did not investigate any other metrics for matrix comparison in this scenario, exploring other options for the covariance matching is an interesting area for future studies. We have not considered the case where there exists class imbalance in the dataset. We would expect the model to separate the classes if the imbalance is weak and the classes are sufficiently different such that the reconstruction loss outweighs the regularization penalty for the mismatch. Extending our prior to accommodate for this by introducing a weighted GMM prior is also a very interesting direction for future work.

heading: 6 Conclusion
text: Recent studies have illustrated the effectiveness of flexible priors in VAEs to learn more meaningful latent representations. Following recent work that highlights the potential of deterministic alternatives to the variational formulation in VAEs, we propose a simple deterministic autoencoding framework with more powerful regularizers to accommodate for expressive multi-modal priors. In particular, we derive a novel deterministic regularization scheme from a strong metric on probability distributions. The proposed approach can be readily applied to effectively shape the latent space of existing autoencoding frameworks towards multi-modal Gaussian priors. Our experimental evaluations show that the proposed training objective yields comparable sampling quality to those of variational autoencoders and achieves better performance in modelling complex discrete data structures.

Review summaries:
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. 

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.

The author(s) motivate and derive a novel deterministic auto-encoder loss that naturally encourages a multi-model latent representation that follows a GMM. The author(s) motivate and derive this loss using the KS test metric and generalize it to multiple dimensions and multiple modes. Their experiments, whereby they ancestrally resample latent codes show an impressive ability to generate realistic samples.

The paper builds on the ideas of Ghosh et al 2020 [10] pointing out the equivalence between Gaussian prior VAEs and deterministic AEs with noise inputation and further equivalence with regularized deterministic AEs. Since the original paper [10] needed to estimate the latent distribution ex post with a GMM, the authors here propose to use GMM prior directly and translate this into a deterministic AE with more complex regularization derived from univariate Kolmogorov-Smirnof distance combined with another term for covariance matching.They further propose a heuristic to find a reasonable heyper-parameter setting for the two regularizers. Finally, they show on a battery of experiments (generations over standard image datasets, unsupervised clustering, generating discrete structures) the favourable properties of their method.