Geometry-Preserving Encoder/Decoder in Latent Generative Models

TL;DR

The paper tackles the challenge of diffusion-based generative modeling in high-dimensional spaces by proposing a geometry-preserving encoder/decoder (GPE) that embeds data onto a low-dimensional latent manifold while maintaining global geometric structure. It introduces a Gromov-Monge embedding-based encoder and trains the decoder separately via reconstruction, yielding provable stability and convergence advantages that differ from the ELBO-based VAE paradigm. Theoretical results establish weak bi-Lipschitz guarantees, stability of encoder training, and improved decoder convergence under geometry-preserving encodings, along with Wasserstein-distance bounds in latent diffusion models. Empirical results across MNIST, CIFAR-10, CelebA, and CelebA-HQ demonstrate faster reconstruction and generation convergence, with GPE often outperforming VAE baselines in both accuracy and efficiency, and with clearly better preservation of the data geometry in the latent space. Overall, GPE offers a scalable, stable alternative to VAEs for latent diffusion and related generative tasks, with potential extensions to multi-modal and large-model settings.

Abstract

Generative modeling aims to generate new data samples that resemble a given dataset, with diffusion models recently becoming the most popular generative model. One of the main challenges of diffusion models is solving the problem in the input space, which tends to be very high-dimensional. Recently, solving diffusion models in the latent space through an encoder that maps from the data space to a lower-dimensional latent space has been considered to make the training process more efficient and has shown state-of-the-art results. The variational autoencoder (VAE) is the most commonly used encoder/decoder framework in this domain, known for its ability to learn latent representations and generate data samples. In this paper, we introduce a novel encoder/decoder framework with theoretical properties distinct from those of the VAE, specifically designed to preserve the geometric structure of the data distribution. We demonstrate the significant advantages of this geometry-preserving encoder in the training process of both the encoder and decoder. Additionally, we provide theoretical results proving convergence of the training process, including convergence guarantees for encoder training, and results showing faster convergence of decoder training when using the geometry-preserving encoder.

Figures (13)

• Figure 1: Illustration showing how GPE (left) and VAE (right) encoders work. The GPE encoder embeds the data distribution into the latent space while maintaining the geometric structure of the data distribution. In contrast, the VAE encoder maps the data distribution onto the latent distribution (or a prior distribution) $\nu$ such that $T_\#\mu \approx \nu$. Thus, there exist mappings that send distinct inputs to nearby points in the latent space, resulting in $\alpha$ close to $0$.
• Figure 2: Visualization of latent generative model framework
• Figure 3: Comparison of embedded data distributions from GPE and VAE. The first column shows the input data distribution $\mu$ in 500D, while the second and third columns show the embedded distributions $T_\#\mu$ from GPE and VAE, respectively, in 2D.
• Figure 4: The plots illustrate the results of minimizing the reconstruction loss in \ref{['eq:rec-cost']} on four datasets, MNIST, CIFAR10, CelebA, and CelebA-HQ using three different encoders, each characterized by a different tolerance values showing the encoder with a smaller GME loss converges faster.
• Figure 5: Convergence of training a decoder $G$ using a geometry-preserving encoder with various datasets: first row (MNIST), second row (CIFAR10), third row (CelebA), and fourth row (CelebA-HQ). Each figure caption shows the iteration number and the time taken in seconds (s) and minutes (m). Different neural network architectures for the decoder were used for each dataset. All experiments were done with the same GPU settings: 2 A40 GPUs.

Theorems & Definitions (53)

• Definition 1: Weak $\alpha$–bi-Lipschitz
• Proposition 1: Bi-Lipschitz on well-separated pairs lee2023monotone
• Theorem 1
• proof
• Theorem 2: lee2023monotone
• Remark 1
• Example 1
• Proposition 2
• proof
• Proposition 3