SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We consider computing the Wasserstein gradient flow of the KL divergence  $\mathcal{F}(\mu) = \mathrm{KL}(\mu \| \mu^{*})$  where we have density access to the target measure  $\mu^{*}$ . To fit into our framework, we set  $f_{t}(x;\mu_{t}) = \nabla \log p^{*}(x) - \nabla \log p_{t}(x)$  which matches (4) with  $b_{t}(x) = \nabla \log p^{*}(x)$  and  $D_{t}(x) = I_{d}$ . Following the experimental setup in Mokrov et al. [2021] and Fan et al. [2021], we take  $\mu^{*}$  to be a mixture of 10 Gaussians with identity covariance and means sampled uniformed in  $[-5,5]^d$ . The initial measure is  $\mu_0^* = \mathcal{N}(0,16I_d)$ . We solve the corresponding Fokker-Planck PDE for a total time of  $T = 5$  and for  $d = 10,\ldots ,60$ . As TIPF parameterization does not scale to high dimensions, we only consider SCVM-NODE in this experiment.
Figure 1 shows the probability flow produced by SCVM-NODE in dimension 60 at different time steps; as we can see, the flow quickly converges to the target distribution.
Figure 1: Probability flow produced by SCVM-NODE for a 60-dimensional mixture of Gaussians at irregular time steps. Samples are projected onto the first two PCA components and kernel density estimation is used to generate the contours.
In Figure 2, we quantitatively compare our method with Mokrov et al. [2021], Fan et al. [2021], and Shen and Wang [2023]. Training time is reported for all methods.
- In the left two columns of Figure 2, we find that even though the adjoint method ADJ [Shen and Wang, 2023] minimizes the self-consistency loss (7) directly, the decay of self-consistency can be much slower than that of SCVM-NODE as the dimension increases. We suspect this is due to the amount of error accumulated in the adjoint method which involves two numerical integration passes to obtain the gradient. Moreover, ADJ requires up to third-order spatial derivatives of the parameterized neural velocity field which can be inaccurate even if the consistency loss is low—in comparison SCVM-NODE only requires one integration pass and the first-order spatial derivative of the network. Despite the bias of the gradient used in SCVM-NODE, it finds more efficient gradient trajectories than ADJ. Additionally, ADJ takes 80 times longer to train than SCVM-NODE in dimension 10, and scaling up to higher dimensions becomes prohibitive.
- The rightmost column of Figure 2 shows SCVM-NODE achieves far lower symmetric KL compared to the JKO methods. The gradient flow computed by JKO methods does not decrease KL divergence monotonically, likely because the optimization at each JKO step has yet to reach the minimum even though we use 2000 gradient updates for each step. For both JKO methods, the running time for each JKO step increases linearly because samples (and for JKO-ICNN also log det terms) need to be pushed through a growing chain of ICNNs; as a result, the total running time scales quadratically with the number of JKO steps. JKO methods also take about 40 times as long evaluation time as
SCVM-NODE in dimension 60 since density access requires solving an optimization problem for every JKO step. On top of the computational advantage and better results, our method also does not have temporal discretization: after being trained, the flow can be accessed at any time  $t$  (Figure 1).
Figure 2: Quantitative comparison for the mixture of Gaussians experiment. The left two columns plot the symmetric KL (at  $t = T$  compared against the target measure) and consistency (7) versus the training iterations for SCVM-NODE (ours) and ADJ [Shen and Wang, 2023]. The rightmost column plots the symmetric KL across time  $t$  (compared against the target measure) for SCVM-NODE and the JKO methods in high dimensions. Training time: for  $d = 10$ , SCVM-NODE takes 7.37 minutes, ADJ takes 585.2 minutes; for  $d = 60$ , SCVM-NODE takes 23.9 minutes, JKO-ICNN takes 375.2 minutes, and JKO-ICNN-PD takes 24.4 minutes.
4.2 Ornstein-Uhlenbeck process
We consider the Wasserstein gradient flow of the KL divergence with respect to a Gaussian with the initial distribution being a Gaussian, following the same experimental setup in Mokrov et al. [2021], Fan et al. [2021]. In this case, the gradient flow at time  $t$  is a Gaussian  $G(t)$  with a known mean and covariance; see Appendix D.1 for details. We quantitatively compare all methods in Figure 3:
- ADJ achieves the best results in dimensions  $d = 5$  and  $d = 10$ . However, this is at the cost of high training time: in dimension 10, ADJ takes 341 minutes to train, while SCVM-TIPF and SCVM-NODE take 23 and 9 minutes respectively for the same 20k training iterations. As such, we omit ADJ in higher-dimensional comparisons.
- Both our SCVM-TIPF and SCVM-NODE achieve good results second only to ADJ in dimensions 5 and 10. In low dimensions, SCVM-TIPF results in lower probability divergences than SCVM-NODE likely due to having exact density access. Although not shown, SCVM-TIPF also satisfies the initial condition well (numbers at  $t = 0$  are comparable to those at  $t = 0.25$  in the left two columns in Figure 3).
- For the two JKO methods, they result in much higher errors for  $t \leq 0.5$  compared to later times: this is expected because the dependency of  $G(t)$  on  $t$  is exponential, so convergence to  $\mu^{*}$  is faster in the beginning, yet a constant step size is used for JKO methods.
- For DFE, the result is highly sensitive to the forward Euler step size  $\Delta t$ . We choose step size  $\Delta t = 0.01$  which empirically gives the best results among  $\{0.1, 0.01, 0.001, 0.0001\}$ . As DFE achieves far lower symmetric KL divergence or  $f$ -divergence compared to alternatives, we only include its Bures-Wasserstein distance in high dimensions where its number is slightly worse than alternatives (bottom right plot of Figure 3).
4.3 Porous medium equation
Following Fan et al. [2021], we consider the porous medium equation with only diffusion:  $\partial_t p_t = \Delta p_t^m$  with  $m > 1$ . Its solution is the Wasserstein gradient flow of  $\mathcal{F}(\mu) = \int \frac{1}{m - 1} p(x)^m \, \mathrm{d}x$  where  $p$  is the density of  $\mu$  with  $\nabla_{W_2}\mathcal{F}(\mu)(x) = \nabla (\frac{m}{m - 1} p^{m - 1}(x))$ —see Appendix D.2 for details. We consider only SCVM-TIPF and JKO methods here.
Figure 3: Quantitative results for the OU process experiment. The left two columns show the metrics (symmetric KL, symmetric  $f$ -divergence, and Bures-Wasserstein distance) versus time  $t$  of various methods computed against the closed formed solution  $G(t)$  in dimension  $d = 5, 10$ . The right column shows the metrics averaged across  $t$  versus dimension  $d$  in higher dimensions.
We show the efficiency of SCVM-TIPF compared to JKO-ICNN in dimension  $d = 1,2,\ldots ,6$ . We exclude JKO-ICNN-PD because it produces significantly worse results. We visualize the density  $p_t$  of the solution from SCVM-TIPF and JKO-ICNN on the top of Figure 4 in dimension 1 compared to  $p_t^*$ . Both methods approximate  $p_t^*$  well with SCVM-TIPF being more precise at small  $t$ ; this is consistent with the observation in Figure 2 where JKO methods result in bigger errors for small  $t$ .
Figure 4: Top: visualization of the densities of  $p_t^*$  and  $p_t$  for the porous medium equation in dimension 1 at varying time steps  $t$  for SCVM-TIPF and JKO-ICNN. Bottom: total variation distance, symmetric  $f$ -divergence, and Wasserstein-2 distances across dimensions at  $t = 0.004$  and  $t = 0.025$  between  $p_t$  and  $p_t^*$  for solving the porous medium equation.
On the bottom row of Figure 4, we plot the total variation (TV) distance, the symmetric  $f$ -divergence, and the Wasserstein-2 distance (details on the TV distance are given in Appendix C.6) between the recovered solution  $p_t$  and  $p_t^*$  for both methods at  $t = 0.004$  and  $t = 0.025$ . Note that the values of all metrics are very low implying that the solution from either method is very accurate, with SCVM-TIPF more precise in TV distance and symmetric  $f$ -divergence, especially for  $d > 3$ . Like with the experiments in previous sections, JKO-ICNN is much slower to train: in dimension 6, training JKO-ICNN took 102 minutes compared to 21 minutes for SCVM-TIPF.
4.4 Time-Dependent Fokker-Planck equation
In this section, we qualitatively evaluate our method for solving a PDE that is not a Wasserstein gradient flow. In this case, JKO-based methods cannot be applied. Consider the OU process from Section 4.2 when the mean  $\beta$  and the covariance matrix  $\Gamma$  become time-dependent as  $\beta_{t}$  and  $\Gamma_t$ . The resulting PDE is a time-dependent Fokker-Planck equation of the form (4) with
$$
f _ {t} (X, \mu_ {t}) = \Gamma_ {t} \left(\beta_ {t} - X\right) - D \nabla \log p _ {t} (X). \tag {11}
$$
In this configuration, when the initial measure  $p_0$  is Gaussian, the solution  $\mu_t$  can again be shown to be Gaussian with mean and covariance following an ODE—see Appendix D.3 for more details. We consider, in dimension 2 and 3, time-dependent attraction towards a harmonic mean  $\beta_t = a(\sin(\pi\omega t), \cos(\pi\omega t))$  using the expression of  $\beta_t$  from Boffi and Vanden-Eijnden [2023], augmented to  $\beta_t = a(\sin(\pi\omega t), \cos(\pi\omega t), t)$  in dimension 3.
We apply both SCVM-TIPF and SCVM-NODE to this problem and compare our results with alternatives. Similar to Figure 3, as shown in Figure 5, both SCVM-TIPF and SCVM-NODE achieve results on par with ADJ, with both SCVM methods being 30 times faster than ADJ in dimension 10. DFE results in good Wasserstein-2 metrics but worse divergences. Visualization of the evolution of a few sampled particles are given in Figure 9 and Figure 10.
In Appendix D.4, we augment (11) with an interaction term to simulate a flock of (infinitely many) birds, resulting in a non-Fokker-Planck PDE that can be readily solved by our method.
Figure 5: Symmetric KL divergence and Wasserstein-2 distances across time for  $d = 2, 3$  between the recovered flows and the ground truth for the time-dependent Fokker-Planck equation.
4.5 Additional qualitative low-dimensional dynamics
To demonstrate the flexibility of our method, we apply our algorithm to model more general mass-conserving dynamics than the ones considered in the previous sections.
Flow splashing against obstacles. We model the phenomenon of a 2-dimensional flow splashing against obstacles using a Fokker-Planck equation (4) where  $b_{t}$  encodes the configuration of three
obstacles that repel the flow (See Appendix D.5 for details). We solve this PDE using SCVM-NODE for  $T = 5$  and visualize the recovered flow in (6). When solving the same PDE using SDE-EM, the flow incorrectly crosses the bottom right obstacle due to a finite time step size (Figure 14). When using DFE, the path of initial samples appears jagged (right of Figure 13); our method has no such issue and results in continuous sample paths (left of Figure 13). Method ADJ suffers from numerical instability and cannot be trained without infinite loss in this example.
Figure 6: A flow splashing against three obstacles (in purple) produced by SCVM-NODE. Particles are colored based on the initial  $y$  coordinates.
Smooth interpolation of measures. To illustrate the flexibility of our method, we demonstrate two ways to formulate the problem of smoothly interpolating a list of measures. First, we model the interpolation as a time-dependent Fokker-Planck equation and use it to interpolate MNIST digits 1, 2, and 3, starting from a Gaussian (Figure 7). Next, we adopt an optimal transport formulation and use it to generate an animation sequence deforming a 3D hand model to a different pose and then to a ball, similar to the setup in Zhang et al. [2022]. Note that the optimal transport formulation is not solvable using competing methods. See Appendix D.6 for more details.
t=0.00
t=0.60
t-1.80
t-340
t=3.00
+360
+420
t-4.80
t-540
Figure 7: Smooth interpolation of measures. Top: interpolating MNIST digits 1 to 3. Bottom: interpolating hand from the initial pose to a different pose and then to a ball.
By extending the concept of self-consistency from Shen et al. [2022], we present an iterative optimization method for solving a wide class of mass-conserving PDEs without temporal or spatial discretization. In all experiments considered, our method achieves strong quantitative results with significantly less training time than JKO-based methods and the adjoint method in high dimensions.
Below we highlight a few future directions. First, as discussed, the two ways to parameterize a probability flow, TIPF, and NODE, both have their specific limitations. Finding a new parameterization that combines the advantages of both TIPF and NODE is an important next step. Secondly, we hope to extend our approach to incorporate more complicated boundary conditions. Finally, given that the proposed algorithm is highly effective empirically, it would be an interesting theoretical step to explore its convergence properties.
Acknowledgements The MIT Geometric Data Processing group acknowledges the generous support of Army Research Office grants W911NF2010168 and W911NF2110293, of Air Force Office of Scientific Research award FA9550-19-1-031, of National Science Foundation grant CHS-1955697, from the CSAIL Systems that Learn program, from the MIT-IBM Watson AI Laboratory, from the Toyota-CSAIL Joint Research Center, from a gift from Adobe Systems, and from a Google Research Scholar award.
SH acknowledges the financial support from the University of Bordeaux (UBGR grant) and the French Research Agency (PostProdLEAP).
SCVM-TIPF	SCVM-NODE	ADJ	JKO-ICNN	JKO-ICNN-PD
O(Std2)	O(STNode)d	O(STNode d3)	O(St2d3)	O(St2d)
Table 1: Time complexity of training for  $S$  iterations for the methods considered, in terms of dimension  $d$ . We assume for simplicity that the same batch size is used, the training is done for  $T$  time steps, and any network forward pass takes  $O(d)$  time. For SCVM-NODE and ADJ,  $N_{\mathrm{ode}}$  denotes the number of ODE integration steps. For SCVM-TIPF, RealNVP is used to build the coupling layer, and we use  $d$  coupling layers, hence the extra multiple of  $d$ . For ADJ, the  $d^3$  term comes from the third-order spatial derivatives. For JKO-ICNN, the  $d^3$  term is due to computing the log-determinant term. For both JKO methods, the quadratic dependency on  $T$  is due to maintaining a growing chain of neural networks of size  $T$  as described in the related works section.
In this appendix, we provide details and further justification of the proposed method. In Appendix A, we provide an interpretation of the update (9) as a gradient descent step with a biased gradient. In Appendix B, we explain the integration-by-parts trick used to prove (3.1). In Appendix C, we provide implementation details of all considered methods. In Appendix D, we provide additional experimental details and results.
Algorithm 1 Self-consistent velocity matching
Input:  $f_{t}(\cdot ,\cdot),\mu_{0}^{*},T,N_{train},B,L.$  Initialize network weights  $\theta$    
for  $k = 1,\dots ,N_{train}$  do  $\theta^{\prime}\gets \theta .$  Sample  $x_{1},\ldots ,x_{B}\sim \mu_{0}^{*},t_{1},\ldots ,t_{L}\sim [0,T].$ $y_{b,l}\gets \Phi_{t_l}^{\theta '}(x_b),\forall b = 1,\ldots ,B,l = 1,\ldots ,L.$ $\theta \leftarrow \theta -\eta \nabla_{\theta}\frac{1}{BL}\sum_{b,l}\left\| v_{t_l}^{\theta}(y_{b,l}) - f_{t_l}(y_{b,l};\mu_{t_l}^{\theta '})\right\| ^2.$    
end for   
Output: optimized  $\theta$
A Biased Gradient Interpretation
Assume  $\mu_t^\theta$  has density  $p_t^\theta$ . Using  $f_t^\theta(x)$  to denote  $f_t(x; \mu_t^\theta)$  from (7), the self-consistency loss can be written as,
$$
L (\theta) := \int_ {0} ^ {T} \int p _ {t} ^ {\theta} (x) \left\| v _ {t} ^ {\theta} (x) - f _ {t} ^ {\theta} (x) \right\| ^ {2} \mathrm {d} x \mathrm {d} t.
$$
Assuming all terms involving  $\theta$  are differentiable with respect to  $\theta$ , the gradient of  $L(\theta)$  with respect to the neural network parameters  $\theta$  can be written as:
$$
\begin{array}{l} \nabla L (\theta) = \int_ {0} ^ {T} \int \nabla_ {\theta} p _ {t} ^ {\theta} (x) \| v _ {t} ^ {\theta} (x) - f _ {t} ^ {\theta} (x) \| ^ {2} \mathrm {d} x \mathrm {d} t (12) \\ + 2 \int_ {0} ^ {T} \int p _ {t} ^ {\theta} (x) J _ {\theta} v _ {t} ^ {\theta} (x) ^ {\top} \left(v _ {t} ^ {\theta} (x) - f _ {t} ^ {\theta} (x)\right) \mathrm {d} x \mathrm {d} t (13) \\ - 2 \int_ {0} ^ {T} \int p _ {t} ^ {\theta} (x) J _ {\theta} f _ {t} ^ {\theta} (x) ^ {\top} \left(v _ {t} ^ {\theta} (x) - f _ {t} ^ {\theta} (x)\right) \mathrm {d} x \mathrm {d} t. (14) \\ \end{array}
$$
Here we use  $J_{\theta}$  to denote the Jacobian with respect to  $\theta$ . On the other hand, the gradient used in the updates (8) is  $\nabla_{\theta}F(\theta ,\theta^{\prime})$  at  $\theta^{\prime} = \theta$ :
$$
\left. \nabla_ {\theta} F \left(\theta , \theta^ {\prime}\right) \right| _ {\theta^ {\prime} = \theta} = 2 \int_ {0} ^ {T} \int p _ {t} ^ {\theta} (x) J _ {\theta} v _ {t} ^ {\theta} (x) ^ {\top} \left(v _ {t} ^ {\theta} (x) - f _ {t} ^ {\theta} (x)\right) \mathrm {d} x \mathrm {d} t. \tag {15}
$$
We see that (15) is exactly the middle term (13). Hence our formulation can be interpreted as doing gradient descent with a biased gradient estimator. It remains a future work direction to
theoretically analyze the amount of bias in (15) and the condition under which the dot product  $\langle \nabla L(\theta),\nabla_{\theta}F(\theta ,\theta^{\prime})|_{\theta^{\prime} = \theta}\rangle \geq 0$  . The central challenge would be to relate  $J_{\theta}v_{t}^{\theta}$  and  $J_{\theta}f_{t}^{\theta}$  ; this depends on the neural network architecture and the type of the PDE.
B Integration-by-Parts Trick
This is a common trick used in score-matching literature [Hyvarinen and Dayan, 2005].
Proof of Proposition 3.1. Fix  $t > 0$ . The form of  $f_{t}$  in (4) is
$$
f _ {t} (x; \mu_ {t}) = b _ {t} (x) - D _ {t} (x) \nabla \log p _ {t} (x).
$$
$$
\mathbf {E} _ {X \sim \mu_ {t} ^ {\theta^ {\prime}}} \left[ v _ {t} ^ {\theta} (X) ^ {\top} f _ {t} (X; \mu_ {t} ^ {\theta^ {\prime}}) \right] = \mathbf {E} _ {X \sim \mu_ {t} ^ {\theta^ {\prime}}} \left[ v _ {t} ^ {\theta} (X) ^ {\top} b _ {t} (X) \right] - \mathbf {E} _ {X \sim \mu_ {t} ^ {\theta^ {\prime}}} \left[ v _ {t} ^ {\theta} (X) ^ {\top} D _ {t} (X) \nabla \log p _ {t} ^ {\theta^ {\prime}} (X) \right].
$$
The second term can be written as
$$
\begin{array}{l} \mathbf {E} _ {X \sim \mu_ {t} ^ {\theta^ {\prime}}} \left[ v _ {t} ^ {\theta} (X) ^ {\top} D _ {t} (X) \nabla \log p _ {t} ^ {\theta^ {\prime}} (X) \right] = \int v _ {t} ^ {\theta} (x) ^ {\top} D _ {t} (x) \nabla \log p _ {t} ^ {\theta^ {\prime}} (x) \mathrm {d} p _ {t} ^ {\theta^ {\prime}} (x) \\ = \int v _ {t} ^ {\theta} (x) ^ {\top} D _ {t} (x) \nabla p _ {t} ^ {\theta^ {\prime}} (x) / p _ {t} ^ {\theta^ {\prime}} (x) \cdot p _ {t} ^ {\theta^ {\prime}} (x) d x \\ = \int v _ {t} ^ {\theta} (x) ^ {\top} D _ {t} (x) \nabla p _ {t} ^ {\theta^ {\prime}} (x) d x \\ = - \int \boldsymbol {\nabla} \cdot \left(D _ {t} (x) ^ {\top} v _ {t} ^ {\theta} (x)\right) p _ {t} ^ {\theta^ {\prime}} (x) d x \\ = - \mathbf {E} _ {X \sim \mu_ {t} ^ {\theta^ {\prime}}} \left[ \nabla \cdot \left(D _ {t} (X) ^ {\top} v _ {t} ^ {\theta} (X)\right) \right], \\ \end{array}
$$
where we use integration-by-parts to get the second last equation and the assumption that  $v_{t}^{\theta}, D_{t}$  are bounded and  $p_t^\theta (x)\to 0$  as  $\| x\| \to \infty$ .
C.1 Network architectures for SCVM.
For TIPF, our implementation follows Dinh et al. [2016]. Each coupled layer uses 3-layer fully connected networks with layer sizes 64, 128, 128 for both scale and translation prediction. We use twice as many coupling layers as the dimension of the problem while each coupling layer updates one coordinate; we found using fewer layers with random masking gives much worse results.
For NODE, we use a 3-layer fully connected network for modeling the velocity field with layer size 256. We additionally add a low-rank linear skip connection  $x \mapsto A(t)x + b(t)$  where  $A(t) = L(t)L^{\top}(t)$  and  $L(t)$  is a  $d \times 20$  matrix to make  $A(t)$  low-rank.
We use SILU activation [Elfwing et al., 2018] which is smooth for all layers for both TIPF and NODE. For NODE, we apply layer normalization before applying activation. We also add a sinusoidal embedding for the time input  $t$  plus two fully connected layers of size 64 before concatenating it with the spatial input.
The numerical integration for NODE is done using Diffrax library [Kidger, 2021] with a relative and absolute tolerance of  $10^{-4}$ ; we did not find considerable improvement when using a lower tolerance.
We use the integration-by-parts trick for SCVM-NODE whenever possible. Since TIPF has tractable log density, we do not use such a trick and optimize (9) directly for SCVM-TIPF which we found to produce better results.
C.2 Hyperparameters.
Unless mentioned otherwise, we choose the following hyperparameters for Algorithm 1. We set  $N_{\mathrm{train}} = 10^5$  or  $2 \times 10^5$ ,  $B = 1000$ ,  $L = 10$  or 20. We use Adam [Kingma and Ba, 2014] with
a cosine decay learning rate scheduler, with initial learning rate  $10^{-3}$ , the number of decay steps same as  $N_{\mathrm{train}}$ , and  $\alpha = 0.01$  (so the final learning rate is  $10^{-5}$ ). Since we are effectively performing gradient descent using a biased gradient, we set  $b_{2} = 0.9$  in Adam (instead of the default  $b_{2} = 0.999$ ), so that the statistics in Adam can be updated more quickly; we found this tweak improves the results noticeably.
C.3 Implementation of JKO methods.
We base our JAX implementation of ICNN on the codebase by the original ICNN author: https://github.com/facebookresearch/w2ot. Compared to the original ICNN implementation by Amos et al. [2017], we add an additional convex quadratic skip connections used by Mokrov et al. [2021], which we found to be crucial for the OU process experiment. For ICNNs, we use hidden layer sizes 64, 128, 128, 64. The quadratic rank for the convex quadratic skip connections is set to 20. The activation layer is taken to be CELU.
To implement the method by Fan et al. [2021], we model the dual potential as a 4-layer fully connected network with layer size 128, with CELU activation. For the gradient flow of KL divergence and generalized entropy (used in Section 4.3), we follow closely the variational formulation and the necessary change of variables detailed in Fan et al. [2021, Corollary 3.3, Corollary 3.4].
In order to compute the log density at any JKO step, following Mokrov et al. [2021], we need to solve a convex optimization to find the inverse of the gradient of an ICNN. We use the LBFGS algorithm from JAXopt [Blondel et al., 2021] to solve the optimization with tolerance  $10^{-2}$  (except for Section 4.3 we use a tolerance of  $10^{-3}$  to obtain finer inverses, but it takes 6x longer compared to  $10^{-2}$ ).
We always use 40 JKO steps, consistent with past works. For each JKO step, we perform 1000 stochastic gradient descent using Adam optimizer with a learning rate of  $10^{-3}$ , except for the mixture of Gaussians experiment, we use 2000 steps—using fewer steps will result in worse results. We have tested with the learning rate schedules used in Fan et al. [2021], Mokrov et al. [2021] and did not notice any improvement.
C.4 Implementation of ADJ method.
We implement the adjoint method carefully following the formulation in Shen and Wang [2023]. The neural network for parameterizing the velocity field is identical to the one used in SCVM-NODE. The ODE integration also uses the same hyperparameters as that of SCVM-NODE. This way we can compare ADJ with SCVM-NODE in a fair manner since they only differ in the gradient estimation.
C.5 Implementation of DFE method.