Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Understanding DDPM Latent Codes Through Optimal Transport

Valentin Khrulkov, Gleb Ryzhakov, Andrei Chertkov, Ivan Oseledets

TL;DR

The paper tackles the fundamental question of what structure the DDPM encoder induces on data by examining its connection to optimal transport. It proves analytically that the encoder coincides with the Monge OT map to $\mathcal{N}(0,I)$ in the Gaussian (Gaussian diffeo) case and provides extensive numerical evidence via solving the Fokker-Planck equation with a tensor-train solver for synthetic distributions, showing near-perfect OT equivalence. It further validates the hypothesis on real image datasets by demonstrating consistent texture-level correspondences across multiple DDPMs and by comparing OT costs to the encoder-based transport costs, including a conditional ImageNet setting. The results offer a principled lens to understand latent codes in diffusion models, suggesting that pixel-space quadratic costs may be suboptimal for visual data and opening avenues for feature-based OT formulations and future theoretical proofs.

Abstract

Diffusion models have recently outperformed alternative approaches to model the distribution of natural images, such as GANs. Such diffusion models allow for deterministic sampling via the probability flow ODE, giving rise to a latent space and an encoder map. While having important practical applications, such as estimation of the likelihood, the theoretical properties of this map are not yet fully understood. In the present work, we partially address this question for the popular case of the VP SDE (DDPM) approach. We show that, perhaps surprisingly, the DDPM encoder map coincides with the optimal transport map for common distributions; we support this claim theoretically and by extensive numerical experiments.

Understanding DDPM Latent Codes Through Optimal Transport

TL;DR

The paper tackles the fundamental question of what structure the DDPM encoder induces on data by examining its connection to optimal transport. It proves analytically that the encoder coincides with the Monge OT map to in the Gaussian (Gaussian diffeo) case and provides extensive numerical evidence via solving the Fokker-Planck equation with a tensor-train solver for synthetic distributions, showing near-perfect OT equivalence. It further validates the hypothesis on real image datasets by demonstrating consistent texture-level correspondences across multiple DDPMs and by comparing OT costs to the encoder-based transport costs, including a conditional ImageNet setting. The results offer a principled lens to understand latent codes in diffusion models, suggesting that pixel-space quadratic costs may be suboptimal for visual data and opening avenues for feature-based OT formulations and future theoretical proofs.

Abstract

Diffusion models have recently outperformed alternative approaches to model the distribution of natural images, such as GANs. Such diffusion models allow for deterministic sampling via the probability flow ODE, giving rise to a latent space and an encoder map. While having important practical applications, such as estimation of the likelihood, the theoretical properties of this map are not yet fully understood. In the present work, we partially address this question for the popular case of the VP SDE (DDPM) approach. We show that, perhaps surprisingly, the DDPM encoder map coincides with the optimal transport map for common distributions; we support this claim theoretically and by extensive numerical experiments.
Paper Structure (27 sections, 1 theorem, 45 equations, 5 figures, 1 table)

This paper contains 27 sections, 1 theorem, 45 equations, 5 figures, 1 table.

Table of Contents

  1. Introduction
  2. Reminder on diffusion models
  3. Denoising diffusion probabilistic models
  4. Stochastic and deterministic sampling.
  5. SDE view on Diffusion Models
  6. Encoder map.
  7. DDPM Encoder and Optimal Transport
  8. Toy example
  9. Reminder on optimal transport
  10. Kantorovich and Monge formulation.
  11. Multivariate normal distribution case
  12. Discussion.
  13. Numerical experiments on synthetic data
  14. Method.
  15. Distributions.
  16. ...and 12 more sections

Key Result

Theorem 3.1

Let $\mu_0 \sim \mathcal{N}(a(0), \Sigma(0))$ be a multivariate normal distribution. Then $E_{\mu_0}$ is the Monge optimal transport map between $\mu_0$ and $\mathcal{N}(0, I)$, i.e.,

Figures (5)

  • Figure 1: A toy example of our approach on a $2d$ multivariate normal distribution $\mu$. The first four plots visualize the diffusion process. The fifth plot demonstrates the trajectories of the probability flow ODE. In this case, the optimal transport map is known analytically and exactly coincides with the mapping $E_{\mu}$, introduced in \ref{['sec:reminder']}. Note that the trajectories of the probability flow ODE are not straight lines even in this simple case.
  • Figure 2: An example of $2$-dimensional distributions considered for numerical comparison of the Monge optimal transport map against the DDPM encoder map.
  • Figure 3: An example of the trajectories of the probability flow ODE and the limiting encoder map for two $2$-dimensional distributions studied in \ref{['sec:numerical']}.
  • Figure 4: Examples of synthetic samples produced with DDIM sampling from the same latent codes. (Top) Three independent DDPMs trained on AFHQ Dog/Cat/Wild. Each row translates into each row in most cases preserving high-level semantics such as pose and texture, thus supporting the claim of our hypothesis. (Bottom) Two DDPMs trained on FFHQ/MetFaces; we note that high-level features such as gender and texture/color are transferred. Interestingly, 'noisy' samples seem also to be shared across models.
  • Figure 5: Examples of synthetic samples produced with DDIM sampling from same latent codes for different classes with a conditional ImageNet model. We observe that samples share texture and pixel-based similarity.

Theorems & Definitions (1)

  • Theorem 3.1