Uniform Transformation: Refining Latent Representation in Variational Autoencoders

TL;DR

The paper tackles irregular latent posteriors in VAEs that hinder disentanglement and reliable sampling by introducing a three-stage Uniform Transformation (UT) module. The UT pipeline combines Gaussian Kernel Density Estimation (G-KDE) clustering to identify child Gaussians, Gaussian Mixture (GM) modeling to represent the posterior, and Probability Integral Transform (PIT) to map the GM to a uniform latent distribution, thereby reducing posterior collapse and improving interpretability. Empirical results on dSprites and MNIST show enhanced disentanglement metrics for several baseline VAEs, with the UT approach yielding sharper reconstructions and better capture of periodic features, while remaining computationally feasible. This framework offers a modular path to more structured latent spaces, potentially extending to richer datasets and downstream tasks where disentanglement and robust sampling are critical.

Abstract

Irregular distribution in latent space causes posterior collapse, misalignment between posterior and prior, and ill-sampling problem in Variational Autoencoders (VAEs). In this paper, we introduce a novel adaptable three-stage Uniform Transformation (UT) module -- Gaussian Kernel Density Estimation (G-KDE) clustering, non-parametric Gaussian Mixture (GM) Modeling, and Probability Integral Transform (PIT) -- to address irregular latent distributions. By reconfiguring irregular distributions into a uniform distribution in the latent space, our approach significantly enhances the disentanglement and interpretability of latent representations, overcoming the limitation of traditional VAE models in capturing complex data structures. Empirical evaluations demonstrated the efficacy of our proposed UT module in improving disentanglement metrics across benchmark datasets -- dSprites and MNIST. Our findings suggest a promising direction for advancing representation learning techniques, with implication for future research in extending this framework to more sophisticated datasets and downstream tasks.

Figures (7)

• Figure 1: Learned Latent Variable Histograms with Irregular Distributions from Replicating the $\beta$-VAE of dSprites Dataset higginsBetaVAELearningBasic2016a. ${\textnormal{z}}_7$ encodes the shape and scale labels together. ${\textnormal{z}}_6$ and ${\textnormal{z}}_8$ encode the rotation label. The encoding correspondence is analyzed by correlation in Appendix \ref{['apx:betavae']}.
• Figure 2: VAE Structure. The reparameterization trick separates the encoder output into mean $\mu_{{\textnormal{z}}}$ and variance $\sigma_{{\textnormal{z}}}^2$, which are then combined with an independent random variable $\xi$ to construct the latent variable ${\textnormal{z}}$.
• Figure 3: Illustration of GM Reasoning: The yellow dashed curve represents the PDF of a normal distribution, softly constrained by Eq. \ref{['eqn:kl']}. The red dashed curves denote the PDFs of child Gaussian distributions, constrained by Eq. \ref{['eqn:reparameter']}. The blue shadow illustrates the GM distribution of the posterior $q({\textnormal{z}})$, where $w^{(k)}$, $\mu^{(k)}$, and $(\sigma^{(k)})^2$ denote the weight, mean, and variance of the $k$th child Gaussian, respectively.
• Figure 4: Schematic of the Uniform Transformation Module integrated into a VAE framework. The module, highlighted within the gray box, transforms latent variables ${\mathbf{z}}$ to $\tilde{{\mathbf{z}}}$, representing the pre- and post-transformation states, respectively.
• Figure 5: Comparative Generative Results on dSprites Dataset using VAE Models with a Fixed Random Seed.