Physics-informed neural networks for sampling

Authors: jingtong  sun
julius  berner
anima  anandkumar
kamyar  azizzadenesheli
Venue: ICLR_cc_2024_Workshop_AI4DiffEqtnsInSci
Published: 2024-02-11 11:20:08

Abstract: We present a framework to sample from high-dimensional unnormalized densities using physics-informed neural networks (PINNs). For various computational science tasks, it is essential to draw samples from a target distribution where the density is known up to a normalizing constant. Without access to any training samples, existing methods based on normalizing flows and diffusion models rely on the simulation of (stochastic) differential equations for training and suffer from mode collapse. Our approach circumvents these issues by solving the underlying continuity and Fokker-Planck equations using PINNs. Motivated by optimal transport and Schrödinger bridges, we further incorporate regularizers based on Hamilton-Jacobi-Bellman equations. Through evaluations on several benchmarks, we demonstrate that our approach can mitigate mode collapse and significantly outperform various baselines.

Main paper: 
 sampling with stochastic processes: 
 We are given a strictly positive, sufficiently smooth function q target : R d → (0, ∞) and our goal is to sample from the density p target = q target Z , where Z := q target (x) dx.(1)To achieve this goal, we want to construct stochastic processes (X t ) t∈[0,T ] controlled by neural networks such that X T ∼ p target .(2)Typically, the process X is assumed to start at a tractable prior density X 0 ∼ p prior , for instance a standard Gaussian p prior = N (0, I). Moreover, we assume that X is governed by a stochastic1 differential equation (SDE)EQUATIONwhere σ : [0, T ] → R d×d is a given diffusion coefficient and the drift µ : R d × [0, T ] → R d is parametrized by neural networks. We also might have additional constraints on the trajectories of X.Recently, several such methods based on SDEs (nonzero σ) as well as ODEs (σ ≡ 0) have been established. The former are, for instance, based on Schrödinger (half-)bridges, diffusion models, or annealed flows (Vargas & Nüsken, 2023; Zhang & Chen, 2022; Berner et al., 2022; Richter et al., 2023; Zhang et al., 2023; Vargas et al., 2023) . The latter leverage continuous-time normalizing flows and combinations with Monte Carlo (MCMC) methods (Wu et al., 2020; Midgley et al., 2022; Matthews et al., 2022; Arbel et al., 2021) , see Appendix A for further related work. However, all the previously mentioned methods rely on simulating (parts of) the process X for training. This requires time-discretizations and typically results in unstable and slow convergence.In this work, we identify the partial differential equations (PDEs) governing the underlying dynamics of many of these methods. Leveraging neural PDE solver, such as physics-informed neural networks (PINNs) (Raissi et al., 2017; Sirignano & Spiliopoulos, 2018) , we show that this leads to simulation-and discretization-free objectives. We compare various objectives and demonstrate strong performance on challenging, high-dimensional distributions. 
learning the evolution: 
 Assuming that X t has a sufficiently smooth density p X (•, t), we first note that the density satisfies the Fokker-Planck equationEQUATIONFor numerical stability, it is convenient to work with the log-density v := log p X satisfying the equationEQUATION)This nonlinear PDE can be viewed as a Hamilton-Jacobi-Bellman (HJB) equation and its derivation follows from the Hopf-Cole transform, see also Evans (2010) . Motivated by Máté & Fleuret (2023) , we want to use the unnormalized density q target in our parametrization of the log-density v w,z , i.e.,EQUATIONwhere w and z are neural networks approximating the log-density and normalizing constant.If v w,z satisfies (5), i.e., ∂ t v w,z = P σ,µ FP (v w,z ), conservation of mass implies that z(T ) = Z and thus the terminal condition v w,z (•, T ) = log p target is satisfied. Penalizing the square of the PDE residual for v w,z in (5), this yields the lossEQUATIONwhere (ξ, τ ) is a suitable random variable distributed on R d × [0, T ]. As commonly done for PINNs, we chose τ ∼ Unif([0, T ]) and ξ ∼ Unif(K) for a sufficiently large compact set K ⊂ R d (that can potentially depend on τ ). We refer to Appendix E.1 for other approaches.Remark 2.1 (Deterministic vs. stochastic evolutions). Setting the diffusion coefficient σ to zero in (3) yields an ordinary differential equation (ODE) with a deterministic evolution. When the drift µ is controlled by a neural network, this is referred to as continuous-time normalizing flow (Papamakarios et al., 2021; Rezende & Mohamed, 2015) . The PDE in (4) reduces to the continuity equationEQUATIONand the log-density v satisfies that∂ t v = -div(µ) -∇v • µ.(9) This is computationally cheaper than the stochastic case σ ̸ = 0, where the computation of secondorder derivatives of v is required for the term 1 2 Tr(σσ ⊤ ∂ xx v). Motivated by the success of diffusion models (Ho et al., 2020; Kingma et al., 2021; Nichol & Dhariwal, 2021; Vahdat et al., 2021; Song & Ermon, 2020) and their interpretation based on SDEs (Song et al., 2020) , we still want to explore whether the additional noise can be beneficial.There are infinitely many combinations of drifts µ and log-densities v w,z that define a stochastic evolution from the prior to the target distribution, i.e., minimize the loss in (7). The next paragraphs show how we can adapt our loss to yield a unique solution.Annealed flows: As already observed in Máté & Fleuret (2023) , we can just define w ≡ 0 to anneal between p prior and q target , bearing similarity with Annealed Importance Sampling (AIS) (Neal, 2001) and Controlled Monte Carlo Diffusions (CMCD) (Vargas & Nüsken, 2023) . Under mild conditions, we can find both an ODE or an SDE governed by a potential µ = ∇Φ that has the prescribed marginals, see Neklyudov et al. (2022) . Accordingly, this yields the lossEQUATION)Diffusion models: Let us now consider a non-zero diffusion coefficient σ in (3). Moreover, we consider the drift µ = σσ ⊤ ∇v w,z -f for a suitable function f to be specified later. Plugging it into the differential operator in (5), we see thatP σ,µ FP (v w,z ) = div(f ) + ∇v w,z • f -1 2 ∥σ ⊤ ∇v w,z ∥ 2 -1 2 Tr(σσ ⊤ ∂ xx v w,z ) = -P σ,f FP (v w,z ). (11) If v w,z satisfies ∂ t v w,z = P σ,µFP (v w,z ), we thus have thatEQUATIONwhere we write ⃗ σ for the time-reversal, i.e., ⃗ σ(t) = σ(T -t). This means that we can identify ⃗ v w,z = log p Y as the log-density of the processEQUATIONas also derived in Berner et al. (2022) . Thus, a viable strategy is to pick f and σ such that p Y (•, T ) ≈ p prior , see Song et al. (2020) for suitable choices, and minimize the lossEQUATION)Optimal transport and Schrödinger bridges: We can also seek the drift µ that additionally minimizes an energy of the formEQUATIONFor σ = 0, this is connected to optimal transport (OT) problems w.r.t. the Wasserstein metric (Benamou & Brenier, 2000) . For nonzero σ, this corresponds to the dynamic Schrödinger bridge (SB) problem (Dai Pra, 1991) . In these cases, the optimal solution can be written as µ := ∇Φ, where Φ solves the Hamilton-Jacobi-Bellman (HJB) equationEQUATIONsee Appendix B. We can add such regularization using the lossEQUATION)3 EXPERIMENTSIn the following section, we evaluate the above losses and regularizers on different benchmarks. Specifically, we consider the following losses:EQUATION)For the benchmarks, we follow Richter et al. (2023) and consider the following Gaussian mixture model and Many-Well distribution.Gaussian mixture model (GMM): We consider the densityEQUATION)Following Zhang & Chen (2022) , we choose m = 9, Σ i = 0.3 I, and(µ i ) 9 i=1 = {-5, 0, 5} × {-5, 0, 5} ⊂ R 2 (23)to obtain well-separated modes. A typical problem in molecular dynamics considers sampling from the stationary distribution of a Langevin dynamics. In our example we shall consider a d-dimensional many-well potential, corresponding to the (unnormalized) densityEQUATION)with m ∈ N combined double wells and a separation parameter δ ∈ (0, ∞), see also Wu et al. (2020) ; Berner et al. (2022) . Note that, due to the many-well structure of the potential, the density contains 2 m modes. For these multimodal examples, we can compute reference solutions by numerical integration since q target factorizes in the dimensions.We compare against the Path Integral Sampler (PIS) (Zhang & Chen, 2022 ) and the Time-Reversed Diffusion Sampler (DIS) (Berner et al., 2022) , including the log-variance loss by Richter et al. (2023) .We refer to Appendix C and Appendix D for further details on our experiments.Our results are summarized in Table 1 . Comparing our losses in ( 18)-( 21), we observe that the ODE method typically works best. Moreover, we can also outperform our baselines on several metrics.In Figure 1 we also see that it accurately covers the modes of the distributions. On the other hand, we observe that the HJB regularization does not provide substantial improvements, and there is no clear advantage of the methods with prescribed density (i.e., SDE-diffusion and ODE-anneal). This indicates that, in general, the non-uniqueness of the ODE/SDE-objectives does not seem to hurt performance. Finally, the additional computational complexity of considering SDE-based (as opposed to ODE-based) methods appears not to pay off. 
conclusion: 
 We provide a framework for using PINNs to sample from unormalized densities. In particular, we propose to learn the controls of SDEs or ODEs as solutions to (systems of) PDEs that govern their densities. First, this provides a unifying PDE perspective on various sampling methods that are based on normalizing flows, diffusion models, optimal transport, and Schrödinger bridges. Moreover, it yields flexible objectives that are free of time-discretizations and simulations.We benchmark our methods on multimodal target distributions with up to 50 dimensions. While some SDE-based methods are still unstable, ODE-based variants yield competitive methods that can outperform various baselines. We anticipate that our methods can be improved even further using combinations with simulation-based losses as well as common tricks for PINNs, see Appendix D. (Richter et al., 2023) 0.087 6.823 0.8453 0.000 DIS-KL (Berner et al., 2022) 1.785 6.854 0.0225 0.009 DIS-LV (Richter et al., 2023) 1 2001) , and Sequential Monte Carlo (SMC) (Del Moral et al., 2006; Doucet et al., 2009) . However, these methods typically only guarantee asymptotic convergence to the target density, with potentially slow convergence rates in practical scenarios (Robert et al., 1999) . Variational methods, such as meanfield approximations (Wainwright et al., 2008) and normalizing flows (Papamakarios et al., 2021) , offer an alternative approach. In these methods, the problem of density estimation is transformed into an optimization problem by fitting a parametric family of tractable distributions to the target density. In the context of normalizing flows, we want to mention works on constructing better loss functions (Felardos et al., 2023) or gradient estimators (Vaitl et al., 2022) .(d = 50, m = 5, δ = 2) PIS-LVIn this work, we provide a comprehensive PDE perspective on SDE-based sampling methods. Our approach is loosely inspired by Máté & Fleuret (2023) , however, extended to diffusion models, optimal transport (OT), and Schrödinger bridges (SBs). Moreover, we consider other parametrizations and do not rely on the ODE for sampling the collocation points (ξ, τ ). For a corresponding mean-field games (MFG) perspective, we refer to Zhang & Katsoulakis (2023) . We also mention path space measure perspectives on SDE-based methods in Vargas & Nüsken (2023) ; Richter et al. (2023) .The PDE for diffusion models has been derived in Berner et al. (2022) For the usage of PINNs for a generalized SB in the context of colloidal self-assembly, we refer to Nodozi et al. (2023) . An orthogonal direction to our approach is using divergence-free neural networks, which automatically satisfy the continuity equation and only require to fit the boundary distributions p target and p prior (Richter-Powell et al., 2022) . We also mention that higher-dimensional Fokker-Planck equations have also been tackled with time-varying Gaussian mixtures (Chen & Majda, 2018) , and there exist SDE-based neural solvers for HJB equations (Richter & Berner, 2022; Nüsken & Richter, 2021b) and combinations with PINNs (Nüsken & Richter, 2021a) .Finally, we want to highlight recent works on simulation-free learning of (stochastic) dynamics using flow matching (Tong et al., 2023; Lipman et al., 2022) and action matching techniques (Neklyudov et al., 2022) . However, these methods rely on samples from the target distribution p target . Similarly, many works on solving SB and OT problems using deep learning require samples from the target distribution (Chen et al., 2021; De Bortoli et al., 2021; Fernandes et al., 2021; Vargas et al., 2021) . 
 
 Appendix: Let us present a sketch of the proof that the optimal drift can be represented as a gradient field, see also Neklyudov et al. (2022) ; Koshizuka & Sato (2022) ; Benamou & Brenier (2000) ; Caluya & Halder (2021) . Let us consider the optimization problemEQUATION)for a sufficiently smooth density p. Introducing a Lagrange multiplier Φ : R d × [0, T ] → R, we can rewrite the problem asEQUATIONUsing integration by parts, we can calculateEQUATIONUniform We can simply chose τ ∼ Unif([0, T ]) and ξ ∼ Unif(K) for a sufficiently large compact set K ⊂ R d . This choice allows us to uniformly explore the domain K, which is particularly interesting at the beginning of the training. Moreover, different from most other methods, we do not need to rely on (iterative) simulations of the SDE in (3). In order to specify K, however, we need prior information to estimate the domain where v is above some minimal threshold. We incur an approximation error if the set K is chosen too small. On the other hand, if it is too large, low probability areas of p target can lead to instabilities and might require clipping. This could be circumvented by picking a distribution ξ that is supported on the whole spatial domain R d . Alternatively, we can use adaptive methods as outlined in the next paragraphs.Along the Trajectories We can also simulate the SDE using the partially learned drift coefficient µ to exploit the learned dynamics. This corresponds to the choices τ ∼ Unif([0, T ]) and ξ ∼ X τ . Note that we just use the SDE/ODE for sampling the collocation points, and we are not backpropagating through the solver (to update the drift µ). In other words, we detach ξ from the computational graph.Moreover, we want to mention improved sampling strategies for PINNs, see, e.g., Tang et al. (2023) ; Chen et al. (2023) . Similar to Quasi-Monte Carlo methods, one could also leverage low-discrepancy samplers for the time coordinate τ , as, e.g., used by Kingma et al. (2021) .It is commonly known that PINNs can be sensitive to hyperparameter settings. We can make use of a plethora of tricks that have been proposed to stabilize their training (Wang et al., 2023) . For instance, for the neural networks, one could additionally consider random weight factorization and Fourier features for the spatial coordinates. Moreover, we can choose the penalty parameter λ for the HJB loss L HJB adaptively based on the residuals and their gradients. We also mention that the computation of divergences and Laplacians using automatic differentiation can be prohibitive in very high dimensions and might require (stochastic) estimators, such as Hutchinson's trace estimator. Alternatvely, we could also explore the OT-Flow architecture for Φ, which has been successfully employed by Onken et al. (2021) ; Koshizuka & Sato (2022) ; Ruthotto et al. (2020) .We can consider time-dependent diffusion coefficients σ, which have been successfully employed for diffusion models. For instance, we can adapt the VP-SDE in Song et al. (2020) Our framework also allows for diffusion coefficients σ which depend on the spatial coordinate x.Finally, we could also learn the diffusion, for instance, using the parametrization σ = diag(exp(s)) for a neural network s.More generally, we could extend our framework to (stochastic) mean-field games (MFG), mean-field control problems, and generalized SBs (Benamou et al., 2017; Zhang & Katsoulakis, 2023; Liu et al., 2022; Lin et al., 2021; Koshizuka & Sato, 2022; Ruthotto et al., 2020) .We will also consider the case where the diffusion coefficient σ is zero. The evolution in (3) then corresponds to an ordinary differential equation (ODE), however X is still a stochastic process due to its random initial condition given by pprior. 


Decision: The reviewer marks this paper as a clear accept and I agree. It is a nice and novel application of PINNs.

Reviews: review: title: title: review
main_review: main_review: This paper lays out a framework for ODE and SDE distributional sampling using PINNs. It presents the loss terms for different evolution models, and compares PINN to path integral samplers. This is an interesting idea, it's great to see PINNs being used in domains outside of physical modeling. This isn't quite my area, so I can't speak to whether there's similar existing work though.
The benchmarks are rather low-dimensional and simple, but that's fine for an exploration in the context of a workshop.