PaperlayerView exact revision on arXiv

Source-linked AI summary

Understanding DDPM Latent Codes Through Optimal Transport

Valentin Khrulkov, Gleb Ryzhakov, Andrei Chertkov, Ivan Oseledets

arXiv:2202.07477v2stat.MLcs.AIcs.LGcs.NAmath.APmath.NA

TL;DR

The mathematical properties of deterministic diffusion-model encoder maps and their latent codes are not yet well understood. This paper studies DDPMs through the probability-flow encoder and optimal transport, finding agreement for multivariate normal distributions and strong numerical support across other tested distributions.

  • Problem

    The mathematical properties of deterministic diffusion-model encoder maps and their latent codes remain insufficiently understood, despite their practical applications.

    Source: §1 · p. 1

  • Method

    The paper studies the DDPM encoder induced by the probability-flow ODE, proving its optimal-transport equality for multivariate normal distributions and testing it numerically on synthetic and real data.

    Source: §1 · pp. 1–2

  • Results

    For multivariate normal distributions, the DDPM encoder is the Monge optimal transport map; experiments report high numerical precision on synthetic and real distributions.

    Source: §1 · p. 2, §3.3 · p. 5, §3.4 · p. 7

  • Takeaways & Limitations

    For visual DDPMs, pixel-based Euclidean distance corresponds to high-level texture- and color-level similarity, as observed on real models.

    Source: §1 · p. 1

    The paper leaves a more precise theoretical analysis for arbitrary distributions to future work, although it reports high-precision numerical support for a large class.

    Source: §3.3 · p. 5

Abstract

from arXiv · show

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.

1 INTRODUCTION

Although diffusion models support useful applications, the mathematical properties of their deterministic encoder maps and latent codes remain insufficiently understood. This work argues that, for DDPMs based on the VP SDE, the encoder map coincides with optimal transport, supporting the claim with theory and experiments.

1 INTRODUCTION

  • Diffusion models have outperformed alternative approaches on natural-image realism and diversity, and DDPMs are used in colorization, inpainting, super-resolution, and semantic editing.The passage says DDPMs often achieve more impressive results than GANs in these applications.Source: §1 · p. 1
  • For the VP SDE, the encoder map coincides with the Monge optimal transport map with large numerical accuracy.The paper supports this claim with extensive empirical evidence on controlled synthetic examples and real datasets.Source: §1 · p. 1
  • The paper addresses limited understanding of the mathematical properties of diffusion-model encoder maps and their latent codes.It identifies likelihood estimation among the practical applications of deterministic mappings.Source: §1 · p. 1
  • The DDPM encoder map coincides with the Monge optimal transport map for multivariate normal distributions, as verified theoretically.The paper also reports synthetic-distribution experiments with negligible errors and qualitative evidence on real image datasets.Source: §1 · p. 2

2 REMINDER ON DIFFUSION MODELS

The section introduces DDPMs, deterministic sampling, and the continuous SDE framework used to define an encoder from data toward the limiting distribution. The probability-flow ODE and Fokker–Planck equation provide the continuous tools for studying that map.

2.1 DENOISING DIFFUSION PROBABILISTIC MODELS

  • DDPMs gradually add noise to data toward a Gaussian distribution, while the generative model learns parameterized Gaussian transitions to reverse the process.The noise predictor estimates the noise component from a sample and timestep; at convergence its score estimate approximates the intermediate density's score function.Source: §2.1 · p. 2
  • DDIM changes the sampling algorithm without retraining DDPM, producing the same marginal distributions deterministically and enabling a latent space with x_T as input.The stochastic procedure has no single latent space; the deterministic approach supplies one.Source: §2.1 · p. 2

2.2 SDE VIEW ON DIFFUSION MODELS

  • The SDE view treats the diffusion as continuous stochastic dynamics, and reversing the probability-flow ODE maps the limiting distribution to the data distribution.The probability-flow ODE has the same marginal distributions as the SDE; the Fokker–Planck equation determines how a density evolves under an SDE.Source: §2.2 · p. 3
  • The analysis removes β(t) by a change of variables and works with the simplified equation because its specific form does not affect the results.The encoder map transports a point along the probability-flow vector field to its final encoding as t approaches infinity, and implicitly depends on intermediate densities.Source: §2.2 · p. 3

3 DDPM ENCODER AND OPTIMAL TRANSPORT

The paper investigates whether the DDPM encoder map to a standard normal is also an optimal transport map. It proves this for multivariate normals and tests the broader hypothesis numerically on synthetic distributions and image datasets.

3 DDPM ENCODER AND OPTIMAL TRANSPORT

  • Although each infinitesimal step of the diffusion trajectory is locally optimal, those steps need not compose into a global optimal map; the authors hypothesize closeness when the target is N(0, I).Their experiments test this hypothesis across synthetic and real distributions.Source: §3 · pp. 3–4

3.1 TOY EXAMPLE

  • For a shifted standard normal, integrating the probability-flow ODE gives Eµ0(x) = x − a, exactly the optimal transport map, despite nonstraight ODE trajectories.Figure 1 illustrates the diffusion process and the probability-flow paths for this toy case.Source: §3.1 · p. 4, Figure 1 · p. 4

3.2 REMINDER ON OPTIMAL TRANSPORT

  • The Monge formulation seeks a map pushing one distribution to another; this work uses quadratic cost, whose optimal cost defines W2.The paper also recalls the Kantorovich transport-plan formulation and its equivalence to the Monge formulation under stated general assumptions.Source: §3.2 · pp. 4–5

3.3 MULTIVARIATE NORMAL DISTRIBUTION CASE

  • For every multivariate normal, the DDPM encoder is the Monge optimal transport map to N(0, I), while exact equality for arbitrary distributions remains unproved.A cited counterexample differs at a single point, though the authors report extremely low differences for finite samples in higher dimensions.Source: §3.3 · p. 5

3.4 NUMERICAL EXPERIMENTS ON SYNTHETIC DATA

  • Across 100 random densities in d=2, 3, and 7, encoder-map and optimal-transport costs agreed up to machine precision; Table 1 reports maximum errors.The evaluation solves the Fokker–Planck equation using low-rank tensor-train approximations, then compares the ODE-induced transport cost with the optimal cost.Source: §3.4 · pp. 5–7, Table 1 · p. 7

3.5 EXPERIMENTS OF IMAGE DATASETS

  • Using the same latent codes, DDPM samples transferred pose and texture across AFHQ, FFHQ/MetFaces, and conditional ImageNet classes; conditional BigGAN showed no similar texture alignment.Shared features also included gender and model artifacts in face samples.Source: §3.5 · pp. 7–9
  • On AFHQ, optimal-transport costs between 500 validation images and their DDIM codes matched the ODE-induced costs to machine precision in all three categories.The authors note that the high-dimensional pixel space and small sample count make this exact match unsurprising in light of most synthetic experiments.Source: §3.5 · p. 9

4 RELATED WORK

Prior work connects optimal transport to diffusion processes, PDEs, computational methods, and normalizing flows. This paper distinguishes its DDIM connection to non-regularized OT from Schrödinger bridge theory and notes related theoretical counterexamples.

4 RELATED WORK

  • Schrödinger bridges relate score-based models to entropy-regularized OT, whereas this paper studies DDIM with non-regularized OT.The passage describes Schrödinger bridges as a continuous analogue of the Sinkhorn algorithm.Source: §4 · p. 9
  • Prior machine-learning work studied computational OT and its relationship to normalizing flows, while another method used Wasserstein gradient flows to solve the Fokker-Planck equation.The cited works are Cuturi (2013), Altschuler et al. (2017), Lin et al. (2019), Onken et al. (2021), and Mokrov et al. (2021).Source: §4 · p. 9
  • Lavenant & Santambrogio (2022) studied a similar question and provided a theoretical counterexample; this paper's analysis is mostly numerical and practically motivated.Its stated motivation is understanding the DDPM encoder.Source: §4 · p. 9

5 CONCLUSION

The conclusion presents theoretical and experimental evidence that the DDPM encoder map is optimal transport. It identifies pixel-space quadratic cost as seemingly suboptimal for visual data and leaves feature-based costs for future work.

5 CONCLUSION

  • The paper concludes that theoretical and experimental evidence supports the DDPM encoder map being the optimal transport map.The authors hope this result will inspire a general proof of the hypothesis.Source: §5 · p. 9
  • For visual data, quadratic cost in pixel space seems suboptimal, and a diffusion process leading to feature-based cost may be worth exploring.The authors leave this analysis to future work.Source: §5 · p. 9
  • The authors expect the result to interest both optimal transport and diffusion-model communities.They hope it will inspire researchers to obtain a general proof of the hypothesis.Source: §5 · p. 9

A BENAMOU-BRENIER FORMULATION OF OPTIMAL TRANSPORT

The Benamou-Brenier formulation describes particles moving between endpoint densities while minimizing work. Its optimal action equals squared Wasserstein-2 distance, but its velocity field generally differs from the probability-flow ODE even when the ODE's limiting flow map solves the Monge problem.

A BENAMOU-BRENIER FORMULATION OF OPTIMAL TRANSPORT

  • The formulation views transport as particles moving from density ρ0 to ρ1 while performing the least amount of work.In the continuous formulation, time-dependent densities and velocity fields obey endpoint constraints and the continuity equation.Source: A BENAMOU-BRENIER FORMULATION OF… · p. 13
  • The infimum of the Benamou-Brenier action is W2(ρ0, ρ1)^2.The passage also states that trajectories generated by the optimal-transport vector field are straight lines.Source: A BENAMOU-BRENIER FORMULATION OF… · p. 13
  • The probability-flow ODE need not solve the Benamou-Brenier problem, although its limiting flow map solves the Monge problem.In general, the vector fields differ; the Benamou-Brenier trajectories are straight, unlike those of the probability-flow ODE.Source: A BENAMOU-BRENIER FORMULATION OF… · p. 13

B PROOF OF THEOREM 3.1

The proof solves the probability-flow ODE for Gaussian distributions by tracking their evolving means and covariances. Its resulting encoder expression is the optimal transport map between the initial distribution and a standard normal target.

B PROOF OF THEOREM 3.1

  • Gaussian intermediate densities remain Gaussian, with ODEs describing how their means and covariances transform.The proof uses this structure to obtain the encoder map.Source: B PROOF OF THEOREM 3.1 · p. 13
  • For the zero-mean case, spectral integration gives a limiting encoder Eµ0(x) = Σ(0)^-1/2x, the optimal transport map to N(0,I).The scalar spectral factor converges to λ^-1/2.Source: B PROOF OF THEOREM 3.1 · pp. 13–14
  • For the general case, the proof treats a(t) as nonzero through an inhomogeneous ODE and combines its solution with the homogeneous solution.The mean and covariance evolve as Σ(t) = I + e^-2t(Σ(0) - I) and a(t) = e^-ta(0).Source: B PROOF OF THEOREM 3.1 · p. 14
  • Combining this piece with the limit of the general solution gives exactly the desired optimal transport map between µ0 and N(0,1).The expression is also familiar from the commonly used Fréchet Inception Distance metric, where OT distance is used to evaluate generative models (Heusel et al., 2017).Source: B PROOF OF THEOREM 3.1 · p. 14

C TENSOR-TRAIN SOLVER FOR THE FOKKER-PLANCK EQUATION

The tensor-train solver advances the Fokker–Planck equation with second-order operator splitting. It solves diffusion and convection separately, represents the evolving density in low-rank tensor-train format, and then uses those tensors to solve the probability-flow ODE.

C TENSOR-TRAIN SOLVER FOR THE FOKKER-PLANCK EQUATION

Loading 2202.07477v2…