Towards Healing the Blindness of Score Matching
Abstract
Score-based divergences have been widely used in machine learning and statistics applications.
Despite their empirical success, a blindness problem has been observed when using these for multi-modal distributions.
In this work, we discuss the blindness problem and propose a new family of divergences that can mitigate the blindness problem.
We illustrate our proposed divergence in the context of density estimation and report improved performance compared to traditional approaches
Introduction
Score-based divergences such as the Fisher Divergence (FD; also known as score-matching divergence)
[10, 9] and Kernel Stein Discrepancy (KSD) [13, 3] are widely used in machine learning and statistics [1, 17].
Their main advantage is that the score function, a derivative of a log-density, can be evaluated without knowledge of the normalization constant of the density and can be applied to problems where other classical divergences (e.g. KL divergence) are intractable.
Unfortunately, this advantage can also be a curse in certain scenarios because the score function only provides local information about the slope of a density, but ignores more global information such as the importance of a point relative to another.
This has led to a blindness problem in many applications of score-based methods where the densities are multi-modal, including in density estimation [21, 16, 11], MCMC convergence diagnosis [8], Bayesian inference [14, 5]; see [20] for a detailed discussion.
To illustrate this problem, we recall the definition of FD and an example from [20].
Given two distributions with differentiable densities p and q supported on a common domain X ⊆ R d , the FD is FD(
[Equation]
where we denote by s p ( x ) ≡ ∇ x log p ( x ) and s q ( x ) ≡ ∇ x log q ( x ) the score functions of p and q respectively.
The classic sufficient conditions [10, 2] for the FD to be a valid statistical divergence (i.e. FD( p || q ) = 0 ⇔ p = q ) are: (i) p and q are differentiable with support X
= R d and (ii) s p , s q are square integrable, i.e. s p − s q ∈ L 2 ( p ) , where we denote f ∈ L 2 ( p ) ≡ (cid:82) X || f ( x ) ||
22 p ( x ) dx < ∞ .
The blindness problem of the FD can be illustrated through the following example due to [20].
Let p and q be a mixtures with the same components but different mixing weights:
[Equation]
where α p ̸ = α p , and g 1 , g 2 are Gaussian densities with variance σ 2 and means − µ and µ respectively.
Then FD( p || q ) → 0 when µ/σ 2 → ∞ regardless of the mixture proportions α p and α q .
To build intuition, we let µ = 5 , σ = 1 , α p = 0 .
2 , α q = 0 .
8 and plot the densities and score functions of p, q [Figure] in Figure 1a and 1b.
We can find the two distributions are very different but their scores are only different around x = 0 , which has a negligible density value under p .
We then fix α p = 0 .
2 and plot the FD( p || q ) as a function of α in Figure 1c.
Here we see the FD is 0 constant function, which shows the FD is ‘blind’ to the value of the mixture weight.
See [14] for a similar example for discrete X .
Understanding the Blindness Problem
In the example above, blindness is a numerical problem since the problem occurs despite the fact that the FD is a divergence in that case (i.e. FD( p || q ) = 0 ⇔ p =
q since (i) and (ii) are satisfied).
When µ/σ 2 → ∞ , although the Gaussian distributions still have the same support, the regions that contain most of the mass of g 1 and g 2 tend to be disjoint, which creates numerical issues.
However, the blindness problem is not simply a numerical problem, as illustrated in the following example.
Consider the case where p and q are mixtures whose identical components have disjoint supports.
For example, let g
and g 2 in Equation 2 have disjoint support sets X 1 , X 2 ⊆ R d respectively with
X 1 ∩ X 2 = ∅ .
Then, g 2 ( x ′ ) =
∇ x g 2 ( x ′ )
=
for x ′ ∈ X 1 and g 1 ( x ′ )
= ∇ x g 1 ( x ′ )
=
for x ′ ∈ X 2 .
In this case, the FD is independent of α q (see Appendix A.1 for a derivation):
[Equation]
Therefore, the FD is not a valid divergence here since FD( p || q ) =
̸⇒ p = q .
This example guides us to further study the topology properties of the distributions’ support required by the FD.
We first extend the Fisher divergence to distributions that have support on the connected space.
Theorem 1 (FD on a connected set) .
Assume two distributions (i) have differentiable densities p and q with support on a common 1 open connected set X ⊆ R d and (ii) s p − s q ∈ L 2 ( p ) .
Then, the FD is a valid divergence i.e. FD( p || q ) = 0 ⇔ p = q .
See Appendix A.2 for a proof.
Theorem 1 generalizes the classic FD that is defined on distributions with X =
R d [10, 2] ( R d is a special case of the connected set).
Secondly, Theorem 2 shows that connectedness of the support is a necessary condition to define a valid FD.
Theorem 2 (FD is ill-defined on disconnected sets) .
Assume two distributions have common support
X consisting of disjoint sets.
Then, the FD is not a valid divergence i.e. FD( p || q ) =
̸⇒ p = q .
See Appendix A.3 for a proof.
Intuitively, the score function only considers the local derivatives and contains no information of the global normalization constant.
If the domain is disconnected, it cannot determine the mass allocation in different domains.
This observation can also be extended to the KSD by viewing KSD as a kernelized FD [13, 3], see Appendix A.5 for a detailed discussion.
Healing the Blindness Problem with the Mixture Fisher Divergence
In this section, we propose a new variant of the FD which is well-defined in the disconnected scenario.
Consider a distribution with density m with support X m =
R d and define the mixtures
[Equation]
where 0 < β < 1 .
We then define the Mixture Fisher Divergence (MFD) as
[Equation]
Theorem 3 shows the MFD is well-defined when p and q have support on a disconnected space.
[Figure]
Theorem 3 (Validity of the MFD) .
Consider two distributions with differentiable densities p, q supported on X p , X q ⊆ R d with s p , s q ∈ L 2 ( p ) and a differentiable density m with support X m =
R d , s m ∈ L 2 ( p ) .
Then MFD is a valid divergence, i.e. MFD( p || q ) = 0 ⇔ FD(˜ p || ˜ q ) ⇔ p = q .
See Appendix A.6 for a proof.
For MFD, we no longer require that p, q have common connected support , since X m = R d results in ˜ p, ˜
q having connected support R d 2 .
The requirements of m ( x ) are mild and hold for simple choices of distribution e.g. a Gaussian.
To avoid the numerical problem mentioned in Section 1, m ( x ) should be chosen to effectively connect the different component distributions.
As an example, for the toy problem described in Figure 1 with components N ( − 5 , 1) and N ( − 5 , 1) we can choose β = 0 .
5 and m ( x ) =
N (0 , 9) that covers both components.
Figureshows the densities and their score functions for ˜ p, ˜ q .
We see that the score functions are different on the high-density region of ˜ p .
Figure 2c also shows the minimal value of the MFD( p || q ) is attained when α q = α p , which indicates that the proposed MFD heals the blindness problem in this example.
Density Estimation with Energy-based Models
Given a dataset X train = { x 1 , · · · , x N } sampled i.i.d.
from a data distribution p d with support X p d ⊆ R d , we would like to learn a model q θ to approximate p d .
We are interested in a family of models which can only be evaluated up to a normalization constant, e.g. an energy-based model q θ ( x ) =
e − f θ ( x ) /Z
( θ ) , where f θ is a neural network and Z ( θ ) =
(cid:82)
e − f θ ( x ) dx .
In this case, the standard Maximum Likelihood Estimation (MLE) is not applicable (since Z ( θ ) cannot be evaluated during training) and an alternative form of the FD [10] can be applied (see Appendix A.4 for a derivation and additional assumptions)
[Equation]
where ∇ x s q θ ( x ) = ∇ 2 x log q θ ( x ) is the Hessian matrix and the constant represents the terms that are independent of θ .
The integration over p d can be approximated by Monte-Carlo using X train .
Because both s q θ and ∇ x s q θ only depend on f θ , the normalizer Z q ( θ ) is not required during training and we only need to estimate Z q ( θ ∗ ) once after training.
Therefore, density estimation with FD in this setting contains two steps: (1) learn θ ∗ using Equation 6; (2) estimate Z ( θ ∗ ) to obtain the normalized density q θ ( x ) .
This scheme can result in blindness in practice [21].
To heal the blindness, we can apply the proposed MFD.
However, if we directly minimize MFD in step (1), the score s ˜
q θ ( x ) = ∇ x log ( β exp( − f θ ( x )) /Z
q ( θ ))
+ (1 − β ) m ( x )) requires estimating
Z q ( θ ) .
This negates the advantage of using score matching because now Z q ( θ ) must be estimated for every gradient step during training (similar to MLE).
To avoid this, we propose to instead directly approximate ˜ p
d with an energy-based model ˜ q θ ( x ) ≡ e − f θ ( x ) /Z
˜
q
( θ ) and ˜ q θ can then be trained using
[Equation]
where the integration over ˜ p
d ( x ) can be approximated using the samples from the mixture ˜ p
d ( x ) =
βp d ( x ) + (1 − β ) m ( x ) .
Therefore, the learning of θ is independent of Z ˜ q
( θ ) .
Optimally we have [Figure] ˜ q θ ∗ ( x ) = ˜ p
d ( x ) =
βp d ( x )+(1 − β ) m ( x ) .
To obtain a model of the underlying true density q ∗
= p d , we need to remove the mixture component from ˜ q θ ∗ , which can be done through a ‘correction step’:
[Equation]
This procedure for obtaining q ∗ is equivalent to q ∗ ( x ) =
arg min q MFD( p d ( x ) || q ( x ))
and when
MFD( p d ( x ) || q ( x ))
= 0 , we have q ∗ ( x ) =
p d ( x ) .
Therefore, density estimation with MFD in this setting contains three steps: (1) learn θ ∗ by minimizing Equation 7; (2) estimate
Z ˜ q
( θ ∗ ) ; and (3) apply the correction step (Equation 8) to obtain q ∗ .
Compared to FD, the additional correction step has negligible computation cost.
Choice of m and β : As we discussed in Section 3, a good m should have support R d and be able to bridge disconnected component distributions.
For a given set of data samples { x 1 , · · · , x N } ∼ p d , we can simply choose m ( x ) =
N (¯ µ, ¯Σ) , where ¯ µ and ¯Σ are the empirical mean and covariance of the available training data: ¯ µ = 1 N (cid:80) Nn =1 x n , ¯Σ =
1 N (cid:80) Nn =1 x n x Tn , which corresponds to an empirical moment matching approximation of p d and can thus cover different components.
The β is treated as a hyper-parameter in our method.
Intuitively, a large beta means that the proportion of data points from p d is small, and the model is learning m .
On the other hand, a small value means we may still have the numerical version of the blindness issue.
In this experiment, we use β = 0 .
8 can find it can empirically heal blindness.
We leave the theoretical study of choosing the β into future work.
Demonstration: We apply the proposed method to train a deep energy-based model and examine the performance against two target densities with multiple isolated components: 1) a weighted mixture of four Gaussians p d ( x ) = 0 .
1 g 1 ( x ) + 0 .
2 g 2 ( x ) + 0 .
3 g 3 ( x ) + 0 .
4 g 4 ( x ) , where g 1 , g 2 , g 3 , g 4 are 2D
Gaussians with identity covariance matrix and mean [ − 5 , − 5] , [ − 5 , 5] , [5 , 5] , [5 , − 5] respectively; and 2) a mixture of 3 concentric circles as proposed in [21].
We use Simpson’s rule for the 2D numerical integration to estimate the normalization constant for both methods.
The model specifications and training details can be found in Appendix B.
In Figure 3 we plot the ground truth and the estimated density with classic FD and the proposed MFD methods.
We also provide the corresponding KL evaluation (see Appendix B) between the ground truth density p d and the estimated model p θ .
We find the proposed MFD method can significantly improve performance and heal the blindness problem.
Related Work
In addition to the mixture construction, conducting a Gaussian convolution on both p d and q θ can also bridge the disjoint components and defines a valid divergence [22].
However, the score function is generally intractable for a deep energy-based model q , see Appendix D for a detailed discussion.
Paper [16] proposes to add Gaussian noise with variance σ 2 only to p d and anneal σ 2 → 0 during training.
This helps alleviate the blindness problem in the early stage of training but when σ 2 ≈ 0 , the blindness phenomenon will be observed again, see Appendix C for an example.
Paper [7] proposes to transform p and q with a common differentiable invertible function before defining the FD, which is also shown to be equivalent to [2].
However, since the invertible transformation is a homeomorphism and will not change the topology of its domain [4, 23], the invertible transformation will not fix the blindness caused by the disconnected support sets in principle.
The blindness problem also exists in other score-based applications.
As discussed in Section 3, directly applying the MFD requires knowing the normalizer, which potentially sheds light on choosing score-based methods.
We leave the case-by-case study of how to heal the blindness to future work.
A Derivations and Proofs A.1 Derivation of Equation
Let two differentiable densities g 1 and g 2 have disjoint supports X 1 ∩ X 2 = ∅ and
[Equation]
The FD between p and q can be written as
[Equation]
Since g 1 and g 2 has disjoint support [Equation] so g 2 will be a zero function on the support of g [Equation]
so g 2 ( x ′ )
=
∇ x g 2 ( x ′ )
=
for x ′ ∈ X 1 .
We then have
[Equation]
and
[Equation]
Similarly, for x ′ ∈ X 2 we have s p ( x ′ )
=
s q ( x ′ )
=
s g 2 ( x ′ ) .
Therefore, the FD is equivalent to
FD
[Equation]
|
[Equation]
13) which is independent of α q .
A.2 Proof of Theorem
The following two lemmas can be found in Folland [6, Corollary 2.41 and Theorem 2.42].
For completeness, we also provide simplified proofs.
Lemma 4.
Suppose f : X → R is differentiable on an open convex set X ⊆ R d and ∇ x f ( x ) =for all x ∈ X , then f is a constant on X .
Proof.
For any two points x 1 , x 2 ∈ X , we denote the the line segment that connects a, b as L x 1 ,x 2 .
Since X is a convex set, then L x 1 ,x 2 ⊆ X .
By the Mean Value Theorem (see Folland [6, Theorem 2.39]), there exists a point x 3 ∈ L x 1 ,x 2 such that f ( x 2 ) − f ( x 1 ) =
∇ x f ( x 3 )( x 2 − x 1 ) .
Since x 3 ∈ S , so ∇ x f ( x 3 ) = 0 thus f ( x 2 ) =
f
( x 1 ) .
Therefore, f has to be a constant function.
Lemma 5.
Suppose f : X → R is differentiable on a connected open set X ⊆ R d and ∇ x f ( x ) =for all x ∈ X , then f is a constant on X .
Proof.
For any point a ∈ X , we define X 1 = { x ∈ X : f ( x ) =
f ( a ) } and X 2 = { x ∈ X : f ( x ) ̸ = f ( a ) } ,
so X =
X 1 ∪ X 2 by construction.
For every x ∈ X 1 , there is a ball B ∈ S centred at x .
Since
B is convex, we have B ∈ X 1 by Lemma 4.
Therefore, every point x ∈ X 1 is an interior point of X 1 , so X 1 is an open set.
The image of X 2 under f : R \ { f ( a ) } is an open set, so X 2 is a open set since f is a continuous function (see Folland [6, Theorem 1.33]).
We thus have both X 1 and X 2 are open sets and X 1 is non-empty (it contains a ).
Since any connected space cannot be written as an union of two disjoint non-empty sets (see Tao [18, Definition 2.4.1]), so X
= X 1 ∪ X 2 indicates X 2 = ∅ .
Therefore, f is a constant function.
We can then prove the Theorem 1.
For two a.c. distributions that are supported on a connected space
X ⊆ R d with differentiable density p and q .
Then FD( p || q ) = 0 ⇔ ∇ x log p ( x ) = ∇ x log q ( x ) for x ∈ S .
We define function f ( x ) = log p ( x ) − log q ( x ) , so f ( x ) differentiable on X and ∇ x f ( x ) = 0 .
By Lemma 5, we have f as a constant function (we denote as c ) so we have p = q exp( c ) .
Since p and q are densities, we have (cid:82) q ( x ) exp( c ) dx = 1 ⇔ c = 0 .
Therefore, FD( p || q ) = 0 ⇔ p = q .
A.3 Proof of Theorem
Since we can always represents a distribution with disjoint support set as a mixture distribution with components supported on several connected subsets, we can then prove the theorem by Proposition 1.
Proposition 1 (FD is ill-defined on disconnected sets) .
Let a set of a.c. distributions have differentiable densities { g 1 , · · · , g K } with mutual disjoint (disconnected) support sets {X 1 , · · · , X K } :
X i (cid:84)
X j = ∅ for any i ̸
=
j and each support X i is connected.
Let two densities p
= (cid:80) k
α kp g k and q = (cid:80) k
α kq g k with positive coefficients (cid:80) k
α kp = 1 and (cid:80) k
=
α kq = 1 .
Then FD( p || q ) = [-
 ⇔α-]0 ⇔ α kp = α kq e c k , where { c 1 , · · · , c K } is a set of constants with constraints (cid:80) k
e c k
= 1 .
We can decompose FD( p || q ) = 12 (cid:80)
Kk =
α kp (cid:82) X k g k
( x ) || s p ( x ) − s q ( x )
|| 22 dx .
Since α kp and g k are positive, FD( p || q ) =
⇒ (cid:82) X k g k
( x ) || s p ( x ) − s q ( x )
|| 22 dx =
for any k , so ∇ x log p ( x ) = ∇ x log q ( x ) for x ∈ (cid:83) Kk =1 X k .
Since X k is connected, by Lemma 5, we have for x ∈ X k , log p ( x ) − log q ( x ) =
c k ⇔ p ( x ) =
q ( x ) e c k ⇔ α kp g k ( x ) =
α kq g k
( x ) e c k ⇔ α kp = α kq e c k , where { c 1 , · · · , c K } is a set of constants.
Since (cid:80) k
α kp = (cid:80) k
α kq e c
k
= 1 and (cid:80) k
α kq = 1 , we then have the constrain (cid:80) k
e c k
= 1 .
A.4 Derivation of Score Matching
Let p d and q θ are differentiable densities with a common support X ⊆ R d and assume q θ is twice differentiable, we can rewrite the FD as [10]
[Equation]
where the constant terms are independent of the model parameters θ .
Using the log-trick, we have
[Equation]
For simplicity, we assume X = R and p d ( x ) s p θ ( x ) vanishes at −∞ and ∞ , using integration by parts, we have
[Equation]
In general, this holds for X
= R d
and lim || x ||→∞ p d ( x ) s q θ ( x ) = 0 or X ⊆
R d is a compact subset of R d and f ( x ) p ( x ) = 0 for x ∈ ∂ X where ∂
X is the piecewise smooth boundary of X (by the divergence theorem [6, Theorem 5.34]), also see [13] for a similar discussion.
Therefore, we have
[Equation]
A.5 Kernelized Stein Discrepancy Extensions
For two a.c. distributions p and q , the Kernelized Stein Discrepancy [13, 3] can be defined as (see [13, Definition 3.2])
[Equation]
where k is an integrally strictly positive kernel (see [13, Definition 3.1]) and x, x ′ are i.i.d.
samples from p ( x ) .
The KSD( p || q ) =
if and only if s p =
s q (see [13, 3]).
Therefore, when p and q are supported on a connected open set, by Lemma 5, we have KSD( p || q ) =
⇔ s p = s q ⇔ p = q .
When p and q are supported on a disconnected space, we have KSD( p || q ) =
̸⇒ p = q .
This is because the KSD can be upper bounded by a (positively) scaled FD [13, Theorem 5.1]:
[Equation]
[Equation] we then have FD( p || q ) =
⇒ KSD( p || q ) = 0 .
When p and q are supported on a disconnected space, we have FD( p || q ) =
̸⇒ p = q (Theorem 2), so KSD( p || q ) =
̸⇒ p = q .
A.6 Proof of Theorem
Since the support of m as X m = R d then ˜ p and ˜ q have the same support X
= R d .
For the score functions, we also have
[Equation]
so s ˜
p ∈
L 2 (˜ p ) and similarly s ˜ q ∈ L 2 (˜ p ) .
Therefore, the FD between ˜ p and ˜ q is a valid divergence i.e.
[Equation]
) = βq ( x ) + (1 − β ) m ( x ) ⇔ p ( x ) =
q ( x ) , thus
MFD( p || q ) = 0 ⇔ FD(˜ p || ˜ q ) = 0 ⇔ p = q .
B Experiment Details
For both experiments, we sample 100k data from p d as our training datasets.
The energy network f θ ( x ) is a 3-layer feedforward network with 200 hidden units and swish activation functions [15].
We train the model for 30k iterations with the Adam optimizer [12] and batch-size 300.
For the numerical integration we use Simpson’s rule provided in the package [19].
We use a Monte-Carlo approximation to estimate the KL divergence evaluations (cid:99) KL( p d ( x ) || q θ ( x ))
= 1 K (cid:80)
Kk =1 log p d ( x k ) − log p θ ( x k ) , where we use K = 10000 .
C Data Noise Annealing Doesn’t Help
In this section, we empirically show that only adding noise to the data and annealing the noise to 0 during training won’t fix the blindness problem in practice.
We use a deep energy-based model with a 3-layer feedforward neural network with 30 hidden units and tanh activation function to learn the toy mixture of two Gaussian distributions described in Section 1.
We train the model with Adam optimizer with a learning rate 3 e − 4 for 10k iterations and batch size 300.
We add convolutional
Gaussian noise to the data samples with a standard deviation of 3.0 and anneal to 0 by multiplying by 0.9999 at each iteration.
The noise at the end of training has a standard deviation less than 0.001.
In
Figure 4 we plot the learned density during training.
We find that when the noise is big the model can identify the correct mixture co-efficient, but when the noise is close to 0, the model fails to capture the correct mixing proportions.
We also plot the density estimation results with vanilla FD and the proposed MFD in Figure 5a and 5b and we find that the density estimation with MFD achieves the best performance.
[Figure]
D Spread Fisher Divergence
For two distributions with densities p d ( x ) and q θ ( x ) with supports X p d , X q θ ⊆
R d , we can choose k (˜ x | x ) =
N( x, σ 2 ) and let
[Equation]
We follow the spread f -divergence [22] and define the Spread Fisher Divergence ( (cid:102) FD ) as
[Equation]
The convolution transform makes ˜ p d and ˜ q θ have support X ˜ p
d = X ˜ q θ = R d (which is a connected space) and (cid:102)
FD k ( p d || q θ )
≡ FD(˜ p d || ˜ q θ ) is a valid discrepancy, i.e. (cid:102) FD k ( p d || q θ ) = 0 ⇔ ˜ p
d = ˜ q θ ⇔ p d = q θ .
The spread Fisher divergence is also well-defined for the singular distributions (distributions that are not a.c. w.r.t Lebesgue measure), see [22] for a detailed discussion.
Similar to the FD, we can rewrite the (cid:102) FD as
[Equation] [Equation]
(29) where the constant terms are independent of the model parameters.
For an energy-based model q θ ( x ) =
e − f θ ( x ) /Z
( θ ) , the spread model ˜ q θ (˜ x ) =
[Equation]
has an intractable score.
Additionally, unlike the mixture construction, if we directly assume ˜ q θ (˜ x ) =
e − f θ (˜ x ) /Z
( θ ) , the underlying ‘correct’ model q θ ( x ) can not be recovered from ˜ q θ (˜ x ) even if we know Z ( θ ) .
Therefore, the (cid:102) FD is not directly applicable in this case.