\section{Predictor step}\label{sec:predictor}


Next, we need an ODE discretization analysis.

\begin{lem}\label{l:predictor-1step}
Suppose the score function satisfies Assumption~\ref{ass:lip_score}.
Assume that $L\ge 1$, $h\lesssim 1/L$, and $T-(t_0+h)\ge \fc{T-t_0}2$. Then 
\[
W_2(q\Podeth{t_0}{h}, q\Podesth{t_0}{h}) \lesssim Ld^{1/2} h^2\, \Bigl(L^{1/2} \vee \frac{1}{{(T-t_0)}^{1/2}}\Bigr) + h\esc\,.
\]
\end{lem}
\begin{proof}
    We have the ODEs
    \begin{align*}
        \dot x_t
        &= x_t + \nabla \ln q_t(x_t)\,, \\
        \dot \Xs_t
        &= \Xs_t + s_{t_0}(\Xs_{t_0})\,,
    \end{align*}
    for $t_0 \le t \le t_0+h$,
    with $x_{t_0} = \Xs_{t_0}\sim q$, $x_{t_0+h}\sim q\Pode$, and $\Xs_{t_0+h}\sim q\Podes$. Then,
    \begin{align*}
        \partial_t \norm{x_t - \Xs_t}^2 
        &= 2\,\langle x_t - \Xs_t, \dot x_t - \dot \Xs_t\rangle \\
        &= 2\pa{\norm{x_t - \Xs_t}^2 + \an{x_t - \Xs_t,\nb \ln q_t(x_t) + s_{t_0}(\Xs_{t_0})}}\\
        &\le \bigl(2+\rc h\bigr)\,\norm{x_t - \Xs_t}^2 + h\,\norm{\nb \ln q_t(x_t) - s_{t_0}(\Xs_{t_0})}^2\,.
\end{align*}
By Gr\"onwall's inequality, noting that $h=O(1)$,
\begin{align}
    \nonumber
    \E[\norm{x_{t_0+h} - \Xs_{t_0+h}}^2]
    &\le \exp\bigl(\bigl(2+\frac{1}{h}\bigr)\,h\bigr) \int_{t_0}^{t_0+h} h\E[\norm{\nb\ln q_t(x_t) - s_{t_0}(\Xs_{t_0})}^2] \,\D t\\
    &\lesssim h \int_{t_0}^{t_0+h} \E[\norm{\nb\ln q_t(x_t) - s_{t_0}(\Xs_{t_0})}^2] \,\D t\,.
    \label{e:pred-gronwall}
\end{align}
We split up the error term as
\begin{align*}
\norm{\nb\ln q_t(x_t) - s_{t_0}(\Xs_{t_0})}^2 
&\lesssim \norm{\nb \ln q_t(x_t) - \nb \ln q_{t_0}(x_{t_0})}^2 + \norm{\nb \ln q_{t_0}(x_{t_0}) - s_{t_0}(\Xs_{t_0})}^2\,.
\end{align*}
By Corollary~\ref{c:sp}, the expectation of the first term is bounded by
\begin{align*}
\E[\norm{\nb\ln q_t(x_t) - \nb \ln q_{t_0}(x_{t_0})}^2]
&\lesssim L^2dh^2 \,\bigl(L \vee \frac{1}{T-t_0}\bigr)\,.
\end{align*}
The second term is bounded in expectation by $\esc^2$. Plugging back into~\eqref{e:pred-gronwall} gives
\begin{align*}
\E[\norm{x_{t_0+h} - \Xs_{t_0+h}}^2]
&\lesssim h^2 \, \Bigl(L^2dh^2 \,\bigl(L \vee \frac{1}{T -t_0}\bigr) + \esc^2\Bigr)\,.
\end{align*}
The Wasserstein distance is bounded by the square root of this quantity, and the lemma follows.
\end{proof}

Lemma~\ref{l:predictor-1step} suggests that focusing on the dependence on $d$, we will be able to take $h \asymp d^{-1/2}$ (we need to keep one factor of $h$ in the bound, as we need to sum up the bound over $1/h$ iterations). 

\begin{rmk}\label{r:sp}
Our improved score perturbation lemma is necessary to obtain this $d^{1/2}$ dependence. The original score perturbation lemma \cite[Lemma C.11--12]{leelutan22sgmpoly} combined with a space discretization bound gives a bound of 
\[
\E\bigl[\norm{\nb \ln q_t^{\rightarrow}(x_t) - \nb \ln q_s^{\rightarrow}(x_s)}^2\bigr] \lesssim 
L^2 dh
\]
in place of Corollary~\ref{c:sp}. Note this is a $\frac12$-H\"older continuity bound rather than a Lipschitz bound.
The bound in Lemma~\ref{l:predictor-1step} then becomes 
\[
W_2(q\Podeth{t_0}{h}, q\Podesth{t_0}{h}) \lesssim Ld^{1/2} h^{3/2} + h\esc\,,
\]
and we would only be able to take $h\asymp d^{-1}$. We also note that our bound has an extra factor of $\max\{L^{1/2}, (T-t_0)^{-1/2}\}$; we do not know if this extra factor is necessary.
\end{rmk}

We now iterate Lemma~\ref{l:predictor-1step} to obtain the following result. Note that we now need to also assume that the score estimate is $L$-Lipschitz.

\begin{lem}\label{lem:pred}
Suppose that both Assumptions~\ref{ass:lip_score} and~\ref{ass:lip_estimate} hold.
Let $h_1,\ldots, h_N > 0$ be a sequence such that
letting $t_N = h_1+\cdots + h_N$, we have $t_N \le 1/L$. 
Let $h_{\max} = \max_{1\le n\le N} h_n$. 
\begin{enumerate}
\item
If $T-(t_0 + t_N) \gtrsim 1/L$, then 
\begin{align*}
W_2(q\Podeth{t_0}{h_1,\ldots, h_N}, 
q\Podesth{t_0}{h_1,\ldots, h_N}) \lesssim 
 L^{3/2} d^{1/2} h_{\max} t_N + \esc t_N
 \le L^{1/2} d^{1/2}h_{\max} +\fc{\esc}{L}\,.
\end{align*}
\item
If $T-t_0\lesssim 1/L$ and $h_{n+1}\le \fc{T-t_0-t_n}2$ for each $n$, then 
\begin{align*}
W_2(q\Podeth{t_0}{h_1,\ldots, h_N}, 
q\Podesth{t_0}{h_1,\ldots, h_N}) \lesssim 
 L^{1/2} d^{1/2} h_{\max} + \esc t_N
 \le  L^{1/2} d^{1/2} h_{\max} + \fc{\esc}{L}\,.
\end{align*}
\end{enumerate}
\end{lem}
\begin{proof}
We abbreviate $\Pode^N  \deq  \Podeth{t_0}{h_1,\ldots, h_N}$ and $\Podes^N  \deq  \Podesth{t_0}{h_1,\ldots, h_N}$. Using the triangle inequality, 
\begin{align*}
        W_2(q\Pode^N, q\Podes^N)
        &\le W_2(q\Pode^N, q\Pode^{N-1} \Podes) + W_2(q\Pode^{N-1}\Podes, q\Podes^N) \\
        &\le O(Ld^{1/2} h_N^{2}\max\{L^{1/2}, {(T-t_0 - t_N)}^{-1/2}\} + h_N \esc) \\
        &\qquad{} +\exp(O(
        L h_N)) \, W_2(q\Pode^{N-1}, q\Podes^{N-1})
    \end{align*}
    where the bound on the first term is by Lemma~\ref{l:predictor-1step}. By induction,
    \begin{align*}
W_2(q\Pode^N, q\Podes^N)
        &\lesssim \sum_{n=1}^N \bigl(L d^{1/2} h_n^2\max\{L^{1/2}, (T-t_0-t_n)^{-1/2}\} + h_n \esc\bigr) \\
        &\qquad\qquad\qquad{} \times \exp(O(L\,(h_{n+1}+\cdots + h_N)))\,.
        \end{align*}
       By assumption, $h_{n+1}+\cdots + h_N\le t_N \le 1/L$. In the first case, we get 
        \begin{align*}
        W_2(q\Pode^N, q\Podes^N)
        &\lesssim L^{3/2} d^{1/2} h_{\max} t_N + \esc t_N\,. 
\end{align*}
In the second case we get
\begin{align*}
W_2(q\Pode^N, q\Podes^N)
        &\lesssim Ld^{1/2} h_{\max} \sum_{n=1}^N \fc{h_n}{{(T-t_0 - t_n)}^{1/2}} + \esc t_N \lesssim L^{1/2} d^{1/2} h_{\max} + \esc t_N
\end{align*}
by interpreting the summation as a Riemann sum, and noting that the condition $h_{n+1}\le \fc{T-t_0 - t_n}{2}$ implies that this is a constant-factor approximation of the integral $\int_{T-t_0-t_N}^{T-t_0} \rc{t^{1/2}} \,\D t \lesssim \sqrt{T-t_0}$.
\end{proof}

\paragraph{Choice of step sizes.} In the first case, we can take all the step sizes to be equal, but in the second case, we may need to take decreasing step sizes. Given a target time $T - t_0 - t_N=\delta$, by taking 
$h_1=h_{\max}$ and then 
\begin{align}
\label{e:stepsize}
h_n=\min\Bigl\{\delta,\,h_{\max},\, \fc{T-t_0-t_{n-1}}2\Bigr\}\,,
\end{align}
we can reach the target time in
\begin{align}\label{e:N-round}
    N = O\Bigl(\frac{1}{Lh_{\max}} + \ln \frac{h_{\max}}{\delta}\Bigr)\quad\text{steps}\,.\end{align}