SlowGuess/NeurIPS_2023_OCR huggingface dataset:
we have  $\mathsf{KL}(\pmb {p}P_{\mathsf{ULD}}^N\parallel \pmb {q})\leq W_{c_{T_{\mathsf{corr}}}}(p\otimes \gamma_d,q\otimes \gamma_d)$ . For  $v = v^{\prime}$  and  $T_{\mathrm{corr}}\lesssim 1 / \sqrt{L}$ , note that
$$
c _ {T _ {\text {c o r r}}} \left(\left(z, v\right), \left(z ^ {\prime}, v\right)\right) \lesssim \frac {1}{L ^ {1 / 2} T _ {\text {c o r r}} ^ {3}} \| z - z ^ {\prime} \| ^ {2},
$$
so the claim follows by Pinsker's inequality.
Next, we define the following processes:  $\mathrm{d}z_{t}^{\circ} = v_{t}^{\circ}\mathrm{d}t$ ,  $\mathrm{d}z_{t} = v_{t}\mathrm{d}t$
$$
\mathrm {d} v _ {t} ^ {\circ} = - \gamma v _ {t} ^ {\circ} \mathrm {d} t - \nabla U \left(z _ {t} ^ {\circ}\right) \mathrm {d} t + \sqrt {2 \gamma} \mathrm {d} B _ {t}, \quad \left(z _ {0} ^ {\circ}, v _ {0} ^ {\circ}\right) \sim \boldsymbol {q},
$$
$$
\mathrm {d} v _ {t} = - \gamma v _ {t} \mathrm {d} t - \nabla U (z _ {t}) \mathrm {d} t + \sqrt {2 \gamma} \mathrm {d} B _ {t}, \quad \left(z _ {0}, v _ {0}\right) \sim \boldsymbol {p}.
$$
It follows that for any integer  $k\geq 0$
$$
\left(z _ {k h} ^ {\circ}, v _ {k h} ^ {\circ}\right) \sim \boldsymbol {q} P _ {\text {U L D}} ^ {k} = \boldsymbol {q}, \quad \left(z _ {k h}, v _ {k h}\right) \sim \boldsymbol {p} P _ {\text {U L D}} ^ {k}.
$$
We couple these processes by using the same Brownian motion and coupling  $q \otimes \gamma_d$  and  $p \otimes \gamma_d$  optimally (in particular,  $v_0 = v_0^\circ$ ).
Before we proceed to bound  $\mathsf{TV}(pP_{\mathsf{ULD}}^N,p\widehat{P}_{\mathsf{ULMC}}^N)$ , we start with the following lemma.
Lemma 10. If  $T_{\mathrm{corr}} \lesssim 1 / \sqrt{L}$ , then for all  $0 \leq t \leq T_{\mathrm{corr}}$
$$
\mathbb {E} \left[ \| z _ {t} - z _ {t} ^ {\circ} \| ^ {2} \right] \lesssim W _ {2} ^ {2} (p, q).
$$
$$
\nabla U (z _ {t}) - \nabla U (z _ {t} ^ {\circ}) = \left(\int_ {0} ^ {1} \nabla^ {2} U (z _ {t} + u (z _ {t} ^ {\circ} - z _ {t})) \mathrm {d} u\right) (z _ {t} - z _ {t} ^ {\circ}) := \mathcal {H} _ {t} (z _ {t} - z _ {t} ^ {\circ}),
$$
and the operator  $\mathcal{H}_t$  satisfies
$$
\left\| \mathcal {H} _ {t} \right\| _ {\mathrm {o p}} \leq L.
$$
Let  $\alpha \coloneqq 2 / \gamma$ . For the vectors  $\delta_t \coloneqq (z_t + \alpha v_t) - (z_t^\circ + \alpha v_t^\circ)$  and  $\eta_t \coloneqq z_t - z_t^\circ$ , we have
$$
\begin{array}{l} \frac {1}{2} \mathrm {d} (\| \delta_ {t} \| ^ {2} + \| \eta_ {t} \| ^ {2}) = - (\delta_ {t}, \eta_ {t}) ^ {\intercal} \left[ \begin{array}{c c} (\gamma - \frac {1}{\alpha}) I _ {d} & \frac {1}{2} (\alpha \mathcal {H} _ {t} - \gamma I _ {d}) \\ \frac {1}{2} (\alpha \mathcal {H} _ {t} - \gamma I _ {d}) & \frac {1}{\alpha} I _ {d} \end{array} \right] (\delta_ {t}, \eta_ {t}) \\ \lesssim \sqrt {L} \left(\| \delta_ {t} \| ^ {2} + \| \eta_ {t} \| ^ {2}\right). \\ \end{array}
$$
$$
\left\| \delta_ {t} \right\| ^ {2} + \left\| \eta_ {t} \right\| ^ {2} \leq e ^ {O (\sqrt {L} t)} \left(\left\| \delta_ {0} \right\| ^ {2} + \left\| \eta_ {0} \right\| ^ {2}\right),
$$
so if we couple the two processes by coupling  $z_0$  and  $z_0^\circ$  optimally and taking  $v_0 = v_0^\circ$ , we obtain
$$
\mathbb {E} [ \| z _ {t} - z _ {t} ^ {\circ} \| ^ {2} ] \lesssim \mathbb {E} [ \| \delta_ {t} \| ^ {2} + \| \eta_ {t} \| ^ {2} ] \leq e ^ {O (\sqrt {L} t)} \mathbb {E} [ \| \delta_ {0} \| ^ {2} + \| \eta_ {0} \| ^ {2} ] \lesssim e ^ {O (\sqrt {L} t)} W _ {2} ^ {2} (p, q) \lesssim W _ {2} ^ {2} (p, q),
$$
recalling that  $t \leq T_{\mathrm{corr}} \lesssim 1 / \sqrt{L}$  by hypothesis.
It remains to bound  $\mathrm{TV}(\pmb{p}\widehat{P}_{\mathsf{ULMC}}^N,\pmb{p}P_{\mathsf{ULD}}^N)$
Lemma 11. If  $T_{\mathrm{corr}} \lesssim 1 / \sqrt{L}$ , then
$$
\begin{array}{l} \operatorname {T V} \left(\boldsymbol {p} \widehat {P} _ {\mathrm {U L M C}} ^ {N}, \boldsymbol {p} P _ {\mathrm {U L D}} ^ {N}\right) \lesssim \sqrt {\operatorname {K L} \left(\boldsymbol {p} P _ {\mathrm {U L D}} ^ {N} \| \boldsymbol {p} \widehat {P} _ {\mathrm {U L M C}} ^ {N}\right)} \\ \lesssim L ^ {3 / 4} T _ {\mathrm {c o r r}} ^ {1 / 2} W _ {2} (p, q) + L ^ {- 1 / 4} T _ {\mathrm {c o r r}} ^ {1 / 2} \varepsilon_ {\mathrm {s c}} + L ^ {3 / 4} T _ {\mathrm {c o r r}} ^ {1 / 2} d ^ {1 / 2} h. \\ \end{array}
$$
Proof. As  $(z, v)$  and  $(z^{\circ}, v^{\circ})$  are driven by the same Brownian motion, by Girsanov's theorem and the data processing inequality we have
$$
\mathsf {K L} (\boldsymbol {p} P _ {\mathsf {U L D}} ^ {N} \| \boldsymbol {p} \widehat {P} _ {\mathsf {U L M C}} ^ {N}) \lesssim \frac {1}{\gamma} \sum_ {k = 0} ^ {N - 1} \int_ {k h} ^ {(k + 1) h} \mathbb {E} [ \| s (z _ {k h}) - \nabla U (z _ {u}) \| ^ {2} ] \mathrm {d} u.
$$
$$
\begin{array}{l} \mathbb {E} \left[ \| s (z _ {k h}) - \nabla U (z _ {u}) \| ^ {2} \right] \lesssim \mathbb {E} \left[ \| s (z _ {k h}) - s (z _ {k h} ^ {\circ}) \| ^ {2} + \| s (z _ {k h} ^ {\circ}) - \nabla U (z _ {k h} ^ {\circ}) \| ^ {2} \right. \\ \left. + \left\| \nabla U \left(z _ {k h} ^ {\circ}\right) - \nabla U \left(z _ {u} ^ {\circ}\right) \right\| ^ {2} + \left\| \nabla U \left(z _ {u} ^ {\circ}\right) - \nabla U \left(z _ {u}\right) \right\| ^ {2} \right] \\ \leq L ^ {2} \mathbb {E} [ \| z _ {k h} - z _ {k h} ^ {\circ} \| ^ {2} ] + \varepsilon_ {\mathrm {s c}} ^ {2} + L ^ {2} \mathbb {E} [ \| z _ {k h} ^ {\circ} - z _ {u} ^ {\circ} \| ^ {2} ] + L ^ {2} \mathbb {E} [ \| z _ {u} ^ {\circ} - z _ {u} \| ^ {2} ] \\ \end{array}
$$
$$
\lesssim L ^ {2} W _ {2} ^ {2} (p, q) + \varepsilon_ {\mathfrak {s c}} ^ {2} + L ^ {2} \mathbb {E} \left[ \| z _ {k h} ^ {\circ} - z _ {u} ^ {\circ} \| ^ {2} \right], \tag {16}
$$
where we applied Lemma 10.
It remains to bound  $\mathbb{E}[\| z_{kh}^{\circ} - z_u^{\circ}\| ^2 ]$ . Note that
$$
\mathbb {E} \left[ \| z _ {k h} ^ {\circ} - z _ {u} ^ {\circ} \| ^ {2} \right] = \mathbb {E} \left[ \left\| \int_ {k h} ^ {u} v _ {s} ^ {\circ} \mathrm {d} s \right\| ^ {2} \right] \leq h \int_ {k h} ^ {u} \mathbb {E} \left[ \| v _ {s} ^ {\circ} \| ^ {2} \right] \mathrm {d} s \leq d h ^ {2},
$$
where in the last step we used the fact that  $v_{s}^{\circ} \sim \gamma_{d}$ . Substituting this into (16), we conclude that
$$
\mathsf {K L} \left(\boldsymbol {p} P _ {\mathsf {U L D}} ^ {N} \| \boldsymbol {p} \widehat {P} _ {\mathsf {U L M C}} ^ {N}\right) \lesssim L ^ {3 / 2} T _ {\text {c o r r}} W _ {2} ^ {2} (p, q) + L ^ {- 1 / 2} T _ {\text {c o r r}} \varepsilon_ {\mathsf {s c}} ^ {2} + L ^ {3 / 2} d h ^ {2} T _ {\text {c o r r}}.
$$
The claimed bound on  $\mathsf{TV}(\pmb{p}P_{\mathsf{ULD}}^N,\pmb{p}\widehat{P}_{\mathsf{ULMC}}^N)$  follows by Pinsker's inequality.
Proof of Theorem 5. This is immediate from Lemma 9 and Lemma 11, recalling that  $T_{\mathrm{corr}} \lesssim 1 / \sqrt{L}$  so that the bound in Lemma 9 dominates the  $W_{2}(p,q)$  term in Lemma 11.
Remark 3. In all other sections of this paper, we abuse notation as follows. Given a distribution  $p$  on  $\mathbb{R}^d$ , we write  $p\widehat{P}_{\mathrm{ULMC}}$  to denote the projection onto the  $z$ -coordinate of  $\pmb{p}\widehat{P}_{\mathrm{ULMC}}$ , i.e., we view  $\widehat{P}_{\mathrm{ULMC}}$  as a Markov kernel on  $\mathbb{R}^d$  rather than on  $\mathbb{R}^d\times \mathbb{R}^d$  (and similarly for  $P_{\mathrm{ULD}}$ ).
E End-to-end analysis
Lemma 12 (TV error after one round of predictor and corrector). Choose predictor step sizes  $h_1, \ldots, h_{N_{\text{pred}}}$  as in Lemma 5 with  $T_{\text{pred}} = h_1 + \dots + h_{N_{\text{pred}}} \leq 1 / L$ . That is, if  $T - t_0 - T_{\text{pred}} \lesssim 1 / L$ , then we ensure that  $h_{n+1} \leq \frac{T - t_0 - h_1 - \cdots - h_n}{2}$  for all  $n$ , and if  $T - t_0 \gtrsim 1 / L$ , then we can take  $h_1 = \dots = h_N$ . Let  $h_{\text{pred}} := \max_{1 \leq n \leq N_{\text{pred}}} h_n$  and abbreviate  $P_{\text{ODE}}^{(N_{\text{pred}})} := P_{\text{ODE}}^{t_0, h_1, \ldots, h_{N_{\text{pred}}}}$  (and similarly for  $\widehat{P}_{\text{ODE}}$ ).
1. Consider running the overdamped Langevin corrector for time  $T_{\mathrm{corr}} \asymp 1 / L$ , step size  $h_{\mathrm{corr}}$  and stationary distribution  $q_{t_0}P_{\mathrm{ODE}}^{(N_{\mathrm{pred}})} = q_{t_0 + T_{\mathrm{pred}}}$ ; set  $N_{\mathrm{corr}} = T_{\mathrm{corr}} / h_{\mathrm{corr}}$ . Then,
$$
\mathsf {T V} (p \widehat {P} _ {\text {O D E}} ^ {(N _ {\text {p r e d}})} \widehat {P} _ {\text {L M C}} ^ {N _ {\text {c o r r}}}, q _ {t _ {0} + T _ {\text {p r e d}}}) \leq \mathsf {T V} (p, q _ {t _ {0}}) + O \left(L \sqrt {d} h _ {\text {p r e d}} + \sqrt {L d h _ {\text {c o r r}}} + \frac {\varepsilon_ {\text {s c}}}{\sqrt {L}}\right).
$$
2. Consider running the underdamped Langevin corrector for time  $T_{\mathrm{corr}} \asymp 1 / \sqrt{L}$ , step size  $h_{\mathrm{corr}}$ , and stationary distribution  $q_{t_0}P_{\mathrm{ODE}}^{(N_{\mathrm{pred}})} = q_{t_0 + T_{\mathrm{pred}}}$ ; set  $N_{\mathrm{corr}} = T_{\mathrm{corr}} / h_{\mathrm{corr}}$ . Then,
$$
\begin{array}{l} \operatorname {T V} (p \widehat {P} _ {\text {O D E}} ^ {(N _ {\text {p r e d}})} \widehat {P} _ {\text {U L M C}} ^ {N _ {\text {c o r r}}}, q _ {t _ {0} + T _ {\text {p r e d}}}) \leq \operatorname {T V} (p, q _ {t _ {0}}) + O \left(L \sqrt {d} h _ {\text {p r e d}} + \sqrt {L d} h _ {\text {c o r r}} + \frac {\varepsilon_ {\text {s c}}}{\sqrt {L}}\right). \end{array}
$$
<sup>9</sup>Again, we can avoid checking Novikov's condition using the approximation argument of [Che+23a].
Proof. By the triangle inequality and the data-processing inequality,
$$
\begin{array}{l} \operatorname {T V} \left(p \widehat {P} _ {\text {O D E}} ^ {(N _ {\text {p r e d}})} \widehat {P} _ {\text {L M C}} ^ {N _ {\text {c o r r}}}, q _ {t _ {0} + T _ {\text {p r e d}}}\right) \\ \leq \mathrm {T V} \left(p \widehat {P} _ {\mathrm {O D E}} ^ {(N _ {\text {p r e d}})} \widehat {P} _ {\mathrm {L M C}} ^ {N _ {\text {c o r r}}}, q _ {t _ {0}} \widehat {P} _ {\mathrm {O D E}} ^ {(N _ {\text {p r e d}})} \widehat {P} _ {\mathrm {L M C}} ^ {N _ {\text {c o r r}}}\right) + \mathrm {T V} \left(q _ {t _ {0}} \widehat {P} _ {\mathrm {O D E}} ^ {(N _ {\text {p r e d}})} \widehat {P} _ {\mathrm {L M C}} ^ {N _ {\text {c o r r}}}, q _ {t _ {0} + T _ {\text {p r e d}}}\right) \\ \leq \mathsf {T V} (p, q _ {t _ {0}}) + \mathsf {T V} (q _ {t _ {0}} \widehat {P} _ {\text {O D E}} ^ {(N _ {\text {p r e d}})} \widehat {P} _ {\text {L M C}} ^ {N _ {\text {c o r r}}}, q _ {t _ {0} + T _ {\text {p r e d}}}). \\ \end{array}
$$
For overdamped Langevin, applying Theorem 4,
$$
\operatorname {T V} \left(q _ {t _ {0}} \widehat {P} _ {\text {O D E}} ^ {(N _ {\text {p r e d}})} \widehat {P} _ {\text {L M C}} ^ {N _ {\text {c o r r}}}, q _ {t _ {0} + T _ {\text {p r e d}}}\right) \lesssim \sqrt {L} W _ {2} \left(q \widehat {P} _ {\text {O D E}} ^ {(N _ {\text {p r e d}})}, q _ {t _ {0} + T _ {\text {p r e d}}}\right) + \varepsilon_ {\mathrm {s c}} / \sqrt {L} + \sqrt {L d h _ {\text {c o r r}}} \tag {17}
$$
For the Wasserstein term, Lemma 5 yields
$$
W _ {2} (q _ {t _ {0}} \widehat {P} _ {\texttt {O D E}} ^ {(N _ {\texttt {p r e d}})}, q _ {t _ {0} + T _ {\texttt {p r e d}}}) = W _ {2} (q _ {t _ {0}} \widehat {P} _ {\texttt {O D E}} ^ {(N _ {\texttt {p r e d}})}, q _ {t _ {0}} P _ {\texttt {O D E}} ^ {(N _ {\texttt {p r e d}})}) \lesssim \sqrt {L d} h _ {\texttt {p r e d}} + \frac {\varepsilon_ {\texttt {S C}}}{L}.
$$
Combining these bounds yields the result for the overdamped corrector. For the underdamped corrector, we modify (17) by replacing the use of Theorem 4 with Theorem 5.
We also need the following lemma on the convergence of the OU process.
Lemma 13. Let  $(q_t^{\rightarrow})_{t\geq 0}$  denote the marginal law of the OU process started at  $q_0^{\rightarrow} = q_\star$ . Then, for all  $T\gtrsim 1$ , it holds that
$$
\operatorname {T V} \left(q _ {T} ^ {\rightarrow}, \gamma^ {d}\right) \lesssim (\sqrt {d} + \mathfrak {m} _ {2}) \exp (- T).
$$
Proof. This follows from [CLL23, Lemma C.4]. Alternatively, using the short-time regularization result of [BGL01, Lemma 4.2] for time  $t_0 \asymp 1$  and the Wasserstein contraction of the OU process,
$$
\mathsf {T V} (q _ {T} ^ {\rightarrow}, \gamma^ {d}) \lesssim \sqrt {\mathsf {K L} (q _ {T} ^ {\rightarrow} \| \gamma^ {d})} \lesssim \frac {W _ {2} (q _ {T - t _ {0}} ^ {\rightarrow} , \gamma^ {d})}{\sqrt {t _ {0}}} \leq \exp (- (T - t _ {0})) W _ {2} (q _ {\star}, \gamma^ {d}).
$$
The result follows from  $W_{2}(q_{\star},\gamma^{d})\leq W_{2}(q_{\star},\delta_{0}) + W_{2}(\delta_{0},\gamma^{d})\leq \mathfrak{m}_{2} + \sqrt{d}$
We now prove our main theorems.
Proof of Theorems 2 and 3. For  $t \in [0, T]$ , let  $p_t \coloneqq \mathrm{law}(\widehat{x}_t)$ . From Lemma 13,
$$
\mathrm {T V} (p _ {0}, q _ {0}) = \mathrm {T V} (q _ {T} ^ {\rightarrow}, \gamma^ {d}) \lesssim (\sqrt {d} + \mathfrak {m} _ {2}) \exp (- T).
$$
We divide our analysis according to the two stages of the algorithm. In the first stage, after iterating Lemma 12 for  $N_0 \asymp LT$  steps,
$$
\begin{array}{l} \mathrm {T V} \left(p _ {T - h _ {\text {p r e d}}}, q _ {T - h _ {\text {p r e d}}}\right) \leq \mathrm {T V} \left(p _ {0}, q _ {0}\right) + O \left(L \sqrt {d} h _ {\text {p r e d}} + \sqrt {L d} h _ {\text {c o r r}} ^ {\mathfrak {p}} + \frac {\varepsilon_ {\mathrm {s c}}}{\sqrt {L}}\right) \times N _ {0} \\ \lesssim (\sqrt {d} + \mathfrak {m} _ {2}) \exp (- T) + L ^ {2} T d ^ {1 / 2} h _ {\text {p r e d}} + L ^ {3 / 2} T d ^ {1 / 2} h _ {\text {c o r r}} ^ {\mathfrak {p}} + L ^ {1 / 2} T \varepsilon_ {\mathfrak {s c}} \\ \end{array}
$$
where  $\mathfrak{p} = \frac{1}{2}$  if we use the overdamped corrector and  $\mathfrak{p} = 1$  if we use the underdamped corrector. Applying the second part of Lemma 12 for the second stage of the algorithm, we then conclude that
$$
\mathsf {T V} (p _ {T - \delta}, q _ {T - \delta}) \lesssim (\sqrt {d} + \mathfrak {m} _ {2}) \exp (- T) + L ^ {2} T d ^ {1 / 2} h _ {\text {p r e d}} + L ^ {3 / 2} T d ^ {1 / 2} h _ {\text {c o r r}} ^ {\mathfrak {p}} + L ^ {1 / 2} T \varepsilon_ {\mathfrak {s c}}.
$$
Finally, we note that if we take  $\delta \asymp \frac{\varepsilon^2}{L^2(d\vee\mathfrak{m}_2^2)}$ , then by [LLT23, Lemma 6.4],  $\mathrm{TV}(q_{T - \delta},q_T)\leq \varepsilon$ ; a triangle inequality thus finishes the proof.
Remark 4. Alternatively, instead of taking geometrically decreasing step sizes and employing early stopping, we could split the algorithm into two stages: for time  $t < T - h_{\mathrm{pred}}$ , we take constant step size  $h_{\mathrm{pred}}$ , and for time  $t > T - h_{\mathrm{pred}}$ , we use a smaller constant step size  $h'$  as required if working with the original score perturbation lemma (see Remark 2).
F Numerical experiments
In this section, we provide preliminary numerical experiments to illustrate our theory. We implement DPUM on a toy model that is not log-concave (mixture of Gaussians).
Setup. The target distribution is a mixture of five Gaussians in dimension 5. On Figures 1 and 2, the red stars represent the means of the Gaussians and the red ellipses around the stars represent the variances of the Gaussians. We start by sampling 500 independent points (in blue) from a standard Gaussian. Then, we run DPUM from the blue dots over 300 iterations and plot the two first coordinates of the dots at iterations 0, 100, 200 and 300. This is a low-dimensional toy example so it does not illustrate our theory, rather we include it as a simple sanity check.
Parameters. We use a closed form formula for the score along the forward process. In other words, the score estimation error is equal to zero. The step size of the predictor is 0.01 and the step size of the corrector is 0.001. The corrector consists in 3 steps of the underdamped Langevin algorithm. In the latter algorithm, we initialize the velocity as a centered Gaussian random variable with standard deviation 0.001 and set the parameter  $\gamma$  to 0.01.
Observations. We observe the expected behavior: although the target distribution is highly non-log-concave, DPUM is able to provide samples from a distribution that is close to the target distribution. In particular, the initial Gaussian distribution splits in clusters that will fit each component of the target mixture of Gaussians. Recall that we experiment without score error but with discretization error: our numerical results illustrate the common wisdom that score knowledge along the forward process can replace convexity assumptions. In particular, we observe that even isolated, low probability components of the Gaussian mixture, are recovered by DPUM.
Figure 1: A realization of DPUM for a mixture of Gaussians.
Figure 2: A realization of DPUM for another mixture of Gaussians.