Mutual Information Collapse Explains Disentanglement Failure in $β$-VAEs

TL;DR

The paper identifies an information-theoretic collapse in $\beta$-VAEs: when $\beta>1$, stationarity dynamics contract the encoder gain, driving latent mutual information $I(\mathbf{X};\mathbf{V})$ to zero and degenerating disentanglement metrics such as MIG and SAP. It formalizes this in a linear-Gaussian framework and proves that the trivial, uninformative fixed point emerges under high regularization. To counteract this, it introduces the $\lambda\beta$-VAE, which adds an auxiliary reconstruction penalty to decouple semantic signal from regularization pressure, thereby preserving latent informativeness over a broader range of $\beta$. The authors validate the theory through controlled linear-Gaussian experiments and deep nonlinear benchmarks (dSprites, Shapes3D, MPI3D-real), showing that $\lambda>0$ stabilizes disentanglement while maintaining reconstruction quality. This two-parameter regularization provides a principled design principle for robust, informative disentangled representations in complex scientific and generative tasks.

Abstract

The $β$-VAE is a foundational framework for unsupervised disentanglement, using $β$ to regulate the trade-off between latent factorization and reconstruction fidelity. Empirically, however, disentanglement performance exhibits a pervasive non-monotonic trend: benchmarks such as MIG and SAP typically peak at intermediate $β$ and collapse as regularization increases. We demonstrate that this collapse is a fundamental information-theoretic failure, where strong Kullback-Leibler pressure promotes marginal independence at the expense of the latent channel's semantic informativeness. By formalizing this mechanism in a linear-Gaussian setting, we prove that for $β> 1$, stationarity-induced dynamics trigger a spectral contraction of the encoder gain, driving latent-factor mutual information to zero. To resolve this, we introduce the $λβ$-VAE, which decouples regularization pressure from informational collapse via an auxiliary $L_2$ reconstruction penalty $λ$. Extensive experiments on dSprites, Shapes3D, and MPI3D-real confirm that $λ> 0$ stabilizes disentanglement and restores latent informativeness over a significantly broader range of $β$, providing a principled theoretical justification for dual-parameter regularization in variational inference backbones.

Figures (8)

• Figure 1: Linear-Gaussian VAE: Ground-truth factors $\mathbf{V}$ generate observations $\mathbf{Y}$; the encoder maps $\mathbf{Y}\mapsto \mathbf{X}$; the decoder yields reconstructions $\hat{\mathbf{Y}}$.
• Figure 2: Reconstruction error, SAP score, and $I_m$ score for the $\beta$-VAE with $(n,m,s)=(100,10,5)$, using fixed-point iteration and AdamW. For $\beta>1$, SAP and $I_m$ collapse toward zero across both optimization procedures, consistent with convergence to the trivial low-information solution in Theorem \ref{['thm:trivial']}.
• Figure 3: Mean reconstruction error, SAP score, and $I_m$ score for the $\lambda\beta$-VAE with $(n,m,s)=(100,10,5)$. Positive $\lambda$ prevents the collapse observed at $\lambda=0$ and preserves informative latent-factor dependence at larger $\beta$, yielding non-degenerate SAP and $I_m$.
• Figure 4: Reconstruction error, SAP score, and MIG score for the $\beta$-VAE across dSprites, Shapes3D, and MPI3D-real. Disentanglement peaks at intermediate $\beta$ values and deteriorates at large $\beta$.
• Figure 5: Mean reconstruction error, SAP, and MIG for the $\lambda\beta$-VAE across dSprites, Shapes3D, and MPI3D-real. Positive $\lambda$ improves reconstruction and preserves disentanglement at large $\beta$, consistent with information preservation.

Theorems & Definitions (7)

• Lemma 3.1: $\beta$-VAE Stationary Point
• Theorem 3.2: Informational Collapse for $\beta>1$
• Corollary 4.1: Metric Degeneracy for $\beta>1$
• Lemma 4.2: $\lambda\beta$-VAE Stationary Point
• Lemma 1.1: Uniform Spectral Bound
• proof
• proof