Solving Poisson Equations using Neural Walk-on-Spheres

TL;DR

This work introduces Neural Walk-on-Spheres (NWoS), a neural PDE solver for high-dimensional Poisson equations that integrates Walk-on-Spheres with a neural network trained via a Feynman-Kac–style loss. NWoS replaces costly time-discretized Brownian simulations with sphere-based sampling and Green’s-function corrections, achieving end-to-end learning of $u$ over the whole domain without requiring spatial derivatives in the loss. The approach yields superior accuracy, speed, and memory efficiency compared to PINNs, the Deep Ritz method, and diffusion/BSDE-based losses, and scales to parametric and PDE-constrained optimization problems. The method shows strong empirical performance across Laplace/Poisson benchmarks, high-dimensional committor problems, and PDE-constrained optimization, highlighting its potential for practical applications in molecular dynamics and related fields.

Abstract

We propose Neural Walk-on-Spheres (NWoS), a novel neural PDE solver for the efficient solution of high-dimensional Poisson equations. Leveraging stochastic representations and Walk-on-Spheres methods, we develop novel losses for neural networks based on the recursive solution of Poisson equations on spheres inside the domain. The resulting method is highly parallelizable and does not require spatial gradients for the loss. We provide a comprehensive comparison against competing methods based on PINNs, the Deep Ritz method, and (backward) stochastic differential equations. In several challenging, high-dimensional numerical examples, we demonstrate the superiority of NWoS in accuracy, speed, and computational costs. Compared to commonly used PINNs, our approach can reduce memory usage and errors by orders of magnitude. Furthermore, we apply NWoS to problems in PDE-constrained optimization and molecular dynamics to show its efficiency in practical applications.

Figures (10)

• Figure 1: Convergence of the relative $L^2$-error when solving the 10$d$ Laplace equation in \ref{['sec:numerics']} using our considered methods.
• Figure 2: Left: Time-discretization of the solution $X^\xi$ to the SDE in \ref{['eq:sde']} with stopping time $\tau(\Omega, \xi)$ in \ref{['eq:stopping_time']} for the domain $\Omega =[0,1]^2$. Right: Realization of the Walk-on-Spheres algorithm in \ref{['sec:nwos']}.
• Figure 3: Peak GPU memory usage of different methods during training with batch size $512$ for the Poisson equation in \ref{['sec:numerics']} in different dimensions $d$.
• Figure 4: Neural Walk-on-Spheres (NWoS): Our algorithm for learning the solution to Poisson equations $\Delta u = f$ on $\Omega\subset \mathop{\mathrm{\mathbb{R}}}\nolimits^d$ and $u|_{\partial \Omega} = g$. In each gradient descent step, we sample a batch of random points $(\xi_0^i)_{i=1}^m$ in the domain $\Omega$ and simulate Brownian motions by iteratively sampling $\xi_k^i$ from spheres $B_{r^i_k}$ inscribed in the domain. To account for the source term $f$, we sample $\gamma_k^i \sim \mathcal{U}(B_{r^i_k})$ to compute an MC approximation $|B_{r^i_k}|f(\gamma_k^i)G_{r_k^i}(\xi_k^i,\gamma_k^i)$ to the solution of the Poisson equation on the sphere $B_{r^i_k}$ using the Green's function $G_{r_k^i}$ in \ref{['sec:source']}. We stop after a fixed number of maximum steps $K$ and either evaluate our neural network $v_\theta$ or the boundary condition $g$ if we reach an $\varepsilon$-shell of $\partial \Omega$. If $v_\theta$ satisfies the PDE, the mean-value property implies that $v_\theta(\xi_0^i)$ is approximated by the expected value of $y^i$ minus the accumulated source term contributions. We thus minimize the corresponding mean squared error over the parameters $\theta$ using gradient descent.
• Figure 5: Qualitative assessment of the solution to the PDE-constrained optimization problem. (Left) Initial function $u_c$ for random parameters $c\in D$. (Middle) Predicted function $u_{\hat{c}}$ for the parameters ${\hat{c}}$ obtained after a few gradient descent steps using the approximation of the solution to the parametric Poisson equation obtained with NWoS. (Right) The groundtruth solution $u_{c^*}$.

Theorems & Definitions (1)

• Remark 4.1