[Paper Review] Variational Autoencoder with Learned Latent Structure

Marissa Connor, Gregory Canal+1 more|arXiv (Cornell University)|Jun 18, 2020

Generative Adversarial Networks and Image Synthesis27 references4 citations

TL;DR

This paper proposes VAELLS, a variational autoencoder that learns a data-driven latent manifold structure using transport operators to replace the fixed prior in traditional VAEs. By modeling nonlinear data manifolds through learnable dynamics matrices and anchor-based priors, VAELLS enables more accurate generation, meaningful interpolation paths, and class-specific transformations, outperforming standard VAEs on MNIST and rotated MNIST benchmarks.

ABSTRACT

The manifold hypothesis states that high-dimensional data can be modeled as lying on or near a low-dimensional, nonlinear manifold. Variational Autoencoders (VAEs) approximate this manifold by learning mappings from low-dimensional latent vectors to high-dimensional data while encouraging a global structure in the latent space through the use of a specified prior distribution. When this prior does not match the structure of the true data manifold, it can lead to a less accurate model of the data. To resolve this mismatch, we introduce the Variational Autoencoder with Learned Latent Structure (VAELLS) which incorporates a learnable manifold model into the latent space of a VAE. This enables us to learn the nonlinear manifold structure from the data and use that structure to define a prior in the latent space. The integration of a latent manifold model not only ensures that our prior is well-matched to the data, but also allows us to define generative transformation paths in the latent space and describe class manifolds with transformations stemming from examples of each class. We validate our model on examples with known latent structure and also demonstrate its capabilities on a real-world dataset.

Motivation & Objective

• To address the mismatch between fixed priors and complex, nonlinear data manifolds in standard VAEs.
• To enable meaningful, nonlinear transformation paths in the latent space that reflect natural data variations.
• To model class-specific manifolds using transformations derived from class examples.
• To improve generative quality and disentanglement by learning a data-adaptive prior via transport operators.
• To provide a framework for interpretable, identity-preserving transformations in latent space.

Proposed method

• Introduces a learnable manifold model using transport operators, where dynamics are defined by a weighted sum of learned basis matrices $\Psi_m$.
• Uses a linear dynamical system $\dot{z} = Az$ to define smooth paths on the latent manifold, with solutions $z_t = \mathrm{expm}(At)z_0$.
• Defines a generative model where $z_1 = \mathrm{expm}(A)z_0 + n$, allowing interpolation and extrapolation via learned $A$ matrices.
• Employs anchor points per class to define a localized prior, with the prior distribution based on the closest anchor point.
• Integrates the transport operator model into the VAE framework by learning the $\Psi_m$ matrices and coefficients $c_m$ jointly with the encoder and decoder.
• Uses a multi-stage optimization with separate learning rates for network, anchors, and $\Psi$ matrices, including warm-up and coefficient inference restarts.

Experimental results

Research questions

• RQ1 Can a learnable latent manifold model improve generative performance compared to standard VAEs with fixed priors?
• RQ2 Can transport operators enable more realistic interpolation and extrapolation paths in the latent space?
• RQ3 Can class-specific transformations be effectively modeled and learned from example data?
• RQ4 Does learning a data-adaptive prior via anchor-based transport operators enhance disentanglement and sample quality?
• RQ5 How does VAELLS perform on datasets with known geometric structure, such as rotated MNIST?

Key findings

• VAELLS achieves improved sample quality and reconstruction accuracy on both rotated MNIST and natural MNIST datasets compared to standard VAEs.
• Latent space sampling in VAELLS produces realistic image variations, such as rotations and style changes, as shown in decoded outputs.
• The model learns eight distinct transport operators on MNIST that correspond to interpretable digit transformations, such as stroke thickening or curvature changes.
• Applying learned transport operators with positive or negative coefficients generates meaningful extrapolations, as visualized in Figures 17 and 18.
• The latent space exhibits better class separation and identity-preserving paths, reducing unintended class transitions during interpolation.
• Anchor points remain stable during training, indicating robustness of the localized prior model.