SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \mathsf {K L} \left(P _ {0, t + \Delta t} \| P _ {0, t}\right) = \mathbb {E} _ {x \sim P _ {0}} \mathsf {K L} \left(P _ {t + \Delta t | 0} (\cdot \mid x), P _ {t | 0} (\cdot \mid x)\right) = \mathsf {K L} \left(\mathcal {N} (0, (t + \Delta t) I) \| \mathcal {N} (0, t I)\right) \\ = \frac {d}{2} \left(- \ln \frac {t + \Delta t}{t} + \frac {t + \Delta t}{t} - 1\right) = \frac {d}{4} \left(\frac {\Delta t}{t}\right) ^ {2} + O \left(\left(\frac {\Delta t}{t}\right) ^ {3}\right). \\ \end{array}
$$
Plugging into (11) gives
$$
\mathbb {E} \left[ \left\| \left[ \partial_ {t} \nabla \ln p _ {t} (y) \right] \right| _ {y = x _ {t}} \right\rVert^ {2} ] \leq \frac {L ^ {2} d}{t}. \tag {12}
$$
For the second term, by assumption we have  $\| \nabla^2\ln p_t\|_{\mathrm{op}}\leq L$ . Then, since  $x_{t}\sim p_{t}$
$$
\mathbb {E} \left[ \| \nabla^ {2} \ln p _ {t} \left(x _ {t}\right) \nabla \ln p _ {t} \left(x _ {t}\right) \| ^ {2} \right] \leq L ^ {2} \mathbb {E} _ {p _ {t}} \left[ \| \nabla \ln p _ {t} \| ^ {2} \right] \leq L ^ {3} d \tag {13}
$$
using the fact  $\mathbb{E}_{\mu}[\| \nabla \ln \mu \| ^2 ]\leq Ld$  for any measure  $\mu$  such that  $\ln \mu$  is  $L$ -smooth, which follows from integration by parts. From (9), (12), and (13), and the elementary inequality  $\langle a,b\rangle \leq \| a\| ^2 +4\| b\| ^2$  we get
$$
\begin{array}{l} \mathbb {E} \left[ \| \partial_ {t} \nabla \ln p _ {t} (x _ {t}) \| ^ {2} \right] \leq \mathbb {E} \left[ \| \left[ \partial_ {t} \nabla \ln p _ {t} (y) \right] | _ {y = x _ {t}} \| ^ {2} \right] + \mathbb {E} \left[ \| \nabla^ {2} \ln p _ {t} (x _ {t}) \nabla \ln p _ {t} (x _ {t}) \| ^ {2} \right] \\ \leq L ^ {2} d \left(L + \frac {1}{t}\right). \\ \end{array}
$$
The above result holds for the dynamics  $\dot{x}_t = -\frac{1}{2}\nabla \ln p_t(x_t)$  for which  $(p_t)_{t\geq 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.
Lemma 2 (Reparameterization). Suppose that  $(x_{t})_{t\geq 0}$  satisfies the probability flow ODE for Brownian motion starting at  $p_0$ ; that is, letting  $p_t = p_0*\mathcal{N}(0,tI)$ , we have  $x_0\sim p_0$ ,  $\dot{x}_t = -\frac{1}{2}\nabla \ln p_t(x_t)$ . Then, if we set
$$
y _ {t} = e ^ {- t} x _ {e ^ {2 t} - 1},
$$
then  $(y_{t})_{t\geq 0}$  satisfies the probability flow ODE for the OU process starting at  $p_0$ ; that is, letting  $q_{t}^{\rightarrow}$  be the density of the OU process at time  $t$ , we have  $y_{0}\sim p_{0} = q_{0}^{\rightarrow}$ ,  $\dot{y}_t = -y_t - \nabla \ln q_t^\rightarrow (y_t)$ .
Proof. By direct calculation, one can check that for any  $y \in \mathbb{R}^d$ , it holds that  $q_t^\rightarrow(y) \propto p_{e^{2t} - 1}(e^ty)$ . The claim follows from the chain rule.
Lemma 3 (Score perturbation for OU). Suppose  $q_{t}^{\rightarrow}$  is the density of the OU process at time  $t$ , started at  $q_{0}^{\rightarrow}$ , and  $y_0\sim q_0^\rightarrow$ ,  $\dot{y}_t = -y_t - \nabla \ln q_t^\rightarrow (y_t)$ . Suppose for all  $t$  and all  $x$  that  $\| \nabla^2\ln q_t^\rightarrow (x)\|_{\mathrm{op}}\leq L$  where  $L\geq 1$ . Then,
$$
\mathbb {E} [ \| \partial_ {t} \nabla \ln q _ {t} ^ {\rightarrow} (y _ {t}) \| ^ {2} ] \lesssim L ^ {2} d \left(L \vee \frac {1}{t}\right).
$$
Proof. Using the relationship  $q_{t}^{\rightarrow}(y) \propto p_{e^{2t} - 1}(e^{t}y)$ ,
$$
\nabla \ln q _ {t} ^ {\rightarrow} (y) = e ^ {t} \nabla \ln p _ {e ^ {2 t} - 1} (e ^ {t} y),
$$
$$
\partial_ {t} \nabla \ln q _ {t} ^ {\rightarrow} (y _ {t}) = \underbrace {e ^ {t} \nabla \ln p _ {e ^ {2 t} - 1} (x _ {e ^ {2 t} - 1})} _ {=: A} + \underbrace {e ^ {t} \partial_ {s} \nabla \ln p _ {s} (x _ {s}) | _ {s = e ^ {2 t} - 1} \cdot 2 e ^ {2 t}} _ {=: B}.
$$
If  $\| \nabla^2\ln q_t^\rightarrow \|_{\mathrm{op}}\leq L$  , then  $\| \nabla^2\ln p_{e^{2t} - 1}\|_{\mathrm{op}}\leq e^{-2t}L.$  By Lemma 1,
$$
\mathbb {E} [ \| \partial_ {s} \nabla \ln p _ {s} (x _ {s}) | _ {s = e ^ {2 t} - 1} \| ^ {2} ] \lesssim e ^ {- 4 t} L ^ {2} d \left(e ^ {- 2 t} L \vee \frac {1}{e ^ {2 t} - 1}\right).
$$
$$
\mathbb {E} \left[ B ^ {2} \right] \lesssim L ^ {2} d \left(L \vee \frac {1}{t}\right).
$$
$$
\mathbb {E} \left[ A ^ {2} \right] \leq e ^ {2 t} \mathbb {E} \left[ \| \nabla \ln p _ {e ^ {2 t} - 1} \left(x _ {e ^ {2 t} - 1}\right) \| ^ {2} \right] \leq e ^ {2 t} e ^ {- 2 t} L d \leq L d.
$$
Finally, we use Lemma 3 to derive a bound on how much the score changes along the trajectory of the probability flow ODE.
Corollary 1. Consider the setting of Lemma 3, and suppose  $0 < s < t$ ,  $h = t - s$ .
1. If  $s, t \gtrsim 1 / L$ , then
$$
\mathbb {E} \left[ \| \nabla \ln q _ {t} ^ {\rightarrow} (x _ {t}) - \nabla \ln q _ {s} ^ {\rightarrow} (x _ {s}) \| ^ {2} \right] \lesssim L ^ {3} d h ^ {2}.
$$
2. If  $\frac{t}{2} \leq s \leq t \lesssim \frac{1}{L}$ , then
$$
\mathbb {E} \left[ \| \nabla \ln q _ {t} ^ {\rightarrow} (x _ {t}) - \nabla \ln q _ {s} ^ {\rightarrow} (x _ {s}) \| ^ {2} \right] \lesssim \frac {L ^ {2} d h ^ {2}}{t}.
$$
Proof. By Lemma 3,
$$
\begin{array}{l} \mathbb {E} \left[ \| \nabla \ln q _ {t} ^ {\rightarrow} (x _ {t}) - \nabla \ln q _ {s} ^ {\rightarrow} (x _ {s}) \| ^ {2} \right] = \mathbb {E} \left[\left\| \int_ {s} ^ {t} \partial_ {u} \nabla \ln q _ {u} ^ {\rightarrow} (x _ {u}) \mathrm {d} u \right\| ^ {2} \right] \\ \leq (t - s) \int_ {s} ^ {t} \mathbb {E} [ \| \partial_ {u} \nabla \ln q _ {u} ^ {\rightarrow} (x _ {u}) \| ^ {2} ] \mathrm {d} u \\ \lesssim h \int_ {s} ^ {t} L ^ {2} d \max  \left\{L, \frac {1}{u} \right\} d u. \\ \end{array}
$$
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\frac{1}{u}\mathrm{d}u) = O(L^2 dh\ln (t / s)) = O(L^2 dh^2 /t)$ .
C Predictor step
Next, we need an ODE discretization analysis.
Lemma 4. Suppose the score function satisfies Assumption 2. Assume that  $L \geq 1$ ,  $h \lesssim 1 / L$ , and  $T - (t_0 + h) \geq \frac{T - t_0}{2}$ . Then
$$
W _ {2} (q P _ {\sf O D E} ^ {t _ {0}, h}, q \widehat {P} _ {\sf O D E} ^ {t _ {0}, h}) \lesssim L d ^ {1 / 2} h ^ {2} \left(L ^ {1 / 2} \vee \frac {1}{(T - t _ {0}) ^ {1 / 2}}\right) + h \varepsilon_ {\sf s c}.
$$
Proof. We have the ODEs
$$
\dot {x} _ {t} = x _ {t} + \nabla \ln q _ {t} (x _ {t}),
$$
$$
\dot {\hat {x}} _ {t} = \widehat {x} _ {t} + s _ {t _ {0}} (\widehat {x} _ {t _ {0}}),
$$
for  $t_0 \leq t \leq t_0 + h$ , with  $x_{t_0} = \widehat{x}_{t_0} \sim q$ ,  $x_{t_0 + h} \sim qP_{\mathsf{ODE}}$ , and  $\widehat{x}_{t_0 + h} \sim q\widehat{P}_{\mathsf{ODE}}$ . Then,
$$
\begin{array}{l} \partial_ {t} \| x _ {t} - \widehat {x} _ {t} \| ^ {2} = 2 \left\langle x _ {t} - \widehat {x} _ {t}, \dot {x} _ {t} - \dot {\widehat {x}} _ {t} \right\rangle \\ = 2 \left(\| x _ {t} - \widehat {x} _ {t} \| ^ {2} + \langle x _ {t} - \widehat {x} _ {t}, \nabla \ln q _ {t} (x _ {t}) + s _ {t _ {0}} (\widehat {x} _ {t _ {0}}) \rangle\right) \\ \leq \left(2 + \frac {1}{h}\right) \| x _ {t} - \widehat {x} _ {t} \| ^ {2} + h \| \nabla \ln q _ {t} (x _ {t}) - s _ {t _ {0}} (\widehat {x} _ {t _ {0}}) \| ^ {2}. \\ \end{array}
$$
By Gronwall's inequality, noting that  $h = O(1)$ ,
$$
\begin{array}{l} \mathbb {E} \left[ \| x _ {t _ {0} + h} - \widehat {x} _ {t _ {0} + h} \| ^ {2} \right] \leq \exp \left(\left(2 + \frac {1}{h}\right) h\right) \int_ {t _ {0}} ^ {t _ {0} + h} h \mathbb {E} \left[ \| \nabla \ln q _ {t} (x _ {t}) - s _ {t _ {0}} (\widehat {x} _ {t _ {0}}) \| ^ {2} \right] d t \\ \lesssim h \int_ {t _ {0}} ^ {t _ {0} + h} \mathbb {E} \left[ \| \nabla \ln q _ {t} \left(x _ {t}\right) - s _ {t _ {0}} \left(\widehat {x} _ {t _ {0}}\right) \| ^ {2} \right] d t. \tag {14} \\ \end{array}
$$
We split up the error term as
$$
\| \nabla \ln q _ {t} (x _ {t}) - s _ {t _ {0}} (\widehat {x} _ {t _ {0}}) \| ^ {2} \lesssim \| \nabla \ln q _ {t} (x _ {t}) - \nabla \ln q _ {t _ {0}} (x _ {t _ {0}}) \| ^ {2} + \| \nabla \ln q _ {t _ {0}} (x _ {t _ {0}}) - s _ {t _ {0}} (\widehat {x} _ {t _ {0}}) \| ^ {2}.
$$
By Corollary 1, the expectation of the first term is bounded by
$$
\mathbb {E} [ \| \nabla \ln q _ {t} (x _ {t}) - \nabla \ln q _ {t _ {0}} (x _ {t _ {0}}) \| ^ {2} ] \lesssim L ^ {2} d h ^ {2} \left(L \vee \frac {1}{T - t _ {0}}\right).
$$
The second term is bounded in expectation by  $\varepsilon_{\mathrm{sc}}^2$ . Plugging back into (14) gives
$$
\mathbb {E} [ \| x _ {t _ {0} + h} - \widehat {x} _ {t _ {0} + h} \| ^ {2} ] \lesssim h ^ {2} \left(L ^ {2} d h ^ {2} \left(L \vee \frac {1}{T - t _ {0}}\right) + \varepsilon_ {\mathfrak {s c}} ^ {2}\right).
$$
The Wasserstein distance is bounded by the square root of this quantity, and the lemma follows.
Lemma 4 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).
Remark 2. Our improved score perturbation lemma is necessary to obtain this  $d^{1/2}$  dependence. The original score perturbation lemma [LLT22, Lemma C.11-12] combined with a space discretization bound gives a bound of
$$
\mathbb {E} \big [ \| \nabla \ln q _ {t} ^ {\rightarrow} (x _ {t}) - \nabla \ln q _ {s} ^ {\rightarrow} (x _ {s}) \| ^ {2} \big ] \lesssim L ^ {2} d h
$$
in place of Corollary 1. Note this is a  $\frac{1}{2}$ -Hölder continuity bound rather than a Lipschitz bound. The bound in Lemma 4 then becomes
$$
W _ {2} \left(q P _ {\text {O D E}} ^ {t _ {0}, h}, q \widehat {P} _ {\text {O D E}} ^ {t _ {0}, h}\right) \lesssim L d ^ {1 / 2} h ^ {3 / 2} + h \varepsilon_ {\mathfrak {s c}},
$$
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.
We now iterate Lemma 4 to obtain the following result. Note that we now need to also assume that the score estimate is  $L$ -Lipschitz.
Lemma 5. Suppose that both Assumptions 2 and 3 hold. Let  $h_1, \ldots, h_N > 0$  be a sequence such that letting  $t_N = h_1 + \dots + h_N$ , we have  $t_N \leq 1 / L$ . Let  $h_{\max} = \max_{1 \leq n \leq N} h_n$ .
1. If  $T - (t_0 + t_N) \gtrsim 1 / L$ , then
$$
W _ {2} (q P _ {\sf O D E} ^ {t _ {0}, h _ {1}, \dots , h _ {N}}, q \widehat {P} _ {\sf O D E} ^ {t _ {0}, h _ {1}, \dots , h _ {N}}) \lesssim L ^ {3 / 2} d ^ {1 / 2} h _ {\mathrm {m a x}} t _ {N} + \varepsilon_ {\sf s c} t _ {N} \leq L ^ {1 / 2} d ^ {1 / 2} h _ {\mathrm {m a x}} + \frac {\varepsilon_ {\sf s c}}{L}.
$$
2. If  $T - t_0 \lesssim 1 / L$  and  $h_{n + 1} \leq \frac{T - t_0 - t_n}{2}$  for each  $n$ , then
$$
W _ {2} (q P _ {\texttt {O D E}} ^ {t _ {0}, h _ {1}, \dots , h _ {N}}, q \widehat {P} _ {\texttt {O D E}} ^ {t _ {0}, h _ {1}, \dots , h _ {N}}) \lesssim L ^ {1 / 2} d ^ {1 / 2} h _ {\max } + \varepsilon_ {\mathfrak {s c}} t _ {N} \leq L ^ {1 / 2} d ^ {1 / 2} h _ {\max } + \frac {\varepsilon_ {\mathfrak {s c}}}{L}.
$$
Proof. We abbreviate  $P_{\mathrm{ODE}}^N \coloneqq P_{\mathrm{ODE}}^{t_0,h_1,\dots,h_N}$  and  $\widehat{P}_{\mathrm{ODE}}^N \coloneqq \widehat{P}_{\mathrm{ODE}}^{t_0,h_1,\dots,h_N}$ . Using the triangle inequality,
$$
\begin{array}{l} W _ {2} \left(q P _ {\mathrm {O D E}} ^ {N}, q \widehat {P} _ {\mathrm {O D E}} ^ {N}\right) \leq W _ {2} \left(q P _ {\mathrm {O D E}} ^ {N}, q P _ {\mathrm {O D E}} ^ {N - 1} \widehat {P} _ {\mathrm {O D E}}\right) + W _ {2} \left(q P _ {\mathrm {O D E}} ^ {N - 1} \widehat {P} _ {\mathrm {O D E}}, q \widehat {P} _ {\mathrm {O D E}} ^ {N}\right) \\ \leq O (L d ^ {1 / 2} h _ {N} ^ {2} \max \{L ^ {1 / 2}, (T - t _ {0} - t _ {N}) ^ {- 1 / 2} \} + h _ {N} \varepsilon_ {\mathfrak {s c}}) \\ + \exp (O (L h _ {N})) W _ {2} (q P _ {\mathrm {O D E}} ^ {N - 1}, q \widehat {P} _ {\mathrm {O D E}} ^ {N - 1}) \\ \end{array}
$$
where the bound on the first term is by Lemma 4. By induction,
$$
\begin{array}{l} W _ {2} \left(q P _ {\mathrm {O D E}} ^ {N}, q \widehat {P} _ {\mathrm {O D E}} ^ {N}\right) \lesssim \sum_ {n = 1} ^ {N} \left(L d ^ {1 / 2} h _ {n} ^ {2} \max  \left\{L ^ {1 / 2}, \left(T - t _ {0} - t _ {n}\right) ^ {- 1 / 2} \right\} + h _ {n} \varepsilon_ {\mathrm {s c}}\right) \\ \times \exp (O (L (h _ {n + 1} + \dots + h _ {N}))). \\ \end{array}
$$
By assumption,  $h_{n + 1} + \dots + h_N \leq t_N \leq 1 / L$ . In the first case, we get
$$
W _ {2} \left(q P _ {\mathrm {O D E}} ^ {N}, q \widehat {P} _ {\mathrm {O D E}} ^ {N}\right) \lesssim L ^ {3 / 2} d ^ {1 / 2} h _ {\max } t _ {N} + \varepsilon_ {\mathfrak {s c}} t _ {N}.
$$
In the second case we get
$$
W _ {2} (q P _ {\sf O D E} ^ {N}, q \widehat {P} _ {\sf O D E} ^ {N}) \lesssim L d ^ {1 / 2} h _ {\mathrm {m a x}} \sum_ {n = 1} ^ {N} \frac {h _ {n}}{(T - t _ {0} - t _ {n}) ^ {1 / 2}} + \varepsilon_ {\sf s c} t _ {N} \lesssim L ^ {1 / 2} d ^ {1 / 2} h _ {\mathrm {m a x}} + \varepsilon_ {\sf s c} t _ {N}
$$
by interpreting the summation as a Riemann sum, and noting that the condition  $h_{n + 1} \leq \frac{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} \frac{1}{t^{1/2}} \, \mathrm{d}t \lesssim \sqrt{T - t_0}$ .
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
$$
h _ {n} = \min  \left\{\delta , h _ {\max }, \frac {T - t _ {0} - t _ {n - 1}}{2} \right\},
$$
we can reach the target time in
$$
N = O \left(\frac {1}{L h _ {\operatorname* {m a x}}} + \ln \frac {h _ {\operatorname* {m a x}}}{\delta}\right) \quad \text {s t e p s}.
$$
D Corrector step
In §D.1 (resp. §D.2), we will show that if  $p, q$  are close in Wasserstein distance, then running the corrector step based on the overdamped (resp. underdamped) Langevin diffusion starting from  $p$  and from  $q$  for some amount of time results in distributions which are close in total variation distance. In the end-to-end analysis in §E, we combine this "total variation to Wasserstein" regularization with the Wasserstein discretization analysis of the predictor step in §C in order to establish our final results.
D.1 Corrector via overdamped Langevin
We will take the potential and score estimate defining the Markov kernels  $P_{\mathrm{LD}}$  and  $\widehat{P}_{\mathrm{LMC}}$  from §2.3 to be  $U$  and  $s$  respectively. Recall that these correspond respectively to running the overdamped Langevin diffusion with stationary distribution  $q \propto \exp(-U)$  and running the discretized diffusion with score estimate  $s$ , both for time  $h$ .
The main result of this section is to show that  $p\widehat{P}_{\mathrm{LMC}}^{N}$  and  $q$  are close in total variation if  $p$  and  $q$  are close in Wasserstein.
Theorem 4 (Overdamped corrector). For any  $T_{\mathrm{corr}} \coloneqq Nh \lesssim 1 / L$ ,
$$
\mathsf {T V} (p \widehat {P} _ {\mathsf {L M C}} ^ {N}, q) \lesssim W _ {2} (p, q) / \sqrt {T _ {\mathsf {c o r r}}} + \varepsilon_ {\mathsf {s c}} \sqrt {T _ {\mathsf {c o r r}}} + L \sqrt {d h T _ {\mathsf {c o r r}}}.
$$
In particular, for  $T_{\mathrm{corr}} \asymp 1 / L$ ,
$$
\mathsf {T V} (p \widehat {P} _ {\mathsf {L M C}} ^ {N}, q) \lesssim \sqrt {L} W _ {2} (p, q) + \varepsilon_ {\mathsf {s c}} / \sqrt {L} + \sqrt {L d h}.
$$
We will bound  $\mathrm{TV}(pP_{\mathrm{LD}}^N,q)$  and  $\mathrm{TV}(p\widehat{P}_{\mathrm{LMC}}^N,pP_{\mathrm{LD}}^N)$  separately. For the former, we use the following short-time regularization result:
Lemma 6 ([BGL01, Lemma 4.2]). If  $T_{\mathrm{corr}} \lesssim 1 / L$ , then
$$
\mathrm {T V} (p P _ {\mathrm {L D}} ^ {N}, q) \lesssim \sqrt {\mathrm {K L} (p P _ {\mathrm {L D}} ^ {N} \| q)} \lesssim W _ {2} (p, q) / \sqrt {T _ {\mathrm {c o r r}}}.
$$
Proof. The first inequality is Pinsker's inequality. The second inequality is a consequence of [BGL01, Lemma 4.2], which gives a bound of  $\mathsf{KL}(pP_{\mathsf{LD}}^N\parallel q)\lesssim L(1 + 1 / (e^{2LT_{\mathsf{corr}}} - 1))W_2^2 (p,q)$ . The claim then follows from simplifying by using  $T_{\mathrm{corr}}\lesssim 1 / L$
For the latter term, we introduce notation for two stochastic processes.
$$
\mathrm {d} x _ {t} ^ {\circ} = - \nabla U \left(x _ {t} ^ {\circ}\right) \mathrm {d} t + \sqrt {2} \mathrm {d} B _ {t}, \quad x _ {0} ^ {\circ} \sim q,
$$
$$
\mathrm {d} x _ {t} = - \nabla U (x _ {t})   \mathrm {d} t + \sqrt {2}   \mathrm {d} B _ {t}  , \hskip 1 4. 2 2 6 3 7 8 p t x _ {0} \sim p  .
$$
Note that for any integer  $k\geq 0$
$$
x _ {k h} ^ {\circ} \sim q P _ {\mathrm {L D}} ^ {k}, \quad x _ {k h} \sim p P _ {\mathrm {L D}} ^ {k}.
$$
Observe that marginally,  $qP_{\mathrm{LD}}^k = q$  for any  $k \geq 0$  because  $q$  is the stationary distribution of the Langevin diffusion. The three processes are coupled by using the same Brownian motion and by coupling  $x_0 = \widehat{x}_0 \sim p$  and  $x_0^\circ \sim q$  optimally.
Before we proceed to bound  $\mathrm{TV}(pP_{\mathsf{LD}}^N,p\widehat{P}_{\mathsf{LMC}}^N)$ , we need the following simple lemma.
Lemma 7. If  $T_{\mathrm{corr}} \lesssim 1 / L$ , then
$$
\mathbb {E} [ \| x _ {t} - x _ {t} ^ {\circ} \| ^ {2} ] \lesssim W _ {2} ^ {2} (p, q)
$$
for all  $0 \leq t \leq T_{\mathrm{corr}}$ .
Proof. By Itô's formula,
$$
\mathrm {d} \big (\| x _ {t} - x _ {t} ^ {\circ} \| ^ {2} \big) = - 2 \left\langle x _ {t} - x _ {t} ^ {\circ}, \nabla U (x _ {t}) - \nabla U (x _ {t} ^ {\circ}) \right\rangle \leq 2 L \left\| x _ {t} - x _ {t} ^ {\circ} \right\| ^ {2}.
$$
$$
\left\| x _ {t} - x _ {t} ^ {\circ} \right\| ^ {2} \leq e ^ {2 L t} \left\| x _ {0} - x _ {0} ^ {\circ} \right\| ^ {2},
$$
so that if we couple the two processes by coupling  $x_0$  and  $x_0^\circ$  optimally, we conclude that
$$
\mathbb {E} \left[ \left\| x _ {t} - x _ {t} ^ {\circ} \right\| ^ {2} \right] \leq e ^ {2 L t} \mathbb {E} \left[ \left\| x _ {0} - x _ {0} ^ {\circ} \right\| ^ {2} \right] = e ^ {2 L t} W _ {2} ^ {2} (p, q) \lesssim W _ {2} ^ {2} (p, q),
$$
recalling that  $t \leq T_{\mathrm{corr}} \lesssim 1 / L$  by hypothesis.
It remains to bound TV  $(p\widehat{P}_{\mathrm{LMC}}^N,pP_{\mathrm{LD}}^N)$
Lemma 8. If  $T_{\mathrm{corr}} \lesssim 1 / L$ , then
$$
\mathsf {T V} (p \widehat {P} _ {\mathsf {L M C}} ^ {N}, p P _ {\mathsf {L D}} ^ {N}) \lesssim \sqrt {\mathsf {K L} (p P _ {\mathsf {L D}} ^ {N} \| p \widehat {P} _ {\mathsf {L M C}} ^ {N})} \lesssim L \sqrt {T _ {\mathsf {c o r r}}} W _ {2} (p, q) + \varepsilon_ {\mathsf {s c}} \sqrt {T _ {\mathsf {c o r r}}} + L \sqrt {d h T _ {\mathsf {c o r r}}}.
$$
Proof. As  $x$  and  $\widehat{x}$  are driven by the same Brownian motion, by Girsanov's theorem and the data processing inequality we have
$$
\mathrm {K L} \left(p P _ {\mathrm {L D}} ^ {N} \| p \widehat {P} _ {\mathrm {L M C}} ^ {N}\right) \lesssim \sum_ {k = 0} ^ {N - 1} \int_ {k h} ^ {(k + 1) h} \mathbb {E} \left[ \| s \left(x _ {k h}\right) - \nabla U \left(x _ {u}\right) \| ^ {2} \right] \mathrm {d} u.
$$
We can decompose the integrand as follows:
$$
\begin{array}{l} \mathbb {E} \left[ \| s (x _ {k h}) - \nabla U (x _ {u}) \| ^ {2} \right] \lesssim \mathbb {E} \left[ \| s (x _ {k h}) - s \left(x _ {k h} ^ {\circ}\right) \| ^ {2} + \| s \left(x _ {k h} ^ {\circ}\right) - \nabla U \left(x _ {k h} ^ {\circ}\right) \| ^ {2} \right. \\ \left. + \left\| \nabla U \left(x _ {k h} ^ {\circ}\right) - \nabla U \left(x _ {u} ^ {\circ}\right) \right\| ^ {2} + \left\| \nabla U \left(x _ {u} ^ {\circ}\right) - \nabla U \left(x _ {u}\right) \right\| ^ {2} \right] \\ \leq L ^ {2} \mathbb {E} [ \| x _ {k h} - x _ {k h} ^ {\circ} \| ^ {2} ] + \varepsilon_ {\mathrm {s c}} ^ {2} + L ^ {2} \mathbb {E} [ \| x _ {k h} ^ {\circ} - x _ {u} ^ {\circ} \| ^ {2} ] + L ^ {2} \mathbb {E} [ \| x _ {u} ^ {\circ} - x _ {u} \| ^ {2} ] \\ \end{array}
$$
$$
\lesssim L ^ {2} W _ {2} ^ {2} (p, q) + \varepsilon_ {\mathfrak {s c}} ^ {2} + L ^ {2} \mathbb {E} [ \| x _ {k h} ^ {\circ} - x _ {u} ^ {\circ} \| ^ {2} ] \tag {15}
$$
where we used Lemma 7 to bound  $\mathbb{E}[\| x_u^\circ -x_u\|^2 ]$
It remains to bound  $\mathbb{E}[\| x_{kh}^{\circ} - x_u^{\circ}\| ^2 ]$  . Note that
$$
\begin{array}{l} \mathbb {E} [ \| x _ {k h} ^ {\circ} - x _ {u} ^ {\circ} \| ^ {2} ] = \mathbb {E} \left[ \left| \left| \int_ {k h} ^ {u} - \nabla U (x _ {s} ^ {\circ}) \mathrm {d} s + \sqrt {2} \left(B _ {u} - B _ {k h}\right) \right| \right| ^ {2} \right] \lesssim h \int_ {k h} ^ {u} \mathbb {E} [ \| \nabla U (x _ {s} ^ {\circ}) \| ^ {2} ] \mathrm {d} s + d h \\ \leq L d h ^ {2} + d h \lesssim d h, \\ \end{array}
$$
where in the last step we used that  $\mathbb{E}[\| \nabla U(x_u^\circ)\|^2 ]\leq Ld$  . Substituting this into (15), we obtain
$$
\mathsf {K L} (p P _ {\mathsf {L D}} ^ {N} \parallel p \widehat {P} _ {\mathsf {L M C}} ^ {N}) \lesssim L ^ {2} T _ {\mathsf {c o r r}} W _ {2} ^ {2} (p, q) + \varepsilon_ {\mathsf {s c}} ^ {2} T _ {\mathsf {c o r r}} + L ^ {2} d h T _ {\mathsf {c o r r}}.
$$
The claimed bound on  $\mathsf{TV}(pP_{\mathsf{LD}}^N,p\widehat{P}_{\mathsf{LMC}}^N)$  follows by Pinsker's inequality.
Proof of Theorem 4. This is immediate from Lemma 6 and Lemma 8, recalling that  $T_{\mathrm{corr}} \lesssim 1 / L$  so that the bound in Lemma 6 dominates the  $W_{2}(p,q)$  term in Lemma 8.
D.2 Corrector via underdamped Langevin
Throughout, we set the friction parameter to
$$
\gamma \asymp \sqrt {L}.
$$
We will take the potential and score estimate defining the Markov kernels  $P_{\mathrm{ULD}}$  and  $\widehat{P}_{\mathrm{ULMC}}$  from §2.3 to be  $U$  and  $s$  respectively. Recall that these correspond respectively to running the underdamped Langevin diffusion with stationary distribution  $q$  and running the discretized diffusion with score estimate  $s$ , both for time  $h$ .
Given probability measures  $p$  and  $q$ , we write  $\pmb{p} \coloneqq p \otimes \gamma_d$  and  $\pmb{q} \coloneqq q \otimes \gamma_d$ , where  $\gamma_d$  is the standard Gaussian measure in  $\mathbb{R}^d$ .
The main result of this section is to show that  $\pmb{p}\widehat{P}_{\mathrm{ULMC}}^{N}$  and  $\pmb{q}$  are close in total variation if  $p$  and  $q$  are close in Wasserstein. Compared to §D.1, the discretization error for the underdamped Langevin diffusion is smaller.
Theorem 5 (Underdamped corrector). For  $T_{\mathrm{corr}} \lesssim 1 / \sqrt{L}$
$$
\mathsf {T V} (\boldsymbol {p} \widehat {P} _ {\mathsf {U L M C}} ^ {N}, \boldsymbol {q}) \lesssim \frac {W _ {2} (p , q)}{L ^ {1 / 4} T _ {\mathsf {c o r r}} ^ {3 / 2}} + \frac {\varepsilon_ {\mathsf {s c}} T _ {\mathsf {c o r r}} ^ {1 / 2}}{L ^ {1 / 4}} + L ^ {3 / 4} T _ {\mathsf {c o r r}} ^ {1 / 2} d ^ {1 / 2} h.
$$
In particular, if we take  $T_{\mathrm{corr}} \asymp 1 / \sqrt{L}$ , then
$$
\mathsf {T V} (\boldsymbol {p} \widehat {P} _ {\mathsf {U L M C}} ^ {N}, \boldsymbol {q}) \lesssim \sqrt {L} W _ {2} (p, q) + \varepsilon_ {\mathsf {s c}} / \sqrt {L} + \sqrt {L d} h.
$$
Although the validity of Girsanov's theorem typically requires Novikov's condition to be satisfied, this can be avoided via an approximation argument as in [Che+23a].
We will bound  $\mathrm{TV}(\pmb{p}P_{\mathrm{ULD}}^N,\pmb{q})$  and  $\mathrm{TV}(\pmb{p}\widehat{P}_{\mathrm{ULMC}}^N,\pmb{p}P_{\mathrm{ULD}}^N)$  separately. For the former, we use the short-time regularization result of [GW12]:
Lemma 9. If  $T_{\mathrm{corr}} \lesssim 1 / \sqrt{L}$ , then
$$
\mathrm {T V} (\boldsymbol {p} P _ {\mathsf {U L D}} ^ {N}, \boldsymbol {q}) \lesssim \sqrt {\mathsf {K L} (\boldsymbol {p} P _ {\mathsf {U L D}} ^ {N} \| \boldsymbol {q})} \lesssim \frac {W _ {2} (p , q)}{L ^ {1 / 4} T _ {\text {c o r r}} ^ {3 / 2}}.
$$
Proof. This is a consequence of [GW12, Corollary 4.7 (1)]. The condition to check therein is their Eq. (3.6), which in our setting is satisfied by the constants  $K_{1} = L$  and  $K_{2} = \gamma$ . The Corollary then states that for the cost function
$$
c _ {T _ {\text {c o r r}}} \left(\left(z, v\right), \left(z ^ {\prime}, v ^ {\prime}\right)\right) := \inf  _ {t \in \left(0, T _ {\text {c o r r}} \right]} \frac {t}{2 \gamma} \left\{\left(\frac {6}{t ^ {2}} + L + \frac {3 \gamma}{2 t}\right) \| z - z ^ {\prime} \| + \left(\frac {4}{t} + \frac {4 L t}{2 7} + \gamma\right) \| v - v ^ {\prime} \| \right\} ^ {2},
$$