January 15, 2026
Diffusion models have shown remarkable empirical success in sampling from rich multi-modal distributions. Their inference relies on numerically solving a certain differential equation. This differential equation cannot be solved in closed form, and its resolution via discretization typically requires many small iterations to produce high-quality samples. More precisely, prior works have shown that the iteration complexity of discretization methods for diffusion models scales polynomially in the ambient dimension and the inverse accuracy \(1/\varepsilon\). In this work, we propose a new solver for diffusion models relying on a subtle interplay between low-degree approximation and the collocation method [1], and we prove that its iteration complexity scales polylogarithmically in \(1/\varepsilon\), yielding the first “high-accuracy” guarantee for a diffusion-based sampler that only uses (approximate) access to the scores of the data distribution. In addition, our bound does not depend explicitly on the ambient dimension; more precisely, the dimension affects the complexity of our solver through the effective radius of the support of the target distribution only.
Diffusion models [2]–[8] are the dominant paradigm in image generation, among other modalities. They sample from high-dimensional distributions by numerically simulating a reverse process driven by a certain differential equation with drift learned from training data. The reverse process is meant to undo some noise process, and by simulating the reverse process sufficiently accurately, one can generate fresh samples from the distribution out of pure noise.
The empirical success of this method has spurred a flurry of theoretical work in recent years to understand the mechanisms by which diffusion models are able to easily sample from rich multimodal distributions. These works draw upon tools from the extensive literature on log-concave sampling and arrive at a remarkable conclusion: diffusion models can efficiently sample from any distribution in high dimensions, even highly non-log-concave ones, provided one has access to a sufficiently accurate estimate of the score of the distribution along a noise process (see Section 1.2 for an overview of this literature).
The earliest such results showed that for any smooth distribution with bounded second moment, one can sample to error \(\varepsilon\) in total variation distance in \(\mathrm{poly}(d,1/\varepsilon)\) iterations, given sufficiently accurate score estimation [9]–[13]. These results focus on samplers that simulate the reverse process in discrete time. The discretization introduces bias, which necessitates taking the discretization steps sufficiently small – i.e. of size \(\mathrm{poly}(\varepsilon,1/d)\) – to ensure the trajectory of the sampler remains sufficiently close to that of the continuous-time reverse process.
While these discretization bounds have been refined significantly by recent work (see Section 1.2), what remains especially poorly understood is the dependence on \(\epsilon\). In the log-concave sampling literature, there is a well-understood taxonomy along this axis: there are (1) “low-accuracy” methods like Langevin Monte Carlo that get iteration complexity scaling polynomially in \(1/\varepsilon\), and (2) “high-accuracy” methods which correct for discretization bias, e.g., via Metropolis adjustment, and get iteration complexity polylogarithmic in \(1/\varepsilon\).
For diffusion models, guarantees of the second kind remain largely unexplored. We therefore ask:
Are there samplers for diffusion modeling that achieve \(O(\mathrm{polylog} 1/\varepsilon)\) iteration complexity?
Here we answer this in the affirmative. We will focus on the following class of target distributions \(q\):
Assumption 1 (Bounded plus noise). Let \(R, \sigma > 0\). There exists a distribution \(q_{\sf pre}\) supported on the origin-centered ball of radius \(R\) in \(\mathbb{R}^d\) for which \(q = q_{\sf pre}\star \mathcal{N}(0,\sigma^2 I)\).
In other words, we consider sampling from a distribution \(q\) which is the convolution of a compactly supported distribution with Gaussian noise. Our bounds will scale polynomially in \(R / \sigma\). A natural example of a distribution satisfying Assumption 1 is a mixture of isotropic Gaussians; more generally, as observed by [10], distributions satisfying Assumption 1 naturally model what diffusion models in practice try to sample from via early stopping. Positivity of \(\sigma\) also makes it possible to prove convergence guarantees in strong divergence-based metrics like total variation distance.
In this paper, we consider sampling from \(q\) by approximately simulating the reverse process. We first show a key structural property: the high-order time derivatives of the score function along the reverse process are bounded, implying that the score function can be pointwise-approximated by a low-degree polynomial in time. The statement is technical and we defer it to the supplement.
Leveraging this result, we show how to adapt a certain specialized ODE solver of [1] to simulate steps along the reverse process, and argue that our simulation remains close to the trajectory of the reverse process at all times, see Algorithm [alg:main]. We thus prove that Algorithm [alg:main] can efficiently sample from \(q\), yielding the following main guarantee:
Theorem 1 (Informal, see Corollary 1). Let \(\varepsilon, R, \sigma > 0\). Suppose \(q\) is a distribution satisfying Assumption 1 for parameters \(R,\sigma\). Given access to Lipschitz score estimates for which the estimation error is \(L^2\)-bounded and has subexponential tails, there is a diffusion-based algorithm (see Algorithm [alg:main]) which outputs samples from a distribution \(\varepsilon\) close to \(q\) in total variation (TV) distance in \(\tilde{O}((R/\sigma)^2\cdot \operatorname{poly}\log 1/\varepsilon)\) iterations.
To the best of our knowledge, this is the first diffusion-based sampler to achieve polylogarithmic, rather than polynomial, dependence on \(1/\varepsilon\), while only assuming access to a sufficiently accurate score. While our result requires a stronger assumption on the distribution of score error than prior work, we note that the analysis in prior work incurs polynomial dependence on \(1/\varepsilon\) even assuming perfect score estimation.
Another appealing feature of our result is that the iteration complexity of our algorithm is independent of the ambient dimension \(d\), instead depending on the radius \(R\) in Assumption 1. While in general \(R\) can scale with \(d\), the Gaussian mixture setting offers a natural setting where \(R\) can be much smaller than \(d\). For example, in the theory of distribution learning, the most challenging regime is precisely when the component centers are at distance \(\Theta(\sqrt{\log k})\) from each other [14]–[17]. This corresponds to the case of \(R = \Theta(\sqrt{\log k})\) and \(\sigma = 1\), for which our sampler achieves iteration complexity scaling only polylogarithmically in \(k\) and \(1/\epsilon\), whereas existing diffusion-based methods would scale polynomially in \(d\) and \(1/\epsilon\) for this example (see Section 1.2 for discussion of a recent result [18] that obtains a dimension-free rate in the Gaussian mixture setting).
We review the extensive line of recent works on diffusion model convergence bounds in Appendix 5, focusing here on the threads most relevant to our result.
1.25ex -1em Dependence on \(\epsilon\). First, mirroring the popularization of higher-order solvers for diffusion sampling in practice [19]–[21], a number of works have sought to improve the \(\epsilon\) dependence for diffusion-based sampling by appealing to ideas from higher-order numerical solvers, e.g., [22], [22], [23] which achieved fixed polynomial acceleration in the \(\epsilon\) dependence under relatively minimal assumptions.
The recent works of Huang et al. [24], [25] and Li et al. [26] used higher-order solvers to achieve arbitrary polynomial acceleration, i.e. a rate scaling like \(O(d^{1+O(1/K)}/\epsilon^{1/K})\) for arbitrarily large constant \(K\). We note however while this might suggest that one could take \(K = \tilde{\Theta}(\log (1/\epsilon))\) and achieve similar polylogarithmic scaling as our work, the analysis in these works implicitly assumes that the number of iterations is at least exponential in \(K\) (see e.g., the Discussion in [26]), so the acceleration they obtain is at best polynomial rather than exponential like in our work. In addition, the assumptions in these works are incomparable to ours. For instance, [24], [25] assumes at least second-order smoothness of the score estimates; [26] does not require any smoothness of the score estimate, but it assumes its Jacobian of the score estimate is close to that of the true score. In contrast, we need to assume first-order smoothness of the score estimate (which is justified because the true score is Lipschitz) and sub-exponential score error. Finally, we note that these three works and ours all assume the data distribution has compact support.
To our knowledge, the only prior work which truly achieves exponentially improved \(\epsilon\) dependence is [27]. However, that work applies in a stronger access model than is commonly studied in this literature, as they assume approximate query access to log density ratios between arbitrary points, in addition to the scores, in order to implement a Metropolis-Hastings filter. We also note the independent and concurrent work of [28] which decomposes the reverse process into sampling from a sequences of posteriors which are strongly log-concave, and notes that by using off-the-shelf high-accuracy guarantees for strongly log-concave sampling, one can achieve a logarithmic dependence on \(\epsilon\) for diffusion-based sampling. However, these off-the-shelf routines require knowledge not just of the scores, but of the log density of the data distribution convolved with noise, similar to [27].
1.25ex -1em Dimension-free bounds. As for our dependence on \(R/\sigma\) instead of \(d\), there is a relevant line of work [29]–[32] on sampling from distributions with low intrinsic dimension. They show that for any distribution whose intrinsic dimension, which they quantify in terms of covering number of the support denoted \(k\), DDPM can sample with iteration complexity \(O(k^4/\varepsilon^2)\). While this provides another example of a diffusion-based sampler whose rate adapts to the underlying geometry of the distribution, it is incomparable to our result: there can be supports which are high-dimensional but bounded in radius, and vice versa. On the other hand, [18] considered the special case of mixtures of isotropic Gaussians and showed, using very different techniques from those of the present work, that diffusion models can sample from these distributions in \(T = \tilde{O}(\mathrm{polylog}(k)/\varepsilon)\) iterations, provided that the means of the Gaussians have norm at most \(T^C\) for arbitrarily large absolute constant \(C\). We remark that this result is incomparable to ours. We obtain polylogarithmic dependence on \(1/\varepsilon\) and our guarantee applies to a wide class of distributions beyond just the special case of mixtures of isotropic Gaussians. On the other hand, the result of [18] has superior scaling in other parameters for Gaussian mixtures: whereas we have a polynomial dependence on the maximum norm of any center in the mixture, they have an arbitrarily small dependence on the radius.
1.25ex -1em Collocation for diffusions. Finally, we note that the collocation method (see Section 2.3) has been studied in the context of diffusions, but primarily as a way to parallelize the steps of the sampler [33]–[35], but not using low-degree polynomial approximation. Those works obtain polylogarithmic round complexity, but the total work done is still polynomial in the target accuracy. In contrast, the key technical ingredient in our work is to establish low degree approximability of the score in time, which allows us to establish that the implementation of collocation by [1] via low-degree polynomials over a small time window of the probability flow ODE converges at an exponential rate, even in a sequential setting.
1.25ex -1em High-accuracy log-concave sampling. The literature on high-accuracy samplers, which is traditionally centered around log-concave distributions and more generally distributions which satisfy a functional inequality, is too extensive to do justice to here, and we refer to Chapter 7 of [36] for a detailed overview. We briefly overview this in Appendix 5. Our techniques do not draw upon this literature but are instead based on the work of [1], which devised the general framework of collocation via low-degree approximation that we employ. Motivated by sampling problems connected to logistic regression, they applied their method to obtain a high-accuracy sampler for densities of the form \(q\propto e^{-f}\), where \(f(x) = \sum_i \phi_i (\langle a_i, x\rangle) + \lambda\|x\|^2\).
In Section 2, we provide technical preliminaries and give a description of our sampler. In this section we also provide intuition for how this sampler can achieve the high-accuracy guarantee of Theorem 1. In Section 3, we sketch the proof of our main structural result showing that the drift of the reverse process is well-approximated by a low-degree polynomial in time, as well as the main steps for completing the proof of Theorem 1. We defer the full proof details to the Appendix.
1.25ex -1em Notation. Let \(\gamma^d\) denote the \(d\)-dimensional standard Gaussian distribution. Given interval \(I\subseteq\mathbb{R}^d_{\ge 0}\), let \(\mathcal{C}(I,\mathbb{R}^d)\) denote the space of continuous maps \(I\to \mathbb{R}^d\). For any distribution \(p\) on \(\mathbb{R}^d\), \(L^2(p)\) denotes the space of squared integrable random vectors in \(\mathbb{R}^d\). We recall the definitions of the total variation distance \(\mathrm{TV}(P,Q)=\tfrac12\int_{\mathbb{R}^d}|p-q|\,\mathrm{d}x\) for densities \(p,q\), and the Wasserstein-2 distance \(W_2^2(P,Q)=\inf_{\pi\in\Pi(P,Q)}\int\|x-y\|^2\,\mathrm{d}\pi(x,y)\) where \(\Pi(P,Q)\) is the set of couplings. Finally, for any vector \(x = (x_1, \ldots, x_d) \in \mathbb{R}^d\), \(\|x\|_{\infty} = \max_i |x_i|\) and if \(x\) is random, \(\|x\|_{p, \infty} = \max_i \left( \mathbb{E} |x_i|^p \right)^{1/p}\).
Here we review the basics of diffusion models; for a more detailed overview, we refer the reader to the surveys of [37] and [38].
Diffusion models are built upon two main components: the forward process and the reverse process. The forward process is a noise process driven by a stochastic differential equation of the form
\[\mathrm{d}x_t = -x_t\,\mathrm{d}t + \sqrt{2}\,\mathrm{d}B_t\,, \qquad x_0 \sim q_{\sf pre}\,,\] where \((B_t)_{t\ge 0}\) is a standard Brownian motion in \(\mathbb{R}^d\). Conditioned on \(x_0\), the process at time \(t\) is distributed as \(e^{-t} x_0 + \sqrt{1 - e^{-2t}} g\) for \(g\sim\gamma^d\); denote this marginal distribution by \(p^\rightarrow_t\), where \(p^\rightarrow_0 = q_{\sf pre}\).
We run the forward process up to time \(T\). Define \[\sigma_{t} = \sqrt{1 - e^{-2(T-t)}} \,. \label{eq:sigmadef}\tag{1}\]
The reverse process is designed to undo this noise process, i.e. transform \(p^\rightarrow_T\) to \(p^\rightarrow_0\). For convenience of the notation, we denote \(p^\rightarrow_{T-t}\) by \(q_t\), so that \(q_T = q_{\sf pre}\). One version of this reverse process is given by the probability flow ODE, which is specified by
\[\mathrm{d}y_t = (y_t \, + \nabla \ln q_{t}(y_t))\,\mathrm{d}t\,,\label{eq:backwardode}\tag{2}\] where \(\nabla \ln q_{t}\) is called the score function. The key property is that if \(y_0 \sim q_0\), then \(y_T \sim q_T = q_{\sf pre}\). In practice, this is run using score estimates \(s_t \approx \nabla \ln q_t\) instead of the actual score functions, and in the theoretical literature it is standard to assume that these are close in \(L_2(q_t)\). In this work, we will make a somewhat stronger assumption:
Assumption 2 (Sub-exponential score error). There is some parameter \(\epsilon_{\sf err} > 0\) such that the error incurred by the score estimate* \(s_t: \mathbb{R}^d\to\mathbb{R}^d\) has sub-exponential norm at most \(\epsilon_{\sf err}\) for all \(0 \le t \le T\). That is, \(\Pr[t]{\|s_t - \nabla \ln q_t\| \geq z} \le 2\exp(-z/\epsilon_{\sf err})\) for all \(y \geq 0\).*
We leave as open whether this can be relaxed to the more standard assumption of \(L_2\)-accurate score estimation but note that even under the assumption of perfect score estimation, it was previously not known how to achieve high-accuracy guarantees.
Additionally, instead of initializing the ODE at \(y_0 \sim q_0\), one initializes at \(y_0 \sim \gamma^d\) for some noise distribution for which \(\pi \approx q_0\) and from which it is easy to sample. For instance, in the example of the standard OU process, we can take \(\pi = \gamma^d\) because the forward process converges exponentially quickly to \(\gamma^d\). As \(\pi\) is easy to sample from, the probability flow ODE can be used to (approximately) generate fresh samples from \(q\).
Finally, we further assume the score estimate is Lipschitz.
Assumption 3 (Lipschitz score estimate). \(s_t\) is \(\tilde{L}\)-Lipschitz: \(\forall x,y\in \mathbb{R}^d, \;\|s_t(x) - s_t(y)\| \leq \tilde{L}\|x - y\|\).
Traditionally, to simulate the continuous-time ODE in discrete time, some numerical method like Euler-Maruyama discretization is used. Such discretization schemes can only approximately simulate the continuous time ODE by using a fixed gradient in a window of time. For example, to approximate the ODE \(x_t = f_t(x_t)dt\) with one gradient step in time window \([0,h]\), the Euler-Maruyama method uses the gradient at the beginning time \(0\), resulting in the update \(\hat{x}_{h} = h f_0(\hat{x}_0) + x_0\), where \(\hat{x}_h\) denotes the discretized process. Here, \(h\) is the step size of the algorithm. Unfortunately, this type of discretization puts a fundamental dimension-dependent limitation on how large the step size can be.
To illustrate the point, consider the vector field \(f_t = x\) for all times \(t\geq 0\). In this case, the discretized step is given by \(\hat{x}_h = h (1)_{i=1}^d + x_0 = ((1+h))_{i=1}^d\) while the solution of the ODE \(x_t = f_t(x_t) dt\) with initial condition \(x_0 =(1)_{i=1}^d\) is given by \(x_t = (e^t)_{i=1}^d.\) This implies \(\|x_h - \hat{x}_h\| = \|(e^h - (1+h))_{i=1}^d\| = \Theta(h^2 \sqrt d)\). Therefore, in order to keep the deviation \(\|\hat{x}_h - x_h\|\) bounded by \(O(1)\), the step size \(h\) has to be \(O(1/d^{1/4})\). This is because as we move along the ODE solution \(x(t)\), the vector field \(f_t(x_t)\) can change drastically from its initial value, so using a fixed vector \(f_0(x_0)\) as the velocity does not allow us to go that far.
However, if we had a prior knowledge about the functional form of \(f_t(x_t)\) along the path \((x_t)_{[0,h]}\), perhaps we could use that to take larger steps to simulate the ODE. One approach would be to leverage higher-order smoothness of \(f_t\) and run a higher-order numerical solver; as discussed in Section 1.2, many prior works pursued this approach [22]–[26], [39], but none of them are known to achieve exponential acceleration in \(\epsilon\).
Instead, we opt for a different strategy. Suppose that we knew \(f_t(x_t)\) lives in a low-dimensional subspace with basis \(\{\phi_j(t)\}_{j=1}^D\), namely \(f_t(x_t) = \sum_{i=1}^D c_i \phi_i\) for some coefficients \((c_i)_{i=1}^D \in \mathbb{R}^D\), then if we could estimate these coefficients, that allows us to predict the path \(x_t\) for large times \(t\). In this work, we show that perhaps surprisingly, under minimal conditions on the target distribution, each coordinate of the score function on the probability flow ODE path, \(f_t(x_t)\), is well-approximated by a low-degree polynomial in \(t\). We then prove that one can estimate these coefficients for each coordinate using the collocation method together with our score estimates. Notably, this enables us to take longer discretized steps, without any explicit dimension dependency, to follow the ODE path. We emphasize that while the collocation method, i.e. Picard iteration, has been used in prior diffusion model theory work [34], [35], [40], the key novelty in our approach is to leverage low-degree polynomial approximation.
The collocation method is a numerical scheme for approximating the solution to an ordinary differential equation through fixed-point iteration. Here, consider a generic initial value problem: \(\mathrm{d}x_t = f_t(x_t)\,\mathrm{d}t\) and \(x_0 = v\) for all \(t\in[0,H]\). This admits the integral representation \(x_t = v + \int^t_0 f_s(x_s)\, \mathrm{d} s\,,\) which can be thought of as the fixed-point solution to the equation \(x = \mathcal{T}(x)\) for the operator \(\mathcal{T}: \mathcal{C}([0,H], \mathbb{R}^d) \to \mathcal{C}([0,H],\mathbb{R}^d)\) given by \(\mathcal{T}(x)_t \triangleq v + \int^t_0 f_s(x_s)\,\mathrm{d}s\) which has a unique fixed point provided there exists \(k\in \mathbb{N}\) such that \(\mathcal{T}^k\) is \(L\)-Lipschitz with \(L < 1\). Authors in [1] proposed to solve this fixed-point equation as follows. Suppose that each coordinate of the time derivative of the solution, i.e. \(t\mapsto f_t(x_t)\), is well-approximated by a low-degree polynomial. Letting \(\phi_1,\ldots,\phi_D: [0,H]\to\mathbb{R}\) be an appropriately chosen basis of polynomials, by polynomial interpolation we can find nodes \(c_1,\ldots,c_D\) such that \(\phi_j(c_i) = \mathbb{1}[i = j]\) for all \(i,j\in[D]\) so that in particular, \[\frac{\mathrm{d}x_t}{\mathrm{d}t} \approx \sum^D_{j=1} f_{c_j}(x_{c_j})\phi_j(t)\,.\]
Writing this in integral form as before, we arrive at the approximate fixed-point equation: \[x \approx \mathcal{T}_\phi(x)\,, \qquad \mathcal{T}_\phi(x)_t \triangleq v + \int^t_0 \sum^D_{j=1} f_{c_j}(x_{c_j})\phi_j(s)\,\mathrm{d}s\,. \label{eq:Picard}\tag{3}\] This suggests a natural algorithm: instead of maintaining the entire continuous-time solution \(x:[0,H]\to\mathbb{R}^d\) over the course of fixed-point iteration, simply maintain the values \((x_{c_j})_{j\in[D]}\) and update these according to Eq. 3 , which amounts to a matrix-vector multiplication. This algorithm is summarized in Algorithm [alg:picard] below.
Define \(A_\phi\in\mathbb{R}^{D\times D}\) by \((A_\phi)_{i,j} = \int^{c_j}_0 \phi_i(s)\,\mathrm{d}s\) Define \(X^{(0)} \in\mathbb{R}^{d\times D}\) to be \(v 1^\top_D\), where \(1_D\) is the all-ones vector.
[1] give conditions under which collocation with appropriately chosen basis polynomials and nodes converges to a sufficiently accurate solution to the ODE. We will need to adapt their guarantees to our setting to account for score estimation error. The full guarantee will be given in Section 3. We require the following condition for the polynomial basis \(\phi\):
Definition 1. We say that \(\phi\) is \(\gamma_\phi\)-bounded* if, for all basis elements \(\phi_j: [0,H]\to\mathbb{R}\) and \(0 \leq x \leq H\), \(\sum^D_{j=1} \Bigl|\int^H_0 \phi_j(s)\,\mathrm{d}s\Bigr| \le \gamma_\phi t\).*
In this work, concretely, we will take \(\phi_j(x) = \widetilde{\phi}_j(2x/H - 1)\), where \[\widetilde{\phi}_j(x) = \frac{\sqrt{1 - c^2_j}\cos(D\cos^{-1} x)}{D(x-c_j)}\,,\] where \(c_j = \cos(\frac{2j-1}{2D}\pi)\) is the \(j\)-th root of the Chebyshev polynomial of degree \(D\). As shown in the proof of [1], \(\widetilde{\phi}_j\) is a degree-\(j\) polynomial satisfying \(\widetilde{\phi}_j(c_i) = \mathbb{1}[i=j]\); moreover, [1] notes that by [41], \[\sum_{j=1}^D\Big|\int_{-1}^1 \widetilde{\phi}_j(x) \mathrm{d}x\Big| \,\leq 2000\,.\] After a change of variable to rescale from \(\widetilde{\phi}_j\) to \(\phi_j\), we conclude that this basis \((\phi_j)\) is \(\gamma_\phi = O(H)\)-bounded.
We are now ready to combine the ingredients in the previous subsections to give a description of our sampler.
\(\hat{x}_{0} \sim \gamma^d\)
The algorithm simply splits up the reverse process into a series of small time windows, each of which is of length \(\tilde{O}(\frac{\sigma^2}{R^2})\), and successfully solves the corresponding initial value problems in each of these windows using collocation via Picard. There will be some accumulation of errors, as the Picard calls do not result in exact solutions, but these errors can be made extremely small because Picard solves the relevant ODEs to extremely high precision.
In this section we describe the main ingredients of our proof, deferring proofs to the Appendices. Recall that key to our approach is to prove a structural result showing that the time derivative of the score function along the reverse process can be controlled.
Our starting point is to relate time derivatives of the score at a point \(y_t\) along the reverse process to higher-order moments of the posterior distribution on the clean sample that would have generated \(y_t\) along the forward process.
We begin with the following calculation.
lemmavecfieldderiv Suppose \(y_t = X_t + \sigma_t \xi\), where \(X_t = e^{-(T-t)} \bar X\) for \(\bar X\sim q_{\sf pre}\) and \(\xi\sim \gamma^d\). We use the notation \(\mathbb{E}^{t,y_t}\) to denote the conditional expectation of \(X_t\) given \(y_t = X_t + \sigma_t \xi\), and \(q^{t,y_t}\) to denote the density for this posterior. Define the posterior mean \(\mu_t(y) \triangleq \mathbb{E}^{t,y}X_t\).
The time derivative of the vector field of probability flow ODE can be calculated as \[\begin{align} &\partial_{t} (y_{t} + \nabla \log q_t(y_{t})) = \left( {y_t + \frac{\mathbb{E}^{t,y_t}X_t - y_t}{\sigma_t^2}} \right) \\ &\qquad+ \frac{1}{\sigma_t^2}\mathbb{E}^{t,y_t} X_t + \frac{1-\sigma_t^2}{2\sigma_t^6}\mathbb{E}^{t,y_t} \left( {\left\langle {y_t,X_t}\right\rangle + \left\Vert {X_t}\right\Vert^2} \right) \left( {X_t - X'_t} \right) + \frac{4}{\sigma_t^4}\mathbb{E}^{t,y_t} \left( {\left\langle {y_t,X_t}\right\rangle - \left\Vert {X}\right\Vert^2} \right)\left( {X_t - X'_t} \right)\\ &\qquad+\frac{1}{\sigma_t^2}(\frac{1}{\sigma_t^2} - 1)y_t - \frac{1}{\sigma_t^4}\mathbb{E}^{t,y_t}X_t + \frac{1}{\sigma_t^4}\mathbb{E}^{t,y_t}\left( {2\left\langle {y_{t}, \frac{ X''_t}{\sigma_t^2}}\right\rangle - \left\langle {X_t, y_t}\right\rangle - \frac{1}{\sigma_t^2}\left\langle {X_t,X''_t}\right\rangle} \right) (X_t - X'_t). \end{align}\] Here, \(X_t, X'_t, X''_t\) denote independent, identically distributed draws from the posterior distribution with expectation operator \(\mathbb{E}^{t,y_t}\).
Iterating this calculation multiple times, we arrive upon the following key calculation expressing the higher-order derivatives of the vector field of the probability flow ODE:
lemmaderform The \(k\)th derivative of the backward ODE vector field can be expanded as \[\begin{align} &\partial_{t}^k \left( {y_{t} + \nabla \ln q_t(y_{t})} \right) \\ &= \sum_{r_1, r_2, r_3, r_4}\sum_{\mathbf{i}\in [k]^{r_3}, \mathbf{k}\in [k]^{r_4}, \mathbf{l}\in [k]^{r_4}} \frac{\left( {1-\sigma_t^2} \right)^{r_1}}{\sigma_t^{r_2}}\mathbb{E}^{t,y_t}\prod_{j=1}^{r_3} \left\langle {y_t, X^{(\mathbf{i}_j)}}\right\rangle \prod_{j=1}^{r_4}\left\langle { X^{(\mathbf{k}_j)}, X^{(\mathbf{l}_j)}}\right\rangle (a_{\mathbf{i}, \mathbf{k}, \mathbf{l}} X + b_{\mathbf{i}, \mathbf{k}, \mathbf{l}} y_t) \end{align}\] where \(r_1,r_2,r_3 \leq k\) and \(r_2 \leq 4k\) and the coefficients \((a_{\mathbf{i},\mathbf{k},\mathbf{l}}, b_{\mathbf{i},\mathbf{k},\mathbf{l}})\) satisfy \(\sum_{\mathbf{i}\in [k]^{r_3}, \mathbf{k}\in [k]^{r_4}, \mathbf{l}\in [k]^{r_4}} \big|a_{\mathbf{i}, \mathbf{k}, \mathbf{l}}\big| + \big|b_{\mathbf{i}, \mathbf{k}, \mathbf{l}}\big| \leq (74k)^k\,.\) Here the random variables \(X^{(i)}\) are independently distributed according to the posterior distribution \(q^{t,y}\), and \(\mathbf{i}, \mathbf{k}, \mathbf{l}\) are arbitrary tuples of indices.
This result can be understood in the same spirit as Tweedie’s formula, which relates the posterior mean of \(X\) to the score function. Higher-order generalizations of Tweedie’s formula in the literature [42] typically relate moments of the posterior to derivatives of the score with respect to the space variable, whereas in Lemma [lem:derformula] we consider derivatives with respect to the time variable.
An important feature of our technique is to bound this multiple integration formula for higher derivatives of the score independent of the ambient dimension \(d\), which will be crucial for obtaining iteration complexity bounds that only depend on the effective radius \(R\) of the distribution. More precisely, the bound on the effective radius allows us to bound expectations of the form \(\mathbb{E}[\prod^{r_4}_{j=1} \langle X^{(\mathbf{k}_j)}, X^{(\mathbf{l}_j)}\rangle]\) pointwise and expectations of the form \(\mathbb{E}[\prod^{r_3}_{j=1} \langle y, X^{(\mathbf{i}_i)}\rangle]\) with high probability by dimension-independent quantities, from which we can deduce bounds on the left-hand side of Lemma [lem:derformula] via applications of Cauchy-Schwarz (see Lemma 4). Leveraging these formulas, we arrive at our first key technical ingredient, a dimension-free bound on the higher-order time derivatives of the score:
lemmabndlem If \(q\) satisfies Assumption 1, \((y_t)_t\) is a solution to 2 , and \(y_0 \sim q_T\), then the \(k^{\text{th}}\) derivative of the vector field along the probability flow ODE is bounded as \[\begin{align} \left\|\partial_{t}^k \left( {y_t + \nabla \log q_{t}(y_t)} \right)\right\|_{p, \infty} \lesssim R \left( {\frac{74k}{\sigma_t^2}} \right)^k\left( {\frac{R}{\sigma_t} + (kp)^{1/2}} \right)^{2k} \end{align}\]
Note that these bounds are singly exponential in the order of the derivative and scale with \((R/\sigma_t)^k\). If these derivative bounds held pointwise, then we could simply appeal to the Taylor remainder theorem to argue that the true score is well-approximated by a low-degree polynomial. But ultimately Lemma [lem:bnd1] only offers a high-probability statement over the distribution induced by the true reverse process. We will need to handle deviations from the true reverse process, which arise because the sampler is initialized at Gaussian instead of the correct marginal, and because in each subsequent window in which we apply collocation, the initialization is also slightly off from the correct marginal due to discretization errors accumulated from previous windows. Next, we describe how to deal with these deviations.
To argue about low-degree approximation along the trajectory of the actual sampler, we will need to argue that the posterior distribution from Lemma [lem:derformula] is robust to perturbations to the conditioning.
Our key tool for doing so is the following coupling lemma:
lemmacouplinglemma Let \(\tilde{y}\) a random variable such that \(\left\Vert {y - \tilde{y}}\right\Vert \leq \delta\) for \(\delta \leq \frac{1}{6(\sqrt d + \sqrt{\ln(1/\epsilon_1)})}\). Then, given \(T - t \geq 1\), under the event \(\mathcal{A}_1\) defined in Equation 12 , we have \[\label{eq:wpeps} \mathrm{TV}(q^{t,y}, q^{t,\tilde{y}}) \leq 8\delta (\sqrt d + \sqrt{\ln(1/\epsilon_1)})\,.\tag{4}\] In particular, Eq. 4 holds with probability at least \(1-\varepsilon_1\).
Combining Lemmas [lem:bnd1] and [lem:klboundlemma-main], we obtain a bound on higher time derivatives of the score along the reverse process when initialized at a distribution that is slightly off from the true marginal:
theoremderivativethm Let \(\delta \in (0,1)\). If \(q\) satisfies Assumption 1, and \((\tilde{y}_t)_t\) is given by initializing the reverse process at a distribution which is \(\delta\)-close in \(W_2\) to \(q_0\) and running the probability flow ODE in continuous time. Then, for \(\delta\) small enough (as rigorously characterized in Lemma 10), with probability at least \(1 - \delta_2\) over the randomness of the initialization, the \(k^{\text{th}}\) derivative of the vector field along the probability flow ODE is bounded at \(\tilde{y}_t\) by \[\begin{align} \left\| \partial_{t}^k \left( {\tilde{y}_t + \nabla \log q_{t}(\tilde{y}_t)} \right)\right\|_{\infty} \lesssim R \left( {\frac{k}{\sigma_{t}^2}} \right)^k\left( {\frac{R}{\sigma_{t}}} \right)^{2k}+ \left( {k\log(74kd/\delta_2)} \right)^{2k}\,. \end{align}\]
This theorem provides the theoretical underpinning for our approach. Because it is a high probability bound over iterates of a process initialized at the correct intermediate points of the true sampler, we can now instantiate the aforementioned Taylor truncation argument to get the desired low-degree approximation along the trajectory of the sampler itself, rather than of the idealized reverse process. We will use this ingredient in conjunction with the technology in [1] to prove the following convergence bound for Picard (Algorithm [alg:picard]), recalling the notations from Section 2.3:
Suppose Picard is run using a \(\gamma_\phi\)-bounded polynomial basis \(\phi\) with nodes \((c_j)\). For any \(h \le \frac{1}{2\tilde{L}\gamma_\phi}\), for the true probability flow ODE \((y_t)\) and arbitrary \(x \in \mathcal{C}([t_0,t_0 + h], \mathbb{R}^d)\), \[\begin{align} &\left\Vert {T_\phi^{\circ m} (y) - y}\right\Vert_{[t_0, t_0+h]} \leq 2\Bigl(\epsilon_{\mathop{\sf ld}}+ \max_{1\le j \le D}\left\Vert {s_{c_j}(y(c_j)) - y(c_j) - \nabla \log q_{c_j}(y(c_j))}\right\Vert\Bigr)\left( {1 + \gamma_\phi} \right) h \tag{5}\\ &\|T_\phi^{\circ m}(x) - T_\phi^{\circ m}(y)\|_{[t_0,t_0 + h]} \leq \frac{1}{2^m} \left\Vert {x-y}\right\Vert_{[t_0,t_0 + h]}\,, \tag{6} \end{align}\] where here \(\|\cdot\|_{[a,b]}\) denotes the sup norm over the interval \([a,b]\), and \(\epsilon_{\mathop{\sf ld}}\) denotes the approximation error (in sup norm) incurred by approximating the score function along the probability flow ODE with a polynomial in the span of \(\phi\).
Eq. 5 above shows that the trajectory of the true probability flow ODE, which is an exact fixed point for exact Picard iteration, is locally an approximate fixed point for the polynomial interpolation-based approximate Picard iteration given by Picard (Algorithm [alg:picard]). Eq. 6 shows that any trajectory which is locally close to this true trajectory contracts towards the true trajectory at an exponential rate, up until an error given by the bound in Eq. 5 .
An important caveat is that the solver can only be run for time which scales inversely in the boundedness \(\gamma_\phi\) of the basis (recall Definition 1), which can be of constant order, and inversely in the Lipschitzness of the vector field of the probability flow ODE. This is why we cannot directly use collocation to solve the probability flow ODE in one shot. Nevertheless, the crucial point is that, as the second inequality above suggests, the error incurred by Picard is decreasing exponentially in the number of iterations, modulo some additional error that needs to be carried around coming from low-degree approximation and score estimation error. We will then chain together multiple applications of the above bound to get our final guarantee.
Finally, we note that while the core idea for the Proposition is from [1], one novelty of our bound is that we account for score estimation error and show that the errors coming from that do not compound exponentially over the Picard iterations. This will be essential for keeping our sampler sufficiently close to the true trajectory of the probability flow ODE.
We now have all the necessary ingredients to prove our main result.
We first show how to get a sampler which is \(W_2\)-close to \(q\). An appealing feature of this guarantee relative to our TV convergence guarantee is that 1) the requirement on the score estimation error is far less stringent, and 2) it only involves running the algorithm CollocationDiffusion, i.e. no postprocessing is needed on top of simply simulating the true reverse process using the ODE solver of [1] to implement discrete steps.
Theorem 2. Let \(\epsilon_{\sf err} \leq O(\epsilon / \ln(d/R))\), \(\sigma > 0\) and \(R \geq 1\). Suppose \(q\) is a distribution satisfying Assumption 1 for parameters \(R,\sigma\). Given access to score estimates \(s_t\) satisfying Assumptions 2 and 3, with probability at least \(1 - O(\epsilon (R/\sigma)^2)\), Algorithm [alg:main] outputs samples from a distribution \(\hat{q}\) for which \(W_2(q,\hat{q}) \le \tilde{O}(\varepsilon\cdot \log^2(1/\epsilon))\) in \(\tilde{O}((R/\sigma)^2\cdot \log(1/\epsilon))\) number of rounds and \(\tilde{O}(\log(1/\epsilon))\) number of Picard iterations per round, for a total of \(\tilde{O}((R/\sigma)^2\cdot\log^2(1/\epsilon))\) calls to the score estimate oracle.
Above, \(\tilde{O}(\cdot)\) hides factors of \(\log d, \log R, \log \tilde{L}, \log\log(1/\epsilon)\). Furthermore, the assumption \(R \geq 1\) is only for simplifying the bounds. For the full theorem please see Appendix 11.
Here we informally sketch the proof of this result, deferring the proof details to the Appendix. The idea will be to track the trajectory of the true probability flow ODE closely at all times and successively run Picard over small time windows.
Suppose that up to some time \(t\) in the simulation of the reverse process, the iterate of the sampler is distributed according to a distribution whose \(W_2\) distance from the true marginal is small. Then consider trying to simulate the reverse process for another \(h\) time steps to approximate \(q_{T-t-h}\). One can do this by simply running CollocationDiffusion, provided \(h\) is small enough, and because the starting distribution is \(W_2\)-close to the true marginal, the Taylor truncation argument in conjunction with the bound on the higher-order derivatives of the vector field ensure that we can safely invoke Proposition [cor:contractionprops] and drive any error coming from the \(W_2\) discrepancy exponentially quickly to zero, leaving behind a fixed amount of score estimation and polynomial approximation error.
Note that the exponential contraction of the error from the \(W_2\) discrepancy is absolutely essential to our result. It allows us to “reset” the amount by which we stray from the curve of the true reverse process every time we run Picard. This is reminiscent of a different recent approach to analyzing the probability flow ODE by [43]. In that work, the authors tried simulating the ODE for short windows of time, appealing to a naive Wasserstein coupling argument in each of those windows. In their setting, the Wasserstein couplings incur some error that they would like to linearly, not exponentially, accumulate over the course of the sampler. The way to do this was to effectively “restart” the Wasserstein coupling at the start of each new time window by injecting a small amount of noise into the sampler. In contrast, rather than inject noise to restart the coupling, we use Picard iteration to suppress the accumulating Wasserstein error.
1.25ex -1em Upgrading to total variation closeness. We can use the above idea from [43] to derive an additional consequence of our main theorem. Specifically, as a consequence of regularizing properties of underdamped Langevin dynamics, our sampling guarantee in Wasserstein distance (Theorem 2) can be converted into a guarantee in total variation distance using standard arguments [34], [43], under the additional assumption that the true vanilla score \(\nabla \ln q\) is Lipschitz.
Corollary 1. Let \(\varepsilon, R, \sigma > 0\). Suppose \(q\) is a distribution satisfying Assumption 1 for parameters \(R,\sigma\) and that \(\nabla \ln q\) is \(L\)-Lipschitz. Given access to score estimates \(s_t\) satisfying Assumption 2, with \(\varepsilon_{\rm sc}\le \tilde{O}(\frac{\epsilon^{2/3} L^{1/2}d^{1/6}}{(R/\sigma)^{2/3}})\), Algorithm [alg:main2] outputs samples from a distribution \(\hat{q}\) for which \(\mathrm{TV}(\hat{q},q) \le \varepsilon\) in \(\tilde{O} \left((R/\sigma)^2\right)\) iterations.
In the high-accuracy regime, it might seem counterintuitive how one can so easily convert to closeness in TV starting from closeness in Wasserstein, given that one has to run underdamped Langevin Monte Carlo without any kind of Metropolis adjustment. The key idea is to simply run underdamped Langevin for a shorter amount of time so that the bias coming from discretization is not too large, an idea from [34]. We make this argument formal in Appendix 12.
In this work, we gave a diffusion-based sampler for sampling from a wide class of distributions that includes Gaussian mixtures as a special case, which both has iteration complexity which scales polylogarithmically in \(1/\varepsilon\) (“high-accuracy”) and does not depend explicitly on the dimension but rather on the effective radius of the distribution. The key technical step was to establish a bound on the higher-order derivatives of the vector field of the probability flow ODE, which in turn allowed us to approximate it by a low-degree polynomial. This approximation result made it possible to exploit a variant of the collocation method, a fixed-point iteration for numerically solving ODEs, proposed by [1] which converges exponentially quickly to the true ODE solution if run over sufficiently small time windows.
Our work has a few limitations that raise interesting questions for future study. Firstly, can we somehow relax our assumption on the underlying distribution \(q\)? Instead of insisting that it is a noised version of a distribution \(q_{\sf pre}\) whose support is compact, we could merely ask that \(q_{\sf pre}\) have bounded moments. We could also ask for TV or KL convergence to \(q\) in a number of steps which scales only logarithmically in \(1/\sigma\), matching what is known in the low-accuracy regime [11], [44]. Lastly, is the assumption that the score errors have sub-exponential tails truly necessary for achieving high-accuracy guarantees for sampling given score access?
We thank Sinho Chewi, Holden Lee, Yin Tat Lee, and Jianfeng Lu for helpful suggestions about discretization bounds. SC is supported by NSF CAREER award CCF-2441635.
1.25ex -1em Additional works on discretization bounds for diffusions.
There have been a large number works in recent years giving general convergence guarantees for diffusion models [9]–[13], [34], [39], [43]–[55]. Of these, one line of work [9]–[13], [55] analyzed DDPM, the stochastic analogue of the probability flow ODE, and showed \(\tilde{O}(d)\) iteration complexity bounds, which were recently established by [12], [13] under only the assumption that the data distribution has bounded second moment and very recently refined by [55] to apply to distributions with bounded first moment and to depend only linearly on \(1/\varepsilon\). Another set of works [22], [34], [39], [43], [52], [54], [56], [57] studied the probability flow ODE and proved similar rates, one notable difference being that under smoothness assumptions it is known how to get sublinear in \(d\) dependence [34], [43], [56], [57]. Recently, [58], [59] showed that sublinear scaling in \(d\) is also possible purely using DDPMs. As it stands however, all of these works work only in the “low-accuracy” regime, that is, the best-known \(\varepsilon\) achieved by these works is still \(\mathrm{poly}(1/\varepsilon)\), even under smoothness assumptions.
1.25ex -1em High-accuracy sampling for log-concave distributions.
One of the central ideas in the line of work on high-accuracy log-concave sampling is to append a Metropolis-Hastings filter step after every step of an otherwise low-accuracy sampler. Combining this with Langevin Monte Carlo or Hamiltonian Monte Carlo gives rise to the popular methods of MALA or MHMC respectively. A line of works [60]–[64], culminating in the recent breakthrough of [65], has established an iteration complexity of \(\tilde{O}(\sqrt{d}\mathrm{poly}(1/\varepsilon))\) for MALA for sampling log-concave distributions.
We will work with a data distribution \(q\) over \(\mathbb{R}^d\). Throughout, \(\bar{X}\) will denote a random vector sampled from \(q\), and \(X_t\) will denote the scaling \(X_t = \sqrt{1 - \sigma^2_t}\bar{X}\). Likewise, we use \(\bar{x}\) and \(x_t\) to denote realizations of these random vectors. We use \(\xi\) to denote a random vector sampled from \(\gamma^d\). Given \(t, y\), we will use \(q^{t,y}\) to denote the posterior distribution on \(X_t\) conditioned on \(\sqrt{1 - \sigma^2_t} \bar{X} + \sigma_t \xi = y\). We will use \(\mathbb{E}^{t,y}\) to denote expectation with respect to \(q^{t,y}\), and write \(\mu_t(y) \triangleq \mathbb{E}^{t,y}X_t\).
We will use the following noise schedule \[\sigma_t \triangleq \sqrt{1-e^{-2(T-t)}},.\] Throughout, we let \(q_t \triangleq \mathrm{law}(\sqrt{1 - \sigma^2_t}\bar{X} + \sigma_t \xi)\), so that \(q_{\sf pre}= q_T\).1
Given \(t\), let \(F^*_t: \mathbb{R}^d\to\mathbb{R}^d\) denote the map \[F^*_t(y) = y + \nabla \ln q_t(y)\,. \label{eq:Fstar}\tag{7}\] When \(t\) is clear from context, we will omit the subscript. Note that \[\begin{align} \nabla \ln q_{t}(y) = \frac{\mathbb{E}^{t,y} X_t - y}{\sigma_t^2} = \frac{\mu_t(y) - y}{\sigma_t^2}\,.\label{eq:scoreformula} \end{align}\tag{8}\]
We denote our estimate for the score function by \(F_t(y)\). When it is clear from the context, we drop the time index \(t\) in \(F_t\) and \(F^*_t\). Moreover, we consider the following processes:
\((y_t)\): the process given by the true probability flow ODE \[\mathrm{d}y_t = F^*_t(y_t)\,\mathrm{d}t\,,\qquad y_0\sim q_0\]
\((\hat{y}_t)\): the process given by our Picard Algorithm [alg:picard] for \(t\in [t_i, t_{i+1}]\), for all \(1\leq i \leq n\), as \[\begin{align} \hat{y}_t \triangleq x_{t_i} + \int^{t - t_i}_0 \sum^D_{i=1} F_{c_i}(X_{:,i})\phi_i(s)\,\mathrm{d}s, \end{align}\] where \(X_{:,i}\) are defined in Algorithm [alg:picard].
\((\tilde{y}_t)\): the process that, in the interval \([t_i, t_{i+1}]\), is given by the true probability flow ODE starting from the algorithm iterate \(\tilde{y}(t_i)\), for all \(1\leq i\leq n-1\), as defined in [alg:main]. Namely, \(\tilde{y}(0) = y_{t_i}\), and for all \(t\in [t_i, t_{i+1}]\): \[\frac{d}{dt} {\tilde{y}}(t) = F^*_t(\tilde{y}(t)).\]
In the course of the proofs, we denote the interval \([t_i, t_{i+1}]\) for a fixed \(1\leq i\leq n-1\) by \([t_0, t_0 + h]\). For an arbitrary curve \(z(t): t \in [t_0, t_0 + h] \rightarrow \mathbb{R}^d\), we further use the following sup norm in our calculations: \[\begin{align} \left\Vert {z}\right\Vert_{[t_0, t_0 + h]} = \sup_{t_0 \leq t\leq t_0 + h} \left\Vert {z(t)}\right\Vert. \end{align}\]
In order to show that the score function \(F^*_t(y_t)\) is a low-degree function of time \(t\), our approach is to bound higher derivatives of \(F^*_t(y_t)\) with respect to time. Having such bound then enables us to prove that \(F^*_t(y_t)\) is close to its \(k\)th order tailor approximation using the Taylor remainder theorem, which we recall below.
For function \(g: [t_0,t_0 + h] \rightarrow \mathbb{R}\), define the low degree approximation \(P_{\leq k} g(s)\) with degree \(k\) as: \[P_{\leq k} g(s) = \sum_{r=0}^{k-1} \frac{d^r g(s)}{ds^r}\Big|_{s=t_0} \frac{\left( {s-t_0} \right)^r}{r!}.\] Then by the Taylor remainder theorem, we have \[\left\Vert {g - P_{\leq k} g}\right\Vert_{[t_0, t_0+h]} \leq \frac{\max_{s\in [t_0, t_0+h]} \Big|\frac{d^{k} g(s)}{ds^k}\Big|}{k!} h^k.\]
With this in mind, we now proceed to prove dimension-independent bounds on the higher derivatives of the score function as stated in the main body (Lemmas [lem:vecfieldderivative]-[lem:bnd1]).
The main building blocks of the bound in Lemma [lem:derformula] are Lemmas 1 and [lem:vecfieldderivative]. Lemma 1 has two parts, stated in Lemmas 2, 3 below which compute the derivative of the posterior mean \(\mu_t(y)\) with respect to time (\(t\)) and space (\(y\)) variables. In Lemma [lem:bnd1] we combine all of these ingredients to obtain a bound on the higher derivatives of the score as a function of time.
We restate Lemma [lem:vecfieldderivative] here for convenience:
For this derivation, we will require the following calculation whose proof is given in Section 7.2 below.
Lemma 1. In the notation of Lemma [lem:vecfieldderivative], for any \(v\in\mathbb{R}^d\), \[\begin{align} & \partial_t \mu_t(y) = \mu_t(y) + \frac{1-\sigma_t^2}{2\sigma_t^4}\mathbb{E}^{t,y}\left( {\left\langle {y,X_t}\right\rangle + \left\Vert {X_t}\right\Vert^2} \right) \left( {X_t - X'_t} \right) + \frac{4}{\sigma_t^2}\mathbb{E}^{t,y}\left( {\left\langle {y,X_t}\right\rangle - \left\Vert {X_t}\right\Vert^2} \right)\left( {X_t - X'_t} \right), \\ & D(\mu_t(y))[v] = \frac{1}{\sigma_t^2}\mathbb{E}^{t,y}\left( {\left\langle {y-X_t, v}\right\rangle} \right)(X_t - X_t')\,, \end{align}\] where the expectations on the right-hand sides are with respect to i.i.d. replicas \(X, X'\) sampled from the posterior distribution given \(y\).
The proof of Lemma [lem:vecfieldderivative] follows from Lemma 1, namely combining Lemmas 2 and 3:
Proof. We can write \[\begin{align} \partial_{t} (y_{t} + \nabla \ln q_t(y_{t})) &= y_{t} + \nabla \ln q_t(y_{t}) + \partial_t \nabla \ln q_t(y_{t})\\ & = y_{t} + \nabla \ln q_t(y_{t}) + \partial_{t} \nabla \ln q_t(y)\Big|_{y=y_{t}} + \nabla^2 q_t(y_{t})(y_{t} + \nabla \ln q_t(y_{t})). \end{align}\] Now using Lemma 3 with \(v = y_{t} + \nabla \ln q_t(y_{t})\), as well as 8 , we get \[\begin{align} &\nabla^2 q_t(y_{t})(y_{t} + \nabla \ln q_t(y_{t})) = \frac{1}{\sigma_t^2}\nabla \left( {\frac{\mathbb{E}^{t,y} X_t - y}{\sigma_t^2}} \right)\Big|_{y=y_{t}} [y_{t} + \nabla \ln q_t(y_{t})]\\ &= \frac{1}{\sigma_t^2}\left( {-y_t +\frac{y_t- \mathbb{E}^{t,y_t}X_t}{\sigma_t^2} + \mathbb{E}^{t,y_t}\frac{1}{\sigma_t^2}\left( {\left\langle {y_{t}-X_t, y_{t} + \nabla \ln q_t(y_{t})}\right\rangle - \left\Vert {y_{t}}\right\Vert^2 +\frac{\left\Vert {y_{t}}\right\Vert^2}{\sigma_t^2}} \right) (X_t - X'_t)} \right)\\ &= \frac{1}{\sigma_t^2}(\frac{1}{\sigma_t^2} - 1)y_t - \frac{1}{\sigma_t^4}\mathbb{E}^{t,y_t}X_t + \frac{1}{\sigma_t^4}\mathbb{E}^{t,y_t}\left( {2\left\langle {y_{t}, \frac{ X''_t}{\sigma_t^2}}\right\rangle - \left\langle {X_t, y_{t}}\right\rangle - \frac{1}{\sigma_t^2}\left\langle {X_t,X''_t}\right\rangle} \right) (X_t - X'_t). \end{align}\] Combining this with Lemma 3, we obtain \[\begin{align} &\partial_{t} (y_{t} + \nabla \ln q_t(y_{t})) = \partial_t y_t + \partial_t(\nabla \ln q_t(y))\Big|_{y=y_t} + \nabla^2 q_t(y_{t})(y_{t} + \nabla \ln q_t(y_{t})) \\ & = \partial_t y_t + \frac{1}{\sigma_t^2}\partial_t(\mathbb{E}^{t,y} X_t)\Big|_{y=y_t} + \nabla^2 q_t(y_{t})(y_{t} + \nabla \ln q_t(y_{t})) \\ & = \left( {y_t + \frac{\mathbb{E}^{t,y_t}X_t - y_t}{\sigma_t^2}} \right) \\ &\qquad+ \frac{1}{\sigma_t^2}\mathbb{E}^{t,y}X_t + \frac{1-\sigma_t^2}{2\sigma_t^6}\mathbb{E}^{t,y}\left( {\left\langle {y,X_t}\right\rangle + \left\Vert {X_t}\right\Vert^2} \right) \left( {X_t - X'_t} \right) + \frac{4}{\sigma_t^4}\mathbb{E}^{t,y}\left( {\left\langle {y,X_t}\right\rangle - \left\Vert {X_t}\right\Vert^2} \right)\left( {X_t - X'_t} \right)\\ &\qquad+\frac{1}{\sigma_t^2}(\frac{1}{\sigma_t^2} - 1)y_t - \frac{1}{\sigma_t^4}\mathbb{E}^{t,y_t}X_t + \frac{1}{\sigma_t^4}\mathbb{E}^{t,y_t}\left( {2\left\langle {y_{t}, \frac{ X''_t}{\sigma_t^2}}\right\rangle - \left\langle {X_t, y_{t}}\right\rangle - \frac{1}{\sigma_t^2}\left\langle {X_t,X''_t}\right\rangle} \right) (X_t - X'_t)\addtocounter{equation}{1}\label{eq:firstder}. \end{align}\tag{9}\] where the last term is the result of taking derivative with respect to \(\sigma_t\) in the denominator. ◻
Below we prove the two parts of Lemma 1 computing the time and spatial derivatives of the posterior mean \(\mu_t(y) = \mathbb{E}^{t,y}X_t\) respectively.
Lemma 2. Given any time \(t\) and \(y\in \mathbb{R}^d\), the time derivative of the expectation of the posterior distribution \(q^{t,y}\) can be calculated as \[\partial_t \mathbb{E}^{t,y}X_t = \mathbb{E}^{t,y}X_t + \frac{1-\sigma_t^2}{2\sigma_t^4}\mathbb{E}^{t,y}\left( {\left\langle {y,X_t}\right\rangle + \left\Vert {X_t}\right\Vert^2} \right) \left( {X_t - X'_t} \right) + \frac{4}{\sigma_t^2}\mathbb{E}^{t,y}\left( {\left\langle {y,X_t}\right\rangle - \left\Vert {X_t}\right\Vert^2} \right)\left( {X_t - X'_t} \right)\,,\] where \(X,X'\) are independent draws from the posterior distribution \(q^{t,y}\).
Proof. We have \[\begin{align} \mathbb{E}^{t,y}X_t = \frac{\int e^{-\frac{\left\Vert {y - x_t}\right\Vert^2}{2\sigma_t^2}} x_t q_{\sf pre}(\bar x) d\bar x}{\int e^{-\frac{\left\Vert {y - x_t}\right\Vert^2}{2\sigma_t^2}} q_{\sf pre}(\bar x)d\bar x} \end{align}\] with random variables \(\xi \sim \gamma^d\), \(X_t = \sqrt{1-\sigma_t^2}\bar X_t\), and \(\bar X \sim q_{\sf pre}\), where \(y = X_t + \sigma_t \xi\) and we denote the values of \(\bar X\) and \(X_t\) by \(\bar x\) and \(x_t = \sqrt{1-\sigma_t^2} \bar x\), respectively. Then \[\begin{align} \partial_t \mathbb{E}^{t,y} X_t &= \frac{\int e^{-\frac{\left\Vert {y - x_t}\right\Vert^2}{2\sigma_t^2}} x_t q_{\sf pre}(\bar x) d\bar x}{\int e^{-\frac{\left\Vert {y - x_t}\right\Vert^2}{2\sigma_t^2}} q_{\sf pre}(\bar x)d\bar x}\\ &+ 4\frac{\int e^{-\frac{\left\Vert {y - x_t}\right\Vert^2}{2\sigma_t^2}} x_t \frac{\left\langle {y-x_t , x_t}\right\rangle}{2\sigma_t^2} q_{\sf pre}(\bar x) d\bar x}{\int e^{-\frac{\left\Vert {y - x_t}\right\Vert^2}{\sigma_t^2}} q_{\sf pre}(\bar x)d\bar x} - 4\frac{\int e^{-\frac{\left\Vert {y - x_t}\right\Vert^2}{2\sigma_t^2}} x_t q_{\sf pre}(\bar x) d\bar x}{\int e^{-\frac{\left\Vert {y - x_t}\right\Vert^2}{2\sigma_t^2}} q_{\sf pre}(\bar x)d\bar x} \frac{\int e^{-\frac{\left\Vert {y - x_t}\right\Vert^2}{2\sigma_t^2}} \frac{\left\langle {y-x_t , x_t}\right\rangle}{2\sigma_t^2} q_{\sf pre}(\bar x) d\bar x}{\int e^{-\frac{\left\Vert {y - x_t}\right\Vert^2}{\sigma_t^2}} q_{\sf pre}(\bar x)d\bar x}\\ &+ \frac{\int e^{-\frac{\left\Vert {y - x_t}\right\Vert^2}{2\sigma_t^2}} x_t \frac{\left\Vert {y - x_t}\right\Vert^2}{2\sigma_t^4}(1-\sigma_t^2) q_{\sf pre}(\bar x) d\bar x}{\int e^{-\frac{\left\Vert {y - x_t}\right\Vert^2}{2\sigma_t^2}} q_{\sf pre}(\bar x)d\bar x} - \frac{\int e^{-\frac{\left\Vert {y - x_t}\right\Vert^2}{2\sigma_t^2}} x_t q_{\sf pre}(\bar x) d\bar x}{\int e^{-\frac{\left\Vert {y - x_t}\right\Vert^2}{2\sigma_t^2}} q_{\sf pre}(\bar x)d\bar x} \frac{\int e^{-\frac{\left\Vert {y - x_t}\right\Vert^2}{2\sigma_t^2}} \frac{\left\Vert {y - x_t}\right\Vert^2}{2\sigma_t^4}(1-\sigma_t^2) q_{\sf pre}(\bar x) d\bar x}{\int e^{-\frac{\left\Vert {y - x_t}\right\Vert^2}{2\sigma_t^2}} q_{\sf pre}(\bar x)d\bar x}, \end{align}\] where the first line follows from taking \(\partial_t\) derivative w.r.t. \(x_t\) that is not in the exponent (note that \(x_t = \sqrt{1-\sigma_t^2} \bar x = e^{-(T-t)} \bar x\) is a function of \(t\)), the second line is the result of taking derivative w.r.t. \(x_t\) that is in the exponent, and the final line is taking derivative w.r.t. the \(1/\sigma_t^2\) in the exponent. Note that taking derivative w.r.t the numerator and then the denominator in these cases results in a covariance term between \(x_t\) and \(\left\langle {y-x_t,x_t}\right\rangle\) and \(\left\Vert {y-x_t}\right\Vert^2\) in the second and third lines, respectively: \[\begin{align} \partial_t \mathbb{E}^{t,y}X_t &= \mathbb{E}^{t,y}X_t + \operatorname{Cov}\left( {\frac{(1-\sigma_t^2)\left\Vert {y-X_t}\right\Vert^2}{2\sigma_t^4}, X_t} \right) + 4\operatorname{Cov}\left( {\frac{\left\langle {y-X_t, X_t}\right\rangle}{\sigma_t^2}, X_t} \right)\\ &= \mathbb{E}^{t,y}X_t + \frac{1-\sigma_t^2}{2\sigma_t^4}\mathbb{E}^{t,y}\left( {\left\langle {y,X_t}\right\rangle + \left\Vert {X_t}\right\Vert^2} \right) \left( {X_t - \mathbb{E}^{t,y}X_t} \right) + \frac{4}{\sigma_t^2}\mathbb{E}^{t,y}\left( {\left\langle {y,X_t}\right\rangle - \left\Vert {X_t}\right\Vert^2} \right)\left( {X_t - \mathbb{E}^{t,y}X_t} \right)\\ &= \mathbb{E}^{t,y}X_t + \frac{1-\sigma_t^2}{2\sigma_t^4}\mathbb{E}^{t,y}\left( {\left\langle {y,X_t}\right\rangle + \left\Vert {X_t}\right\Vert^2} \right) \left( {X_t - X'_t} \right) + \frac{4}{\sigma_t^2}\mathbb{E}^{t,y}\left( {\left\langle {y,X_t}\right\rangle - \left\Vert {X_t}\right\Vert^2} \right)\left( {X_t - X'_t} \right), \end{align}\] where \(X'_t\) is an independent replica of \(X_t\). ◻
Lemma 3. The directional derivative of the posterior mean with respect to \(y\) is given by \[D(\mathbb{E}^{t,y}X_t)[v] = \frac{1}{\sigma_t^2}\mathbb{E}^{t,y}\left( {\left\langle {y-X_t, v}\right\rangle} \right)(X_t - X_t') = \frac{1}{\sigma_t^2}\mathbb{E}^{t,y}\left( {\left\langle {y-X_t, v}\right\rangle + \frac{\left\Vert {y}\right\Vert^2}{\sigma_t^2} - \left\Vert {y}\right\Vert^2} \right)(X_t - X'_t)\,.\]
Proof. We have \[\begin{align} D \left( {\mathbb{E}^{t,y} X_t} \right)[v] = &-\frac{\int e^{-\frac{\left\Vert {y - x_t}\right\Vert^2}{2\sigma_t^2}} \left\langle {\frac{y-x_t}{\sigma_t^2}, v}\right\rangle x_t q_{\sf pre}(\bar x) d\bar x}{\int e^{-\frac{\left\Vert {y - x_t}\right\Vert^2}{2\sigma_t^2}} q_{\sf pre}(\bar x)d\bar x} + \frac{\int e^{-\frac{\left\Vert {y - x_t}\right\Vert^2}{2\sigma_t^2}} x_t q_{\sf pre}(\bar x) d\bar x}{\int e^{-\frac{\left\Vert {y - x_t}\right\Vert^2}{2\sigma_t^2}} q_{\sf pre}(\bar x)d\bar x}\frac{\int e^{-\frac{\left\Vert {y - x_t}\right\Vert^2}{2\sigma_t^2}} \left\langle {\frac{y-x_t}{\sigma_t^2}, v}\right\rangle q_{\sf pre}(\bar x) d\bar x}{\int e^{-\frac{\left\Vert {y - x_t}\right\Vert^2}{2\sigma_t^2}} q_{\sf pre}(\bar x)d\bar x}\\ &= \operatorname{Cov}\left( {\left\langle {\frac{y-X_t}{\sigma_t^2}, v}\right\rangle, X_t} \right)\\ &= \mathbb{E}^{t,y} \frac{1}{\sigma_t^2}\left( {\left\langle {y-X_t, v}\right\rangle + \frac{\left\Vert {y}\right\Vert^2}{\sigma_t^2} - \left\Vert {y}\right\Vert^2} \right) (X_t - \mathbb{E}^{t,y} X_t)\\ &= \frac{1}{\sigma_t^2}\mathbb{E}^{t,y} \left( {\left\langle {y-X_t, v}\right\rangle + \frac{\left\Vert {y}\right\Vert^2}{\sigma_t^2} - \left\Vert {y}\right\Vert^2} \right) (X_t - X'_t), \end{align}\] where we used the fact that \(y\) is constant w.r.t the covariance calculation, hence adding \(\frac{\left\Vert {y}\right\Vert^2}{\sigma_t^2} - \left\Vert {y}\right\Vert^2\) does not change the value of the covariance. ◻
In this section, we prove the bound in Lemma [lem:derformula] on the higher derivatives of the score function, restated here for convenience:
Proof. To control the higher derivatives \(\partial_{t}^k \left( {y_{t} + \nabla \ln q_t(y_{t})} \right)\), Lemma [lem:vecfieldderivative] motivates us to analyze what happens after taking derivative from terms of the form \[\begin{align} &\frac{(1 - \sigma_t^2)^{r_1}}{\sigma_t^{r_2}} \mathbb{E}^{t,y_t}\prod_{j=1}^{r_3} \left\langle {y_t, X^{(\mathbf{i}_j)}_t}\right\rangle \prod_{j=1}^{r_4}\left\langle { X^{(\mathbf{k}_j)}_t, X^{(\mathbf{l}_j)}_t}\right\rangle X_t, \\ &\frac{(1 - \sigma_t^2)^{r_1}}{\sigma_t^{r_2}}\mathbb{E}^{t,y_t}\prod_{j=1}^{r_3} \left\langle {y_t, X^{(\mathbf{i}_j)}_t}\right\rangle \prod_{j=1}^{r_4}\left\langle { X^{(\mathbf{k}_j)}_t, X^{(\mathbf{l}_j)}_t}\right\rangle y_t,\label{eq:convterms} \end{align}\tag{10}\] for some integers \(r_1, r_2, r_3, r_4, r \leq k\); such terms naturally appear after taking \(k\) derivatives. In particular, we have \[\begin{align} \partial_{t} \mathbb{E}^{t,y_t}X_t &= \partial_{t} \mathbb{E}^{t, y} X_t\Big|_{y=y_t} + D(\mathbb{E}^{t,y} X_t)[y_t + \nabla \ln q_t(y_t)]\\ &= \mathbb{E}^{t,y}X_t + \frac{1-\sigma_t^2}{2\sigma_t^4}\mathbb{E}^{t,y}\left( {\left\langle {y,X_t}\right\rangle + \left\Vert {X_t}\right\Vert^2} \right) \left( {X_t - X'_t} \right) + \frac{4}{\sigma_t^2}\mathbb{E}^{t,y}\left( {\left\langle {y,X_t}\right\rangle - \left\Vert {X_t}\right\Vert^2} \right)\left( {X_t - X'_t} \right)\\ &+\frac{1}{\sigma_t^2}(\frac{1}{\sigma_t^2} - 1)y_t - \frac{1}{\sigma_t^4}\mathbb{E}^{t,y_t}X_t + \frac{1}{\sigma_t^4}\mathbb{E}^{t,y_t}\left( {2\left\langle {y_{t}, \frac{ X''_t}{\sigma_t^2}}\right\rangle - \left\langle {X_t, y_{t}}\right\rangle - \frac{1}{\sigma_t^2}\left\langle {X_t,X''_t}\right\rangle} \right) (X_t - X'_t), \end{align}\] and \[\begin{align} \partial_{t}(y_{t}) = y_{t} + \nabla \ln q_t(y_{t}) = y_{t} + \frac{\mathbb{E}^{t,y_t}X_t - y_{t}}{\sigma_t^2}. \end{align}\] The key here is to keep track of the total number of \(X_t\)’s and \(y_t\)’s that appear in Equation 10 . Let’s denote this number by \(\kappa = 2(r_3 + r_4)\). Then, first note that each step of taking derivatives, we are either taking derivative of some \(y_t\) or some \(X_t\) in Equation 10 . In either case, the total number of \(X_t\) and \(y_t\)’s increase by at most \(2\), and the power of \(\sigma_t^2\) in the denominator, i.e. \(r_2\), increases by at most \(4\). Hence, \(\kappa \leq 2k\). Taking derivative wrt \(\sigma_t\): \[\partial_{t} \frac{(1-\sigma_t^2)^{r_1}}{\sigma_t^{r_2}} = -r_2 \frac{(1-\sigma_t^2)^{r_1 + 1}}{\sigma_t^{r_2 + 2}} + \frac{-2r_1(1-\sigma_t^2)^{r_1 - 1}}{\sigma_t^{r_2 - 1}},\label{eq:wrtsigma}\tag{11}\] which increases \(r_1\) by at most one and \(r_2\) by at most two. Therefore, \(r_1 \leq k\) and \(r_2 \leq 4k\). On the other hand, taking derivative with respect to either \(y_t\) or \(X_t\) creates at most \(34 = 1 + 2 + 4 \times 2\times 2 + 2\times 2 + 2 + 1 + 4 \times 2\) new terms (counting term with coefficient \(4\) four times) and taking derivative wrt \(\sigma_t\) (according to Equation 11 ) creates at most \(r_2 + 2r_1 \leq 6k\), which means overall we have at most \((2\times 34 + 6)k = 74k\) new terms generated after taking \(k\) derivatives. Therefore, overall, after taking \(k\) derivatives we have at most \((74k)^k\) terms generated. ◻
We formally state and prove our upper bound on the derivatives \(\partial_{t}^k \left( {y_{t} + \nabla \ln q_t(y_{t})} \right)\) in Lemma [lem:bnd1], restated here for convenience:
Proof. The proof directly follows from Lemma [lem:derformula] in conjunction with Lemma 4 below. ◻
In Lemma 4 we bound the \(p\)-norm of the terms that appear in Lemma [lem:derformula].
Lemma 4. We have \[\begin{align} &\left\Vert {\frac{(1 - \sigma_t^2)^{r_1}}{\sigma_t^{r_2}}\prod_{j=1}^{r_3} \left\langle {y_t, X_t^{(\mathbf{i}_j)}}\right\rangle \prod_{j=1}^{r_4}\left\langle {X_t^{(\mathbf{k}_j)}, X_t^{(\mathbf{l}_j)}}\right\rangle X_t}\right\Vert_{p, \infty} \;\vee \;\left\Vert {\frac{(1 - \sigma_t^2)^{r_1}}{\sigma_t^{r_2}}\prod_{j=1}^{r_3} \left\langle {y, X_t^{(\mathbf{i}_j)}}\right\rangle \prod_{j=1}^{r_4}\left\langle {X_t^{(\mathbf{k}_j)}, X_t^{(\mathbf{l}_j)}}\right\rangle y_t}\right\Vert_{p, \infty} \\ &\qquad \leq R \left( {\frac{1}{\sigma_t^2}} \right)^k\left( {\frac{R}{\sigma_t} + (kp)^{1/2}} \right)^{2k}, \end{align}\] where by \(\left\Vert {v}\right\Vert_{p, \infty}\) we mean the infinity norm of the \(p\)th moment of vector \(v\).
Proof. First, from the fact that the radius of support of \(q_0\) is bounded by \(R\), we have \[\begin{align} \left\Vert {\frac{(1 - \sigma_t^2)^{r_1}}{\sigma_t^{r_2}}\prod_{j=1}^{r_3} \left\langle {y_t, X_t^{(\mathbf{i}_j)}}\right\rangle \prod_{j=1}^{r_4}\left\langle {X_t^{(\mathbf{k}_j)}, X_t^{(\mathbf{l}_j)}}\right\rangle X_t}\right\Vert_{p,\infty} \leq \frac{1}{\sigma_t^{r_2}} \left\Vert {\prod_{j=1}^{r_3} \left\langle {y_t, X_t^{(\mathbf{i}_j)}}\right\rangle}\right\Vert_p R^{2r_4 + 1}. \end{align}\] Now using the derivation in Lemma 4.4 in [66], we get \[\begin{align} \left\Vert {\prod_{j=1}^{r_3} \left\langle {y_t, X_t^{(\mathbf{i}_j)}}\right\rangle}\right\Vert_p \leq R^{2r_3} \left( {1 + \frac{(r_3 p)^{1/2}\sigma_t}{R}} \right)^{r_3}. \end{align}\] Plugging this above and using \(R\geq 1\): \[\begin{align} \mathrm{LHS} \leq \frac{R^{2r_3 + 2r_4 + 1}}{\sigma_t^{r_2}} \left( {1 + \frac{(r_3 p)^{1/2}\sigma_t}{R}} \right)^{r_3} &\leq R\left( {\frac{R}{\sigma_t^2}} \right)^{2k} \left( {1 + \frac{(k p)^{1/2}\sigma_t}{R}} \right)^{k}\\ & \leq R\left( {\frac{1}{\sigma_t}} \right)^{2k}\left( {\frac{R}{\sigma_t}} \right)^{2k} \left( {1 + \frac{(k p)^{1/2}\sigma_t}{R}} \right)^{2k}\\ & \leq R \frac{1}{\sigma_t^{2k}}\left( {\frac{R}{\sigma_t} + (kp)^{1/2}} \right)^{2k} \end{align}\] Similarly for the other term, using Cauchy-Schwarz \[\begin{align} \left\Vert {\frac{(1 - \sigma_t^2)^{r_1}}{\sigma_t^{r_2}}\prod_{j=1}^{r_3} \left\langle {y, X_t^{(\mathbf{i}_j)}}\right\rangle \prod_{j=1}^{r_4}\left\langle {X_t^{(\mathbf{k}_j)}, X_t^{(\mathbf{l}_j)}}\right\rangle y}\right\Vert_{p,\infty} &\leq \frac{1}{\sigma_t^{r_2}} \left\Vert {\prod_{j=1}^{r_3} \left\langle {y, X_t^{(\mathbf{i}_j)}}\right\rangle}\right\Vert_{2p} R^{2r_4}\left\Vert {y}\right\Vert_{2,\infty}\\ &\leq \frac{1}{\sigma_t^{r_2-1}} R^{2r_3} \left( {1 + \frac{(2r_3 p)^{1/2}\sigma_t}{R}} \right)^{r_3} R^{2r_4}\\ &\leq \frac{1}{\sigma_t^{2k}}\left( {\frac{R}{\sigma_t} + (kp)^{1/2}} \right)^{2k}. \qedhere \end{align}\] ◻
We prove Theorem [thm:derivative] by translating the inequality in Lemma [lem:bnd1] from \(y\) to a close by \(\tilde{y}\) via a coupling between the distribution of \(X_t\) conditioned on \(y\) and \(\tilde{y}\). Note that the argument of Lemma [lem:bnd1] is not immediate for \(\tilde{y}\) because it does not follow the true reverse process as \(y\) does.
First, recall that \(\bar X \sim q_{\sf pre}\) and \(\xi \sim \gamma^d\) and \(y_t= X_t + \sigma_t \xi\), \(X_t = \sqrt{1-\sigma_t^2}\bar X\). Now for every \(\varepsilon_1 > 0\), consider the event \[\begin{align} \mathcal{A}_1 \mathrel{\vcenter{:}}= \left \{ \left\Vert {y_t-X_t}\right\Vert \leq \sqrt d + \sqrt{\ln(1/\epsilon_1)} \right \},\label{eq:a1event} \end{align}\tag{12}\] and note that \(\mathbb{P}(\mathcal{A}_1) \geq 1 - \epsilon_1\), since \(\sigma_t \leq 1\).
Proof. Since \(T - t \geq 1\), we have \[\begin{align} \sigma_t^2 = 1- e^{-2(T-t)} \geq 1 - e^{-2} \geq 0.5. \end{align}\] Let \(p(x,y)\) denote the joint distribution of \((X_t,y_t)\), for \(y_t = X_t + \xi\), \(X_t = \sqrt{1-\sigma_t^2} \bar X\), \(\bar X\sim q_{\sf pre}\), and \(\xi \sim N(0, \sigma_t^2 I)\). Note that \(q^{t,y} = p(\cdot \mid y)\). We have \[\begin{align} \Big|\ln\left( {\frac{p(y\mid x)}{p(\tilde{y}\mid x)}} \right)\Big| &= \frac{1}{2\sigma_t^2}\Big|\left\Vert {y-x}\right\Vert^2 - \left\Vert {\tilde{y} - x}\right\Vert^2\Big|\\ &\leq \Big|\left\Vert {y-x}\right\Vert^2 - \left\Vert {\tilde{y} - x}\right\Vert^2\Big|\\ &\leq 2\left\Vert {y-\tilde{y}}\right\Vert\left\Vert {y-x}\right\Vert + \left\Vert {y-\tilde{y}}\right\Vert^2\\ &\leq 2\delta (\sqrt d + \sqrt{\ln(1/\epsilon_1)}) + \delta^2\\ &\leq 3\delta (\sqrt d + \sqrt{\ln(1/\epsilon_1)}), \end{align}\] and by the assumed bound on \(\delta\), we have \(2\delta (\sqrt d + \sqrt{\ln(1/\epsilon_1)}) + \delta^2 \leq 0.5\), which implies \[(1 - 3\delta (\sqrt d + \sqrt{\ln(1/\epsilon_1)})) p(y\mid x) \leq p(\tilde{y}\mid x) \leq p(y\mid x) (1 + 3\delta (\sqrt d + \sqrt{\ln(1/\epsilon_1)})).\] Multiplying \(p(x)\) on both sides, we get \[\begin{align} (1 - 3\delta (\sqrt d + \sqrt{\ln(1/\epsilon_1)})) p(y\mid x)p(x) \leq p(\tilde{y}\mid x)p(x) \leq p(y\mid x)p(x) (1 + 3\delta (\sqrt d + \sqrt{\ln(1/\epsilon_1)})). \end{align}\] But from Bayes’ rule, it is easy to see that this implies \[\begin{align} (1 - 3\delta (\sqrt d + \sqrt{\ln(1/\epsilon_1)}))^2 q^{t,y}(x) \leq q^{t,\tilde{y}}(x) \leq q^{t,y}(x) (1 + 3\delta (\sqrt d + \sqrt{\ln(1/\epsilon_1)}))^2. \end{align}\] Given that \((1 + 3\delta (\sqrt d + \sqrt{\ln(1/\epsilon_1)}))^2 \leq 1 + 8\delta (\sqrt d + \sqrt{\ln(1/\epsilon_1)})\) by the assumed bound on \(\delta\), we conclude that \[\begin{align} \mathrm{TV}(q^{t,y}, q^{t,\tilde{y}}) &\leq 8\delta (\sqrt d + \sqrt{\ln(1/\epsilon_1)}). \qedhere \end{align}\] ◻
Here we show a similar bound on the derivatives of the vector field for the curve \(\tilde{y}\) instead of \(y\), where \(y\) and \(\tilde{y}\) are defined in Section 6. Throughout this section, we assume \(\tilde{y}_t\) and \(y_t\) are close by \(\delta\) in the window \([t_0, t_0 + h]\): \[\begin{align} \left\Vert {y_t - \tilde{y}_t}\right\Vert_{[t_0, t_0 + h]} \leq \delta.\label{eq:closenesscondition} \end{align}\tag{13}\] Later on in the proof of Theorem 3, we will show by induction that it is valid to assume Equation 13 holds. When it is clear, we will drop the \(t\) indices from the variables for clarity. First, we restate Theorem [thm:derivative] and show the short proof based on the rest of this section. We then get into the details of proving the intermediate lemmas.
Proof. For simplicity of exposition we will drop the time index \(t\) from \(y_t, X_t\) in the rest of this section, and write \(y = X + \sigma_t\xi\), where \(\bar X\sim q_{\sf pre}\), \(X = \sqrt{1-\sigma_t^2} \bar X\) and \(\xi \sim \gamma^d\). Let \(p(.|y)\) be the distribution of \(X\) conditioned on \(y\) and let \(\tilde{X}\) be a sample from the conditional \(p(.|\tilde{y})\). For index \(i\) let \(X^{(i)}\) and \(\tilde{X}^{(i)}\) be an independent and identical sample of \(X\) and \(\tilde{X}\), respectively. Now using Lemma 10 with \(\delta_1 = \frac{\delta_2}{d(60k)^k}\), we get \[\mathbb{P}\left( {\Big|\mathbb{E}_{\{\tilde{X}^{(i)}\}}\frac{(1 - \sigma_t^2)^{r_1}}{\sigma_t^{r_2}}\prod_{j=1}^{r_3} \left\langle {\tilde{y}, \tilde{X}^{(\mathbf{i}_j)}}\right\rangle \prod_{j=1}^{r_4}\left\langle {\tilde{X}^{(\mathbf{k}_j)}, \tilde{X}^{(\mathbf{l}_j)}}\right\rangle X\Big|_\infty \geq eR\left( {\frac{1}{\sigma_t^2}} \right)^k\left( {\frac{R}{\sigma_t}} \right)^{2k}+ e \left( {k\log(74kd/\delta_2)} \right)^k} \right) \leq \delta_2\,.\] Taking union bound over all the terms, we get with probability at least \(1-\delta_2\) \[\begin{align} \left\Vert {\partial_{t}^k \left( {\tilde{y}_{t} + \nabla \log q_t(\tilde{y}_{t})} \right)}\right\Vert_{\infty} &\leq e\left( {R + \delta + \sigma_t + 1} \right)\left( {\frac{kh}{\sigma_t^2}} \right)^k\left( {\frac{R}{\sigma_t}} \right)^{2k}+ e \left( {k\log(74kd/\delta_2)} \right)^{2k}\,. \qedhere \end{align}\] ◻
Next, we derive a pessimistic bound on \(\mathbb{E}\Big|\prod_{j=1}^{r_3} \left\langle {\tilde{y}, X^{(\mathbf{i}_j)}}\right\rangle\Big|\) for all ordered tuples of indices \(\mathbf{i}\), where each \(X^{(i)}\) denotes an independent sample from the posterior distribution \(q^{t,y}\), by obtaining a pessimistic bound on the term \(\left\Vert {\prod_{j=1}^{r_3} \left\langle {y, X^{(\mathbf{i}_j)}}\right\rangle}\right\Vert_p\).
Lemma 5. For all tuples of indices \(\mathbf{i}\), we have \[\begin{align} \mathbb{E}\Big|\prod_{j=1}^{r_3} \left\langle {\tilde{y}, X^{(\mathbf{i}_j)}}\right\rangle\Big| \leq \left( {R^2\left( {1 + \frac{(r_3 p)^{1/2}\sigma_t}{R}} \right) + \delta R} \right)^{r_3} \end{align}\]
Proof. Using Cauchy-Schwarz, \[\begin{align} \left\Vert {\prod_{j=1}^{r_3} \left\langle {\tilde{y}, X^{(\mathbf{i}_j)}}\right\rangle}\right\Vert_p &\leq \sum_{S \subseteq [k]}^{r_3} \left\Vert {\prod_{j \in S}\left\langle {y,X^{(\mathbf{i}_j)}}\right\rangle}\right\Vert_p \left\Vert {\prod_{j \notin S}\left\langle {\tilde{y} - y,X^{(\mathbf{i}_j)}}\right\rangle}\right\Vert_p \\ &\leq \sum_{\ell=1}^{r_3} \binom{r_3}{\ell} \left( {R^2\left( {1 + \frac{(r_3 p)^{1/2}\sigma_t}{R}} \right)} \right)^{\ell} \left( {\delta R} \right)^{r_3 - \ell} \\ &\leq \left( {R^2\left( {1 + \frac{(r_3 p)^{1/2}\sigma_t}{R}} \right) + \delta R} \right)^{r_3}. \qedhere \end{align}\] ◻
Lemma 6. For \(\nu \triangleq \left( {3\delta (\sqrt d + \sqrt{\ln(1/\epsilon_1)})} \right)^{1/p}\), we have for all tuples of indices \(\mathbf{i}\) that \[\left\Vert {\mathrm 1\Big\{{\mathcal{A}_1}\Big\}\prod_{j=1}^{r_3} \left\langle {\tilde{y},\tilde{X}^{(\mathbf{i}_j)}}\right\rangle}\right\Vert_p \leq \left( {R^2\left( {1 + \frac{(2r_3 p)^{1/2}\sigma_t}{R}} \right) + \delta R + \nu R\left( {\sqrt{2p d} + R + \delta} \right)} \right)^{r_3}\,. \label{eq:deltabasedbound}\qquad{(1)}\]
Proof. First, note that we have the following bound: \[\begin{align} \left\Vert {\prod_{j=1}^{r_3} \left\langle {\tilde{y},\tilde{X}^{(\mathbf{i}_j)}}\right\rangle}\right\Vert_p^p &= \mathbb{E}\prod_{j=1}^{r_3} \left\langle {\tilde{y},\tilde{X}^{(\mathbf{i}_j)}}\right\rangle^p\nonumber\\ &\leq \mathbb{E}\prod_{j=1}^{r_3}\left\Vert {\tilde{y}}\right\Vert^p \left\Vert {\tilde{X}^{(\mathbf{i}_j)}}\right\Vert^p\nonumber\\ &\leq \mathbb{E}\prod_{j=1}^{r_3}\left( {\left\Vert {y - \tilde{y}}\right\Vert + \left\Vert {\xi}\right\Vert + \left\Vert { X^{\mathbf{i}_j}}\right\Vert} \right)^p \left\Vert {\tilde{X}^{(\mathbf{i}_j)}}\right\Vert^p\nonumber\\ &\leq \mathbb{E}\prod_{j=1}^{r_3}\left( {\delta + \left\Vert {\xi}\right\Vert + R} \right)^p R^p\nonumber\\ &\leq \mathbb{E}\sum_{\ell=1}^{r_3 p}\binom{r_3 p}{\ell}\left\Vert {\xi}\right\Vert^\ell \left( {R + \delta} \right)^{r_3 p-\ell} R^{r_3 p}\nonumber\\ &\leq \sum_{\ell=1}^{r_3 p}\binom{r_3 p}{\ell}\left( {\sqrt{\ell d}} \right)^\ell \left( {R + \delta} \right)^{r_3 p-\ell} R^{r_3p} \nonumber\\ &= \left( {R\left( {\sqrt{p d} + R + \delta} \right)} \right)^{r_3p},\label{eq:firstder95new} \end{align}\tag{14}\] which implies \[\begin{align} \left\Vert {\prod_{j=1}^{r_3} \left\langle {\tilde{y},\tilde{X}^{(\mathbf{i}_j)}}\right\rangle}\right\Vert_p \leq \left( {R\left( {\sqrt{p d} + R + \delta} \right)} \right)^{r_3}. \end{align}\] Therefore, applying Lemma [lem:klboundlemma-main] with \(\epsilon_1 = \delta\), on event \(\mathcal{A}_1\) we have \[\begin{align} TV(q^{t,y},q^{t,\tilde{y}}) \leq 3\delta (\sqrt d + \sqrt{\ln(1/\epsilon_1)}), \end{align}\] which means for all \(m\), \[\begin{align} TV(X^{(m)}, \tilde{X}^{(m)}) \leq 3\delta (\sqrt d + \sqrt{\ln(1/\epsilon_1)}). \end{align}\] Therefore, the optimal coupling between \(X^{(m)}\) and \(\tilde{X}^{(m)}\) conditioned on \(y\) satisfies \[\begin{align} \mathbb{P}\left( {X^{(m)} = \tilde{X}^{(m)}} \right) \geq 1 - 3\delta (\sqrt d + \sqrt{\ln(1/\epsilon_1)}) \label{eq:tvbound} \end{align}\tag{15}\] Now equipped with this coupling, for a subset \(\mathcal{U}\subset [r_3]\), let \(\mathcal{A}_2 = \mathcal{A}_2(\mathcal{U})\) be the event that \(X^{(\mathbf{i}_j)} = \tilde{X}^{(\mathbf{i}_j)}\) for all \(j \in \mathcal{U}\). Then we can write \[\begin{align} &\mathbb{E}\left( {\mathrm 1\Big\{{\mathcal{A}_1\cap \mathcal{A}_2(\mathcal{U})}\Big\} \prod_{j=1}^{r_3} \left\langle {\tilde{y},\tilde{X}^{(\mathbf{i}_j)}}\right\rangle} \right)^p \\ &= \mathbb{E} \left( {\mathrm 1\Big\{{\mathcal{A}_1 \cap \mathcal{A}_2(\mathcal{U})}\Big\}\prod_{j\in \mathcal{U}} \left\langle {\tilde{y},\tilde{X}^{(\mathbf{i}_j)}}\right\rangle} \right)^p \left( {\mathrm 1\Big\{{\mathcal{A}_1 \cap \mathcal{A}_2(\mathcal{U})}\Big\}\prod_{j\notin \mathcal{U}} \left\langle {\tilde{y},\tilde{X}^{(\mathbf{i}_j)}}\right\rangle} \right)^p. \\ &\leq \sqrt{\mathbb{E}\left( {\mathrm 1\Big\{{\mathcal{A}_1 \cap \mathcal{A}_2(\mathcal{U})}\Big\}\prod_{j\in \mathcal{U}} \left\langle {\tilde{y},\tilde{X}^{(\mathbf{i}_j)}}\right\rangle} \right)^{2p}} \sqrt{\mathbb{E}\left( {\mathrm 1\Big\{{\mathcal{A}_1 \cap \mathcal{A}_2(\mathcal{U})}\Big\}\prod_{j\notin \mathcal{U}} \left\langle {\tilde{y},\tilde{X}^{(\mathbf{i}_j)}}\right\rangle} \right)^{2p}}\addtocounter{equation}{1}\label{eq:couplingseparation} \end{align}\tag{16}\] But on one hand, from Equation 15 and the independence of the couplings, we have \[\mathbb{P}(\mathcal{A}_2(\mathcal{U}) | \mathcal{A}_1) \leq \left( {3\delta (\sqrt d + \sqrt{\ln(1/\epsilon_1)})} \right)^{r_3 - |\mathcal{U}|}\,.\] Therefore, similar to the derivation in 14 \[\begin{align} &\mathbb{E}\left( {\mathrm 1\Big\{{\mathcal{A}_1 \cap \mathcal{A}_2(\mathcal{U})}\Big\}\prod_{j\notin \mathcal{U}} \left\langle {\tilde{y},\tilde{X}^{(\mathbf{i}_j)}}\right\rangle} \right)^p\\ &\leq \mathbb{E}\mathrm 1\Big\{{\mathcal{A}_1}\Big\}\sum_{\ell=1}^{(r_3 - |\mathcal{U}|)p}\binom{(r_3 - |\mathcal{U}|)p}{\ell}\left\Vert {\xi}\right\Vert^\ell \left( {R + \delta} \right)^{(r_3 - |\mathcal{U}|)p-\ell} R^{(r_3 - |\mathcal{U}|)p} \mathbb{P}_{|y}(\mathcal{A}_2(\mathcal{U}) | \mathcal{A}_1) \\ &\leq \mathbb{P}_{|y}(\mathcal{A}_2(\mathcal{U}) | \mathcal{A}_1)\sum_{\ell=1}^{(r_3 - |\mathcal{U}|)p}\binom{(r_3 - |\mathcal{U}|)p}{\ell}\left\Vert {\xi}\right\Vert^\ell \left( {R + \delta} \right)^{(r_3 - |\mathcal{U}|)p-\ell} R^{(r_3-|\mathcal{U}|)p} \\ &\leq \left( {3\delta (\sqrt d + \sqrt{\ln(1/\epsilon_1)})} \right)^{|\mathcal{U}|}\left( {R\left( {\sqrt{p d} + R + \delta} \right)} \right)^{(r_3 - |\mathcal{U}|) p} \\ &= \nu^{p(r_3 - |\mathcal{U}|)}\left( {R\left( {\sqrt{p d} + R + \delta} \right)} \right)^{(r_3 - |\mathcal{U}|) p}, \end{align}\] where \(\nu \triangleq \left( {3\delta (\sqrt d + \sqrt{\ln(1/\epsilon_1)})} \right)^{1/p}.\) This implies \[\left\Vert {\mathrm 1\Big\{{\mathcal{A}_1 \cap \mathcal{A}_2(\mathcal{U})}\Big\}\prod_{j\notin \mathcal{U}}^{r_3} \left\langle {\tilde{y},\tilde{X}^{(\mathbf{i}_j)}}\right\rangle}\right\Vert_p \leq \left( {3\delta (\sqrt d + \sqrt{\ln(1/\epsilon_1)})} \right)^{1/p}\left( {R\left( {\sqrt{p d} + R + \delta} \right)} \right)^{r_3 - |\mathcal{U}|}\,. \label{eq:notinU}\tag{17}\] For the first term, we can use the coupling and substitute \(\tilde{X}^{(\mathbf{i}_j)}\)’s by \(X^{(\mathbf{i}_j)}\) and then use Lemma 5: \[\begin{align} \mathbb{E}\left( {\mathrm 1\Big\{{\mathcal{A}_1 \cap \mathcal{A}_2(\mathcal{U})}\Big\}\prod_{j\in \mathcal{U}} \left\langle {\tilde{y},\tilde{X}^{(\mathbf{i}_j)}}\right\rangle} \right)^p &= \mathbb{E}\left( {\mathrm 1\Big\{{\mathcal{A}_1 \cap \mathcal{A}_2(\mathcal{U})}\Big\}\prod_{j\in \mathcal{U}} \left\langle {\tilde{y},X^{(\mathbf{i}_j)}}\right\rangle} \right)^p \nonumber\\ &\leq \mathbb{E}\left( {\prod_{j\in \mathcal{U}} \left\langle {\tilde{y},\tilde{X}^{(\mathbf{i}_j)}}\right\rangle} \right)^p \nonumber\\ &\leq \left( {R^2\left( {1 + \frac{(r_3 p)^{1/2}\sigma_t}{R}} \right) + \delta R} \right)^{p |\mathcal{U}|}.\label{eq:firsttermbound} \end{align}\tag{18}\] Combining Equations 18 and 17 and plugging back into Equation 16 : \[\begin{align} &\mathbb{E}\left( {\mathrm 1\Big\{{\mathcal{A}_1}\Big\} \prod_{j=1}^{r_3} \left\langle {\tilde{y},\tilde{X}^{(\mathbf{i}_j)}}\right\rangle} \right)^p \\ &= \sum_{\mathcal{U} \subseteq [r_3]}\mathbb{E}\left( {\mathrm 1\Big\{{\mathcal{A}_1\cap \mathcal{A}_2(\mathcal{U})}\Big\} \prod_{j=1}^{r_3} \left\langle {\tilde{y},\tilde{X}^{(\mathbf{i}_j)}}\right\rangle} \right)^p\\ &\leq \sum_{\mathcal{U} \subseteq [r_3]}\left( {R^2\left( {1 + \frac{(2r_3 p)^{1/2}\sigma_t}{R}} \right) + \delta R} \right)^{p|\mathcal{U}|}\left( {\nu R\left( {\sqrt{2p d} + R + \delta} \right)} \right)^{(r_3 - |\mathcal{U}|) p}\\ &\leq \left( {R^2\left( {1 + \frac{(2r_3 p)^{1/2}\sigma_t}{R}} \right) + \delta R + \nu R\left( {\sqrt{2p d} + R + \delta} \right)} \right)^{r_3 p}. \qedhere \end{align}\] ◻
Lemma 7 (Replica bounds for the coupled variables). Suppose we have \[\begin{align} \delta \leq R \wedge \frac{1}{\left( {\sqrt{d} + \sqrt{\ln(1/\epsilon_1)}} \right)(\sqrt{2pd}/R + 1)^{p}}. \end{align}\] Then for \(\tilde{y}\) as defined in Section 6, let \(\tilde{X}^{(i)}\) be independent samples from the posterior distribution \(q^{t,\tilde{y}}\). Then, for integers \(r_2, r_3, r_4\) that satisfy \(r_2 \leq 2(r_3 + r_4) \leq 2k\), under event \(\mathcal{A}_1\) defined previously, we have for all tuples of indices \(\mathbf{i}, \mathbf{k}, \mathbf{l}\), \[\begin{align} &\left\Vert {\mathrm 1\Big\{{\mathcal{A}_1}\Big\}\frac{(1 - \sigma_t^2)^{r_1}}{\sigma_t^{r_2}}\prod_{j=1}^{r_3} \left\langle {\tilde{y}, \tilde{X}^{(\mathbf{i}_j)}}\right\rangle \prod_{j=1}^{r_4}\left\langle {\tilde{X}^{(\mathbf{k}_j)}, \tilde{X}^{(\mathbf{l}_j)}}\right\rangle \tilde{X}}\right\Vert_{p, \infty} \lesssim \frac{R}{\sigma_t^{2k}}\left( {\frac{R}{\sigma_t} + (kp)^{1/2}} \right)^{2k},\\ &\left\Vert {\mathrm 1\Big\{{\mathcal{A}_1}\Big\}\frac{(1 - \sigma_t^2)^{r_1}}{\sigma_t^{r_2}}\prod_{j=1}^{r_3} \left\langle {\tilde{y}, \tilde{X}^{(\mathbf{i}_j)}}\right\rangle \prod_{j=1}^{r_4}\left\langle {\tilde{X}^{(\mathbf{k}_j)}, \tilde{X}^{(\mathbf{l}_j)}}\right\rangle \tilde{y}}\right\Vert_{p, \infty} \lesssim \frac{R + 1}{\sigma_t^{2k}}\left( {\frac{R}{\sigma_t} + (kp)^{1/2}} \right)^{2k}\,, \end{align}\] where by \(\left\Vert {v}\right\Vert_{p, \infty}\) we mean the infinity norm of the \(p\)th moment of the vector \(v\).
Proof. Note that given the condition on \(\delta\), then in the bound of Equation ?? in Lemma 6 the second term is dominated by the first term, which implies \[\begin{align} \left\Vert {\mathrm 1\Big\{{\mathcal{A}_1}\Big\}\prod_{j=1}^{r_3} \left\langle {\tilde{y},\tilde{X}^{(\mathbf{i}_j)}}\right\rangle}\right\Vert_p &\lesssim \left( {R^2\left( {1 + \frac{(r_3 p)^{1/2}\sigma_t}{R}} \right) + \delta R} \right)^{r_3}\\ &\lesssim \left( {R^2\left( {1 + \frac{(r_3 p)^{1/2}\sigma_t}{R}} \right)} \right)^{r_3}. \end{align}\]
Therefore, from \(r_2 \leq 2(r_3 + r_4)\) and using the fact that support of \(q_0\) has bounded radius \(R\): \[\begin{align} &\left\Vert {\mathrm 1\Big\{{\mathcal{A}_1}\Big\}\frac{(1 - \sigma_t^2)^{r_1}}{\sigma_t^{r_2}}\prod_{j=1}^{r_3} \left\langle {\tilde{y}, \tilde{X}^{(\mathbf{i}_j)}}\right\rangle \prod_{j=1}^{r_4}\left\langle {\tilde{X}^{(\mathbf{k}_j)}, \tilde{X}^{(\mathbf{l}_j)}}\right\rangle \tilde{X}}\right\Vert_{p,\infty}\\ &\leq \frac{1}{\sigma_t^{r_2}} \left\Vert {\mathrm 1\Big\{{\mathcal{A}_1}\Big\}\prod_{j=1}^{r_3} \left\langle {\tilde{y}, \tilde{X}^{(\mathbf{i}_j)}}\right\rangle}\right\Vert_p R^{2r_4 + 1}\\ & \lesssim R\left( {\frac{1}{\sigma_t^2}} \right)^{r_3 + r_4} \left( {\frac{R^2}{\sigma_t^{2}}} \right)^{r_4} \left( {\frac{R^2}{\sigma_t^2} + \frac{R}{\sigma_t}(r_3p)^{1/2}} \right)^{r_3} \\ & \leq \frac{R}{\sigma_t^{2k}}\left( {\frac{R}{\sigma_t}} \right)^{2r_4} \left( {\left( {\frac{R}{\sigma_t} + (r_3p)^{1/2}} \right)^2} \right)^{r_3} \\ &\leq \frac{R}{\sigma_t^{2k}}\left( {\frac{R}{\sigma_t} + (kp)^{1/2}} \right)^{2k}. \end{align}\] where we used the fact that \(r_3 + r_4 \leq k\). Similarly for the other term, using Cauchy Swartz \[\begin{align} &\left\Vert {\mathrm 1\Big\{{\mathcal{A}_1}\Big\} \frac{(1 - \sigma_t^2)^{r_1}}{\sigma_t^{r_2}}\prod_{j=1}^{r_3} \left\langle {\tilde{y}, \tilde{X}^{(\mathbf{i}_j)}}\right\rangle \prod_{j=1}^{r_4}\left\langle {\tilde{X}^{(\mathbf{k}_j)}, \tilde{X}^{(\mathbf{l}_j)}}\right\rangle y}\right\Vert_{p,\infty}\\ &\leq \frac{1}{\sigma_t^{r_2}} \left\Vert {\prod_{j=1}^{r_3} \left\langle {\tilde{y}, \tilde{X}^{(\mathbf{i}_j)}}\right\rangle}\right\Vert_{2p} R^{2r_4}\left( {\left\Vert {\tilde{y} - y}\right\Vert_{2,\infty} + \left\Vert {y}\right\Vert_{2,\infty}} \right)\\ &\leq \frac{1}{\sigma_t^{r_2}} \left\Vert {\prod_{j=1}^{r_3} \left\langle {\tilde{y}, \tilde{X}^{(\mathbf{i}_j)}}\right\rangle}\right\Vert_{2p} R^{2r_4}\left( {R + \sigma_t + \delta} \right)\\ &\lesssim \frac{R + 1}{\sigma_t^{2k}}\left( {\frac{R}{\sigma_t} + (kp)^{1/2}} \right)^{2k}.\qedhere \end{align}\] ◻
Lemma 8 (High probability control on higher derivatives – fixed time). In the setting of Lemma 7 let \[\begin{align} g(t) = \tilde{y}_{t} + \nabla \log q_t(\tilde{y}_{t}), \end{align}\] where \(\tilde{y}(t)\) is defined in Eq. 7 . Then, for constant \(c_1\) and arbitrary coordinate \(i\), \[\begin{align} &\mathbb{P}\left( {\Big|\partial^k_{t} g_i(t)\Big| \geq c_1(60k)^k (R + 1)\left( {\frac{1}{\sigma_t^2}} \right)^k \left( {\frac{R}{\sigma_t} + \sqrt{\ln(1/\epsilon_1)}} \right)^{2k}} \right) \leq \epsilon_1. \end{align}\]
Proof. Using Lemmas 7 and [lem:derformula], there is a constant \(c_1\) such that \[\begin{align} \left\Vert {\partial^k_{t} g_i(t)}\right\Vert_{p} &= \left\Vert {\partial^k_{t} (y_{t, i} + \partial_i \log q_t(y_{t}))}\right\Vert_{p}\\ &\leq (c_1/e)(60k)^k (R+1) \left( {\frac{1}{\sigma_t^2}} \right)^k\left( {\frac{R}{\sigma_t} + (kp)^{1/2}} \right)^{2k}. \end{align}\] Moreover, for fixed \(j \leq d\), let \(\mathcal{A}_3 \subseteq \mathcal{A}_1\) be the event inside \(\mathcal{A}_1\) where \[\Big|\partial^k_{t} g_i(t)\Big| \geq c_1(60k)^k (R + 1)\left( {\frac{1}{\sigma_t^2}} \right)^k\left( {\frac{R}{\sigma_t} + \sqrt{\ln(1/\epsilon_1)}} \right)^{2k}.\] Taking expectation with respect to \(y\) \[\begin{align} &\mathbb{E}_y \mathbb{E}_{\{\tilde{X}^{(i)}\}} \Big|\partial^k_{t} g_i(t)\Big|^p \\ &\geq \mathbb{P}(\mathcal{A}_3) \left( {c_1(60k)^k(R + 1)\left( {\frac{1}{\sigma_t^2}} \right)^k\left( {\frac{R}{\sigma_t} + \sqrt{\ln(1/\epsilon_1)}} \right)^{2k}} \right)^p. \end{align}\] Now setting \(p = \ln(1/\epsilon_1)/k\), we get \[\begin{align} \mathbb{P}(\mathcal{A}_3) &\leq \left( {\frac{\frac{R}{\sigma_t}+ (kp)^{1/2}}{e\frac{R}{\sigma_t} + e \sqrt{\ln(1/\epsilon_1)}}} \right)^{2kp}= \left( {\frac{1}{e}} \right)^{2\ln(1/\epsilon_1)} \leq \epsilon_1\,.\qedhere \end{align}\] ◻
Lemma 9. In the same setting of Lemma 8, for \[\begin{align} \delta \leq R \wedge \frac{1}{\left( {\sqrt{d} + \sqrt{\ln(1/\epsilon_1)}} \right)(\sqrt{2\ln(1/\epsilon)d}/(R\sqrt k) + 1)^{\ln(1/\epsilon_1) / k}}\,, \end{align}\] for some constant \(c_1\), we have \[\begin{align} \mathbb{P}\left( {\sup_{\tilde{y}: \left\Vert {\tilde{y} - y}\right\Vert\leq \delta} \Big|\partial_t^k g_i(t)\Big|_{s=0}\Big| \geq c_1(60k)^k (R + 1)\left( {\frac{1}{\sigma_t^2}} \right)^k \left( {\frac{R}{\sigma_t} + \sqrt{\ln(1/\epsilon_1)}} \right)^{2k}} \right) \leq \epsilon_1. \end{align}\]
Proof. Note that the \(\tilde{y}\) can be an arbitrary mapping of \(y\) in the neighborhood \(\left\Vert {y - \tilde{y}}\right\Vert\leq \delta\). Hence, we pick \(\tilde{y}\) to be \[\begin{align} \tilde{y}(y) \triangleq \text{argmax}_{\tilde{y}: \left\Vert {\tilde{y} - y}\right\Vert\leq \delta} \Big|\partial_t^k g_i(t)\Big|_{t=0}\Big|. \end{align}\] The result then follows from Lemma 8. ◻
Lemma 10. Let the distance between the two curves, namely the true and algorithm curves, be upper bounded by \[\begin{align} \tilde{\delta} \triangleq R \wedge \frac{1}{\left( {\sqrt{d} + \sqrt{\ln(1/\epsilon_1)}} \right)(\sqrt{2\ln(1/\epsilon_1) d}/(R \sqrt k) + 1)^{\ln(1/\epsilon_1) / k}}\label{eq:smallenoughbound} \,, \end{align}\qquad{(2)}\] that is \(\sup_{t\in [t_0, t_0 + h]} \norm{y(t) - \tilde{y}(t)} \le \tilde{\delta}\), where \(\tilde{y}(t)\) is defined in Lemma 9 and \(y\) is defined analogously with \(y(0) = y\). Then \[\begin{align} \mathbb{P}\left( {\sup_{t\in [t_0, t_0 + h]} \Big|\partial_t^k g_i(t)\Big|_{t=0}\Big| \geq C_{1}} \right) \leq 2\epsilon_1. \end{align}\] where \[C_{1} \triangleq c_1(60k)^k (R + 1)\left( {\frac{1}{\sigma_t^2}} \right)^k \left( {\frac{R}{\sigma_t} + \sqrt{\ln\left( {\frac{R + \sqrt d + \sqrt{\ln(1/\epsilon_1)}}{\epsilon_1}} \right)}} \right)^{2k}.\]
Proof. Suppose we take a cover \(\mathcal{C}\) of the interval \([t_0, t_0 + h]\) with accuracy \(\frac{\tilde{\delta}}{6\left( {R + \sqrt d + \sqrt{\ln(1/\epsilon_1)}} \right)}\). Then, using Lemma 9 with a union bound, particularly since \(2\tilde{\delta}\) satisfies the condition of this lemma, we get (using \(h \leq 1\)) \[\begin{align} \mathbb{P}\left( {\sup_{\tilde{y}: \left\Vert {\tilde{y} - y}\right\Vert\leq 2\tilde{\delta}, y\in \mathcal{C}} \Big|\partial_t^k g_i(t)\Big|_{t=0}\Big| \geq C_1} \right) \leq \epsilon_1.\label{eq:mainprob} \end{align}\tag{19}\]
On the other hand, for every \(s\in [t_0, t_0+h]\), if we assume \(\bar s \in \mathcal{C}\) be the closest element of the cover to \(s\), then from Lemma 16 and the choice of accuracy of the cover: \[\begin{align} \left\Vert { y(\bar t) - y(t)}\right\Vert \leq \tilde{\delta}. \end{align}\] Combining this with the distance between the curves and triangle inequality, we get \[\begin{align} \left\Vert {\tilde{y}(t) - y(\bar t)}\right\Vert \leq 2\tilde{\delta}. \end{align}\] Therefore, we showed that every point on the curve \(\tilde{y}\) in time interval \([t_0, t_0 + h]\) is close to a point in the cover \(\mathcal{C}\). Combining this with Equation 19 implies with probability at least \(1-2\epsilon_1\) \[\begin{align} \mathbb{P}\left( {\sup_{t\in [t_0, t_0 + h]} \Big|\partial_t^k g_i(t)\Big|_{t=0}\Big| \geq C_{1}} \right) &\leq 2\epsilon_1\,. \qedhere \end{align}\] ◻
In this section we exploit the framework we developed for obtaining a low-degree approximation of the score function on the probability flow ODE using the Picard iteration.
Here we recall the basic setup for Picard iteration with polynomial approximation, as developed in [1].
Definition 2. Given vector field \(F(x):\mathbb{R}^d \rightarrow \mathbb{R}^d\), motivated by the Picard iteration and following [1], define the operator \(T(x)\) that acts on a curve \(x(t): [t_0, t_0 + h] \rightarrow \mathbb{R}^d\) as \[\begin{align} T(x)(t) = x(t_0) + \int_{t_0}^{t_0 + h} F(x(s)) ds. \end{align}\] Furthermore, suppose we have a basis \(\{\phi_j\}_{j=1}^D\) of smooth one dimensional functions and points \(\{c_j\}_{j=1}^D\) such that \(1\leq \forall i,j \leq D\), \(\phi_j(c_i) = 0\) if \(i\neq j\) and \(\phi_i(c_i) = 1\). Then, define the approximation \(T_{\phi}\) of the \(T\) operator corresponding to the basis \(\{\phi_j\}_{j=1}^D\) by \[\begin{align} T_\phi(x)(t) = \int_{t_0}^{t_0 + h} \sum_{j=1}^D F(x(c_j))\phi_j(s) ds.\label{eq:approxop} \end{align}\qquad{(3)}\]
In particular, we pick \(\phi_j\)’s to be a basis for one dimensional polynomials of degree less than \(D\); we choose them as the Lagrange multiplier polynomials for points \(c_j\). This way, the integral in ?? can be computed in closed form.
Definition 3. We say the basis \(\phi = \{\phi_j\}_{j=1}^D\) is \(\gamma_\phi\) bounded if \[\begin{align} \sum_j \Bigl|\int_{t_0}^{t_0 + h} \phi_j(s) \, ds\Bigr| \leq \gamma_\phi h. \end{align}\]
Next, we generalize the approach of [1] and show an important Lipschitz property of \(T_{\phi}\) with respect to the supremum norm \(\|.\|\), when the vector field \(F\) is close to a Lipschitz one.
Lemma 11. Suppose the basis \(\phi\) is \(\gamma_\phi\) bounded, and the vector field \(F_s\) is \(\tilde{L}\) Lipschitz. Then, for arbitrary \(x,y \in \mathcal{C}([t_0,t_0 + h], \mathbb{R}^d)\) with \(x(t_0) = y(t_0)\), \[\begin{align} &\|T_\phi(x) - T_\phi(y)\|_{[t_0,t_0 + h]} \leq \tilde{L}\|x - y\|_{[t_0,t_0 + h]} \gamma_\phi h,\nonumber\\ &\|T_\phi^{\circ \ell}(x) - T_\phi^{\circ \ell}(y)\|_{[t_0,t_0 + h]} \leq \left( {\tilde{L}\gamma_\phi h} \right)^\ell \left\Vert {x-y}\right\Vert_{[t_0,t_0 + h]}.\label{eq:lcombinations} \end{align}\qquad{(4)}\]
Proof. For every \(1\leq s\leq h\), using \(\tilde{L}\) Lipschitz property of our estimate for the score function, \[\begin{align} \|F_{t}(x(t)) - F_t(y(t))\| &\leq \tilde{L}\|x(t) - y(t)\|\\ &\leq \tilde{L}\|x - y\|_{[t_0,t_0 + h]}. \end{align}\] Hence \[\begin{align} \|T_\phi(x) - T_\phi(y)\|_{[t_0,t_0 + h]} &= \sup_{t_0\leq h'\leq t_0 + h} \Big\|\int_{t_0}^{t_0 + h'} \sum_{j=1}^D \Big(F_{c_j}(x(c_j)) - F_{c_j}(y(c_j))\Big)\phi_j(s) ds\Big\|\\ &\leq \sup_{1\leq j\leq D} \left\Vert {F_{c_j}(x(c_j)) - F_{c_j}(y(c_j))}\right\Vert_\infty \sum_j \Big|\int_{t_0}^{t_0 + h} \phi_j(s) ds\Big|\\ &\leq \tilde{L}\|x - y\|_{[t_0,t_0 + h]}\gamma_\phi h. \end{align}\] the second line in ?? follows from applying the first line in ?? \(\ell\) times. ◻
Definition 4. (Low-degree vector field)Let \(y_t\) be the solution to the ODE \(\dot{y}_t= F^*_t(y_t)\). We say the vector field \(F^*\) is low-degree along \(y_t\)* if it accepts the following low degree approximation: \[\begin{align} \left\Vert {F^* \circ y - P_{\leq D} (F^* \circ y)}\right\Vert_{[t_0, t_0 + h]} \leq \epsilon_{\sf ld}, \end{align}\] where the curve \(P_{\leq D} (F^*\circ y)\) is an approximation of \(F^* \circ y\) whose coordinates are degree at most \(D\) polynomials in the time variable \(t\in [t_0, t_0 + h]\).*
Lemma 12. For \(\dot{y}_t= F^*_t(y_t)\), given that \(F^*\) is a low-degree vector field based on Definition 4, we have \[\begin{align} \left\Vert {T_\phi^{\circ m} y - y}\right\Vert \leq \frac{\left( {\tilde{L}\gamma_\phi h} \right)^{m} - 1}{\left( {\tilde{L}\gamma_\phi h} \right) - 1} (\epsilon_{\mathop{\sf ld}}+ \max_{j=1}^D \left\Vert {F(y(c_j)) - F^*(y(c_j))}\right\Vert)\left( {1 + \gamma_\phi} \right) h. \end{align}\]
Proof. We can write \(T_\phi(y) = y_{t_0} + S_\phi(F\circ y)\), where for arbitrary curve \(z\in \mathcal{C}([t_0,t_0 + h],\mathbb{R}^d)\) define \[S_\phi(z)(.) = \sum_{j=1}^D z(c_j) \int_{t_0}^{.} \phi_j(s)ds.\] Note that because \(P_{\leq D} \left( {F^* \circ y} \right)\) is degree at most \(D\) and \(\phi_j\)’s are the Lagrange multiplier polynomials at the Chebyshev points \((c_i)_{j=1}^D\), then \[\begin{align} P_{\leq D} \left( {F^* \circ y} \right) = \sum_{j=1}^D P_{\leq D} \left( {F^* \circ y} \right)(c_j) \phi_j, \end{align}\] which from the definition of \(S_{\phi}(z)\), implies \[\begin{align} \int_{t_0}^{.} P_{\leq D} \left( {F^* \circ y} \right)(s) ds = S_\phi \left( {P_{\leq D} \left( {F^*\circ y} \right)} \right). \end{align}\] Now combining this with Lemma 11 and using the low-degree Definition 4 \[\begin{align} &\left\Vert {S_\phi(F \circ y) -\int_{t_0}^{.} P_{\leq D} \left( {F^* \circ y} \right)(s) ds}\right\Vert_{[t_0, t_0 + h]} \nonumber\\ &=\left\Vert {S_\phi(F \circ y) - S_\phi \left( {P_{\leq D} \left( {F^*\circ y} \right)} \right)}\right\Vert_{[t_0, t_0+h]}\nonumber\\ &=\left\Vert {\sum_{j=1}^D \left( {F\circ y\Big|_{c_j} - P_{\leq D} \left( {F^*\circ y} \right)\Big|_{c_j}} \right) \left( {y_{t_0} + \int_{t_0}^{.} \phi_j(s)ds} \right)}\right\Vert_{[t_0, t_0 + h]}\nonumber\\ &\leq (\max_{j=1}^D \left\Vert {F^*(y(c_j)) - P_{\leq D}\left( {F^*(y(c_j))} \right)}\right\Vert +\left\Vert {F(y(c_j)) - F^*(y(c_j))}\right\Vert) \sum_{j=1}^D \left\Vert {y_{t_0} + \int_{t_0}^{.} \phi_j(s)ds}\right\Vert_{[t_0, t_0+h]}\nonumber\\ &\leq (\epsilon_{\mathop{\sf ld}}+ \max_{j=1}^D \left\Vert {F(y(c_j)) - F^*(y(c_j))}\right\Vert) \sum_{j=1}^D \left\Vert {y_{t_0} + \int_{t_0}^{.} \phi_j(s)ds}\right\Vert_{[t_0, t_0+h]}\nonumber\\ &\leq (\epsilon_{\mathop{\sf ld}}+ \max_{j=1}^D\left\Vert {F(y(c_j)) - F^*(y(c_j))}\right\Vert) \gamma_\phi h.\label{eq:avali} \end{align}\tag{20}\] But from the definition of \(\epsilon_{\sf ld}\), \[\left\Vert {P_{\leq D} \left( {F^* \circ y} \right) - F^*\circ y}\right\Vert_{[t_0, t_0 + h]} \leq \epsilon_{\sf ld},\] which using the identity \(y -y_{t_0}= \int_{t_0}^{.} F\circ y(s)ds\) implies \[\begin{align} \left\Vert {\int_{t_0}^{.} P_{\leq D} \left( {F^* \circ y} \right)(s)ds - (y-y_{t_0})}\right\Vert_{[t_0, t_0+h]} &= \left\Vert {\int_{t_0}^{.} P_{\leq D} \left( {F^* \circ y} \right)(s)ds - \int_{t_0}^{.} F^*\circ y (s) ds}\right\Vert_{[t_0, t_0 + h]} \nonumber\\ &\leq \epsilon_{\sf ld} h \label{eq:dovomi} \end{align}\tag{21}\] Combining 20 and 21 \[\begin{align} \left\Vert {T_\phi(y) - y}\right\Vert_{[t_0, t_0 + h]} &\leq \left\Vert {T_\phi(y) - y_{y_{t_0}} - \int_{t_0}^{.} P_{\leq D} \left( {F^* \circ y} \right) ds}\right\Vert_{[t_0, t_0 + h]} \\ &+ \left\Vert { \int_{t_0}^{.} P_{\leq D} \left( {F^* \circ y} \right)ds - (y - y_{t_0})}\right\Vert_{[t_0, t_0 + h]}\\ &\leq (\epsilon_{\mathop{\sf ld}}+ \max_{j=1}^D\left\Vert {F(y(c_j)) - F^*(y(c_j))}\right\Vert)\left( {1 + \gamma_\phi} \right) h. \end{align}\] Now applying Lemma 11 \(i-1\) times \[\begin{align} \left\Vert {T_\phi^{\circ i} (y) - T_\phi^{\circ (i-1)} (y)}\right\Vert &\leq \left( {\tilde{L}\gamma_\phi h} \right)^{i-1} \left\Vert {T_\phi (y) - y}\right\Vert_{[0,h]} \nonumber\\ &\leq \left( {\tilde{L}\gamma_\phi h} \right)^{i-1} (\epsilon_{\mathop{\sf ld}}+ \max_{j=1}^D\left\Vert {F(y(c_j)) - F^*(y(c_j))}\right\Vert)\left( {1 + \gamma_\phi} \right) h\label{eq:irelation} \end{align}\tag{22}\] Summing 22 for \(i=1,\dots,m\): \[\begin{align} \left\Vert {T_\phi^{\circ m} (y) - y}\right\Vert &\leq \sum_{i=1}^m \left\Vert {T_\phi^{\circ i} y - T_\phi^{\circ (i-1)} y}\right\Vert\\ &\leq \frac{\left( {\tilde{L}\gamma_\phi h} \right)^{m} - 1}{\left( {\tilde{L}\gamma_\phi h} \right) - 1} (\epsilon_{\mathop{\sf ld}}+ \max_{j=1}^D\left\Vert {F(y(c_j)) - F^*(y(c_j))}\right\Vert)\left( {1 + \gamma_\phi} \right) h\,. \qedhere \end{align}\] ◻
Corollary 2 (Effect of approximate Picard iterations). For the exact probability flow ODE \(\dot{y}_t= F^*_t(y_t)\), given that \(F^*\) is a low-degree vector field based on Definition 4 and given \(\tilde{L}\gamma_\phi h \leq \frac{1}{2}\), for arbitrary curve \(x \in \mathcal{C}([t_0,t_0 + h], \mathbb{R}^d)\) we have \[\begin{align} &\left\Vert {T_\phi^{\circ m} (y) - y}\right\Vert_{[t_0, t_0+h]} \leq 2(\epsilon_{\mathop{\sf ld}}+ \max_{j=1}^D\left\Vert {F(y(c_j)) - F^*(y(c_j))}\right\Vert)\left( {1 + \gamma_\phi} \right) h\\ &\|T_\phi^{\circ m}(x) - T_\phi^{\circ m}(y)\|_{[t_0,t_0 + h]} \leq \frac{1}{2^m} \left\Vert {x-y}\right\Vert_{[t_0,t_0 + h]}. \end{align}\]
Lemma 13. Recall the definition of \(F\) and \(F^*\) from Section 6. For initial point \(\left\Vert {\bar y_{t_0} - y_{t_0}}\right\Vert \leq \epsilon_p\), define \(y_s\), \(\tilde{y}_s\), and \(\bar y_s\) for \(t_0 \leq s \leq t_0 + h\) as \[\begin{align} &\dot{y}_s = F^*_s(y_s),\\ &\dot{\tilde{y}}_s = F^*_s(\tilde{y}_s),\\ &\bar y(s) = \bar y_{t_0} + (s-t_0)F_{t_0}( \bar y_{t_0}). \end{align}\] Then, under the low degree Assumption 4 for the curve \(\tilde{y}\) in the time interval \([t_0, t_0 + h]\), and picking step size \(h = O(\frac{1}{1 + R^2})\) and assuming \(\tilde{L}\leq \frac{1}{2}\) and \(\sigma_t \geq 1\), \[\begin{align} \left\Vert {y - T_\phi^{\circ m}(\bar y)}\right\Vert_{[t_0,t_0+h]} &\leq 2\epsilon_p(1 + 2(1+\gamma_\phi)h) + (\epsilon_{\mathop{\sf ld}}+ \max_{j=1}^D\left\Vert {F(y(c_j)) - F^*(y(c_j))}\right\Vert)\left( {1 + \gamma_\phi} \right) h \\ &+ \frac{h}{2^m} \left( {\tilde{L}\epsilon_p + \epsilon_{\sf err} + \sqrt d R (R^2 + \ln(1/\epsilon_1)) + (2 + 4R^2) \epsilon_p } \right). \end{align}\]
Proof. First, since we have the assumption \(T-t_0 - h\geq 1\), we have \(L_s \leq (2 + 4R^2), \forall s\in [t_0, t_0 + h]\). Hence, since we know \(\gamma_\phi = O(1)\) for the Chebyshev basis, the assumption \(h \leq O(\frac{1}{1+R^2})\) satisfies the precondition of Corollary [cor:contractionprops] on \(h\). On the other hand, from Lemma 15, again from the condition \(h = O(\frac{1}{1+R^2})\), we get \(\left\Vert {y - \tilde{y}}\right\Vert_{[t_0, t_0+h]} \leq 2\epsilon_p\). Now combining this with Corollary [cor:contractionprops]: \[\begin{align} \left\Vert {y - T_\phi^{\circ m}(\bar y)}\right\Vert_{[t_0, t_0 + h]} &\leq \left\Vert {y - \tilde{y}}\right\Vert_{[t_0, t_0 + h]} + \left\Vert {\tilde{y} - T_\phi^{\circ m}(\tilde{y})}\right\Vert_{[t_0, t_0 + h]} + \left\Vert {T_\phi^{\circ m} (\bar y) - T_\phi^{\circ m}(\tilde{y})}\right\Vert_{[t_0, t_0 + h]}\\ &\leq 2\epsilon_p + 2(\epsilon_{\mathop{\sf ld}}+ \max_{j=1}^D\left\Vert {F(\tilde{y}(c_j)) - F^*(\tilde{y}(c_j))}\right\Vert)\left( {1 + \gamma_\phi} \right) h + \frac{1}{2^m} \left\Vert {\bar y-\tilde{y}}\right\Vert_{[t_0,t_0 + h]}\\ &\leq 2\epsilon_p + 2(\epsilon_{\mathop{\sf ld}}+ 2\left\Vert {y - \tilde{y}}\right\Vert_{[t_0, t_0 + h]} + \max_{j=1}^D\left\Vert {F( y(c_j)) - F^*( y(c_j))}\right\Vert)\left( {1 + \gamma_\phi} \right) h \\ &+ \frac{1}{2^m} \left\Vert {\bar y-\tilde{y}}\right\Vert_{[t_0,t_0 + h]}\\ &\leq 2\epsilon_p(1 + 2(1+\gamma_\phi)h) + 2(\epsilon_{\mathop{\sf ld}} + \max_{j=1}^D\left\Vert {F( y(c_j)) - F^*( y(c_j))}\right\Vert)\left( {1 + \gamma_\phi} \right) h \\ & + \frac{1}{2^m} \left\Vert {\bar y-\tilde{y}}\right\Vert_{[t_0,t_0 + h]} \addtocounter{equation}{1}\label{eq:aavali} \end{align}\tag{23}\] But from Lemma 16, with probability at least \(1-\epsilon_1\), for all \(t_0 \leq s\leq t_0 + h\) \[\begin{align} \left\Vert {F^*_s(\tilde{y}_s)}\right\Vert \leq 6\left( {R + \sqrt d + \sqrt{\ln(1/\epsilon_1)}} \right). ~\label{eq:scoreupperbound2} \end{align}\tag{24}\]
Hence, using Lemma 14 and 15, \[\begin{align} \frac{d}{ds}\left\Vert {\bar y_s - \tilde{y}_s}\right\Vert &\leq \frac{\langle F_{t_0}(\tilde{y}_{t_0}) - F^*_s(\tilde{y}_s), \bar y_t- y_t\rangle}{\left\Vert {\bar y_t- y_t}\right\Vert}\\ &\leq \left\Vert {F_{t_0}( \tilde{y}_{t_0}) - F^*_t(\tilde{y}_t)}\right\Vert\\ &\leq \left\Vert {F_{t_0}(\tilde{y}_{t_0}) - F_{t_0}(y_{t_0})}\right\Vert +\left\Vert {F_{t_0}(y_{t_0}) - F^*_{t_0}(y_{t_0})}\right\Vert \\& +\left\Vert {F^*_{t_0}( y_{t_0}) - F^*_{t}( y_{t})}\right\Vert + \left\Vert {F^*_{t}(y_{t}) - F^*_{t}(\tilde{y}_{t})}\right\Vert \\ &\leq \tilde{L}\epsilon_p + \epsilon_{\sf err} + \sqrt d R (R^2 + \ln(1/\epsilon_1)) + (2 + 4R^2) e^{\frac{5R^2}{\sigma_t^2} h}\epsilon_p . \end{align}\] which implies \[\begin{align} \left\Vert {\bar y - \tilde{y}}\right\Vert_{[t_0, t_0+h]} \leq h\left( {\left\Vert {F_{t_0}(\tilde{y}_{t_0}) - F^*_{t_0}(\tilde{y}_{t_0})}\right\Vert + 12\left( {R + \sqrt d + \sqrt{\ln(1/\epsilon_1)}} \right)} \right). \end{align}\] Plugging this back into 23 \[\begin{align} \left\Vert { y - T_\phi^{\circ m}(\bar y)}\right\Vert &\leq 2\epsilon_p(1 + 2(1+\gamma_\phi)h) + (\epsilon_{\mathop{\sf ld}}+ \max_{j=1}^D\left\Vert {F(y(c_j)) - F^*(y(c_j))}\right\Vert)\left( {1 + \gamma_\phi} \right) h \\ &+ \frac{h}{2^m} \left( {\tilde{L}\epsilon_p + \epsilon_{\sf err} + \sqrt d R (R^2 + \ln(1/\epsilon_1)) + (2 + 4R^2) e^{\frac{5R^2}{\sigma_t^2} h}\epsilon_p } \right). \end{align}\] Using the assumption \(h = O(\frac{1}{1 + R^2})\) completes the proof. ◻
In this section, we put everything together to show our main end-to-end result, namely how to sample from the target distribution which is \(\tilde{O}(\epsilon_{\sf err} R^2)\)-close in TV to a distribution which is \(\tilde{O}(\epsilon_{\sf err})\)-close in \(W_2\) to the true distribution \(q\).
Theorem 3 (Formal version of Theorem 2). For error parameters \(\epsilon_1, \epsilon_{\sf err} > 0\) which satisfy the bounds \(\epsilon_{\sf err} \leq \epsilon_1/\ln(d/R)\) and \(\epsilon_{\sf err} \leq \left( {(R/\sigma) \wedge (1/\sqrt d)} \right)/\left( {\ln(\sigma/R) + \ln(d)} \right)^2\), given step size \(h = \frac{\gamma}{k(1+(R/\sigma)^2)\left( {\ln(1/\epsilon_1) + \ln(d)} \right)}\) for small enough constant \(\gamma\), where \(k \triangleq \ln(1/\epsilon_{\sf err})\) is the degree of the polynomial approximation, with probability at least \[\begin{align} 1 - O\Big(\epsilon_1 T(R/\sigma)^2\ln(1/\epsilon_{\sf err})(\ln(1/\epsilon) + \ln(d))\Big), \end{align}\] for \(T = \ln(R/\sigma) + \ln(d) + \ln(1/\epsilon_{\sf err})\), Algorithm [alg:main] outputs a sample whose distribution is close in 2-Wasserstein distance to the target measure \(q\) to within error \[\begin{align} O(Tk\epsilon_{\sf err}\log(1/\epsilon_1)) = O\Big(\left( {\ln\left( {R/\sigma} \right) + \ln\left( {d} \right) + \ln(1/\epsilon_{\sf err})} \right) \ln(1/\epsilon_{\sf err}) \ln(k/\epsilon_1)\epsilon_{\sf err}\Big), \end{align}\] in \[O\Big(T(R/\sigma)^2\ln(1/\epsilon_{\sf err})\left( {\ln(1/\epsilon_1) + \ln(d)} \right)\Big)\] number of rounds and \(m = \ln\left( {R/\sigma} \right) + \ln(d) + \ln\ln(1/\epsilon_1) + \ln(1/\epsilon_{\sf err}) + \ln(\tilde{L})\) number of Picard iterations in each round.
Proof. We can scale the distribution, sample from the scaled distribution, then scale back. Note that this scaling procedure does not change the ratio \(R/\sigma\). Therefore, without loss of generality, we assume \(\sigma_0 = \sigma = 1\), since our bounds only depend on the quantity \(R/\sigma\). We run the forward process up to time \[\begin{align} T \mathrel{\vcenter{:}}= \Theta(\ln(R) + \ln(d) + \ln(1/\epsilon_{\sf err})).\label{eq:tdefinition} \end{align}\tag{25}\]
Now we prove inductively that \(\left\Vert {y - T_\phi^{\circ m}(\bar y)}\right\Vert_{[t^{(i)},t^{(i)}+h]} = O((t^{(i)} + h)k\epsilon_{\sf err}\log(k/\epsilon_1))\) for all \(i\). First, note that we can bound the Wasserstein distance of the target distribution \(q_{\sf pre}\) and \(\mathcal{\gamma}^d\) as \[W_2(q_{\sf pre}, \gamma^d) \leq W_2(q_{\sf pre}, \delta_{\{0\}}) + W_2(\delta_{\{0\}}, \gamma^d),\] where \(\delta_{\{0\}}\) is the point mass at the origin. But \[\begin{align} &W_2(q_{\sf pre}, \delta_{\{0\}}) = \mathbb{E}_{q_{\rm base}} \left\Vert {Y}\right\Vert^2 \leq R,\\ &W_2(\gamma^d, \delta_{\{0\}}) = \sqrt d. \end{align}\] Therefore \[W_2(q_{\sf pre}, \gamma^d) \leq R + \sqrt d.\] Now from the Wasserstein contraction property of the OU process and the choice of \(T\) in 25 , we get \[W_2(q_T, \gamma^d) \leq \left( {R + \sqrt d} \right)e^{-T} = O(\epsilon_{\sf err})\,.\] To prove the step of induction for the interval \([t^{(i)}, t^{(i)} + h] = [t_0, t_0 + h]\), we know from the hypothesis of induction that for all previous intervals \((t^{(j)}, t^{(j)} + h)\) for \(j < i\): \[\left\Vert {y - T_\phi^{\circ m}(\bar y)}\right\Vert_{[t^{(j)},t^{(j)}+h]} \leq O((t^{(j)}+h)k\epsilon_{\sf err} \log(k/\epsilon_1)),\] which from the definition of the iterates of the algorithm, i.e. \(\hat{y} = T_\phi^{\circ m}(\bar y)\) implies \[\left\Vert {y - \hat{y}}\right\Vert_{[0,t_0]} \leq O(t_0 k\epsilon_{\sf err} \log(k/\epsilon_1)).\] Using Lemma 15, this further implies, given the choice of \(h\), \[\left\Vert {\tilde{y} - y}\right\Vert_{[t_0, t_0 + h]} = O(t_0k\epsilon_{\sf err} \log(k/\epsilon_1)).\] Hence, we can use \(\epsilon_p = O(t_0 k\epsilon_{\sf err} \log(k/\epsilon_1))\) in Lemma 13.
Now given the assumptions \(\epsilon_{\sf err} \leq (1/\ln(d/R)) \epsilon_1\) and \(\epsilon_{\sf err} \leq \left( {R \wedge (1/\sqrt d)} \right)/\left( {\ln(1/R) + \ln(d)} \right)^2\), and since \(\sigma_t \geq \sigma_0 \geq 1\) it is easy to check that \(\tilde{\delta} = O(t_0k\epsilon_{\sf err}\log(k/\epsilon_1))\) satisfies the assumption of Lemma 10.
Hence, we can apply Lemma 10 with \[k = \ln(1/\epsilon_{\sf err}),\] using the fact that \(\sigma_t \geq \frac{1}{2}\) and \(R\geq 1\) there is constant \(c_2\) such that for the interval \([t_0, t_0 + h] \mathrel{\vcenter{:}}= [t^{(i)}, t^{(i)} + h]\), \[\begin{align} \mathbb{P}\left( {\sup_{s\in [t_0, t_0 + h]} \Big|\partial_s^k g^{(\tilde{y}(s))}_i(s)\Big|_{s=0}\Big| \geq \left( {c_2k^2R^2 \left( {\ln(1/\epsilon_1) + \ln(d)} \right)} \right)^k} \right) \leq 2\epsilon_1. \end{align}\] Moreover, using this with Fact [fact:tailorremainder] and by picking step size \[\begin{align} h = \frac{\gamma}{k(1+R^2)\left( {\ln(1/\epsilon_1) + \ln(d)} \right)} \end{align}\] for small enough constant \(\gamma\), we get with probability at least \(1-2\epsilon_1\) \[\begin{align} \left\Vert {g - P_{\leq k} g}\right\Vert_{[t_0, t_0+h]} \leq \left( {h c_2 e kR^2 \left( {\ln(1/\epsilon_1) + \ln(d)} \right)} \right)^k \leq \epsilon_{\sf err}.\label{eq:errld} \end{align}\tag{26}\] Therefore, we can now use Lemma 13 with \[\begin{align} m \mathrel{\vcenter{:}}= \ln\left( {R} \right) + \ln\left( {d} \right) + \ln\ln(1/\epsilon_1) + \ln(1/\epsilon_{\sf err}) + \ln(\tilde{L}) \end{align}\] number of Picard iterations; using the fact that \(\gamma_\phi = O(1)\) for the Chebyshev basis \((\phi_0,\ldots,\phi_k)\) for the polynomials defined in Section 2.3, setting \(\epsilon_{\sf ld} = \epsilon_{\sf err}\) in Assumption 4 based on Equation 26 , and using Assumption 2, we get \[\begin{align} \left\Vert {y - T_\phi^{\circ m}(\bar y)}\right\Vert_{[t_0,t_0+h]} &= O\left( {\epsilon_p + h\left( {\epsilon_{\sf err} + \sum_{j=1}^k \left\Vert {F(y(c_j)) - F^*(y(c_j))}\right\Vert} \right)} \right) \nonumber\\ &+ \frac{h}{2^m}O\left( {\tilde{L}\epsilon_p + \epsilon_{\sf err} + \sqrt d R (R^2 + \ln(1/\epsilon_1)) + (2 + 4R^2) \epsilon_p } \right)\nonumber\\ &\leq O(\epsilon_p + h k\epsilon_{\sf err} \log(k/\epsilon_1)) \nonumber\\ &+ \frac{h}{2^m}O\left( {\tilde{L}\epsilon_p + \epsilon_{\sf err} + \sqrt d R (R^2 + \ln(1/\epsilon_1)) + (2 + 4R^2) \epsilon_p } \right)., \label{eq:middleeq} \end{align}\tag{27}\] where the last line follows from the choice of \(m\).
Now using the fact that \(h = O(\frac{1}{1+R^2})\), we have \[\begin{align} \frac{h}{2^m} (2+4R^2)\epsilon_p = O(\frac{1}{2^m} Tk\epsilon_{\sf err} \log(1/\epsilon_1)) = O(\frac{T/h}{2^m}hk\epsilon_{\sf err} \log(1/\epsilon_1)) = O(hk\epsilon_{\sf err} \log(k/\epsilon_1)), \end{align}\] where we used the fact that \[\begin{align} T/h = O\left( {(1+R^2)\ln(1/\epsilon_{\sf err})^2\left( {\ln(1/\epsilon_1) + \ln(d)} \right)} \right) = O(2^m). \end{align}\] Similarly, it is not hard to check that with the choice of \(m\) we have \[\begin{align} \frac{h}{2^m}\left( {\tilde{L}\epsilon_p + \epsilon_{\sf err} + \sqrt d R (R^2 + \ln(1/\epsilon_1))} \right) = O(hk\epsilon_{\sf err} \log(k/\epsilon_1)) \end{align}\]
Therefore, we can upper bound 27 as \[\begin{align} \mathbb{E} \left\Vert {y - T_\phi^{\circ m}(\bar y)}\right\Vert_{[t_0,t_0+h]} \leq O(\epsilon_p) + O(hk \epsilon_{\sf err}\log(k/\epsilon_1)) = O((t_0+h) k\epsilon_{\sf err}\log(k/\epsilon_1)). \end{align}\] which proves the step of induction. Therefore, overall we showed \[\begin{align} \left\Vert {y - \hat{y}}\right\Vert_{[0,T]} = O(Tk\epsilon_{\sf err} \log(1/\epsilon_1)) = O\Big(\left( {\ln\left( {R} \right) + \ln\left( {d} \right) + \ln(1/\epsilon_{\sf err})} \right) \ln(1/\epsilon_{\sf err}) \log(k/\epsilon_1)\epsilon_{\sf err}\Big), \end{align}\] after \(O\Big(TR^2\ln(1/\epsilon_{\sf err})\left( {\ln(1/\epsilon_1) + \ln(d)} \right)\Big)\) number of rounds, where in each round we apply the Picard iteration for \(m = \ln\left( {R} \right) + \ln(d) + \ln\ln(1/\epsilon_1) + \ln(1/\epsilon_{\sf err}) + \ln(\tilde{L})\) number of times. Note that this Wasserstein guarantee holds only with probability at least \(1-O(\epsilon_1 T/h) = 1 - O\Big(\epsilon_1 TR^2\ln(1/\epsilon_{\sf err})(\ln(1/\epsilon) + \ln(d))\Big)\) after applying a union bound. ◻
In this section we show how to use underdamped Langevin Monte Carlo to upgrade the Wasserstein bound achieved by our sampler into a total variation bound, thus proving Corollary 1. The pseudocode for this is provided in Algorithm [alg:main2], where we propose a slight modification of Algorithm [alg:main] that relies on running the underdamped Langevin Monte Carlo algorithm (see Appendix 12.1) for a short period of time at the end of Algorithm [alg:main]. This step is often referred to as a corrector step in the diffusion model literature. The resulting Algorithm [alg:main2] admits a TV guarantee, stated in Corollary 1.
\(\eta \gets \Theta(\min(\epsilon^{5/3}/(L^{1/4}d^{1/2}),\epsilon\sigma^2/R^2))\), where \(L \le \mathrm{poly}(d/\eta)\) is an upper bound on the Lipschitz constant of \(q\). \(x\gets\)CollocationDiffusion(\((s_t),\varepsilon\)) Run underdamped Langevin Monte Carlo (see Appendix 12.1) with friction parameter \(\Theta(\sqrt{L})\) and step size \(h = \epsilon^{2/3}/(d^{1/3}M(\eta)^{1/3}L^{1/2})\) for \(M(\eta)\) steps, where \(M(\eta)\) denotes the number of steps used to run CollocationDiffusion in the previous step. Let the resulting sample be \(x'\).
In this section, for the sake of completeness, we briefly review underdamped Langevin Monte Carlo, which is only used in this final phase of our algorithm to convert from Wasserstein closeness to TV closeness.
Given an estimate \(s\) of the log-density of a distribution \(q\), and a friction parameter \(\gamma\), underdamped Langevin Monte Carlo with step size \(h\) and score estimate \(s\) is given by \[\begin{align} \mathrm{d}x_t &= v_t \, \mathrm{d}t\\ \mathrm{d}v_t &= (s(x_{\lfloor t/h\rfloor h}) - \gamma v_t)\,\mathrm{d}t + \sqrt{2\gamma}\,\mathrm{d}B_t\,, \end{align}\] where \(B_t\) is a standard Brownian motion.
We will use the following result of [43], which is a consequence of the short-time regularization of [67]:
Theorem 4 (Theorem A.5 of [43], restated). Let \(q \propto e^{-H}\) be a distribution over \(\mathbb{R}^d\) for which \(\nabla H\) is \(L\)-Lipschitz. Let \(p\) be an arbitrary distribution over \(\mathbb{R}^d\). Suppose \(s: \mathbb{R}^d\to\mathbb{R}^d\) satisfies \(\|s - \nabla H\|_{L_2(q)}^2 \le \varepsilon_{\rm sc}^2\). Let \(T\lesssim 1/\sqrt{L}\).
If \(p_N\) denotes the distribution given by running underdamped Langevin Monte Carlo initialized at \(p\) for \(T / h\) steps and step size \(h\) with friction parameter \(\Theta(\sqrt{L})\), then \[\mathrm{TV}(p_N, q) \lesssim \frac{W_2(p,q)}{L^{1/4}T^{3/2}} + \frac{\varepsilon_{\rm sc}T^{1/2}}{L^{1/4}} + L^{3/4}T^{1/2}d^{1/2}h\,. \label{eq:reg}\qquad{(5)}\]
A priori it might appear that the third term in Eq. ?? forces a choice of \(h = O(d^{-1/2}\epsilon)\), translating to an iteration complexity of \(\Omega(d^{1/2}/\epsilon)\). Here we use an idea of [34]: because in our application of Theorem 4, \(W_2(p,q)\) can be made quite small, we can take \(T\) to be small to get a much better bound on the iteration complexity \(T/h\).
Proof of Corollary 1. For \(\eta\) to be tuned, let \(M(\eta) = (R/\sigma)^2\cdot \operatorname{polylog}(1/\epsilon,d,R,\tilde{L})\) denote the number of iterations of Algorithm [alg:main] needed to achieve Wasserstein error \(\eta\) in Theorem 2. We will take \(h = \epsilon^{2/3} / (d^{1/3} M(\eta)^{1/3} L^{1/2})\) and \(T = M(\eta)h\) in Theorem 4 to conclude that, starting from the distribution given by Algorithm [alg:main], if we run underdamped Langevin Monte Carlo for \(M(\eta)\) iterations, we will produce a distribution \(p_N\) for which \[\mathrm{TV}(p_N, q) \lesssim \frac{\eta L^{1/4} d^{1/2}}{M(\eta)^{2/3}\epsilon^{2/3}} + \frac{\varepsilon_{\rm sc}M(\eta)^{1/3} \epsilon^{1/3}}{L^{1/2} d^{1/6}} + \epsilon\,, \label{eq:shorttime95tv}\tag{28}\] conditioned on the event that Algorithm [alg:main] succeeds in achieving Wasserstein error \(\eta\). The latter happens with probability \(1 - \eta(R/\sigma)^2\), so the overall TV error of Algorithm [alg:main2] is given by Eq. 28 plus \(\eta(R/\sigma)^2\).
If we take \(\eta = \min(\epsilon^{5/3}/(L^{1/4}d^{1/2}), \epsilon\sigma^2/R^2\), then \(\eta(R/\sigma)^2 \le \epsilon\), \(M(\eta) \le \tilde{O}(R/\sigma)^2 \cdot \operatorname{polylog}(1/\epsilon,d,\tilde{L},L)\), and the first term on the right-hand side is bounded by \(\epsilon\). By the assumption that \(\varepsilon_{\rm sc}\le \tilde{O}(\frac{\epsilon^{2/3} L^{1/2}d^{1/6}}{(R/\sigma)^{2/3}})\), the second term on the right-hand side is bounded by \(\epsilon\). By replacing \(\epsilon\) with \(c\epsilon\) for sufficiently small constant \(c\) in the above, we obtain the claimed bound. ◻
In this section we prove some useful properties of the score function. First, we bound the smoothness of the score function \(\nabla \log q_t\).
Lemma 14. (Bounding the operator norm of the true score) For \(t \geq 1\), we have \[\begin{align} &\left\Vert {\nabla \left( {y_{t} + \nabla \log q_t(y_{t})} \right)}\right\Vert \leq e^{-(T-t)}\left( {2 + 4R^2} \right),\\ &\left\Vert {\nabla^2 \log q_t(y_{t})}\right\Vert_{\sf op} \leq 2 + 4e^{-(T-t)}R^2, \end{align}\] and for \(2(T - t) < 1\), \[\begin{align} \left\Vert {\nabla \left( {y_{t} + \nabla \log q_t(y_{t})} \right)}\right\Vert \leq \frac{5R^2}{(T-t)^2}. \end{align}\]
Proof. It is easy to check the Hessian of \(\log q_t\) can be written in the following form: \[\begin{align} -\nabla^2 \log q_t(y) = -\frac{1}{\sigma_t^4} \operatorname{Cov}(q^{t,y}) + \frac{1}{\sigma_t^2} I. \end{align}\] But using the radius \(R\) assumption on the support of \(q_{\sf pre}\), we can bound the covariance as \[\begin{align} \operatorname{Cov}(q^{t,y}) \leq \frac{1-\sigma_t^2}{\sigma_t^4} R^2 I. \end{align}\] Therefore \[\begin{align} -\nabla \left( {y_{t} + \nabla \log q_t(y_{t})} \right) = \frac{1-\sigma_t^2}{\sigma_t^2} I - \frac{1-\sigma_t^2}{\sigma_t^4} R^2 I.\label{eq:coretmp} \end{align}\tag{29}\] Now for \(t \geq 1\), we have \(\sigma_t^2 = (1-e^{-2(T-t)}) \geq 1-1/e^2 \geq 0.5\). Therefore, for \(t\geq 1\), \[\begin{align} \left\Vert {\nabla \left( {y_{t} + \nabla \log q_t(y_{t})} \right)}\right\Vert \leq (1-\sigma_t^2)\left( {2 + 4R^2} \right) = e^{-2(T-t)}\left( {2 + 4R^2} \right). \end{align}\] Similarly \[\begin{align} \left\Vert {\nabla^2 \log q_t(y_{t})}\right\Vert \leq 2 + (1-\sigma_t^2)\left( {4R^2} \right) = 2 + 4e^{-2(T-t)}R^2. \end{align}\] On the other hand, for \(s = 2(T - t) < 1\), using \(e^{-s} \leq 1-s+\frac{s^2}{2} \leq 1-\frac{s}{2}\) we have \(\sigma_t^2 = 1-e^{-(T-t)} \geq T - t\). Then from the assumption \(R \geq 1\) and Equation 29 , we have \[\begin{align} \left\Vert {\nabla \left( {y_{t} + \nabla \log q_t(y_{t})} \right)}\right\Vert &\leq \left( {\frac{2}{T - t}+ \frac{4R^2}{(T-t)^2}} \right) \leq \frac{5R^2}{(T-t)^2}\,. \qedhere \end{align}\] ◻
We can then use this bound on the smoothness to control the extent to which two processes evolving according to the same probability flow ODE diverge over time:
Lemma 15. (Distance between true ODE solutions starting from close points) For \(y_{t}, \tilde{y}_{t}\) that evolve according to probability flow ODE, i.e., \[\begin{align} &\frac{d}{dt}y_{t} = y_{t} + \nabla \log q_t(y_{t})\\ &\frac{d}{dt}\tilde{y}_{t} = \tilde{y}_{t} + \nabla \log q_t(\tilde{y}_{t}), \end{align}\] with initial condition satisfying \(\left\Vert {y_{t_0} - \tilde{y}_{t_0}}\right\Vert \leq \epsilon\), then, for time window \(s\) such that \(T-(t_0+s) \geq 1\), we have \[\begin{align} \left\Vert {y_{t_0+s} - \tilde{y}_{t_0+s}}\right\Vert \leq e^{\frac{5R^2}{\sigma_t^2} s}\epsilon \end{align}\]
Proof. Note that for \(t_0 \leq t \leq t_0 + s\), we have \(t = T-t \geq 1\) from our assumption, therefore from Lemma 14: \[\begin{align} \frac{d}{dt}\left\Vert {y_{t} - \tilde{y}_{t}}\right\Vert &= \frac{\langle y_{t} - \tilde{y}_{t}, \nabla \log q_t(y_{t}) - \nabla \log q_t(\tilde{y}_{t})\rangle}{\left\Vert {y_{t} - \tilde{y}_{t}}\right\Vert}\\ &= \frac{\int_{r=0}^1 \left( {y_{t} - \tilde{y}_{t}} \right)^\top\nabla \left( {y_{t} + q_t(y_{t} + r(\tilde{y}_{t} - y_{t}))} \right)\left( {y_{t} - \tilde{y}_{t}} \right)dr}{\left\Vert {y_{t} - \tilde{y}_{t}}\right\Vert}\\ &\leq e^{-t}(2+4R^2)\frac{\left\Vert {y_{t} - \tilde{y}_{t}}\right\Vert^2}{\left\Vert {y_{t} - \tilde{y}_{t}}\right\Vert}\\ &= e^{-t}(2+4R^2)\left\Vert {y_{t} - \tilde{y}_{t}}\right\Vert\,, \end{align}\] which implies \[\begin{align} \frac{d}{dt}\ln\left( {\left\Vert {y_{t} - \tilde{y}_{t}}\right\Vert} \right) \leq e^{-t}(2+4R^2).\label{eq:timeder95new} \end{align}\tag{30}\] Integrating Equation 30 from \(t = t_0\) to \(t = t_0 + s\) we get the desired result for the first part. Similarly for the case when \(T - t_0 < 1\): \[\begin{align} \frac{d}{dt}\ln\left( {\left\Vert {y_{t} - \tilde{y}_{t}}\right\Vert} \right) \leq \frac{5R^2}{t^2}, \end{align}\] which implies (using inequality \(\sigma_t^2 = 1-e^{-t}\leq t\) for \(t = T - (t_0+s)\)) \[\begin{align} \left\Vert {y_{t_0+s} - \tilde{y}_{t_0+s}}\right\Vert &\leq e^{\frac{5R^2}{T-(t_0+s)} s}\epsilon \leq e^{\frac{5R^2}{\sigma_t^2} s}\epsilon\,. \qedhere \end{align}\] ◻
Lemma 16. (Distributional guarantees for probability flow ODE) Along the backward ODE \[\begin{align} \dot{y}_t= F^*_t(y_t)\label{eq:odedeff} \end{align}\qquad{(6)}\] with \(F^*_t(y) = y + \nabla \log q_t(y)\) we have with probability \(1-\epsilon_1\), \[\begin{align} &\forall t\in [t_0, t_0+h], \left\Vert {y_t}\right\Vert \leq e\left( {R + \sqrt d + \sqrt{\ln(1/\epsilon_1)}} \right),\\ &\forall t\in [t_0, t_0+h], \left\Vert {F^*_t(y_t)}\right\Vert \leq 6\left( {R + \sqrt d + \sqrt{\ln(1/\epsilon_1)}} \right),\\ &\forall s_1, s_2, \left\Vert {y_{s_1} - y_{s_2}}\right\Vert\leq 6(s_2 - s_1)\left( {R + \sqrt d + \sqrt{\ln(1/\epsilon_1)}} \right)\,\\ &\left\Vert { F^*_{t_0}(y_{t_0}) - F^*_{t_0 + h}(y_{t_0 + h})}\right\Vert \lesssim \sqrt d R (R^2 + \ln(1/\epsilon)). \end{align}\]
Proof. Recall we can write \(y_t = X_t + \sigma_t \xi\) where \(F^*_t(y) = y_t + \frac{\mathbb{E}^{t,y} X - y_t}{\sigma_t^2}\). Now for time \(t_0\), using \(\sigma_{t_0}^2 \geq 0.5\) for \(t_0 \geq 1\) as we showed in Lemma 14, we have \[\begin{align} \left\Vert {F^*_{t_0}(y_{t_0})}\right\Vert &= \mathbb{E}\left\Vert {y_{t_0} + \frac{\mathbb{E}^{t,y_{t_0}} X - y_{t_0}}{\sigma_{t_0}^2}}\right\Vert\\ &\leq \left( {1+\frac{1}{\sigma_{t_0}^2}} \right)\left\Vert {y_{t_0}}\right\Vert + \frac{1}{\sigma_{t_0}^2}R\\ &\leq \frac{3}{2} \left\Vert {y_{t_0}}\right\Vert + \frac{R}{2}\label{eq:vecfieldbound} \end{align}\tag{31}\] Now because \(\left\Vert {\xi}\right\Vert\) is subgaussian, we get with probability at least \(1-\epsilon_1\), \[\begin{align} \left\Vert {y_{t_0}}\right\Vert \leq R + \sqrt d + \sqrt{\ln(1/\epsilon_1)}. \end{align}\] On the other hand, from the definition ?? \[\begin{align} \frac{d}{ds}\left\Vert {y_s}\right\Vert \leq \left\Vert {F^*_s(y_s)}\right\Vert \leq \frac{3}{2}\left\Vert {y_s}\right\Vert + \frac{R}{2}, \end{align}\] so up to time \(s\), we get \[\begin{align} \left\Vert {y_{s}}\right\Vert \leq (R + \sqrt d + \sqrt{\ln(1/\epsilon_1)}) e^{s\left( {R + \sqrt d + \sqrt{\ln(1/\epsilon_1)}} \right)}, \end{align}\] which implies from the fact that \(h \leq \frac{1}{R + \sqrt d + \sqrt{\ln(1/\epsilon_1)}}\), \[\begin{align} \left\Vert {y_s}\right\Vert \leq e\left( {R + \sqrt d + \sqrt{\ln(1/\epsilon_1)}} \right). \end{align}\] The second part follows from 31 . For the third part we have \[\begin{align} \left\Vert {y_{s_1} - y_{s_2}}\right\Vert &= \left\Vert {\int_{s_1}^{s_2} F^*_s(y_s) ds}\right\Vert\\ &\leq \int_{s_1}^{s_2} \left\Vert {F^*_s(y_s)}\right\Vert ds\\ &\leq 6(s_2 - s_1)\left( {R + \sqrt d + \sqrt{\ln(1/\epsilon_1)}} \right). \end{align}\]
For the last part, using Lemma [lem:bnd1]: \[\begin{align} \left( {\mathbb{E}\left\Vert {\partial_t F^*_t(y_t)}\right\Vert^p} \right)^{1/p} \lesssim \sqrt d R (R + \sqrt p)^2. \end{align}\] Integrating from \(t_0\) to \(t_0 + h\): \[\begin{align} \left( {\mathbb{E}\left\Vert { F^*_{t_0}(y_{t_0}) - F^*_{t_0 + h}(y_{t_0 + h})}\right\Vert^p} \right)^{1/p} \lesssim h \sqrt d R (R + \sqrt p)^2, \end{align}\] which implies with probability at least \(\epsilon_1\) we have \[\begin{align} \left\Vert { F^*_{t_0}(y_{t_0}) - F^*_{t_0 + h}(y_{t_0 + h})}\right\Vert \lesssim \sqrt d R (R + \sqrt{\ln(1/\epsilon)})^2. \end{align}\] This completes the proof. ◻
Note that the convention in some prior works has been to parametrize time such that \(q = q_0\), but we eschew this as our choice makes the notation in our proofs cleaner.↩︎