Regularizing Score-based Models with Score Fokker-Planck Equations
Abstract
Score-based generative models learn a family of noise-conditional score functions corresponding to the data density perturbed with increasingly large amounts of noise.
These perturbed data densities are tied together by the Fokker-Planck equation (FPE), a PDE governing the spatial-temporal evolution of a density undergoing a diffusion process.
In this work, we derive a corresponding equation characterizing the noise-conditional scores of the perturbed data densities (i.e., their gradients), termed the score FPE .
Surprisingly, despite impressive empirical performance, we observe that scores learned via denoising score matching (DSM) do not satisfy the underlying score FPE.
We mathematically analyze two implications of satisfying the score FPE and a potential explanation for why the score FPE is not satisfied in practice.
At last, we propose to regularize the DSM objective to enforce satisfaction of the score FPE, and show its effectiveness on synthetic data and MNIST.
Score-based generative models
[16] unifies denoising score matching [14] and diffusion probabilistic models [13, 4] via a stochastic process x ( t ) of continuous time t ∈ [0 , T ] driven by the forward SDE
[Equation]
where f ( · , t ): R D → R D , g ( · ): R → R and w t is a standard Wiener process.
Under some moderate conditions [1], one can obtain a reverse time SDE from T to[-
-]
[Equation]
where ¯ w t is a standard Wiener process in reverse time.
Let q t ( x ) denote the ground truth marginal density of x ( t ) following Eq.
We can train a neural network s θ =
s θ ( x , t ) to approximate ∇ x log q t ( x ) by minimizing the denoising score matching (DSM) loss [18, 16]:
[Equation]
[Figure]
where q 0 t ( x ( t ) | x (0)) is the transition kernel from x (0) to x ( t ) .
After s θ ( x , t ) ≈ ∇ x log q t ( x ) is learned, we replace ∇ x log q t ( x ) in Eq.
2 with s θ and get a parametrized reverse-time SDE for stochastic process ˆ x θ ( t )
[Equation]
Let p SDE t, θ denote the marginal distribution of ˆ x θ ( t ) .
We can design f and g in Eq.
2 so that q T ( x ) approximates a simple prior π , and hence, can generate samples ˆ x θ (0) ∼ p SDE 0 , θ by numerically solving Eq.
4 backward with an initial sample from the prior ˆ x θ ( T ) ∼ π .
Intuitively, ˆ x θ (0) should be close to a sample from the data distribution.
The Fokker-Planck equation for a score vector field
It is well known that the evolution of the ground truth density q t ( x ) associated to Eq.
1 is governed by the Fokker-Planck equation (FPE) [10] (details in Appx.
E).
As there is a one-to-one mapping between densities and their scores, we can derive an equivalent system of PDEs that the ground truth scores ∇ x log q t ( x ) must satisfy.
We call it as a score Fokker-Planck equation , for short score FPE .
Corollary 1 (score FPE) .
Assume that the ground truth density q t ( x ) is sufficiently smooth for ( x , t ) ∈ R D × [0 , T ] .
Then its score s ( x , t ) :
= ∇ x log q t ( x ) satisfies the following system of PDEs
[Equation]
This result shows that the time-conditional scores s θ ( x , t ) learned by score-based models (via Eq. 3) are highly redundant.
In principle, given a ground truth score at an initial time t 0 , we can theoretically recover scores for all times t ≥ t 0 by solving the score FPE.
We explain its intuition by considering a special case when f ≡ 0 and g ≡ 1 .
That is, x ( t ) is obtained by adding Gaussian noise.
It is well-known that the densities q t and q t 0 are related in a convolutional way as q t = q t 0 ∗ N (0 , t ) , and q t can be analytically obtained from q t [-
 [MR92]-]0 [8] (e.g., by applying a Fourier transform and dividing).
Hence, all scores can in principle be obtained analytically from the score at a single time-step, without any further learning.
In Appx.
B we empirically support this idea.
Theoretically, with sufficient data and model capacity, (denoising) score matching ensures the optimal solution to Eq.
3 should satisfy Eq.
5 as it approximates the ground truth score well.
However, we observe that pre-trained s θ learned via Eq.
3 do not numerically satisfy the score FPE.
We hereby introduce an error term ϵ s θ =
ϵ s θ ( x , t ) in order to quantify how s θ deviates from the score FPE
[Equation]
We further define the following averaged residuals of DSM and the score FPE for t ∈ [0 , 1] :
[Equation]
Fig.
1 plots these residuals for score models pre-trained via DSM on MNIST and CIFAR-10.
Despite achieving low r DSM-like score-matching loss across all t (green curve), pre-trained score models fail to satisfy the score FPE equation especially for small t (blue, orange curves).
In Sec. 3, we theoretically analyze these findings and provide new insights into the score FPE.
In Sec. 4, we propose a novel score matching objective with score FPE as a regularizer (Eq.
and examine its effectiveness on synthetic dataset and MNIST.
We refer to Appx.
C for implementation details and Appx.
E for proofs.
Theoretical implications and interpretations of score FPE
In this section, we first study two implications of satisfying the score FPE.
More precisely, we show in Sec. 3.1 that controlling ϵ s θ can implicitly enforce conservativity of s θ .
Moreover, if the score FPE is satisfied, we prove in Sec.
3.2 the equivalence of s θ , ground truth score s and ∇ x log p SDE t, θ holds under some conditions, where p SDE t, θ is defined in Sec.
1 as the marginal density of parametrized diffusion process.
In Sec. 3.3, we investigate the connection between higher-order score matching [9, 7] and score FPE.
3.1 Conservativity
The ground truth score s ( x , t ) = ∇ x log q t ( x ) is a conservative vector field.
That is, it can be expressed as a gradient of some real-valued function.
However, scores learned in practice do not satisfy this property [12].
Below we prove that we can implicitly enforce conservativity by minimizing the time-averaged residual of the score FPE.
Proposition 1.
If there is a t θ ∈ [0 , T ] so that s θ ( x , t θ ) =
∇ x log q t θ ( x ) for all x ∈ R D , then there is a real-valued function Ψ θ : R D × [0 , T ] → R given by Ψ θ ( x , t )
= log q t θ ( x ) + (cid:82) t t θ (cid:104) 12 g 2 ( τ ) div x
( s θ ) + 12 g 2 ( τ ) ∥ s θ ∥ 22 − ⟨ f , s θ ⟩ − div x ( f ) (cid:105) dτ so that for all ( x , t ) ∈ R D × [0 , T ]
[Equation]
Consider a model s θ , and assume that for a large enough timestep t θ , it captures exactly the perturbed density ( s θ ( x , t θ )
=
∇ x log q t θ ( x ) ) which is close to the prior (normal distribution) because t θ ≈ T .
Intuitively, s θ is “nearly” conservative as the score of the prior is known to be conservative (because its score has a closed form as the gradient of log-density).
Indeed, Prop .1 supports the intuition by saying that the the estimated score should nearly be conservative if it approximately satisfies the score
FPE.
Actually, Fig.
1 shows that (cid:82) t t θ (cid:13)(cid:13)
[-ϵ s θ 
-][Equation]
2 dτ is numerically small.
Therefore, s θ ( x , t ) is close to the gradient of a scalar function Ψ θ ( x , t ) ; namely, it is conservative.
3.2 Equivalence between s θ , s and ∇ x log p SDE t, θ
We now investigate another implication of satisfying the score FPE which connects the score s θ with the ground truth s and ∇ x log p SDE t, θ .
The following proposition states that all aforementioned scores are identical if we train to reach a zero residual of score FPE for all ( x , t ) (under some technical assumptions ensuring the system of PDEs has a unique solution).
Proposition 2.
Suppose we know that in some suitable function space, 0 is the unique strong solution to the PDEs
∂ t v − ∇ x (cid:2) 12 g 2 ( t ) div x ( v ) + 12 g 2 ( t ) (cid:0) ∥
v ∥ 22 + 2 ⟨ v , s ⟩ (cid:1) − ⟨ f , v ⟩ (cid:3) =
with zero initial condition v ( x , 0) ≡ 0 and zero boundary condition.
If there is some θ 0 so that ϵ s θ 0 ( x , t ) = 0 for all ( x , t ) , then s θ 0 ≡ s .
Moreover, suppose that the PDEs ∂ t v + ∇ x (cid:2) 12 g 2 ( t ) div x ( v ) + 12 g 2 ( t ) ∥ v ∥ 22 + ⟨ f , v ⟩ (cid:3) = 0 with zero initial and boundary condition has 0 as the unique strong solution, then ϵ s θ 0 ≡ 0 implies s θ 0 ≡ ∇ x log p SDE t, θ 0 .
We hypothesize that satisfying the score FPE has a smoothing effect when f ( x , t ) is linear in x .
Suppose the assumptions of Prop. 2 hold and hence, s θ 0 ≡ ∇ x log p SDE t, θ 0 .
As linearly transforming normal distributions (by f ) remains normal, [7] proves that p SDE t, θ 0 turns out to be a Gaussian distribution for any t .
In practice, the assumptions are not likely to be met exactly, i.e. the residual will not be exactly zero ϵ s θ 0 ≡ 0 .
In this case, we hypothesize that learning θ to reduce ∥ ϵ s θ ∥ can reduce the gap (cid:13)(cid:13)(cid:13) s θ − ∇ x log p SDE t, θ (cid:13)(cid:13)(cid:13) , and may further modify the direction s θ toward high density region of Gaussian ( smoothing effect ). [Figure] [Figure]
[Table]
3.3 Higher-order score matching
Higher-order derivatives of score can yield additional information about the data distribution.
We prove a property stating that error bounds of higher-order score matching can further control the residual (cid:13)(cid:13)(cid:13)(cid:82) t
ϵ s θ ( x , τ ) dτ (cid:13)(cid:13)(cid:13) 2 for all t ∈ [0 , T ] .
This may explain why the scores learned via Eq.
3 are not sufficient to satisfy the score FPE as their higher-order scores may deviate from the ground truth.
Proposition 3.
Denote C ( t ) :
= 1 2 (cid:82) t 0 g 2 ( τ ) dτ .
Assume the following error estimates hold for higher-order score matching that for all ( x , t ) ∈ R D × [0 , T ] [Equation]
[Equation]
Then we have that for all ( x , t ) ∈ R D × [0 , T ]
[Equation]
Experimental results
We have demonstrated that score models learned via J DSM (Eq. 3) do not satisfy the score FPE, a property that ground truth scores should satisfy a priori .
Therefore, we devise a novel loss function, consisting of J DSM and a score FPE-regularizer R FP ( θ ) :
[Equation] 1 D [Equation] [Equation]
[Equation]
where γ ≥ 0 is a hyper-parameter.
Since ϵ s θ in R FP is generally expensive to calculate, we propose to exploit the finite difference method [3] for ∂ t s θ and Hutchinson’s trace estimator [5] for div x ( s θ ) to reduce the computational efforts (details in Appx. D).
The effectiveness of J FP is examined on synthetic dataset (Gaussian mixture models) and MNIST.
Synthetic dataset We consider a Gaussian mixture model as the training data distribution.
Fig.illustrates (a) ground truth density, and the density produced by probability flow ODE [16] of scores trained with (b) λ = 0 .
(i.e, conventional score matching training) and (c) λ
= 0 .
The score trained with score FPE-regularizer can approximate the data density well, improving over vanilla score-matching.
We hypothesize score FPE-regularizer may improve density estimation with the probability flow ODE, as it enforces a known self-consistency property of the ground truth score.
MNIST We evaluate the proposed J FP ( θ ; λ ( · ) , γ ) on MNIST with different γ ’s. Table 1 reports negative log-likelihood (NLL) in bits/dim (bpd) across three instantiations of the forward SDE (see
Appx.
A) and two choices of weighting functions λ ( · ) ’s ([16] and [15]).
We observe a general improvement in NLL with γ = 1 .
(see Appx.
F for a demonstration of generated samples).
Conclusion
We introduce the score FPE and theoretically study its relation with score matching, conservativity and density induced by parametric reverse diffusion.
Moreover, we propose to penalize on residual of score FPE and show its effectiveness on simple dataset.
However, it is unclear how the dynamics of score FPE affects, for instance, training of a larger scale dataset or variational lower bound.
A Instantiation of SDE and score FPE
[16] categorizes the forward SDE into three types based on the behavior of the variance during evolution.
Here we focus on two of them, which are Variance Explosion (VE) SDE and Variance
Preserving (VP) SDE.
VE SDE It has a zero drift term f = 0 and diffusion term g ( t ) = (cid:113) dσ 2 ( t ) dt with some function σ ( t ) .
Hence, the forward SDE (Eq. 1) becomes
[Equation]
A typical instance of VE SDE is Score Matching of Langevin dynamics (SMLD) [14], where σ ( t ) :
= σ min (cid:16) σ max σ min (cid:17) t for t ∈ (0 , 1] .
In our implementation, we follow the conventional setup of ( σ min , σ max ) :
= (0 . 01 , 50) .
In [6], they proposed a variant of VE SDE attempting to resolve the unbounded score problem [2], which is called Reciprocal VE (RVE).
Let ϵ > 0 be a fixed constant.
RVE SDE also has zero drift term but with a different parametrization for diffusion
[Equation]
VP SDE Let β be a non-negative function of t .
VP SDE has a linear drift term f ( x , t ) =
− 12 β ( t ) x
and diffusion term g ( t ) =
(cid:112) β ( t ) .
Thus, the forward SDE is
[Equation]
A classic example of VP SDE is Denoising Diffusion Probabilistic Modeling (DDPM)
[13, 4], where β ( t ) :
= β min + t ( β max − β min ) for t ∈ [0 , 1] .
We adopt the common setup of ( β min , β max ) :
= (0 . 1 , 20) in our implementation.
We summarize the aforementioned instantiations of SDE and their associated score FPE in Table 2.
[Table]
B How scores satisfy score FPE?
We experimentally demonstrate how score functions should satisfy the score FPE in two different aspects.
We consider the data distribution as a Gaussian mixture model (GMM) of the density 15 N (cid:0) ( − 5 , − 5) , I (cid:1) + 45 N (cid:0) (5 , 5) , I (cid:1) on R 2 whose samples are illustrated in Fig. 3a.
The diffusion process is taken as VE SDE (Eq. 9).
The ground truth score of GMM, denoted as s GMM , can be expressed explicitly with a closed formula throughout the diffusion (as the diffusion process is linear in x ).
First of all [Table] we examine if s GMM satisfies the score FPE by computing r FP ( t ; s GMM ) and plot it in
Fig.
5a (blue curve).
We can see that the score FPE residual of the ground truth is almost zero, which empirically supports Corollary 1.
Second, as we explain in Sec. 2, we can solve score FPE for the score at any time if we are merely given a score at a single time moment.
Namely, once we find a solution ˜ s
to the following initial value problem of system of PDEs, we know a score at all time.
[Equation]
Solutions of Eq.
10 can be parametrized as neural network ˜ s GMM θ [11].
We then can solve the PDEs by learning parameters θ to reduce both the residuals of the initial condition ˜ s GMM θ
[Equation]
and evolution ϵ ˜ s GMM θ :
[Equation]
t [Equation] s θ ∥
[Equation]
That is, ˜ s GMM θ is learned with the score FPE-guided objective function:
[Equation] [Equation]
We demonstrate generated samples by the learnt ˜ s GMM θ in Fig.
3b and plot its score FPE residual r FP ( t ; ˜ s GMM θ ) in Fig.
5a (orange curve).
Interestingly, it also generates quite satisfactory samples and
Fig.
4b show it estimates the ground truth score (Fig. 4a) well.
This supports our argument.
On the other hand, we compare with a score s GMM θ learned from the denoising score matching (Eq. 3).
We observe that from Fig.
5b that s GMM θ does not satisfy score FPE even though it works decently on generation (Fig. 3c) and score estimation (Fig. 4c).
[Figure] [Figure]
and is guided
[Figure]
C Explanation of Implementation
In Fig.
1, we train a score on MNIST for 200 epochs with a learning rate 1 e − 3 and batch size 32 by using an identical neural network structure to the repository 1 but modify the forward SDE as VE
SDE or VP SDE (see Appx.
A).
The network structure of synthetic dataset in Appx.
B is similar to the aforementioned one but we simply replace all convolutional layers with fully connected layers.
We train for 2 , 000 epochs with a learning rate 1 e − 3 and batch size 500 .
For the case of CIFAR10, we use the pre-trained score models provided by [16] 2 instead of training them from scratch.
The VE SDE and VP SDE are taken as NCSN++ cont.
and DDPM++ cont., respectively.
The neural network setup in Fig.
2 is the same as toy model structures provided in the repository of [7] 3 .
We found out setting the weight of score FPE to λ = 0 .
001 can generally work well for the toy dataset.
D Techniques to reduce computation costs for score FPE
The computation of ϵ s θ ( x , t ) in R FP ( θ ) is expensive and hard to scale-up to higher dimensional data.
We thus propose two techniques which can help reduce the computation costs.
We recall that ϵ s θ ( x , t ) is defined as ∂ t s θ −∇ x (cid:104) 12 g 2 ( t )
div x ( s θ )
+ 12 g 2 ( t )
∥ s θ ∥ 22 −⟨ f , s θ ⟩− div x ( f ) (cid:105) .
We explain the details of how the computation of ∂ t s θ and div x
( s θ ) can be respectively relieved by the finite difference method and Hutchinson’s trace estimator.
D. ∂ t s θ term
Lemma 1.
[3] Let α : [0 , 1] → R D be a vector-valued function which is continuously differentiable up to third order derivatives.
Denote h s and h d as hyper-parameters of step sizes.
Then we have the following estimate of α ′ ( t ) :
[Equation]
In particular, if h s = h d =: h , then the estimate becomes
[Equation]
In our implementation for MNIST, we consider α ( · ) :
= s θ ( · , x ) and set ( h s , h d ) = (0 . 001 , 0 . 0005) for the approximation of ∂ t s θ .
D. div x ( s θ ) term
Hutchinson’s trace estimator [5] stochastically estimates the trace tr ( A ) of any square matrix A .
Its idea is choose a distribution p v so that E v ∼ p v [ v ] = 0 and E v ∼ p v [ vv T ]
= I .
Hence, tr ( A ) = tr ( A E v ∼ p v [ vv T ]) =
E v ∼ p v [ tr ( Avv T )]
=
E v ∼ p v [ tr ( vAv T )]
=
E v ∼ p v [ vAv T ] .
By i.i.d.
sampling { v j } Mj =1 from p v , we can use the following unbiased estimator
[Equation]
to estimate tr ( A ) .
We notice that div x ( s θ ( x , t ))
=
tr (cid:0) ∇ x s θ (cid:1) .
Thus, we can apply Hutchinson’s trick and replace div x ( s θ ) term with the estimation
[Equation]
In the implementation, p v is usually taken as a standard normal distribution or a Rademacher distribution whose random vector has components equal to +1 or − 1 with equal probability.
We follow the convention in [17] which sets M = 1 and shows its effectiveness.
D. Projection of ϵ s θ ( x , t )
We further propose another potential trick to reduce the computation cost of differentiation.
[Equation]
Using automatic differentiation to compute the gradient in ϵ s θ ( x , t ) (the (II) part in Eq. 12) is generally cumbersome for high dimensional data.
We propose to use random projection to relieve the computation of gradient (multi-dimension) to directional derivative (one-dimensional).
Thus, we can further apply the finite difference trick introduced in Appx.
D.1 to reduce the computation efforts.
We first recall a fundamental property before rigorously formulating the technique.
Lemma 2.
Let M :=
M ( x , t ): R D × [0 , T ] → R be a continuously differentiable function of x .
For any v ∈ R D ,
[Equation]
where D v M ( x , t ) means the directional derivative of M in x along the direction v which is defined as:
[Equation]
For simplicity, let us denote M ( x , t ) :
= 12 g 2 ( t ) div x
[Equation]
s [Equation] g [Equation]
t ) ∥ s θ ∥ 22 − ⟨ f , s θ ⟩ − div x ( f ) and let v ∈ R D be arbitrary vector.
We project ϵ s θ ( x , t ) along direction v and apply the Lemma 2,
[Equation]
Notice that both ∂ t s θ and ddh M ( x + h v , t ) (cid:12)(cid:12)(cid:12) h =0 are one-dimensional differentiation, which can be estimated via Lemma 1 and hence, we can avoid automatic differentiation.
We hereafter propose an estimated score FPE-regularizer which may replace R FP with
[Equation]
where p v is a distribution of random vector v ∈ R D .
We observe that the performance may degrade by using ˆ R FP , which may due to the inaccurate approximation to the exact score FPE.
Therefore, a further study is required to have lower computation costs while preventing the deterioration in the performance.
E Proofs and discussion E.1 Proof of Corollary
We prove the result with a more general forward SDE
[Equation]
15) where F ( · , t ): R D
→ R D and G ( · , t ): R D → R D × D .
We know that the density q t ( x ) satisfies the Fokker-Planck equation [10]
[Equation]
16) where ˜ F ( x , t ) :
= F ( x , t ) − 12 ∇ · [ G ( x , t ) G ( x , t ) T ] − 12 G ( x , t ) G ( x , t ) T ∇ x log q t ( x ) .
We further denote A ( x , t ) :
= F ( x , t ) − 12 ∇ · [ G ( x , t ) G ( x , t ) T ] and B ( x , t ) :
= − 12 G ( x , t ) G ( x , t ) T .
Now ˜ F ( x , t ) =
A ( x , t ) + B ( x , t ) s ( x , t ) , and we have
[Equation]
[Equation] [Equation]
Since log q t ( x ) is sufficiently smooth
[Equation]
we can swap the order of differentiations and get ∂ t s = ∂ t ∇ x log q t ( x ) = ∇ x ∂ t log q t ( x ) .
Hence, the statement is proved.
■
Remark 1.
In Eq.
1 where G does not depend on x , namely G ( x , t ) ≡ g ( t ) I , then ˜ F ( x , t )
=
f ( x , t ) − 12 g 2 ( t ) ∇ x log q t ( x ) and
[Equation]
E.2 Proof of Proposition
Integrating the following equation w.r
[Equation]
t. time from τ = t θ to τ = t with t ∈ [0 , T ] fixed,
[Equation]
leads to
[Equation]
where the swap of integration and differentiation is valid if the integrand is sufficiently smooth.
With the assumption [Equation] we obtain that for all t ∈ [0 , T ]
[Equation]
By taking the norm of the above equation [Equation] one can obtain
[Equation]
From which we obtain
[Equation]
The upper and lower bound of integral can be respectively written as max { t, t θ } and min { t, t θ } [Equation] and whence, the proposition is proved.
■
E.3 Proof of Proposition
Lemma 3.
Let s θ be a score obtained from denoising score matching (Eq. 3) and write ˆ s θ := ∇ x log p SDE t, θ .
Then
[7] Eq.
4 associates with the following forward SDE whose marginal density is ˆ s θ :
[Equation]
ˆ s θ satisfies the following score FPE:
[Equation]
Proof of Lemma 3.
Consider
[Equation]
and
[Equation]
in Eq.
[Equation]
and apply Corollary 1, the lemma is then established.
■
Proof of Proposition 2.
We recall Eq. 6, which indicates
[Equation]
First [Equation] we subtract Eq. 27 by the above equation and get
[Equation]
[Equation]
Consider when θ = θ 0 and let u θ 0 :
=
ˆ s θ 0 − s θ 0 .
Then the PDEs become
[Equation]
Here, u θ 0 is a solution to the PDEs.
It is noticed that this system of PDEs has a zero initial condition and zero boundary condition as both s θ 0 and ˆ s θ 0 share the same initial/boundary condition.
Thus, from the assumption of the uniqueness of solution, we know that u θ 0 ≡ 0 , and hence, ˆ s θ 0 ≡ s θ 0 .
We repeat the same trick to subtract Eq.
5 by Eq.
28 from which we can obtain s θ 0 ≡ s .
■
E.4 Proof of Proposition
By subtracting the following two equations
[Equation]
we obtain
[Equation]
Notice that ∥ s θ ∥ 22 − ∥ s ∥ 22 = ∥
s θ − s ∥ 22 + 2 ⟨ s θ − s , s ⟩ .
Integrating over time from τ = 0 to τ = t , we obtain
[Equation]
By applying the ℓ 2 -norm and Cauchy-Schwartz inequality while noting the relation ∥
A ∥ 2 ≤ ∥ A ∥ F for a general square matrix A , the statement is proved.
■
F Demonstration of generated MNIST examples
In Fig.
6, we show examples generated with different choices of the SDE and the weight of score
FPE-regularizer, γ , on MNIST.
[Figure]