Learning the Infinitesimal Generator of Stochastic Diffusion Processes

We appreciate the reviewer's insightful evaluation and valuable comments. In what follows, we aim to address the highlighted weaknesses and respond to the reviewer's questions.

Weaknesses:

1. Thank you for emphasizing the realistic setting of imperfect partial knowledge. This motivated us to show theoretically and empirically that indeed our IG learning method with estimated Dirichlet form is guaranteed to work. Concerning the sampling, our assumption to have i.i.d. data from the invariant distribution $\pi$ can be relaxed as we reply in Q6.
2. Concerning the experiments, please see our general reply.
3. We agree that more discussion of Fig. 1 is needed to enhance the understanding of the experiment section. If accepted, an extra page will allow us to improve it, as well as to include additional content from the attached PDF and this rebuttal in the revision.

Questions:

• Answers to both Q1 and Q2 are positive. Please see our general reply for the discussion.
• Q3: TO methods apply only to equally spaced data and the sampling frequency $1/\Delta t$ must be high enough to distinguish all relevant time-scales. Otherwise, since TO eigenvalues are $e^{\lambda_i \Delta t}$, small spectral gaps complicate learning (see Thm. 3 [22]). Conversely, our IG method, which uses gradient information, is time-scale independent, handles irregularly spaced measurements, and does not rely on time discretizations (see reply to Q6). While it results in better generalization, as shown in Fig. 1 c)-d)-e) where our IG estimator captures ground truth significantly better than TO for the same sample size, this generalization across all time scales incurs quadratic computational complexity w.r.t. state dimension and not statistical accuracy (see also reply to Q7). Lastly, with our additional discussion on imperfect knowledge, our IG method can be safely applied in a fully data-driven regime.
• Q4: For i.i.d. data from $\pi$, at this point, it is hard to answer this question, while for trajectory data, important improvement can be expected, see reply to Q6. While the standard assumptions (SD) and (KE) are the same for TO and IG error bounds, regularity assumptions in these two settings are not easily comparable. This motivates the development of non-standard regularity assumptions for TO and IG learning that could expose the benefit of the prior knowledge. This is an interesting problem in its own right, and it will be the subject of our future work.
• Q5: As illustrated in Fig. 3 in the attached pdf, estimators are quite robust w.r.t. value of the shift parameter $\mu$. Theoretically, it is important that $\mu$ is not too small.
• Q6: Thank you for raising this question, we will emphasize these important aspects of our method in the revision.
  • Recalling the risk functional in eq. (8), we see that the “label” of the model $\chi_{\mu}(x) \approx G^*\phi(x)$ is the action of the resolvent. Since this “label” is not computable, we “fight fire (resolvent) with fire (generator)”, i.e. we use the energy norm of eq. (9) to rewrite regularized problem (8) as in line 222. Crucially, this allows us to obtain estimators via energy covariance $W_{\mu}$ in (14) that completely captures infinitesimal nature of the learning problem without needing time-lagged observations. This contrasts with TO methods, where choosing the time-lag $\Delta t$ is the major bottleneck in real applications, [18,42].
  • Thank you for this reference (Meng et al. 2024) that we will cite in the revision. Note that this method is exactly the Galerkin projection of the Yoshida approximation of IG on a finite-dimensional RKHS (dictionary of functions). It essentially corresponds to solving (8) where $\phi$ is finite-dimensional, expectation is used instead of energy and $\chi_{\mu}$ is replaced with $\mu^2\chi_{\mu}(x)-\mu\phi(x)$. To do this, authors estimate resolvent via equation in line 167 which requires approximating an operator integral via equally spaced data with high sampling frequency (small $\Delta t$), essentially suffering from the same bottleneck as TO methods. Since in (Meng et al. 2024) the provided learning theory and the empirical evaluation is limited, guarantees and limitations of this method are not fully clear to us.
  • In practical applications to obtain data from an invariant distribution $\pi$ one uses the trajectory data after some burn-in time needed to ensure that ergodic mean approximates well $\pi$. Then, the problem is reduced to studying only the dependence as is done e.g. in [21] using standard tools of $\beta$-mixing and the method of blocks. This allows one to obtain non-parametric learning bounds for TO methods, c.f. [22], where the effective sample size suffers from the multiplicative effect of the time-lag to achieve approximate independence. Recalling Q3, TO methods “waste” a lot of data negatively impacting statistical accuracy (we expect similar limitations for the method of (Meng et al. 2024)). Contrary to this, our method can be applied to data with larger time-lags (even irregularly spaced) so that effective sample size is close to the true one.
• Q7: Compared to FEM, the key difference lies in the approximation error $|\lambda_k - \widehat{\lambda}_k| \leq c h^{2p} |\lambda_k|$, where $p$ is the polynomial degree used to construct finite elements, $h$ is the mesh size, and $c$ depends on the eigenfunctions' smoothness. As the number of mesh elements grows exponentially with $d$, i.e., $\sim h^{-d}$, reducing $h$ is the major bottleneck mitigated by computationally demanding adaptive higher order methods. On the other hand, our IG method that requires less or no knowledge has a quadratic impact of the $d$ only on the computational complexity. Indeed, sample complexity depends on the effective dimension of the equilibrium distribution on the domain that can be much lower than $d$.

I thank the authors for the extensive rebuttal in the general reply and comprehensively addressing my review comments. Also, it is great to see the authors put in the effort to deliver adjusted theory as well as run new experiments.

Many important aspects of the paper got clearer: limits of FEM and TO approaches, fully-data driven regime, energy functional rationale and shift parameter. Accordingly, I have increased my overall score.

I expect this rebuttal content to be incorporated as part of a final submission.

We appreciate the reviewer's insightful evaluation and valuable comments. In what follows, we aim to address the highlighted weaknesses and respond to the reviewer's questions.

Weaknesses:

Major:

• Indeed, the topic of learning IG of a stochastic process with kernel-based methods, and, in particular, development of nonparametric learning rates contains several complex concepts. More so, since the paper is the first one to formalize a consistent way to learn IG from data. We tried our best to present the sophisticated theory and mitigate the complexities, believing that this work paves the path to efficient and reliable algorithms for important applications in ML (see also general reply).
• In addition to our general reply, we would like to emphasise here that, from molecular dynamics to weather models, the ability to reliably learn eigenfunctions of differential operators that generate the dynamics is of capital importance for building trustworthy and physics informed AI. Hence, we strongly believe this work is of high relevance, if not for the most general ML community, definitely for the AI-for-Science community. With that respect, we would like to further stress that our method is the first of its kind in machine learning to have error bounds that theoretically demonstrate the superiority of an ML approach for spectral estimation of differential operators over classical numerical methods like FEM in high-dimensional settings. While empirical evidence supporting this has been well documented in practice, particularly in fields such as molecular dynamics [42], to the best of our knowledge, this had never been theoretically proven until now. Following reviewer's suggestions, in the revised manuscript we will better emphasize possible applications and further improve readability and accessibility of our results.

Minor:

• Thank you for spotting the typos, we will correct them in the revision.

Questions:

As we elaborate in the general reply, indeed we can show, both theoretically and empirically (Figures 1 and 2 in the attached pdf), that our method works in a fully data-driven regime. While empirically this might not come as a big surprise, we believe that adding a theoretical proof motivated by the comments of reviewers is a very nice additional contribution. If the reviewer wishes to have more details on this aspect, we are happy to elaborate.

We appreciate the reviewer's insightful evaluation and valuable comments. In what follows, we aim to address the highlighted weaknesses and respond to the reviewer's questions.

Weaknesses:

Main contributions and broader impact. In the general reply we tried our best to provide more details on positioning of our main contributions in the larger context of solving SDEs, and help the reviewer evaluate the possible impact of the presented results.

Experiments. While we would like to stress out that we originally provided quantitative comparison to the cited methods ([1] is in fact just a special case of [2], while [3] is not applicable to Langevin dynamics) in Fig. 1 a) which implies the failure of this method to capture the true dynamics, we follow the reviewer’s suggestion and provide additional experiments that make this point more clear for the reader.

Choice of hyper-parameters. As suggested, we provide the empirical evaluation of the robustness of hyper-parameter choice, as indicated in the general answer. In summary, the robustness of the choice of kernel’s hyper-parameters and Tikhonov regularization is similar to general kernel methods for (vector-valued) regression. On the other hand when estimating IG via its resolvent the choice of additional shift hyper-parameter $\mu$, as expected from the theory, has a much smaller impact on the performance, see Fig. 3 of the attached pdf. Finally, we wish to stress out that we do go beyond classical validation via prediction by introducing novel quantities, the empirical spectral biases $\widehat s_i$ of eq. (24), which are (theoretically and empirically) good metrics to use for fine-tuning the model.

Notations. Thank you for your suggestion. We will prepare a smaller summary table that only introduces necessary notations for presenting the IG learning framework and the method, while omitting extensive notation related to proving the generalization bounds. We believe this is feasible, since, if the manuscript is accepted, the additional page can be used to present this table in the main body of the paper.

(KE) assumption. The kernel embedding (KE) assumption originates from the study of mini-max optimal excess risk bounds for classical kernel regression [12]. Whenever the kernel is bounded, (SD) and (KE) do not impose any additional constraints, they are just used to describe the interplay between the data distribution and the chosen RKHS. This is a formal way to quantify the impact of the kernel choice to the learning rate. We will additionally clarify this in the revised manuscript.

Dirichlet form. To complement the discussion on this issue from the general reply, we would like to stress out that many diffusion processes have an infinitesimal generator that can be expressed in the form (IG). In addition to the Langevin and the CIR processes already studied, we can mention:

• the Wright-Fisher diffusion (in dimension one), which can be defined in the context of population genetics and can be adapted to model interest rates, see [G],
• the geometric Brownian motion which models the price process of a financial asset,
• the multi-dimensional Brownian motion (a=0) that corresponds to the heat equation,
• the transport processes associated with advection-diffusion equation, see [H]
• the process associated with Poisson's equation in electrostatics, see [I].

Minor remarks. Thank you for pointing out the typo, we will fix it. Concerning the claim in line 178, it relates to the first equality in Equation (7). There, one can see that knowing drift and diffusion terms, we can evaluate $[L f] (x)$ and approximate the energy as empirical mean of $\mu |f (x) |^2 + f(x)[L f] (x)$ via samples $x$ from the invariant distribution.

Additional references:

[A] Fukushima, Masatoshi, Yoichi Oshima, and Masayoshi Takeda. Dirichlet forms and symmetric Markov processes. Vol. 19. Walter de Gruyter, 2011

[B] Jean Jacod. Discretization of processes. Springer, 2012

[C] Danielle Florens-Zmirou. Approximate discrete-time schemes for statistics of diffusion processes. Statistics: A Journal of Theoretical and Applied Statistics 20.4, 1989

[D] Fan, J., & Zhang, C. (2003). A reexamination of diffusion estimators with applications to financial model validation. Journal of the American Statistical Association, 98(461), 118-134.

[E] Yuri Kutoyants. Statistical inference for ergodic diffusion processes. Springer Science & Business Media, 2013. Yuri Kutoyants. Parameter estimation for stochastic processes, 1984.

[F] Fabienne Comte,Valentine Genon-Catalot Nonparametric drift estimation for iid paths of stochastic differential equations. The Annals of Statistics, 48(6), 2020

[G] Martin Grothaus, Max Sauerbrey. Dirichlet form analysis of the Jacobi process. Stochastic Processes and their Applications, 157, 2023.

[H] Antoine Lejay, Lionel Lenôtre, Géraldine Pichot. Analytic expressions of the solutions of advection-diffusion problems in one dimension with discontinuous coefficients. SIAM Journal on Applied Mathematics, 79(5), 2019.

[I] Sethu Hareesh Kolluru. Preliminary Investigations of a Stochastic Method to solve Electrostatic and Electrodynamic Problems. Phd Thesis, 2008.

We appreciate the reviewer's insightful evaluation and valuable comments. In what follows, we aim to address the highlighted weaknesses and respond to the reviewer's questions.

Weaknesses:

Thank you for pointing this out. Indeed, clearer motivation beyond the Dirichlet form should help the reader understand our approach better. While some aspects of this are addressed in the general reply, in the following we try to address your concerns in detail.

Questions:

1. As we have discussed in the general reply, a large class of processes can be written in Dirichlet form and, importantly, for many of them statistical estimation of the Dirichlet coefficient (related to the diffusion coefficient) from data is much easier than recovering the drift term. This means that our method is possible to apply even when the challenging drift term is not estimated. That said, we are, to our knowledge, the only existing method that is able to learn the infinitesimal generator model from a single trajectory that is theoretically consistent with the true dynamics. Finally, as we show in the general reply, note that our method is robust under imperfectly estimated partial knowledge. This is not only proven, but also empirically demonstrated in Fig. 1 and 2 of the attached pdf.

2. Indeed, there is the consistency of our IG method and the properly executed TO methods. This is theoretically backed by Thm. 2 in [22] for TO and our Thm. 2 for IG. However, note that TO methods apply only to equally spaced data and the sampling frequency $1/\Delta t$ must be high enough to distinguish all relevant time-scales. Otherwise, since TO eigenvalues are $e^{\lambda_i \Delta t}$, small spectral gaps complicate learning (see Thm. 3 [22]). Conversely, our IG method, which uses gradient information, is time-scale independent, handles irregularly spaced measurements, and does not rely on time discretizations. This results in better generalization, as shown in Fig. 1 c)-d)-e) where our IG estimator captures true phase transitions between meta-stable states significantly better than TO for the same sample size. Lastly, with our additional discussion on imperfect knowledge, our IG method can be safely applied in a fully data-driven regime as TO methods.

3. Thank you for spotting this. Indeed, in App. F we didn’t specify that the RBF Gaussian kernel with specified length-scales was used in all cases. As reported in the generall reply, the robustness of kernel hyperparameters and Tikhonov regularization is essentially the same as for TO kernel estimators, while the IG estimator is very robust against the new shift hyperparameter, see Fig. 3 of the attached pdf.

4. Thank you for this question. Besides original comparisons of the time-scales (IG eigenvalues) and meta-stable states (IG eigenfunctions), providing the additional experiments, we demonstrate our method’ consistency with true dynamics from the perspective of equilibrium distribution recovery and long term forecasting of conditional cumulative density functions in Fig. 1 and 2 of the attached pdf, respectively. This might come as a surprise, but after a closer look, the cited works, only showcase empirical evaluation of the leading eigenfunctions, which due to the reported effect of spuriousness in Fig, 1a requires picking proper eigenpairs from an abundance of wrongly estimated ones. Indeed, if one builds such models without expert knowledge needed to remove spurious estimation, the obtained dynamics is inconsistent with the true one, see caption of Fig. 1 of the attached pdf.