SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} u _ {t} (\theta) = \frac {- 1}{p _ {t} (\theta) \cdot \sigma_ {t} \cdot t} \left(g _ {0, \sigma_ {t} ^ {2}} * \mathcal {N} [ t \hat {\mu}, (t \hat {\sigma}) ^ {2} ]\right) (\theta) + \frac {\theta}{p _ {t} (\theta) \cdot t} \left(\mathcal {N} [ 0, \sigma_ {t} ^ {2} ] * \mathcal {N} [ t \hat {\mu}, (t \hat {\sigma}) ^ {2} ]\right) (\theta) \\ = \frac {- 1}{p _ {t} (\theta) \cdot \sigma_ {t} \cdot t} \frac {(\theta - t \hat {\mu}) \cdot \sigma_ {t} ^ {2}}{\sigma_ {t} ^ {2} + (t \hat {\sigma}) ^ {2}} \cdot \mathcal {N} [ t \hat {\mu}, (\sigma_ {t} ^ {2} + (t \hat {\sigma}) ^ {2}) ] (\theta) + \frac {\theta}{p _ {t} (\theta) \cdot t} \mathcal {N} [ t \hat {\mu}, (\sigma_ {t} ^ {2} + (t \hat {\sigma}) ^ {2}) ] (\theta) \\ = \frac {\left(\sigma_ {t} ^ {2} + (t \hat {\sigma}) ^ {2}\right) \theta - (\theta - t \hat {\mu}) \cdot \sigma_ {t}}{p _ {t} (\theta) \cdot t \cdot \left(\sigma_ {t} ^ {2} + (t \hat {\sigma}) ^ {2}\right)} \cdot p _ {t} (\theta) \\ = \frac {\left(\sigma_ {t} ^ {2} + (t \hat {\sigma}) ^ {2} - \sigma_ {t}\right) \theta + t \hat {\mu} \cdot \sigma_ {t}}{t \cdot \left(\sigma_ {t} ^ {2} + (t \hat {\sigma}) ^ {2}\right)}. \tag {17} \\ \end{array}
$$
By choosing a vector field  $v_{t}$  of the form (17) with learnable parameters  $\hat{\mu}, \hat{\sigma}^2$ , we can thus define a continuous flow that is restricted to a one-dimensional Gaussian.
B Mass covering properties of flows
In this supplement, we investigate the mass covering properties of continuous normalizing flows trained using mean squared error and in particular prove Theorem 1. We first recall the notation from the main part. We always assume that the data is distributed according to  $p_1(\theta)$ . In addition, there is a known and simple base distribution  $p_0$  and we assume that there is a vector field  $u_t: [0,1] \times \mathbb{R}^d \to \mathbb{R}^d$  that connects  $p_0$  and  $p_1$  in the following sense. We denote by  $\phi_t$  the flow generated by  $u_t$ , i.e.,  $\phi_t$  satisfies
$$
\partial_ {t} \phi_ {t} (\theta) = u _ {t} \left(\phi_ {t} (\theta)\right). \tag {18}
$$
Then we assume that  $(\phi_1)_*p_0 = p_1$  and we also define the interpolations  $p_t = (\phi_t)_*p_0$ .
We do not have access to the ground truth distributions  $p_t$  and the vector field  $u_t$  but we try to learn a vector field  $v_t$  approximating  $u_t$ . We denote its flow by  $\psi_t$  and we define  $q_t = (\psi_t)_*q_0$  and  $q_0 = p_0$ . We are interested in the mass covering properties of the learned approximation  $q_1$  of  $p_1$ . In particular, we want to relate the KL-divergence  $\mathrm{D}_{\mathrm{KL}}(p_1||q_1)$  to the mean squared error,
$$
\mathrm {M S E} _ {p} (u, v) = \int_ {0} ^ {1} \mathrm {d} t \int p _ {t} (\mathrm {d} \theta) \left(u _ {t} (\theta) - v _ {t} (\theta)\right) ^ {2}, \tag {19}
$$
of the generating vector fields. The first observation is that without any regularity assumptions on  $v_{t}$  it is impossible to obtain any bound on the KL-divergence in terms of the mean squared error.
Lemma 1. For every  $\varepsilon >0$  there are vector field  $u_{t}$  and  $v_{t}$  and a base distribution  $p_0 = q_0$  such that
$$
\mathrm {M S E} _ {p} (u, v) <   \varepsilon \text {a n d} \mathrm {D} _ {\mathrm {K L}} \left(p _ {1} \| q _ {1}\right) = \infty . \tag {20}
$$
In addition we can construct  $u_{t}$  and  $v_{t}$  such that the support of  $p_{1}$  is larger than the support of  $q_{1}$ .
Proof. We consider the uniform distribution  $p_0 = q_0 \sim \mathcal{U}([-1, 1])$  and the vector fields
$$
u _ {t} (\theta) = 0 \tag {21}
$$
$$
v _ {t} (\theta) = \left\{ \begin{array}{l l} \varepsilon & \text {f o r} 0 \leq \theta <   \varepsilon , \\ 0 & \text {o t h e r w i s e .} \end{array} \right. \tag {22}
$$
As before, let  $\phi_t$  denote the flow of the vector field  $u_t$  and similarly  $\psi_t$  denote the flow of  $v_t$ . Clearly  $\phi_t(\theta) = \theta$ . On the other hand
$$
\psi_ {t} (\theta) = \left\{ \begin{array}{l l} \min  (\theta + \varepsilon t, \varepsilon) & \text {i f} 0 \leq \theta <   \varepsilon , \\ \theta & \text {o t h e r w i s e .} \end{array} \right. \tag {23}
$$
$$
\psi_ {1} (\theta) = \left\{ \begin{array}{l l} \varepsilon & \text {i f} 0 \leq \theta <   \varepsilon , \\ \theta & \text {o t h e r w i s e .} \end{array} \right. \tag {24}
$$
This implies that  $p_1 = (\phi_1)_*p_0 \sim \mathcal{U}([-1,1])$ . On the other hand  $q_1 = (\psi_1)_*q_0$  has support in  $[-1,0] \cup [\varepsilon, 1]$ . In particular, the distribution of  $q_1$  is not mass covering with respect to  $p_1$  and  $\mathrm{D}_{\mathrm{KL}}(p_1||q_1) = \infty$ . Finally, we observe that the MSE can be arbitrarily small
$$
\mathrm {M S E} _ {p} (u, v) = \int_ {0} ^ {1} \mathrm {d} t \int p _ {t} (\mathrm {d} \theta) | u _ {t} (\theta) - v _ {t} (\theta) | ^ {2} = \int_ {0} ^ {1} \int_ {0} ^ {\varepsilon} \frac {1}{2} \varepsilon^ {2} = \frac {\varepsilon^ {3}}{2}. \tag {25}
$$
Here we used that the density of  $p_t(\mathrm{d}\theta)$  is  $1/2$  for  $-1 \leq \theta \leq 1$ .
We see that an arbitrary small MSE-loss cannot ensure that the probability distribution is mass covering and the KL-divergence is finite. On a high level this can be explained by the fact that for vector fields  $v_{t}$  that are not Lipschitz continuous the flow is not necessarily continuous, and we can generate holes in the distribution. Note that we chose  $p_0$  to be a uniform distribution for simplicity, but the result extends to any smooth distribution, in particular the result does not rely on the discontinuity of  $p_0$ .
Next, we investigate the mass covering property for Lipschitz continuous flows. When the flows  $u_{t}$  and  $v_{t}$  are Lipschitz continuous (in  $\theta$ ) this ensures that the flows  $\psi_{1}$  and  $\phi_{1}$  are continuous in  $x$  and it is not possible to create holes in the distribution as shown above for non-continuous vector fields. We show a weaker bound in this setting.
Lemma 2. For every  $0 \leq \delta \leq 1$  there is a base distribution  $p_0 = q_0$  and the are Lipschitz-continuous vector fields  $u_t$  and  $v_t$  such that  $\mathrm{MSE}_p(u, v) = \delta$  and
$$
\mathrm {D} _ {\mathrm {K L}} \left(p _ {1} \| q _ {1}\right) \geq \frac {1}{2} \mathrm {M S E} _ {p} (u, v) ^ {1 / 3}. \tag {26}
$$
Proof. We consider  $p_0, q_0$  and  $u_t$  as in Lemma 1, and we define
$$
v _ {t} (\theta) = \left\{ \begin{array}{l l} 2 \theta & \text {f o r} 0 \leq \theta <   \varepsilon , \\ 2 \varepsilon - \theta & \text {f o r} \varepsilon \leq \theta <   2 \varepsilon , \\ 0 & \text {o t h e r w i s e .} \end{array} \right. \tag {27}
$$
Then we can calculate for  $0 \leq \theta \leq e^{-2}\varepsilon$  that
$$
\psi_ {t} (\theta) = \theta e ^ {2 t}. \tag {28}
$$
Similarly we obtain for  $\varepsilon \leq \theta \leq 2\varepsilon$  (solving the ODE  $f' = 2f$ )
$$
\psi_ {t} (\theta) = 2 \varepsilon - (2 \varepsilon - \theta) e ^ {- 2 t}. \tag {29}
$$
We find
$$
\psi_ {1} (0) = 0, \psi_ {1} \left(e ^ {- 2} \varepsilon\right) = \varepsilon , \psi_ {1} (\varepsilon) = 2 - \varepsilon e ^ {- 2}; \psi_ {2} (2 \varepsilon) = 2 \varepsilon . \tag {30}
$$
Next we find for the densities of  $q_{1}$  that
$$
q _ {1} \left(\psi_ {1} (\theta)\right) = q _ {0} (\theta) \left| \psi_ {1} ^ {\prime} (\theta) \right| ^ {- 1} = \frac {1}{2} \left\{ \begin{array}{l l} e ^ {- 2} & \text {f o r} 0 \leq \theta \leq e ^ {- 2} \varepsilon , \\ e ^ {2} & \text {f o r} \varepsilon \leq \theta \leq 2 \varepsilon . \end{array} \right. \tag {31}
$$
Together with (30) this implies that the density of  $q_{1}$  is given by
$$
q _ {1} (\theta) = \frac {1}{2} \left\{ \begin{array}{l l} e ^ {- 2} & \text {f o r} 0 \leq \theta \leq \varepsilon , \\ e ^ {2} & \text {f o r} 2 \varepsilon - \varepsilon e ^ {- 2} \leq \theta \leq 2 \varepsilon . \end{array} \right. \tag {32}
$$
Note that  $p_1(\theta) = 1 / 2$  for  $-1\leq \theta \leq 1$  and therefore
$$
\int_ {0} ^ {\varepsilon} \ln \frac {p _ {1} (\theta)}{q _ {1} (\theta)} p _ {1} (\mathrm {d} \theta) = \int_ {0} ^ {\varepsilon} \ln \left(e ^ {2}\right) \frac {1}{2} \mathrm {d} \theta = \varepsilon , \tag {33}
$$
$$
\int_ {2 \varepsilon - \varepsilon e ^ {- 2}} ^ {2 \varepsilon} \ln \frac {p _ {1} (\theta)}{q _ {1} (\theta)} p _ {1} (\mathrm {d} \theta) = \int_ {2 \varepsilon - \varepsilon e ^ {- 2}} ^ {2 \varepsilon} \ln (e ^ {- 2}) \frac {1}{2} \mathrm {d} \theta = - \varepsilon e ^ {- 2}. \tag {34}
$$
Moreover we note
$$
\int_ {\varepsilon} ^ {2 \varepsilon - \varepsilon e ^ {- 2}} q _ {1} (\mathrm {d} \varepsilon) = \int_ {e ^ {- 2} \varepsilon} ^ {\varepsilon} q _ {0} (\mathrm {d} \varepsilon) = \frac {1}{2} \varepsilon \left(1 - e ^ {- 2}\right) = \int_ {\varepsilon} ^ {2 \varepsilon - \varepsilon e ^ {- 2}} p _ {1} (\mathrm {d} \varepsilon), \tag {35}
$$
which implies (by positivity of the KL-divergence) that
$$
\int_ {\varepsilon} ^ {2 \varepsilon - \varepsilon e ^ {- 2}} \ln \left(\frac {p _ {1} (\theta)}{q _ {1} (\theta)}\right) p _ {1} (\mathrm {d} \theta) \geq 0. \tag {36}
$$
We infer using also  $p_1(\theta) = q_1(\theta) = 1 / 2$  for  $\theta \in [-1,0]\cap [2\varepsilon ,1]$  that
$$
\mathrm {D} _ {\mathrm {K L}} \left(p _ {1} \mid \mid q _ {1}\right) = \int \ln \left(\frac {p _ {1} (\theta)}{q _ {1} (\theta)}\right) p _ {1} (\mathrm {d} \theta) \geq \varepsilon \left(1 - e ^ {- 2}\right). \tag {37}
$$
On the other hand we can bound
$$
\int_ {0} ^ {1} \mathrm {d} t \int p _ {t} (\mathrm {d} \theta) | v _ {t} (\theta) - u _ {t} (\theta) | ^ {2} = \frac {1}{2} \int_ {0} ^ {1} \mathrm {d} t \int_ {0} ^ {2 \varepsilon} | u _ {t} (\theta) | ^ {2} = \int_ {0} ^ {\varepsilon} s ^ {2} \mathrm {d} s = \frac {\varepsilon^ {3}}{3}. \tag {38}
$$
$$
\mathrm {D} _ {\mathrm {K L}} \left(p _ {1} \| q _ {1}\right) \geq \frac {1}{2} \left(\mathrm {M S E} _ {p} (u, v)\right) ^ {1 / 3}. \tag {39}
$$
In particular, it is not possible to bound the KL-divergence by the MSE even when the vector fields are Lipschitz continuous.
Let us put this into context. It was already shown in [42] that we can, in general, not bound the forward KL-divergence by the mean squared error and our Lemmas 1 and 2 are concrete examples. On the other hand, when considering SDEs the KL-divergence can be bounded by the mean squared error of the drift terms as shown in [43]. Indeed, in [42] the favorable smoothing effect was carefully investigated.
Here we show that we can alternatively obtain an upper bound on the KL-divergence when assuming that  $u_{t}$ ,  $v_{t}$ , and  $p_0$  satisfy additional regularity assumptions. This allows us to recover the mass covering property from bounds on the means squared error for sufficiently smooth vector fields. The scaling is nevertheless still weaker than for SDEs.
We now state our assumptions. We denote the gradient with respect to  $\theta$  by  $\nabla = \nabla_{\mu}$  and second derivatives by  $\nabla^2 = \nabla_{\mu \nu}^2$ . When applying the chain rule, we leave the indices implicit. We denote by  $|\cdot|$  the Frobenius norm  $|A| = \left(\sum_{ij} A_{ij}^2\right)^{1/2}$  of a matrix. The Frobenius norm is submultiplicative, i.e.,  $|AB| \leq |A| \cdot |B|$  and directly generalizes to higher order tensors.
Assumption 1. We assume that
$$
\left| \nabla u _ {t} \right| \leq L, \left| \nabla v _ {t} \right| \leq L, \left| \nabla^ {2} u _ {t} \right| \leq L ^ {\prime}, \left| \nabla^ {2} v _ {t} \right| \leq L ^ {\prime}. \tag {40}
$$
We require one further assumption on  $p_0$ .
Assumption 2. There is a constant  $C_1$  such that
$$
\left| \nabla \ln p _ {0} (\theta) \right| \leq C _ {1} (1 + | \theta |). \tag {41}
$$
We also assume that
$$
\mathbb {E} _ {p _ {0}} | \theta | ^ {2} <   C _ {2} <   \infty . \tag {42}
$$
Note that (41) holds, e.g., if  $p_0$  follows a Gaussian distribution but also for smooth distribution with slower decay at  $\infty$ . If we assume that  $|\nabla \ln p_0(\theta)|$  is bounded the proof below simplifies slightly. This is, e.g., the case if  $p_0(\theta) \sim e^{-|\theta|}$  as  $|\theta| \to \infty$ .
We need some additional notation. It is convenient to introduce  $\phi_t^s = \phi_t \circ (\phi_s)^{-1}$ , i.e., the flow from time  $s$  to  $t$  (in particular  $\phi_t^0 = \phi_t$ ) and similarly for  $\psi$ . We can now restate and prove Theorem 1.
Theorem 2. Let  $p_0 = q_0$  and assume  $u_t$  and  $v_t$  are two vector fields whose flows satisfy  $p_1 = (\phi_1)_*p_0$  and  $q_1 = (\psi_1)_*q_0$ . Assume that  $p_0$  satisfies Assumption 2 and  $u_t$  and  $v_t$  satisfy Assumption 1. Then there is a constant  $C > 0$  depending on  $L, L', C_1, C_2$ , and  $d$  such that (for  $\mathrm{MSE}_p(u,v) < 1$ )
$$
\mathrm {D} _ {\mathrm {K L}} \left(p _ {1} \mid \mid q _ {1}\right) \leq C \operatorname {M S E} _ {p} (u, v) ^ {\frac {1}{2}}. \tag {43}
$$
Remark 1. We do not claim that our results are optimal, it might be possible to find similar bounds for the forward KL-divergence with weaker assumptions. However, we emphasize that Lemma 2 shows that the result of the theorem is not true without the assumption on the second derivative of  $v_{t}$  and  $u_{t}$ .
Proof. We want to control  $\mathrm{D}_{\mathrm{KL}}(p_1||q_1)$ . It can be shown that (see equation above (25) in [43] or Lemma 2.19 in [42])
$$
\partial_ {t} \mathrm {D} _ {\mathrm {K L}} \left(p _ {t} \| q _ {t}\right) = - \int p _ {t} (\mathrm {d} \theta) \left(u _ {t} (\theta) - v _ {t} (\theta)\right) \cdot \left(\nabla \ln p _ {t} (\theta) - \nabla \ln q _ {t} (\theta)\right). \tag {44}
$$
Using Cauchy-Schwarz we can bound this by
$$
\partial_ {t} \mathrm {D} _ {\mathrm {K L}} \left(p _ {t} \mid \mid q _ {t}\right) \leq \left(\int p _ {t} (\mathrm {d} \theta) \left| u _ {t} (\theta) - v _ {t} (\theta) \right| ^ {2}\right) ^ {\frac {1}{2}} \left(\int p _ {t} (\mathrm {d} \theta) \left| \nabla \ln p _ {t} (\theta) - \nabla \ln q _ {t} (\theta) \right| ^ {2}\right) ^ {\frac {1}{2}}. \tag {45}
$$
We use the relation (see (3))
$$
\ln \left( \right.p _ {t} \left(\phi_ {t} \left(\theta_ {0}\right)\right) = \ln \left(p _ {0} \left(\theta_ {0}\right)\right) - \int_ {0} ^ {t} (\operatorname {d i v} u _ {s}) \left(\phi_ {s} \left(\theta_ {0}\right)\right) \mathrm {d} s, \tag {46}
$$
which can be equivalently rewritten (setting  $\theta = \phi_t\theta_0$ ) as
$$
\ln \left(p _ {t} (\theta)\right) = \ln \left(p _ {0} \left(\phi_ {0} ^ {t} \theta\right)\right) - \int_ {0} ^ {t} \left(\operatorname {d i v} u _ {s}\right) \left(\phi_ {s} ^ {t} \theta\right) \mathrm {d} s. \tag {47}
$$
We use the following relation for  $\nabla \phi_s^t$
$$
\nabla \phi_ {s} ^ {t} (\theta) = \exp \left(\int_ {t} ^ {s} \mathrm {d} \tau \left(\nabla u _ {\tau}\right) \left(\phi_ {\tau} ^ {t} (\theta)\right)\right). \tag {48}
$$
This relation is standard and can be directly deduced from the following ODE for  $\nabla \phi_s^t$
$$
\partial_ {s} \nabla \phi_ {s} ^ {t} (\theta) = \nabla \partial_ {s} \phi_ {s} ^ {t} (\theta) = \nabla \left(u _ {s} \left(\phi_ {s} ^ {t} (\theta)\right)\right) = \left(\left(\nabla u _ {s}\right) \left(\phi_ {s} ^ {t} (\theta)\right)\right) \cdot \nabla \phi_ {s} ^ {t} (\theta). \tag {49}
$$
We can conclude that for  $0 \leq s, t \leq 1$  the bound
$$
\left| \nabla \phi_ {s} ^ {t} (\theta) \right| \leq e ^ {L} \tag {50}
$$
holds. We find
$$
\begin{array}{l} \left| \nabla \ln \left(p _ {t} (\theta)\right) \right| = \left| \nabla \ln \left(p _ {0}\right) \left(\phi_ {0} ^ {t} \theta\right) \cdot \nabla \phi_ {0} ^ {t} (\theta) - \int_ {0} ^ {t} \left(\nabla \operatorname {d i v} u _ {s}\right) \left(\phi_ {s} ^ {t} \theta\right) \cdot \nabla \phi_ {s} ^ {t} (\theta) d s \right| \tag {51} \\ \leq | \nabla \ln (p _ {0}) (\phi_ {0} ^ {t} \theta) | e ^ {L} + L ^ {\prime} e ^ {L}, \\ \end{array}
$$
and a similar bound holds for  $q_{t}$ . In words, we have shown that the score of  $p_t$  at  $\theta$  can be bounded by the score of  $p_0$  of theta transported along the vector field  $u_{t}$  minus a correction which quantifies the change of score along the path. We now bound using the definition  $p_t = (\phi_t)_*p_0$  and the assumption (41)
$$
\begin{array}{l} \int p _ {t} (\mathrm {d} \theta) | \nabla \ln p _ {0} \left(\phi_ {0} ^ {t} (\theta)\right) | ^ {2} = \int p _ {0} (\mathrm {d} \theta_ {0}) | \nabla \ln p _ {0} \left(\phi_ {0} ^ {t} \phi_ {t} \left(\theta_ {0}\right)\right) | ^ {2} = \mathbb {E} _ {p _ {0}} | \nabla \ln p _ {0} \left(\theta_ {0}\right) | ^ {2} \tag {52} \\ \leq \mathbb {E} _ {p _ {0}} \left(C _ {1} \left(1 + | \theta_ {0} |\right) ^ {2}\right) \leq 2 C _ {1} ^ {2} \left(1 + \mathbb {E} _ {p _ {0}} | \theta_ {0} | ^ {2}\right) \leq 2 C _ {1} ^ {2} \left(1 + C _ {2} ^ {2}\right). \\ \end{array}
$$
Similarly we obtain using  $q_0 = p_0$
$$
\int p _ {t} (\mathrm {d} \theta) | \nabla \ln q _ {0} \left(\psi_ {0} ^ {t} \theta\right) | ^ {2} = \int p _ {0} (\mathrm {d} \theta_ {0}) | \nabla \ln q _ {0} \left(\psi_ {0} ^ {t} \phi_ {t} \theta_ {0}\right) | ^ {2}. \tag {53}
$$
In words, to control the score of  $q$  integrated with respect to  $p_t$  we need to control the distortion we obtain when moving forward with  $u$  and backwards with  $v$ . We investigate  $\psi_0^t\phi_t(\theta_0)$ . We now show
$$
\left. \partial_ {h} \psi_ {t} ^ {t + h} \phi_ {t + h} ^ {t} (\theta) \right| _ {h = 0} = u _ {t} (\theta) - v _ {t} (\theta). \tag {54}
$$
First, by definition of  $\phi$ , we find
$$
\left. \partial_ {h} \phi_ {t + h} ^ {t} (\theta) \right| _ {h = 0} = \partial_ {h} \phi_ {t + h} \phi_ {t} ^ {- 1} (\theta) | _ {h = 0} = u _ {t} \left(\phi_ {t} \phi_ {t} ^ {- 1} (\theta)\right) = u _ {t} (\theta). \tag {55}
$$
To evaluate the second contribution we observe
$$
\begin{array}{l} 0 = \partial_ {h} \theta | _ {h = 0} = \partial_ {h} \psi_ {t + h} ^ {t + h} (\theta) | _ {h = 0} = \partial_ {h} \psi_ {t + h} \psi_ {t + h} ^ {- 1} (\theta) | _ {h = 0} \\ = \left(\partial_ {h} \psi_ {t + h}\right) \psi_ {t} ^ {- 1} (\theta) | _ {h = 0} + \psi_ {t} \left(\partial_ {h} \psi_ {t + h} ^ {- 1}\right) (\theta) | _ {h = 0} = v _ {t} \left(\psi_ {t} \psi_ {t} ^ {- 1} (\theta)\right) + \partial_ {h} \psi_ {t} \psi_ {t + h} ^ {- 1} (\theta) | _ {h = 0} \tag {56} \\ = v _ {t} (\theta) + \partial_ {h} \psi_ {t} ^ {t + h} (\theta) | _ {h = 0} \\ \end{array}
$$
Now (54) follows from (55) and (56) together with  $\phi_t^t = \psi_t^t = \mathrm{Id}$ . Using (54) we find
$$
\partial_ {t} \left(\psi_ {0} ^ {t} \phi_ {t}\right) \left(\theta_ {0}\right) = \partial_ {h} \left(\psi_ {0} ^ {t} \psi_ {t} ^ {t + h} \phi_ {t + h} ^ {t} \phi_ {t}\right) \left(\theta_ {0}\right) | _ {h = 0} = \left(\nabla \psi_ {0} ^ {t}\right) \left(\phi_ {t} \left(\theta_ {0}\right)\right) \cdot \left(\left(u _ {t} - v _ {t}\right) \left(\phi_ {t} \left(\theta_ {0}\right)\right)\right). \tag {57}
$$
Using (50) we conclude that
$$
\begin{array}{l} \left| \psi_ {0} ^ {t} \phi_ {t} \left(\theta_ {0}\right) - \theta_ {0} \right| \leq \left| \int_ {0} ^ {t} \partial_ {s} \psi_ {0} ^ {s} \phi_ {s} \left(\theta_ {0}\right) \mathrm {d} s \right| \leq \int_ {0} ^ {t} | (\nabla \psi_ {0} ^ {s}) \left(\phi_ {s} \left(\theta_ {0}\right)\right) | \cdot | u _ {s} - v _ {s} | \left(\phi_ {s} \left(\theta_ {0}\right)\right) \mathrm {d} s \tag {58} \\ \leq e ^ {L} \int_ {0} ^ {t} | u _ {s} - v _ {s} | (\phi_ {s} (\theta_ {0})) d s. \\ \end{array}
$$
We use this and the assumption (41) to continue to estimate (53) as follows
$$
\begin{array}{l} \int p _ {t} (\mathrm {d} \theta) | \nabla \ln q _ {0} (\psi_ {0} ^ {t} \theta) | ^ {2} = \int p _ {0} (\mathrm {d} \theta_ {0}) | \nabla \ln q _ {0} (\psi_ {0} ^ {t} \phi_ {t} (\theta_ {0})) | ^ {2} \\ \leq C _ {1} ^ {2} \int p _ {0} (\mathrm {d} \theta_ {0}) (1 + | \psi_ {0} ^ {t} \phi_ {t} (\theta_ {0}) |) ^ {2} \\ \leq C _ {1} ^ {2} \int p _ {0} (\mathrm {d} \theta_ {0}) \left(1 + | \psi_ {0} ^ {t} \phi_ {t} (\theta_ {0}) - \theta_ {0} | + | \theta_ {0} |\right) ^ {2} \\ \leq 3 C _ {1} ^ {2} + 3 C _ {1} ^ {2} \int p _ {0} (\mathrm {d} \theta_ {0}) \left(| \psi_ {0} ^ {t} \phi_ {t} (\theta_ {0}) - \theta_ {0} | ^ {2} + | \theta_ {0} | ^ {2}\right) \\ \leq 3 C _ {1} ^ {2} \left(1 + \mathbb {E} _ {p _ {0}} | \theta_ {0} | ^ {2}\right) + 3 C _ {1} ^ {2} e ^ {2 L} \int p _ {0} (\mathrm {d} \theta_ {0}) \left(\int_ {0} ^ {t} \mathrm {d} s | u _ {s} - v _ {s} | \left(\phi_ {s} \left(\theta_ {0}\right)\right)\right) ^ {2}. \tag {59} \\ \end{array}
$$
Here we used  $(a + b + c)^2 \leq 3(a^2 + b^2 + c^2)$  in the second to last step. We bound the remaining integral using Cauchy-Schwarz as follows
$$
\begin{array}{l} \int p _ {0} (\mathrm {d} \theta_ {0}) \left(\int_ {0} ^ {t} | u _ {s} - v _ {s} | (\phi_ {s} (\theta_ {0}))\right) ^ {2} \leq \int p _ {0} (\mathrm {d} \theta_ {0}) \left(\int_ {0} ^ {t} \mathrm {d} s | u _ {s} - v _ {s} | ^ {2} (\phi_ {s} (\theta_ {0}))\right) \left(\int_ {0} ^ {t} \mathrm {d} s 1 ^ {2}\right) \\ \leq t \int_ {0} ^ {t} d s \int p _ {0} (d \theta_ {0}) | u _ {s} - v _ {s} | ^ {2} (\phi_ {s} (\theta_ {0})) \\ = t \int_ {0} ^ {t} d s \int p _ {s} (d \theta_ {s}) | u _ {s} - v _ {s} | ^ {2} (\theta_ {s}) \\ \leq \int_ {0} ^ {1} \mathrm {d} s \int p _ {s} (\mathrm {d} \theta_ {s}) | u _ {s} - v _ {s} | ^ {2} (\theta_ {s}) = \operatorname {M S E} _ {p} (u, v). \tag {60} \\ \end{array}
$$
The last displays together imply
$$
\int p _ {t} (\mathrm {d} \theta) | \nabla \ln q _ {0} \left(\psi_ {0} ^ {t} \theta\right) | ^ {2} \leq 3 C _ {1} ^ {2} \left(1 + \mathbb {E} _ {p _ {0}} \left| \theta_ {0} \right| ^ {2} + e ^ {2 L} \operatorname {M S E} _ {p} (u, v)\right). \tag {61}
$$
Now we have all the necessary ingredients to bound the derivative of the KL-divergence. We control the second integral in (45) using (51) (and again  $\left(\sum_{i=1}^{4} a_i\right)^2 \leq 4 \sum a_i^2$ ) as follows,
$$
\begin{array}{l} \int p _ {t} (\mathrm {d} \theta) | \nabla \ln p _ {t} (\theta) - \nabla \ln q _ {t} (\theta) | ^ {2} \tag {62} \\ \leq 2 \cdot 2 ^ {2} \cdot L ^ {\prime 2} e ^ {2 L} + 4 e ^ {2 L} \int p _ {t} (\mathrm {d} \theta) \left(| \nabla \ln q _ {0} (\psi_ {0} ^ {t}) \theta | ^ {2} + | \nabla \ln p _ {0} (\phi_ {0} ^ {t}) \theta) | ^ {2}\right). \\ \end{array}
$$
Using (52) and (61) we finally obtain
$$
\begin{array}{l} \int p _ {t} (\mathrm {d} \theta) | \nabla \ln p _ {t} (\theta) - \nabla \ln q _ {t} (\theta) | ^ {2} \leq 8 \cdot L ^ {\prime 2} e ^ {2 L} + C _ {1} ^ {2} e ^ {2 L} \left(2 0 \left(1 + C _ {2} ^ {2}\right) + 1 2 \operatorname {M S E} _ {p} (u, v)\right) \tag {63} \\ \leq C (1 + \mathrm {M S E} _ {p} (u, v)) \\ \end{array}
$$
for some constant  $C > 0$ . Finally, we obtain
$$
\begin{array}{l} \mathrm {D} _ {\mathrm {K L}} \left(p _ {1} \right\| q _ {1}) = \int_ {0} ^ {1} \mathrm {d} t \partial_ {t} \mathrm {D} _ {\mathrm {K L}} \left(p _ {t} \right\| q _ {t}) \\ \leq \left(C \left(1 + \operatorname {M S E} _ {p} (u, v)\right)\right) ^ {\frac {1}{2}} \int_ {0} ^ {1} \mathrm {d} t \left(\int p _ {t} (\mathrm {d} \theta) \left| u _ {t} (\theta) - v _ {t} (\theta) \right| ^ {2}\right) ^ {\frac {1}{2}} \tag {64} \\ \leq \left(C \left(1 + \mathrm {M S E} _ {p} (u, v)\right)\right) ^ {\frac {1}{2}} \left(\int_ {0} ^ {1} \mathrm {d} t \int p _ {t} (\mathrm {d} \theta) | u _ {t} (\theta) - v _ {t} (\theta) | ^ {2}\right) ^ {\frac {1}{2}} \\ \leq \left(C \left(1 + \mathrm {M S E} _ {p} (u, v)\right)\right) ^ {\frac {1}{2}} \mathrm {M S E} _ {p} (u, v) ^ {\frac {1}{2}}. \\ \end{array}
$$
C SBI Benchmark
In this section, we collect missing details and additional results for the analysis of the SBI benchmark in Section 4.
C.1 Network architecture and hyperparameters
For each task and simulation budget in the benchmark, we perform a mild hyperparameter optimization. We sweep over the batch size and learning rate (which is particularly important as the simulation budgets differ by orders of magnitudes), the network size and the  $\alpha$  parameter for the time prior defined in Section 3.3 (see Tab. 2 for the specific values). We reserve  $5\%$  of the simulation budget for validation and choose the model with the best validation loss across all configurations.
C.2 Additional results