Learning Disentangled Discrete Representations
David Friede Christian Reimers Heiner Stuckenschmidt Mathias Niepert
Recent successes in image generation, model-based reinforcement learning, and
text-to-image generation have demonstrated the empirical advantages of discrete
latent representations, although the reasons behind their benefits remain
unclear. We explore the relationship between discrete latent spaces and
disentangled representations by replacing the standard Gaussian variational
autoencoder (VAE) with a tailored categorical variational autoencoder. We show
that the underlying grid structure of categorical distributions mitigates the
problem of rotational invariance associated with multivariate Gaussian
distributions, acting as an efficient inductive prior for disentangled
representations. We provide both analytical and empirical findings that
demonstrate the advantages of discrete VAEs for learning disentangled
representations. Furthermore, we introduce the first unsupervised model
selection strategy that favors disentangled representations.