SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Lingxiao Li
lingxiao@mit.edu
Univ. Bordeaux, Bordeaux INP, CNRS, IMB
Justin Solomon
jsolomon@mit.edu
We present a discretization-free scalable framework for solving a large class of mass-conserving partial differential equations (PDEs), including the time-dependent Fokker-Planck equation and the Wasserstein gradient flow. The main observation is that the time-varying velocity field of the PDE solution needs to be self-consistent: it must satisfy a fixed-point equation involving the probability flow characterized by the same velocity field. Instead of directly minimizing the residual of the fixed-point equation with neural parameterization, we use an iterative formulation with a biased gradient estimator that bypasses significant computational obstacles with strong empirical performance. Compared to existing approaches, our method does not suffer from temporal or spatial discretization, covers a wider range of PDEs, and scales to high dimensions. Experimentally, our method recovers analytical solutions accurately when they are available and achieves superior performance in high dimensions with less training time compared to alternatives.
Mass conservation is a ubiquitous phenomenon in dynamical systems arising from fluid dynamics, electromagnetism, thermodynamics, and stochastic processes. Mathematically, mass conservation is formulated as the continuity equation:
$$
\partial_ {t} p _ {t} (x) = - \boldsymbol {\nabla} \cdot \left(v _ {t} p _ {t}\right), \forall x, t \in [ 0, T ], \tag {1}
$$
where  $p_t: \mathbf{R}^d \to \mathbf{R}$  is a scalar quantity such that the total mass  $\int p_t(x)$  is conserved with respect to  $t$ ,  $v_t: \mathbf{R}^d \to \mathbf{R}^d$  is a velocity field, and  $T > 0$  is total time. We will assume, for all  $t \in [0, T]$ ,  $p_t \geq 0$  and  $\int p_t(x) \, \mathrm{d}x = 1$ , i.e.,  $p_t$  is a probability density function. We use  $\mu_t$  to denote the probability measure with density  $p_t$ . Once a pair  $(p_t, v_t)$  satisfies (1), the density  $p_t$  is coupled with  $v_t$  in the sense that the evolution of  $p_t$  in time is characterized by  $v_t$  (Section 3.1).
We consider the subclass of mass-conserving PDEs that can be written succinctly as
$$
\partial_ {t} p _ {t} (x) = - \boldsymbol {\nabla} \cdot \left(f _ {t} \left(x; \mu_ {t}\right) p _ {t}\right), \forall x, t \in [ 0, T ], \tag {2}
$$
where  $f_{t}(\cdot ;\mu_{t}):\mathbf{R}^{d}\to \mathbf{R}^{d}$  is a given function depending on  $\mu_t$  , with initial condition  $\mu_0 = \mu_0^*$  for a given initial probability measure  $\mu_0^*$  with density  $p_0^*$
Different choices of  $f_{t}$  lead to a large class of mass-conserving PDEs. For instance, given a functional  $\mathcal{F}:\mathcal{P}_2(\mathbf{R}^d)\to \mathbf{R}$  on the space of probability distributions with finite second moments, if we take
$$
f _ {t} (x; \mu_ {t}) := - \nabla_ {W _ {2}} \mathcal {F} (\mu_ {t}) (x), \tag {3}
$$
where  $\nabla_{W_2}\mathcal{F}(\mu):\mathbf{R}^d\to \mathbf{R}^d$  is the Wasserstein gradient of  $\mathcal{F}$ , then the solution to (2) is the Wasserstein gradient flow of  $\mathcal{F}$  [Santambrogio, 2015, Chapter 8]. Thus, solving (2) efficiently allows us to optimize in the probability measure space. If we take
$$
f _ {t} (x; \mu_ {t}) := b _ {t} (x) - D _ {t} (x) \nabla \log p _ {t} (x), \tag {4}
$$
where  $b_{t}$  is a velocity field and  $D_{t}(x)$  is a positive-semidefinite matrix, then we obtain the time-dependent Fokker-Planck equation [Risken and Risken, 1996], which describes the time evolution of the probability flow undergoing drift  $b_{t}$  and diffusion with coefficient  $D_{t}$ .
A popular strategy to solve (2) is to use an Eulerian representation of the density field  $p_t$  on a discretized mesh or as a neural network [Raissi et al., 2019]. However, these approaches do not fully exploit the mass-conservation principle and are usually limited to low dimensions. Shen et al. [2022], Shen and Wang [2023] recently introduced the notion of self-consistency for the Fokker-Planck equation and more generally McKean-Vlasov type PDEs. This notion is a Lagrangian formulation of (2). They apply the adjoint method to optimize self-consistency. In this work, we extend their notion of self-consistency to mass-conserving PDEs of the general form (2). Equipped with this formulation, we develop an iterative optimization scheme called self-consistent velocity matching. With the probability flow parameterized as a neural network, at each iteration, we refine the velocity field  $v_t$  of the current flow to match an estimate of  $f_t$  evaluated using the network weights from the previous iteration. Effectively, we minimize the self-consistency loss with a biased but more tractable gradient estimator.
This simple scheme has many benefits. First, the algorithm is agnostic to the form of  $f_{t}$ , thus covering a wider range of PDEs compared to past methods. Second, we no longer need to differentiate through differential equations using the adjoint method as in Shen and Wang [2023], which is orders of magnitude slower than our method with worse performance in high dimensions. Third, this iterative formulation allows us to rewrite the velocity-matching objectives for certain PDEs to get rid of computationally expensive quantities such as  $\nabla \log p_t$  in the Fokker-Planck equation (Proposition 3.1). Lastly, our method is flexible with probability flow parameterization: we have empirically found that the two popular ways of parameterizing the flow—as a time-varying pushforward map [Biloš et al., 2021] and as a time-varying velocity field [Chen et al., 2018]—both have merits in different scenarios.
Our method tackles mass-conserving PDEs of the form (2) in a unified manner without temporal or spatial discretization. Despite using a biased gradient estimator, in practice, our method decreases the self-consistency loss efficiently (second column of Figure 2). For PDEs with analytically-known solutions, we quantitatively compare with the recent neural JKO-based methods [Mokrov et al., 2021, Fan et al., 2021, Alvarez-Melis et al., 2021], the adjoint method [Shen and Wang, 2023], and the particle-based method [Boffi and Vanden-Eijnden, 2023]. Our method faithfully recovers true solutions with quality on par with the best previous methods in low dimensions and with superior quality in high dimensions. Our method is also significantly faster than competing methods, especially in high dimensions, at the same time without discretization. We further demonstrate the flexibility of our method on two challenging experiments for modeling flows splashing against obstacles and smooth interpolation of measures where the comparing methods are either not applicable or have noticeable artifacts.
Classical PDE solvers for mass-conserving PDEs such as the Fokker-Planck equation and the Wasserstein gradient flow either use an Eulerian representation of the density and discretize space as a grid or mesh [Burger et al., 2010, Carrillo et al., 2015, Peyre, 2015] or use a Lagrangian representation, which discretizes the flow as a collection of interacting particles simulated forward in time [Crisan and Lyons, 1999, Westdickenberg and Wilkening, 2010]. Due to spatial discretization, these methods struggle with high-dimensional problems. Hence, the rest of the section focuses solely on recent neural network-based methods.
Physics-informed neural networks. Physics-informed neural networks (PINNs) are prominent methods that solve PDEs using deep learning [Raissi et al., 2019, Karniadakis et al., 2021]. The main idea is to minimize the residual of the PDE along with loss terms to enforce the boundary conditions and to match observed data. Our notion of self-consistency is a Lagrangian analog of the residual in PINN. Our velocity matching only occurs along the flow of the current solution where interesting dynamics happen, while in PINNs the residual is evaluated on collocation points that occupy the entire domain. Hence our method is particularly suitable for high-dimensional problems where the dynamics have a low-dimensional structure.
Neural JKO methods. Recent works [Mokrov et al., 2021, Alvarez-Melis et al., 2021, Fan et al., 2021] apply deep learning to the time-discretized JKO scheme [Jordan et al., 1998] to solve Wasserstein gradient flow (3). By pushing a reference measure through a chain of neural networks parameterized as input-convex neural networks (ICNNs) [Amos et al., 2017], these methods avoid discretizing the space. Mokrov et al. [2021] optimize one ICNN to minimize Kullback-Leibler (KL) divergence plus a Wasserstein-2 distance term at each JKO step. This method is extended to other functionals by Alvarez-Melis et al. [2021]. Fan et al. [2021] use the variational formulation of  $f$ -divergence to obtain a faster primal-dual approach.
An often overlooked problem of neural JKO methods is that the total training time scales quadratically with the number of JKO steps: to draw samples for the current step, initial samples from the reference measure must be passed through a long chain of neural networks, along with expensive quantities like densities. However, using too few JKO steps results in large temporal discretization errors. Moreover, the optimization at each step might not have fully converged before the next step begins, resulting in an unpredictable accumulation of errors. In contrast, our method does not suffer from temporal discretization and can be trained end-to-end. It outperforms these neural JKO methods with less training time in experiments we considered.
Velocity matching. A few recent papers employ the idea of velocity matching to construct a flow that follows a learned velocity field. di Langosco et al. [2021] simulate the Wasserstein gradient flow of the KL divergence by learning a velocity field that drives a set of particles forward in time for Bayesian posterior inference. The velocity field is refined on the fly based on the current positions of the particles. Boffi and Vanden-Eijnden [2023] propose a similar method that applies to a more general class of time-dependent Fokker-Planck equations. These two methods approximate probability measures using finite particles which might not capture high-dimensional distributions well. Liu et al. [2022], Lipman et al. [2022], Albergo and Vanden-Eijnden [2022] use flow matching for generative modeling by learning a velocity field that generates a probability path connecting a reference distribution to the data distribution. Yet these methods are not designed for solving PDEs.
Most relevant to our work, Shen et al. [2022] propose the concept of self-consistency for the Fokker-Planck equation, later extended to McKean-Vlasov type PDEs [Shen and Wang, 2023]. They observe that the velocity field of the flow solution to the Fokker-Planck equation must satisfy a fixed-point equation. They theoretically show that, under certain regularity conditions, a form of probability divergence between the current solution and the true solution is bounded by the self-consistency loss that measures the violation of the fixed-point equation. Their algorithm minimizes such violation using neural ODE parameterization [Chen et al., 2018] and the adjoint method. Our work extends the concept of self-consistency to a wider class of PDEs in the form of (2) and circumvents the computationally demanding adjoint method using an iterative formulation. We empirically verify that our method is significantly faster and reduces the self-consistency loss more effectively in moderate dimensions than that of Shen and Wang [2023] (Figure 2).
3 Self-Consistent Velocity Matching
3.1 Probability flow of the continuity equation
A key property of the continuity equation (1) is that any solution  $(p_t, v_t)_{t \in [0,T]}$  (provided  $p_t$  is continuous with respect to  $t$  and  $v_t$  is bounded) corresponds to a unique flow map  $\{\Phi_t(\cdot): \mathbf{R}^d \to \mathbf{R}^{d'}\}_{t \in [0,T]}$  that solves the ordinary differential equations (ODEs) [Ambrosio et al., 2005, Proposition 8.1.8]
$$
\Phi_ {0} (x) = x, \frac {\mathrm {d}}{\mathrm {d} t} \Phi_ {t} (x) = v _ {t} \left(\Phi_ {t} (x)\right), \forall x, t \in [ 0, T ], \tag {5}
$$
and the flow map satisfies  $\mu_t = (\Phi_t)_{\#} \mu_0$  for all  $t \in [0, T]$ , where  $(\Phi_t)_{\#} \mu_0$  to denote the pushforward measure of  $\mu_0$  by  $\Phi_t$ . Moreover, the converse is true: any solution  $(\Phi_t, v_t)$  of (5) with Lipschitz continuous and bounded  $v_t$  is a solution of (1) with  $\mu_t = (\Phi_t)_{\#} \mu_0$  [Ambrosio et al., 2005, Lemma 8.1.6]. Thus the Eulerian viewpoint of (1) is equivalent to the Lagrangian viewpoint of (5). We next exploit this equivalence by modeling the probability flow using the Lagrangian viewpoint so that it automatically satisfies the continuity equation (1).
3.2 Parametrizing probability flows
Our algorithm will be agnostic to the exact parameterization used to represent the probability flow. As such, we need a way to parameterize the flow to access the following quantities for all  $t \in [0,T]$ :
-  $\Phi_t: \mathbf{R}^d \to \mathbf{R}^d$ , the flow map, so  $\Phi_t(x)$  is the location of a particle at time  $t$  if it is at  $x$  at time 0.
-  $v_{t}:\mathbf{R}^{d}\to \mathbf{R}^{d}$ , the velocity field of the flow at time  $t$ .
-  $\mu_t \in \mathcal{P}(\mathbf{R}^d)$ , the probability measure at time  $t$  with sample access and density  $p_t$ .
We will assume all these quantities are sufficiently continuous and bounded to ensure the Eulerian and Lagrangian viewpoints in Section 3.1 are equivalent. This can be achieved by using continuously differentiable activation functions in the network architectures and assuming the network weights are finite similar to the uniqueness arguments given in [Chen et al., 2018]. We can now parameterize the flow modeling either the flow map  $\Phi_t$  or the velocity field  $v_t$  as a neural network.
Time-dependent Invertible Push Forward (TIPF). We first parameterize a probability flow by modeling  $\Phi_t: \mathbf{R}^d \to \mathbf{R}^d$  as an invertible network for every  $t$ . The network architecture is chosen so that  $\Phi_t$  has an analytical inverse with a tractable Jacobian determinant, similar to [Biloš et al., 2021]. We augment RealNVP from Dinh et al. [2016] so that the network for predicting scale and translation takes  $t$  as an additional input. To enforce the initial condition, we need  $\Phi_0$  to be the identity map. This condition can be baked into the network architecture [Biloš et al., 2021] or enforced by adding an additional loss term  $\mathbf{E}_{X \sim \mu_0^*} \| \Phi_0(X) - X \|^2$ . For brevity, we will from now on omit in the text this additional loss term. The velocity field can be recovered via  $v_t(x) = \partial_t \Phi_t (\Phi_t^{-1}(x))$ . To recover the density  $p_t$  of  $\mu_t = (\Phi_t)_{\#} \mu_0$ , we use the change-of-variable formula  $\log p_t(x) = \log p_0^* (\Phi_t^{-1}(x)) + \log \det |J \Phi_t^{-1}(x)|$  (see (1) in Dinh et al. [2016]).
Neural ODE (NODE). We also parameterize a flow by modeling  $v_{t}:\mathbf{R}^{d}\to \mathbf{R}^{d}$  as a neural network; this is used in Neural ODE [Chen et al., 2018]. The network only needs to satisfy the minimum requirement of being continuous. The flow map and the density can be recovered via numerical integration:  $\Phi_t(x) = x + \int_0^t v_s(\Phi_s(x))\mathrm{d}s$  and  $\log p_t(\Phi_t(x)) = \log p_0^* (x) - \int_0^t\nabla \cdot v_s(\Phi_s(x))\mathrm{d}s$ , a direct consequence of (1) also known as the instantaneous change-of-variable formula [Chen et al., 2018]. To obtain the inverse of the flow map, we integrate along  $-v_{t}$ . With NODE, the initial condition  $\mu_0 = \mu_0^*$  is obtained for free.
We summarize the advantages and disadvantages of TIPF and NODE as follows. While the use of invertible coupling layers in TIPF allows exact access to samples and densities, TIPF becomes less effective in higher dimensions as many couple layers are needed to retain good expressive power, due to the invertibility requirement. In contrast, NODE puts little constraints on the network architecture, but numerical integration can have errors. Handling the initial condition is trivial for NODE while an additional loss term or special architecture is needed for TIPF. As we will show in the experiments, both strategies have merits.
3.3 Formulation
We now describe our algorithm for solving mass-conserving PDEs (2). A PDE of this form is determined by  $f_{t}(\cdot ;\mu_{t}): \mathbf{R}^{d} \to \mathbf{R}^{d}$  plus the initial condition  $\mu_0^*$ . If a probability flow  $\mu_t$  with flow map  $\Phi_t$  and velocity field  $v_{t}$  satisfies the following self-consistency condition,
$$
v _ {t} (x) = f _ {t} \left(x; \mu_ {t}\right), \forall x \text {i n}
$$
then the continuity equation of this flow implies the corresponding PDE (2) is solved. Conversely, the velocity field of any solution of (2) will satisfy (6). Hence, instead of solving (2) which is a condition on the density  $p_t$  that might be hard to access, we can solve (6) which is a more tractable condition. Shen et al. [2022], Shen and Wang [2023] develop this concept for the Fokker-Planck equation and McKean-Vlasov type PDEs; here we generalize it to a wider class of PDEs of the form (2).
Let  $\theta$  be the network weights that parameterize the probability flow using TIPF or NODE. The flow's measure, velocity field, and flow map at time  $t$  are denoted as  $\mu_t^\theta$ ,  $v_t^\theta$ ,  $\Phi_t^\theta$  respectively. One option to solve (6) would be to minimize the self-consistency loss
$$
\min  _ {\theta} \int_ {0} ^ {T} \mathbf {E} _ {X \sim \mu_ {t} ^ {\theta}} \left[ \left\| v _ {t} ^ {\theta} (X) - f _ {t} \left(X; \mu_ {t} ^ {\theta}\right) \right\| ^ {2} \right] d t. \tag {7}
$$
This formulation is reminiscent of PINNs [Raissi et al., 2019] where a residual of the original PDE is minimized. Direct optimization of (7) is challenging: while the integration over  $[0, T]$  and  $\mu_t^\theta$  can be approximated using Monte Carlo, to apply stochastic gradient descent, we must differentiate through  $\mu_t^\theta$  and  $f_t$ : this can be either expensive or intractable depending on the network parameterization. The algorithm by Shen and Wang [2023] minimizes (7) with the adjoint method specialized to Fokker-Planck equations and McKean-Vlasov type PDEs; extending their approach to more general PDEs requires a closed-form formula for the time evolution of the quantities within  $f_t$ , which at best can only be obtained on a case-by-case basis.
Instead, we propose the following iterative optimization algorithm to solve (7). Let  $\theta_{k}$  denote the network weights at iteration  $k$ . We define iterates
$$
\theta_ {k + 1} := \theta_ {k} - \eta \nabla_ {\theta} | _ {\theta = \theta_ {k}} F (\theta , \theta_ {k}), \tag {8}
$$
$$
F (\theta , \theta_ {k}) := \int_ {0} ^ {T} \mathbf {E} _ {X \sim \mu_ {t} ^ {\theta_ {k}}} \left[ \left\| v _ {t} ^ {\theta} (X) - f _ {t} \left(X; \mu_ {t} ^ {\theta_ {k}}\right) \right\| ^ {2} \right] d t. \tag {9}
$$
Effectively, in each iteration, we minimize (9) by one gradient step where we match the velocity field  $v_{t}^{\theta}$  to what it should be according to  $f_{t}$  based on the network weights  $\theta_{k}$  from the previous iteration. This scheme can be interpreted as a gradient descent on 7 using the biased gradient estimate  $\nabla_{\theta}F(\theta ,\theta_{k})$  —see Appendix A for a discussion. We call this iterative algorithm self-consistent velocity matching.
If  $f_{t}$  depends on the density of  $\mu_{t}$  only through the score  $\nabla \log p_{t}$  (corresponding to a diffusion term in the PDE), then we can apply an integration-by-parts trick [Hyvärinen and Dayan, 2005] to get rid of this density dependency by adding a divergence term of the velocity field. Suppose  $f_{t}$  is from the Fokker-Planck equation (4). Then the cross term in (9) after expanding the squared norm has the following alternative expression.
Proposition 3.1. For every  $t \in [0, T]$ , for  $f_t$  defined in (4), assume  $v_t^\theta, D_t$  are bounded and continuously differentiable, and  $\mu_t^{\theta'}$  is a measure with a continuously differentiable density  $p_t^{\theta'}$  that vanishes in infinity and not at finite points. Then we have
$$
\mathbf {E} _ {X \sim \mu_ {t} ^ {\theta^ {\prime}}} \left[ v _ {t} ^ {\theta} (X) ^ {\top} f _ {t} \left(X; \mu_ {t} ^ {\theta^ {\prime}}\right) \right] = \mathbf {E} _ {X \sim \mu_ {t} ^ {\theta^ {\prime}}} \left[ v _ {t} ^ {\theta} (X) ^ {\top} b _ {t} (X) + \boldsymbol {\nabla} \cdot \left(D _ {t} ^ {\top} (X) v _ {t} ^ {\theta} (X)\right) \right]. \tag {10}
$$
We provide the derivation in Appendix B. With Proposition 3.1, we no longer need access to  $\nabla \log p_t$  when computing  $\nabla_{\theta}F$ . This is useful for NODE parameterization since obtaining the score would otherwise require additional numerical integration.
Our algorithm is summarized in Algorithm 1 in the appendix. We use Adam optimizer [Kingma and Ba, 2014] to modulate the update (8). For sampling time steps  $t_1, \ldots, t_L$  in  $[0, T]$ , we use stratified sampling where  $t_l$  is uniformly sampled from  $[(l-1)T/L, lT/L]$ ; such a sampling strategy results in more stable training in our experiments. We implemented our method using JAX [Bradbury et al., 2018] and FLAX [Heek et al., 2020]. See Appendix C for the implementation details.
We show the efficiency and accuracy of our method on several PDEs of the form (2). We start with three Wasserstein gradient flow experiments (Section 4.1, Section 4.2, Section 4.3). Next, we consider the time-dependent Fokker-Planck equation that simulates attraction towards a harmonic mean in Section 4.4. Finally, in Section 4.5, we apply our framework to generate complicated low-dimensional dynamics including flows splashing against obstacles and smooth interpolation of measures. We will use SCVM-TIPF and SCVM-NODE to denote our method with TIPF and NODE parameterization respectively. We use JKO-ICNN to denote the method by Mokrov et al. [2021], JKO-ICNN-PD to denote the method by Fan et al. [2021] (PD for "primal-dual"), ADJ to denote the adjoint method by Shen and Wang [2023], SDE-EM to denote the Euler-Maruyama method for solving the SDE associated with the Fokker-Planck equation, and DFE ("discrete forward Euler") to denote the method by Boffi and Vanden-Eijnden [2023]. We implemented all competing methods in JAX—see more details in Appendix C—and we compare quantitatively against these methods when possible.
In Table 1, we compare the time complexity of training the described methods, where we show that SCVM-TIPF and SCVM-NODE have low computational complexity among all methods.
Evaluation metrics. For quantitative evaluation, we use the following metrics. To compare measures with density access, following Mokrov et al. [2021], we use the symmetric Kullback-Leibler (symmetric KL) divergence, defined as  $\mathrm{SymKL}(\rho_1,\rho_2)\coloneqq \mathrm{KL}(\rho_1\parallel \rho_2) + \mathrm{KL}(\rho_2\parallel \rho_1)$ , where  $\mathrm{KL}(\rho_1\parallel \rho_2)\coloneqq \mathbf{E}_{X\sim \rho_1}[\log \mathrm{d}\rho_1(X) / \mathrm{d}\rho_2(X)]$ . Sample estimates of KL can be negative which complicates log-scale plotting, so when this happens, we consider an alternative  $f$ -divergence  $D_{f}(\rho_{1}\parallel \rho_{2})\coloneqq \mathbf{E}_{X\sim \rho_{2}}[(\log \rho_{1}(X) - \log \rho_{2}(X))^{2} / 2]$  whose sample estimates are always non-negative. We similarly define the symmetric  $f$ -divergence  $\mathrm{Sym}D_{f}(\rho_{1},\rho_{2})\coloneqq D_{f}(\rho_{1}\parallel \rho_{2}) + D_{f}(\rho_{2}\parallel \rho_{1})$ . For particle-based methods, we use kernel density estimation (with Scott's rule) to obtain the density function before computing symmetric KL or  $f$ -divergence. We also consider the Wasserstein-2 distance [Bonneel et al., 2011] and the Bures-Wasserstein distance [Kroshnin et al., 2021]; these two measures only require sample access. All metrics are computed using i.i.d. samples. See Appendix C.6 for more details.
4.1 Sampling from mixtures of Gaussians