[Paper Review] Score-based Generative Modeling Secretly Minimizes the Wasserstein Distance

Dohyun Kwon, Ying Fan+1 more|arXiv (Cornell University)|Dec 13, 2022

Generative Adversarial Networks and Image Synthesis4 citations

TL;DR

This paper establishes that score-based generative models secretly minimize the Wasserstein distance between the generated and data distributions by proving an upper bound on the Wasserstein distance in terms of the score matching loss. Using optimal transport theory, it shows the Wasserstein distance is bounded by the square root of the loss up to constants and an offset, offering theoretical justification for the models' empirical success in distribution approximation.

ABSTRACT

Score-based generative models are shown to achieve remarkable empirical performances in various applications such as image generation and audio synthesis. However, a theoretical understanding of score-based diffusion models is still incomplete. Recently, Song et al. showed that the training objective of score-based generative models is equivalent to minimizing the Kullback-Leibler divergence of the generated distribution from the data distribution. In this work, we show that score-based models also minimize the Wasserstein distance between them under suitable assumptions on the model. Specifically, we prove that the Wasserstein distance is upper bounded by the square root of the objective function up to multiplicative constants and a fixed constant offset. Our proof is based on a novel application of the theory of optimal transport, which can be of independent interest to the society. Our numerical experiments support our findings. By analyzing our upper bounds, we provide a few techniques to obtain tighter upper bounds.

Motivation & Objective

• To theoretically understand why score-based generative models perform well in practice beyond just minimizing KL divergence.
• To investigate whether score-based models implicitly minimize the Wasserstein distance, despite not being explicitly trained for it.
• To establish a rigorous upper bound on the Wasserstein distance using the score matching objective function.
• To provide practical insights into tightening the theoretical upper bound through model design and training choices.

Proposed method

• The authors use optimal transport theory to analyze the time evolution of probability measures in score-based models via continuity equations.
• They derive the time derivative of the Wasserstein distance between the data and model distributions using velocity fields derived from the Fokker-Planck and reverse SDE dynamics.
• A key inequality is established: the rate of change of the Wasserstein distance is bounded by a function of the score matching loss and Lipschitz constants.
• The bound is expressed as $ -\frac{d}{dt}W_2(p_t, q_t) \leq (L_f + L_s g^2)W_2(p_t, q_t) + g^2 b^{1/2} $, where $ b(t) $ is the expected squared score error.
• Theoretical upper bounds on $ W_2(p_0, q_0) $ are derived in terms of the score matching loss $ J_{SM} $, showing it scales with $ \sqrt{J_{SM}} $.
• Numerical experiments validate the theoretical bounds and explore techniques to tighten them, such as adjusting time horizon $ T $ and noise schedule.

Experimental results

Research questions

• RQ1 Does score-based generative modeling implicitly minimize the Wasserstein distance between the generated and data distributions?
• RQ2 Can the Wasserstein distance be upper bounded using the score matching loss function, despite the model not optimizing it directly?
• RQ3 How does the choice of time horizon $ T $ and noise schedule affect the tightness of the theoretical upper bound on the Wasserstein distance?
• RQ4 What are the practical techniques to improve the tightness of the theoretical bound in real-world training?

Key findings

• The Wasserstein distance between the data and model distributions is upper bounded by the square root of the score matching loss $ J_{SM} $, up to multiplicative constants and a fixed offset.
• The bound is derived using optimal transport theory and the time derivative of the Wasserstein metric along the trajectories of the Fokker-Planck and reverse SDE processes.
• The authors prove $ J_{SM} \leq J_{DSM} $ under suitable assumptions, showing that the more practical $ J_{DSM} $ loss can be used as a proxy for bounding the Wasserstein distance.
• Numerical experiments confirm that as the score matching loss decreases, both the KL divergence and the Wasserstein distance also decrease, supporting the theoretical claim.
• The upper bound tightens significantly when the time horizon $ T $ is large (e.g., $ T \geq 100 $), with $ I(T)W_2(p_T, q_T) $ decaying almost exponentially.
• The model is shown to be robust to initial perturbations in the data distribution, as the Wasserstein distance remains bounded under such changes.