Learning Disentangled Discrete Representations
This work addresses the challenge of unsupervised disentanglement in machine learning, offering a novel method and model selection strategy, though it is incremental in building on existing VAE frameworks.
The paper tackles the problem of learning disentangled representations by replacing Gaussian VAEs with categorical VAEs, showing that discrete latent spaces mitigate rotational invariance and act as an inductive prior for disentanglement, with empirical and analytical evidence supporting this approach.
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.