Convergence in KL and Rényi Divergence of the Unadjusted Langevin Algorithm Using Estimated Score
Abstract
We study Inexact Langevin Algorithm (ILA) for sampling using an estimated score function when the target distribution satisfies log-Sobolev inequality (LSI), motivated by Score-based Generative Modeling (SGM).
We prove convergence in Kullback-Leibler (KL) divergence under a sufficient assumption on the error of score estimator called bounded Moment Generating Function (MGF) assumption.
Our assumption is weaker than the previous assumption which requires the error has finite L ∞ norm everywhere.
Under the L ∞ error assumption, we also prove convergence in Rényi divergence, which is stronger than KL divergence.
On the other hand, under L p error assumption for any 1 ≤ p < ∞ which is weaker than bounded MGF assumption, we show that the stationary distribution of Langevin dynamics with an L p -accurate score estimator can be far away from the desired distribution.
Thus having an L p -accurate score estimator cannot guarantee convergence.
Our results suggest controlling mean squared error which is the form of commonly used loss function when using neural network to estimate score function is not enough to guarantee the upstream algorithm will converge, hence in order to get a theoretical guarantee we need a stronger control over the error in score matching.
Despite requiring an exponentially decaying error probability, we give an example to demonstrate the bounded MGF assumption is achievable when using Kernel Density Estimation (KDE)-based score estimator.
Introduction
Score-based Generative Modeling (SGM) is a family of sampling methods which have state-of-the-art performance in many applied areas (Song and Ermon, 2019;
Ho et al., 2020; Song et al., 2021), but the theoretical understanding of the methods is still lacking.
The basic idea of SGM is a two-step procedure.
The first step is to estimate the score function from the data (for example via score matching using neural network), the second step is to get new samples from the estimated score via
Langevin dynamics.
There have been extensive studies on convergence rate of Langevin dynamics and related discrete time algorithms in the setting of knowing exact score function (Vempala and
Wibisono, 2019; Chewi et al., 2022).
In contrast, theoretical analysis in the setting of using estimated score is limited.
De Bortoli et al.
(2021) studied the convergence in total variation that requires L ∞ error assumption.
Although L ∞ is sufficient to guarantee the convergence, it is too strong to hold in practice since it requires a finite and uniform error at every point x .
Lee et al.
and De Bortoli (2022) studied the convergence in total variation and Wasserstein distance of order one respectively under L 2 error assumption.
The convergence results are achieved after running the algorithm for a moderate amount of time.
The reason why we cannot run the algorithm for infinitely long time is that the stationary distribution of Langevin dynamics with an L 2 -accurate score estimator can be far away from the desired distribution (Balasubramanian et al., 2022; Lee et al., 2022).
This suggests having an L 2 -accurate score estimator is not sufficient to guarantee a distribution which is close to the desired one after running the algorithm for a long time.
Therefore there is a need to find a sufficient assumption on the error of score estimator which is achievable in practice and meanwhile can guarantee convergence.
Our contribution We prove convergence in KL divergence under a sufficient assumption on the error of score estimator called bounded Moment Generating Function (MGF) assumption which is weaker than L ∞ but stronger than L p for 1 ≤ p < ∞ .
We also prove convergence in Rényi divergence under L ∞ error assumption, which is stronger than KL divergence.
Contrary to previous results (Lee et al., 2022; De Bortoli, 2022), our convergence result is stable.
Moreover our results hold in KL divergence, which is stronger than TV.
We also show that L p error assumption for any 1 ≤ p < ∞ is not sufficient to guarantee long-term convergence even in TV.
Our analysis in this paper suggests controlling mean squared loss in score matching is not enough, we may need a stronger control over the error when constructing the loss function for neural network to train the score function, such as MGF of the error (mean of exponential of squared error).
We show the bounded MGF assumption is achievable by giving an example when using Kernel Density Estimator (KDE) to approximate the score function in Gaussian data.
Since our original submission to this workshop, we have generalized our analysis to score-based generative models and proved a similar stable convergence result in KL divergence (more details in Wibisono and Yang (2022)).
Notations In this paper, we use H ν ( ρ ) to denote KL divergence of ρ w.r.t.
ν , J ν ( ρ ) to denote relative Fisher information of ρ w.r.t.
ν , R q,ν ( ρ ) to denote Rényi divergence of order q of ρ w.r.t.
ν .
We provide a review on the definitions in Appendix A.
Preliminaries
Given the target distribution ν ∝ e − f in R n , the Unadjusted Langevin Algorithm (ULA) with step size h > 0 is the following discrete-time algorithm
[Equation]
where z k ∼ N (0 , I ) are independent standard Gaussians in R n .
This is a discretization of the Langevin dynamics in continuous time, which is a stochastic process that converges to the target distribution.
The convergence properties of the Langevin dynamics have been well studied under various assumptions such as strong log-concavity or weaker isoperimetric inequality such as logSobolev inequlity (LSI), which allows for some non-log-concavity.
There are also convergence guarantees in discrete time with small bias under LSI and smoothness with small step size (Vempala and Wibisono, 2019; Balasubramanian et al., 2022).
In practice, for example in SGM, we may not know true score function ∇ log ν ( x ) , so we replace it by an estimated score function s ( x ) in Eq.
(1) and we run Inexact Langevin Algorithm (ILA) with s ( x ) as follows
[Equation]
We study the convergence of above algorithm.
One step of the algorithm is the solution x k +
=
X h of the following SDE at time h starting from X 0 = x k
[Equation]
where W t is the standard Brownian motion in R n .
We present the following lemma on Eq.
(3), which will be used later.
Lemma 1.
Let ρ t be the law of Eq.
(3), then we have the following bound for time derivative of KL divergence
[Equation]
Proof of Lemma 1 is in Appendix B.
Main results 3.1 Convergence under bounded MGF assumption
We first introduce the following assumptions.
Assumption 1 (LSI) .
The target distribution ν satisfies LSI with constant α > 0 , which means for any probability distribution ρ
[Equation]
Assumption 2 ( L -smoothness) .
f = − log ν is L -smooth, which means ∇ f is L -Lipschitz.
Assumption 3 (Lipschitz score estimator) .
The score estimator s ( x ) is L s -Lipschitz.
Assumption 4 (MGF error assumption) .
The error of score estimator s ( x ) has finite moment generating function of some order r , i.e.
[Equation]
Later in this paper we will compare the above assumption with the following.
Assumption 5 ( L ∞ error assumption) .
The error of score estimator s ( x ) has finite L ∞ norm at every x , i.e.
[Equation]
Assumption 6 ( L p error assumption) .
The error of score estimator s ( x ) has finite moment of order p for some 1 ≤ p < ∞ , i.e.
[Equation]
Our main result is the following, which shows biased convergence rate of ILA under the MGF error assumption.
Theorem 1 (Convergence of KL divergence under MGF error assumption) .
Suppose Assumption 1and 4 with r = 9 α hold.
If 0 < h < min( α 12 L s L , 12 α ) , then after k iterations of ILA (2)
[Equation]
where C [-
-]1 = O ( nL 2 s L α ) and C 2 = 163 .
Proof of Theorem 1 is in Appendix C.
The convergence rate is similar to Theorem 1 in Vempala and
Wibisono (2019) in the setting of knowing exact score function but has an extra non-vanishing term induced by the error of score estimator.
So in order to have a small asymptotic error, we need an accurate score estimator.
Theorem 1 directly implies the following corollary.
Corollary 1.
Suppose Assumption 1-3 hold.
For any ε > 0 , if the score estimator s ( x ) has a small MGF error
[Equation]
then running ILA (2) with step size h ≤
[Equation]
for at least
[Equation]
iterations reaches H ν ( ρ k ) ≤ ε .
3.2 Comparing different error assumptions
Because L ∞ implies MGF error assumption, the convergence result in Theorem 1 also holds under Assumption 5 in place of Assumption 4.
In this case, we provide a simpler proof in Appendix D.
In addition, under L ∞ error assumption, we prove convergence in Rényi divergence of order q ≥ 1 , which is stronger than KL divergence.
We conjecture the convergence holds under the MGF error assumption and leave it for future work.
Theorem 2 (Convergence of Rényi divergence under L ∞ error assumption) .
Suppose Assumption 1-3 and 5 hold.
Let q ≥ 1 .
If 0 < h < min( α 12 LL s q , q 4 α ) , then after k iterations of ILA (2)
[Equation]
where C 1 = O ( hL 2 s Lq 2 α ) and C
= O ( q 2 α ) .
Proof of Theorem 2 is in Appendix E.
L ∞ -accurate score estimator requires a finite and uniform error at every point x , which is too strong and may not hold in practice.
By the construction of loss function in score estimation, the estimated score is only L 2 -accurate with high probability (Proposition 9 in Block et al.
However, the stationary distribution of Langevin dynamics with L 2 -accurate score can be far away from the true distribution in TV (Lee et al., 2022; Balasubramanian et al., 2022).
So having an L 2 -accurate score estimator cannot guarantee the sampling algorithm will converge to the desired distribution.
We further claim that any finite higher moment (Assumption 6) cannot guarantee the convergence either.
This is to say, the stationary distribution of Langevin dynamics with an L p -accurate score estimator where 1 ≤ p < ∞ can also be far away from the true distribution.
We use the example in
Balasubramanian et al.
(2022) to illustrate this.
Example 1 ( L p bound is not sufficient) .
Let ν = 34 N ( − m, 1) + 14 N ( m, 1) and µ = 12 N ( − m, 1) + 12 N ( m, 1) .
For all m ≥ 180 and p ≥ 1 , the following holds
[Equation]
However, the total variation between µ and ν is large, TV ( µ, ν ) ≥ 1800 .
The example above shows if the target distribution is ν , but we run Langevin dynamics with an
L p -accurate score estimator s = ∇ log µ , then the limiting distribution will be µ , hence TV ( µ, ν ) is the asymptotic bias.
We provide more details of Example 1 in
Appendix F. Therefore, in order to obtain a convergence guarantee with an achievable sufficient assumption, we need the error of score estimator in between L p and L ∞ .
The MGF error assumption is weaker than L ∞ but stronger than L p error bound for any 1 ≤ p < ∞ .
Indeed, the MGF error assumption implies P ν
( ∥∇ log ν ( x ) − s ( x ) ∥ 2 ≥ c n ) → 0 exponentially as c n → ∞ , which is strong.
It is not clear whether the assumption holds via score matching using neural network.
But here we show an example where the assumption is achieved when using a simple KDE-based score estimator in Gaussian data.
Example 2.
Let ν = N ( µ, σ 2 ) , then f = − log ν is 1 /σ 2 -strongly convex, thus ν satisfies LSI with constant 1 /σ 2 .
The score function is ∇ log ν ( x ) = − ( x − µ ) /σ 2 .
We estimate the score function via Gaussian KDE with bandwidth h > 0 , so the score estimator is
[Equation]
and s ( x ) is [Equation] Lipschitz.
Let r = 9 σ 2 .
If h < σ √ [Equation] [Equation] 1 , then the MGF is finite,
[Equation]
Discussion
In this paper we study the convergence of Inexact Langevin Algorithm (ILA) using estimated score function, motivated by SGM.
We prove convergence result in KL divergence under MGF error assumption, which is weaker than L ∞ error assumption.
In addition, under L ∞ assumption, we provide convergence in Rényi divergence, which is stronger than KL divergence.
On the other hand, we show L p error assumption for any 1 ≤ p < ∞ is not sufficient to guarantee convergence.
Our results suggest controlling L 2 loss (mean squared error) in score matching is not enough to guarantee the upstream algorithm will converge, therefore in order to get a theoretical guarantee one may need a stronger control over the error when constructing the loss function for neural network to estimate the score function, such as MGF (mean of exponential of squared error).
Using such loss functions will give us a more accurate score estimator especially in low density regions, since they punish more on those large errors which are usually from low density areas (Song and Ermon, 2019).
Although the MGF error assumption requires an exponentially decaying tail, we show it is achievable when using a simple KDE score estimator in Gaussian data.
It is intriguing to explore if the KDE score estimator satisfies the assumption for a larger class of distributions, such as strongly log-concave distributions.
It is also interesting to investigate if the score matching estimator by neural network satisfies this assumption.
A A review on notations and definitions
In this section, we review notations and definitions of Kullback-Leibler (KL) divergence, relative
Fisher information, Rényi divergence and Rényi information.
Let ρ, ν be two probability distributions in R n denoted by their probability density functions w.r.t.
Lebesgue measure on R n .
Assume ρ and ν have full support on R n .
Definition 1 (KL divergence) .
The Kullback-Leibler (KL) divergence of ρ w.r.t.
ν is
[Equation]
Definition 2 (Relative Fisher information) .
The relative Fisher information of ρ w.r.t.
ν is
[Equation]
Definition 3 (Rényi divergence) .
For q ≥ 0 , q ̸ = 1 , the Rényi divergence of order q of ρ w.r.t.
ν is
[Equation]
where
[Equation]
Definition 4 (Rényi information) .
For q ≥ 0 , the Rényi information of order q of ρ w.r.t.
ν is
[Equation]
B Proof of Lemma
Proof.
The continuity equation corresponding to Eq.
(3) is
[Equation]
[Equation]
x ]
[Equation]
+
[Equation]
ρ
[Equation]
( x ) .
Therefore
[Equation]
integrating by parts
[Equation]
C Proof of Theorem
We restate the full theorem here.
Theorem 1.
Suppose the target distribution ν ∝ e − f satisfies LSI with constant α > 0 and f is
L -smooth.
Let s ( x ) : R n → R n be an L s -Lipschitz score estimator and the error satisfies Assumption 4.
Let 0 < h < min( α 12 L s L , 12 α ) , after k iterations of ILA (2)
[Equation]
To prove Theorem [Equation] we will use the following auxiliary result.
Lemma 2.
If the score estimator s ( x ) is L s -Lipschitz and t ≤ 13 L s , then
[Equation]
Lemma 3.
Suppose the assumptions in Theorem 1 hold, then along each step of ILA (2) ,
[Equation]
Proof of Lemma 2.
By L s -Lipschitzness
[Equation]
It is more convenient for our subsequent analysis to have a bound in terms of s ( x t ) rather than s ( x 0 ) , so we use
[Equation]
Rearranging it gives
[Equation]
since x t =
[Equation]
+
[Equation]
by triangle inequality
[Equation]
since t ≤ 1 3 L s ⇐ = t ≤ h
[Equation]
α 12 L s L and α ≤ L
Therefore
[Equation]
So we can bound ∥ s ( x t ) − s ( x 0 ) ∥ 2 as follows ∥
[Equation]
by plugging in Eq.
[Equation] [Equation]
since t
[Equation]
Proof of Lemma 3.
For simplicity suppose k = 0 , then one step of ULA using estimated score with step size h > 0 is
[Equation]
This is the solution to the following SDE at
[Equation]
where W t is the standard Brownian motion in R n .
Let M ( x ) =
∥∇ log ν ( x ) − s ( x ) ∥ 2 .
By Lemma 1,
[Equation]
[Table] [Equation] [Table] [Equation]
[Table]
[Equation]
[Table] [Equation] [Table] [Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
by [Equation] [Table]
[Equation]
[Table] [Equation] [Table]
[Equation]
since t 2 ≤ h 2 ≤ α144 L
[Equation]
L
[Equation]
nL
[Equation]
t by LSI
[Equation]
{+where ( i ) is because t 2 ≤ h 2 ≤ α 2 
 L 2 s L 2 ≤ 1 144 L 2 s and ( ii ) is by Lemma 12 in Vempala and Wibisono (2019).+}
We then bound the second term by applying Donsker–Varadhan variational characterizations of KL divergence E P [ f ( x )]
≤ log E Q e f ( x ) + H Q ( P ) for f ( x ) =
α M ( x ) , P = ρ t and Q = ν
[Equation]
Therefore
[Equation]
by MGF error assumption
[Equation]
since t ∈ (
[Equation]
h )
This is equivalent to
[Equation]
since e − 14 αh ≤ 1 and e c − 1 ≤ 2 c for c = α 4 h ∈ (0 , 1) .
Renaming ρ 0 as ρ k and ρ h as ρ k +1 ,
[Equation]
The proof of Theorem 1 then follows Lemma 3.
Proof of Theorem 1.
Applying the recursion contraction k times
[Equation]
D Proof of convergence of KL divergence under L ∞ error assumption
We state the theorem here.
Theorem 3.
Suppose the target distribution ν ∝ e − f satisfies LSI with constant α > 0 and f is
L -smooth.
Let s ( x ) : R n → R n be an L s -Lipschitz score estimator with error bound in L ∞ ,
i
[Equation]
e.
[Equation]
Let 0 < h < min( α 12 L s L , 12 α ) , then after k iterations of ILA (2)
[Equation]
for some constants C 1 and C 2 .
We will use the following auxiliary result.
Lemma 4.
Suppose the assumptions in Theorem 3 hold, then along each step of ILA (2)
[Equation]
Proof of Lemma
[Equation]
By Lemma
[Equation]
∥ s ( x t ) − s ( x 0 ) ∥ 2 ≤ 18 L 2 s t 2 ∥ s ( x t ) − ∇ log ν ( x t ) ∥ 2 + 18 L 2 s t 2 ∥∇ log ν ( x t ) ∥ 2 + 6 L 2 s t ∥ z 0 ∥≤ 18 L 2 s t 2 M 2 ∞ + 18 L 2 s t 2 ∥∇ log ν
( x t ) ∥
2 + 6 L 2 s t ∥ z 0 ∥by L ∞ error assumption
[Equation]
The condition t ≤ 13 L s in Lemma 2 is satisfied since t ≤ h ≤ α 12 L s L and α ≤ L .
Therefore,
[Equation]
The second inequality is by Lemma 12 in Vempala and Wibisono (2019).
Therefore, the time derivative of KL divergence is bounded by
[Equation]
by LSI
[Equation]
since t ∈ [Equation] [Equation] h )
[Equation]
since h 2 ≤ α144 L
[Equation]
L
[Equation]
This is equivalent to
[Equation]
Integrating from 0 to h
[Equation]
where the last inequality is from e c − 1 ≤ 2 c for c = α 2 h ∈ (
[Equation]
Rearranging gives
H
[Equation]
since e − α 2 h ≤ 1 . Rename ρ
[Equation]
Proof of Theorem 3.
Applying the recursion (5) k times
[Equation]
we obtain
[Equation]
where ( i ) is because 1 − e − c ≥ 34 c for 0 <
c = α 2 h < 14 .
E Proof of Theorem
We restate the full theorem here.
Theorem 2.
Suppose the assumptions in Theorem 3 hold and let 0 < h < min( α 12 LL s q , q 4 α ) , then after k iterations of ILA (2)
[Equation]
for some constants C 1 and C 2 .
We will use the following auxiliary result.
Lemma 5.
Let φ t = ρ t ν and ψ t =
φ q − 1 t E ν
[ φ qt ] =
φ q − 1 t F q,ν ( ρ t ) , then
[Equation]
This is a generalized version of Proposition 15 in Chewi et al.
to the setting of estimated score.
Lemma 6.
Suppose the assumptions in Theorem 2 hold, then along each step of ILA (2)
[Equation]
Proof of Lemma
[Equation]
[Equation] = [Equation] ∥∇ (cid:0)
[Equation]
(cid:105) , because
[Equation]
Then we compute A 2 . Note that ∇ (cid:0)
[Equation]
, therefore
[Equation]
Applying ⟨ x, y ⟩ ≤ 12 q ∥ x ∥ 2 + q 2 ∥ y ∥ 2 to x = (cid:0) ρ
[Equation]
v (cid:1) − 12 ∇ (cid:0) ρ
[Equation]
v (cid:1)
[Equation]
{+2 and y = (cid:0) ρ t v (cid:1) q − 12 (cid:0) s ( x 
 ) −∇ log ν ( x t ) (cid:1) , (cid:10)(cid:0) ρ t v [Equation] 12 ∇+}
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
Therefore
[Equation]
Let φ t = ρ t ν and ψ t = φ q − 1 t E ν [ φ qt ] =
φ q − 1 t F q,ν ( ρ t ) , then
[Equation]
Proof of Lemma 6.
Lemma 5 states
[Equation]
All we need is to bound E ρ 0 t (cid:2) ψ t ( x t )
∥ s ( x
[Equation]
) − ∇ log ν ( x
[Equation]
) ∥ 2 (cid:3) . E ρ 0 t (cid:2) ψ t ( x t )
{+∥ s ( x 
+}
) −∇ {+log ν ( x t ) ∥ 2 (cid:3) ≤ 2 E ρ 0 t (cid:2) ψ t ( x t ) ∥ s ( x 0 ) − s ( x t ) ∥ 2 (cid:3) (cid:124) (cid:123)(cid:122) (cid:125) A+
 E ρ t (cid:2)+}
[Equation]
where
[Equation]
by Lemma
[Equation]
So we have
[Equation]
By Lemma 16 in Chewi et al.
under the assumption of ∇ log ν being L -Lipschitz,
[Equation]
Putting together,
[Equation]
since t 2 ≤ h [Equation] [Equation] L
[Equation]
where the last inequality is from Lemma 5 in Vempala and Wibisono (2019) under the assumption of ν satisfying LSI [Equation]
Therefore
[Equation]
Integrating from 0 to h ,
[Equation]
where the last inequality is because e c − 1 ≤ 2 c for c = α q h ∈ (0 , 1) .
Rearranging and renaming ρ 0 ≡ ρ k , ρ h ≡ ρ k
+1 , we obtain the recursive contraction
[Equation]
Proof of Theorem 2.
Applying the recursion in Lemma 6 k times, we have
[Equation]
where ( i ) is because 1 − e − c ≥ 34 c for c ∈
[Equation]
[-(cid:16) 
-]
F Details of Example
Proof.
Consider ν as target distribution and we estimate ∇ log ν by ∇ log µ .
We will show L paccuracy goes to 0 as m → ∞ but the total variation between ν and µ is lower bounded by a positive number.
For convenience, let ν 0 =
N ( − m, 1) , ν 1 = N ( m, 1) and rewrite ν = 34 ν 0 + 14 ν 1 , µ = 12 ν 0 + 12 ν 1 .
The lower bound of TV ( ν, µ ) follows from Devroye et al.
; Balasubramanian et al.
Moreover, Proposition 1 in Balasubramanian et al.
(2022) calculates
[Equation]
Therefore,
[Equation]
p [Equation] [Equation] ν p 0 ν p(cid:0)
{+ν 
 + ν 1 (cid:1) 2 p −since 3 ν 
 + ν 1 ≥ ν 0 + ν≤ 4 p − 1 m p (cid:16) (cid:90) R − ν pν p − 1+ (cid:90)+}
R + ν pν p − 1(cid:17) . Since (cid:90) R − ν pν p − 1= exp (cid:0) 2 p (2 p − 1) m 2 (cid:1) √ 2 π (cid:90) R − exp (cid:16) − (cid:0) x − (2 p − 1) m (cid:1) 2 (cid:17) dx = exp (cid:0) 2 p (2 p − 1) m 2 (cid:1) P N (0 , 1) { Z ≤ (2 p − 1) m } = exp (cid:0) 2 p (2 p − 1) m 2 (cid:1) P N (0 , 1) { Z ≥ − (
[Equation]
[Table]
[Equation]
and similarly
[Equation]
Therefore,
[Equation]