Integration Matters for Learning PDEs with Backwards SDEs

We thank the reviewer for their thoughtful and constructive review. We greatly appreciate the recognition of the importance of comparing PINNs and BSDE methods, and the novelty of our approach.

Response to W1/Q1/Q2/Q3 (mathematical statements and assumptions):

Thanks for bringing this up. We would like to clarify a few points. First, our main theorems (Thm. 4.1 and 4.2) are fully rigorous and include all the necessary assumptions in the theorem statement. On the other hand, the discussions that follow the formal theorem statements (including Eq. (4.3) and Line 903, 905) do omit exact technical assumptions, as we opted for clarity/brevity of exposition—we wanted to convey the message of how our main theorems could be interpreted without getting bogged down into too many technical details which we felt distracted from the main message. Nevertheless, given the feedback, we will add these details back into the draft.

To answer the reviewer’s questions, and to give a sense of how we plan to edit the paper (recall we are not allowed to upload a new draft), we’ll derive Eq. (4.3) more carefully below. This will not only illustrate the kinds of assumptions we make, but also fill in some of the analysis details (e.g., how we go from Riemann sum to integral) that the reviewer is asking about. The details presented here will be added to our revision.

The starting point for Eq. (4.3) is from the proof of Theorem 4.1 (Appendix C.1), where we show that:

where $L_{d,u_\theta}$ is a finite constant that depends on the dimension $d$, in addition to the Lipschitz constant of $D^2 u_\theta$. Note that the latter is bounded, since we assume that $u_\theta \in C^{2,1}$ ($C^{2,1}$ is defined in Line 136).

Hence at this point, we have that:

$\left|L_{\mathrm{EM},\tau}(\theta) - \frac{1}{N\tau^2} \sum_{n=0}^{N-1} \mathbb{E}_{\hat{X}_n}\left[ \tau^2\left( (R[u_\theta]( \hat{X}_n, t_n))^2 + \frac{1}{2} \mathrm{tr}( (H(\hat{X}_n, t_n) \nabla^2 u(\hat{X}_n, t_n))^2 ) \right) \right]\right| \leq L_1\tau^{1/2},$

where $L_1 = L_{d,u_\theta}$ (this was necessary to get openreview to render the previous equation).

Our next step is to argue that the map $(x, t) \mapsto F_\theta(x, t) := (R[u_\theta](x, t))^2 + \frac{1}{2} \mathrm{tr}( (H(x, t) \cdot \nabla^2 u_\theta(x, t))^2 )$ is Lipschitz for some constant $L_{F_\theta}$. This is what the role of “basic regularity assumptions on $R[u_\theta]$, $H$, and $u_{\theta}$” (Line 150) plays; there is not a unique set of assumptions that render the map $F_\theta(x, t)$ Lipschitz. Nevertheless, we specify one sufficient condition now: (a) $H, f, h[u_\theta]$ are all bounded, and (b) $u_\theta$ and all its derivatives up to 2nd order are Lipschitz and bounded. This is certainly not the weakest assumption; one could instead assume growth conditions of all the relevant functions and their derivatives as $(x, t)$ to go infinity, and then combine this with a high probability bound on the trajectories $X_t$ on $t \in [0, T]$. In the interest of simplicity for this rebuttal, we opt for our simpler stated assumption.

With the Lipschitz condition on $F_\theta$ in place, we can now apply order $1/2$ strong convergence of EM. Note that reference [34] (cited on Line 147)---specifically Theorem 10.2.2 of [34]---gives conditions on the drift $f$ and diffusion $g$ of the forward SDE for strong convergence to occur; these conditions are roughly Lipschitz $f, g$ and some growth conditions of $f, g$ as a function of $x$. We then combine strong convergence with Lipschitz $F_\theta$ to argue: $|\mathbb{E}[ F(\hat{X}_n, t_n) - F(X(t_n), t_n) ]| \leq L_F \mathbb{E}\| \hat{X}_n - X(t_n) \| \leq L_F C \tau^{1/2}.$

(we had to drop the subscripts on $\theta$ and write $X_{t_n}$ (from the paper) as $X(t_n)$ to get the equation to render). Note that if in the previous step we derive a Lipschitz constant $L_F(x)$ which depends on $x$, the definition of strong convergence involving second moments becomes important as it allows us to decouple the Lipschitz constant from the bound on $\| \hat{X}_n - X(t_n) \|$ via Cauchy-Schwarz. Combining this together with the previous bound, we have

$\left| L_{\mathrm{EM},\tau}(\theta) - \frac{1}{N}\sum_{n=0}^{N-1} \mathbb{E}[ F(X(t_n), t_n)) ]\right| \leq L_1 \tau^{1/2} + L_F C \tau^{1/2}.$

Our final step is to treat the second term on the LHS as a Riemann sum approximation to a Riemann integral. From standard approximation properties of the left Riemann sum, we have $\left| \frac{1}{T} \int_0^T \mathbb{E}[F(X(t), t)] dt - \frac{1}{N}\sum_{n=0}^{N-1} \mathbb{E}[ F(X(t_n), t_n))] \right| \leq \frac{L_F T}{2N} = \frac{L_F \tau}{2}.$

This verifies the claimed Eq. (4.3): $\left| L_{\mathrm{EM},\tau}(\theta) - \frac{1}{T}\int_0^T \mathbb{E}[ F(X(t), t)) ] dt \right| \leq L_1 \tau^{1/2} + L_F C \tau^{1/2} + L_F \tau/2 = O(\tau^{1/2}).$

We hope this clarifies the conditions under which Eq. (4.3) hold, and makes precise what is assumed in Lines 903, 905, 909. We note that Eq. (4.8) is also derived using similar arguments; we will also add its derivation to the paper.

Response to W2/Q4 (PDE background):

(Note that this question was also brought up by Reviewer sGAy, so we have duplicated our response in both rebuttals to make the rebuttals self-contained).

The non-linear PDEs we consider (cf. Eq. (3.1)) are the standard kind of PDEs considered for BSDE methods (see e.g., references [4-6]). There are many types of applications, including in quantitative finance (e.g., multi-asset options pricing via Black-Scholes, local/stochastic volatility extensions from the work of Duprie and Heston), optimal control (e.g., Hamilton-Jacobi-Bellman equations for finding value functions, Hamilton-Jacobi-Issacs for reach-avoid problems), reaction-diffusion in chemistry (e.g., Allen-Cahn and Fisher-KPP equations), committer/exit-time probabilities for transition path analysis (e.g. [R1]), and so on. We will include a paragraph about these applications to our revised draft.

We also note that another reason we consider the specific elliptic/parabolic form (3.1) given in the paper (i.e., $h[u]$ does not depend on $\nabla^2 u$) is to favor the standard EM-BSDE formulation, since the form (3.1) means that the EM-BSDE approach does not need to compute Hessians of $u$ in the backwards SDE, whereas the Heun-BSDE does. In general, $h[u]$ can also depend on $\nabla^2 u$ as noted in e.g., [Pham et al., 19] referenced by the reviewer, allowing for fully nonlinear second order PDEs.

[R1]: E, Vanden‑Eijnden (2006), Towards a Theory of Transition Paths. J. Stat. Phys.

Response to W3/Q5 (other benchmarks/references):

Thanks for pointing us to these references, and we will further incorporate them into our related work. For ease of reference, we will refer to these papers as [C1] to [C5]. In regards to direct comparison, our work focuses on the impact that the most widely used integration scheme in practice (Euler-Maruyama) has on the performance of BSDEs. Almost all the references above are focused on enhancements to the BSDE loss, and are also compatible with our proposed Heun-BSDE scheme. For example, we could modify our Heun-BSDE method to also utilize a control variate [C3] (i.e., split $u = u_1 + u_2$ each with their own properties), or to use Heun-BSDE in operator splitting settings [C1], or to apply to fully non-linear PDEs [C5], comparing these methods to their EM-BSDE + technique baseline. While this is an interesting direction which we will certainly incorporate in our revision, we believe the merits of our work can be judged on the basis of our current experiments which are focused on a direct comparison of EM vs. Heun, without extra confounding factors.

We next specifically discuss higher order Itô integration schemes and the paper [C4]. The next step up in complexity from a standard Euler scheme is a trapezoidal integration scheme; Heun integration is the explicit form of this. Direct application of Heun to integrate an SDE converges to the Stratonovich solution, which forms the basis of the study in our work. However, this is not the only approach. There are e.g., implicit Crank-Nicholson (CN) schemes and also higher order explicit Runge-Kutta (RK) approaches, as discussed in [C4]. We will certainly consider an empirical comparison in our future draft. However, we note several points:

• First, the Heun-BSDE + Stratonovich formulation presented has several algorithmic advantages compared to other approaches: (a) it is an explicit method, and hence does not require solving implicit equations at every step, (b) it is very simple to implement, unlike e.g., higher order RK methods for Itô which require non-trivial implementations for handling diffusion terms that are not commutative---which is essentially always the case for backwards SDEs where the diffusion term depends on $\nabla u$. Thus, part of our contribution is also to advocate for the advantages of reinterpreting the SDEs as Stratonovich, in addition to using more sophisticated integrators.

• Second, the schemes presented in [C4] are not directly comparable as described, since they are designed for the recursive BSDE loss formulations, where a separate network is trained at every discretization point. On the other hand, in this paper we focus on self-consistency losses (see the discussion starting on Line 116 for the difference), which are more practical as only one network is trained. The integration schemes can be applied in the self-consistency setting as well (which we will investigate in the revision), but this is not what is specifically considered in [C4].

Thank you for your reply. I think this is an excellent suggestion that will greatly clarify the presentation. Overall, the rebuttal has addressed my concerns. I have revised my rating and recommend the paper for acceptance.

We thank the reviewer for their thoughtful and constructive review. We greatly appreciate the recognition of the usefulness of our numerical analysis perspective, with impact possibly beyond the ML community.

Response to W1/Q2/Q3 (Computational Costs):

(Note that similar questions were also brought up by Reviewers LXC6 and Fvez, so we have duplicated our response to make the rebuttals self-contained).

We agree that the computational costs of the Heun integration presented in the original draft are key limitations for practical implementation. Hence, we spent a bit of time investigating how we could improve our implementation. We are happy to report that we have improved our implementation—for all methods—and the result is that the computational overhead of the Heun implementation is much more reasonable. Due to the given rebuttal time constraints, we were unfortunately only able to re-run our new implementation for all methods on the 100D HJB case, but we expect the computational benefits from the implementation changes we made (which we describe below) to carry over to all PDE cases. We will of course re-run all methods over all PDE cases for the revised draft.

In the original Heun-BSDE implementation, the forward and backward SDEs rollouts were performed simultaneously, evaluating $L_{Heun_{k,\tau}}(\theta)$ at each step. As a result, the loss gradient calculation required back-propagation through the entire rollout; we had to resort to gradient checkpointing when combined with the Hessian calculation of $u$ in order to avoid running out of memory. In our new implementation, we first perform a forward SDE rollout to generate trajectories, $\mathcal{X}=\lbrace(x^i_t,t)\mid t\in[0,T],i=1,\dots,N\rbrace$, then calculate the loss in one pass by batching across $\mathcal{X}$. Additionally, we further optimize and calculate the loss on $\mathcal{B}\subset\mathcal{X}$, which importantly decouples the batch size from trajectory length. Both EM-BSDE and FS-PINNs methods can be (and were) optimized in the same manner, with the exception of the no-resetting method (EM-BSDE (NR)) (see Lines 229-232) which requires additional considerations due to its propagation of $Y_t$. In the following table, we report the average performance and overhead across three runs for each method with a batch size of 256:

These results align closely with the RL2 errors reported in Table 1, albeit with significantly lower runtimes (cf. Table 4). We also note some key details:

• The updated trajectory rollout schemes for FS-PINNs and EM-BSDE are more efficient with nearly zero overhead. This is reflected in the runtime of FS-PINNs vs. PINNs.
• The Heun-BSDE scheme requires approximately double the compute for Heun-based trajectory rollouts and loss calculation versus FS-PINNs, which is reflected in its overhead.
• For the specific elliptic/parabolic PDE we consider in (3.1), EM-BSDE does not require the computationally expensive Hessian computation of $u$ in the backwards SDE, while Heun-BSDE (and (FS)-PINNs) does, leading to significantly faster runtimes. This does not necessarily hold true for all PDEs however. In general, $h[u]$ can also depend on $\nabla^2u$ as noted in e.g., [Pham et al., 19] referenced by Reviewer DQqj, allowing for fully nonlinear second order PDEs. Alternatively, when solving first-order PDEs, the Stratonovich formulation $\mathrm{d}X_t=f(X_t,t)\mathrm{d}t+\sigma\circ\mathrm{d}B_t$ $\mathrm{d}Y_t=[\partial_tu(X_t,t)+\nabla_xu(X_t,t)\cdot f(X_t,t)]\mathrm{d}t+Z_t^T\circ\mathrm{d}B_t$ does not introduce a Hessian term, while the Ito formulation $\mathrm{d}X_t=f(X_t,t)\mathrm{d}t+\sigma\mathrm{d}B_t$ $\mathrm{d}Y_t=\left[\partial_tu(X_t,t)+\nabla_xu(X_t,t)\cdot f(X_t,t)+\frac{\sigma^2}{2}\mathrm{Tr}\left[\nabla^2_{xx}u(X_t,t)\right]\right]\mathrm{d}t+Z_t^T\mathrm{d}B_t$ requires the Hessian term for calculation. In these cases, we actually expect Heun-BSDE to either match or be faster than EM-BSDE.

In addition, we ran the original EM-BSDE solver to equal compute time as the original Heun-BSDE method on the 100D HJB case, and as predicted by our theoretical analysis, we found that EM-BSDE does not improve significantly with more iterations (see table).

While EM-BSDE converges faster than all other methods, the RL2 error quickly plateaus, leading to little improvement despite further optimization.

In the final revision, we will update all relevant results with the updated implementation discussed above and further include our insights on Hessian computations. Additionally, we agree that a study on dimensionality and run-time normalized results would provide valuable insights for the paper. Therefore, we plan to include results from both studies as plots for added clarity and completeness.

Response to W2/W3 (Empirical Analysis and Generality of PDEs considered):

First, we hypothesize that EM-BSDE vs Heun-BSDE performance on the 100D BZ case in Table 1 is due to the optimization challenges which arise from the high-dimensional, coupled nature of the problem, which may be confounding performance results (cf. Lines 252-256). This is evidenced by the lack of convergence for PINNs, FS-PINNs, and EM-BSDE (NR), in addition to recovering the expected ordering at lower dimensions (10D BZ).

(Note that the following response’s question was also brought up by Reviewer DQqj, so we have duplicated our response to make the rebuttals self-contained).

Additionally, the non-linear PDEs we consider (cf. Eq. (3.1)) are the standard kind of PDEs considered for BSDE methods (see e.g., references [4-6]). There are many types of applications, including in quantitative finance (e.g., multi-asset options pricing via Black-Scholes, local/stochastic volatility extensions from the work of Duprie and Heston), optimal control (e.g., Hamilton-Jacobi-Bellman equations for finding value functions, Hamilton-Jacobi-Issacs for reach-avoid problems), reaction-diffusion in chemistry (e.g., Allen-Cahn and Fisher-KPP equations), committer/exit-time probabilities for transition path analysis (e.g. [R1]), and so on. We will include a paragraph about these applications to our revised draft.

We also note that another reason we consider the specific elliptic/parabolic form (3.1) given in the paper (i.e., $h[u]$ does not depend on $\nabla^2u$) is to favor the standard EM-BSDE formulation, since the form (3.1) means that the EM-BSDE approach does not need to compute Hessians of $u$ in the backwards SDE, whereas the Heun-BSDE does.

The empirical examples tested in the paper were deliberately chosen for their simplicity and existence of analytical solutions in order to focus on isolating the integration bias introduced by EM-BSDE. For additional examples of PDEs, we are currently working to apply these methods to more complex non-linear optimal control problems (e.g., a high-dimensional task involving collision avoidance for many Dubin’s cars), and will report on our results in the final paper.

[R1]: E, Vanden‑Eijnden (2006), Towards a Theory of Transition Paths. J. Stat. Phys.

Response to Q1 (FS-PINNs Performance):

The FS-PINNs performance is in some sense not surprising, since the RL2 metric we evaluate performance over uses the forward SDE trajectories to compute the collocation points (cf. Lines 239-240, Eq. (6.2)). Hence, the FS-PINNs sampling distribution in training matches the performance metric better versus the PINNs sampling distribution in training, which uses a normal distribution fit to match the spatial distribution of the forward SDE trajectories. This explains the difference in performance.

We note that our choice of RL2 metric is deliberate—using the forward SDE as the sampling distribution is a much more meaningful measure than using e.g., uniform or isotropic Gaussian over a high-dimensional spatial domain. Furthermore, for problems such as solving Hamilton-Jacobi-Bellman (HJB) equations, the forward SDE trajectories actually have semantic meaning for the problem (e.g., either open-loop or closed-loop trajectory rollouts for HJB), which is substantially more useful for downstream tasks than the typical uniform distributions used in standard PINNs approaches.

Response to Q4 (PINNs Speed):

With the exception of EM-BSDE, all models require Hessian computations in the tested cases. The perceived speed advantage of PINNs stems from the other methods requiring trajectory rollouts for sampling, which impacted speed in our original implementation. We address this issue in our new implementation described in (Computational Costs).

Response to Q6 (f32 vs f64):

We observed this as well and believe that one plausible explanation is that float32 works as a form of implicit regularization due to the stochasticity introduced by lower precision (but importantly, the benefits are not enough to overcome the issues with EM-BSDE vs Heun-BSDE, cf. Table 3). Equally likely, there could also be subtle differences in how either jax, the Cuda layer, or the underlying GPU hardware handles f32 vs f64 which causes the discrepancy—these types of numerical issues are unfortunately not uncommon in modern NN training pipelines. To further investigate, we ran a similar test for both EM-BSDE variants on the 100D HJB benchmark which showed minimal differences between f32 and f64 (see table below, cf. Table 3). Interestingly however, there does still remain a persistent gap between f32 and f64, which we leave to future investigation.

We thank the reviewer for their thoughtful and constructive review. We greatly appreciate the recognition of both the originality of our work and its careful and thorough numerical analysis.

Response to W1/Q2 (Comparisons to [6]):

Thank you for raising this point. While we do not specifically use the term “diffusion loss” from [6] in the baselines of our experiments, the EM-BSDE implementation of the multi-step self consistency losses $L_{BSDE,H}(\theta)$ evaluated in Section 6.2 is precisely the method described in [6]; we discuss this in both Footnote 2 and Lines 205-207. In our revision, we will improve the clarity of the writing in the paper describing the connection between the two losses. Given this connection, Figure 4 demonstrates exactly the comparison that the reviewer is inquiring about: tuning the skip length does improve performance consistent with the empirical results in [6], but ultimately does not solve the bias issue in EM-BSDE (cf. Lines 303-305). Again, we will make this connection more explicitly clear in our revision.

Response to W2/Q2 (Computational Costs):

(Note that similar questions were also brought up by Reviewers LXC6 and sGAy, so we have duplicated our response to make the rebuttals self-contained).

We agree that the computational costs of the Heun integration presented in the original draft are key limitations for practical implementation. Hence, we spent a bit of time investigating how we could improve our implementation. We are happy to report that we have improved our implementation—for all methods—and the result is that the computational overhead of the Heun implementation is much more reasonable (e.g., <2x overhead over PINNs/FS-PINNs). Due to the given rebuttal time constraints, we were unfortunately only able to re-run our new implementation for all methods on the 100D HJB case, but we expect the computational benefits from the implementation changes we made (which we describe below) to carry over to all PDE cases. We will of course re-run all methods over all PDE cases for the revised draft.

We now describe the implementation changes we made. In the original Heun-BSDE implementation, the forward and backward SDEs rollouts were performed simultaneously, evaluating $L_{Heun_{k,\tau}}(\theta)$ at each step. As a result, the loss gradient calculation required back-propagation through the entire rollout; we had to resort to gradient checkpointing when combined with the Hessian calculation of $u$ in order to avoid running out of memory. In our new implementation, we first perform a forward SDE rollout to generate trajectories, $\mathcal{X}=\lbrace (x^i_t,t)\mid t \in [0,T],i=1,\dots,N \rbrace$, then calculate the loss in one pass by batching across $\mathcal{X}$. Additionally, we further optimize and calculate the loss on $\mathcal{B}\subset \mathcal{X}$, which importantly decouples the batch size from trajectory length. Both EM-BSDE and FS-PINNs methods can be (and were) optimized in the same manner, with the exception of the no-resetting method (EM-BSDE (NR)) (see Lines 229-232) which requires additional considerations due to its propagation of $Y_t$. In the following table, we report the average performance and overhead across three runs for each method with a batch size of 256:

These results align closely with the RL2 errors reported in Table 1, albeit with significantly lower runtimes (cf. Table 4). We also note some key details:

• The updated trajectory rollout schemes for FS-PINNs and EM-BSDE are more efficient with nearly zero overhead. This is reflected in the runtime of FS-PINNs vs. PINNs.
• The Heun-BSDE scheme requires approximately double the compute for Heun-based trajectory rollouts and loss calculation versus FS-PINNs, which is reflected in its overhead.
• For the specific elliptic/parabolic PDE we consider in (3.1), EM-BSDE does not require the computationally expensive Hessian computation of $u$ in the backwards SDE, while Heun-BSDE (and (FS)-PINNs) does, leading to significantly faster runtimes. This does not necessarily hold true for all PDEs however. In general, $h[u]$ can also depend on $\nabla^2 u$ as noted in e.g., [Pham et al., 19] referenced by Reviewer DQqj, allowing for fully nonlinear second order PDEs. Alternatively, when solving first-order PDEs (e.g. deterministic HJB problems), the Stratonovich formulation $\mathrm{d}X_t=f(X_t,t)\mathrm{d}t+\sigma \circ \mathrm{d} B_t$ $\mathrm{d}Y_t=[\partial_tu(X_t,t)+\nabla_xu(X_t,t)\cdot f(X_t,t)]\mathrm{d} t + Z_t^T\circ \mathrm{d}B_t$ does not introduce a Hessian term, while the Ito formulation $\mathrm{d}X_t=f(X_t,t)\mathrm{d}t+\sigma \mathrm{d} B_t$ $\mathrm{d}Y_t=\left[\partial_tu(X_t,t)+\nabla_xu(X_t,t)\cdot f(X_t,t)+\frac{\sigma^2}{2}\mathrm{Tr}\left[\nabla^2_{xx}u(X_t,t)\right]\right]\mathrm{d} t + Z_t^T \mathrm{d}B_t$ requires the Hessian term for calculation. In these cases, we actually expect Heun-BSDE to either match or be faster than EM-BSDE.

In addition, we ran the original EM-BSDE solver to equal compute time as the original Heun-BSDE method on the 100D HJB case, and as predicted by our theoretical analysis, we found that EM-BSDE does not improve significantly with more iterations (see table).

While EM-BSDE converges faster than all other methods, the RL2 error quickly plateaus, leading to little improvement despite further optimization.

Overall, the updated results align with the paper’s original claims while demonstrating the potential for practical application. In the final revision, we will update all relevant results with the updated implementation discussed above and further include our insights on Hessian computations. Additionally, we will include runtime-normalized performance plots (e.g. error vs wall-clock time) with each method for added clarity and completeness.

Response to Q1 (Literature review):

We thank the reviewer for the helpful clarifications regarding literature discussions. We agree that the abstract and introduction could be more precise in the positioning of the PINNs and BSDE literature. Specifically,

• We acknowledge that as noted by the reviewer, [6] demonstrates that BSDE solvers can exceed the performance of standard PINNs solvers in high dimensional settings, with the use of e.g., the proposed diffusion loss. However, as shown in [6], standard (EM-)BSDE methods do indeed struggle compared with PINNs; our statement about existing (EM-)BSDE solvers in the abstract was intended to refer to this phenomenon. It is this gap in performance, in addition to trying to answer why the diffusion loss empirically improves performance, that motivated the present analysis. We apologize for the confusion, and in the final version we will revise the abstract and introduction to clarify this point.
• We agree that the characterization of early PINNs as “scalable alternatives" for high-dimensional PDEs should be clarified. As pointed out by the reviewer, the original PINNs works [1,2] cited primarily focused on low-dimensional PDEs. However, we do believe that the motivation for using PINNs like methods for solving high-dimensional PDEs did arise quite early in the development. For example, the initial work on Deep Galerkin Methods (DGMs) from [Sirignano and Spiliopoulos, 2018] was motivated by looking for meshfree algorithms for high-dimensional PDEs. Furthermore, as another example, practitioners from the control theory community started to use PINNs methods to solve Hamilton-Jacobi reachability PDEs [Bansal and Tomlin, 2020] not long after the DGM work, with the explicit motivation to avoid the curse of dimensionality from classic grid-based PDE solvers. Nevertheless, we do agree that the work of [Hu et al., 2023] does seem to be the first work that explicitly studies the extent to which PINNs methods can be scaled to high dimension (e.g., 100k dimensions) through the use of algorithmic modifications to the training. We will certainly update our related work to reflect this chronology more accurately.
• We agree that [3] is a distinct method from PINNs and should not be grouped under the PINNs umbrella– thanks for catching this mistake. We will correct this in our revision.

Response to Q3 (Stratonovich-based BSDE theory for deep-learning based methods):

This is an interesting question, and we believe the answer to be yes. We answer the question in the context of the self-consistency losses we consider, but similar answers should also hold for the recursive formulations (where separate models are estimated at every timestep, see Line 116 for the distinction). Our result Theorem (4.2) combined with Eq. (4.8) is essentially the starting point for such a Stratonovich theory, by showing that the error of the $L_{Heun,\tau}(\theta)$ loss scales as the PDE residual term $R[u_\theta]$ (integrated along trajectories of $X_t$) plus an $O(\tau^{1/2})$ term that arises from the discretization. Hence, by choosing the discretization small enough, the difference between the Heun loss and the PDE residual integral can be made small enough. What is left to derive is that having small PDE residual error implies control of e.g., the true backwards SDE paths $(Y_t)$ and approximate path $(\hat{Y}^\theta_{nk})$ generated by the model $u_\theta$, or other measures of quantify solution quality that is useful for downstream applications. This part, while technical in nature, should at a high level follow similar arguments as the cited papers for the Ito setting.

We thank the reviewer for their thoughtful and constructive review. We greatly appreciate the recognition of the importance of studying alternatives to PINNs approaches for solving PDEs, in addition to the well-motivated approach that we take in this work.

Response to W1 (Computational Costs):

(Note that similar questions were also brought up by Reviewers Fvez and sGAy, so we have duplicated our response to make the rebuttals self-contained).

We agree that the computational costs of the Heun integration presented in the original draft are key limitations for practical implementation. Hence, we spent a bit of time investigating how we could improve our implementation. We are happy to report that we have improved our implementation—for all methods—and the result is that the computational overhead of the Heun implementation is much more reasonable (e.g., <2x overhead over PINNs/FS-PINNs). Due to the given rebuttal time constraints, we were unfortunately only able to re-run our new implementation for all methods on the 100D HJB case, but we expect the computational benefits from the implementation changes we made (which we describe below) to carry over to all PDE cases. We will of course re-run all methods over all PDE cases for the revised draft.

We now describe the implementation changes we made. In the original Heun-BSDE implementation, the forward and backward SDEs rollouts were performed simultaneously, evaluating $L_{Heun_{k,\tau}}(\theta)$ at each step. As a result, the loss gradient calculation required back-propagation through the entire rollout; we had to resort to gradient checkpointing when combined with the Hessian calculation of $u$ in order to avoid running out of memory. In our new implementation, we first perform a forward SDE rollout to generate trajectories, $\mathcal{X}=\lbrace (x^i_t,t)\mid t \in [0,T],i=1,\dots,N \rbrace$, then calculate the loss in one pass by batching across $\mathcal{X}$. Additionally, we further optimize and calculate the loss on $\mathcal{B}\subset \mathcal{X}$, which importantly decouples the batch size from trajectory length. Both EM-BSDE and FS-PINNs methods can be (and were) optimized in the same manner, with the exception of the no-resetting method (EM-BSDE (NR)) (see Lines 229-232) which requires additional considerations due to its propagation of $Y_t$. In the following table, we report the average performance and overhead across three runs for each method with a batch size of 256:

These results align closely with the RL2 errors reported in Table 1, albeit with significantly lower runtimes (cf. Table 4). We also note some key details:

• The updated trajectory rollout schemes for FS-PINNs and EM-BSDE are more efficient with nearly zero overhead. This is reflected in the runtime of FS-PINNs vs. PINNs.
• The Heun-BSDE scheme requires approximately double the compute for Heun-based trajectory rollouts and loss calculation versus FS-PINNs, which is reflected in its overhead.
• For the specific elliptic/parabolic PDE we consider in (3.1), EM-BSDE does not require the computationally expensive Hessian computation of $u$ in the backwards SDE, while Heun-BSDE (and (FS)-PINNs) does, leading to significantly faster runtimes. This does not necessarily hold true for all PDEs however. In general, $h[u]$ can also depend on $\nabla^2 u$ as noted in e.g., [Pham et al., 19] referenced by Reviewer DQqj, allowing for fully nonlinear second order PDEs. Alternatively, when solving first-order PDEs (e.g. deterministic HJB problems), the Stratonovich formulation $\mathrm{d}X_t=f(X_t,t)\mathrm{d}t+\sigma \circ \mathrm{d} B_t$ $\mathrm{d}Y_t=[\partial_tu(X_t,t)+\nabla_xu(X_t,t)\cdot f(X_t,t)]\mathrm{d} t + Z_t^T\circ \mathrm{d}B_t$ does not introduce a Hessian term, while the Ito formulation $\mathrm{d}X_t=f(X_t,t)\mathrm{d}t+\sigma \mathrm{d} B_t$ $\mathrm{d}Y_t=\left[\partial_tu(X_t,t)+\nabla_xu(X_t,t)\cdot f(X_t,t)+\frac{\sigma^2}{2}\mathrm{Tr}\left[\nabla^2_{xx}u(X_t,t)\right]\right]\mathrm{d} t + Z_t^T \mathrm{d}B_t$ requires the Hessian term for calculation. In these cases, we actually expect Heun-BSDE to either match or be faster than EM-BSDE.

In addition, we ran the original EM-BSDE solver to equal compute time as the original Heun-BSDE method on the 100D HJB case, and as predicted by our theoretical analysis, we found that EM-BSDE does not improve significantly with more iterations (see table).

While EM-BSDE converges faster than all other methods, the RL2 error quickly plateaus, leading to little improvement despite further optimization.

Overall, the updated results align with the paper’s original claims while demonstrating the potential for practical application. In the final revision, we will update all relevant results with the updated implementation discussed above and further include our insights on Hessian computations. Additionally, we will include runtime-normalized performance plots (e.g. error vs wall-clock time) with each method for added clarity and completeness.