1 Introduction↩︎

Score-based diffusion models are a generative modeling framework based on the fact that the process of a stochastic differential equation (SDE) admits a time-reversal process, which is given by another SDE depending on the score of the (forward) distribution [1], [2]. The forward process is used for training, and the reverse process is used for sampling. Since the main interest is in the distribution itself and not in particular paths of the SDE, a whole family of SDEs, indexed by a stochasticity parameter \(\gamma=\gamma(t)\), including an ordinary differential equation (the probability flow ODE, corresponding to \(\gamma(t) = 0\)), can be used for sampling. A major motivation of this work is to provide insight into the experimental evidence that stochasticity can either benefit or harm performance [1]–[3].

We estimate the evolution, along the sampling process, of the Kullback-Leibler divergences (KL divergence), also known as relative entropies, between the generated and the target distributions, as they depend on the free stochasticity parameter. At each sampling instant, the KL evolution is affected both by accumulated errors (from the starting sampling distribution and/or previous steps of sampling) and by the error resulting from the approximation of the score function at the current instant. The interaction between these two kinds of error can be used to understand the influence (beneficial or detrimental) of stochasticity on the final performance, extending previous approaches to the study of stochasticity in diffusion model sampling [4]–[6]. We conduct numerical experiments to further investigate the effect of this balance of errors. We also present theoretical error estimates in a general setting, including estimates implicitly discussed in the literature, and use numerical experiments to illustrate the obtained theoretical results. Experiments are performed on toy data sets, where all quantities can be accurately computed; on typical image benchmarks such as MNIST and CIFAR-10; and also on a geological 3D problem, where, besides visual accuracy, it is important to match the distribution of rock statistics. Finally, we provide a detailed analysis of a fully analytical example.

Diffusion models are currently the state of the art for generative image and video synthesis [7], [8], and are also being used successfully in many scientific and technological applications, such as natural language processing [9], text-to-speech synthesis [10], medical imaging [11], time series analysis [12], protein modeling and design [13], [14], bioinformatics [15], global weather forecasting [16], porous media synthesis [17], and more [18]. Despite their remarkable success in applications, the precise evaluation of the performance of a generative model presents challenges such as the independent quantification of sample quality and distribution coverage [19], [20]. Informally, sample quality can be defined as the plausibility of generated samples in the context of original data, and distribution coverage as the amount to which the original distribution is captured by the generated distribution. Some examples of poor sample quality are low-resolution generated images and hallucinations [21], while examples of poor distribution coverage include mode collapse and lack of variability [17], [22], [23].

The KL divergence is a standard information-theoretic measure of discrepancy between probability distributions. Due to its asymmetry, it provides two distinct notions of approximation to compare a generated distribution \(\tilde{p}\) with a target distribution \(p\) through the divergences \(H(\tilde{p}|p)\) and \(H(p|\tilde{p})\). Both divergences measure generated sample quality and target distribution coverage, but \(H(\tilde{p}|p)\) places more weight on the former, whereas \(H(p|\tilde{p})\) emphasizes the latter. These properties are referred to in variational inference as zero-forcing and zero-avoiding behavior; see, for instance, [24] and [25]. The difference can be seen from the definition \(H(p_1|p_2):=\int p_1 \log (p_1/p_2)\,dx\): the integrand is large where \(p_2(x)\ll p_1(x)\), so that these regions contribute more to \(H(p_1|p_2)\) than to \(H(p_2|p_1)\). Different applications may be more sensitive either to sample quality or to distribution coverage, making both divergences relevant for evaluating a generative model from a theoretical perspective. Indeed, standard metrics for image generation, such as Inception Score (IS) or Fréchet Inception Distance (FID), typically measure both aspects in a single score [20].

Motivated by this, we investigate the impact of stochastic sampling on model performance through both divergences. In particular, the knowledge available for the target distributions \(p_t\) facilitates the use of log-Sobolev inequalities (LSI) for a more precise control of \(H(\tilde{p}_t|p_t)\) (associated with sample quality) from above. The analysis of \(H(\tilde{p}_t|p_t)\) complements the previous literature on KL divergence bounds for diffusion models, which typically considers only \(H(p|\tilde{p})\) for obtaining fully data-independent bounds (see Section 1.1). A more careful study of \(H(\tilde{p}|p)\) is also motivated by the observed occurrence of hallucinations in diffusion models [21]. In Section 4.3.3, we show an example where stochasticity can be beneficial for one divergence while being detrimental to the other.

The framework we consider is the one introduced by [2] and advanced by [1], with the diffusion of the data distribution defined by the (forward) SDE \[\label{sde} dX_t = f(X_t,t)\,dt + g(t)\,dW_t,\tag{1}\] in \(\mathbb{R}^n,\) on a time interval \([0, T],\) with probability density \(p_t(x),\) where \(f=f(x, t)\) and \(g=g(t)\) are given and \(\{W_t\}_{t\geq 0}\) is an \(n\)-dimensional Wiener process. The associated reverse-time SDE takes the form \[\label{rev32sde} d\tilde{X}_\tau = \left(-\bar{f}(\tilde{X}_\tau, \tau) + \frac{1}{2}\bar{g}^2(\tau)(1+\gamma(\tau))\nabla\log \bar{p}_\tau(\tilde{X}_\tau)\right)\,d\tau + \sqrt{\gamma(\tau)} \bar{g}(\tau) \,dW_\tau,\tag{2}\] in the reversed time variable \(\tau=T-t,\) where \(\bar{f}(x, \tau):=f(x,T-\tau)\), \(\bar{g}( \tau):=g(T-\tau)\), \(\bar{p}_\tau(x):=p_{T-\tau}(x),\) and where \(\gamma(\tau) \geq 0\) is arbitrary. It is well known [1], [26] that, under compatible initial conditions, the reverse-time SDE 2 has the same densities as the forward SDE 1 under the change of variables \(\tau=T-t\). This result is condensed in Lemma 1, with a short proof provided for the sake of completeness. Note that 2 gives a family of equations, indexed by the free functional parameter \(\gamma(\tau)\geq0\).

If the (Stein) score function \(\nabla\log p(t,\cdot)\) of the forward process given by 1 is known exactly or approximately, for all \(0<t\leq T,\) and if at time \(t=T\) the distribution \(p_T\) is close to a Gaussian distribution or any other distribution which is easy to sample from, the reverse-time SDE 2 can then be used as a generic sampler for \(p_0.\) In practice, the score function \(\nabla\log p(t,\cdot)\) can be approximated by a neural network \(s_\theta(\cdot, t)\) trained with denoising score-matching [1], [27]. Then, sampling is performed via the approximate reverse-time SDE \[\label{rev32sde32approx} d\tilde{X}_\tau = \left(-\bar{f}(\tilde{X}_\tau, \tau) + \frac{1}{2}\bar{g}^2(\tau)(1+\gamma(\tau))\bar{s}_\theta(\tilde{X}_\tau, \tau)\right)\,d\tau + \sqrt{\gamma(\tau)} \bar{g}(\tau) \,dW_\tau,\tag{3}\] where \(\bar{s}_\theta(x, \tau) = s_\theta(x, T - \tau)\) is the reverse-time approximate score. In the machine-learning literature [2], two common choices for \(\gamma\) correspond to setting \(\gamma\equiv0\) (Probability Flow ODE), and \(\gamma\equiv1\) (Anderson/Song’s reverse-time SDE).

In summary, the diffusion model algorithm consists of

1. choosing an adequate forward process 1 for which \(p_T\) is close to a known and easy-to-sample distribution \(q\);

2. training a score model \(s_\theta(x, t)\) to approximate \(\nabla\log p_t(x),\) for all \(0<t\leq T,\) by minimizing the global loss function of denoising score-matching;

3. sampling the starting condition \(\tilde{X}_0\) from the prior distribution \(q(x)\approx p_T(x);\) and

4. numerically integrating, from \(\tau=0\) to \(\tau=T,\) the approximate reverse-time equation 3 with the starting condition \(\tilde{X}_0.\)

This pipeline generates a final sample \(\tilde{X}_T\sim\tilde{p}_T\), with the aim that \(\tilde{p}_T\) be as close as possible to the unknown data distribution \(p_0\). Its main sources of error are

1. approximation errors of the score function;

2. approximation errors of the starting prior distribution \(p_T\);

3. numerical integration errors of the (approximate) reverse-time equation.

In this work, we investigate the effect of the free parameter \(\gamma(t)\) on the correction and propagation of errors from sources [errorscore] and [errorprior]. Thus, for theoretical results, we consider the exact solution of either 2 or 3 , while for the numerics we simply use Euler and Euler-Maruyama schemes with a very fine time mesh. Regarding discretization, it is well known that deterministic sampling has the advantage of achieving good performance even with a small number of steps, which makes it much faster in practice since the main computational burden of sampling is in the evaluation of the neural score function. Computational efficiency has also motivated the development of fast stochastic samplers [28]–[31]. In the literature, an implicit assumption often found is that the only hurdle to stochastic sampling comes from discretization, which has been shown in [1] to not be the case for all trained models. Therefore, the search for more efficient discretization schemes for both deterministic and stochastic sampling can benefit from a more careful study of whether stochasticity is advantageous in a continuous setting.

In Section 3.1, we consider the exact-score reverse-time SDE 2 with an approximate starting prior, isolating the source [errorprior] of error. In this case, Theorem 1 explicitly shows the advantage of stochasticity, in the form of KL divergence bounds which are decreasing in \(\gamma(t)\). A previous result has been obtained in [4], restricted to the choices of \(\gamma(\tau)=\gamma=1\) and \(\gamma(\tau)=\gamma=0\). As we discuss in Section 3.1.1, increasing \(\gamma\) makes the problem stiff, adding an effective faster time scale to the dynamics. This poses a challenge to the numerical approximation of the problem, limiting the choice of \(\gamma\) in practice.

Next, in Section 3.2, we consider the more realistic setting of an approximate score function \(s(x, t) = \nabla\log p_t(x) + \epsilon(x,t),\) where \(\epsilon=\epsilon(x,t)\) is the associated error, addressing both errors [errorprior] and [errorscore]. Here, stochasticity still produces a correction term in the evolution of the generated distribution, but it also amplifies the score error \(\epsilon\) by a multiplicative factor, making its overall impact more subtle. In Proposition 2, we explicitly state the conditions to ensure that the formula for the time derivative of the KL divergence holds. Then, if \(\bar{p}_\tau\) satisfies a log-Sobolev inequality with constant \(C_\text{LSI}(\tau),\) as in ?? , Theorem 2 yields the KL divergence estimate \[H(\tilde{p}_\tau|\bar{p}_\tau) \leq e^{-\int_0^\tau \alpha(s)\,ds} H(\tilde{p}_0|\bar{p}_0) + \frac{1}{2}\int_0^\tau \bar{g}(s)^2(1+\gamma(s))E_{\tilde{p}_s}\left[\epsilon_s\cdot\nabla\log h_s\right]e^{-\int_s^\tau \alpha(r)\,dr}\,ds,\] where \(\alpha(s)=C_\text{LSI}(s)^{-1} \gamma(s) \bar{g}(s)^2\) and \(h_s=\tilde{p}_s/\bar{p}_s\). When \(\gamma\) is strictly positive, a more explicit estimate is given by equation ?? , which eliminates the explicit dependence on \(\nabla\log h\) and depends on \(\epsilon\) only through its \(L^2(\tilde{p})\) norm.

In this setting, we obtain in Section 3.2.1 a qualitative criterion for the final impact of the stochastic term by minimizing the time derivative of the KL divergence. In broad terms, stochasticity is beneficial at sampling time \(\tau\) when the ratio between accumulated errors and the current score error surpasses a certain threshold. This would make the performance of stochastic sampling highly dependent on the time-profile of the score error: an error that increases too rapidly near the end of the sampling process would make deterministic sampling preferable, whereas an error concentrated near the beginning or middle of the process would allow the model to benefit from stochasticity. This generalizes to the non-asymptotic case (finite \(\gamma\), intermediate \(\tau\) values) the main conclusion in [5].

Section 4 shows numerical experiments on several data sets. First, we compare the estimated values of the KL divergences with the theoretical bounds, on two toy data sets with different available LSI constants, sampling with both exact and learned score functions. Then, for toy and also more complex data sets such as MNIST and CIFAR-10, we relate the insights of Section 3.2.1 to the time-profile of the score error of trained models, identifying factors that contribute to the observed differences in the effect of \(\gamma\) across trained models. Finally, in Section 4.3.4 we investigate the case of variable \(\gamma=\gamma(t)\), performing a grid-search over intervals to include stochasticity, considering step functions \(\gamma(t)\). Based on the results and the theoretical considerations, we believe that, in general, the best choice is to perform a grid-search near the end of sampling, searching for intervals \([t_1, t_2]\) with \(t_1>\epsilon\) and \(t_2\ll T\).

Finally, Section 5 presents the analysis of a fully analytical example, based on a single Gaussian data set with linear perturbations in the score. In this setting, the reverse-time SDE can be explicitly solved, all the relevant quantities can be computed, and it is possible to solve the optimal control problem for the choice of \(\gamma(t)\) minimizing each KL divergence. Remarkably, even in this simple setting the effect of \(\gamma\) is dependent on the relation between the score error and the prior error, which can also be thought of as the error accumulated from the initial part of sampling.

1.1 Related work↩︎

Some recent articles also obtain KL divergence bounds for score-based diffusion models. For instance, [32]–[34] provide complexity bounds for \(H(p_\text{data}|p_\text{generated})\), accounting for errors arising from the prior distribution approximation, score approximation and discretization errors. These works consider a forward Ornstein-Uhlenbeck diffusion, and the particular choice of \(\gamma\equiv1\). Working with \(H(p_\text{data}|p_\text{generated})\) makes the expectation over the score error to be taken with respect to \(p_\text{data}\) instead of \(p_\text{generated}\), which produces bounds that are entirely independent of the generated distribution. Although our motivation is different, there are similarities in the techniques, and, in particular, the formula in our Lemma 3 corresponds to the formula in Lemma C.1 of [32].

Bounds for \(H(p_\text{data}|p_\text{generated})\) are also obtained in [35], for more general forward processes in the context of stochastic interpolants. For a general reverse SDE with positive stochasticity parameter, it is shown that this divergence is bounded by the score-matching error, whereas a similar result is not guaranteed to hold for ODE sampling. A similar result has been previously obtained in [36] for \(\gamma\equiv1\). The papers [6], [37] extend the work of [35] by proposing a choice of \(\gamma(t)\) which minimizes their bound. In our formulation, this corresponds to setting \(\gamma\equiv 1\), as discussed in more detail in Remark 7.

For ODE sampling, bounds for both KL divergences can be found in [38]. These bounds are formulated in terms of the implicit score of the generated distribution and the learned score, and are used for establishing the convergence of a self-consistent training procedure instead of evaluating the impact of the score error on model performance.

Some recent works specifically analyze the effect of stochasticity at sampling. In [4], under the hypothesis of exact score functions, log-Sobolev inequalities for \(p_t\) are used to obtain a bound for \(H(p_\text{generated}|p_\text{data})\) which shows the decrease in this divergence along sampling when \(\gamma\equiv1\), while when \(\gamma\equiv0\) the divergence remains constant. In [5], an asymptotic analysis of \(H(p_\text{data}|p_\text{generated})\) with \((\text{score error})\to0\) is performed to compare the ODE with the asymptotic SDE \(\gamma\to\infty\). Their conclusion reflects a particular aspect of the analysis in Section 3.2.1: if the score errors are located in a region in the beginning or middle of the sampling process, maximum stochasticity is optimal, while stochasticity is altogether detrimental if the score errors are located exactly at the end of sampling.

In comparison, our work provides more complete KL bounds than those available in the literature, accounting for both divergences, general forward processes, and arbitrary stochasticity functions \(\gamma(t)\). We also rely less on formal arguments, explicitly stating the conditions for the results to hold. Moreover, the theoretical insights on the evolution of the KL divergences provide a connection between the time-profile of the score error and the effect of stochasticity in a non-asymptotic setting, which is closer to practice.

2 Background↩︎

We present here some background results to be used in the next sections. For completeness, proofs are provided in Appendix 7.1, since these results appear in the literature in a slightly different form.

2.1 Reverse-time formula↩︎

Consider a (forward) SDE on the interval \([0,T]\), of the form \[\label{sde322} dX_t = f(X_t,t)\,dt + g(t)\,dW_t,\tag{4}\] where \(f\) and \(g\) are smooth and \(f\) is globally Lipschitz with respect to \(x,\) uniformly on \(t\in [0, T]\), so that the corresponding initial-value problem admits a unique solution. Let \(p_t\) be its associated density, which solves the Fokker-Planck (FP) equation \[\label{fp} \frac{\partial p}{\partial t} = -\nabla\cdot(fp) +\frac{1}{2}g^2\Delta p.\tag{5}\] As mentioned in Section 1, we extensively use the following result from [1]. Earlier formulations can be found in [2], [26].

Since ?? is an SDE in reverse time that has the same densities as the forward SDE 4 , it is called a reverse-time SDE associated with 4 . For the sake of completeness, we include a short proof of this result, which can also be found in [1], in Appendix 7.1. The proof uses the same formal computation together with a result of [39] that guarantees the well-posedness of ?? .

2.2 Logarithmic Sobolev Inequalities (LSI)↩︎

We recall the definition of the logarithmic Sobolev inequality (LSI). The LSI is known to be important for the exponential convergence towards the asymptotic invariant measure for the Langevin equation and plays a key role in our estimates for the reverse-time SDE, as well.

More explicitly, we use the following consequence of the LSI.

For probability distributions in \(\mathbb{R}\), the LSI is equivalent to a more explicit integrability characterization [40]. In the general case, subgaussianity is a necessary condition, and there are widely known sufficient conditions, such as the Bakry-Emmery log-concavity criterion [41] and the Holley-Stroock bounded perturbation lemma [42]. For a Gaussian convolution of a given measure, a sufficient condition is that the measure has compact support [43], as we use in Section 3.

3 KL divergence evolution and bounds↩︎

It is well known that the probability densities \(q_t(x)\) associated with the time-independent Langevin equation \[\label{langevin} dX_t = \nabla\log \pi(X_t)\,dt + \sqrt{2}\,dW_t\tag{6}\] converge exponentially to the stationary distribution \(\pi\) when \(t\to\infty\), under some conditions on \(\pi\) [44], [45]. In particular, under the assumption that \(\pi\) satisfies an LSI with constant \(C\), it is obtained, in [45], an exponential bound on the KL divergence, given by \[\label{langevin32entropy32decrease} H(q_t|\pi)\leq e^{-\frac{2t}{C}}H(q_0|\pi).\tag{7}\] Thus, deviations in the initial distribution are corrected over time by the Langevin diffusion.

Considering now the reverse-time SDE ?? instead of the Langevin equation, we can aggregate the terms with \(\gamma\) and write it as \[\begin{gather} \label{rev32sde323} d\tilde{X}_\tau = \left(-\bar{f}(\tilde{X}_\tau, \tau) + \frac{1}{2}\bar{g}^2(\tau)\nabla\log \bar{p}_\tau(\tilde{X}_\tau)\right)\,d\tau \\ + \left( \frac{1}{2}\bar{g}^2(\tau)\gamma(\tau)\nabla\log \bar{p}_\tau(\tilde{X}_\tau)\,d\tau+ \sqrt{\gamma(\tau)} \bar{g}(\tau) \,dW_\tau\right). \end{gather}\tag{8}\] Notice that the first term on the right-hand side corresponds to the Probability Flow ODE, i.e. the reverse-time SDE with \(\gamma\equiv0\), and the second is closely related to the Langevin equation 6 . Inspired by this form, in [1] it is proposed that the reverse-time SDE should be seen as the Probability Flow ODE with the addition of a Langevin-like error-correcting term. This suggests that a bound similar to 7 might be found for the reverse-time SDE.

This idea can be formulated and studied more precisely by considering the KL divergences between two densities \(\tilde{p}_\tau, \bar{p}_\tau\) of the reverse-time SDE, one with a given initial density \(\tilde{p}_0\) and the other with \(\bar{p}_0=p_T\), matching the forward process. By Lemma 1, we know that the densities \(\bar{p}_\tau\) correspond to the densities of the forward process for all \(\tau\in[0, T)\). As discussed in Section 1.1, the KL divergences are usual metrics to evaluate the performance of the diffusion algorithm. In Section 3.1, we obtain a result analogous to 7 , with \(\bar{p}_\tau\) assuming the role of the stationary distribution \(\pi\). Section 3.2 studies the more general case of an approximated reverse-time SDE, in which \(\nabla\log \bar{p}_\tau(x)\) is replaced by \(s(x,\tau) := \nabla\log \bar{p}_\tau(x) + \epsilon(x, \tau)\), for an approximation error \(\epsilon(x, \tau)\).

3.1 Exact score functions↩︎

In this section, we consider only the approximation error in the initial condition, that is, when \(\tilde{p}_\tau\) and \(\bar{p}_\tau\) are solutions of the same equation ?? , but with different initial conditions \(\tilde{p}_0\) and \(\bar{p}_0=p_T\). The proofs in this and the next subsection rely on a general formula for the time derivative of the KL divergence between solutions of Fokker-Planck equations sharing a diffusion coefficient, which we state and prove in Appendix 7.2 (Lemma 3).

Proposition 1 shows that the KL divergence decreases as \(\tau\) increases whenever the stochasticity parameter \(\gamma(\tau)\) is positive. If \(\gamma\equiv0\), which is the case for the Probability Flow ODE, the divergence remains constant for \(\tau\in(0,T)\). However, this result alone is not sufficient to quantify the final divergence decrease due to the dependence of the right hand side of ?? on the unknown \(\tilde{p}_\tau\). This can be mitigated by the use of a logarithmic Sobolev inequality.

The following result is stated for forward SDEs with a homogeneous linear drift, in which case the Gaussian estimates on \(\bar{p}_\tau\) required by Proposition 1 are automatically satisfied. A homogeneous linear drift is also the standard choice in practice, since in this case the transition probabilities become Gaussian [1], [2], significantly simplifying the training process.

If \(p_0\) is sub-Gaussian, it is known that \(p_t\) satisfies an LSI at least for large \(t\) [43]. Moreover, when \(p_0\) has compact support, such as with image data, the distribution \(p_t\) satisfies an LSI for all \(t\in(0,T)\), and \(\nabla\log\bar{p}_\tau\) is automatically Lipschitz. This is summarized by the following corollary.

More general conditions for ensuring that \(\nabla\log p_t\) is Lipschitz can be found in [47].

3.1.1 Remark: the functional parameter \(\gamma\) at discretization↩︎

Since the bound ?? shows a KL divergence decay which is exponential on the stochasticity parameter \(\gamma\), what is the disadvantage of setting \(\gamma\) arbitrarily large?

The answer is related to the numerical approximation. A simple Euler-Maruyama discretization step of the reverse-time SDE ?? is given by \[\begin{align} X_{\tau_{i+1}} = & X_{\tau_i} +\left(-\bar{f}(X_{\tau_i}, \tau_i) + \frac{1}{2}\bar{g}^2(\tau_i)(1+\gamma(\tau_i))\nabla\log \bar{p}_{\tau_i}(X_{\tau_i})\right)\Delta\tau_i + \sqrt{\gamma(\tau_i)} \bar{g}(\tau_i) \sqrt{\Delta\tau_i}\xi \nonumber \\ = & X_{\tau_i} +\left(-\bar{f}(X_{\tau_i}, \tau_i) + \frac{1}{2}\bar{g}^2(\tau_i)\nabla\log \bar{p}_{\tau_i}(X_{\tau_i})\right)\Delta\tau_i \nonumber \\ & + \frac{1}{2}\bar{g}^2(\tau_i)\nabla\log \bar{p}_{\tau_i}(X_{\tau_i})(\gamma(\tau_i)\Delta\tau_i) + \bar{g}(\tau_i) \sqrt{\gamma(\tau_i)\Delta\tau_i}\xi \label{discr32rev32sde} \end{align}\tag{9}\] where \(\xi\sim\mathcal{N}(0,1)\), and \(\Delta \tau_i:=\tau_{i+1} - \tau_i\). Thus, comparing 9 with 8 , one can see that \(\gamma(\tau_i)\Delta\tau_i\) gives an effective discretization step size for the Langevin-like part of the reverse-time SDE, which, together with the standard step \(\Delta\tau_i,\) makes the problem stiff for large \(\gamma.\) Setting \(\gamma\) too large requires a very fine mesh, which will inevitably increase the computational cost and potentially produce a significant accumulation of discretization errors. In Section 4, we show the trade-off between error correction and discretization errors in numerical experiments; see Figure 2.

3.2 Approximate score functions↩︎

The previous bounds can be generalized to the case where the reverse-time SDE depends on an approximate score \(s(x, \tau)\). This models the setting of \(s(x, \tau)\) being a neural network trained by score-matching, as in the diffusion model algorithm described in Section 1. By writing \(s(x, \tau) = \nabla\log \bar{p}_\tau(x) + \epsilon_\tau(x)\), with an associated error function \(\epsilon_\tau(x)\) representing the difference between the exact and approximate scores, the approximated version of equation ?? is given by \[\label{rev32sde32232pert} d\tilde{X}_\tau = \left(-\bar{f}(\tilde{X}_\tau, \tau) + \frac{1}{2}\bar{g}^2(\tau)(1+\gamma(\tau))\left(\nabla\log \bar{p}_\tau(\tilde{X}_\tau) + \epsilon_\tau(\tilde{X}_\tau)\right)\right)\,d\tau + \sqrt{\gamma(\tau)} \bar{g}(\tau) \,dW_\tau,\tag{10}\] with the corresponding Fokker-Planck equation \[\label{pert32rev95fp} \frac{\partial \tilde{p}}{\partial \tau} = -\nabla\cdot\left(\left(-\bar{f} + \frac{1}{2}\bar{g}^2(\gamma+1)(\nabla\log \bar{p} + \epsilon)\right)\tilde{p}\right) +\frac{1}{2}\gamma \bar{g}^2\Delta \tilde{p}.\tag{11}\] Then, we can use Lemma 3 (Appendix 7.2) to obtain straightforward generalizations of Proposition 1 and Theorem 1.

A corresponding result for the framework of stochastic interpolants can be found in [35], Lemma 2.22.

Corollary 2 shows that, for stochastic sampling, it is possible to bound the KL divergence \(H(\bar{p}_\tau|\tilde{p}_\tau)\) by a multiple of the score error plus the initial error, as has been observed in the literature [35], [36]. This is, in fact, a consequence of the divergence decay phenomenon for the SDE, which for \(H(\tilde{p}|\bar{p})\) can be better quantified with the use of a log-Sobolev inequality, in the following generalization of Theorem 1.

3.2.1 Minimizing the KL divergence derivative at each \(\tau\)↩︎

Since the formula for the KL divergence time derivative obtained in Proposition 2 is linear in \(\gamma\), we can readily obtain the values of \(\tau\) for which a positive \(\gamma\) can produce a decrease in the divergence. Observe that this instantaneous minimum may not correspond exactly to the perfect-foresight minimum, since, for reverse-times \(\tau_1<\tau_2\), the density \(\tilde{p}_{\tau_2}\) also depends on \(\gamma(\tau_1)\). However, the problem of finding the perfect-foresight minimum is infinite-dimensional and typically intractable, except in very specific settings, such as the analytical example considered in Section 5. In this section, we explore the information provided by the instantaneous optimum, while keeping in mind its limitations.

More precisely, writing equation ?? as \[\frac{d}{d\tau}H(\tilde{p}_\tau|\bar{p}_\tau) = \frac{1}{2}\bar{g}^2\int\epsilon\cdot \left(\nabla\log \frac{\tilde{p}_\tau}{\bar{p}_\tau}\right)\tilde{p}_\tau\; \,dx +\frac{1}{2}\gamma \bar{g}^2 \int\left( \epsilon\cdot \nabla\log \frac{\tilde{p}_\tau}{\bar{p}_\tau} - \left|\nabla\log\frac{\tilde{p}_\tau}{\bar{p}_\tau}\right|^2\right) \tilde{p}_\tau \,dx,\] we can see that \(dH/d\tau\) is minimized at each \(\tau\) by the bang-bang optimum \[\label{optimal95g95instantaneous} \gamma^*(\tau) = \begin{cases} \gamma_{\max}, & \text{when}\quad E_{\tilde{p}_\tau}[\epsilon\cdot \nabla\log (\tilde{p}_\tau/\bar{p}_\tau) - |\nabla\log (\tilde{p}_\tau/\bar{p}_\tau)|^2]<0 \\ 0, & \text{when}\quad E_{\tilde{p}_\tau}[\epsilon\cdot \nabla\log (\tilde{p}_\tau/\bar{p}_\tau) - |\nabla\log (\tilde{p}_\tau/\bar{p}_\tau)|^2]>0, \end{cases}\tag{13}\] where \(\gamma_{\max}\) is the stipulated maximum \(\gamma\) value, limited by the discretization issues discussed in 3.1.1. We can identify two main factors at play:

• Alignment of \(\epsilon_\tau\) and \(\nabla\log (\tilde{p}_\tau/\bar{p}_\tau)\): As observed in Remark 5, the alignment between \(\epsilon_\tau(x)\) and \(\nabla\log (\tilde{p}_\tau/\bar{p}_\tau)(x)\), in terms of the inner product \(E_{\tilde{p}_\tau}[\epsilon_\tau\cdot\nabla\log (\tilde{p}_\tau/\bar{p}_\tau)]\), is a major factor in the accumulation of errors. This is also true here, since a low enough alignment, even with large \(E[|\epsilon_\tau|^2]\), can produce a negative expectation in 13 , favoring the SDE. However, it is not clear how to evaluate \(\epsilon_\tau\cdot \nabla\log (\tilde{p}_\tau/\bar{p}_\tau)\) in practice, since it depends on the unknown \(\nabla\log \tilde{p}_\tau\).

• Relative magnitudes of \(\epsilon_\tau\) and \(\nabla\log (\tilde{p}_\tau/\bar{p}_\tau)\): If \(|\epsilon_\tau|(x)\) is smaller on average than \(|\nabla\log (\tilde{p}_\tau/\bar{p}_\tau)|(x)\), the sign of the expectation in 13 must be negative. More precisely, we have that \[\int\left( \epsilon_\tau\cdot \nabla\log \frac{\tilde{p}_\tau}{\bar{p}_\tau} - \left|\nabla\log\frac{\tilde{p}_\tau}{\bar{p}_\tau}\right|^2\right) \tilde{p}_\tau \,dx \leq \int\left( |\epsilon_\tau| - \left|\nabla\log\frac{\tilde{p}_\tau}{\bar{p}_\tau}\right|\right)\left|\nabla\log\frac{\tilde{p}_\tau}{\bar{p}_\tau}\right| \tilde{p}_\tau \,dx <0\] if \(\int|\epsilon_\tau|\tilde{\mu}_\tau dx<\int|\nabla\log \tilde{p}_\tau/\bar{p}_\tau|\tilde{\mu}_\tau dx\), where \(\tilde{\mu}_\tau := \tilde{p}_\tau|\nabla\log \tilde{p}_\tau/\bar{p}_\tau|\) ¹. Since \(\nabla\log \tilde{p}_\tau/\bar{p}_\tau\) is the difference between the implicit score of \(\tilde{p}_\tau\) and the true score of \(\bar{p}_\tau\), it can be thought as a measure of the accumulated error. In short, if, at time \(\tau\), the average magnitude of the score error is lower than a threshold given by the accumulated error, stochasticity is beneficial.

This last point shows that the time-localization of the magnitude of the score error can be a crucial factor for the effect of stochasticity, extending the conclusion obtained in [5] to the non-asymptotic case, i.e. for intermediate \(\tau\) and \(\gamma\) values. In particular, score errors concentrated in the beginning of sampling would make stochasticity beneficial, whereas score errors with magnitude increasing too fast near the end of sampling could make it detrimental. These aspects are investigated in the experiments of Section 4.3.2.

4 Numerical experiments↩︎

In this section, we explore the relation between stochasticity at sampling and model performance in several experimental scenarios. We consider two toy data sets for which the exact score functions can be computed and more complex data sets such as MNIST, CIFAR-10, and a three-dimensional porous media data set [17]. Sampling with exact score functions, on the toy data sets, is explored in Section 4.2. Then, Section 4.3 explores sampling with learned score models on both toy and complex data sets, covering the setting of Section 3.2.

The experiments were performed using the diffusion model package ‘diffsci’, with the code available at https://github.com/Lacadame/DiffSci/tree/main/stochasticity_paper.

4.1 Experimental details↩︎

4.1.1 Sampling parameters↩︎

We performed experiments using the three forward processes considered in [1]: EDM, Variance Exploding (VE) and Variance Preserving (VP), which are given by the forward SDEs \[\begin{align} dX_t & = \sqrt{2t}\,dW_t , \\ dX_t & = \,dW_t,\\ dX_t & = -\frac{1}{2}(\beta_1 t+\beta_2)X_t \,dt + \sqrt{\beta_1 t+\beta_2}\,dW_t . \end{align}\] Note that the three equations satisfy the linear hypothesis in Theorems 1 and 2. To avoid confusion, in this section we will refer only to the forward time variable \(t\), and consider the densities \(p_t\) and \(\tilde{p}_t\) in forward time. In this way, diffusion sampling starts at an "initial time" \(t=T\) and ends at \(t=0\).

In order to maintain discretization errors independent of the choice of \(\gamma\), we perform numerical integration (at sampling) with Euler and Euler-Maruyama methods for ODEs and SDEs, respectively. Since the main focus of this work is not discretization, as discussed in the Introduction, by default we use \(1000\) discretization steps for CIFAR-10 and \(500\) for the remaining data sets, following the EDM choice of steps \(\{t_i\}\) in [1]. In these settings, we observed qualitatively similar results from the three forward processes. This agrees with findings that major differences between the processes arise only at low \(n_\text{steps}\), with EDM yielding the lowest discretization error [1]. Therefore, we only show experiments with the EDM process.

4.1.2 Datasets↩︎

The toy data distributions consist of a mixture of Gaussians or a mixture of uniform distributions, for which exact score functions can be analytically obtained. This allows sampling with exact scores and also computing the score error function \(\epsilon_t(x) = s(x, t) - \nabla\log p_t(x)\) considered in Section 3.2. To maximize the estimation accuracy of the KL divergences, we focus on a one-dimensional setting. More explicitly, we consider the toy data distributions \(p_0\) and \(q_0\) in \(\mathbb{R}\), consisting of mixtures \[\begin{align} p_0 & \sim \lambda \mathcal{N}(\mu_1, \sigma_1^2) + (1 - \lambda)\mathcal{N}(\mu_2, \sigma_2^2), \tag{14} \\ q_0 & \sim \lambda \mathcal{U}(a_1, b_1) + (1 - \lambda)\mathcal{U}(a_2, b_2), \tag{15} \end{align}\] where \(\mathcal{N}(\mu, \sigma^2)\) is a Gaussian distribution of mean \(\mu\) and variance \(\sigma^2\), and \(\mathcal{U}(a, b)\) is a Uniform distribution in the interval \([a, b]\subset\mathbb{R}\). For linear forward processes, the score functions \(\nabla\log p_t\) can be analytically computed, and are Lipschitz continuous (see Appendix 7.3 for the mixture of Gaussians, and Corollary 1 for the mixture of Uniforms), satisfying the conditions on \(\bar p_\tau\) in Theorems 1 and 2; the conditions on \(\tilde{p}_\tau\) are discussed in Remark 1. For training score models, we consider \(4000\) samples of the toy distributions. At sampling, we consider \(10^5\) generated samples.

We also consider the more complex data sets MNIST [48], CIFAR-10 [49], and a geological data set consisting of 3D volumes of micro-CT scans of porous media [17]. Due to the high dimensionality of the data, it is not possible to reliably estimate KL divergences. Instead, we evaluate performance via FID scores, for the benchmark data sets, and KL divergences of geological statistics, for the porous media data set.

4.1.3 Estimation of KL divergences and bounds↩︎

KL divergences are computed in two ways. The initial divergence \(H(q_T|\rho_T)\), where \(q\) and \(\rho\) correspond either to \(\tilde{p}\) or \(\bar p\), is obtained by numerical integration of the densities, which can be calculated analytically for the toy data sets. Intermediate divergences \(H(q_t|\rho_t)\) are computed from samples of \(q_t\) and \(\rho_t\) via the discretization of the empirical densities by a histogram². To reduce the dependence on the parameter number of histogram bins, we consider the mean of the values obtained with bins in the range \(\{n_\text{bins}-20, ..., n_\text{bins}\}\). We found this estimator to have a precision on the order of \(5\times 10^{-4}\).

The bounds of Section 3 are computed with a different LSI constant for each toy data set. For the mixture of Uniforms, we use the constant for compactly-supported measures stated in Corollary 1, which are from [43]. For the mixture of Gaussians we use the constants obtained in [50], detailed in the following. The \(\chi^2\)-divergence of one-dimensional Gaussian distributions \(\mu_0 \sim \mathcal{N}(m_0,\sigma_0^2)\) and \(\mu_1 \sim \mathcal{N}(m_1,\sigma_1^2)\) has the closed form \[\chi^2(\mu_0|\mu_1)=\dfrac{\sigma_1^2}{\sigma_0\sqrt{2\sigma_1^2-\sigma_0^2}}\exp\left(\dfrac{(m_0-m_1)^2}{2\sigma_1^2-\sigma_0^2}\right)-1\] provided that \(2\sigma_1^2>\sigma_0^2\). Then, from [50] we know that the LSI constant \(C_p\) for the mixture of Gaussians \(\mu_p = p \mu_0 + (1-p) \mu_1\) satisfies \[\label{GaussianLSB} C_p\leq\min\{C_0, C_1\},\tag{16}\] with \[\begin{align} C_0 & = \max\{\sigma_0^2 (1 + (1-p)\lambda_p), \; \sigma_1^2 (1+p\lambda_p \chi_0)\}, \\ C_1 & = \max\{\sigma_1^2 (1 + p\lambda_p), \; \sigma_0^2 (1+(1-p)\lambda_p \chi_1)\}, \end{align}\] where \[\begin{align} \lambda_p := \left\{ \begin{array}{cl} 2 &, \;p=\frac{1}{2}, \\ \frac{\log p - \log(1-p)}{ 2p - 1} &,\;p\neq \frac{1}{2}, \end{array} \right. \end{align}\] \[\chi_0 := \chi^2(\mu_0||\mu_1) + 1 \qquad\text{and}\qquad \chi_1 := \chi^2(\mu_1||\mu_0) + 1.\] A plot of LSI constants across diffusion time \(t\) can be seen in Figure 1.

4.2 Exact score functions↩︎

First, we generate distributions with exact analytical scores, in the setting of Section 3.1. Figure 1 shows generated distributions and KL divergence evolution, for sampling from a perturbed prior distribution at time \(T=80\). We artificially perturb the mean and the variance of the prior distribution so that the decrease in the divergences is more perceptible and better quantifiable.

We can see a substantial decrease in the divergences for the SDE, higher than that ensured by the bound for \(H(\tilde{p}| p)\), and very similar across both divergences. Moreover, the bound for the mixture of Gaussians (MoG) is tighter than the one for the mixture of Uniforms (MoU), due to the different known LSI constants, whereas the decrease in the divergences is remarkably similar. This might reflect the fact that the LSI constant of \(p_t\) is valid for any \(q_t\), but in the time-derivative formula it is applied to a very specific \(q_t=\tilde{p}_t\), which is close to \(p_t\). Observe that the small mismatch between initial and ODE divergence values is due to estimation errors for the latter, since by Proposition 1 they should be the same. Observe also that the final jump in \(H(\tilde{p}|p)\) for the MoU can be attributed to the sensitivity of the KL divergence estimator for \(t\approx0\), since \(\tilde{p}_0\) may not be absolutely continuous with respect to \(p_0\) for this data set.

Figure 2 illustrates the relation between \(\gamma\) and the number of discretization steps, discussed in Section 3.1.1. Both panels show the same final KL divergences with different horizontal axes: on the left, we see an optimal interval for \(\gamma\) for each choice of \(n_\text{steps}\), on the right, we see that the increase in errors for large \(\gamma\) is a function of the "effective step size" \(\gamma/n_\text{steps}\propto\gamma\Delta t\) discussed in 3.1.1. More precisely, equation 9 gives two effective step sizes: \(\Delta \tau_i\) for the Flow ODE part and \(\gamma(\tau_i)\Delta \tau_i\) for the Langevin-like part, but for large \(\gamma\) the errors from the latter dominate. In the plot on the right, we see that for smaller \(n_\text{steps}\) the Flow ODE effective step size still plays a significant role, with the curves not following exactly the Langevin effective step size scaling.

4.3 Learned score functions↩︎

Here, we consider sampling with learned score functions, in the setting of Section 3.2, for the mixture of Gaussians and the more complex data sets. For the mixture of Gaussians, we train a multilayer perceptron (MLP) of layers \([128, 128, 128]\) with SiLU activations and sinusoidal time-embeddings, with the EDM training configurations [1]. For MNIST, we train a UNet with 45M parameters, while for CIFAR-10 we use pretrained models from [1], and for the porous media example we use pretrained models from [17].

The mixture of Gaussians allows a direct computation of the score error \(\epsilon_t(x)\), since the exact score \(\nabla\log p_t(x)\) is also available. For the more complex data sets, the true score function is unknown, but we can still compare the score matching (SM) error \(E_{\bar{p}_t}[|\epsilon_t|^2]\) between different models. This can be done via the denoising score matching (DSM) errors, since SM and DSM errors differ only by a model-independent constant [27].

4.3.1 Prior error correction↩︎

First, we evaluate the ability of reverse SDEs with approximate scores to correct errors from the prior distribution. As in Section 4.2, we consider perturbed prior distributions at time \(T=80\), here perturbing only its standard deviation from the default choice \(\sigma(T)=T\) matching the forward process. Figure 3 (b) shows a performance comparison across constant values of \(\gamma\), for a model trained on the mixture of Gaussians and a model trained on MNIST. We can see that above a certain \(\gamma\) the SDE essentially eliminates the error from the prior, with final performance independent of the prior choice. Figure 3 (c) shows, for the mixture of Gaussians, the evolution of the KL divergences and of the bounds from Section 3, in the case of \(\sigma(T)=96\) and \(\gamma=1\). We see an initial correction of the prior distribution error by the SDE, followed by the accumulation of score errors near the end of the sampling process, which is observed for the ODE as well. Observe that the bound for \(H(\tilde{p}|p)\), from Theorem 2 (equation ?? ) tracks the initial decrease in the divergences, unlike the bound for \(H(p|\tilde{p})\), from Corollary 2. The first bound is computed with an instantaneous optimal \(\delta=\delta(\tau)\) obtained by minimizing the bound on \(dH/d\tau\) 12 at each time step \(\tau_i\).

4.3.2 Time-profile of error magnitude↩︎

We now move to a more realistic setting, where the initial error is very small, and focus on score approximation errors. Motivated by the analysis in Section 3.2.1, we test how the time-profile of error magnitude interacts with the effect of \(\gamma\) on final performance. Figure 4 (b) shows three illustrative examples of score models with the same network architecture, trained with the same hyperparameters and having approximately the same validation loss values. Comparing the error profiles of model 1 and 2, we see that model 1 has a higher \(L^2\) norm of the error for \(t>4\) and lower in the intermediate interval \([0.01, 0.5]\). Based on the reasoning from Section 3.2.1, we would expect stochastic correction to be more effective for model 1, and this is indeed observed in the \(\text{KL}\times\gamma\) plot. For model 3, the error near the end of sampling (\(t<0.1\)) is smaller than for the others, so from Section 3.2.1 we would expect it to produce a smaller amplification of errors for large \(\gamma\), and this is also observed in the \(\text{KL}\times\gamma\) plot. As an illustration of the well-known fact that deterministic sampling is more robust to the reduction of the number of steps, Figure 4 (c) shows the impact of \(\gamma\) across several \(n_\text{steps}\) values. This happens even for model 3, which benefits from stochasticity at large \(n_\text{steps}\), and should be more pronounced when using a fast deterministic sampler instead of Euler’s method.

Complementing these examples, we investigate the relationships outlined in the previous paragraph across all the epoch checkpoints for three training runs of the same model. We extract, from the \(\text{KL}\times\gamma\) and error profile curves, scalars capable of quantifying specific aspects of these curves, and measure correlations. For the influence of stochasticity, we consider the relative improvement \(I:=(H_0 - H_{\min})/ H_{\min}\) and the relative deterioration \(D:=(H_{\gamma_{\max}} - H_{\min})/ H_{\min}\), where \(H_0\) denotes a final KL divergence at \(\gamma=0\), \(H_{\gamma_{\max}}\) a final KL divergence at \(\gamma=\gamma_{\max}\), and \(H_{\min}\) the lowest across all \(\gamma\in[0, \gamma_{\max}]\). For the error profiles, we consider the ratios of integrated errors \[E_\text{early/mid}:=\frac{\int_{t_\text{early}}^T E_{p_t}[|\epsilon_t|^2]dt}{\int_{t_\text{mid}}^{t_\text{early}}E_{p_t}[|\epsilon_t|^2]dt} \qquad \text{and}\qquad E_\text{late/tot}:=\frac{\int_0^{t_\text{late}} E_{p_t}[|\epsilon_t|^2]dt}{\int_0^T E_{p_t}[|\epsilon_t|^2]dt},\] with threshold parameters \(t_\text{early}\), \(t_\text{mid}\) and \(t_\text{late}\). Figure 4 (d) shows scatter plots of \(I \times E_\text{early/mid}\) and \((\log D) \times E_\text{late/tot}\), for both KL divergences, along with values of Pearson and Spearman coefficients and the obtained regression lines. We see a moderate positive correlation, where the remaining unexplained variance could be associated with more complex interactions of the time-profile of the error. The values of \(t_\text{early}\), \(t_\text{mid}\) and \(t_\text{late}\) were chosen to maximize the sum of Pearson and Spearman correlation coefficients on a different set of checkpoints. A plot of the sensitivity of the correlations on the threshold values, for both the main and the held-out checkpoints, can be seen in Figure 15. The correlation values are stable along small perturbations of the thresholds, but the most significant choices still may depend on the training hyperparameters and the data set. Thus, we regard this experiment more as a quantitative proof-of-concept than as providing general scalar metrics for any trained model. A more systematic exploration is left for future work.

Next, we consider examples of models on benchmark data sets. Since the DSM losses only give relative information on the error, we plot the relative error differences to the best model, which is taken as baseline. For the MNIST data set (Figure 5 (b)), we see the same pattern as before: model 1 has a larger error for \(t\in[10, 40]\), and shows performance improvement by positive \(\gamma\). For model 2, with a smaller error for large \(t\), this effect is not observed, while for model 3, with a smaller error for small \(t\), the improvement reappears. The same pattern is found for the CIFAR-10 data set with pretrained models differing only in their training preconditioning [1] (Figure 5 (d)), where the larger initial (\(t>1\)) error of the VP model is associated with an improvement for positive \(\gamma\).

Additionally, Figure 6 shows examples of latent diffusion models trained on 3D volumes of porous media, in an application presented in [17]. Performance in this case is measured by the distribution of scalar geophysical properties, for which the KL divergences can be computed. This case illustrates the deterioration aspect: the larger error of model 1 for \(t<1\) is associated with the error amplification for large \(\gamma\) observed in all the metrics.

4.3.3 Example: different impacts on the KL divergences↩︎

Closer inspection of the \(\text{KL}\times\gamma\) curves for model 1, in Figure 4 (b), reveals that, for a certain interval of \(\gamma\), the SDE can produce lower \(H(\tilde{p}|p)\) but higher \(H(p|\tilde{p})\) when compared to the ODE. Figure 7 shows the generated distributions and the evolution of the KL divergences for one such case. We can see a higher presence of outliers (samples in the low-density region of the data distribution) for the ODE, and lower mode coverage (lack of samples in some regions of the data distribution) for the SDE. This agrees with the usual interpretation of the KL divergences, mentioned in the Introduction.

4.3.4 Variable \(\gamma(t)\)↩︎

The dependence of the relative improvement and relative deterioration on the score error within certain intervals suggests that the benefits from stochasticity can be maximized, and the detrimental effects mitigated, by setting a positive \(\gamma\) only within certain intervals. Motivated by this and by the bang-bang structure of the instantaneous optimum obtained in Section 3.2.1, we perform grid searches for the best interval \([S_{\min}, S_{\max}]\) to set a constant positive \(\gamma\).

Figure 8 shows FID values for the three MNIST models discussed in Section 4.3.2, across several intervals \([S_{\min}, S_{\max}]\) and two choices of \(\gamma\). First, notice that the minimum FID scores for models 2 and 3 are substantially lower than their minimum FID scores for constant \(\gamma\) (Figure 5 (b)). We also see that stochasticity only in the beginning of sampling is always detrimental, with the largest deterioration for model 1, which has the largest initial score errors (Figure 5 (b)). This suggests that the SDE is amplifying score errors in the beginning of sampling, but it is not active in the end to correct the accumulated errors. In the case of \(\gamma=5\), we also observe a deterioration when stochasticity is used only over the very end of sampling, which is also consistent with the positive correlation between the relative deterioration \(D\) and the late/total error ratio \(E_\text{late/tot}\) found in the previous section.

Corresponding experiments for the mixture of Gaussians and CIFAR-10 data sets can be found in Figures 16 and 17. Although the worst stochasticity intervals vary significantly across models and data sets, the pattern for the most beneficial intervals appears to be more stable. For models which benefit significantly from stochasticity, we always see an optimal stochasticity interval located close to (but not at) the end of sampling, with \(S_{\max}\leq1\) and \(S_{\min}>0\). The idea of grid-searching for optimal stochasticity intervals is not new [1], but the existence of an optimal stochasticity interval with \(S_{\max}\leq1\) and \(S_{\min}>0\) might be a more general feature of diffusion sampling, given that it also aligns with the time-profile analysis explored in Section 3.2.1. If true, it would allow grid-searching for an optimal interval across a limited range of intervals, given a trained model. However, the experimental evidence provided here is still preliminary.

5 A fully analytical example↩︎

In this section, we consider an analytical example for which we can explicitly compute the sampling distributions, KL divergences, the first bound of Theorem 2, and all other relevant quantities. This toy setting allows a direct analysis of the effect of the parameter \(\gamma\), including an optimal control solution for \(\gamma(\tau)\) at the end of the section. We will consider linear perturbations in the score and the prior distribution, and analyze how these errors interact with each other. Although linear perturbations do not accurately reflect the neural network errors found in practice, this is the only setting for which the reverse-time SDE can be analytically solved. In the following analysis, observe that the error attributed to the prior can also be seen as the error arising from inaccurate scores at the early phase of sampling.

5.1 One-dimensional case↩︎

We start with the one-dimensional case, assuming that the “data” distribution is a Gaussian \(X_0 \sim \mathcal{N}(\mu_0, \sigma_0^2)\) and considering a forward equation of the form \[dX_t = g(t)\,dW_t,\] on a time interval \(0\leq t \leq T\). The forward process is a Gauss-Markov process \(X_t \sim \mathcal{N}(\mu_0, \sigma(t)^2),\) with score function \[\label{scorefullyanalytic} \nabla \log p(x, t) = -\frac{x - \mu_0}{\sigma(t)^2}, \qquad \sigma(t)^2 = \sigma_0^2 + \sigma_{\text{diff}}(t)^2, \quad \sigma_{\text{diff}}(t)^2 = \int_0^t g(s)^2 ds.\tag{17}\] We consider an approximate score of the form \[s_\theta(x, t) = -\frac{x - \mu_\theta}{\sigma_\theta(t)^2}, \qquad \sigma_\theta(t)^2 = \alpha_\theta \sigma(t)^2,\] with constant parameters \(\mu_\theta\in\mathbb{R}\) and \(\alpha_\theta > 0.\) The error to the exact score, introduced in Section 3.2, becomes \[\epsilon_t(x) = s_\theta(x, t) - \nabla \log p(x, t) = \frac{\mu_\theta - \mu_0}{\sigma_\theta(t)^2} + \left(1 - \frac{\sigma(t)^2}{\sigma_\theta(t)^2} \right)\frac{x - \mu_0}{\sigma(t)^2}.\]

This approximate score yields the family of linear approximate reverse equations \[\label{AE:pertubedequation} d\tilde{X}_\tau = -\frac{1}{2}\bar{g}^2(\tau)(1+\gamma(\tau))\left( \frac{\tilde{X}_\tau - \bar\mu_\theta(\tau)}{\bar\sigma_\theta(\tau)^2}\right) \,d\tau + \sqrt{\gamma(\tau)} \bar{g}(\tau) \,dW_\tau,\tag{18}\] in the interval \(0\leq\tau\leq T\). Since this is a linear equation, it can be solved explicitly [51], yielding the solution \[\tilde{X}_\tau = e^{-A(\tau)} \tilde{X}_0 + e^{-A(\tau)} \int_0^\tau \bar\mu_\theta(s) a(s) e^{A(s)} \;\mathrm{d}s + e^{-A(\tau)} \int_0^\tau \sqrt{\gamma(s)}\bar g(s) e^{A(s)} \;\mathrm{d}W_s,\] where \[A(\tau) = \int_0^\tau a(s)\;\mathrm{d}s, \qquad a(\tau) = \frac{1}{2} \frac{\bar g(\tau)^2 (1 + \gamma(\tau))}{\bar\sigma_\theta(\tau)^2}.\] Finally, we consider a prior, for the reverse process, which is also Gaussian, \(\tilde{X}_0 \sim \mathcal{N}(\mu_T, \sigma_T^2),\) with \(\sigma_T^2 = \beta_T \sigma(T)^2,\) for a given parameter \(\beta_T > 0\). Thus, we obtain \[\tilde{\mu}_\theta(\tau) = e^{-A(\tau)}\mu_T + e^{-A(\tau)} \int_0^\tau \bar\mu_\theta(s) a(s) e^{A(s)} \;\mathrm{d}s\] and \[\tilde{\sigma}_\theta(\tau)^2 = e^{-2A(\tau)}\sigma_T^2 + e^{-2A(\tau)} \int_0^\tau \gamma(s) \bar g(s)^2 e^{2A(s)}\;\mathrm{d}s.\]

5.1.1 Constant \(\gamma(t)=\gamma\)↩︎

For more explicit calculations, here we assume \[\label{parametersanalytic} \gamma(t) = \gamma, \quad \forall t \geq 0,\tag{19}\] with a given constant \(\gamma \geq 0.\) Thus, the free (constant) parameters for the approximate reverse process are the shape-changing parameters \(\alpha_\theta, \beta_T > 0,\) the translation parameters \(\mu_\theta, \mu_T \in \mathbb{R},\) the stochasticity parameter \(\gamma \geq 0,\) and the starting time \(T > 0.\) Besides these, we also have the parameters \(\mu_0, \sigma_0\) defining the data distribution.

Then, the mean-square error in the score becomes \[\mathbb{E}_{\bar p_\tau}[\tilde{\epsilon}_\tau^2] = \frac{(\mu_\theta - \mu_0)^2}{\alpha_\theta^2\bar\sigma(\tau)^4} + \left(1 - \frac{1}{\alpha_\theta} \right)^2 \frac{1}{\bar\sigma(\tau)^2}.\] Since the variance of the forward diffusion increases with time, the error in the score decreases, which coincides with the behavior in the numerical example, where the score is approximated by a neural network trained with denoising score-matching (see e.g. Figure 3).

With constant \(\gamma\), the mean and the variance of the Gauss-Markov approximate reverse process become \[\label{meananalytic} \tilde{\mu}_\theta(\tau) = \mu_\theta + \left(\frac{\bar\sigma(\tau)^2}{\bar\sigma(0)^2}\right)^{\frac{1}{2} \frac{(1 + \gamma)}{\alpha_\theta}}(\mu_T - \mu_\theta)\tag{20}\] and \[\begin{align} \label{varianceanalytic} \tilde{\sigma}_\theta(\tau)^2 & = \begin{cases} \bar\sigma(\tau)^2 \left( \frac{\gamma\alpha_\theta}{(1 + \gamma) - \alpha_\theta} + \left(\beta_T - \frac{\gamma\alpha_\theta}{(1 + \gamma) - \alpha_\theta}\right)\left(\frac{\bar\sigma(\tau)^2}{\bar\sigma(0)^2}\right)^{\frac{(1 + \gamma)}{\alpha_\theta} - 1} \right), & 1 + \gamma \neq \alpha_\theta, \\ \bar\sigma(\tau)^2\left( \beta_T +\gamma\ln\left(\frac{\bar\sigma(0)^2}{\bar\sigma(\tau)^2}\right)\right), & 1 + \gamma = \alpha_\theta, \end{cases} \end{align}\tag{21}\] where the apparent singularity as \(\gamma \rightarrow \alpha_\theta - 1\) in the first expression is a removable singularity.

With that, we compute both KL divergences \[\begin{align} \tag{22} H(\tilde{p}_\tau|\bar{p}_\tau) & = \frac{1}{2}\left(\ln {\frac{\bar\sigma(\tau)^2}{\tilde{\sigma}_\theta(\tau)^2}}+{\frac{\tilde{\sigma}_\theta(\tau)^{2}+(\mu_{0}-\tilde{\mu}_\theta(\tau))^{2}}{\bar\sigma(\tau)^{2}}}-1 \right), \\ \tag{23} H(\bar{p}_\tau|\tilde{p}_\tau) & = \frac{1}{2}\left(\ln {\frac{\tilde{\sigma}_\theta(\tau)^2}{\bar\sigma(\tau)^2}}+{\frac{\bar\sigma(\tau)^{2}+(\mu_{0}-\tilde{\mu}_\theta(\tau))^{2}}{\tilde{\sigma}_\theta(\tau)^{2}}}-1 \right). \end{align}\]

In Figure 9, we can see the evolution of \(H(\tilde{p}_\tau|\bar{p}_\tau)\) and of the two bounds of Theorem 2, for a particular choice of prior and score errors, for different values of \(\gamma\). Note that the second bound can, in some cases, be tighter than the first one, even though its derivation uses one further inequality. This is possible since the extra inequality is applied before the LSI, as discussed in Remark 6.

We now use formulas 22 and 23 to better understand the effect of \(\gamma\) on the KL divergences. We summarize, below, the findings discussed in more detail in Appendix 7.4. Assuming, for simplicity, that the mean is preserved, i.e. \(\tilde{\mu}_\theta(\tau) = \mu_0\), stemming from the choices \(\mu_T = \mu_\theta = \mu_0\), the KL divergences reduce, after sampling (corresponding to \(t = 0\) or, equivalently, \(\tau = T\)), to \[H(\tilde{p}_T | p_0) = \frac{v(\gamma) - \ln v(\gamma) - 1}{2} \qquad \textrm{and} \qquad H(p_0 | \tilde{p}_T) = \frac{\frac{1}{v(\gamma)} + \ln v(\gamma) - 1}{2}\] where \[v(\gamma) = \frac{\tilde{\sigma}_\theta(T)^2}{\bar\sigma(T)^2} = \begin{cases} \frac{\gamma\alpha_\theta}{(1 + \gamma) - \alpha_\theta} + \left(\beta_T - \frac{\gamma\alpha_\theta}{(1 + \gamma) - \alpha_\theta}\right)\zeta_T^{\frac{(1 + \gamma)}{\alpha_\theta} - 1}, & 1 + \gamma \neq \alpha_\theta, \\ \beta_T - \gamma\ln\zeta_T, & 1 + \gamma = \alpha_\theta, \end{cases}\] with \(\zeta_T = \bar\sigma(T)^2/\bar\sigma(0)^2 = \sigma_0^2 / ( \sigma_0^2 + \sigma_{\text{diff}}(T)^2 ) < 1.\)

If the score is exact and only the prior is off, meaning \(\alpha_\theta = 1\) and \(\beta_T \neq 1,\) we have \(v(\gamma) = 1 + (\beta_T - 1)\zeta_T^{\gamma} \neq 1\), and \(v(\gamma)\) converges monotonically and exponentially towards \(1\) as \(\gamma \rightarrow \infty\). In this case, both KL divergences decay exponentially to zero, as \(\gamma\) increases, as expected. This means stochastic sampling is always better than ODE sampling when the score is exact, as discussed before.

With \(\alpha_\theta \neq 1,\) the behavior is more involved. For \(\gamma = 0,\) we simply have \(v_0:=v(0) = \beta_T \zeta_T^{(1 - \alpha_\theta)/\alpha_\theta},\) while, asymptotically in \(\gamma,\) we have \(v_\infty := \lim_{\gamma\rightarrow\infty} v(\gamma) = \alpha_\theta.\) Looking at the ratio \(v_0/v_\infty = (\beta_T/\alpha_\theta)\zeta_T^{(1 - \alpha_\theta)/\alpha_\theta}\), we see that \(v_0 < v_\infty\) if, and only if, either \(\zeta_T < (\alpha_\theta/\beta_T)^{\alpha_\theta / (1 - \alpha_\theta)},\) when \(\alpha_\theta < 1\), or \(\zeta_T > (\beta_T/\alpha_\theta)^{\alpha_\theta / (\alpha_\theta - 1)},\) when \(\alpha_\theta > 1\). Taking into account that \(\zeta_T < 1\) and looking at \(v'(\gamma)\), we see that, when \(\beta_T \leq \alpha_\theta < 1,\) the variance ratio \(v(\gamma)\) increases monotonically from \(v_0\) to \(v_\infty,\) approaching \(v_\infty=\alpha_\theta\) from below, so both KL divergences decrease monotonically with \(\gamma.\) Similarly, when \(\beta_T \geq \alpha_\theta > 1,\) the variance ratio \(v(\gamma)\) decreases monotonically from \(v_0\) toward \(v_\infty = \alpha_\theta,\) so both KL divergences decrease monotonically toward \(h_\infty.\) In both cases, increasing \(\gamma\) is beneficial, and \(\gamma \rightarrow \infty\) is optimal. These two regimes occur when the training and the prior errors are aligned, i.e. either both underestimate or both overestimate the exact variances, provided also that the error in the prior is relatively more pronounced.

Analogous behaviors occur when considering only errors in the mean, and they are also aligned, i.e. when either \(\mu_T \leq\mu_\theta\leq \mu_0\) or \(\mu_T \geq\mu_\theta\geq \mu_0\). Then, increasing \(\gamma\) is beneficial, and \(\gamma \rightarrow \infty\) is optimal (see Appendix 7.4.3). In more general scenarios, the behavior is more nuanced, and the KL divergences may not be monotonic, and may achieve local or global maxima or minima at finite values of the diffusion parameter. In some cases, the error in the prior may compensate the error in the score, and sampling with the reverse ODE may be better. In other cases, the best sampling occurs at an intermediate value of the stochasticity parameter. Figure 10 illustrates a few different scenarios, where the 2 leftmost plots of the first row are the aligned cases mentioned above. Interestingly, we can also see that the performance of the SDEs seems to be more stable to variations of the parameters, as they converge quickly in \(\gamma\) to the value of \(\alpha_\theta\), independently of the other parameters.

With the same scenarios, Figure 11 shows the evolution of the mean \(\tilde{\mu}_\theta(t)\) and the standard deviation \(\tilde{\sigma}_\theta(t)\) in the phase-space plane \((\mu, \sigma)\) of the analytical solution, given by 20 and 21 . In the first two columns, variance error from the prior and the score are aligned, either underestimating or overestimating the true value, and the error accumulation makes stochastic sampling better. In the remaining plots, the two errors are in different directions, yielding non-monotonic curves in \(\gamma\), but only in the last one the prior error and score errors balance out in a way that the reverse ODE is always better.

Interestingly, the alignment of the prior error with the score error is related to the alignment of \(\epsilon\) and \(\nabla\log\tilde{p}_\tau/\bar{p}_\tau\), discussed in Remark 5 and also Section 3.2.1. For \(\mu_0=0\), at the initial time \(\tau=0\) we have \[E_{\tilde{p}_0}\big[\epsilon_0\cdot\nabla\log(\tilde{p}_0/\bar{p}_0)\big] = \frac{1}{\alpha_\theta\sigma(T)^4}\left(\mu_\theta \mu_T + \mu_T^2(\alpha_\theta-1) + (\alpha_\theta-1)(\beta_T-1)\sigma(T)^2\right).\] The first term is negative if the means are in opposite directions, i.e. \(\mu_\theta\mu_T<0\), and the third term is negative if the variances are in opposite directions, i.e. \((\alpha_\theta-1)(\beta_T-1)<0\). The second term gives the more subtle interaction between score variance and prior mean, and it can become large when both the error in the prior mean and the error in the score variance are large, which intuitively corresponds to the under-correction of the prior error on the mean by too wide score fields. Nevertheless, for typical values of \(\sigma(T)\) and \(\mu_T\) the third term becomes dominant, making the variance alignment the dominant factor. This is exemplified in the insets of Figure 11 showing the sign of \(\epsilon\nabla\log\tilde{p}_\tau/\bar{p}_\tau\). When the variances of the errors are misaligned, we see that the score error initially reduces the KL (negative \(\epsilon\nabla\log\tilde{p}_\tau/\bar{p}_\tau\)), until the accumulated error \(\nabla\log\tilde{p}_\tau/\bar{p}_\tau\) switches to the direction of the score error.

5.1.2 Variable \(\gamma(t)\)↩︎

We may also consider a time-varying \(\gamma = \gamma(t)\) and look for the function that minimizes each terminal KL divergence. This can be formulated as an optimal control problem with terminal cost, imposing an upper bound \(\gamma(t)\in [0, \gamma_{\textrm{max}}]\) on the control variable. As detailed in Appendix 7.5, the control problem in this specific analytical example turns out to be linear in \(\gamma\), and applying Pontryagin’s maximum principle reveals that the optimal solution is a “bang-bang” control: the optimal \(\gamma(t)\) takes only two possible values, alternating between the probability flow ODE and stochastic sampling with maximal stochasticity parameter. Moreover, the solution admits at most one switch. Remarkably, this structure is independent of the specific terminal cost considered, being a consequence of the underlying linear dynamics of the state space arising from the linear perturbation on the score.

Figure 10 also shows, in the dashed lines, KL divergences for the optimal \(\gamma^*(t)\), where the switching time \(\tau^*\) is obtained by numerically searching in the interval \([0, T]\). We can see that the optimal \(\gamma^*(t)\) yields lower KL divergences than any constant \(\gamma\) when the mean errors are in different directions but the variance errors are not. This global bang-bang optimum can also be compared with the instantaneous optimum discussed in Section 3.2.1, which results from minimizing the KL divergence time derivative at each instant \(\tau\). This is done in Figure 12 for the previous choices of parameters, where we see that in some cases the final divergences of the instantaneous optimum agree with those of the global optimum, but that is not always the case. This indicates some limitations of the reasoning developed in Section 3.2.1, at least in this simplified setting.

5.2 Multidimensional case↩︎

We now extend this example to the multivariate case, with the “data” distribution being a Gaussian \(X_0 \sim \mathcal{N}(\mu_0, \Sigma_0)\) in \(\mathbb{R}^n\). To simulate a \(d\)-dimensional “data” manifold, we let the covariance be block-diagonal: \[\Sigma_0 = \begin{pmatrix} {\sigma_0'}^2 I_d & 0 \\ 0 & {\sigma_0''}^2 I_{n-d} \end{pmatrix},\] with the variances satisfying \({\sigma_0''}^2 \ll {\sigma_0'}^2\). The first \(d\) components represent the data manifold, embedded on a high-dimensional ambient space with dimension \(n \gg d.\)

The forward equation is essentially the same, \(dX_t = g(t)\,dW_t,\) so the forward process remains Gaussian with covariance \(\Sigma(t) = \Sigma_0 + \sigma_{\text{diff}}(t)^2 I_n\), where \(\sigma_{\text{diff}}(t)^2\) is as given in 17 .

The exact score function is linear and splits into two independent blocks: \[\nabla \log p(x, t) = -\Sigma(t)^{-1}(x - \mu_0) = \begin{pmatrix} -\frac{x' - \mu'_0}{{\sigma'}(t)^2} \\ -\frac{x'' - \mu''_0}{{\sigma''}(t)^2} \end{pmatrix},\] where \({\sigma'}(t)^2 = {\sigma_0'}^2 + \sigma_{\text{diff}}(t)^2\) and \({\sigma''}(t)^2 = {\sigma_0''}^2 + \sigma_{\text{diff}}(t)^2\).

We consider an approximate score \(s_\theta(x, t)\) modeled as a linear perturbation of the exact score. We introduce a diagonal error matrix \(\Lambda_\theta = \text{diag}(\alpha' I_d, \alpha'' I_{n-d})\) and approximate parameters \(\mu_\theta\), yielding \[s_\theta(x, t) = - \Lambda_\theta^{-1} \Sigma(t)^{-1} (x - \mu_\theta) = \begin{pmatrix} -\frac{x' - \mu'_\theta}{\alpha' {\sigma'}(t)^2} \\ -\frac{x'' - \mu''_\theta}{\alpha'' {\sigma''}(t)^2} \end{pmatrix}.\] This formulation allows for different score accuracies on the manifold versus the ambient space, depending on \(\mu_\theta', \alpha'\) and \(\mu_\theta'', \alpha''\). This modeling allows for simulating and investigating the situation in which learned scores first align to the manifold and later learn finer details of the distribution, as suggested in a recent work [52].

The form of the perturbed score yields a linear reverse SDE. The prior is taken without knowledge of the manifold and is an isotropic Gaussian \(\tilde{X}_0 \sim \mathcal{N}(\mu_T, \sigma_T^2 I_n)\), and we write \(\sigma_T^2 = \beta'_T\sigma'(T)^2,\) for a parameter \(\beta'_T\).

Since the matrices \(\Sigma_0\) and \(\Lambda_\theta\) commute (they are both diagonal), the system decouples into two independent blocks: the manifold block \(\tilde{X}'\) and the ambient block \(\tilde{X}''\). The corresponding variances evolve according to the scalar formulas derived for the 1D case, shown in equation 21 . For the manifold block (assuming \(1+\gamma \neq \alpha'\)), we have \[\tilde{\sigma}'_\theta(\tau)^2 = \bar{\sigma}'(\tau)^2 \left( \frac{\gamma\alpha'}{(1 + \gamma) - \alpha'} + \left(\beta'_T - \frac{\gamma\alpha'}{(1 + \gamma) - \alpha'}\right)\left(\frac{\bar{\sigma}'(\tau)^2}{\bar{\sigma}'(0)^2}\right)^{\frac{(1 + \gamma)}{\alpha'} - 1} \right).\] An analogous formula holds for the ambient variance \(\tilde{\sigma}''_\theta(\tau)^2\) using parameter \(\alpha''\) and the dependent parameter \(\beta''_T:=\beta'_T\sigma'(T)^2/\sigma''(T)^2\).

Since the marginals are independent, the total KL divergence is the sum of the divergences of the two orthogonal subspaces, which are then multiples of the scalar components: \[H(\tilde{p}_\tau|\bar{p}_\tau) = d h'(\tau) + (n-d) h''(\tau),\] where \(h'(\tau)\) and \(h''(\tau)\) are the corresponding scalar Gaussian KL divergence contribution, given by 22 , with the appropriate parameters.

Since \(n-d \gg d\) in typical generative modeling scenarios, the global KL divergence is heavily weighted on the ambient term, but that can be counteracted by a relatively smaller error \(h''(\tau)\), on the direction normal to the manifold. Nevertheless, this reveals a fundamental trade-off: tune \(\gamma\) to \(\alpha'\) to improve the distribution within the manifold, as characterized here by \(\sigma_0'\), or tune \(\gamma\) to \(\alpha''\) to improve the location of the manifold, as characterized here by \(\sigma_0''\). Figure 13 exemplifies this trade-off: in the first two plots, the two optimal constant-\(\gamma\) values are close, whereas in the last two plots they are not.

As for the optimal time-dependent stochasticity parameter \(\gamma=\gamma(t)\), the control problem has the same essential structure, with the same bang-bang solution. However, since the switching function is now a sum of two terms, one for each block, with different exponential rates, it can change sign at most twice, yielding an optimal schedule with at most two switches, e.g., deterministic, then stochastic, then deterministic again, reflecting the competing requirements of the two groups. If we extend the analysis allowing the approximate score to have \(d\) different score error factors \(\alpha_i\), the same structure persists, with at most \(d\) switches, as detailed in Appendix 7.5.4.

6 Conclusions↩︎

In this work, we investigated the effect of stochasticity on the generated distribution through the lens of KL divergence evolution along sampling. At each instant, the time derivative of the divergences in stochastic sampling (Proposition 2) reveals a fundamental trade-off between the correction of accumulated errors and the amplification of current score errors. From this general property, more specific conclusions are derived. When sampling with exact scores, only the error correction part is present, and log-Sobolev inequalities can be used to bound the divergence decay from above (Theorem 1). The same technique produces bounds in the case of approximate scores (Theorem 2), but the additional term associated with the score error prevents the derivation of an explicit divergence decay.

Minimizing the time derivative of the divergences at each instant with respect to \(\gamma\) might not correspond to a minimum for the final KL divergence but provides a complementary perspective to the bounds. In particular, we show that stochasticity reduces the divergence at a given instant when the norm of the score error at that instant is smaller than a threshold depending on the magnitude of the accumulated errors (Section 3.2.1). We further investigate this point in numerical experiments, and evaluate correlations between the time-profile of the model’s score error and the impact of stochasticity on its performance (Section 4.3.2). Moreover, grid searches over intervals that incorporate stochasticity, with models trained on different data sets, indicate that a family of intervals near the end of sampling produces the best results (Section 4.3.4). This is consistent with the analysis in Section 3.2.1, and we conjecture that it might be a more general property of diffusion models, although the exact optimal interval could depend on the data set and the architecture considered.

Acknowledgment↩︎

The authors acknowledge the support of ExxonMobil Exploração Brasil Ltda. and Agência Nacional de Petróleo, Gás Natural e Biocombustíveis (ANP) through the grant no. 23789-1.

7 Additional proofs↩︎

We include, in this appendix, the details of a number of results mentioned in the main part of the paper.

7.1 Section 2↩︎

We start with the proofs of the results mentioned in Section 2, which are stated in slightly different forms than what is found in the literature.

7.2 Section 3↩︎

The proofs of Propositions 1 and 2 use the following lemma, whose formula appears in [32]. We restate it at the level of solutions of the Fokker-Planck equations, with explicit conditions under which the formula holds.

The verification of the hypotheses of Propositions 1 and 2 for the density of the (approximate) reverse-time SDE is discussed in Remark 1 and deferred to future work. In the remainder of this section, we collect the verifications used for \(\bar p\) in the linear-drift case.

7.3 Section 4↩︎

As discussed in the last section (proof of Theorem 1), \(p_t\) can be written as the Gaussian convolution \(p_t=p_0^t\ast \mathcal{N}(0, s(t)^2\sigma(t)^2I)\), where \(p_0^t(x):=s(t)^{-n}p_0(x/s(t))\). If \(p_0\) is the mixture of Gaussians \(p_0=\sum_{l=1}^n w_l\mathcal{N}(\mu_l, \sigma_l^2I)\), we have \(p_0^t=\sum_{l=1}^n w_l\mathcal{N}(s(t)\mu_l, s(t)^2\sigma_l^2I)\) and thus \[\begin{align} p_t & = p_0^t\ast \mathcal{N}(0, s(t)^2\sigma(t)^2I) \\ & = \sum_{l=1}^n w_l\mathcal{N}(s(t)\mu_l, s(t)^2\sigma_l^2I) \ast \mathcal{N}(0, s(t)^2\sigma(t)^2I) \\ & = \sum_{l=1}^n w_l\mathcal{N}(s(t)\mu_l, s(t)^2(\sigma_l^2+ \sigma(t)^2)I), \end{align}\] is also a mixture of spherical Gaussians. Therefore, the following lemma ensures that \(\nabla\log p_t\) is Lipschitz for all \(t\in[0,T]\).

7.4 Section 5: constant \(\gamma\)↩︎

The sample quality and sample coverage KL divergences 22 and 23 can be written more succinctly as \[h_q(\gamma) = H(\tilde{p}_T | p_0) = \frac{1}{2} \left( v(\gamma) - \ln v(\gamma) - 1 + m(\gamma)^2 \right),\] and \[h_c(\gamma) = H(p_0 | \tilde{p}_T) = \frac{1}{2} \left(\frac{1}{v(\gamma)} + \ln v(\gamma) - 1 + \frac{m(\gamma)^2}{v(\gamma)}\right),\] where \[v(\gamma) = \frac{\tilde{\sigma}_\theta(T)^2}{\bar\sigma(T)^2} = \begin{cases} \frac{\gamma\alpha_\theta}{1 + \gamma - \alpha_\theta} + \left(\beta_T - \frac{\gamma\alpha_\theta}{1 + \gamma - \alpha_\theta}\right)\zeta_T^{\frac{1 + \gamma}{\alpha_\theta} - 1}, & 1 + \gamma \neq \alpha_\theta, \\ \beta_T - \gamma\ln\zeta_T, & 1 + \gamma = \alpha_\theta, \end{cases}\] and \[m(\gamma) := \frac{\mu_0 - \tilde{\mu}_\theta(T)}{\bar\sigma(T)} = \frac{\mu_0 - \mu_\theta - (\mu_T - \mu_\theta)\zeta_T^{\frac{1}{2}\frac{1+\gamma}{\alpha_\theta}}}{\sigma_0}\] with \(\zeta_T = \bar\sigma(T)^2/\bar\sigma(0)^2 = \sigma_0^2 / ( \sigma_0^2 + \sigma_{\text{diff}}(T)^2 ) < 1.\)

7.4.1 Exact score↩︎

When the score for the approximate reverse problem is exact, i.e. when \(\alpha_\theta = 1\) and \(\mu_\theta = \mu_0,\) we have \[v(\gamma) = 1 + \left(\beta_T - 1\right)\zeta_T^\gamma, \qquad m(\gamma) = \frac{\mu_0 - \mu_T}{\sigma_0}\zeta_T^{\frac{1 + \gamma}{2}},\] for any \(\gamma \geq 0.\) Since \(\zeta_T < 1,\) they both converge monotonically to \(1\) and \(0\), respectively. Moreover, the ratio \(m(\gamma)^2/v(\gamma)\) is either zero for all \(\gamma\geq 0\), which happens when \(\mu_T = \mu_0\), or decreases monotonically to \(0,\) as \(\gamma \rightarrow \infty\) (indeed, with \(\mu_T \neq \mu_0\), just note that the reciprocal is of the form \(v(\gamma)/m(\gamma)^2 = A\zeta^{-\gamma} + B\), with \(A\) and \(B\) independent of \(\gamma\) and \(A\) strictly positive). Thus, both KL divergences \(h_q(\gamma)\) and \(h_c(\gamma)\) are monotonically decreasing to \(0\) as \(\gamma \rightarrow \infty.\) This means, as expected, that, when the score is exact, stochastic sampling is better than reverse ODE sampling.

7.4.2 Exact mean with aligned variances↩︎

When \(\alpha_\theta \neq 1\) but \(\mu_T = \mu_\theta = \mu_0,\) the mean of the approximate reverse problem is exact, \(\tilde{\mu}(\tau) = \mu_0,\) so that \(m(\gamma) = 0,\) for all \(\gamma > 0.\) In this case, \(h_q(\gamma)\) and \(h_c(\gamma)\) depend only on \(v(\gamma),\) and their critical points occur when either \(v'(\gamma) = 0\) or \(v(\gamma) = 1.\) For the analysis of the critical points, we change the variable to \(u = (1 - \alpha_\theta)/\alpha_\theta + \gamma/\alpha_\theta,\) which increases linearly with \(\gamma = \gamma(u) = \alpha_\theta u + \alpha_\theta - 1\), and write \[\begin{align} \tilde{v}(u) = v(\gamma(u)) & = \frac{\gamma(u)}{u} + (\beta_T - \frac{\gamma(u)}{u}) \zeta_T^u \\ & = \alpha_\theta + \frac{\alpha_\theta - 1}{u} + (\beta_T - \alpha_\theta - \frac{\alpha_\theta - 1}{u}) \zeta_T^u \\ & = \alpha_\theta + (\beta_T - \alpha_\theta)\zeta_T^u + (\alpha_\theta - 1)\frac{1 - \zeta_T^u}{u} \\ & = \alpha_\theta + (\beta_T - \alpha_\theta)\zeta_T^u + (\alpha_\theta - 1)\psi(u), \end{align}\] where \[\psi(u) = \frac{1 - \zeta_T^u}{u}.\] Notice that both \(\zeta_T^u\) and \(\psi(u)\) are positive, decreasing, and convex in \(u\).

When \(0 < \beta_T \leq \alpha_\theta < 1,\) the second term is non-positive and the third term is negative, so that \(v(\gamma) = \tilde{v}(u(\gamma)) < \alpha_\theta < 1,\) for all \(\gamma \geq 0,\) and increasing in \(\gamma.\) Thus, both \(h_q(\gamma)\) and \(h_c(\gamma)\) are decreasing in \(\gamma,\) and we have \[\label{kl95limits95analytical} \begin{align} h_q(\gamma)\to & \left( \alpha_\theta - \ln \alpha_\theta - 1\right)/2 > 0, \\ h_c(\gamma)\to & \left( \alpha_\theta^{-1} + \ln \alpha_\theta - 1\right)/2 > 0, \end{align}\tag{29}\] as \(\gamma \rightarrow \infty.\)

If \(\beta_T \geq \alpha_\theta > 1,\) the second term is nonnegative and the third term is positive, so that \(v(\gamma) = \tilde{v}(u(\gamma)) > \alpha_\theta > 1,\) for all \(\gamma \geq 0,\) and decreasing in \(\gamma.\) Thus, both \(h_q(\gamma)\) and \(h_c(\gamma)\) are decreasing in \(\gamma\) and have the same limits 29 .

7.4.3 Exact variances↩︎

When \(\beta_T = \alpha_\theta = 1\), we have \(v(\gamma) = 1,\) for all \(\gamma \geq 0,\) and thus \(h_q(\gamma) = h_c(\gamma) = m(\gamma)^2/2.\) In this case, we have a trichotomy depending on the values of \(A:=\mu_0-\mu_\theta\) and \(B:=\mu_T-\mu_\theta\). If \(AB\leq0\), corresponding to the orderings \(\mu_T\leq\mu_\theta\leq\mu_0\) or \(\mu_T\geq\mu_\theta\geq\mu_0\), then \(m(\gamma)^2\) is decreasing, hence both \(h_q(\gamma)\) and \(h_c(\gamma)\) decrease monotonically to \((\mu_0 - \mu_\theta)^2/2\sigma_0^2\), in which case SDE sampling is always better than ODE sampling and improves with increased \(\gamma\). If \(0<A/B<\zeta_T^{1/2}\), then \(m(\gamma)^2\) has a unique zero at a finite nonzero value of \(\gamma,\) and thus perfect generation (zero values for both KL divergences) is achieved at this value of \(\gamma.\) Stochastic sampling with the stochasticity parameter near this optimal value outperforms both ODE sampling and highly stochastic sampling, while some values of \(\gamma\) away from the optimum may perform worse than ODE sampling. Finally, if \(\zeta_T^{1/2}\leq A/B\), then \(m(\gamma)^2\) is increasing, in which case ODE sampling is always better than SDE sampling.

7.4.4 Aligned variances and means↩︎

When we have \(0 < \beta_T \leq \alpha_\theta < 1\) together with \(\mu_T \leq \mu_\theta \leq \mu_0\), then \(v(\gamma)\) increases monotonically to \(1\) while \(\mu(\gamma)\) decreases monotonically to \(0\) and then \(m(\gamma)^2/v(\gamma)\) also decreases monotonically to zero. In this case, both KL divergences decrease monotonically as \(\gamma \rightarrow \infty\), meaning that stochastic sampling is always better than reverse ODE sampling, improving with increased \(\gamma.\) If we have different combinations of either \(0 < \beta_T \leq \alpha_\theta < 1\) or \(\beta_T \geq \alpha_\theta > 1,\) together with either \(\mu_T \leq \mu_\theta \leq \mu_0\) or \(\mu_T \geq \mu_\theta \geq \mu_0,\) then we do not have control of the term \(m(\gamma)^2/v(\gamma)\), and we can only guarantee that the KL divergence \(h_q(\gamma)\) decreases monotonically as \(\gamma \rightarrow \infty.\)

7.4.5 Other cases↩︎

In other cases, however, the different terms have different and competing behaviors and the KL divergences may have local or global minima or maxima at finite nonzero values of \(\gamma,\) as illustrated in Figure 10.

7.5 Section 5: optimal control↩︎

Following the setup of Section 5, we now consider a stochasticity parameter \(\gamma(\tau) \in [0, \gamma_{\max}]\). From equations 22 and 23 , the final KL divergences depend on \(\gamma\) through the terminal mean \(\tilde{\mu}_\theta(T)\) and variance \(\tilde{\sigma}_\theta(T)^2\). To formulate an optimal control problem with terminal cost, we consider the state space \((m,u)\), where \(m(\tau):=\tilde{\mu}_\theta(\tau)\) and \(u(\tau):=\tilde{\sigma}_\theta(\tau)^2\).

7.5.1 Reduction to clock time↩︎

Since the approximate reverse SDE 18 is linear, the evolution of the mean and variance of its solution \(\tilde{p}_\tau\) is given by the ODEs \[\begin{align} \dot{m}(\tau) &= -\frac{\bar g(\tau)^2(1+\gamma(\tau))}{2\alpha_\theta\bar\sigma(\tau)^2}\left(m - \mu_\theta\right), &m(0) &= \mu_T,\\ \dot{u}(\tau) &= -\frac{\bar g(\tau)^2(1+\gamma(\tau))}{\alpha_\theta\bar\sigma(\tau)^2}\,u(\tau) + \gamma(\tau)\,\bar g(\tau)^2, \qquad &u(0) &= \beta_T\,\bar\sigma(0)^2. \end{align}\]

We consider the clock time \[s(\tau):=\bar\sigma(0)^2 -\bar\sigma(\tau)^2,\] and the variance (in the new variable) \(\phi(s) := \bar\sigma(0)^2 - s\). Since \(ds/d\tau = -d(\bar\sigma(\tau)^2)/d\tau = \bar g(\tau)^2\), in clock time the dynamics become \[\begin{align} m'(s) &= -\frac{1+\gamma(s)}{2\alpha_\theta\,\phi(s)}\left(m - \mu_\theta\right), &m(0) &= \mu_T, \tag{30} \\ u'(s) &= -\frac{(1+\gamma(s))\,u(s)}{\alpha_\theta\,\phi(s)} + \gamma(s), \qquad &u(0) &= \beta_T\,\phi(0), \tag{31} \end{align}\] making it independent of the specific choice of \(g\).

It is convenient to work with the normalized variance ratio \[\label{eq:vdef} v(s) := \frac{u(s)}{\phi(s)},\tag{32}\] which satisfies \[\label{eq:v95s} v'(s) = \frac{1}{\phi(s)}\left[\,v(s)\left(1 - \frac{1+\gamma(s)}{\alpha_\theta}\right) + \gamma(s)\,\right].\tag{33}\]

7.5.2 Optimal control formulation↩︎

Let \(\Phi(m(s_T), u(s_T))\) denote either of the terminal KL divergences. We formulate the problem of minimizing the terminal cost as a standard optimal control problem for \(s\in[0, s_T]\), with states \(m\) and \(u\), and control \(\gamma \in [0, \gamma_{\max}]\).

The control Hamiltonian is given by \[\label{eq:ham} H(s, m, u, \gamma, p_m, p_u) = - p_m\frac{1+\gamma}{2\alpha_\theta\phi(s)}\left(m - \mu_\theta\right) + p_u\left(-\frac{(1+\gamma)u}{\alpha_\theta\phi(s)} + \gamma\right),\tag{34}\] with costates \(p_m(s)\) and \(p_u(s)\) satisfying \[\label{eq:costate} p_m'(s) = \frac{(1+\gamma(s))\,p_m(s)}{2\alpha_\theta\phi(s)}, \qquad p_u'(s) = \frac{(1+\gamma(s))\,p_u(s)}{\alpha_\theta\phi(s)},\tag{35}\] and terminal transversality conditions \[\label{eq:trans} p_m(s_T) = \partial_m \Phi(m(s_T), u(s_T)), \qquad p_u(s_T) = \partial_u \Phi(m(s_T), u(s_T)).\tag{36}\] Observe that, by 30 and 35 , we have that \[\label{eq:conserved95C} \frac{d}{ds}\left[p_m(s)\big(m(s) - \mu_\theta\big)\right] = p_m'(s)\big(m(s) - \mu_\theta\big) + p_m(s)m'(s) = 0.\tag{37}\] We therefore denote \(C := p_m(s)\big(m(s) - \mu_\theta\big)\).

The Hamiltonian is linear in \(\gamma\), so by Pontryagin’s minimum principle the optimal control is bang-bang with switching function \[\label{eq:switch} \psi(s) := p_u(s)\left(1 - \frac{u(s)}{\alpha_\theta\phi(s)}\right) - \frac{p_m(s)\big(m(s) - \mu_\theta\big)}{2\alpha_\theta\phi(s)} = p_u(s)\left(1 - \frac{v(s)}{\alpha_\theta}\right) - \frac{C}{2\alpha_\theta\phi(s)},\tag{38}\] and optimal control \[\gamma^*(s) = \begin{cases} 0 & \text{if } \psi(s) > 0, \\ \gamma_{\max}& \text{if } \psi(s) < 0, \\ \text{singular} & \text{if } \psi(s) = 0 \text{ on a nontrivial interval}. \end{cases}\]

Observe that the costate ODE 35 is linear and homogeneous in \(p_u\), so \(p_u\) never vanishes unless it is identically zero. In particular, it never changes sign.

7.5.3 At most one switch↩︎

Interestingly, this property is not restricted to KL divergences as a terminal cost.

7.5.4 Multidimensional case (at most \(N\) switches)↩︎

In a general \(N\)-dimensional model, the state space expands to \((\mathbf{m}, \mathbf{u}) = (m_1, \dots, m_N, u_1, \dots, u_N)\) with decoupled dynamics across dimensions, and the control Hamiltonian is given by \[H(s, \mathbf{m}, \mathbf{u}, \gamma, \mathbf{p_m}, \mathbf{p_u}) = \sum_{i=1}^N \Bigg[ - p_{m_i}\frac{1+\gamma}{2\alpha_{\theta,i}\phi(s)}\left(m_i - \mu_{\theta,i}\right) + p_{u_i}\left(-\frac{(1+\gamma)u_i}{\alpha_{\theta,i}\phi(s)} + \gamma\right) \Bigg].\] By the same costate dynamics as in the unidimensional case, the quantities \(C_i = p_{m_i}(s)\big(m_i(s) - \mu_{\theta,i}\big)\) are conserved for each dimension \(i \in \{1, \dots, N\}\). The switching function therefore becomes a sum of the unidimensional switching components: \[\label{eq:switch95multi95N} \psi(s) = \sum_{i=1}^N \underbrace{\left[ p_{u_i}(s)\left(1 - \frac{v_i(s)}{\alpha_{\theta, i}}\right) - \frac{C_i}{2\alpha_{\theta, i}\phi(s)} \right]}_{:= \psi_i(s)}.\tag{41}\]

Then we have the following extension of the unidimensional one-switch result.

8 Scalar counterexample to the LSI↩︎

For a probability measure on \(\mathbb{R}\), the LSI is equivalent to the following integrability characterization, given by [40].

Finally, note that a mollification of this \(p(x)\) would have the same properties while being smooth (as a solution of the reverse-time SDE should be).

9 Additional figures↩︎

We conclude with additional figures mentioned in the main text.

References↩︎

---

1. An analogous condition can be obtained for \(H(\bar p_\tau|\tilde{p}_\tau)\), taking \(\tilde{\mu}_\tau := \bar{p}_\tau(x)|\nabla\log \tilde{p}_\tau(x)/\bar{p}_\tau(x)|\).↩︎

2. This estimator has limited precision, but it was found to be more reliable than the alternatives, such as using a KDE estimator for the densities or directly using a non-parametric estimator for the divergences.↩︎