SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We implement DFE following the algorithm outlined in Boffi and Vanden-Eijnden [2023]. We use 5000 particles. For score estimation, we use the same network architecture as in NODE. At each time step, we optimize the score network 100 steps. We found the result of DFE depends tremendously on the time step size  $\Delta t$ . For the OU process experiment in dimension 60, when  $\Delta t = 0.1, 0.01, 0.001, 0.0001$ , the resulting Bures-Wasserstein distance at the final time to the target measure is 28.11, 0.31, 0.46, 9.21 respectively. Surprisingly, a smaller  $\Delta t$  can result in bigger errors. We choose  $\Delta t = 0.01$  since it gives the best results.
C.6 Evaluation metrics
For all our experiments, calculations of all metrics are repeated 20 times on 1000 samples from each distribution. Our plots show both the average and the standard deviation calculated over these 20 repetitions.
When estimating symmetric KL divergence using samples, due to the finite sample size and the numerical error in estimating the log density, the estimated divergence can be very close to zero or even negative (when this occurs we take absolute values). This explains why the standard deviation regions touch the  $x$ -axis in the log-scale plots in Figure 3.
To compute Bures-Wasserstein distance [Kroshnin et al., 2021], we first fit a Gaussian to the samples of either distribution and then compute the closed-form Wasserstein-2 distance between the two Gaussians.
For the porous medium equation (Section 4.3), the total variation distance is used in Figure 4 and Figure 8 to compare the estimated and ground-truth solutions. It is approximated by the  $L_{1}$  distance between the densities calculated over 50000 samples uniformly distributed on the compact  $[-1.25x_{\mathrm{max}}, 1.25x_{\mathrm{max}}]$  with  $x_{\mathrm{max}} = C / \left(\beta(t + t_0)^{\frac{-2\alpha}{d}}\right)$  being the bound of the support of  $p_t^*$ .
D.1 Ornstein-Uhlenbeck process
The OU process is the Wasserstein gradient flow of the KL divergence with respect to a Gaussian  $\mu^{*} = \mathcal{N}(\beta, \Gamma^{-1})$  where  $\beta \in \mathbf{R}^d$  and  $\Gamma$  is a  $d \times d$  positive-definite matrix. When the initial distribution is  $\mu_0^* = \mathcal{N}(0, I_d)$ , the gradient flow at time  $t$  is known to be a Gaussian distribution  $G(t)$  with mean  $(I_d - e^{-t\Gamma})\beta$  and covariance  $\Gamma^{-1}(I_d - e^{-2t\Gamma}) + e^{-2t\Gamma}$ . We set the total time  $T = 2$ .
D.2 Porous medium equation
This flow has as closed-form solution given by the Barenblatt profile Vázquez [2007] when initialized accordingly:
$$
p _ {t} ^ {*} (x) = \left(t + t _ {0}\right) ^ {- \alpha} \left(C - \beta \| x \| ^ {2} \left(t + t _ {0}\right) ^ {\frac {- 2 \alpha}{d}}\right) _ {+} ^ {\frac {1}{m - 1}}, \tag {16}
$$
where  $t_0 > 0$  is the starting time,  $\alpha = \frac{m}{d(m - 1) + 2}$ ,  $\beta = \frac{(m - 1)\alpha}{2dm}$ , and  $C > 0$  is a free constant. Similar to Fan et al. [2021], we choose  $m = 2$  and total time  $T = 0.025$ . The initial measure follows a Barenblatt distribution supported in  $[-0.25, 0.25]^d$  ( $C$  is chosen accordingly) with  $t_0 = 10^{-3}$ . We use Metropolis-Hastings to sample from  $\mu_0$ .
Figure 8: Metrics (TV, Symmetric  $f$ -divergence and Wasserstein-2 distance) across time for dimensions 3 and 6 between the estimated  $\mu_t$  and the ground-truth  $\mu_t^*$  when solving the Porous Medium Equation.
D.3 Time-Dependant Fokker-Planck equation
We consider a time-dependent Fokker-Planck equation of the form (4) with the velocity field
$$
f _ {t} (X, \mu_ {t}) = \Gamma_ {t} (X - \beta_ {t}) - D _ {t} \nabla \log p _ {t} (X). \tag {17}
$$
When the initial measure  $p_0$  is Gaussian, the solution  $\mu_t$  can again be shown to be Gaussian with mean  $m_t$  and covariance  $\Sigma_t$  solutions of the differential equations:
$$
\left\{ \begin{array}{l} m _ {t} ^ {\prime} = - \Gamma_ {t} \left(m _ {t} - \beta_ {t}\right) \\ \Sigma_ {t} ^ {\prime} = - \Gamma_ {t} \Sigma_ {t} - \Sigma_ {t} \Gamma_ {t} ^ {\top} + 2 D _ {t}. \end{array} \right. \tag {18}
$$
In practice, we experiment with constant  $\Gamma_t = \mathrm{diag}(1,3)$  and  $D_{t} = \sigma^{2}I_{d}$ . We also experience in dimension 3 by considering and  $\Gamma_t = \mathrm{diag}(1,3,1)$ . We set  $a = 3$ ,  $\omega = 1$ ,  $\sigma = \sqrt{0.25}$  and pick as initial distribution  $p_0$  a Gaussian with mean  $b_{0}$  and covariance  $\sigma^2 I_d$ . We set the total time to  $T = 10$ .
We plot in Figure 9, for dimension 2, snapshots at different time steps of particles following the flow given by our method with TIPF parametrization. We only show SCVM-TIPF because SCVM-NODE gives visually indistinguishable trajectories. We also plot in Figure 10 the evolution of particles simulated by Euler-Maruyama (EM-SDE) discretization of the Fokker-Planck equation. Corresponding animated GIFs be found at this link.
Figure 9: Evolution of particles (in blue) following the flow learned with SCVM-TIPF for the time-dependent OU process (Section 4.4). In red is the moving attraction trap.
Figure 10: Evolution of particles (in blue) obtained by SDE-EM discretization for the time-dependent OU process (Section 4.4). In red is the moving attraction trap.
D.4 Flock of birds
We model the dynamics of a flock of birds by augmenting the time-dependent Fokker-Planck equation (11) with an interaction term:
$$
f _ {t} (X, \mu_ {t}) = \Gamma_ {t} (\beta_ {t} - X) + \alpha_ {t} (X - \mathbf {E} [ \mu_ {t} ]) - D \nabla \log p _ {t} (X).
$$
This is similar to the harmonically interacting particles experiment in Boffi and Vanden-Eijnden [2023], but we use a population expectation  $\mathbf{E}[\mu_t]$  instead of an empirical one in modeling the repulsion from the mean. Since  $f_{t}$  needs to access  $\mathbf{E}[\mu_t]$ , the resulting PDE is not a Fokker-Planck equation (4) and hence not solvable using the method in Boffi and Vanden-Eijnden [2023] but can be solved with our method by estimating  $\mathbf{E}[\mu_t]$  using Monte Carlo samples from  $\mu_t$ . We use a similar setup as in Section 4.4, except we now use an "infinity sign" attraction  $\beta_{t} = a(\cos (2\pi \omega t),0.5\sin (2\pi \omega t))$  along with a sinusoidal  $\alpha_{t} = 2\sin (\pi wt)$ . Depending on the sign of  $\alpha_{t}$ , particles are periodically attracted towards or repulsed from their mean. Both SCVM-TIPF and SCVM-NODE produce similar visual results as shown in Figure 11 and Figure 12.
We use a constant  $\Gamma_t = I_d$  and a constant diffusion matrix  $D = \sigma^2 I_d$ . We set  $a = 3$ ,  $\omega = 0.5$ , and  $\sigma = \sqrt{0.25}$ . We pick as initial distribution  $p_0$  a Gaussian with mean  $(0,0)$  and covariance  $\sigma^2 I_d$ . We set the total time to  $T = 10$ .
We respectively show in Figure 11 and Figure 12 simulations of particles following the flow learned with SCVM-TIPF and SCVM-NODE. Corresponding animated GIFs be found at this link.
Figure 11: Evolution of particles following the flow trained with TIPF parametrization on the flock of birds PDE (Section 4.5). In red shows the moving attraction mean.
Figure 12: Evolution of particles following the flow trained with NODE parametrization on the flock of birds PDE (Section 4.5). In red shows the moving attraction mean.
D.5 Flow splashing against obstacles
We use the following formulation for modeling the flow. Each obstacle is modeled as a line segment. The endpoints of the three obstacles are:
$$
((0, 3), (3, 0. 5)), ((1, 0), (1. 5, 0)), ((- 2, - 4), (6, 0)).
$$
We model the dynamics as a Fokker-Planck equation where  $f_{t}$  of the form (4) is defined as
$$
b _ {t} (x) = \left(q _ {\text {s i n k}} - x\right) + 2 0 \sum_ {i = 1} ^ {3} \frac {x - \pi_ {O _ {i}} (x)}{\| x - \pi_ {O _ {i}} (x) \|} p _ {\mathcal {N} (0, 0. 0 4)} \left(\| x - \pi_ {O _ {i}} (x) \|\right),
$$
$$
D _ {t} (x) = I _ {2},
$$
where  $q_{\mathrm{sink}} = (4,0)$ , and  $\pi_{O_i}(x)$  is the projection of  $x$  onto obstacle  $i$  represented as a line segment, and  $p_{\mathcal{N}(0,0.04)}$  is the density of an 1-dimensional Gaussian with variance 0.04.
The initial distribution is chosen to be  $\mathcal{N}(0, 0.25I_2)$ . We train SCVM-NODE for  $10^4$  with an initial learning rate of  $10^{-4}$ . Training takes 5.4 minutes. The time step size for SDE-EM used to produce Figure 14 is 0.005. Corresponding animated GIFs be found at this link.
Figure 13: Trajectory of 200 random particles across time using the same setup as in Figure 6. Left are sample paths obtained by our method, and right are sample paths obtained by DFE [Boffi and Vanden-Eijnden, 2023].
Figure 14: Same setup as in Figure 6 but with SDE-EM. We see the paths of the particles are not continuous. Moreover, the particles spill over the obstacle on the bottom right due to a finite time step size. In comparison, SCVM-NODE does not have such a problem.
D.6 Smooth interpolation of measures
Suppose we are to smoothly interpolate  $M$  measures  $\nu_{1},\ldots ,\nu_{M}$  with densities  $q_{1},\ldots ,q_{M}$ , and we want the flow to approximate  $\nu_{i}$  at time  $r_i$ . To achieve this goal, we present two formulations using different choices of  $f_{t}$ . We use SCVM-NODE in this section.
Measure interpolation using time-dependent Fokker-Planck equations. We model the dynamics as a Fokker-Planck equation where  $f_{t}$  of the form (4) is taken to be
$$
b _ {t} (x) = \sum_ {i = 1} ^ {M} \phi (t - r _ {i}) (\nabla \log q _ {i} (x) - \nabla \log p _ {t} (x))
$$
where  $\phi (t)$  is defined as the continuous bump function
$$
\phi (t) = \left\{ \begin{array}{l l} 1. 0 & | t | <   0. 5 h \\ (0. 6 h - | t |) / (0. 1 h) & | t | <   0. 6 h \\ 0. 0 & \text {o t h e r w i s e}, \end{array} \right.
$$
for bandwidth  $h = 1.0$ .
Below we provide details for the MNIST interpolation result in the top row of Figure 7. We use the first three images of 1, 2, 3 from the MNIST dataset. To construct  $\nu_{i}$  from a digit image, we use a mixture of Gaussians where we put one equally-weighted Gaussian with covariance  $0.02^{2}I_{2}$  on the pixels with values greater than 0.5 (images are first normalized to have values in [0, 1]). The initial distribution is  $\mathcal{N}((0.5, 0.5), 0.04I_{2})$ . To train SCVM-NODE, we use an initial learning rate of  $10^{-4}$  with cosine decay for a total of  $10^{5}$  iterations. This takes roughly 1 hour to train.
Measure interpolation using optimal transport. An alternative way to interpolate measures using our framework is to use optimal transport to define  $f_{t}$ . Recall  $\mu_t$  denotes the probability flow at time  $t$ . We then define
$$
f _ {t} (x) = \sum_ {i = 1} ^ {M} \phi (t - r _ {i}) \nabla_ {W _ {2}} W _ {2} ^ {2} (\mu_ {t}, \nu_ {i}),
$$
where  $W_{2}^{2}$  is the squared Wasserstein-2 distance and  $\nabla_{W_2}W_2^2$  is its Wasserstein gradient. In practice, we compute  $\nabla_{W_2}W_2^2$  using sample access and we employ the debiased Sinkhorn divergence [Genevay et al., 2018, Feydy et al., 2019] implemented in the JAX OTT library [Cuturi et al., 2022]. This formulation differs from the one in Zhang et al. [2022] in that here we prescribe the precise PDE based on  $f_{t}$ , whereas in Zhang et al. [2022] an optimal transport loss is used to fit the keyframes along with many regularizers on the velocity field  $v_{t}$  to promote the smoothness and other desirable properties. In contrast, we do not use any regularizer on  $v_{t}$ .
To train SCVM-NODE to produce the hand-hand-ball animation sequence in the bottom row of Figure 7, we first sample 20000 points from the interior of the three meshes (a hand mesh, a hand mesh in a different pose, and a ball mesh), and we set  $\nu_{i}$  to be the empirical measure of the corresponding point cloud. Note that different from the first formulation using Fokker-Planck equations, in the optimal transport formulation, throughout we only require sample access from each  $\nu_{i}$ . We use an initial learning rate of  $10^{-4}$  with cosine decay for a total of  $5\times 10^{4}$  iterations. This takes 4.5 hours, which is significantly longer than the training time in Zhang et al. [2022] (reported to be 15 minutes). We leave further improvement of our method to interpolate measures faster as future work.
To render the animation, we sample 20000 points and render the point cloud at each time step using metaballs along with smoothing, similar to the procedure in Zhang et al. [2022]. We did not use the barycentric interpolation postprocessing step in Zhang et al. [2022] which makes sure the key measures  $v_{i}$ 's are fit exactly in the resulting animation. We also did not use unbalanced optimal transport, which as reported in Zhang et al. [2022] can make the fingers of the hand more separated, but requires careful parameter tuning.