A Simulation-Free Deep Learning Approach to Stochastic Optimal Control

We thank the reviewer for their comments. We do believe that the confusion is on the reviewer, however. Let us try to clarify our statements.

Weakness 1 (confusion about Girsanov theorem): Lemma 1 is a standard application of Girsanov theorem. It does indeed involves a change of measure, but the result is well-known: in a nutshell, Girsanov theorem allows one to express expectations over the solution of one SDE as expectations over another SDE with a modified drift, provided that the Girsanov factor is added to remove the bias exactly, precisely as stated in our Lemma 1. (The reviewer may also want to take a quick look at the wikipedia page: https://en.wikipedia.org/wiki/Girsanov_theorem )

To avoid confusion we will relabel this lemma as Girsanov theorem.

Weakness 2 (proof of Lemma 1): We don't think it makes sense for us to add a proof of Girsanov theorem since it can be found in any standard textbook, such as e.g. Øksendal book Stochastic Differential Equations: An Introduction with Applications (Springer, 2003). We will add a reference to this book.

Weakness 3 (Proposition 1): Here too, the reviewer seems confused for no reason. Proposition 1 is a direct consequence of Girsanov theorem, as our proof shows. We also note that, as pointed out by reviewer 2UrK and mentioned by us in text, this proposition is not really new: it is a variant of a result first derived in Yang & Kushner (1991). Our paper mainly generalizes this result to control parameterized by deep neural networks and show how to apply it in practice in this context.

Weakness 4 (Experimental results): Most of the examples we used have analytical solutions indeed: they are however standard test cases for SOC numerical methods, and are used precisely because they allow one to benchmark the numerical solution. We also stress that these examples are by no means trivial to solve numerically, as they are high-dimensional: in fact, standard solution methods either fail or scale badly as the dimension of the problem increases.

Regarding the baseline comparison with other methods, it is natural for us to consider mainly the vanilla method (i.e. the one requiring backpropagation trough the solution of the SDE), since it is the method of choice. Several other methods have appeared recently, in particular approaches that require solving an adjoint equation, but these schemes are much more complicated to implement than the one we propose, so we decided to leave a detailed comparison with these to future work.

Finally, let us stress that the approach of Han et al. mentioned by the reviewer is a scheme to solve a PDE by rewriting its solution as the solution to a backward stochastic differential equation (BSDE), rewriting this BSDE as a control problem, and then applying the vanilla method to this control problem (with backpropagation through the SDE solution required). So benchmarking our approach against the vanilla method is in effect a comparison with the scheme proposed by Han et al.

We thank the reviewer for their constructive comments. Here is our reply and we hope this helps address the concerns of the reviewer.

Weakness 1 (base distribution): We do not assume that the base distribution is a point mass, i.e. Propositions 1 and 2 hold also when the initial condition for the SDE is a nonsingular measure $\mu_0$. The point mass assumption is only used in Proposition 3, where we consider the application of our method to the construction of Föllmer processes (i.e. Schrödinger briges between a point mass an a target), and in Proposition 5, we we consider the problem of fine-tuning a difusion. In particular, the numerical examples discussed in Sec. 3.1 and 3.2 involve Gaussian initial data $X_0 \sim N(0,\frac12 d)$.

Weakness 2 (application scenarios): In terms of applications, our method can be applied to any scenario that fits in the formalism described in Section 2.1. Concretely, in the original submission, we illustrated the method on two types of examples: the control of Ornstein-Uhlenbeck processes, and the sampling from a target distribution with unknown normalization factor (see Sec. 3.3). In the revision, we will include another, more challenging, example (see Section 3.4), coming from statistical lattice field theory. More precisely, we show how our method can help fine-tuning a generative model to sample from the so-called $\varphi^4$ model.

Weakness 3 (simulation-free): There are indeed several ways to understand this terminology. Here, we mean that we need to simulate the data but we do not need to backpropagate through trajectories when computing the gradient. In the original submission, we had written that our simulation-free approach "avoids the extra cost of the differentiation through the solution of the SDE". In the revision, we have added the following sentence in the introduction: "We refer to this approach as simulation-free." We hope this will avoid any confusion. We would also consider changing the title if you think this could help.

Question 1 (simulation-free): See our answer to Weakness 3 above.

Question 2 (relationship between propositions): Thank you for giving us the opportunity to clarify the connections between our theoretical results. There are 3 groups of 2 results.

• The first two results give the foundation for our algorithm. Proposition 1 gives an expression for the gradient of the loss function $L(\theta)$ which involves only $\partial_\theta u^\theta$ and does not require differentiating through the trajectory $X^\theta$. Proposition 2 reformulates this gradient in terms of an alternative loss, $L(\theta,\bar\theta)$ which makes the computation readily implementable by automatic differentiation. This result is the foundation for our algorithm; see Eqs. (11) & (12).
• The next two propositions are about sampling from a distribution using a Föllmer process: Proposition 3 rephrases the problem as an SOC problem, and Proposition 4 explains how to correct the generated samples by reweighting in order to obtain unbiased estimators for expectations.
• The last two results, Propositions 5 and 6, are generalizations of Propositions 3 and 4, respectively, to finte-tune a generative model.

Question 3 (Algorithm 1): In the revision, we have clarified how we carry out all the computations in Algorithm 1. Here's a table of our method as compared to the existing methods.

[1] Domingo-Enrich, C., Drozdzal, M., Karrer, B., & Chen, R. T. (2024). Adjoint matching: Fine-tuning flow and diffusion generative models with memoryless stochastic optimal control. arXiv preprint arXiv:2409.08861.

[2] Domingo-Enrich, C., Han, J., Amos, B., Bruna, J., & Chen, R. T. (2023). Stochastic optimal control matching. arXiv preprint arXiv:2312.02027.

[3] Berner, J., Richter, L. & Ullrich, K. (2022). An optimal control perspective on diffusion-based generative modeling. arXiv preprint arXiv:2211.01364

[4] Vargas, F., Padhy, S., Blessing, D. Nüsken, N. (2023). Transport meets Variational Inference: Controlled Monte Carlo Diffusions. arXiv preprint arXiv:2307.01050

Question 4 (Fig. 2 right panel): To confirm the convergence of the result shown in the right plot of Figure 2, we ran both algorithms for a much longer computation time and we plotted the results with confidence intervals. The new plot can be found in the anonymous GoogleDrive and wil also be included in the revised version of our paper. We found out, from the left panel of Figure 2, that both algorithms converged because the $L^2$ error on the control is very small. In this specific problem, although we initialize the neural network's parameters randomly, the initial control (before the first episode) is already quite good, with a control error of order $10^{-1}$ and this is why there is not room for much improvement in the empirical loss function plotted in the right panel of Figure 2.

Question 5 (Fig. 4): The grid only appears in the plot, not in the computations. This is a ten-dimensional numerical example where the conditional distributions of $x_1,\ldots,x_9$ all depend on $x_0$. The plots in Figure 4 are made by squeezing these 9 dimensions (i.e., $x_1,\ldots,x_9$) into one single dimension and plotting the histogram of the new dimension with the dimension of $x_0$. The simulations and the neural network's inputs are in continuous space, which allows to deal efficiently with high dimensions.

We thank again for the careful review and the invaulable suggestions made by the reviewer. Given these additional details we have provided about our work, we respectfully ask the reviewer to consider raising the score to better reflect our work's contributions.

We thank the reviewer for their constructive comments.

Weakness 1 (novelty): Just before Proposition 1, we clearly stated that "A similar formula first appeared in Yang & Kushner (1991) and it is also given in Ribera Borrell et al. (2024)" so we do not claim this is our contribution. We will rephrase how we introduce Propostion 1 to clarify this.

The main novelty of our work is the implementation of an eficient algorithm using this formula for the solution of SOC problems via deep learning, with application to recent problems such as the construction of Föllmer processes and the fine-tuning of generative models. These problems have recently generated a lot of interest in the machine learning community, and for this reason we believe that our results will be of interest to, and generate more work in, this community.

Weakness 2 (applicability): We agree that our method is not directly applicable to settings where the volatility is also controlled. However many SOC problems, including those arising in machine learning, have only a controlled drift (e.g. fine-tuning diffusion models [1], Schrodinger Bridge problems [2]). Furthermore, we believe there are ways to extend our method to controlled volatility, for example by using Malliavin calculus, following the lines of (Gobet & Munos, 2005). We mention this point in the revision after Proposition 1.

Question 1 (typo): There was indeed a typo in the first term, which should be multiplied by $1/2$. We have corrected it in the revision. Our code implements the correct formula, which is updated in the revised draft.

Question 2 (Figure 2): Actually we ran both algorithms with the same mini-batch size. To confirm that there is no mistake in the result shown in the right plot of Figure 2, we ran both algorithms for a much longer computation time and we plotted the results with confidence intervals. The new plot can be found in the anonymous GoogleDrive.

Question 3 (variance of estimator): The norm of our gradient estimator might have a larger variance than the one of the vanilla method, but it is actually not so easy to tell a priori wherher that is the case or not. In addition, and more to the point, what controls the actual performance of the algorithm is the variance in the loss function or the $L^2$ error. From the point of view of these metrics, our experiments (including the new ones available in the anonymous GoogleDrive) indicate that our method performs very well, and better than the vanilla method. If you have a specific concern in mind, please feel free to give us more concrete way to convince you.

[1] Domingo-Enrich, C., Drozdzal, M., Karrer, B., & Chen, R. T. (2024). Adjoint matching: Fine-tuning flow and diffusion generative models with memoryless stochastic optimal control. arXiv preprint arXiv:2409.08861.

[2] Liu, G. H., Lipman, Y., Nickel, M., Karrer, B., Theodorou, E. A., & Chen, R. T. (2023). Generalized Schrödinger Bridge Matching. arXiv preprint arXiv:2310.02233.

We thank the reviewer for their constructive comments. Here is our response to the concerns of the reviewer.

Weakness 1 (Girsanov theorem): While we agree that our approach relies on some assumptions about the form of the control problem, we would like to stress that there are many interesting problems that fall in the framework covered by our assumptions. In our paper we focus on applications involving the construction of Schrödinger bridges and the fine-tuning of diffusion models since they are extremely important in recent applications of generative AI. We would also like to stress that Gaussian noise is ubiquitous in a wide range of applications, including again diffusion models for instance. It might possible to extend our method to more general forms of noises and controlled dynamics but we leave this for future work.

Weakness 2 (high dimension and time horizon): While the computational time increases with dimensionality and time horizon, as with any numerical methods, we have not observed particular challenges with our method due to these aspects. To demonstrate this, we have added one more experiment of the Quadratic OU type, with $d=50$ and $T=4$. Please check the results in this anonymous GoogleDrive.

Weakness 3 (complexity in high dimension): We are not sure what type of "complexity" you are referring to. Computing the gradient of the SOC objective via our Proposition 2 is much less complex than doing so via backpropagating through the whole trajectory (even in high dimension) because it can be done timestep-per-timestep. We have clarified this in the new version of Algorithm 1, which explains in detail how the computations are done. Please us know if you would like further clarification.

Question 1 (high variance): We have not observed such issues in our experiments. The norm of the gradient might have a large variance but this does not necessarily impact the performance of the algorithm. What really matters is the variance in the loss function or the $L^2$ error. From the point of view of these metrics, our experiments show that our method performs very well, with lower variance than the vanilla method. If you have a specific concern in mind, please feel free to give us more concrete way to convince you.

Question 2 (target control): We are not sure what you mean by "reference process". If you refer to the "reference control" denoted by $v$ in Lemma 1, note that this control is set to the current control, $v = u^{\bar\theta}$, in Propositions 1 and 2: that is, in our algorithm, $v = u^{\bar\theta}$ can be viewed simply a version of the neural network $u^\theta$ with respect to which we do not compute the gradients. In practice, we have a single neural network but we use stopgrad on the terms involving $v=u^{\bar\theta}$ to detach them from the compuational graph and prevent their gradients from being computed. This $v=u^{\bar\theta}$ is also used to generate trajectories. There was a typo in the code and we corrected it in the revised version.

We also stress that $v = u^{\bar\theta}$ does not need to be close to the target (optimal) control, at least not initially: it simply converges towards this target during training.

Given these additional explanations about our work and the experimental evaluations we have added, we respectfully ask the reviewer to consider raising the score to better reflect our work's contributions.