SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Futoshi Futami*  
Osaka University / RIKEN AIP  
futami.futoshi.es@osaka-u.ac.jp
Masahiro Fujisawa*†  
RIKEN AIP  
masahiro.fujisawa@riken.jp
We provide novel information-theoretic generalization bounds for stochastic gradient Langevin dynamics (SGLD) under the assumptions of smoothness and dissipativity, which are widely used in sampling and non-convex optimization studies. Our bounds are time-independent and decay to zero as the sample size increases, regardless of the number of iterations and whether the step size is fixed. Unlike previous studies, we derive the generalization error bounds by focusing on the time evolution of the Kullback-Leibler divergence, which is related to the stability of datasets and is the upper bound of the mutual information between output parameters and an input dataset. Additionally, we establish the first information-theoretic generalization bound when the training and test loss are the same by showing that a loss function of SGLD is sub-exponential. This bound is also time-independent and removes the problematic step size dependence in existing work, leading to an improved excess risk bound by combining our analysis with the existing non-convex optimization error bounds.
Stochastic optimization, including stochastic gradient descent (SGD), is central to realizing practical large-scale or deep-learning models. There are currently considerable active discussions on accurately determining the generalization performance of models trained by SGD or its variants. In particular, stochastic gradient Langevin dynamics (SGLD) [13, 39, 29], a noisy variant of SGD, has garnered much attention in this type of study since it provides a useful theoretical framework for generalization error analysis based on the Langevin diffusion context [29]. Our study aims to contribute to a more accurate understanding and evaluation of the generalization performance for SGLD.
There are two main approaches to generalization analysis in SGLD. One is the information-theoretic analysis proposed by Russo and Zou [30] and Xu and Raginsky [40], by which a generalization error bound is derived using the mutual information (MI) between the learned parameters and the training dataset. Recently, some extensions using gradient information have been made to investigate the generalization properties of SGLD, for example, upper-bounding the MI with the norm of gradients [28] and the sum of gradient variances [25, 35, 37, 36]. Information-theoretic generalization bounds are applicable to a wide range of noisy iterative algorithms such as differentially private SGD [11] and stochastic gradient Hamiltonian Monte Carlo [5] modified to include a noisy momentum.
The other approach is stability analysis, by which the effects of changes in the learning algorithm due to the addition or removal of a single training data point on the generalization performance are investigated. Raginsky et al. [29] derived the non-asymptotic generalization and excess risk bound of SGLD via the exponential ergodicity of Langevin diffusion. Starting with the study by Raginsky
\*Equal contribution. \*Corresponding author.
et al. [29], there have been many attempts to improve the generalization analysis in SGLD from the stability perspective, such as those by Zhang et al. [43], Mou et al. [23] and Li et al. [20].
Unfortunately, these existing generalization bounds are time-dependent; namely, they diverge with increasing number of iterations unless the step size is adjusted so that the order of bound values is  $\mathcal{O}(n^{-1})$  or  $\mathcal{O}(n^{-1/2})$ , where  $n$  is the sample size (see Section 5 for details). Farghly and Rebeschini [10] attempted to avoid this problem by Wasserstein stability analysis through reflection coupling [9] under the smoothness and dissipativity [14] assumptions commonly used in sampling and non-convex optimization communities [29, 44]. Although their bounds bypass the divergence problem when taking a supremum over time, the geometry introduced for the reflection coupling yields an unnatural dependence on step size, resulting in a vacuous bound as the step size decreases (see Table 1).
In this paper, we provide novel generalization bounds for SGLD under smooth and dissipative loss functions obtained by the information-theoretic approach. We focus on the upper bound of the MI, namely, the Kullback-Leibler (KL) divergence between the distributions of parameters learned from different training datasets. We then analyze its time evolution caused by the update of the SGLD algorithm through the Fokker-Planck (FP) equation (Lemma 1). On the basis of this analysis, we obtain time-independent generalization error bounds that decay to zero as  $n \to \infty$  regardless of the number of iterations or whether the step size is fixed (Theorem 4 and Corollary 1). Conventional information-theoretic generalization bounds [29, 28, 35, 36] are derived by bounding the MI between the parameters at all iterations and the training dataset. Therefore, these bounds grow linearly with the number of iterations, resulting in a time-dependent generalization error bound. Our analysis based on time evolution eliminates both this linearity issue and the unnatural dependence on step size (the inverse of step size) in the time-independent bound of Farghly and Rebeschini [10].
Another contribution is providing the first information-theoretic generalization bound and the excess risk bound when the same loss is used for training and the generalization performance evaluation. In the conventional information-theoretic approach, deriving generalization error bounds under this setting was challenging owing to the unknown tail behavior of a loss function of SGLD. We overcome this difficulty with our discovery that a smooth and dissipative loss function of SGLD is sub-exponential.
2.1 Problem settings and stochastic gradient Langevin dynamics
We represent random variables in capital letters, such as  $X$ , and deterministic values in lowercase letters, such as  $x$ , and express the Euclidean inner product and distance as  $\cdot$  and  $\| \cdot \|$ . Let  $\mu$  be an unknown generating distribution on the instance space  $\mathcal{Z}$  and  $w \in \mathcal{W} \subseteq \mathbb{R}^d$  be the  $d$ -dimensional parameters such as weights of neural networks, where  $\mathcal{W}$  is the space of the parameters. We consider a loss function  $l: \mathcal{W} \times \mathcal{Z} \to \mathbb{R}$  and the following optimization problem:
$$
\min  _ {w \in \mathcal {W}} L _ {\mu} (w) := \mathbb {E} _ {Z} [ l (w, Z) ] = \int_ {\mathcal {Z}} l (w, Z) \mathrm {d} \mu (z),
$$
which cannot be computed since  $\mu$  is unknown. Instead, we typically minimize the empirical risk estimated using the dataset  $S\coloneqq \{Z_i\}_{i = 1}^n$ :
$$
\min _ {w \in \mathcal {W}} L _ {\mathcal {S}} (w) := \frac {1}{n} \sum_ {i = 1} ^ {n} l (w, Z _ {i}),
$$
where  $\{Z_i\}_{i=1}^n$  are independent and identically distributed (i.i.d.) samples from  $\mu$ , i.e.,  $Z_i \stackrel{\text{i.i.d.}}{\sim} \mu$ .
Stochastic gradient Langevin dynamics. In this paper, we use the SGLD algorithm [39] to solve the empirical risk minimization. SGLD utilizes the gradient information of the loss function; however, some loss functions, such as the 0-1 loss, are not differentiable. In this case, it is common to use the differentiable surrogate loss function  $f: \mathcal{W} \times \mathcal{Z} \to \mathbb{R}$  (e.g., the cross-entropy loss) and minimize the following empirical risk:  $F_{S}(w) := \frac{1}{n} \sum_{i=1}^{n} f(w, Z_{i})$ . Given a mini-batch  $B \subset [n] := \{1, \dots, n\}$  with  $k = |B| \leq n$ , we define its mini-batch version as
$$
F (w, B) := \frac {1}{k} \sum_ {i \in B} f (w, Z _ {i}).
$$
The SGLD algorithm updates the parameters using the following recursion:
$$
W _ {t + 1} = W _ {t} - \eta_ {t} \nabla F (W _ {t}, B _ {t}) + \sqrt {2 \beta_ {t} ^ {- 1} \eta_ {t}} \xi_ {t}, W _ {0} \sim P _ {W _ {0}},
$$
where  $P_{W_0}$  is a given initial distribution,  $\nabla F(w,B)$  is a stochastic gradient,  $t$  is the number of iterations,  $\eta_t$  is the step size,  $\beta_t$  is the inverse temperature, and  $(B_t)_{t=0}^{\infty}$  is an i.i.d. sequence of random variables distributed uniformly on  $\{B \subset [n] : |B| = k\}$ . In addition,  $(\xi_t)_{t=0}^{\infty}$  is an i.i.d. sequence of standard Gaussian random variables, i.e.,  $\xi_t \sim \mathcal{N}(0,\mathbf{I}_d)$ , where  $\mathbf{I}_d$  is the  $d$ -dimensional identity matrix. The output parameters  $W \in \mathcal{W}$  obtained using SGLD can be seen as the samples from a conditional distribution  $P_{W|S} : \mathcal{Z}^n \to \mathcal{W}$ . We express the  $t$ -th output of the SGLD algorithm as  $W_t$ .
2.2 Expected generalization error and its bounds
The focus of this paper is the expected generalization error, defined as
$$
\operatorname {g e n} (\mu , P _ {W \mid S}; L) := \mathbb {E} _ {S, W} [ L _ {\mu} (W) - L _ {S} (W) ], \tag {1}
$$
where the expectation is taken over the joint distribution of  $(S,W)$ , i.e.,  $\mu^n\otimes P_{W|S}$ .
Information-theoretic generalization bounds. Russo and Zou [30] and Xu and Raginsky [40] have shown that Eq. (1) can be bounded by the MI between the input dataset  $S$  and the output parameters  $W$  under the following sub-Gaussian assumption.
Assumption 1 (sub-Gaussian losses). A loss function  $l(w,Z)$  is sub-Gaussian under  $Z\sim \mu$  for all  $w\in \mathcal{W}$ , that is, there is a positive constant  $\sigma_g^2$  such that  $\log \mathbb{E}_Z[\exp (\lambda (l(w,Z) - \mathbb{E}l(w,Z)))]\leq \lambda^2\sigma_g^2 /2$  for all constant  $\lambda \in \mathbb{R}$ .
For example, bounded or Lipschitz-continuous loss functions satisfy this assumption. Assumptions regarding the tail behavior of the loss function distributions as in the above are necessary for the information-theoretic generalization error analysis. Bu et al. [3] have investigated information-theoretic generalization bounds with another tail-behavior assumption such as sub-exponential losses.
We introduce the following standard information-theoretic generalization bound.
Theorem 1 (Russo and Zou [30] and Xu and Raginsky [40]). Suppose that Assumption 1 holds. Then, we have
$$
\left| \operatorname {g e n} \left(\mu , P _ {W \mid S}; L\right) \right| \leq \sqrt {\frac {2 \sigma_ {g} ^ {2}}{n} I (W ; S)}, \tag {2}
$$
under a training dataset  $S = \{Z_{i}\}_{i=1}^{n}$  and the algorithm's output  $W$ , where  $I(W;S)$  is the MI between  $W$  and  $S$ .
In the SGLD context,  $I(W,S)$  of Eq. (2) can be upper-bounded in a form that incorporates the gradient variance [28, 25, 35, 36]. Given the output of the  $T$ -th iterate of the SGLD algorithm,  $W_{T}$ , the following upper bound can be obtained.
Theorem 2 (Modified bound of Pensia et al. [28]). Let  $f(\cdot, z)$  be an  $L$ -Lipschitz continuous function, namely, there is a constant  $L > 0$  such that  $\| f(w, z) - f(\bar{w}, z) \| \leq L \| w - \bar{w} \|$  holds for all  $w, \bar{w} \in \mathcal{W}$  and all  $z \in \mathcal{Z}$ . Then, we obtain
$$
I \left(W _ {T}; S\right) \leq \sum_ {t = 0} ^ {\mathrm {i})} \frac {d}{2} \log \left(1 + \frac {\beta_ {t} \eta_ {t}}{d} \operatorname {V a r} [ \nabla f \left(W _ {t}, B _ {t}\right) | W _ {t} ]\right) \leq \sum_ {t = 0} ^ {T} \frac {d}{2} \log \left(1 + \frac {\beta_ {t} \eta_ {t} L ^ {2}}{d}\right), \tag {3}
$$
where  $\operatorname{Var}[\nabla f(W,B)|W] := \mathbb{E}\mathbb{E}_B[||\nabla_Wf(W,B) - \mathbb{E}_B[\nabla_Wf(W,B)]||^2 |W]$  is the conditional variance.
Note that bound (i) can be obtained by the data-processing inequality [6].
The merit of such information-theoretic generalization bounds is that we can evaluate the bound value using the empirically estimated gradient variance per iteration. However, unfortunately, from Eq. (3), this bound is time-dependent; namely, the bound value can diverge unless the gradient variance or  $\beta_{t}\eta_{t}$  approaches 0 as  $T\to \infty$ . This is due to the data-processing inequality when deriving upper bound (i) in Eq. (3). By the data-processing inequality, we obtain  $I(W_{T};S)\leq I(W^{(T)};S)$ , where
$W^{(T)} \coloneqq (W_0, W_1, \dots, W_T)$  denotes the joint random variables appearing in all the iterations in the algorithm. Since  $W^{(T)}$  is treated simultaneously, the bound is inevitably linear in  $T$ .
Another limitation of the information-theoretic approach appears in the setting where training losses  $(f)$  are also used for performance evaluation, which is often employed in sampling and non-convex optimization studies of SGLD [29, 41]. In this setting, the generalization error is defined as
$$
\operatorname {g e n} \left(\mu , P _ {W | S}; F\right) := \mathbb {E} _ {S, W} \left[ F _ {\mu} (W) - F _ {S} (W) \right], \tag {4}
$$
where  $F_{\mu} \coloneqq \mathbb{E}_{Z}[f(w,Z)]$ . We cannot conduct the information-theoretic analysis for Eq. (4) because the tail behavior of the distribution of the training loss is unclear.
Time-independent generalization bounds for Eq. (4). To solve the above problems, Farghly and Rebeschini [10] provided the generalization error bounds of Eq. (4) from the stability perspective under the following assumptions widely used in the non-convex optimization analysis of SGLD [29, 41, 19].
Assumption 2 (Smoothness). For each  $z \in \mathcal{Z}$ ,  $f(\cdot, z)$  is differentiable and  $M$ -smooth. That is, there is a positive constant  $M$  for all  $w$ ,  $\bar{w} \in \mathcal{W}$  and all  $z \in \mathcal{Z}$  such that
$$
\| \nabla f (w, z) - \nabla f (\bar {w}, z) \| \leq M \| w - \bar {w} \|.
$$
Assumption 3 (Dissipativity [14]). For each  $z \in \mathcal{Z}$ ,  $f(\cdot, z)$  is  $(m, b)$ -dissipative. That is, there are positive constants  $m$  and  $b$  for all  $w \in \mathcal{W}$  and  $z \in \mathcal{Z}$  such that
$$
m \| w \| ^ {2} - b \leq \nabla f (w, z) \cdot w.
$$
The discussion regarding loss functions that satisfy Assumption 3 is presented in Appendix B.
Hereafter, we eliminate the time dependence of the step size and temperature by setting  $\eta_t = \eta$  and  $\beta_t = \beta$ . With this notation, Farghly and Rebeschini [10] derived the following generalization bound.
Theorem 3 (Farghly and Rebeschini [10]). Suppose that Assumptions 2 and 3 hold. Assume that the initial law of  $W_0$  has a finite fourth moment  $\sigma$ . Then, if  $\eta \leq 1/2m$ , for any  $T \in \mathbb{N}$ , we have
$$
\left| \operatorname {g e n} \left(\mu , P _ {W _ {T} | S}; F\right) \right| <   C _ {1} \left(\eta T \wedge \frac {n \left(C _ {2} + 1\right)}{n - k}\right) \left(\frac {k}{n \eta^ {1 / 2}} + \eta^ {1 / 2}\right), \tag {5}
$$
where  $(x\wedge y) = \min \{x,y\}$ , and  $C_1$  and  $C_2$  are the positive constant terms w.r.t.  $\{M,m,b,d,\beta ,\sigma \}$  and  $\{M,m,b,d,\beta \}$ , respectively.
Farghly and Rebeschini [10] utilized the Wasserstein stability on the basis of the contraction property of Langevin diffusion under reflection coupling. The important technique to derive the above bound is that we only focus on  $W_{T}$  differently from  $W^{(T)}$  of the information-theoretic approach when deriving the contraction property. In this way, the resulting bounds do not suffer from divergence as  $T \to \infty$ ; however, it still has a problem. That is, Eq. (5) depends on the factor  $\eta^{-1/2}$ , which implies that it becomes vacuous or even diverges with decreasing  $\eta (= \eta_{T})$  as  $T \to \infty$ .
In this paper, we propose new generalization bounds to address the drawbacks of the information-theoretic and stability-based approaches. Specifically, the proposed bounds are established on the basis of the two expected generalization errors outlined in Eqs. (1) and (4), which remain time-independent and do not diverge as the step size decreases.
3 Time-independent generalization error bound for SGLD
Here, we explain our time-independent bound of  $\mathrm{gen}(\mu, P_{W|S}; L)$  for SGLD. We first introduce the main result (Section 3.1) and then summarize its proof outline (Sections 3.2 and 3.3). Finally, in Section 3.4, we provide a detailed discussion on our bound with concrete examples.
This assumption holds not only for (strongly) convex losses but also for many practically used non-convex loss functions [24]. For example, it applies to non-convex loss functions with  $l_{2}$  constraints and likelihood functions that satisfy Poincaré inequality [2, 32].
3.1 Main result
Our key idea is to derive the generalized error bound using the FP equation. To use the FP equation, we impose the following regularity condition for  $P_{W_0}$ .
Assumption 4 (Regularity of the initial distribution). The initial distribution of  $W_0$ :  $P_{W_0}$  is a Gaussian distribution with a finite variance  $s^2 > 0$ , which is independent of  $\eta$  and  $T$ .
Our analysis is also grounded in the time evolution of the FP equation using the logarithmic Sobolev inequality (LSI) [2] associated with  $\pi$  described as follows. We state that  $\pi$  satisfies the LSI with constant  $c_{\mathrm{LS}}$ , if for any  $\rho \ll \pi$ , the following relation holds:
$$
\mathrm {K L} (\rho | \pi) \leq c _ {\mathrm {L S}} \mathbb {E} \| \nabla \log \rho - \nabla \log \pi \| ^ {2}.
$$
Raginsky et al. [29] showed the existence of  $c_{\mathrm{LS}}$  under Assumptions 2, 3, and  $\beta \geq 2 / m$ . Note that  $c_{\mathrm{LS}}$  is expressed by the problem-dependent constant (see Appendix F.1 for details).
We now introduce our generalization error bound.
Theorem 4. Suppose that Assumptions 1, 2, 3, and 4 are satisfied. Then, for any  $\beta \geq 2 / m$  and  $\eta \in (0,1\wedge \frac{m}{5M^2}\wedge 4\beta c_{\mathrm{LS}})$  and any  $T\in \mathbb{N}$ , we have
$$
| \operatorname {g e n} (\mu , P _ {W _ {T} | S}; L) | \leq \sqrt {\frac {2 c _ {1} \sigma_ {g} ^ {2}}{n} \left(1 \wedge \frac {\eta T}{4 \beta c _ {\mathrm {L S}}}\right) \left(V _ {\nabla} + c _ {2}\right)}, \tag {6}
$$
where  $c_{1}, c_{2}$ , and  $V_{\overline{\Omega}}$  are the positive constant terms w.r.t.  $\{M, m, b, d, \beta, s^{2}\}$ .
The above theorem shows  $|\mathrm{gen}(\mu, P_{W_T|S}; L)| = \mathcal{O}(\sqrt{(\eta T \wedge 1)/n})$ , which implies time independence since it does not diverge even if  $T \to \infty$  and thus converges as  $n \to \infty$ .
In Eq. (6), the term  $V_{\nabla}$  corresponds to stability, which is expressed as the upper bound of the difference of the expected conditional gradients with respect to changes in training datasets at each iteration. This shows a certain similarity to existing information-theoretic generalization bounds, such as Theorem 2, expressed by the variance of gradients with respect to the training datasets. This similarity is discussed in detail in Section 3.4. Additionally, detailed information on the explicit expression of  $c_{1}, c_{2}$ , and  $V_{\nabla}$  can be found in Appendix F.
3.2 Proof outline of Theorem 4
In this section, we present how to derive our bound in Theorem 4. Our aim here is to share the ideas behind our analysis and an outline of the proof, providing the detailed proof in Appendix F.
We adopt the information-theoretic approach and focus on the MI in Eq. (2). By using the Jensen inequality, we have the following upper bound of the MI:
$$
I \left(W _ {T}; S\right) \leq \mathbb {E} _ {S, S ^ {\prime}} \mathrm {K L} \left(P _ {W _ {T} | S} \mid P _ {W _ {T} | S ^ {\prime}}\right), \tag {7}
$$
where  $S$  and  $S^{\prime}$  are random variables drawn independently from  $\mu^n$ , and  $\mathrm{KL}(P_{W_T|S}|P_{W_T|S'})$  is the KL divergence from  $P_{W_T|S'}$  to  $P_{W_T|S}$ . Note that this KL divergence indicates the stability of the learned parameter from two datasets,  $S$  and  $S'$ . We also note that  $P_{W_T|S'}$  can be regarded as the data-dependent prior. Thus, this KL divergence is tighter than that of the data-independent prior, which is often used in the probably approximately correct (PAC)-Bayes bound<sup>3</sup>.
The key idea is to analyze the time evolution of the KL divergence, which is summarized in the following lemma:
Lemma 1. Suppose that the same assumptions in Theorem 4 hold. Then, for any  $t \in \mathbb{N}$ , we have
$$
\mathrm {K L} \left(P _ {W _ {t} | S} \mid P _ {W _ {t} | S ^ {\prime}}\right) \leq e ^ {- \frac {\eta}{4 \beta c _ {\mathrm {L S}}}} \mathrm {K L} \left(P _ {W _ {t - 1} | S} \mid P _ {W _ {t - 1} | S ^ {\prime}}\right) + \eta V _ {\Delta} + c _ {3} \eta , \tag {8}
$$
where  $V_{\Delta}$  and  $c_{3}$  is the constant term w.r.t.  $\{M,m,b,d,\beta ,s^2\}$
2 The Gaussian assumption can be relaxed, e.g., to a Gaussian mixture, in the theorems and corollaries shown in this paper. The detailed discussions are provided in Appendix F.3. We can confirm this from the fact that  $\mathbb{E}_{S,S'}\mathrm{KL}(P_{W_T|S}|P_{W_T|S'}) = \mathbb{E}_S\mathrm{KL}(P_{W_T|S}|P_{W_T}) - \mathbb{E}_{S'}\mathrm{KL}(P_{W_T|S'}|P_{W_T})$ , where  $P_{W_T}$  is a data-independent prior distribution.
We will discuss the details of Lemma 1 in Section 3.3. By recursively applying Eq. (8) from  $t = 0$  to  $T$ , we obtain
$$
\mathrm {K L} (P _ {W _ {T} | S} | P _ {W _ {T} | S ^ {\prime}}) \leq \frac {1 - e ^ {- \frac {\eta T}{4 \beta c _ {\mathrm {L S}}}}}{1 - e ^ {- \frac {\eta}{4 \beta c _ {\mathrm {L S}}}}} \eta (V _ {\nabla} + c _ {3}) \stackrel {\mathrm {(i)}} {\leq} 4 \beta c _ {\mathrm {L S}} \left(1 \wedge \frac {\eta T}{4 \beta c _ {\mathrm {L S}}}\right) \frac {1}{1 - \frac {\eta}{4 \beta c _ {\mathrm {L S}}}} (V _ {\nabla} + c _ {3}),
$$
which is based on the fact that  $\mathrm{KL}(P_{W_0|S}|P_{W_0|S'}) = 0$ . Note that bound (i) can be obtained from
$$
e ^ {- \frac {\eta}{4 \beta c _ {\mathrm {L S}}}} <   1 - \frac {\eta}{4 \beta c _ {\mathrm {L S}}} + \frac {\eta^ {2}}{1 6 \beta^ {2} c _ {\mathrm {L S}} ^ {2}} \text {w h e n} \frac {\eta}{4 \beta c _ {\mathrm {L S}}} \leq 1, \text {a n d} e ^ {- \frac {\eta T}{4 \beta c _ {\mathrm {L S}}}} \geq 1 - \frac {\eta T}{4 \beta c _ {\mathrm {L S}}}.
$$
3.3 Proof outline of Lemma 1 under the continuous Langevin diffusion
Here, we organize our ideas for the proof of Lemma 1 that are important in the derivation of Theorem 4. For simplicity, we now provide an intuitive explanation and an outline of the proof under the continuous Langevin diffusion setting. Note that the results of Theorem 4 and Lemma 1 are based on the SGLD setting, and their proofs are shown in Appendix F.
The Langevin diffusion is defined as
$$
\mathrm {d} W _ {t} = - \nabla F (W _ {t}, S) \mathrm {d} t + \sqrt {2 \beta^ {- 1}} \mathrm {d} H _ {t}, \tag {9}
$$
where  $\mathrm{d}H_{t}$  is the standard Brownian motion in  $\mathbb{R}^d$ . Note that, in this context,  $t$  expresses the continuous time and the full-batch gradient  $\nabla F(W_t,S)$  is used. The stationary distribution of Eq. (9) is given as the Gibbs distribution  $\pi (\mathrm{dw})\propto \exp (-\beta F(w,S))$ .
With some abuse of notation, let us denote  $P_{W_t|S}$  as the conditional distribution obtained using Eq. (9) and express its density as  $\rho_t$ . Then, the FP equation [2] for Eq. (9) can be obtained as
$$
\frac {\partial \rho_ {t} (w , t)}{\partial t} = \nabla \cdot \left(\frac {1}{\beta} \nabla \rho_ {t} + \rho_ {t} \nabla F (w, S)\right). \tag {10}
$$
Similarly, we can define the Langevin diffusion when we use dataset  $S'$  and the conditional distribution using that diffusion as  $P_{W_t|S'}$  with the density  $\gamma_t$ , obtaining the FP equation in the form of  $\rho_t$  replaced by  $\gamma_t$  in Eq. (10).
Now we analyze the time evolution of  $\mathrm{KL}(P_{W_t|S}|P_{W_t|S'}) = \mathrm{KL}(\rho_t|\gamma_t)$  at time  $t$ , i.e.,  $\partial \mathrm{KL}(\rho_t|\gamma_t) / \partial t$ . By utilizing the FP equations of  $\rho_t$  and  $\gamma_t$  and the Cauchy-Schwartz inequality, we obtain the following upper bound:
$$
\frac {\partial \mathrm {K L} (\rho_ {t} | \gamma_ {t})}{\partial t} \leq - \frac {1}{2 \beta} \mathbb {E} \| \nabla \log \rho_ {t} - \nabla \log \gamma_ {t} \| ^ {2} + \frac {\beta}{2} \mathbb {E} \| \nabla F (W _ {t}, S) - \nabla F (W _ {t}, S ^ {\prime}) \| ^ {2}. \tag {11}
$$