Google AI‘s ICML 2019 Best Paper Puts Unsupervised Disentanglement to the Test
Introduction
Unsupervised learning, the ability to extract knowledge from raw unlabeled data, is a holy grail of artificial intelligence research. Unsupervised learning holds the promise of AI systems that can acquire complex world models and reasoning capabilities without expensive human-annotated datasets. Within unsupervised learning, a key challenge is disentangled representation learning – finding representations that separate out meaningful factors of variation in the data.
At the 2019 International Conference on Machine Learning (ICML), one of the world‘s premier AI research gatherings, Google AI made a splash with their best paper award for "Challenging Common Assumptions in the Unsupervised Learning of Disentangled Representations" [1]. The paper takes a hard look at state-of-the-art unsupervised disentanglement approaches and uncovers some concerning gaps between theory and practice. Let‘s dive in and unpack what they found and what it means for the future of unsupervised learning.
What is Disentanglement?
Disentangled representation learning aims to find a representation of data where the underlying factors of variation are separated into independent components [2]. For example, in a dataset of face images, a disentangled representation might have separate dimensions that encode head pose, lighting, facial expression, etc.
Why is this useful? Disentangled representations promise to be more interpretable, controllable, and generalizable:
-
Interpretability: Each dimension corresponds to a semantically meaningful factor, making the learned features easier for humans to understand.
-
Controllability: The factors can be independently manipulated to generate new data with specific properties.
-
Generalizability: Separating the factors may improve robustness and transfer to new tasks and domains.
The hope is that unsupervised learning of disentangled representations, without relying on expensive labels, can provide a foundation for more flexible and capable AI systems. Some potential applications include:
- Generating realistic synthetic data by controlling the underlying factors
- Few-shot learning by leveraging the disentangled factors
- fair machine learning by ensuring the decisions don‘t rely on sensitive factors
- Making reinforcement learning more sample-efficient and robust
Common Approaches and Assumptions
The most popular current approaches for unsupervised disentanglement are based on variations of the variational autoencoder (VAE) [3]. VAEs jointly train an encoder network that maps data to a latent representation and a decoder that reconstructs the data from the latent code. To encourage disentanglement, the latent distribution is regularized to be as close as possible to a factorized prior, often an isotropic Gaussian.
Various regularization terms have been proposed to strengthen the pressure for disentanglement:
- β-VAE [4] scales up the KL divergence term in the VAE objective, limiting the capacity of the latent code and encouraging independence between dimensions.
- FactorVAE [5] and β-TCVAE [6] penalize the total correlation (TC) between latent dimensions.
- DIP-VAE [7] matches the moments of the aggregate posterior to a factorized prior.
The common intuition behind these approaches is that forcing the aggregate posterior to be factorized should make the individual dimensions of the latent representation independent.
Another common evaluation practice is to measure disentanglement with metrics that rely on known ground truth factors, like the BetaVAE score [4], Mutual Information Gap (MIG) [8], and DCI [9]. These metrics generally capture how well each ground truth factor is isolated in a single dimension of the learned representation.
Challenging the Conventional Wisdom
Google‘s ICML paper [1] set out to rigorously test how well these unsupervised disentanglement approaches and evaluation metrics actually work. They conducted a large-scale experimental study, training over 12,000 models covering 6 state-of-the-art methods on 7 datasets. Importantly, they open-sourced the code for their experimental setup and the trained models, providing a reproducible foundation for the research community to build on. Let‘s look at some of the key findings:
1. Mismatch between Representation and Aggregate Posterior
One of the most striking results is that while the unsupervised methods successfully make the aggregate posterior factorized, this does not always translate to the representation itself being disentangled. They find cases where the representation dimensions are significantly correlated even though the aggregate posterior is not. The implication is that the common practice of reading out the mean of the representation during evaluation (as opposed to sampling from the aggregate posterior) may be flawed.
2. Failure to Consistently Disentangle
The disentanglement scores of unsupervised models proved to be highly sensitive to hyperparameters and random seeds, much more so than the model architecture. With the exact same hyperparameters, different random initializations could produce drastically different disentanglement scores. This suggests that the models fail to reliably learn disentangled representations in a way that is agnostic to initialization.
The table below shows how the FactorVAE disentanglement score varies drastically across random seeds on the dSprites dataset:
| Model | Min. Score | Max. Score | Mean Score | Std. Dev. |
|---|---|---|---|---|
| β-VAE | 0.63 | 0.95 | 0.83 | 0.08 |
| FactorVAE | 0.38 | 0.89 | 0.68 | 0.12 |
| DIP-VAE-I | 0.29 | 0.83 | 0.62 | 0.13 |
| DIP-VAE-II | 0.29 | 0.86 | 0.63 | 0.14 |
3. No Consistent Recipes for Model Selection
A key challenge for unsupervised learning is model selection – how to choose the best hyperparameters and random seed when we don‘t have access to ground truth labels. Unfortunately, the study found no consistent recipes for model selection that work across different datasets and methods. Hyperparameters that worked well on one dataset often didn‘t transfer to another. And selecting based on unsupervised metrics like reconstruction error didn‘t correlate well with disentanglement scores.
4. Disentangled Representations Don‘t Help Downstream Tasks
Perhaps most concerningly, the study did not find evidence that unsupervised disentangled representations actually help downstream tasks in terms of sample efficiency or robustness. When used as features for tasks like object classification and regression, the disentangled representations didn‘t perform any better than entangled ones. This calls into question the practical value of disentanglement, at least with current unsupervised approaches.
The Road Ahead
The Google AI study paints a humbling picture of the current state of unsupervised disentanglement learning. It shows that commonly held assumptions and evaluation practices can be misleading, and that our current approaches have significant limitations in learning disentangled representations that are truly reliable and useful.
However, this work is far from the end of the story. The rigorous methodology, comprehensive experiments, and open-source contributions provide an invaluable foundation for the research community to build on. By identifying the current pitfalls and challenges, Google has provided a roadmap for what we need to solve to make disentanglement learning work.
Some key open problems and directions for future work include:
- Developing principled approaches for unsupervised model selection and hyperparameter tuning
- Improving the reliability and reproducibility of disentanglement learning across random initializations
- Advancing our theoretical understanding of what disentanglement means and how to measure it
- Exploring applications where disentangled representations provide measurable benefits
- Expanding beyond purely unsupervised approaches to incorporate other forms of weak supervision
It‘s also important to view disentanglement in the broader context of representation learning. Recent advances in other unsupervised paradigms like contrastive learning [10] and self-supervised learning [11] provide alternative approaches to extracting useful features from raw data. A holistic view of representation learning may require going beyond a narrow conception of disentanglement.
Unsupervised learning remains a grand challenge for AI, but one with immense potential benefits. By revealing the gaps in our current approaches, Google‘s work is an important step in the long journey towards AI systems that can learn rich, versatile representations from raw unlabeled experience more like humans do. As AI pioneer Yann LeCun wrote, "Self-supervised learning is the future" [12]. To get there, careful science and a willingness to challenge assumptions will be essential. Google AI has provided a valuable lesson for the field in how to do that.
