Unpicking Data at the Seams: Understanding Disentanglement in VAEs

TL;DR

This paper formalizes disentanglement in VAEs by defining it as a factorization of the data-manifold density $p_\mu$ into independent 1-D seam factors, tied to axis-aligned latent traversals. It shows an exact relation between diagonal posterior covariances and decoder derivatives, producing two key constraints (C1-C2) on the decoder Jacobian that enforce seam-aligned, independent components via the Jacobian SVD, with singular-vector paths mapping to data seams. Under these conditions, seam factorization yields disentanglement and, in the linear/Gaussian setting, identifiability of ground-truth factors up to permutation and sign; empirical results on linear models and dSprites corroborate the theory, showing diagonal posteriors promote axis-aligned, independent components, while increasing $\beta$ broadens the posterior and enhances disentanglement at the cost of reconstruction fidelity.

Abstract

A generative latent variable model is said to be disentangled when varying a single latent co-ordinate changes a single aspect of samples generated, e.g. object position or facial expression in an image. Related phenomena are seen in several generative paradigms, including state-of-the-art diffusion models, but disentanglement is most notably observed in Variational Autoencoders (VAEs), where oft-used diagonal posterior covariances are argued to be the cause. We make this picture precise. From a known exact link between optimal Gaussian posteriors and decoder derivatives, we show how diagonal posteriors "lock" a decoder's local axes so that density over the data manifold factorises along independent one-dimensional seams that map to axis-aligned directions in latent space. This gives a clear definition of disentanglement, explains why it emerges in VAEs and shows that, under stated assumptions, ground truth factors are identifiable even with a symmetric prior.

Figures (7)

• Figure 1: Disentanglement: full vs diagonal posteriors ($\bm{\Sigma_x}$). Right singular vectors ${\bm{v}}^i\!\in\!{\mathcal{Z}}$ (blue) of the decoder's Jacobian define singular vector paths (dashed blue); left singular vectors ${\bm{u}}^i$ define seams (dashed red). 1-D densities over seams factorise the manifold density. (left) with full posteriors, s.v. paths are not axis-aligned; the axis‑traversal image in ${\mathcal{X}}$ (green) does not follow the seam. (right) under C1-C2 induced by diagonal posteriors, s.v. paths axis‑align and the traversal image follows the seam everywhere, and 1-D densities over seams are independent, achieving disentanglement (D\ref{['def:disentanglement']}).
• Figure 2: An LVAE breaks rotational symmetry. (l) both full-$\Sigma_x$ and diagonal-$\Sigma_x$ VAEs fit the data, i.e. learn ground truth parameters ${\bm{U}}_*$, ${\bm{S}}_*$; (c) only in diagonal-$\Sigma_x$ VAEs do right singulars ${\bm{v}}^i$ of the Jacobian align with standard basis vectors ${\bm{z}}_i\!\in\!{\mathcal{Z}}$ (i.e. ${\bm{V}}_*\!\to\!{\bm{I}}$); (r) images of ${\bm{z}}_i$: full-$\Sigma_x$ VAEs map ${\bm{z}}_i$ to arbitrary directions (red), but diagonal-$\Sigma_x$ VAEs learn (later epochs darker) to map ${\bm{z}}_i$ to the data's independent components (black, i.e. blue $\to$ black).
• Figure 3: (left) Empirical support for C1: rolinek2019variational show that VAEs with diagonal $\Sigma_x$ have increased orthogonality in the decoder Jacobian. (right) Seam factorisation: For $g\!:\!{\mathcal{Z}}\!\to\!{\mathcal{X}}$ in \ref{['thm:seam-factorisation-general']}, Jacobian ${\bm{J}}_{z^*}$ and manifold ${\mathcal{M}}_g\!\subseteq\!{\mathcal{X}}$, singular-vector paths${\mathcal{V}}^k_{z^*}\!\subseteq\!{\mathcal{Z}}$ (D\ref{['def:sv-path']}, dashed blue) following right singular vectors ${\bm{v}}^{k}\!$ of ${\bm{J}}_{z^*\!}$ (solid blue), map to seams${\mathcal{M}}_{g,z^*}^k\!\subseteq\! {\mathcal{M}}_g$ (D\ref{['def:seams']}, dashed red) following left singular vectors ${\bm{u}}^{k}$ at $g(z^*)\!\in\!{\mathcal{X}}$ (solid red). ${\bm{z}}_k\!\in\!{\mathcal{Z}}$ are standard basis vectors.
• Figure 4: Pushing forward $\bm{p(z)}$ from singular vector paths to seams: s.v. path ${\mathcal{V}}_{z^*}^i \!\!\subseteq\!{\mathcal{Z}}$ through $z^*$ (dash blue), following right singular vectors ${\bm{v}}^i$ of ${\bm{J}}_{z}$ (solid blue), maps to seam ${\mathcal{M}}_{f,z^*}^i \!\!\subseteq\!{\mathcal{M}}_f$ through $f(z^*)$ (dash red), following left singular vectors ${\bm{u}}^i$ (solid red). By \ref{['thm:seam-factorisation-general']}, 1-D marginals $p_i(z_i)$ over ${\mathcal{V}}_{z^*}^i$ factorising $p(z)$ map to 1-D seam densities over ${\mathcal{M}}_{f,z^*}^i$factorising$p_\mu$. (l) For linear $f$withoutC1, e.g. full-covariance LVAE, ${\mathcal{V}}_{z^*}^i$ and ${\mathcal{M}}_{f,z^*}^i$ are straight lines; $p_i(z_i)$ and seam densities are independent. (r) For general c.i.d.a.e. $f$withC1-C2, e.g. Gaussian VAE with P\ref{['prp:non-cancellation']}, seam densities are independent (by C2). (Essential properties of disentanglement (D\ref{['def:disentanglement']}) are underlined.)
• Figure 5: (t) $\smaller {\bm{J}}^\top_z{\bm{J}}$; (b) $\smaller (x\hbox{[}1.0]{$-$} d(z)){\bm{\mathsfit{H}}}_z$

Theorems & Definitions (29)

• Lemma 4.0: Disentanglement constraints
• Lemma 5.0: Factorisation over seams
• Theorem 5.1: Disentanglement $\bm{\Leftrightarrow}$ C1-C2
• Corollary 6.0: LVAE Identifiability
• proof
• Remark 6.1: ${\bm{V}}_{\mathsmaller{\bm{W}}}$ immaterial
• Lemma 6.1: Seams are Intrinsic
• Lemma 6.1: A matching decoder finds seams
• Theorem 6.2: Gaussian VAE Identifiability
• Remark 6.3