Back to blog

Disentangled Representations: Why VAEs Can't Disentangle (and What Actually Works)

Imagine you're building a face generation tool. You want a slider for "smile," another for "age," another for "lighting." Each slider controls exactly one thing without affecting the others. That's the promise of disentangled representations: a latent space where each dimension maps to one interpretable factor of variation.

This idea shows up everywhere. In drug discovery, you want to adjust a molecule's binding affinity without changing its toxicity. In image editing, StyleGAN's appeal is that you can change hair color without also changing the background. In fairness, you need to separate "relevant features" from "sensitive attributes" like race or gender.

The problem is that Locatello et al. proved in 2019 that achieving this is fundamentally impossible without additional assumptions. Not "hard." Not "requires more data." Mathematically impossible with standard unsupervised methods. I worked through the proof during my deep learning course at Purdue, and it's one of those results that changes how you think about the whole field.

The Identifiability Problem

Say you have a generative model that maps latent variables $z$ to observations $x$ through some function $f$, so $x = f(z)$. The model learns $f$ from data. The question is: can you recover the "true" latent factors from the learned representation?

The answer is no, and here's why. For any orthogonal rotation matrix $Q$, the pair $(f, z)$ and $(f \circ Q^T, Qz)$ produce exactly the same observations. The model $f \circ Q^T$ applied to the rotated latents $Qz$ gives you:

$$(f \circ Q^T)(Qz) = f(Q^T Q z) = f(z) = x$$

Same output. Different latent space. And if your latent prior is a standard Gaussian $\mathcal{N}(0, I)$, it's rotationally symmetric, so $Qz$ has exactly the same distribution as $z$. There is literally no statistical signal that distinguishes the "true" representation from any rotated version.

This means a VAE or any unsupervised generative model has infinitely many equally valid latent spaces, related by arbitrary rotations. One of them might be nicely disentangled. The rest aren't. And the model has no way to prefer the disentangled one.

What This Looks Like in Practice

Imagine you train a VAE on face images. You hope dimension 1 captures "smile" and dimension 2 captures "head rotation." But any rotation of those two dimensions is equally valid:

  • (0.7 * smile + 0.7 * rotation, -0.7 * smile + 0.7 * rotation) is a valid representation
  • (-smile, rotation) is a valid representation
  • Any 2D rotation of (smile, rotation) is a valid representation

They all reconstruct images equally well. They all have the same ELBO. The loss function simply cannot distinguish between them.

This is why different random seeds produce different "disentanglement" results in papers. The model isn't learning different things about the data. It's landing on different rotations of the same underlying representation. The one that happens to align with human-interpretable factors is luck, not learning.

beta-VAE and Why It Sometimes Seems to Work

beta-VAE (Higgins et al., 2017) adds a weight $\beta > 1$ to the KL divergence term:

$$\mathcal{L} = \mathbb{E}[\log p(x|z)] - \beta \cdot D_{KL}(q(z|x) \| p(z))$$

The higher $\beta$ pushes the posterior closer to the prior, which encourages independence between latent dimensions. And it does produce representations that look more disentangled on benchmarks like dSprites and CelebA.

But here's the catch: the Locatello paper showed that across 12,000 trained models with different hyperparameters and random seeds, there was no reliable relationship between the inductive bias (beta value, model architecture) and the degree of disentanglement achieved. Good disentanglement was mostly a function of the random seed. Some seeds got lucky.

The fundamental impossibility result doesn't go away just because you increase $\beta$. What $\beta > 1$ does is reduce the capacity of the latent space (by pushing it toward the prior), which can accidentally encourage axis-alignment in some cases. But it doesn't solve the rotation ambiguity in any principled way.

Sparsity: The Fix That Actually Works

So how do you get identifiable (provably unique) representations? You need to break the rotational symmetry. One way that works is sparsity.

The idea: require that the Jacobian of the generative function $f$ is sparse. Meaning each observation dimension depends on only a few latent dimensions, not all of them.

Why does this help? Consider what happens when you rotate the latent space. If the original Jacobian $J_f(z)$ has at most $s$ nonzero entries per row (each output pixel depends on at most $s$ latent factors), then the rotated version $J_{f \circ Q^T}(z') = J_f(z) \cdot Q^T$ will generally have more nonzero entries per row. Matrix multiplication with a non-diagonal matrix fills in zeros.

The proof is elegant. If $Q$ is not a permutation-and-scaling matrix, then $J_f \cdot Q^T$ has strictly more nonzeros per row than $J_f$. So if you constrain the Jacobian to be sparse, the only valid transformations $Q$ are permutations and scalings, not arbitrary rotations. That makes the representation identifiable up to the ordering and scale of dimensions, which is the best you can hope for.

# Conceptual: sparse Jacobian encourages identifiability
# In practice, you'd add a sparsity penalty to the decoder Jacobian

def sparsity_loss(decoder, z_sample):
    """Penalize non-sparse Jacobian of the decoder."""
    z_sample.requires_grad_(True)
    x_reconstructed = decoder(z_sample)

    # Compute Jacobian
    jacobian = torch.autograd.functional.jacobian(decoder, z_sample)

    # L1 penalty encourages sparsity
    return torch.abs(jacobian).mean()

Real-World Examples

Disentanglement matters beyond academic benchmarks. A few places where it shows up:

Drug discovery. You want a latent space where one dimension controls a molecule's binding affinity and another controls its toxicity. If these are entangled, you can't optimize one without accidentally changing the other. With a disentangled representation, you can search along the binding affinity axis while holding toxicity fixed.

Image editing. Tools like StyleGAN's latent space let you edit specific attributes (age, expression, lighting) because the representation is partially disentangled. When it's not, editing "hair color" also changes "background" or "pose" because those factors are mixed in the same dimensions.

Fairness in ML. If a model's representation entangles "relevant features" with "sensitive attributes" (race, gender), the model can't be fair even if you remove the sensitive attribute at prediction time. Disentangling the representation is a prerequisite for removing the influence of protected attributes.

Transfer learning. A disentangled representation is more likely to transfer across tasks, because each dimension captures a distinct concept. An entangled representation might capture a task-specific combination of features that's useless for a different task.

The Broader Lesson

What I found most interesting about the identifiability result isn't the math itself (though the proof via orthogonal rotations is clean). It's the implication for how we evaluate representation learning.

If unsupervised disentanglement is impossible without extra assumptions, then every paper claiming to achieve it must be relying on some implicit assumption, whether they acknowledge it or not. That assumption might be the architecture (convolutional structure imposes spatial locality), the optimizer (SGD with specific hyperparameters), or the data (natural images have statistical structure that breaks some symmetries).

Understanding what those implicit assumptions are, and when they hold, is more useful than finding yet another metric for "disentanglement score" on dSprites. The Locatello paper pushed the field in this direction, and recent work on identifiable VAEs, causal representation learning, and sparse independent component analysis has followed.

The takeaway isn't that disentangled representations are hopeless. It's that you need to be explicit about what inductive bias is doing the work. "Train a VAE and hope for the best" is not a strategy. Constraining the Jacobian, using multiple environments or data modalities, or incorporating weak supervision are strategies. The impossibility result doesn't close the door. It tells you which doors are actually open.