$\alpha$-TCVAE: On the relationship between Disentanglement and Diversity

TL;DR

This paper introduces $oldsymbol{ extalpha}$-TCVAE, a variational autoencoder that maximizes disentanglement and diversity of generated data through a novel total correlation bound, demonstrating improved representations and generative performance on complex datasets and downstream tasks.

Contribution

The paper proposes $oldsymbol{ extalpha}$-TCVAE with a new TC lower bound grounded in information theory, enhancing disentanglement and diversity in generative models, especially for complex datasets.

Findings

$oldsymbol{ extalpha}$-TCVAE learns more disentangled representations than baselines.

It generates more diverse observations without losing visual fidelity.

It improves downstream performance in RL tasks, notably in the Ant Maze environment.

Abstract

While disentangled representations have shown promise in generative modeling and representation learning, their downstream usefulness remains debated. Recent studies re-defined disentanglement through a formal connection to symmetries, emphasizing the ability to reduce latent domains and consequently enhance generative capabilities. However, from an information theory viewpoint, assigning a complex attribute to a specific latent variable may be infeasible, limiting the applicability of disentangled representations to simple datasets. In this work, we introduce $\alpha$-TCVAE, a variational autoencoder optimized using a novel total correlation (TC) lower bound that maximizes disentanglement and latent variables informativeness. The proposed TC bound is grounded in information theory constructs, generalizes the $\beta$-VAE lower bound, and can be reduced to a convex combination of the known variational information bottleneck (VIB) and conditional entropy bottleneck (CEB) terms. Moreover, we present quantitative analyses that support the idea that disentangled representations lead to better generative capabilities and diversity. Additionally, we perform downstream task experiments from both representation and RL domains to assess our questions from a broader ML perspective. Our results demonstrate that $\alpha$-TCVAE consistently learns more disentangled representations than baselines and generates more diverse observations without sacrificing visual fidelity. Notably, $\alpha$-TCVAE exhibits marked improvements on MPI3D-Real, the most realistic disentangled dataset in our study, confirming its ability to represent complex datasets when maximizing the informativeness of individual variables. Finally, testing the proposed model off-the-shelf on a state-of-the-art model-based RL agent, Director, significantly shows $\alpha$-TCVAE downstream usefulness on the loconav Ant Maze task.