\section{Score perturbation}\label{sec:score_perturb}


In this section, we prove a score perturbation lemma which refines that of~\cite{leelutan22sgmpoly}.
This improved lemma is necessary in order to obtain $O(\sqrt d)$ dependence for the probability flow ODE\@.

\begin{lem}[Score perturbation]
\label{l:sp}
Suppose $p_t = p_0*\mc N(0,tI)$ and $x_0\sim p_0$, $\dot x_t = -\frac{1}{2}\,\nb \log p_t(x_t)$. 
     Suppose that $\norm{\nabla^2 \ln p_{(t-\rc{2L})\vee 0}(x)}_{\msf{op}} \le L$ for all $x$.
     Then
    \begin{align*}
    \E
    [\ve{\pl_t\nb \log p_{t}(x_t)}^2] &\le
    L^2d \,\bigl(L + \rc t\bigr)\,.
\end{align*}
\end{lem}
\begin{proof}
    Without loss of generality, we may assume $t\le \rc{2L}$, as otherwise, noting that $p_t = p_{t-\rc{2L}} * N(0,\rc{2L}I)$, 
    we may replace $p_0$ with $p_{t-\rc{2L}}$ and $t$ with $\rc{2L}$.
    Suppose $p_0(x)=e^{-V(x)}$.
    Let $P_{0,t}$ denote the joint distribution of $(X_0, X_t)$ where $X_t = X_0 + \sqrt t\,Z$ with $Z\sim \mc N(0,I)$ independent of $X_0$, and let $P_{0\mid t}(\cdot \mid x_t)$ denote the conditional distribution of $X_0$ given $X_t = x_t$.
    We first note that since
    \begin{align*}
        \ln p_t(y)
        &= \ln \int \exp\bigl( - V(x) - \frac{1}{2t}\,\norm{y-x}^2\bigr)\,\D x\,,
    \end{align*}
    we have the following calculations:
    \begin{align*}
        \nb \log p_t(y)
        &= -\rc t \E_{P_{0\mid t}(\cdot\mid y)} (y-\cdot) = -\E_{P_{0\mid t}(\cdot\mid y)} (\nb V)\,,\\
        \nb^2 \log p_t(y) &=
        \Cov_{P_{0\mid t}(\cdot\mid y)}(\nb V) - \E_{P_{0\mid t}(\cdot\mid y)}(\nb^2 V)\,.
    \end{align*}
    Using $\dot x_t=-\frac{1}{2}\,\nb \log p_t(x_t)$, we calculate
\begin{align}
\label{e:score-chain}
    \pl_t \nb \log p_t(x_t)
    &=[\pl_t \nb \log p_t(y)]|_{y=x_t}
    -\frac{1}{2}\, \nb^2\log p_t(x_t) \,\nb \log p_t(x_t)\,.
\end{align}
We bound each term above separately.
For the first term, a quick calculation shows that $\pl_t \nb \ln p_t(y) = -\Cov_{P_{0\mid t}(\cdot\mid y)}\bigl(\fc{\ve{y-\cdot}^2}{2t^2}, \nb V\bigr)$ is finite a.s.: by Cauchy{--}Schwarz, it suffices to show $\E_{P_{0|t}(\cdot|y)}[\ve{y-\cdot}^4]$ and $\E_{P_{0|t}(\cdot|y)}[\ve{\nb V}^2]$ are finite for all $y$, and this follows because for $t\le \frac{1}{2L}$ the measure $P_{0\mid t}(\cdot \mid y)$ is strongly log-concave and $\norm{\nabla V}^2$ can be bounded by a quadratic.
Because $\nb V$ is $L$-Lipschitz,
    \begin{align}
    \nonumber 
        \ve{[\pl_t \nb \log p_t(y)]|_{y=x_t}}^2 &= \norm{[\pl_t \E_{P_{0|t}(\cdot|y)}(\nb V)]|_{y=x_t}}^2\\
            \nonumber 
        &=
        \Bigl\lVert 
        \lim_{\De t\to 0} \rc{\De t}\, 
        \bigl[\E_{P_{0|t+\De t}(\cdot|y)}[\nb V]-
        \E_{P_{0|t}(\cdot|y)}[\nb V]
        \bigr]\bigr|_{y=x_t}
        \Bigr\rVert^2\\
        \label{e:dt-nb-pt}
        &\le L^2 \liminf_{\De t\to 0} 
        \rc{{(\De t)}^2}\,
        W_1^2\bigl(P_{0\mid t+\De t}(\cdot\mid x_t),
        P_{0\mid t}(\cdot\mid x_t)\bigr)\,.
    \end{align}

    Now $P_{0\mid t}(\cdot \mid y)$ has density $p_{0\mid t}(x\mid y)\propto p_0(x)\,e^{-\fc{\ve{x-y}^2}{2t}}$ so if $\norm{\nb^2 \log p_0}_{\msf{op}}\le L$ and $t\le \rc{2L}$, then $P_{0\mid t}$ is
    $\rc{2t}$-strongly log-concave.
    By Talagrand's transport cost inequality,
    \begin{align*}
        W_1^2\bigl(P_{0\mid t+\De t}(\cdot\mid x_t),
        P_{0\mid t}(\cdot\mid x_t)\bigr)
        &\le 
        4t \KL\bigl(P_{0\mid t+\De t}(\cdot\mid x_t)\bigm\Vert P_{0\mid t}(\cdot\mid x_t)\bigr)\,.
    \end{align*}
    Plugging this back in~\eqref{e:dt-nb-pt} and using Fatou's lemma and the chain rule for KL,
    \begin{align}
    \nonumber
        \E[\ve{[\pl_t \nb \log p_t(y)]|_{y=x_t}}^2]
        &\le 
        L^2 \E \liminf_{\De t\to 0}\rc{{(\De t)}^2}\,
        4t \KL\bigl(P_{0\mid t+\De t}(\cdot\mid x_t)\bigm\Vert P_{0\mid t}(\cdot\mid x_t)\bigr)
        \nonumber \\
        &\le 
        L^2 \liminf_{\De t\to 0}\rc{{(\De t)}^2}\,
        4t \E \KL\bigl(P_{0\mid t+\De t}(\cdot\mid x_t)\bigm\Vert P_{0\mid t}(\cdot\mid x_t)\bigr) \\
        &\le 
        L^2 \liminf_{\De t\to 0}\rc{{(\De t)}^2}\, 4t  \KL(P_{0,t+\De t} \mmid P_{0,t})\,.\label{e:dt-nb-pt2}
    \end{align}
    Now
    \begin{align*}
        \KL(P_{0,t+\De t} \mmid P_{0,t})
        & = \E_{x\sim P_0} \KL\bigl(P_{t+\De t\mid 0}(\cdot \mid x), P_{t\mid 0}(\cdot \mid x)\bigr) 
        = \KL\bigl(\mc N(0,(t+\De t)I)\bigm\Vert \mc N(0,t I)\bigr)
        \\
        &= \fc d2\,\Bigl(-\log \frac{t+\De t}{t} + \fc{t+\De t}{t}-1 \Bigr)
        =\fc d4\,\bigl(\frac{\De t}{t}\bigr)^2 + O\Bigl(\bigl(\frac{\De t}{t}\bigr)^3\Bigr)\,.
    \end{align*}
    Plugging into~\eqref{e:dt-nb-pt2} gives
    \begin{align}\label{e:sp-1}
        \E[\ve{[\pl_t \nb \log p_t(y)]|_{y=x_t}}^2]
        &\le \fc{L^2d}{t}\,.
    \end{align}

    For the second term,
    by assumption we have $\norm{\nabla^2 \log p_t}_{\msf{op}} \le L$.
    Then, since $x_t \sim p_t$,
    \begin{align}
        \E[\norm{
        \nb^2\log p_t(x_t) \, \nb \log p_t(x_t)
        }^2]
        &\le 
        L^2
        \E_{p_t}[\norm{\nabla \log p_t}^2]
        \le L^3 d
        \label{e:sp-2}
    \end{align}
    using the fact $\E_\mu[\norm{\nabla \log \mu}^2] \le Ld$ for any measure $\mu$ such that $\log \mu$ is $L$-smooth, which follows from integration by parts.
    From~\eqref{e:score-chain},~\eqref{e:sp-1}, and~\eqref{e:sp-2}, and the elementary inequality $\langle a,b \rangle \le \norm a^2 + 4\,\norm b^2$, we get
    \begin{align*}
        \E[\ve{\pl_t \nb \log p_t(x_t)}^2]
        &\le 
        \E[\ve{[\pl_t \nb \log p_t(y)]|_{y=x_t}}^2]
        + \E[\norm{
        \nb^2\log p_t(x_t) \,\nb \log p_t(x_t)
        }^2] \\
        &\le L^2 d \, \bigl( L + \frac{1}{t}\bigr)\,. \qedhere
    \end{align*}
\end{proof}

The above result holds for the dynamics $\dot x_t = -\frac{1}{2}\,\nabla \ln p_t(x_t)$ for which ${(p_t)}_{t\ge 0}$ follows the heat flow; this corresponds to the variance-exploding SGM\@.
In this paper, since we wish to consider the SGM based on the variance-conserving Ornstein{--}Uhlenbeck (OU) process, we can apply the following reparameterization lemma.

\begin{lem}[Reparameterization]\label{l:reparam}
    Suppose that ${(x_t)}_{t\ge 0}$ satisfies the probability flow ODE for Brownian motion starting at $p_0$; that is, letting $p_t = p_0*\mc N(0, tI)$, we have $x_0\sim p_0$, $\dot x_t = -\frac{1}{2}\,\nb \log p_t(x_t)$. Then, if we set
    \[
    y_t = e^{-t}\, x_{e^{2t}-1}\,,
    \]
    then ${(y_t)}_{t\ge 0}$ satisfies the probability flow ODE for the OU process starting at $p_0$; that is, letting $\quo_t$ be the density of the OU process at time $t$, we have $y_0\sim p_0 = \quo_0$, $\dot y_t = -y_t - \nb \ln \quo_t(y_t)$.
\end{lem}
\begin{proof}
    By direct calculation, one can check that for any $y\in\R^d$, it holds that $\quo_t(y) \propto p_{e^{2t}-1}(e^t y)$.
    The claim follows from the chain rule.
\end{proof}

\begin{lem}[Score perturbation for OU]\label{l:sp-ou}
    Suppose $\quo_t$ is the density of the OU process at time $t$, started at $\quo_0$, and $y_0\sim \quo_0$, $\dot y_t = -y_t-\nb \log \quo_t(y_t)$. 
    Suppose for all $t$ and all $x$ that $\norm{\nb^2\log \quo_{t}(x)}_{\msf{op}} \le L$, where $L\ge 1$.
    Then,
    \begin{align*}
    \E
    [\ve{\pl_t\nb \log \quo_{t}(y_t)}^2]
    &\lesssim 
    L^2 d \,\bigl(L \vee \frac{1}{t}\bigr)\,.
    \end{align*}
\end{lem}
\begin{proof}
    Using the relationship $\quo_t(y) \propto p_{e^{2t}-1}(e^t y)$,
    \begin{align*}
        \nb \log\quo_t(y)
        &= e^t\, \nb \log p_{e^{2t}-1}(e^t y)\,,
        \\
        \pl_t\nb  \log\quo_t(y_t)
        &= 
        \ub{e^t\, \nb \ln p_{e^{2t}-1}(x_{e^{2t}-1})}{\eqqcolon A} + 
        \ub{e^t \,\pl_s \nb \ln p_s (x_s)|_{s=e^{2t}-1}\cdot 2e^{2t}}{\eqqcolon B}\,.
    \end{align*}
    If $\norm{\nb^2 \ln \quo_t}_{\msf{op}} \le L$, then $\norm{\nb^2 \ln p_{e^{2t}-1}}_{\msf{op}} \le e^{-2t}L$. By Lemma~\ref{l:sp},
    \begin{align*}
        \E[\norm{\pl_s \nb \ln p_s(x_s)|_{s=e^{2t}-1}}^2]
        &\lesssim e^{-4t} L^2 d \,\bigl(e^{-2t} L \vee \frac{1}{e^{2t} -1}\bigr)\,.
    \end{align*}
    Hence 
    \begin{align*}
        \E[B^2]
        \lesssim L^2 d \,\bigl(L \vee \frac{1}{t}\bigr)\,.
    \end{align*}
    Next, 
    \begin{align*}
        \E[A^2]
        &\le e^{2t} \E[\norm{\nb \ln p_{e^{2t}-1}(x_{e^{2t}-1})}^2]
        \le e^{2t}\, e^{-2t} Ld
        \le Ld\,.
    \end{align*}
    The result follows.
\end{proof}

Finally, we use Lemma~\ref{l:sp-ou} to derive a bound on how much the score changes along the trajectory of the probability flow ODE.
\begin{cor}\label{c:sp}
Consider the setting of Lemma~\ref{l:sp-ou}, and suppose $0<s<t$, $h=t-s$.
\begin{enumerate}
    \item If $s,t\gtrsim 1/L$, then
\[
\E\bigl[\norm{\nb \ln q_t^{\rightarrow}(x_t) - \nb \ln q_s^{\rightarrow}(x_s)}^2\bigr]\lesssim L^3dh^2\,.
\]
\item 
If $\fc{t}2\le s\le t\lesssim \rc L$, then
\[
\E\bigl[ \norm{\nb \ln q_t^{\rightarrow}(x_t) - \nb \ln q_s^{\rightarrow}(x_s)}^2\bigr] \lesssim 
\frac{L^2 dh^2}{t}\,.
\]
\end{enumerate}
\end{cor}
\begin{proof}
By Lemma~\ref{l:sp-ou},
\begin{align*}
\E\bigl[ \norm{\nabla \ln q_t^{\rightarrow}(x_t) - \nabla \ln q_s^{\rightarrow}(x_s)}^2\bigr]
    &= \E\Bigl[\Bigl\lVert \int_s^t \pl_u \nb \ln \quo_u(x_u)\,\D u\Bigr\rVert^2\Bigr]\\
    &\le (t-s) \int_s^t \E[\norm{\pl_u \nb \ln \quo_u(x_u)}^2]\,\D u\\
    &\lesssim h \int_s^t L^2 d\max\bigl\{L, \rc{u}\bigr\} \,\D u\,.
\end{align*}
In the first case, this is bounded by $O(L^3 dh^2)$. In the second case, this is bounded by $O(L^2 dh \int_s^t\rc{u}\,\D u) = O(L^2 dh \log(t/s)) = O(L^2 dh^2/t)$.
\end{proof}