Geometry of Score Based Generative Models

Score-based (or Diffusion) models are a new type of generative models in the field of computer vision and machine learning, achieving state-of-the-art results in image synthesis (Dhariwal & Nichol, 2021) and log likelihood (Kingma et al., 2021). They have recently gained popularity due to interesting applications such as text to image generation (DALL-E (Ramesh et al., 2022), (Rombach et al., 2022) and Imagen (Saharia et al., 2022)), image super-resolution, image editing (Meng et al., 2022), etc. Score-based generative models have enjoyed diverse perspectives from different fields. Originally, diffusion models (DDPM) were developed from expected lower bound (ELBO) maximization on data log likelihood (Ho et al., 2020). Song & Ermon (2019) showed that we can learn gradient of log likelihood (called score functions) and use it to generate images. Song et al. (2021) showed that the epsilon function in DDPM is in fact the scaled version of the score function. They further generalized these models to a continuous setting as stochastic differential equations (Song et al., 2021). More recent works have connected score-based generative models with Schrodinger bridge problem (De Bortoli et al., 2021) and control theoretic perspectives (Chen et al., 2022; Huang et al., 2021).

In this work, we present a completely different view-point on score-based generative models: geometric perspective. To the best of our knowledge, we are the first to explore the geometric connection of these generative models. Applying the solid mathematical framework in the area of Wasserstein gradient flow (Jordan et al., 1998; Ambrosio et al., 2005; Wibisono, 2018; Salim et al., 2020; Korba et al., 2020), we show that the forward and backward process of adding noise and generating image from the noise are in fact equivalent to moving on a gradient-flow-path in a metric space of probability distributions following the Wasserstein gradient flow equation.

While our understanding of score-based generative models has matured over time, few important questions remain unanswered. For example, why is it a good idea to choose forward and reverse variance the same? Can we choose reverse variance differently? Are score-based generative models same as the energy based models (Xie et al., 2016; Gao et al., 2020; Du et al., 2021)? Furthermore, new models have been proposed like Wavefit (Koizumi et al., 2022), which tries to generalize the diffusion sampling to a proximal gradient type of update. How can we explain this type of algorithms? In this work, we demonstrate that the geometric connection investigated in this work helps to answer these questions from a geometric point of view.

In addition to conceptual advantages and novel perspectives, geometric framework enables us to design practical algorithms with faster sampling capability. Score-based generative models work remarkably well when the number of sampling-steps is large (i.e. small step-size). However, the sampling time is also large for such finer schemes. As we decrease the number of sampling steps, the samples move away from the gradient-flow-path incurring error in each step, and resulting into high overall error. To minimize such error and achieve high quality samples even with small number of sampling steps, we propose to project back intermediate samples to the gradient-flow-path after every step. To achieve this, we propose an efficient estimation of Wasserstein gradient to descend towards the flow-path. As demonstrated in the result section, our proposed method significantly reduces error for smaller number of sampling steps. Below we summarize our contributions. All complete proofs are included in the appendix.

1. 1.

   To the best of our knowledge, this is the first work to theoretically prove the connection between score-based generative models and Wasserstein gradient flow. We establish this relationship through Theorems 1 and 2.

2. 2.

   This connection sheds light on several interesting questions: 1) the reverse variance in score-based generative models, 2) the connection between score-based model and energy based model, and 3) the use of proximal gradient algorithms as proposed in recent works.

3. 3.

   Based on these insights, we propose a new algorithm which generalizes the score-based model and allows for significantly faster sampling, which would otherwise be very difficult to achieve. To achieve this, we also propose an efficient Wasserstein gradient estimation algorithm.

Early works on diffusion models were based on matching the forward and reverse joint distributions through bounds on log likelihood (Ho et al., 2020; Sohl-Dickstein et al., 2015). (Song & Ermon, 2019) proposed a score-based generative model motivating from the Langevin dynamics and estimating the score function. Later (Song et al., 2021) showed that the two approaches are actually equivalent and it can be generalized further in continuous time setting through the stochastic differential equations. On the more theoretical directions, score-based optimization has been shown to be equivalent to likelihood maximization through Feynman-Kac theorem (Chen et al., 2022; Huang et al., 2021). Other notable works are interpreting the forward diffusion and generation as solving the Schrodinger Bridge problem (De Bortoli et al., 2021). Many approaches have been proposed to speed up the sampling process through clever ways to solve differential equations (Lu et al., 2022).

In their seminal work, Jordan, Kinderlehrer, Otto (JKO) proved the connection between Wasserstein gradient flow and the the diffusion systems guided by Fokker-Planck equations (Jordan et al., 1998). This result has been vastly generalized and formalized by (Villani, 2003, 2009) and (Ambrosio et al., 2005) giving birth to the theory of Wasserstein gradient flow and optimization on space of probability measures. Several notable works have followed in machine learning (Wibisono, 2018), (Korba et al., 2020), (Salim et al., 2020), for example for sampling, generative models, etc.

Instead of taking the velocity vector to be negative of Wasserstein gradient, we consider an accelerated flow where at any time, $t$, the velocity is equal to the negative Wasserstein gradient scaled by a time-varying $\beta_{t}$.

Using this accelerated Wasserstein Gradient flow, we can establish a connection with the forward process in score-based generative model. We start from the Fokker-Planck equation corresponding to the stochastic differential equation of the forward diffusion process given by eq.(6):

where the initial measure is $\mu_{0}=\rho$. Following this SDE, we know that it will end up in the final measure, $\mu_{T}=\pi$. Next theorem shows that the forward Fokker Planck equation and the accelerated Wasserstein Gradient descent are equivalent.

This implies that the forward diffusion process considered in the diffusion generative model, DDPM (Ho et al., 2020; Song et al., 2021) can be equivalently thought as an accelerated Wasserstein gradient flow starting from an initial measure $\mu_{0}=\rho$ corresponding to the data distribution and following the negative gradient towards the target measure $\mu_{T}=\mathcal{N}(0,I)$.

Next we investigate the geometric interpretation of the generation process or the reverse SDE.

Next theorem establishes the equivalence between the reverse SDE and the Wasserstein Gradient flow.

The reverse SDE or the score-based model is trying to reverse the forward process by tracing the path followed in the forward process in the opposite direction. One important implication of this theorem is that since we are moving towards the target measure $\mu_{T}=\mathcal{N}(0,I)$ in the forward process, the reverse is actually simply moving away from $\mathcal{N}(0,I)$, which is realized as the accelerated Wasserstein gradient flow with the functional $-\text{KL}(.||\pi)$. The gradient flow path with constant velocity is the geodesic. Since we are considering gradient flow path with acceleration, it is not exactly the geodesic, but similar path traced by gradient flow. We will call it gradient-flow-path in rest of the paper.

We have shown that both the forward and reverse diffusion process involved in score-based generative models are gradient flows on the space of probability measures. We gain geometric insights because of this geometric interpretation.

Once the connection between the gradient flow and score-based generative model is established, we can interpret the generation as a process walking on the gradient-flow-path. If we follow the flow-path with small steps, we can reliably reach the initial data distribution, as demonstrated by the success of the score-based generative models and diffusion models. However, this is not a great idea if we increase the step size. The score-based increment is a linear approximation and therefore it accrues more error as we increase the step size. It has been experimentally observed that the samples get poorer as we increase the step size in score-based models (Dhariwal & Nichol, 2021; Ho et al., 2020; Song et al., 2021). From a geometric point-of-view, we are taking Wasserstein gradient steps using forward-backward strategy. While this strategy works well when the step size is small, it converges to a biased measure for large step-size. Bias associated with the forward-backward strategy for large step size has been studied in the context of Wasserstein gradient flow (Wibisono, 2018). In our case, this issue is further exacerbated by the fact that the functional $-\text{KL}(\mu||\pi)$ we are trying to minimize is actually concave with respect to $\mu$.

To mitigate this issue, we propose an intuitive and geometric idea: projection. As shown in Fig.3, as we try to sample in score-based models with large step-size, the error gets large and the trajectory deviates away from the gradient-flow-path. We propose to resolve this problem by projecting again to the gradient-flow-path before taking another step.

Score-based generative model first trains a score model, $s$ such that $s_{\theta}^{*}(x_{\tau},\tau)=\nabla\log\mu_{\tau}(x_{\tau})$ using score matching strategy. Once, the score model is trained, the discretized Eurler-Maruyama step (eq.(9)) is used for generation of samples, where score function, $s^{*}$ replaces $\nabla\log\mu_{\tau}(x_{\tau})$:

It can also be interpreted as predict and diffuse steps, where predict step is $x_{\tau+\delta\tau}^{pred}=x_{\tau}+(2\beta_{\tau}s_{\theta}^{*}(x_{\tau},\tau)-\beta_{\tau}\nabla\log\pi(x_{\tau}))\delta\tau$. Since generation is trying to trace the gradient-flow-path, after each of these predict-diffuse, we should obtain samples from the measure $\mu_{\tau+\delta\tau}$ on the gradient-flow-path. Because of discretization error and bias, these samples do not lie on the gradient-flow-path, in fact they deviate away. To pull these samples towards the measure, $\mu_{\tau+\delta\tau}$ on the gradient-flow-path, we use the fact that we have a way to sample from the measures on gradient-flow-path. Using the SDE equation, we can write closed form conditional distribution of $\mu_{\tau+\delta\tau}$ as follows:

Sampling from this conditional distribution is given by the following equation:

Comparing eq.(30) with eq.(29), we note that pulling $x_{\tau+\delta\tau}^{pred}$ close to $x_{\tau+\delta\tau}^{mean}$ may be enough to pull the samples $x_{\tau+\delta\tau}$ in eq.(29) towards gradient-flow-path assuming that $\beta_{\tau+\delta\tau}$ is close to $\beta_{\tau}$. In terms of measures, we consider the measure associated with samples $x_{\tau+\delta\tau}^{pred}$ and target samples $x_{\tau+\delta\tau}^{mean}$. Let’s define the measure corresponding to samples $x^{pred}_{\tau+\delta\tau}$ as $\mu^{pred}_{\tau+\delta\tau}$ and the measures corresponding to means, $x_{\tau+\delta\tau}^{mean}$ as $\mu^{mean}_{\tau+\delta\tau}$. We can sample from the measure $\mu^{mean}_{\tau+\delta\tau}$ by first sampling from $x_{\tau=T}\sim\mu_{\tau=T}$, and passing them through $x_{\tau+\delta\tau}^{mean}$. Our strategy to project the samples in eq.(29) onto gradient-flow-path is to project the pred measure, $\mu^{pred}_{\tau+\delta\tau}$ to the mean measure $\mu^{mean}_{\tau+\delta\tau}$. We achieve this through Wasserstein gradient descent in the space of probability measures. For that, we need an efficient way to estimate Wasserstein gradient, which we describe in next subsection.

To demonstrate efficacy of our algorithm, we train and generate samples on three datasets: 1) Celeb-A dataset, 2) LSUN-church dataset, and 3) SVHN dataset, where all images are of size $64\times 64$. Our neural network architecture for both score model and projection model uses standard U-net architecture with attention due to (Dhariwal & Nichol, 2021). We use publicly available code from (Song et al., 2021) as score-based model. In these experiments, we demonstrate that as we decrease the number of sampling steps, the quality of samples decreases in Score-based generative model, but we maintain quality to a reasonable level even when the number of sampling steps is reduced to as low as 20. We use FID metric (Heusel et al., 2017) to measure the sample qualities.

In Table 1, we compare the FID score of generated images from score-based method and our Predict-Project method. We outperform the score-based method by a large margin in all cases except SVHN (N=100). This underperformance could be because our model is not trained well in SVHN (see appendix) due to lack of time. For qualitative comparison, please see Fig. (4) and Fig. (2). These results our claim that projecting to the gradient-flow-path improves sample quality, especially when the number of sampling-step is low.

We presented a novel geometric perspective on score-based generative models (also called diffusion generative models) by showing that they are in fact gradient flows in a space of probability measures. The geometric insight gained from this connection helped us answer and clarify some critical open questions. We also demonstrated that it can help us design faster sampling algorithm. We believe that this connection will help diffuse knowledge between Wasserstein gradient flow field and score-based generative models field in the future inspiring interesting solutions to problems in both areas. Similarly, connection with energy-based models, proximal algorithms and reverse SDE could help design better algorithms in general and generative models in specific. Energy-based models, for example, can be combined with score-based models in the light of geometric understanding.