Monte Carlo Neural PDE Solver for Learning PDEs via Probabilistic Representation

TL;DR

This work presents MCNP Solver, an unsupervised neural PDE solver that leverages the probabilistic Feynman-Kac representation to treat macroscopic PDE dynamics as ensembles of particle trajectories. By combining Heun's method for convection with a PDF-based diffusion expectation, MCNP forms a neural Monte Carlo loss that learns a function-to-function operator capable of generalizing across initial conditions without data supervision. The method achieves superior accuracy and efficiency compared with existing unsupervised baselines across convection-diffusion, Allen-Cahn, Navier-Stokes, and mesh-free fractional diffusion tasks, including a disk geometry, and demonstrates notable GPU speedups over traditional solvers. The work also discusses limitations and future directions, such as extending to higher dimensions and higher-order PDEs, and provides publicly available code for reproducibility.

Abstract

In scenarios with limited available data, training the function-to-function neural PDE solver in an unsupervised manner is essential. However, the efficiency and accuracy of existing methods are constrained by the properties of numerical algorithms, such as finite difference and pseudo-spectral methods, integrated during the training stage. These methods necessitate careful spatiotemporal discretization to achieve reasonable accuracy, leading to significant computational challenges and inaccurate simulations, particularly in cases with substantial spatiotemporal variations. To address these limitations, we propose the Monte Carlo Neural PDE Solver (MCNP Solver) for training unsupervised neural solvers via the PDEs' probabilistic representation, which regards macroscopic phenomena as ensembles of random particles. Compared to other unsupervised methods, MCNP Solver naturally inherits the advantages of the Monte Carlo method, which is robust against spatiotemporal variations and can tolerate coarse step size. In simulating the trajectories of particles, we employ Heun's method for the convection process and calculate the expectation via the probability density function of neighbouring grid points during the diffusion process. These techniques enhance accuracy and circumvent the computational issues associated with Monte Carlo sampling. Our numerical experiments on convection-diffusion, Allen-Cahn, and Navier-Stokes equations demonstrate significant improvements in accuracy and efficiency compared to other unsupervised baselines. The source code will be publicly available at: https://github.com/optray/MCNP.

Figures (12)

• Figure 1: Illustration of Feynman-Kac law and the random walks of particles when hitting different boundaries.A: $M$ particles $\{{\boldsymbol{\xi}}_{s}^m\}_{m=1}^M$ begin at the location $\boldsymbol{x}$, and conduct the random walk according to Eq. \ref{['eq:bsde']} from $t=T$ to $t=0$. When $t=0$, the particles query the value at $u_0$, and their average value can be considered as the approximation of $u_T(x)$ (Eq. \ref{['eq:proba_rep_u']}). B: Monte Carlo methods can naturally encode boundary conditions in the random walk of particles. For periodical/Dirichlet/Neumann boundary conditions, the random walks of particles need to be pulled back/stopped/reflected when hitting the boundaries, respectively.
• Figure 2: Illustration of neural Monte Carlo loss. We construct the training loss via the relationship between ${u}_{t}$ and $u_{t+\Delta t}$ given by the Feynman-Kac law. Given a grid point at $\boldsymbol{x}_p$, we split the random walk in Eq. \ref{['eq:bsde_time_inver']} as the convection and diffusion parts. For the convection process, we utilize the Heun's method to simulate the particle moving from time $t+\Delta t$ to $t$ driven by the drift term $\boldsymbol{\beta}$. For the diffusion process, we calculate the mathematical expectation in Eq. \ref{['eq:sde_g']} through the probability density function (PDF) of neighbouring grid points to eliminate sampling using Monte Carlo methods, where $p_{t}(\boldsymbol{\xi}_{p,t}^{d, i})$ denotes the transition probability for the particle moving from $\boldsymbol{\xi}_{p,t}^{c}$ to $\boldsymbol{\xi}_{p,t}^{d, i}$ driven by the diffusion effect. Please note that we omit external forcing $f$ in the figure for simplification.
• Figure 3: The Choice of Neighbourhood Radius $r$ A: We first choose the smallest possible value of $r$ that satisfies Eq. \ref{['eq:r']}. If we have $\sum_i p_t(\boldsymbol{\xi}_{p,t}^{d,i}) \delta \approx 1$, using Eq. \ref{['eq:expectation_nfk']} directly to approximate the corresponding mathematical expectation can reduce computational cost while ensuring accuracy. $\delta$ denotes the volume of each cell (as shown in the red box) in the coordinate system. B: When the grid size is close to the radius $r$, we may encounter the scenario that $\sum_i p_t(\boldsymbol{\xi}_{p,t}^{d,i}) \delta \ll 1$. To address this issue, we interpolate the coordinate system to a high-resolution one to satisfy the normalization condition in Eq. \ref{['eq:r']}. $\Bar{\delta}$ denotes the volume of each cell in the high-resolution coordinate system.
• Figure 4: Simulation of 1D convection-diffusion equation. The prediction result (Left) and point-wise error (Right) of MCNP-10 for an example in E1-E6. The x-axis and y-axis represent spatial coordinates and the predicted values (point-wise error).
• Figure 5: Simulation of 1D Allen-Cahn equation. The prediction result (Left) and point-wise error (Right) of MCNP-100 for an example in E1-E6. The x-axis and y-axis represent spatial coordinates and the predicted values (point-wise error).