Solving Poisson Equations using Neural Walk-on-Spheres

Authors: hong  chul
julius  berner
anima  anandkumar
Venue: ICLR_cc_2024_Workshop_AI4DiffEqtnsInSci
Published: 2024-02-11 05:23:06

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, (backward) stochastic differential equations, and neural cache. In several challenging, high-dimensional numerical examples, we demonstrate the superiority of NWoS in terms of 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 the context of PDE-constrained optimization as well as molecular dynamics to show its efficiency in practical applications

Main paper: 
 introduction: 
 Partial Differential Equations (PDE) are foundational to our modern scientific understanding in a wide range of domains. While decades of research have been devoted to this topic, numerical methods to solve PDEs still remain expensive for many PDEs. In recent years, deep learning has helped to accelerate the solution of PDEs (Azzizadenesheli et al., 2023; Zhang et al., 2023b; Cuomo et al., 2022) as well as tackle PDEs, which had been completely out of range for classical methods (Han et al., 2018; Scherbela et al., 2022; Nüsken & Richter, 2021b) .In this work, we focus on high-dimensional Poisson equations on general domains. Several deep learning methods (Raissi et al., 2019; Sirignano & Spiliopoulos, 2018; E et al., 2017; Han et al., 2017; Nüsken & Richter, 2021a; Han et al., 2020 ) are amendable to the numerical solution of Poisson equations. However, previous methods suffer from unnecessary high computational costs, bias, or instabilities, see Appendix A for a detailed comparison. To overcome these challenges, we propose a novel approach based on so-called Walk-on-Spheres (WoS) methods (Muller, 1956) . WoS rewrites the solution as an expectation over Brownian motions stopped at the boundary of the domain. Leveraging the isotropy of Brownian motion, WoS accelerates the random walk by iterative sampling from spheres around the current position until reaching the boundary, see Figure 2 .Our approach: We develop Neural Walk-on-Spheres (NWoS), a version of WoS that can be combined with neural networks to learn the solution to (parametric families of) Poisson equations on the whole domain. Our method amortizes the cost of WoS during training so that the solution, as well as gradients, can be evaluated in fractions of seconds afterward (and at arbitrary points in the domain). The resulting objective is more efficient and scalable than competing methods, without the need to balance penalty terms for the boundary condition or compute spatial derivatives. In particular, we demonstrate a significant reduction of GPU memory usage in comparison to PINNs and up to orders of magnitude better performance for a given time and compute budget, see Table 1 and Figure 1 .We provide an in-depth comparison of competing deep learning approaches for the solution of elliptic PDEs in Appendix A. These include physics-informed neural networks (PINNs) (Raissi et al., 2019; Sirignano & Spiliopoulos, 2018) , the Deep Ritz method (Jin et al., 2017) , and the diffusion loss (Nüsken & Richter, 2021a) , see also Table 1 . The diffusion loss can be viewed as an interpolation between PINNs and losses based on backward SDEs (BSDEs) (Han et al., 2017; E et al., 2017; Beck et al., 2019) . Methods based on BSDEs and the Feynman-Kac formula (Beck et al., 2018; Berner et al., 2020a; Richter & Berner, 2022) have been investigated for the solution of parabolic PDEs, where the SDE is stopped at a given terminal time. Due to costly simulation time, they cannot be applied efficiently to elliptic problems. 
neural walk-on-spheres (nwos) method: 
 In this section, we derive our method. We consider the Poisson equation with Dirichlet boundary condition on an open, connected, and sufficiently regular domain Ω ⊂ R d , given by ∆u = f, on Ω, u = g, on ∂Ω.(1)It is well-known that the solution admits the stochastic representationEQUATIONwhere ξ is a random variable distributed on Ω, dX t = √ 2 dB t is a scaled Brownian motion with initial condition X ξ 0 = ξ, and τ := τ (Ω, ξ) = inf{t ∈ [0, ∞) : X ξ t ̸ = Ω} is the first exit time, see Appendix A for further details. Directly leveraging this representation results in long runtimes, since it requires to simulate the Brownian motion until reaching the boundary. To tackle this issue, we cast the solution of the PDE in (1) into nested subproblems of solving Poisson equations on subdomains. Specifically, let Ω 0 ⊂ Ω be an open sub-domain containing2 ξ 0 := ξ and let τ 0 := τ (Ω 0 , ξ) be the corresponding stopping time. Analogously to the stochastic representation in (2), we obtain thatu(ξ) = E u(X ξ τ0 ) -F ξ τ0 ξ .(3) Note that this is a recursive definition since the solution u to the PDE in (1) appears again in the expectation. To resolve the recurrence, we define the random variable ξ 1 ∼ X ξ τ0 and choose another open sub-domain Ω 1 ⊂ Ω containing ξ 1 . Considering the stopping time τ 1 := τ (Ω 1 , ξ 1 ), we can calculate the value of u appearing in the expectationu(X ξ τ0 ) ∼ u(ξ 1 ) = E u(X ξ1 τ1 ) -F ξ1τ1 ξ 1 . We can now iterate this process for k ∈ N and combine the result with (3) to obtain In the above, we used the strong Markov property of the SDE solution as well as the tower property of the conditional expectation, see also Hermann et al. (2020) .EQUATIONPicking Ω k := B r k (ξ k ) to be a ball of radius r k ∈ (0, ∞) around ξ k in the k-th step, the isotropy of Brownian motion ensures thatξ k+1 ∼ X ξ k τ k ∼ U(∂B r k (ξ k )).In other words, we can just sample ξ k+1 uniformly from a sphere of radius r k around the previous value ξ k . To terminate after finitely many steps, we pick the maximal radius in each step, i.e., r k := dist(ξ k , ∂Ω), and stop when reaching an ε-shell, i.e., when r κ < ε for a prescribed ε ∈ (0, ∞). This allows us to "walk" from sphere to sphere until (approximately) reaching the boundary, such that we can estimate the first term in (4). Specifically, the value g(X ξ τ ) in ( 4) is approximated by the boundary value g( ξκ ), where ξκ := arg minx∈∂Ω ∥x -ξ κ ∥is the projection to the boundary, see Figure 2 and Figure 4 in the appendix. We note that the bias from introducing the stopping tolerance ε can be estimated as O(ε) (Mascagni & Hwang, 2003) . Moreover, for well-behaved, e.g., convex, domains Ω, the average number of steps κ behaves like O(log(ε -1 )) (Motoo, 1959; Binder & Braverman, 2012) . This shows that ε can be chosen sufficiently small without incurring too much additional computational cost. We note that this leads to much faster convergence than time-discretizations of the Brownian motion. In particular, to have a comparable bias, we would need to take steps of size O(ε), requiring Ω(ε -2 ) steps to converge.Finally, we note that the second term in (4) can be estimated using Green's functions G r k on the domains Boggio (1905) ; Gazzola et al. (2010) and Appendix B.Ω k = B r k (ξ k ), seeBased on the previous derivations, we can establish a variational formulation, where the minimizer is guaranteed to approximate the Poisson equation in (1) on the whole domain Ω. Specifically, we defineL NWoS [v] := E v(ξ) -WoS(ξ) 2 ,(5) where the single-trajectory WoS method WoS(ξ) with random initial point ξ is given byWoS(ξ) := g( ξκ ) - κ-1 k=0 |B r k (ξ k )|f (ξ k )G r k (γ k , ξ k ).In the above, γ k ∼ U(B r k (ξ k )), and the random variables κ, ξ k , ξκ , and r k are defined as in Section 2. From the stochastic formulation in (4) and Proposition 3.5 by Hermann et al. (2020) i k ∼ U(B r i k ) to compute a MC approximation |B r i k |f (ξ i k )G(ξ i k , γ i k )to the solution of the Poisson equation on the sphere B r i k using the Green's function G in Appendix B. 2 We stop after a fixed number of maximum steps κ and either evaluate our neural network v θ or the boundary condition g if we reach an ε-shell of ∂Ω. 3 If v θ satisfies the PDE, the mean-value property implies that v θ (ξ i 0 ) equals the expected value of y i minus the source term contributions. We thus minimize the corresponding mean squared error over the parameters θ. the minimizer of (5) approximates the solution u in (1) in the uniform norm up to error O(ε), where ε is the stopping tolerance. We refer to Appendix C for further theoretical results. We also remark that, in theory, the loss requires only a single WoS trajectory per sample of ξ since the minimizer of the regression problem in (5) averages out the noise. 
experiments: 
 In this section, we compare the performance of NWoS, PINN, DeepRitz, diffusion loss, and Neural Cache (Li et al., 2023) on various baselines across dimensions from 10d to 50d. We do not consider the FK and BSDE losses since they incur prohibitively long runtimes for simulating the SDEs with sufficient precision. To compare against the baselines, we consider benchmarks from the works proposing the Deep Ritz and diffusion loss Jin et al. (2017); Nüsken & Richter (2021a) . To have a fair comparison, we use a fixed runtime of 25d + 750 seconds and a GPU memory budget of 2GiB for training. Moreover, we ran a grid search over a series of hyperparameter configurations for each method. Then, we performed 5 independent runs for the best configurations w.r.t. the relative L 2 -error. More details on the hyperparameters and our implementations3 can be found in Appendix E.We present our results in Table 2 . We first note that we improve the Deep Ritz method as well as the diffusion loss by almost an order of magnitude compared to the results reported by Jin et al. (2017) ; Nüsken & Richter (2021a) . Still, our NWoS approach can outperform all other methods on our considered benchmarks. In addition to these results, we highlight that the efficient objective of NWoS also leads to faster convergence, see Figure 1 and Appendix H. We provide ablation studies in Appendix G and additional numerical evidence for the Poisson equation in 100d and 500d in Appendix I.We also highlight that the efficient objective of NWoS also leads to fast convergence and less memory footprint, see Figure 1 . The PDE-constrained optimization problem shows that this allows NWoS to Table 2: Relative L 2 -error (and standard deviations over 5 independent runs) of our considered methods, estimated using MC integration on 10 6 uniformly distributed (unseen) points in Ω.Method Problem Laplace (10d) Commitor (10d) Poisson Rect. (10d) Poisson (50d) PINN 7.42e -4 ± 1.84e -4 4.10 -3 ± 1.11e -3 1.35e -2 ± 1.57e -3 7.70e -3 ± 2.25e -3 Deep Ritz8.43e -4 ± 6.29e -5 6.15e -3 ± 5.30e -4 1.06e -2 ± 6.20e -4 1.05e -3 ± 1.70e -4 Diffusion loss 1.57e -4 ± 7.74e -6 4.48e -2 ± 6.93e -3 9.69e -2 ± 1.03e -2 5.96e -4 ± 1.06e -5 Neural Cache 3.99 -4 ± 4.08e -5 1.26e -3 ± 5.82e 50e -3 ± 7.38e -4  NWoS (ours) 4.29e -5 ± 2.02e -6 6.56e -4 ± 2.42e -5 2.60e -3 ± 9.99e -5 4.82e -4 ± 1.32e -5   0 scale to parametric problems, where a whole family of Poisson equations is solved simultaneously. We observe that for this 5-dimensional problem (two spatial dimensions and three-parameter dimensions), NWoS converges within 20 minutes to a relative L 2 -error of 0.79% (averaged over D × Ω). More detailed results are shown in Appendix F.-5 4.98e -2 ± 1.80e -2 1.63e -2 ± 1.42e -2 WoS 1.08e -3 ± 1.34e -6 1.99e -3 ± 9.79e -6 2.32e -1 ± 2.09e -1 4.In addition to solving PDE problems, we perform experiments on solving a trajectory optimization problem with PDE constraint, where the solution is to find the optimal control to achieve the target space while satisfying the PDE constraint. More detailed explanations are shown in Appendix F. It shows that NWoS can be extended to parametric problems, where a whole family of Poisson equations is solved simultaneously. We observe that for this 5-dimensional problem (two spatial dimensions and three-parameter dimensions), NWoS converges within 20 minutes to a relative L 2 -error of 0.79% (averaged over D × Ω). The trained network can then be used to solve the optimization problem directly (where we use L-BFGS) without requiring an inner loop for the PDE solver. The results show a promising relative ℓ 2 -error of 0.039 % for estimating the parameters c * leading to an accurate prediction of the minimizer, see Figure 3 . 
conclusion: 
 We have developed Neural Walk-on-Spheres, a novel way of solving high-dimensional Poisson equations using neural networks. Specifically, we provide a variational formulation with theoretical guarantees that amortizes the cost of the standard Walk-on-Spheres algorithm to learn solutions on the full underlying domain. The resulting estimator is more efficient than competing methods (PINNs, Deep Ritz method, Neural Cache, WoS, and diffusion loss) while achieving better performance at lower computational costs as well as faster convergence. We show that NWoS also performs better on a series of challenging, high-dimensional problems and parametric PDEs. This also highlights its potential for applications where such problems are prominent, e.g., in molecular dynamics and PDE-constraint optimization. 
 
 Appendix: We start by defining our problem as well as describing previous deep learning methods for its solution. Our goal is to approximate the solution4 u ∈ C(Ω) to elliptic PDEs with Dirichlet boundary conditions of the form In the following, we will summarize existing neural PDE solvers for these PDEs, see Table 1 for an overview. On a high level, they propose different variational formulations min v∈V L[v] with the property that every minimizer over a suitable function space V ⊂ C(Ω) is a solution u to the PDE in (1). The space V is then typically approximated by a set of neural networks with a given architecture, such that the minimization problem can be tackled using variants of stochastic gradient descent.EQUATIONLet us start with methods based on strong or weak formulations of the PDE in (6).Physics-informed neural networks (PINNs). In its basic form, the loss of PINNs (Raissi et al., 2019) or Deep Galerkin methods (Sirignano & Spiliopoulos, 2018) , is given byEQUATIONwhereL bnd [v] := E (v(ζ) -g(ζ)) 2 .In the above, β ∈ (0, ∞) is a penalty parameter, and ξ and ζ are a suitable random variables distributed on Ω and ∂Ω, respectively. While improved sampling methods have been investigated, see, e.g., Tang et al. (2023) ; Chen et al. (2023) , the default choice is to pick uniform distributions. Given a set of samples, the expectations are approximated with standard MC estimators.By minimizing the point-wise residual of the PDE, PINNs have gained popularity as a universal and simple method. However, PINNs are sensitive to hyperparameter choices, such as β, and suffer from training instabilities or high variance (Wang et al., 2021; Krishnapriyan et al., 2021; Nüsken & Richter, 2021b) . Moreover, the objective in (7) requires the evaluation of the derivatives appearing in P [v] . While this can be done exactly using automatic differentiation, it leads to high computational costs, see Figure 1 .Deep Ritz method. For the Poisson equation in (1) one can avoid this cost by leveraging weak variational forms, see, e.g., Evans (2010) . Rather than directly optimizing the regression loss in (7), the Deep Ritz method (E et al., 2017) proposes to minimize the objectiveEQUATIONUnder suitable assumptions, the minimizer again corresponds to the solution to the PDE in (1). However, the objective only requires computing the gradient ∇v instead of the Laplacian ∆v. Using backward mode automatic differentiation, this reduces the number of backward passes from d + 1 to one, see also the reduced cost in Figure 1 . Moreover, we note that the loss in (8) allows for weak solutions that are not twice differentiable. We refer to Chen et al. (2020) , for an extensive comparison of the Deep Ritz method to PINNs for elliptic PDEs with different boundary conditions.Finally, we mention that both methods suffer from the fact that the interior losses only consider local, pointwise information at x ∈ Ω. At the beginning of the training, the interior loss might thus not be meaningful. Specifically, the boundary condition g first needs to be learned via the boundary loss L bnd , and then propagate from the boundary ∂Ω to an interior point x via the local interior loss. There exist some heuristics to mitigate this issue by, e.g., progressively learning the solution, see Penwarden et al. ( 2023) for an overview. The next section describes more principled ways of including boundary information in the loss and directly informing the interior points of the boundary condition.From weak solutions, we will now proceed to stochastic representations of elliptic PDEs in (6). To this end, consider the5 solution X ξ to the stochastic differential equationEQUATION)where W is a standard d-dimensional Brownian motion. Moreover, we define the stopping time τ as the first exit time of the stochastic process X ξ from the domain Ω, i.e.,EQUATION)An application of Itô's lemma to the process u(X ξ t∧τ ) shows that we almost surely have thatu(X ξ τ ) = u(X ξ 0 ) + τ 0 P[u](X ξ t ) dt + S u τ ,where S u τ is the stochastic integralS u τ := τ 0 (σ ⊤ ∇u)(X ξ t ) • dW t .Using the fact that X ξ 0 = ξ and assuming that u solves the elliptic PDE in (6), we arrive at the formulaEQUATIONwhere we used the abbreviationF ξ τ := τ 0 f (X ξ t ) dt.Since the stochastic integral S τ u has zero expectation, see, e.g., Baldi (2017, Theorem 10 .2), we can rewrite (11) as a stochastic representation, i.e.,EQUATIONwhich goes back to Kakutani's Theorem (Kakutani, 1944) and is a special case of the Feynman-Kac formula.While the above representation allows us to establish MC methods for the pointwise approximation of u at a given point x ∈ Ω, it also allows us to derive variational formulation for learning u on the whole domain Ω. Based on the above results, we can derive the following three losses.Feynman-Kac loss. The Feynman-Kac loss is given byEQUATIONand follows from the fact that the solution to a quadratic regression problem as in ( 13) is given by the conditional expectation in (12). Notably, this variational formulation does neither require a derivative of the function v nor an extra boundary loss L bnd .Figure 4 : Left: Illustration of the solution X ξ to the SDE in ( 9) and its stopping time τ (Ω, ξ) in ( 10) for the domain Ω = [0, 1] 2 . Right: Realization of the Walk-on-Spheres algorithm in Section 2.BSDE loss. Since the formula in (11) holds if and only if u solves the PDE in (6), we can derive the BSDE lossL BSDE [v] := E v(ξ) -g(X ξ τ ) + F ξ τ + S v τ 2 .(14) Compared to the Feynman-Kac loss in (13), the BSDE loss requires computing the gradient of v at every time discretization of the SDE X ξ in order to compute S τ v . However, due to (11), S τ v acts as a control variate and causes the variance of the MC estimator of ( 14) to vanish at the optimum, see Richter & Berner (2022) for details.For the previous two losses, boundary information is directly propagated along the trajectory of the SDE X ξ to the interior. However, simulating a batch of realizations of the SDE until they reach the boundary ∂Ω, i.e., until the stopping time τ , can incur prohibitively high costs.Diffusion loss. The diffusion loss (Nüsken & Richter, 2021a ) circumvents long simulation times by stopping the SDE at s = τ ∧ T , i.e., at the minimum of a prescribed time T ∈ (0, ∞) and the stopping time τ . Since the trajectories might not reach the boundary, it is required to supplement the loss with a boundary loss. This yields the variational formulationL Diff [v] := E v(ξ) -v(X ξ s ) + F ξ s + S v s 2 + βL bnd [v].Note that this can be viewed as an interpolation between the BSDE loss (for s → ∞) and the PINN loss (for s → 0 and rescaling by s -2 ). In the same way, it also balances the advantages and disadvantages of both losses, see also Table 1 .To compute the second term in (4), we need to accumulate values of the formEQUATIONwith a given ball B = B r (z). By ( 12), we observe that v is just the solution of a Poisson equation on the ball B with zero Dirichlet boundary condition evaluated at z.By (12), we observe that v is just the solution of a Poisson equation on the ball B with zero Dirichlet boundary condition evaluated at z. We can thus use classical results by Boggio (1905) , see also Gazzola et al. (2010) , to write the solution in terms of Green's functions. Specifically, we have thatEQUATIONwhere γ ∼ U(B r (z)) andG r (y, z) := 1 2π log r ∥y-z∥ , d = 2, Γ(d/2-1) 4π d/2 ∥y -z∥ 2-d -r 2-d , d > 2,In practice, we can now approximate the expectation in ( 16) using an MC estimate.To compute integrals of the form (15), we look at a special case of a Poisson equation on a ball B = B r (z) with zero Dirichlet boundary condition, i.e., ∆v = f, on B, v = 0, on ∂B. Analogously to (12), we obtain thatEQUATIONwhere τ (B, z) is the corresponding stopping time, see (10). However, since we simplified the domain to a simple ball, we can write the solution in terms of Green's functions. Specifically, we have thatEQUATIONwhereG r (y, z) := 1 2π log r ∥y-z∥ , d = 2, Γ(d/2-1) 4π d/2∥y -z∥ 2-d -r 2-d , d > 2. We note that ( 18) is equivalent to (16).While this is a classical result by Boggio (1905) , see also Gazzola et al. (2010) , we will sketch a proof in the following. We consider the Laplace equation ∆Φ x = δ x for given x ∈ R d in the distributional sense. It is well known that the fundamental solution Φ x is given byΦ x (y) = 1 2π log ∥y -x∥, d = 2, -∥y-x∥ 2-d (d-2)ω d , d > 2,whereω d = |∂B 1 (0)| = 2π d/2 Γ(d/2) = 4π d/2 (d -2)Γ(d/2 -1)is the surface measure of the d-dimensional unit ball B 1 (0). Under suitable conditions, it further holds that the solution to ( 17) is given byEQUATIONfor every x ∈ B, where the corrector function ϕ x satisfies the Laplace equation ∆ϕ x = 0, on B, ϕ x = Φ x , on ∂B. see Evans (2010, Chapter 2.2). Based on (12) and the fact that Φ z is constant at the boundary of B = z + B r (0), we can compute the value of the corrector function ϕ z , i.e.,ϕ z (y) = E[Φ z (X y τ (B,y) )] = 1 2π log r, d = 2, -r 2-d (d-2)ω d , d > 2.This shows that the value of the Green's function at the center z of the ball B is given by Φ z (y) -ϕ z (y) = -G r (y, z), which, together with (19), establishes the claim.For numerical stability, we directly compute the quantity Gr (γ, z) := |B r (z)|G r (γ, z) in practice, as needed in ( 16). The volume of the hyper-sphere |B r (z)| is given by|B r (z)| = π d/2 Γ(d/2 + 1) r d ,such that we obtainGr (γ, z) :=    r 2 2 log r ∥γ-z∥ , d = 2, r d d(d-2) ∥γ -z∥ 2-d -r 2-d , d > 2.Having established the learning problem in Section 2.2, we can analyze both approximation and generalization errors. For the former, Hermann et al. (2020) bounded the size of neural networks v θ to approximate the solution u up to a given accuracy. In particular, the number of required parameters θ only scales polynomially in the dimension d and the reciprocal accuracy, as long as the functions f , g, and dist(•, ∂Ω) can be efficiently approximated by neural networks.One can then leverage results by Berner et al. (2020b) to show that also the generalization error does not underlie the curse of dimensionality when minimizing the empirical risk, i.e., a MC approximation of ( 5), over a suitable set of neural networks v θ . Specifically, the number of required samples of ξ to guarantee that the empirical minimizer approximates the solution u up to a given accuracy also scales only polynomially in the underlying dimension and the reciprocal accuracy.In this section, we discuss several strategies to trade off accuracy and computational cost, and to reduce the variance of MC estimators of L NWoS in 5.For sufficiently regular geometries, the probability for a walk to take more than k steps, is exponentially decaying in k (Binder & Braverman, 2012) . However, if a single walk in our batch needs significantly more steps, it slows down the overall training. We thus introduce a deterministic maximum number of steps κ ∈ N. However, we do not want to introduce non-negligible bias by just projecting to the closest point on the boundary.Instead, we want to enforce the mean-value property on subdomains of Ω based on our recursion in Section 2. We thus propose to use the model v instead of the boundary condition g if the walk does not converge after κ steps, i.e., we definey ξ,v := v(ξ κ ), d(ξ κ , ∂Ω) > ε, g( ξκ ), else.We can then replace the second term in (5) byWoS(ξ, v) := y ξ,v - κ-1 k=0 |B r k (ξ k )|f (ξ k )G r k (γ k , ξ k ).This helps to reduce the bias when d(ξ κ , ∂Ω) is non-negligible while exploiting the faster convergence assuming that we obtain increasingly good approximations v θ ≈ u during training of a neural network v θ . Our approach bears similarity to the diffusion loss, see Appendix A, however, we do not need to use a time-discretization of the SDE.Boundary Loss: We find empirically that an additional boundary loss can further improve the performance of our method. While it is theoretically not required to converge to the correct solution, it can especially help for a smaller number κ of maximum steps (see the previous paragraph). In general, we thus sample a fraction of the points on the boundary ∂Ω and optimizeL NWoS [v] + βL bnd [v],where L bnd is defined6 as in (7).Variance-Reduction: While not necessarily needed for the objective in (5), we can still compute multiple WoS trajectories N ∈ N per sample of ξ to reduce the variance. This leads to the estimatorL NWoS [v] := 1 m m i=1 v(ξ i ) - 1 N N n=1 WoS n (ξ i ) ,where ξ i are i.i.d. samples of ξ and WoS n (ξ i ) are i.i.d. samples of WoS(ξ i ), i.e., N trajectories with the same initial point ξ i , see ( 5). Note that we vectorize the WoS simulations across both the initial points as well as the trajectories, making our NWoS method highly parallelizable and scalable to the large batch sizes.We further introduce control variates to reduce the variance of estimating WoS(x), where we focus on a fixed x ∈ Ω for the ease of presentation. Control variates seek to reduce the variance by using the estimatorE [WoS(x)] ≈ E[δ] + 1 N N i=1 WoS n (x) -δ n ,where δ n are i.i.d. samples of a random variable δ with known expectation. Motivated by Sawhney & Crane (2020) , we use an approximation of the first-order term of a Taylor series of u in the direction of the first WoS step. We assume that ∇v θ provides an increasingly accurate approximation of the gradient ∇u during training and propose to use where ξ n 1 is the first step of WoS n (x). In particular, ξ n 1 ∼ U(∂B r1 (x)) and thus E[δ] = 0 holds for any function v θ . While we need to compute the gradient ∇v θ (x) for the control variate, we mention that this operation can be detached from the computational graph. In particular, we do not need to compute the derivative of ∇v θ (x) w.r.t. to the parameters θ as is necessary for PINNs, the Deep Ritz method, the diffusion loss, and the BSDE loss. In Appendix G, we empirically show that the overhead of using the control variate is insignificant.δ n := ∇v θ (x) • (ξ n 1 -x),Buffer: Motivated by Li et al. (2023) , we can use a buffer to cache training points (ξ (i) , WoS(ξ (i) )).Since we only update the buffer after a given number of training steps, this accelerates the training. Note that this is not possible for the other methods since they require evaluation of the current model.In every buffer update, we average over further trajectories for a fraction of points to improve their accuracy. However, different from Li et al. (2023) , we also evict a fraction of points from the buffer and replace them with WoS estimates on newly sampled points in the domain Ω to balance the diversity and accuracy of the training data.We implement all methods in PyTorch and provide pseudocode in Algorithms 1 and 2. The experiments have been conducted on A100 GPUs.For all our training, we use the Adam optimizer and limit the runtime to 25d + 750 seconds for a fair comparison. In every step, we sample uniformly distributed samples (ξ, ζ) in the domain Ω and on the boundary ∂Ω to approximate the expectations of the loss and boundary terms. Moreover, we employ an exponentially decaying learning rate, which reduces the initial learning rate by two orders of magnitude throughout training. We choose a feedforward neural network with residual connections, 6 layers, and width 256, and GELU activation function. We also perform the grid search for the boundary loss penalty term, i.e., β ∈ {0.5, 1, 5, 50, 100, 500, 1000, 5000}.We further include the batch size m ∈ {2 i } 17 i=7 in our grid-search. To have a fair comparison, we set a fixed GPU memory budget of 2GiB for training, leading to different maximal batches for network sizes and methods, see also Figure 1 . Let us detail the hyperparameter choices specific to each method in the following.Neural Walk-on-Spheres (NWoS): For all experiments with NWoS, we set ε = 10 -4 . If using a boundary loss, we sweep over {0.1, 0.2, 0.3, 0.4, 0.5} in the grid search to find the optimal proportion of the batch size for the boundary loss. We choose to do a grid search on control variate, neural target, and buffer to ensure that a combination with the best performance is chosen. We fix the buffer size to 10 times that of the batch size, and we sweep the frequency to update buffers over the grid of {10, 100, 1000}. We also sweep the max step in {0, 1, 5, 10, 50, 100} and the number of trajectories in {1, 10, 100, 200, 300, 400, 500, 1000}.Input: neural network v θ with initial parameters θ, optimizer method step for updating the parameters, WoS method WoS in Algorithm 2, maximum number of iterations T , batch size m, n, b for points in domain, boundary and buffer, buffer B of size B, boundary function g, buffer update period L, boundary term coefficient β Output:optimized parameters θ B ← {(x i , g(x i ))} B i=1 ▷ Sample points in ∂Ω and ▷ initialize the buffer for k ← 0, . . . , T do if k mod L == 0 then x Ω ← sample from ξ ⊗m ▷ Sample points in Ω x B ← sample from B ▷ Samples points in B x ← [x Ω , x B ] ▷ Concatenate points [y Ω , y B ] ← vmap[WoS(x, v θ )] ▷ WoS B ← update with (x B , y B ) ▷ Update the buffer B ← replace with (x Ω , y Ω ) ▷ Replace the buffer end if x ∂Ω ← sample from ζ ⊗n ▷ Sample points in ∂Ω (x B , y B ) ← sample from B ▷ Sample points in B L Ω = MSE(v θ (x B ), y B ) ▷ Domain loss L ∂Ω = MSE(v θ (x ∂Ω ), g(x ∂Ω )) ▷ Boundary loss L NWoS = L Ω + β L ∂Ω ▷ Compute loss θ ← step γ, ∇ θ L NWoS ▷ SGD step end forAlgorithm 2 Walk-on-Spheres (WoS) Input: neural network v θ at step k, source term f , boundary term g, point for evaluation x, maximum number of steps κ, stopping tolerance ε, number of trajectories N Output: estimator v of solution v to PDE in (1) at 2), as well as using our improvements from Appendix D, i.e., the control variate and a neural network evaluation for trajectories that did not converge in the given maximum number of steps κ. Right: Time for training one iteration for all considered methods. For NWoS we present the comparison for different maximum numbers of steps κ.x v ← 0 for i ← 1, . . . , N do ▷ Batched in implementation s ← 0 for k ← 0, . . . , κ -1 do r ← dist(x, ∂Ω) ▷ Compute distance to ∂Ω if r < ε then Break ▷ Reach boundary end if γ ← sample from U(B r (x)) ▷ Estimate source s ← s + |B r (x)|f (x)G r (x, γ) u ← sample from U(∂B r (x)) if k = 0 & use control variate then s ← s -∇ x v θ (x) • u ▷ Variance reduction end if x ← x + u ▷ Walk to next point end for if dist(x, ∂Ω) < ε then ▷ Estimate solution at x v ← s + g(x) else v ← s + v θ (x) end if end for v ← 1 N v ▷ TakeDiffusion loss: For the diffusion loss (Nüsken & Richter, 2021a) , also set ε = 10 -4 . Moreover, we perform grid search for max step s ∈ {1, 5, 10, 50} and time step ∆t ∈ {10 -3 , 10 -4 , 10 -5 }. We use 10% of the batch size for boundary points.PINNs: For PINNs (Raissi et al., 2019; Sirignano & Spiliopoulos, 2018) , we use automatic differentiation to compute the Laplacian ∆v θ . We use 10% of the batch size for boundary points.Deep Ritz: For the Deep Ritz method (E et al., 2017) , we additionally experiment with the network architecture proposed in their paper. We sweep the number of blocks in {4, 6, 8}, the number of layers in {2, 4}, and the hidden dimension in {64, 128, 256}. Moreover, we observed better performance using the GELU activation function. We use 10% of the batch size for boundary points.Neural Cache: For the neural cache method (Li et al., 2023) , we also search cache size over the grid of {10000, 20000, 100000, 1000000}, training period (i.e. period to update the cache) over {1, 10, 100, 1000, 5000, 10000}, the number of trajectories over {1, 10, 20, 30, 40, 50, 100, 500, 1000}. WoS: For Walk-on-Spheres (Muller, 1956) , we directly approximate the solution at the evaluation points. We batch trajectories to saturate the memory budget and present the best result for different configurations within the given runtime. Specifically, we pick the number of trajectories in the grid {1, 10, 100, 1000, 10000, 100000} and the maximum step in {0, 1, 10, 100, 1000}.Let us describe our considered PDEs in the following.The first PDE is a Laplace equation on a square domain given byf (x) = 0, g(x) = d/2 i=0 x 2i x 2i+1 , x ∈ Ω = (0, 1) d .To test our models, we compare against the analytic solution as u(x) = (x) = 1 d d i=1 sin(2πx i ), f (x) = - 4π 2 d d i=1sin(2πx i ).We choose c = 0.25 1 d and d = 10 and note that the analytic solution is given by u(x) = 1 d d i=1 sin(2πx i ).Committor Function:The fourth equation deals with committor functions from molecular dynamics. These functions specify likely transition pathways as well as transition rates between (potentially metastable) regions or conformations of interest (Vanden-Eijnden et al., 2006; Lu & Nolen, 2015) . They are typically high-dimensional and known to be challenging to compute. To compare NWoS, we consider the setting in Nüsken & Richter (2021a) . The task is to estimate the probability of a particle hitting the outer surface of an annulusΩ = {x ∈ R d : a < ∥x∥ < b}, a, b ∈ (0, ∞),before the inner surface. The problem can then be formulated as the Laplace equation given byf (x) = 0, g(x) = 1 {∥x∥=b} , x ∈ Ω.For this specific Ω, a reference solution can be computed asu(x) = a 2 -∥x∥ 2-d a 2 a 2 -b 2-d a 2 .We further use the setting by Nüsken & Richter (2021a) and choose a = 1, b = 2, and d = 10.We want to solve the optimization problem Figure 5 shows the comparison of NWoS with DeepRitz, NSDE, and PINN with different maximum number of steps K, see Appendix E. Taking into account the logarithmic scaling of the plot, the training time of NWoS is significantly faster than both NSDE and PINN while slightly slower than Deep Ritz for higher maximum steps. In particular, the best results in Table Table 2 used K = 10 and outperformed Deep Ritz. However, choosing K we can also trade-off between high-fidelity solutions and fast training.min u∈H 1 0 (Ω),m∈L 2 (Ω) 1 2 Ω (u -u d ) 2 dx + α 2 Ω m 2 dx constraint toIn this section, we illustrate the convergence of the relative L 2 error for other PDEs.Figure 6 and 7 all demonstrate that neural WoS achieves the fastest convergence to the best optimum in comparison to all baseline methods provided time constraints.We further evaluate our method on the Poisson equation in 100d and 500d as proposed by E & Yu (2018) . Table 3 demonstrates that our method is better than the baselines. However, we discover empirically that, for this benchmark, a simple projection to the boundary achieves the highest Table 3 : Relative L 2 -error (and standard deviations over 5 independent runs) of our considered methods, estimated using MC integration on 10 6 uniformly distributed (unseen) points in Ω.Poisson (500d) PINN 1.49e -3 ± 3.21e -5 2.42 -2 ± 6.06e -4 Deep Ritz 1.77e -2 ± 1.94e -4 9.92e -3 ± 2.56e -5 Diffusion loss 6.71e-4 ± 1.31e -5 9.47e -3 ± 3.81e -5 Projection 2.92e -4 ± 5.17e -7 1.19e -5 ± 1.67e -8 NWoS (ours) 6.22e -4 ± 1.18e -5 9.14 -3 ± 6.31e -5accuracy. This can be motivated by the smoothness of the solution and the fact that uniformly distributed evaluation samples concentrate at the boundary in high dimensions.NWoS is currently only applicable to Poisson equations with Dirichlet boundary conditions. While this PDE appears frequently in applications, we also believe that future work can extend our method. While not necessary for NWoS, we note that the gradient of the model and an additional boundary loss can still be used to improve performance, see Appendix D.Since ξ is a random variable, the sub-domain Ω0 is random, and the statement is to be understood for each realization.Our PyTorch code can be found at https://github.com/bizoffermark/neural_wos.For simplicity, we assume that a sufficiently smooth strong solution exists.We assume that there is a unique solution, see, e.g., Le Gall (2016) for corresponding conditions.Note that L bnd can be interpreted as a special case of LNWoS where the WoS method directly terminates since the initial points are sampled on the boundary.While d = 100 is considered byJin et al. (2017), we find that a simple projection outperforms all models in sufficiently high dimensions for this benchmark. 


Decision: I appreciate reviewers' comments on this paper. Both agree that this paper proposes novel approach for solving high dimension Poisson Equation. Authors are required to address the comments and questions of reviewers in their final version. I vote accept for this paper.

Reviews: review: title: title: An interesting work, should be interesting for the community, minor revision.
main_review: main_review: The paper applies a novel strategy for solving high-dimensional Poison equations based on the neural walk-on-spheres technique. The approach and the results are interesting and significant, as I see. Also, I should admire the authors for staying loyal to the unanimous sharing of information, even their code. The paper looks interesting and novel to me. Also, the authors carefully explained the details of the work in most of the sections. They, then, compared the results to other approaches. I have below minor comments:
-	It would help clarify the approach if authors also mention the challenges of the method or cases in which it will have problems. 
-	Is it possible to utilize neural WoS for inverse problems?
-	As far as I know, PINNs performance is really a function of choosing the right collocation points, and also a suitable training strategy. It would be better if the configurations that have been used for PINNs were discussed more carefully in appendix E (and maybe for other methods).
title: title: sound work
main_review: main_review: The proposed estimator effectively integrates boundary conditions without resorting to penalty terms, surpassing the efficiency of alternative methods like PINNs, the Deep Ritz method, and diffusion loss. The paper introduces innovative loss functions for neural networks, leveraging the Green function and recursive solution of Poisson equations within domain spheres, eliminating the need for computing spatial gradients in the loss calculation. The detailed implementation information provided in the appendix offers valuable guidance for readers. This paper demonstrates remarkable technical and mathematical rigor, presenting novel approaches and fully laying on the community's interest. I strongly recommend accepting this paper.