VAE Disentanglement

In VAE Resynthesis, the directions in space representing interpretably properties had to be estimated using labeled training data. Other works attempts to improve the characteristics of the latent space so that its coordinate directions correspond to real-world properties.

When each dimension represents an independent real-world factor, the latent space is described as disentangled. For example, when modeling face images, we might hope to uncover head pose or hair color as independent factors.

Methods to encourage disentanglement typically add regularization terms to the loss function based on either:

• The posterior $q(z|x,\theta )$ over the latent variables $z$
• The aggregated posterior

The regularized loss function is

$$L_{\text{new}}=-\text{ELBO}[\theta ,\varphi ]+{\lambda }_1{\mathbb{E}}_{Pr(x)}[r_1[q(z|x,\theta )]]+{\lambda }_2{r}_2[q(z|\theta )]$$

• $r_1[\bullet ]$ is a function of the posterior and weighted by ${\lambda }_1$
• $r_2[\bullet ]$ is a function of the aggregated posterior and is weighted by ${\lambda }_2$

For example, the beta VAE upweights the second term in the ELBO:

$$\text{ELBO}[\theta ,\varphi ]\approx \log [Pr(x|z^{\ast },\varphi )]-\beta \cdot D_{KL}[q(z|x,\theta )||Pr(z)]$$

where $\beta >1$ determines how much more the deviation from the prior $Pr(z)$ is weighted relative to the reconstruction error. Since the prior is usually a multivariate normal with a spherical covariance matrix, its dimensions are independent. Hence, up-weighting this term encourages the posterior distributions to be less correlated.

Another variant is the total correlation VAE, which adds a term to decrease the total correlation between variables in the latent space and maximizes the mutual information between a small subset of the latent variables and the observations.