An Efficient High-dimensional Gradient Estimator for Stochastic Differential Equations

Overparameterized stochastic differential equation (SDE) models have achieved remarkable success in various complex environments, such as PDE-constrained optimization, stochastic control and reinforcement learning, financial engineering, and neural SDEs. These models often feature system evolution coefficients that are parameterized by a high-dimensional vector $\theta \in \mathbb{R}^n$, aiming to optimize expectations of the SDE, such as a value function, through stochastic gradient ascent. Consequently, designing efficient gradient estimators for which the computational complexity scales well with $n$ is of significant interest. This paper introduces a novel unbiased stochastic gradient estimator--the generator gradient estimator--for which the computation time remains stable in $n$. In addition to establishing the validity of our methodology for general SDEs with jumps, we also perform numerical experiments that test our estimator in linear-quadratic control problems parameterized by high-dimensional neural networks. The results show a significant improvement in efficiency compared to the widely used pathwise differentiation method: Our estimator achieves near-constant computation times, increasingly outperforms its counterpart as $n$ increases, and does so without compromising estimation variance. These empirical findings highlight the potential of our proposed methodology for optimizing SDEs in contemporary applications.

Paper

Similar papers

Peer review

Reviewer 5TVX4/10 · confidence 4/52024-06-14

Summary

The article studies classical stochastic control in the continuous d-dimensional SDE setting (with and without jumps). The authors suggest to compute the gradients of the value function (expected total reward in the non-discounted finite time regime), which can be applied (in a model-based setting) to policy gradient type algorithms popular in RL. While in a straight-forward fashion computing the derivatives of the expectations would lead to the simulation of n (dimension of approximation space for the control) expectations that require the simulation of one SDE each, the authors propose a clever interchange of derivatives that reduces the number of SDEs to a polynomial in d. Very careful proofs are given, the findings are applied to simple linear SDE example with quadratic cost function.

Strengths

The article introduces a creative switching trick (switching order of derivatives and derivatives with expectations) that allows to reduce the simulation dimension compared to a straight-forward estimator. The rest of the paper consists of a rigorous verification of interchanging limits (derivatives) and expectations. The proofs are careful (mainly based on differentiability properties of SDE flows), I spotted only few typos.

Weaknesses

I do not agree with the comparison to RL (=sample based stochastic control) in Section 1.1. The entire point of RL is to provide model-free estimators, in particular, the policy gradient theorem cancels out the model transitions. Here, the situation is different. The model ($\mu$ and $\sigma$) is used to write down the SDEs (2.7) required to simulate the estimator. This is not a problem but shifts the article more towards the stochastic control community. I cannot see a clear interest in the ML/AI community in the questions raised, the control setting of SDEs (jump SDEs) is natural in the probability/math finance community but not at all in ML/AI. Even if I try I struggle to come up with examples relevant for the ML/AI community (the authors neither). The article is written for a purely mathematical audience with research level skills in stochastic process theory, the article could appear almost unchanged in journals such as SPA or SIAM Journal on Control and Optimzation. I think NeurIPS is not the right place for publication. I am willing to increase my scores if a more relevant example than the quadratic problem could be provided that shows value of the switching trick to the ML/AI community. The appendix is chaotic. It is tricky to follow the proofs because parts of the proofs are deferred to other sections of the appendix. Given that the proof only justifies interchange of limits and expectations it looks much harder than it is.

Questions

1. You will typically have to solve the three SDES using numerical schemes. Can you provide additional assumptions on the model (the coefficients) that, combined with SDE numeric results (let's say in the Brownian case only), provides a result for the estimation error of the gradient? 2. The additional derivatives on the model coefficients make me doubt the SDE scheme (2.7) works well if the coefficients are not so nice. Many control problems in applications are Heston type models, which are already problematic for numerics with the square-root variance coefficient. I know they do not satisfy the assumptions imposed, did you still try your simulation scheme? 3. I do not understand why you included the jumps without an application in mind. This is normal for a maths journal, but a bit strange in ML/AI. Could you please provide an example? 4. A few minor points: (i) The generator of a Markov process is the action plus the domain. Could you please be a bit more careful when you rework the article? (ii) You might want to mention that uniformly bounded derivatives imply linear growth which is the standard property required for global existence. (iii) I am probably missing where you discuss properties on the model that imply Assumption 2. If not I find assumptions on the derivatives of the mean rewards disturbing that I do not know how to check. (iv) Typo in Assumption 1, 2. (v) At least for the larger expectations it would help to be more consistent with the use of brackets. (vi) Could you please put references to the quadratic loss examples? Since this is the only application it would be good to see where and how it is used. (vii) the use of colons is not very helpful, for instance in (2.6) it is a bit complicated to understand what object is defined (I guess you want to define Z). (viii) The chain rule is missing in (D.3)

Rating

4

Confidence

4

Soundness

4

Presentation

2

Contribution

3

Limitations

The authors did not discuss the limitations or the practical problems of the results.

Reviewer Eiq97/10 · confidence 4/52024-07-05

Summary

The paper proposes a novel algorithm for the gradient of machine learning objectives that are based on SDE paths, w.r.t. parameters. The method is scalable for big parameters, as the computational complexity of the estimation is not related to the number of parameters. Empirical results show the scalability advantage over classical pathwise differentiation.

Strengths

PDE constrained optimization has remained a very difficult question due to the high cost of simulating PDEs, let alone taking derivatives. Therefore, despite its many good theoretical promises, it has not been deployed. This paper attacks the question by reducing the computational cost with a monte-carlo estimation of the derivative, largely advances the promise of PDE constrained optimization. The writing is also very good, and the intuitions behind the methods are well illustrated.

Weaknesses

1. Despite the fact that the computational complexity is irrelevant to the number of parameters, I suspect that as the number grows, the method would require more samples for an accurate monte-carlo estimate of the gradients. Is it possible to either: 1. Provide a theoretical analysis on the convergence rate of the estimation w.r.t. the number of parameters. 2. Provide an empirical analysis on the convergence rate of the estimation w.r.t. the number of parameters.

Questions

See weakness. Also, applying the Feynman-Kac twice requires the objective to be very smooth, can you provide some insights on the applications where this holds true?

Rating

7

Confidence

4

Soundness

4

Presentation

4

Contribution

4

Limitations

As written in questions, the smoothness requirements induced by Feynman Kac etc. will limit the use cases. As written in weaknesses, it is still unclear the convergence w.r.t. high number of parameters.

Reviewer rGnk7/10 · confidence 4/52024-07-13

Summary

This paper formulates an efficient, unbiarsed, and finite variance gradient estimator for an objective function that looks like the stochastic optimal control cost function (it is ubiquitous across various applications). The problem concerns overparametrized SDEs (the paramete dimension n r is of a much higher dimension than the state space d). The algorithm runs at a complexity invariant to the large n. The authors use the Feynman-Kac PDE for v_{\theta} and transform it to an equivalent PDE to solve. They then employ a pathwise differentiation estimator to estimate the 2nd order and 1st order derivatives in the PDE while exploiting certain properties of the Hessian.

Strengths

1. The applications of gradient estimators and the motivation behind this problem is well stated and supported. 2. The methodology is clear, especially when finding probablistic representations of the gradient estimator.

Weaknesses

1. The authors discuss the limitations in the checklist but it would have been more clear to have a limitations section clearly outlined as claimed in the checklist. 2. The assumptions may be restrictive as mentioned. 3. The work could be improved if there was a numerical simulation. The example of LQR is an illustrative example and it may be of didactic use and potentially relatively easy to implement a scenario when the parametrization is of a much higher dimension. 4. The spatial derivatives \partial v_0 are still difficult to compute, but still better than \partial v_{\theta}. What are some error bounds you can get by using the later derived probabilistic representations that can use Monte Carlo

Questions

1. Are there any error bounds from using the probabilistic representations for the spatial derivatives (see #4 in Weaknesses) Happy to revise score if necessary.

Rating

7

Confidence

4

Soundness

4

Presentation

4

Contribution

3

Limitations

Yes they have mentioned some, rest may be in the above Weakness section unless there are some misunderstandings.

Reviewer 5TVX2024-08-08

Thanks for answering my questions!

Thanks for answering my questions! I see in the other reviews good interest in PDE constraint optimization. I cannot judge this direction and it seems this is already largely reflected in the scores of other reviewers. From the point of view of RL and also Math Finance I am still not convinced. ML has not proved much relevance in finance and I do not believe it will in the future. There are plenty of numerical approaches towards affine processes (CIR was really tortured with all possible techniques). Without seeing a proper comparison I cannot believe that the (super inefficient!) policy gradient approach can add any value in this model-based setting. It could in the model-free setting which is not possible with the present approach. I still do not understand why there is no proof for Assumption 2 (in the appendix). If you can prove it, prove it. And identify clearly a set of necessary conditions. Also I do not see a reason not to include a numerical analysis, this would make the paper much stronger. The contribution of only the nice (!) main idea and justifying some changes of limits that everyone in the ML community would believe anyways is a bit too small for NeurIPS, given a practical need is not completely clear. I will keep my score.

Authorsrebuttal2024-08-10

Comments to the reviewer's feedback

We thank the reviewer for your quick feedback. We have the following response. **Application of Machine learning in finance:** We respectfully disagree with the comment that machine learning methods have "not proved much relevance in finance." In fact, the use of artificial neural networks for financial SDE modeling and optimization dates back to the early 1990s. Survey paper [9] summarizes **more than 150 papers** that use neural network-parameterized SDEs to model prices in option pricing and hedging. Additionally, recent papers surveyed in [9] actively explore the use of deep network architectures. This underscores the relevance of our proposed methodology in cutting-edge finance applications. Furthermore, based on the private consulting activities of some of the authors, neural networks and other machine learning techniques are very much alive and of significant interest in financial applications. Given these academic interests and industry experiences, we believe it is hard to argue otherwise. *We respectfully ask why you think that ML has not proven to be relevant in finance.* **Affine processes:** We do not intend to use our method to optimize simple affine processes. Our methodology is designed for estimating sensitivity or optimizing values driven by high-dimensionally parameterized jump-diffusions. Our method excels in modern ML settings where deep neural networks are used in the SDE parameters. **Application in sciences and engineering:** Other than financial applications, our methodology is well-motivated and an important contribution considering its application in science and engineering. For example, neural SDEs aim to learn the behavior of natural (physical) random processes with parameterized SDEs. Highly cited papers in this field include [7] and [5] (with and without jump). Neural SDEs typically optimize $\min_\theta E[g(X_\theta^x(T))]$ subject to an SDE $$ dX_\theta^x(t) = \mu_\theta(t, X_\theta^x(t)) dt + \sigma_\theta(t, X_\theta^x(t)) dB(t) + \chi_\theta(t,X_\theta^x(t))dN(t)$$ where $\mu_\theta,\sigma_\theta,\chi_\theta$ are neural network parameterized and $N(t)$ is a Poisson process. Gradient methods for approximating stochastic optimal control with a parameterized policy class [3] are also important applications. We remark that many neural network-parameterized SDEs are not widely used primarily because of the lack of scalable gradient estimators. Our contribution directly addresses this issue. Given these numerous applications across multiple disciplines, we find it **hard to believe** that NeurIPS, where practitioners and theorists, scientists and engineers exchange ideas and advance machine learning theory and applications, is not a suitable venue to publish this work. **Rigorous validation of limit interchange:** We respectfully disagree with the claim that the interchange of limits is what "everyone in the ML community would believe anyways." Wrong intuition regarding the interchange of limits can lead to serious errors. A researcher might explore the following gradient estimation idea: From semigroup theory, write $E_x[g(X_{\theta}(t))] = (e^{tL_{\theta}}g)(x)$. Formal differentiation (taking the limit quotient) yields: $$\partial_\theta (e^{tL_{\theta}}g)(x) \stackrel{?}{=} t e^{tL_{\theta}} \partial_\theta L_{\theta} g(x) = t E_x[\partial_{\theta}L_{\theta} g(X(t))].$$ Although this can work for Markov jump processes, it is incorrect, unfortunately, for diffusions. The correct expression is precisely the representation in our Theorem 1. The reason readers find the ideas in Section 2 natural is that we have invested research effort into identifying a simple and intuitive justification, all backed by rigorous proof. Nevertheless, given this helpful discussion with the reviewer, we think including this incorrect reasoning in the paper could help to illustrate the need for caution when working with generators of SDEs. **Assumption 2:** We have considered including sufficient conditions on the model primitives to imply Assumption 2 but decided against it for the following reason: The primary focus of this paper is to rigorously justify the **validity** of the generator gradient estimator. Proving Assumption 2 using sufficient conditions could detract from this goal. The differentiability and moment bounds for the value are well-established in the literature. Also, justifying Assumption 2 from model primitives would introduce another set of assumptions typical in jump-diffusion analysis but potentially confusing for users, as we have explained in the previous response. If you still have concerns about establishing Assumption 2, we can provide a proof in the Appendix with smooth and bounded rewards (so that there will be no interference with the variance's growth power). [9] Ruf, Johannes, and Weiguan Wang. Neural networks for option pricing and hedging: a literature review. Journal of Computational Finance.

Reviewer 5TVX2024-08-12

CIR at zero

Thank you for providing this additional numerical study! I am wondering how to read your results for small $\theta$. If $\theta> 1/2$ the model is essentially smooth as the paths avoid the irregular point zero and the coefficients are smooth away from zero. To check empirically if the estimator also works in non-smooth scenarios you are interested in small $\theta$ to challenge the non-smoothness. The careful interpretation would be the estimator potentially does not work well in non-smooth settings. Azt least in small dimension it seems to be much worse than FD. It could happen (I am only speculating) that for non-smooth coefficients in high dimension the entire advantage of the estimator gets lost in gigantic variances. The rebuttal statements the estimator has wide applicability also for non-smooth coefficients might be a bit far fetched. Might be true or not, we do not have good evidence in both directions.

Authorsrebuttal2024-08-12

Clarification

Even though the CIR process will not hit 0 when $\theta> 1/2$, the process can be arbitrarily close to 0 with positive probability. So, this is still a non-trivial case of non-smooth volatility. On the other hand, the ReLU SDE is also one with a non-differentiable drift. In these two cases (CIR with $\theta>1/2$ and ReLU drift), the derivative processes are well defined and integrable. As explained before, our theory can be extended to establish the validity of the generator gradient estimator in these context. We also empirically validate this consistency. However, when $\theta<1/2$, as we have explained, the derivative processes are not globally defined. So, as we have explained, we cannot guarantee the consistency of our estimator. This is reflected by a surge of variance as $\theta\downarrow 0$. Nevertheless, the use case of our proposed methodology usually involves SDEs parameterized by deep neural networks, which are typically smooth (or of ReLU-type) in $X$. Therefore, in these applications, the smoothness should not be a concern to users of our methodology. Moreover, we would like to further assure the reviewer of the validity by mentioning that we have experimented with the neural network parameterized LQ control example using ReLU activation functions. We observe similar computation and variance performance.

Reviewer rGnk2024-08-13

response

Thank you for the clarifications here and in your above rebuttal. With regards to point 3., I think it would be beneficial to have this demonstration. I will keep the score as it is

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC