Disentangled representation learning

1 Grandaddy example: β-VAE

The β-VAE augments the standard variational autoencoder objective

$$\mathcal{L}={\mathbb{E}}_{q(z\mid x)}[-\log p(x\mid z)]+\mathrm{KL}(q(z\mid x)\parallel p(z))$$

by weighting the KL term with a factor β > 1:

$${\mathcal{L}}_{\beta \text{-VAE}}={\mathbb{E}}_{q(z\mid x)}[-\log p(x\mid z)]+\beta \mathrm{KL}(q(z\mid x)\parallel p(z)).$$

This stronger bottleneck (larger β) encourages each latent dimension to carry only the minimal information needed, which pushes them toward independence—and often yields interpretable axes like “rotate” or “zoom.”

2 Fancier: total-correlation in β-TCVAE

β-TCVAE decomposes the VAE’s KL into three parts—mutual information, dimension-wise KL, and total correlation (TC)—and then penalizes TC more heavily. Concretely:

$$\mathrm{TC}(q(z))=\mathrm{KL}(q(z)\parallel \prod _{i}q(z_i)),$$

and the objective becomes

$${\mathcal{L}}_{\beta \text{-TCVAE}}={\mathbb{E}}_{q(z\mid x)}[-\log p(x\mid z)]+\underset{\text{info term}}{\underbrace{\alpha I_q(x;z)}}+\underset{\text{disentanglement}}{\underbrace{\beta \mathrm{TC}(q(z))}}+\underset{\text{marginals}}{\underbrace{\gamma \sum _i\mathrm{KL}(q(z_i)\parallel p(z_i))}}.$$

By dialling up β on the TC term, we more directly push the joint q(z) toward a product of its marginals, which sharpens disentanglement.

3 InfoGAN

InfoGAN (Chen et al. 2016) is a GAN-based disentangling method that lives in that same family. It augments the usual GAN min–max with a mutual-information term to coax one part of the latent code to line up with a single factor:

$$\underset{G}{min}\underset{D}{max}V(D,G)-\lambda I(c;G(z,c)),$$

where

• $V(D,G)$ is the standard GAN loss,
• $z$ is noise,
• $c$ is a “code” you hope will become interpretable (e.g. rotation, thickness), and
• $I(c;G(z,c))$ encourages $c$ to actually control something in the output.

By maximizing $I(c;G(z,c))$, InfoGAN pushes those chosen code dimensions to capture distinct, meaningful factors—just like β-VAE or β-TCVAE push their KL/TC penalties toward independence.

4 Other examples

In general, it seems you pick a backbone (VAE, GAN, diffusion, …), add one of these disentangling penalties. You have succeeded if traversing a single latent smoothly transforms just one aspect of the output. A handful of dimensions are then devoted to “rotation,” “colour,” or “thickness” etc, and the rest of the model behaves normally, doing whatever you were doing before, except now our face generator has a colour knob.

5 Causal

Interesting, and why I got involved in this field in the first place, after seeing Yang et al. (2021). See Causal abstraction for more.

6 Questions

• How many factors can we disentangle?
• How much can this be made unsupervised?
• Are there underexploited tools in this toolkit for Mechanistic interpretability for more.
• Are there underexploited tools in Developmental interpretability for more.

7 References