Uniform Transformation: Refining Latent Representation in Variational Autoencoders

TL;DR

This paper introduces a three-stage Uniform Transformation module for VAEs that reconfigures irregular latent distributions into uniform ones, significantly improving disentanglement and interpretability of latent representations.

Contribution

The paper proposes a novel adaptable three-stage UT module combining G-KDE clustering, GM modeling, and PIT to address irregular latent distributions in VAEs, enhancing their performance.

Findings

Improved disentanglement metrics on dSprites and MNIST datasets.

Effective reconfiguration of latent space into uniform distribution.

Enhanced interpretability of latent representations.

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.