Interpreting diffusion score matching using normalizing ﬂow
Abstract
Scoring matching (SM), and its related counterpart, Stein discrepancy (SD) have achieved great success in model training and evaluations.
However, recent research shows their limitations when dealing with certain types of distributions.
One possible ﬁx is incorporating the original score matching (or Stein discrepancy) with a diffusion matrix, which is called diffusion score matching (DSM) (or diffusion Stein discrepancy (DSD)) .
However, the lack of the interpretation of the diffusion limits its usage within simple distributions and manually chosen matrix.
In this work, we plan to ﬁll this gap by interpreting the diffusion matrix using normalizing ﬂows.
Speciﬁcally, we theoretically prove that DSM (or DSD) is equivalent to the original score matching (or score matching) evaluated in the transformed space deﬁned by the normalizing ﬂow, where the diffusion matrix is the inverse of the ﬂow’s Jacobian matrix.
In addition, we also build its connection to Riemannian manifolds, and further extend it to continuous ﬂows, where the change of DSM is characterized by an ODE.
1. Introduction
Recently, score matching (Hyv ¨ arinen, 2005) and its closely related counterpart, Stein discrepancy (Gorham, 2017) have made great progress in both understanding their theoretical properties and practical usage.
Particularly, unlike Kullback–Leibler (KL) divergence which can only be used for distributions with known normalizing constant, SM (or SD) can be evaluated for unnormalized densities, and requires fewer assumptions for the probability distributions (Fisher et al., 2021).
Such useful properties enable them to be widely applied in training energy-based model (EBM)
(Song et al., 2020a; Grathwohl et al., 2020; Wenliang et al., 2019), state-of-the-art score-based generative model (Song & Ermon, 2019; Song et al., 2020b), statistical tests (Liu et al., 2016;
Chwialkowski et al., 2016) and variational inference (Hu et al., 2018; Liu & Wang, 2016).
Despite their elegant statistical properties, recent work (Barp et al., 2019) demonstrated their failure when dealing with certain type of distributions (e.g. heavy-tailed distributions).
For instance, when the data and the model are heavy tailed distributions, the model can fail to recover the true mode even in one dimensional case.
The root of this problem is that the SM (or SD) objective is highly non-convex and does not correlate well with likelihood.
To ﬁx it, Barp et al.
(2019) proposed a variant called diffusion score matching (and diffusion Stein discrepancy ), where a diffusion matrix is introduced.
However, the author did not provide us an interpretation of this diffusion matrix.
In fact, the diffusion used by the author (Barp et al., 2019) is manually chosen for toy densities.
Such lack of interpretation hinders further development of a proper training method of the diffusion matrix.
In this paper, we aim to give an interpretation based on normalizing ﬂows, which sheds light on developing training method for the diffusion.
We summarize our contributions as follows:
• We theoretically prove that DSM (or DSD) is equivalent to the original SM (or SD) performed in the transformed space deﬁned by the normalizing ﬂow.
The diffusion matrix is exactly the same as the inverse of the ﬂow’s Jacobian matrix.
• We further show that its connection to Riemannian manifold.
Speciﬁcally, we show the diffusion matrix is closely related to the Riemannian metric tensor.
• We further extend DSM to their continuous version.
Namely, we derive an ODE to characterize its instantaneous change.
We hope that by building these connections, a broad range of techniques from normalizing ﬂow communities can be leveraged to develop training methods for the diffusion matrix.
2. Background: Diffusion Stein discrepancy 2.1. Score matching and Stein discrepancy
Let P be the space of Borel probability measures on R D ,
Q ∈ P to be a probability measure, the objective for model learning is to ﬁnd a sequence of probability measures { P θ : θ ∈ Θ } ⊂ P that approximates Q in an appropriate sense.
One common way to achieve this is by deﬁning a discrepancy measure D : P × P
→ R , which quantiﬁes the differences between two probability measures.
Thus, the optimal parameters θ ∗ can be obtained by θ ∗ = argmin D
( Q || P θ ) .
The choice of discrepancy depends on the properties of the probability measures, the efﬁciency and its robustness.
The one we are focused on is called
Fisher divergence .
Assuming for probability measures Q and P θ , we have corresponding twice differentiable densities q ( x ) , p θ ( x ) .
The Fisher divergence (Johnson, 2004) is deﬁned as
[Equation]
where s p ( x ) = ∇ x log p θ ( x ) is called the score of p θ , and s q is deﬁned accordingly.
Despite that q is often used for underlying data densities with the intractable s q , s q in fact acts as a constant for parameter θ .
Thus, one can use integrationby-part to derive the following:
[Equation]
with C q a constant w.r.t. θ .
This equivalent objective SM ( q, p θ ) is referred as score matching (Hyv ¨ arinen, 2005).
Another discrepancy measure we are interested is called
Stein discrepancy , which is deﬁned as
[Equation]
where f : R D → R D is a test function, and H is an appropriate test function family, e.g. reproducing kernel Hilbert space (Liu et al., 2016;
Chwialkowski et al., 2016) or Stein class (Gorham, 2017; Liu et al., 2016).
Recent work (Hu et al., 2018) proved a connection between Stein discrepancy and Fisher divergence by showing the optimal test function:
[Equation]
thus we can show Stein discrepancy is equivalent to Fisher divergence up to a multiplicative constant.
Barp et al. ; Gorham et al. (2019) further extend the score matching and Stein discrepancy by incorporating a diffusion matrix m ( x ) : R D → R D × D . It starts from deﬁning diffusion Fisher divergence
[Equation]
where m ( x ) is a matrix-valued function. Expanding eq.and applying integration by parts (with m short-handing m ( x ) and s p short-handing s p ( x ) ):
[Equation] where C q,m depends on both q and m ( x ) . Similar to the derivation of SM ( q, p θ ) , this also returns an alternative diffusion
[Equation]
m ( q, p θ ) . Similarly, Diffusion Stein discrepancy (DSD) is deﬁned as
[Equation]
It can be shown that as long as m ( x ) is invertible, F m ( q, p θ ) and DSD m ( q, p θ ) are valid divergences.
These two extensions have demonstrated superior performances when dealing with certain type of distributions.
In the following, we give a motivating example similar to Barp et al.
2.2. Motivating example: Student-t distribution
Let assume q , p θ to be 1 dimensional student-t distribution.
The target is to approximate q by p θ .
The training set is 300 i.i.d data sampled from q with mean 0 and scale 0 .
We assume the scale parameter for p θ is the same as q , and the only trainable parameter θ is the mean.
The degree of freedom is 5 for both q , p θ .
The left panel of ﬁgure 1 shows the score matching loss computed for different θ .
We can observe that for original SM ( q, p θ ) loss, it is highly non-convex, and the loss value does not correlate well with likelihood.
Indeed, we can see the true location θ = 0 is protected by two high ’walls’.
In other words, unlike maximum likelihood estimator, a parameter θ that is closer to the ground truth does not necessarily produce low SM loss.
One important consequence is that unless the initialized θ is within the narrow valid region, the gradient-based optimization will never recover the truth.
On the other hand, the middle panel of ﬁgure 1 shows that if we chose m ( x ) = (1 + ( x − θ ) 2 0 . 6 ) (Manual Flow) as the diffusion matrix, the corresponding DSM m ( q, p θ ) loss is convex.
The ground truth can be recovered by minimizing
DSM with a proper gradient-based optimizer.
However, this DSM m ( q, p θ ) is only a surrogate objective for learning θ because the diffusion matrix m contains θ .
Thus, the dropped term C q,m in eq.6 is no longer a constant.
Although one can treat the θ in m as constant during training and ignore its contribution when taking the derivative, this is equivalent to use different losses after each θ update.
We leave its convergence analysis for the future work.
[Figure] The selection of the diffusion matrix is crucial to the success of the estimator.
Unfortunately, the interpretation of this matrix is unclear, not mentioning a selection algorithm.
In the following, we aim to shed lights on this problem by connecting this diffusion with normalizing ﬂows.
3. Diffusion matrix as normalizing ﬂow 3.1. Interpreting DSM/DSD using normalizing ﬂow
Let assume we have two densities q X ( x ) , p X ( x ) deﬁned on R D and are twice differentiable.
We further deﬁne an differentiable invertible transformation T ( x ) : R D → R D : [Equation]
[Equation]
with the corresponding induced densities q Y ( y ) and p Y ( y ) .
We can prove the following theorem:
Theorem 3.1.
For twice differentiable densities q X ( x ) , p X ( x ) and an invertible differentiable transformation T : R D → R D , the diffusion Fisher divergence (Eq.5) is equivalent to the original Fisher divergence in y space:
[Equation]
(9) where y = T ( x ) , and p Y , q Y are corresponding densities after the transformation.
The diffusion matrix m ( x ) is the inverse of the Jacobian matrix ( ∇ x T ( x )) −
Proof.
From the change of variable formula, the corresponding densities p Y ( y ) , q Y ( y ) can be deﬁned as:
[Equation]
Then the Fisher divergence F ( q Y , p Y ) is formulated as:
[Equation]
where the last step comes from changing the variable to x = T − 1 ( y ) and noticing that ∇ y T − 1 ( y )
= ( ∇ x T ( x )) −from the inverse function theorem.
This objective coincides with the diffusion Fisher divergence (Eq.5).
Importantly, F m ( q X , p X ) is a valid divergence (i.e. F m ( p X , q X ) =[-
-]
iff.
p X = q X ) when m ( x ) is an invertible matrix for every x .
As normalising ﬂow transformations naturally give invertible Jacobian matrices, we can easily extablish the connection F ( q Y , p Y ) =
F m ( q X , p X ) with m ( x ) = ( ∇ x T ( x )) − 1 .
We also include the likelihood plots afte the transformation in Appendix D.
Similarly, we can prove the connections between DSD (Eq.7) and normalizing ﬂow.
The proof is in appendix A.
Theorem 3.2.
For twice differentiable densities q X ( x ) , p X ( x ) , an invertible differentiable transformation T ( x ) : R D → R D and differentiable test function in suitable test function family H : f : R D → R D ∈ H , the diffusion
Stein discrepancy (Eq.7) is equivalent to the original Stein discrepancy
[Equation]
where g ( y ) =
f ( T − 1 ( y )) , H (cid:48) is the corresponding function space for g , p Y and q Y are transformed densities by
T ( · ) .
The diffusion matrix m ( x ) = ( ∇ x T ( x )) − 1 .
Based on the above two theorems, we formally establish the connections between the diffusion Fisher divergence / DSD with normalizing ﬂows.
This gives us an interpretation of the diffusion matrix as the inverse of the Jacobian matrix deﬁned by the ﬂow.
3.2. Better ﬂow design
Based on the interpretation, we try to give a better design for the diffusion matrix m . Here, we design a ﬂow T G ( · ) that transforms the Student-t distribution p θ to a standard Gaussian N (0 , 1) , which we named as Gaussian ﬂow : [Equation]
Here F G and F θ are cumulative density functions for
N (0 , 1) and p θ respectively. We plot the corresponding DSM m ( q, p θ ) loss in the right panel of Figure 1. Both the manually designed ﬂow and Gaussian ﬂow can recover the ground truth θ regardless of initialization.
However, Gaussian ﬂow allows faster convergence during training.
The fast convergence regions is the region where the gradient of the DSM w.r.t θ has a magnitude greater than 1 .
The
Gaussian ﬂow has a much wider region compared to manual ﬂow.
The length of the region is 4 . 44 and 10 .
56 respectively (more than 2 times).
For high dimensional distributions, this area of the region can scale up with O (2 D ) , which can have signiﬁcant impact on convergence speed.
Another advantage of this systematic design of the diffusion matrix is its robustness, which is further discussed in Appendix E.
3.3. Interpreting DSM using Riemannian manifold
Assume we have a Riemannian manifold ( M , g ) with Riemannian metric tensor g .
For each point a ∈ M , we assume it has a local coordinates x a = [ x 1 a , . . .
x Da ] .
We can prove the following proposition:
Proposition 3.1.
Deﬁne two probability measures Q , P on the Riemannian manifold ( M , g ) as deﬁned above.
We denote the corresponding densities (in terms of local coordinates x ) w.r.t.
Riemannian manifold as ˜ p ( x ) =
d P d M ( x ) and ˜ q ( x ) = d
[Equation] M ( x ) .
Then, the Fisher divergence from ˜ q to ˜ q is
[Equation]
13) where p ( x ) =
d P d M ( x ) d M ( x ) d x , q ( x ) is deﬁned similarly, and ∆ ( x ) =
s p ( x ) − s q ( x ) .
G ( x ) is an symmetric positive deﬁnite matrix representing the Riemannian metric tensor.
Particularly, if G ( x ) = m ( x ) − T m ( x ) − 1 , then F M (˜ q, ˜ p ) is equivalent to the diffusion Fisher divergence (Eq.5) with diffusion matrix m ( x ) .
The proof is in appendix B.
This result is more general than theorem 3.1.
Speciﬁcally, theorem 3.1 only proves a sufﬁcient condition for the diffusion Fisher divergence to be a valid discrepancy.
Namely, if we have an invertible ﬂow, the diffusion matrix m ( x ) must be invertible.
However, the converse is not true.
On the other hand, proposition 3.1 only requires m ( x ) to be invertible, which is more general.
Indeed, from the topological point of view, if we have an invertible and differentiable ﬂow T , then the transformed space (Riemannian manifold) is actually diffeomorphic to the original space (e.g. R D ).
Thus, this ﬂow can be viewed as a special case of Gemici et al.
But in general, Riemannian manifold may not be diffeomorphic to R D , which explains why theorem 3. is only a sufﬁcient condition.
3.4. Continuous DSM with ODE ﬂow
Previous sections assume a deterministic transformation
T ( x ) .
Recent work has shown promising results for continuous ﬂows characterised by an ODE (Chen et al., 2018;
Grathwohl et al., 2018).
[Equation]
where g ( x ( t )) is a deterministic drift that is uniformly Lipschitz continuous w.r.t.
x .
We deﬁne p t and q t to be the corresponding densities for x ( t ) .
Inspired by Chen et al.
(2018), we can characterise the instantaneous change of the score matching loss d F ( q t ,p t ) dt by the following proposition:
Proposition 3.2.
Let p t ( x ( t )) , q t ( x ( t )) be two probability density functions, where x ( t ) is characterized by an ODE deﬁned in eq.14.
Assume g
x ( t )) is uniformly Lipschitz continuous w.r.t.
x ( t ) .
Then, the instantaneous change of score matching loss follows:
[Equation]
where ∆( x ) = s p t ( x ) − s q t ( x ) .
The proof is in appendix C.
4. Conclusion
In this paper, we discuss the connections of the diffusion score matching and diffusion Stein discrepancy to normalizing ﬂows.
Speciﬁcally, we theoretically prove that the diffusion Fisher divergence (or DSD) is equivalent to performing the original Fisher divergence (or Stein discrepancy ) on the transformed densities.
The diffusion matrix m ( x ) is deﬁned by the inverse of the ﬂow’s Jacobian matrix.
We also establish the connection of diffusion Fisher divergence with densities deﬁned on Riemannian manifolds.
In the end, we extend the diffusion Fisher divergence by continuous ﬂow, and derive an ODE characterizing its instantaneous changes.
By building the connections, we hope to shed lights on developing training method for the diffusion matrix to enable the practical usage for large models.
A. Proof of theorem 3.
Proof.
Let’s ﬁrst deﬁne the Stein operator as
[Equation]
for the test function g ( y ) and density p Y ( y ) .
Thus, the
Stein discrepancy can be rewritten as
[Equation]
In the following [Equation] we will focus on the Stein operator.
From the change of variable formula y =
T ( x ) [Equation] we have
[Equation]
Now we can rewrite the Stein operator:
[Equation]
[Equation]
The second equality is from the chain rule and deﬁnition of divergence operator ∇ T .
For the layout of the matrix calculus, we follow the column vector layout as the following: for a function h : R D → R , and f : R D → R N , we have
[Equation]
Now, we focus on 1 term: ∇ [Equation] 1 ( y )) [Equation]
∇ [Equation] ∇ [Equation]
T − 1 ( y )] = Tr [ ∇ x T ( x ) ∇ y ( ∇ x T ( x )) − 1 ] = Tr [ ∇ x T ( x ) ∇ y T − 1 ( y ) ∇ x ( ∇ x T ( x )) − 1 ] =
Tr [ ∇ x ( ∇ x T ( x )) − 1 ] (21) where we use the inverse function theorem ∇ y T − 1 ( y ) =
( ∇ x T ( x )) − 1 .
In addition, we deﬁne ∇ T x ( ∇ x T ( x )) − 1 =
Tr [ ∇ x ( ∇ x T ( x )) − 1 ] .
So we can set m ( x )
=
( ∇ x T ( x )) − 1 , we can obtain:
[Equation]
which is exactly the same as the inner part of DSD (Eq.7).
So with change of variable formula, we can easily show
[Equation]
B. Proof of proposition 3.
With the deﬁnition of the Riemannian manifold ( M , g ) , for any point a ∈ M with local coordinates x ∈ R D , and two vectors u , v from its tangent plane T a M , we can represents u , v using the basis ( ∂∂x i ) a as
[Equation]
The inner product deﬁned by the metric g can be expressed as
[Equation]
where g ij ( x ) is the ij − th element of matrix G ( x ) and (cid:104)· , ·(cid:105) g is the inner product deﬁned by Riemannian metric g .
We assume the measure M ( x ) is absolutely continuous w.r.t.
Lebesgue measure, then we have the following change of variable formula
[Equation]
Then we can represents the densities ˜ p , ˜ q under Lebessgue measure
[Equation]
and q ( x ) is deﬁned accordingly.
The score matching loss for ˜ p and ˜ q is
[Equation]
Now let’s deﬁne ∇ log ˜ p ( x ) .
From the basics of Riemannian manifold, for a point a ∈ M with local coordinate x , and
X is a vector ﬁeld on M , we have the following deﬁnition
[Equation]
Written in terms of matrix form, assume X
=
[ X 1 , . .
, X D ] T , and g ij ( x ) is the element of symmetric positive deﬁnite matrix G ( x ) , we have
[Equation]
Therefore, we have
[Equation] =
˜∆
[Equation]
[Equation]
By change of variable formula, it is also easy to show that
[Equation]
Therefore, we have ||∇ log
[Equation]
x [Equation] Substitute back to F M (˜ q, ˜ p ) (Eq.28), we can obtain the result.
Particularly, compare to diffusion Fisher divergence (Eq.5), we can observe that if G ( x ) = m ( x ) − T m ( x ) − 1 , the F M (˜ q, ˜ p ) is equivalent to diffusion Fisher divergence .
Indeed, as m ( x ) ∈ R D
× D is an invertible matrix, then m ( x ) − T m ( x ) − 1 must be symmetric positive deﬁnite, which satisﬁes the requirements for G ( x ) .
C. Proof of proposition 3.
An ODE ﬂow is deﬁned by the solution of an ODE:
[Equation]
with g ( x ) the deterministic drift term.
Let us consider the forward Euler discretisation of the ODE, which gives x
[Equation]
With δ ≈
we see that T δ is an invertible transformation.
Now consider y =
x ( t + δ ) and x ( t ) = x .
This again pushes p X ( x ) and q X ( x ) to p Y ( y ) and q Y ( y ) , respectively.
Therefore we can reuse results from theorem 3.1 and derive
[Equation]
where m ( x ) = ( ∇ x T δ ( x )) − 1 Notice that T δ ( x ) = x when δ = 0 .
This means we can compute the change of score matching at time t as: [Equation]
[Equation]
x )
[Equation]
)∆
[Equation]
x )] ,
[Equation]
37) with ∆( x ) = ∇ x log p X ( x ) − ∇ x log q X ( x ) . As ∇ x T δ ( x ) = I +
[Equation]
∇ x g ( x ) , simple calculation shows that
[Equation]
which leads to
[Equation]
As this quantiﬁes the instantaneous changes [Equation] replacing p t and q t for p Y , p X , q Y and q X gives the instantaneous change of score matching loss [Equation]
D. Additional plots
From the motivating example and theorem 3.1, we know
[Equation]
Therefore, by simple calculus, the corresponding transformation y =
T ( x ) can be deﬁned as
[Equation]
Let’s deﬁne the transformed densities p Y
( y ) and q Y ( y )
as
[Equation]
Therefore, we can plot the log likelihood for the original densities p , q and transformed densities p Y , q Y as Figure
[Equation]
10 x 17. L o g li k e li h oo d
Log likelihood for p and q log p ( x ) log q ( x )      3 y Log likelihood for p Y log p Y ( y )      3 y
Log likelihood for q Y log q Y ( y )
In this case, we set p ( x ) has the mean − 2 .
5 with the same scale 0 .
3 as q , whereas q has mean 0 .
For the transformation
T , we set θ = − 2 .
5 with b = 0 .
E. Robustness of Gaussian ﬂow
Here, we investigate the robustness of the diffusion matrix w.r.t.
the degree-of-freedom (DoF) of studnet-t distribution.
We adopt the similar settings as the motivating example (Section 2.2) where q and p θ are Student-t distribution with same scale parameter.
We vary their DoF together to investigate the changes in DSM loss.
Because the manual ﬂow lacks a proper interpretation, so it is difﬁcult to adapted to the change of DoF.
Thus, we use the same m ( x ) = 1 + ( x − θ ) 2 0 .for
all DoF.
On the other hand, Gaussian ﬂow is designed to transform from Student-t to standard Gaussian.
So it can be easily adapted to the change of DoF by using the corresponding F θ .
Figure 3 plots the DSM losses with both manual ﬂow and
Gaussian ﬂow.
From the top panel, we can clearly observe that manual ﬂow only works with DoF = 5 .
For other DoF, the corresponding DSM fails to recover the ground truth θ .
On the other hand, the bottom panel shows that Gaussian ﬂow is robust to the change of DoF, and consistently gives the correct ground truth θ with wide fast convergence region.
We emphasize again that for both ﬂows, the resulting
DSM m ( q, p θ ) is not equivalent to F m ( q, p θ )
(Eq. 5) as
C q,m in Eq.
6 is dependent on θ .
However, unlike the manually designed ﬂow by Barp et al.
(2019), the Gaussian ﬂow returns surrogate losses that have only one global optimum at the desired solution in all cases considered.
Future work will evaluate the Gaussian ﬂow construction of DSM objectives beyond student-t case for further understandings.