[Paper Review] Learning Manifold Dimensions with Conditional Variational Autoencoders

Yijia Zheng, Tong He+2 more|arXiv (Cornell University)|Feb 23, 2023

Generative Adversarial Networks and Image Synthesis10 citations

TL;DR

The paper proves that VAE and CVAE global minima can learn the true data manifold dimension and extends the result to CVAEs with discrete/continuous conditioning, supported by synthetic and real data experiments.

ABSTRACT

Although the variational autoencoder (VAE) and its conditional extension (CVAE) are capable of state-of-the-art results across multiple domains, their precise behavior is still not fully understood, particularly in the context of data (like images) that lie on or near a low-dimensional manifold. For example, while prior work has suggested that the globally optimal VAE solution can learn the correct manifold dimension, a necessary (but not sufficient) condition for producing samples from the true data distribution, this has never been rigorously proven. Moreover, it remains unclear how such considerations would change when various types of conditioning variables are introduced, or when the data support is extended to a union of manifolds (e.g., as is likely the case for MNIST digits and related). In this work, we address these points by first proving that VAE global minima are indeed capable of recovering the correct manifold dimension. We then extend this result to more general CVAEs, demonstrating practical scenarios whereby the conditioning variables allow the model to adaptively learn manifolds of varying dimension across samples. Our analyses, which have practical implications for various CVAE design choices, are also supported by numerical results on both synthetic and real-world datasets.

Motivation & Objective

• Demonstrate that globally optimal VAEs recover the true data manifold dimension for data on low-dimensional manifolds.
• Extend the dimensionality recovery results to CVAEs with continuous and discrete conditioning variables (unions of manifolds).
• Analyze practical CVAE design choices (decoder variance handling, weight sharing) and their impact on learning manifold dimensions.
• Provide numerical evidence on synthetic and real-world datasets to validate theoretical claims.

Proposed method

• Define kappa-simple VAE and CVAE with Gaussian encoder/decoder and Gaussian prior.
• Prove that when data lie on an r-dimensional manifold in Rd, the global VAE minimum uses exactly r active latent dimensions and achieves reconstruction error tending to zero as decoder variance gamma -> 0.
• Extend the analysis to CVAEs showing active dimensions reduce to r-t when conditioning variable c has effective dimension t.
• Discuss adaptive active dimensions for unions of manifolds and scenarios with discrete or continuous conditioning variables.
• Examine CVAE design choices such as conditional vs unconditional priors, gamma initialization, and encoder/prior weight sharing, with theoretical and empirical implications.
• Provide synthetic and real-data experiments (MNIST, Fashion-MNIST) to corroborate theory.

Experimental results

Research questions

• RQ1 Can the global minima of VAEs recover the true manifold dimension r from data lying on an r-dimensional manifold?
• RQ2 How do conditioning variables (continuous or discrete) in CVAEs affect the learned manifold dimension and reconstruction capabilities?
• RQ3 What practical CVAE design choices (decoder variance handling, prior/encoder weight sharing) influence the ability to learn manifold dimensions?
• RQ4 Can CVAEs adaptively learn varying manifold dimensions across samples or regions (unions of manifolds)?

Key findings

• Global minima of kappa-simple VAEs recover the manifold dimension r; the number of active latent dimensions almost surely equals r.
• Reconstruction error scales as O(gamma) while the loss includes (d−r) log gamma, indicating dimensionality learning behavior as gamma -> 0.
• In CVAEs, conditioning variables reduce the required active latent dimensions to r−t when the conditioning has effective dimension t.
• For unions of manifolds, CVAEs with appropriate architecture (e.g., attention in decoder) can adaptively learn varying active dimensions across regions.
• Certain design choices (ignoring conditional priors, choice of gamma initialization, and weight sharing) significantly influence optimization and dimension recovery, with theoretical and empirical support.
• Empirical results on synthetic data and MNIST/Fashion-MNIST align with theoretical predictions, showing AD (active dimensions) matches r or r−t as appropriate.