Paper: Deep Stochastic Mechanics (ICLR-2024)

Reviewer comment (Experiments):
The approach is limited to very simple/small systems (harmonic oscillators and interacting distinguishable particles of only a handful of particles).

Author response:
Even in the “simple” settings considered in our paper, existing algorithms (e.g. Crank-Nicolson) encounter significant limitations and scaling issues both in terms of computation time and memory requirements. The core contribution of our submission is demonstrating an approach that sidesteps these challenges. We detail in our future work section our interest in taking this work and expanding it to more complex settings, but this initial work was an essential first step and an important contribution in its own right.

Reviewer comment (Novelty):
The discussion is quite limited, and bypasses a vast amount of literature in the field, including time-dependent variational wave functions based on neural network parameterizations.

Author response:
Thank you for your suggestion! We are happy to cite and include more discussion of the two mentioned papers (Carleo et al., Schmitt et al.). We agree that the variational approach and its time-dependent variant are old and well-studied methods. Both Carleo et al. and Schmitt et al. are neural network quantum state methods that are applicable to discrete quantum systems, e.g. spin systems, whereas our method is for continuous quantum systems. We are not aware of a time-dependent variant of a continuous neural network quantum state. In addition, unlike our approach, the neural network quantum states discussed in these two papers require expensive MCMC sampling in order to generate samples, which must be done for each training step and each time step in the time evolution.

Reviewer comment (Novelty):
The example statement in the paper about approaches like FermiNet (Pfau et al., 2020) or PauliNet (Hermann et al., 2020) reformulating the problem as maximization of an energy functional and claiming that such an approach cannot be adapted to the time-dependent wave function setting is misleading, as the static variational reformulation does not date to 2020 but is as old as quantum mechanics itself.

Author response:
> Again, this is problematic, because you seem to believe that VMC approaches based on neural networks are highly specialized. They are not (e.g. FermiNet, that you cite as a prominent example of VMC approach for ground states in table 1. Not sure why, by the way, only FermiNet is mentioned in there and not, say, PauliNet or the many other neural-network based VMC methods... but ok).

When comparing VMC with us, we need the following adjustments to make it fair:
1. From FermiNet, remove all pairwise interactions from model features or add them to our approach.
2. Either remove Slater-determinant-based architecture in FermiNet and replace it with our simple architecture or use it within our approach.
3. For all VMC approaches, replace task-dependent efficient MCMC samplers with some generic sampler, such as Langevin SDE. Their efficiency, in general, is only possible as a result of task adaptation, as all interaction potentials are non-convex. Therefore, it implies exponential dependence in the spectral gap (e.g., see Eq. 4.3 in https://arxiv.org/pdf/1702.03849.pdf).
4. …

We could remove all external knowledge specific to the potential, or incorporate it problem-wise into our method. From a No-free-lunch perspective, we are always bound to fail compared to methods that rely on some external additional perspective. To make the comparison fair, we either need to level the field by removing all externalities from the corresponding methods or do the amount of research like in all referenced by your papers to integrate it into our approach for each problem individually. 

And by no means we want to diminish the importance of VMC! We agree that VMC is a powerful framework, and we must discuss it with citations of all the works you mentioned. We are very grateful for such a fruitful discussion and agree that the current work's limitations need to be discussed in Relevant Work and Limitations Section, which we will add in the revision.

Reviewer comment (Novelty):
Neural variational parameterizations of the wave function are routinely used to solve the time-dependent Schrödinger's equation, based on Dirac and Frenkel's time-dependent variational principle, and have been implemented in several works, starting from examples such as https://www.science.org/doi/10.1126/science.aag2302 and https://journals.aps.org/prl/abstract/10.1103/PhysRevLett.125.100503, among many others.

Author response:
Thank you for your suggestion! We are happy to cite and include more discussion of the two mentioned papers (Carleo et al., Schmitt et al.). We agree that the variational approach and its time-dependent variant are old and well-studied methods. Both Carleo et al. and Schmitt et al. are neural network quantum state methods that are applicable to discrete quantum systems, e.g. spin systems, whereas our method is for continuous quantum systems. We are not aware of a time-dependent variant of a continuous neural network quantum state. In addition, unlike our approach, the neural network quantum states discussed in these two papers require expensive MCMC sampling in order to generate samples, which must be done for each training step and each time step in the time evolution.

Reviewer comment (Experiments):
The approach is shown for very small systems, and in cases that are solved easily with many other approaches (including using a discrete basis and the time-dependent variational principle with a neural network wave function). A question is whether the authors have a sense of the scaling in terms of number of particles, and if they can show a case that goes beyond what can be simulated exactly.

Author response:
This question is related to **W1**. In our paper, there is an experiment in Section 5.4 with many non-interacting particles. In particular, our approach shows linear convergence time for a specific case of the harmonic oscillator (Figure 4). In contrast, simulation of this setting (without prior knowledge) using traditional solvers would require a grid of size $N^d$. Our approach bypasses this issue, and our empirical findings support it. 

There are experiments in the original submission that do correspond to interacting particles. **For this rebuttal, we ran new experiments for larger d to show scaling properties more clearly.** The following Tables illustrate time and memory usage for the Crank–Nicolson (C-N) solver for d=2, 3, 4 and training time/time per epoch and memory usage of our DSM method for d=2, 3, 4, and 5. Also, Figure 11 in **the revised supplementary material D.5.1** shows predicted densities for different dimensions. Note that while we don’t have the baseline for d=5, we believe DSM predictions are still reasonable.

Reviewer comment (Experiments):
Seems the neural network employed in this study is a single layer neural network. I am wondering how a single layer neural network could learn dynamics of a complicated wave function. Also the illustrations in the experiment look pretty simple. The authors are encouraged to tackle more complicated cases using the proposed method.

Author response:
Even in the “simple” settings considered in our paper, existing algorithms (e.g., finite-difference based or PINNs) encounter significant limitations and scaling issues both in terms of computation time and memory requirements. The core contribution of our submission is demonstrating an approach that sidesteps these challenges. The most interesting system to study, of course, is many-particle systems with interactions. We show that our method works for a simple interacting system (Sec. 5.4), and we use a more sophisticated network (for example, we add skip-connections) in this case. In Appendix D5, we provide more implementation details for this setting. Also, in our reply (2/2) to the reviewer JVMX, we show how our DSM approach scales with the problem dimension for the interacting system, given that potential can be computed linearly.

Regarding the use of a single-layer neural network, we intentionally choose a simplified architecture to emphasize the foundational aspects of our framework. We aim to prove that our framework works, not due to some factors like a smart choice of architecture. As our research progresses, we are open to exploring more sophisticated task-specialized neural architectures and algorithmic enhancements to tackle more complex simulations. Finally, works like [1] show that even single-layered neural networks with non-linear activations, when used in Neural SDEs, are powerful universal approximators.

[1] Veeravalli T. et al.,  A Constructive Approach to Function Realization by Neural Stochastic Differential Equations

Reviewer comment (Evaluation):
The O(Nd) computational complexity need more elaboration. Is it the computational complexity of training the neural networks on a single trajectory? For learning a stochastic process, we may need to sample many trajectories in order to learn hidden low dimensional representations of process. I am not sure whether it is fair comparisons with other methods as listed in the table 1.

Author response:
Thank you for noticing that. This comparison is unfair for a numerical solver indeed, as it does not have a well-defined notion of one iteration. By an iteration, we mean one training epoch, so we report computational complexity per epoch. We are sorry about this confusion. We will clarify this by modifying the table as detailed below (in fact, it’s modified in the revised version). That said, our reported complexities are important for understanding scaling properties of different methods in higher-dimensional settings. Note that all complexities are given under the assumption that the computational complexity of the potential is $O(d)$. If it’s more than that, then complexities of all methods needs to be corrected by this higher complexity additively.

To modify the table, we will add a column, “Overall complexity”, and we will place “N/A” in the “Iteration complexity” column for a numerical solver. In the “Overall complexity” of the numerical solver, assuming that we aim for a precision $\varepsilon = \mathcal{O}(\frac{1}{\sqrt{N}})$, we will place $\mathcal{O}(\frac{1}{\epsilon^{d+2}})$. In the “Iteration complexity” column, we will keep other methods as they are. Finally, in “Overall complexity”, we will specify that for all other methods existing theoretical results allow only to place lower bounds on complexity. In particular, we will multiply iterations complexity by a factor $\mathcal{O}(\mathrm{poly}(\frac{1}{\epsilon}))$, coming from a general non-convex stochastic gradient descent bound. Such scaling $\mathrm{poly}(\cdot)$ (but not exponential) is broadly generic for non-convex stochastic gradient convergence to the minima; and it is true for a broad class of non-convex stochastic gradient methods [1].

Reviewer comment (Reproducibility):
For training losses defined in eq. 11 to eq. 15, because they need integration operation, I am wondering in practice how the integration is done and what is the window length of the integration during the training process at each iteration/epoch.

Author response:
We consider the simplest integrator of SDE, the Euler-Maruyama scheme. We split time $(0, 1]$ uniformly in $N+1$ parts, where $N=1000$. So, $dt = 0.001$. (We mention these details in Appendix C). A higher-order integrator or an adaptive integrator could achieve the desired accuracy with a much lower $N$. We leave it for future work.

Reviewer comment (Reproducibility):
In page 5, the authors mentioned for each iteration, they will sample new trajectories using eq. 7. How do we handle the cold start problem at the very beginning of the training process. I mean at the beginning of the training process, neural networks have not learned dynamics of the stochastic process. So the trajectories we sampled may be invalid or they will mis-guide the learning of the target stochastic process.

Author response:
We note that despite the sampling of imperfect neural SDE, our samples for X(0) are perfect. This means that the training dynamic is not misguided at initial timestamps. This fact is used to sharpen our bounds in Theorem 4.2. From the theorem, it can be seen that improvements even in ‘misguided’ regions actually improve the bound on the strong convergence, which is an interesting observation by itself. It indicates that learning is robust to misguiding. The theoretical mechanism of incremental accumulation of the errors, instead of blowing up, is explained in Lemma F.5. It is worth additional research on whether using a mechanism of adaptive/reweighting might improve the bounds and further accelerate the convergence, which is out of the score of the current work.

Reviewer comment (Theory):
The framework does not seem to allow estimation of non-diagonal observables, eg momentum.

Author response:
Our framework allows estimation of the momentum in the following ways:

1. The first way is estimating $\mathbb{E} \mathbf{P}(t)$ as $ \mathbb{E} v(X(t), t) = \mathbb{E}(v(X(t), t) \pm u(X(t), t)$. “Quantum fluctuations” by Nelson proves that this expectation of the momentum is the expectation of the osmotic velocity. However, this approach works only for the first moment of the momentum; higher order and arbitrary expectations (i.e.,  $\mathbb{E} g(P(t))$ for arbitrary nonlinear g) can’t be obtained that way.

2. Another way works if a potential is quadratic (e.g., a harmonic oscillator). The Fourier transform allows re-writing the Shrödinger equation in momentum coordinates, which also has a quadratic potential and allows diffusion interpretation. Therefore, a diffusion can be obtained as $\mathrm{d}P(t) = (\nu u(P(t), t) + v(P(t), t))\mathrm{d}t + \sqrt{\frac{\nu \hbar}{2}}\mathrm{d}W^{P}(t)$. This allows us to estimate all momentum $f(P(t))$ for any non-linear $f$. This, however, implies that we can’t obtain an arbitrary nonlinear pointwise observable $g(X(t))$ at the same time.

3. Moreover, an arbitrary non-linear observable can be obtained by recovering $p(x, t)$, $S(x, t)$ from samples using learned $u(x, t), v(x, t)$. This procedure is described in Appendix G.6. However, it is as computationally expensive as classical methods.

4. Finally, in Appendix G.4, we describe how to possibly measure observables of a form $g(X(t’), X(t))$ for arbitrary nonlinear g and  $t’ < t$ by incorporating a measurement process. It requires applying our method at least twice (on an interval [0, t’], and then on [t’, t]). This requires further research on whether deep learning can effectively capture initial conditions at the moment of the measurement $t’$ as an input to neural networks to avoid applying our approach $B+1$ (see definition in Appendix G.4) times but only twice. We leave this direction for future work.

We are ready to add more details regarding momentum estimation to our paper.

Reviewer comment (Experiments):
While I understand the complexity argument for evaluating the loss in O(d), it is unclear to me whether this method can scale better than traditional approaches. In particular, experiment of figure 4 studies the harmonic oscillator for d=1..9. However, if I understand the setting, the Hamiltonian is separable in the dimensions and so the problem has a natural scaling linear in d. I am therefore unsure about the promise for a more favourable scaling than traditional methods.

Author response:
Traditional methods would demand $O(N^d)$ computations/memory, while a PINN would require at least $O(d^2)$ computations/memory. It holds regardless of a problem they are applied to, given no external knowledge of the potential function properties. At the very least, our method does not introduce a computational overhead beyond what is needed to compute the harmonic potential (which is $O(d)$ computations in this non-interacting particles case). This advantageous behavior is observed without any reliance on prior knowledge about the separability of the Hamiltonian. Our method does not rely on a grid in X, offering a significant advantage in terms of computational efficiency. Additionally, compared to PINN, our DSM approach avoids computing the Laplacian, which ensures that no overhead is added in the asymptotics beyond what is needed to compute the potential itself. That establishes lower bounds on the computational complexity of our method; and this bound is sharp for this problem. 

To explore the dimension dependency fully, we need to account for factors that arise from the error of discretizing SDE, a universal approximation bound, and stochastic gradient descent dynamics that might also give additional dependencies on the dimension. Addressing these aspects requires further research beyond the scope of the current work. However, recent works on diffusion models established a few interesting results on the dimension-free universal approximation of particular classes of score functions with deep learning, which relates to our work since $u(x, t)$ is a score function [1].

[1] Lee H. et al., Convergence of score-based generative modeling for general data distributions

Reviewer comment (Novelty):
The use of stochastic formulation of quantum mechanics for quantum mechanics simulation is not necessarily novel, however I have not seen references to it. I am not an expert in this field, but one example I found is [Quantum Dynamics with Trajectories, Robert E. Wyatt].

Author response:
We thank you for pointing out this reference! In the revised version of our paper, we will incorporate a citation to Wyatt's paper to ensure proper acknowledgment of the related work. In fact, our work interpolates between the Bohmian picture, explored in Wyatt’s paper, and the Nelsonian interpretation through parameter $\nu$. Many results from our work can be seen as generalizations and further development of these ideas with deep learning. For example, the equidistribution principle (EP) from Bohmian interpretation is generalized to the Nelsonian picture in our Lemma F.6. Moreover, Wyatt’s work primarily focuses on reconstructing a density from sampled trajectories, but not on efficiently solving the equations guiding those trajectories. From that point of view, Wyatt’s work is complementary to ours, as it proposes how one can reconstruct the density once you have a sampler of trajectories. This extension might open a way for solving a problem of estimating all possible non-diagonal observables if the convergence of the reconstruction can be ensured up to high order derivatives in a suitable sense (e.g., in Sobolev space $W^{p}_{2}$).

Reviewer comment (Writing):
Below eq 5, one of the two u's should be v.

Author response:
It’s fixed; thank you!

Reviewer comment (Experiments):
Can you benchmark the method against traditional solvers for an interacting problem in higher dimension and show a more favourable scaling?

Author response:
There are experiments in the original submission that do correspond to interacting particles. **For this rebuttal, we ran new experiments for larger d to show scaling properties more clearly.** The following Tables illustrate time and memory usage for the Crank–Nicolson (C-N) solver for d=2, 3, 4 and training time/time per epoch and memory usage of our DSM method for d=2, 3, 4, and 5. Also, Figure 11 in **the revised supplementary material D.5.1** shows predicted densities for different dimensions. Note that while we don’t have the baseline for d=5, we believe DSM predictions are still reasonable.


Table 1. Time to get a solution and memory usage of the Crank-Nicolson method for different problem dimensions.
|      |       d=2      |       d=3      |       d=4      |
| ----------- | ----------- |----------- | ----------- |
| Time (s)      | 0.75       |    35.61   |  2363    |  
| Memory (Gb)   | 7.4        |10.6        | 214       |

Table 2. Training time, time per epoch, and memory usage of our DSM method for different problem dimensions.
|      |       d=2      |       d=3      |       d=4      |       d=5   |
| ----------- | ----------- |----------- | ----------- |  ----------- | 
| Training time (min)      | 29.5       |    60.3   |  97.5    |  154.0 |
| Time per epoch (s/ep)      | 0.52       |    1.09   |  1.16    |  1.24 |
| Memory (Gb)   | 17.0      | 22.5      |  28.0    | 33.5 |

Memory: In particular, DSM memory usage and time per epoch grow linearly in d (according to our theory and evident in our numerical results) in contrast to the C-N solver, whose memory usage grows exponentially since discretization matrices are of size  $N^d \times N^d$.
We cannot run the C-N method for d > 4 with the amount of memory available in our computer system. The results show that our method is far more memory efficient for larger d.

Compute time: While our DSM total compute times (including training time) are longer than C-N compute times for small d, the trend as d increases suggests the computational efficiency of our DSM method scales much better with d than the C-N method. In this example, the potential function has computational complexity $O(d^2)$ as it has all pairwise interactions. It implies that the training time complexity is bounded to be quadratic in the best case for all algorithms. Computation complexity of other DSM loss terms that involve computationally heavy differential operators remains to be linear. We assume domination of the differential operators computations explains why time per epoch does not seem to grow quadratically. These findings support the claim that our DSM method training time scaling is more favorable than scaling of the classical methods.

Reviewer comment (Theory):
How would one estimate the expectation of a non-diagonal operator?

Author response:
We replied to this in **W1** and partially in **W3**.