Learning invariant representations of time-homogeneous stochastic dynamical systems

We gratefully thank the reviewer for their review and comments, which we now address.

Weaknesses:

• As you suggested, we have included an expanded version of Algorithm 1 in the supplementary section. Furthermore, upon completion of the review process, we intend to publicly share the GitHub repository containing not only the method's implementation but also the implementation of each experiment.

• Thank you for raising this important point. Please see our global reply where we address this point in detail, as well as our answer to reviewer hdvQ for further comments on the computational complexity

• Experiments:

  • The grid used for the search of the best learning rate was reported in appendix F: 100 equally spaced values in the interval $(10^{-6}, 10^{-2})$. To aid reproducibility, in the updated version we also added a table with the learning rate found by the search for each (experiment, method) pair. 
  • We believe this is a valuable suggestion, and while we keep the batch-size for DPNets and VAMPNets as high as possible to not incur any bias upon estimation of the score function (see our discussion in Sec 4), we are in the process of re-running the Auto-Encoder methods with smaller batch-sizes. We will give an update on this point as soon as the new results are available.
  • Given the number and diversity of the baselines (each with a different set of hyperparameters) the evaluation methodology we devised is to fix every parameter defining the model expressivity (architecture, feature-dimension) and train the best possible model within a budget of epochs. Hence we optimized learning rates and (thanks to your suggestion) batch sizes. Taking a closer look at the fluid flow example in Azencot et al. (2022), one can see that our setting is remarkably close to theirs: almost identical underlying PDE, same Reynolds number, as well as similar size of the dataset. With the exception of a different activation function, we also use a similar MLP architecture (though the one in our experiment is slightly bigger, with 3.8 million parameters compared to theirs which is 2.49 millions). Finally, our implementation agrees with their findings, for which in this experiment they expect results in line with dynamical AEs (see our Fig. 1).
  • The experiment in Ghorbani et al. (2022) considers different proteins. Their method (which is equivalent to VAMPNets) fails on the larger scale Chignolin experiment as we report in Sect. 5. On the other hand applying our method to their smaller scale dataset should be doable, although the hyperparameter tuning may be time consuming. While we cannot guarantee to have it done in the time span before the end of the rebuttal period, we will definitely include this experiment in the revised version. Questions:

• The answer is yes. Please see our global reply for an in-depth discussion of the computational complexity of using covariances.

• That formula is just for definition purposes. In practice we simply compute it as $(\widehat{\Psi}^\top_w)^\dagger \widehat{F}’ =( \widehat{\Psi}^\top_w \widehat{\Psi}_w)^\dagger \widehat{\Psi}_w^\top  \widehat{F}’$, which is equivalent to solving $\ell$ least squares with $r\times r$ matrix $\widehat{\Psi}^\top_w \widehat{\Psi}_w$. Finally, notice that this is performed only once when fitting the operator regression model. Once computed it can be stored and used to predict future states through simple matrix multiplications. Additionally, low rank estimators, as well as Nyström based methods are available to considerably speed up this initial fitting cost even when $r$ or $n$ are very large (see Meanti et al. 2023). 

• In Section 4, just before the “Operation regression” paragraph we actually state: In Alg. 1 we report the training procedure for DPNets, which can be also applied to the other scores (8) and (14) in step 4, respectively. When using mini-batch SGD methods to compute ∇ ̂ Pγ m(w) or ∇ ̂ Sγ m(w) via backpropagation at each step, m ≤ n should be taken sufficiently large to mitigate the impact of biased stochastic estimators”. As you point out, the algorithm is for DPNets-relaxed, while in the main text we stated that it is for DPNets. We have now corrected this.

• $\mathcal{P}^\gamma$ was evaluated using torch.pinv with the “hermitian” flag set to True. The problem of using pseudo-inverses lies in the backpropagation step, as the pseudo-inverse is known not to be differentiable whenever the rank is not stable over the parameter space (see Golub & Pereyra, 1973). On the other hand, on the forward pass, the torch implementation seemed to be stable enough, and when we weren’t able to train DPNets (unrelaxed) and VAMPNets it was a failure of the back-propagation step. 

• This is right, we are now stating it explicitly, thank you.

• 40 prediction steps, that is the full test set.

Remarks:

• That’s right. We have now fixed it, thank you.
• Absolutely. Fixed it as well in the newly uploaded draft.

I thank the authors for their detailed response, revisions, and additional experiments. I have read the other reviews and the discussions that followed (as of this writing).

My questions and concerns have been mostly addressed. Thank you for the reference, although I'll note that in reviewing the reference and in the context of the submission, they mention when training in large batch regimes, it may help to tune other hyperparameters as well, as they can differ from their small batch regimes (e.g. learning rate, momentum, etc.). In fact, in Section 4.7 of the reference (Shallue, 2019) they explicitly state:

"We also found that the best effective learning rate should be chosen by jointly tuning the learning rate and momentum, rather than tuning only the learning rate." [emphasis mine]

So it may help to also tune the momentum if the authors also believe that will be productive. But I do recognize and sympathize that hyperparameter tuning can be computationally expensive.

But my current thoughts are: after reading the response and the other reviews, and due to the limited gradations of the scoring system (score would go up if there were sufficient gradations), I am inclined to have my score remain unchanged. I acknowledge the hard work put into the rebuttal and revisions, and am very grateful and extremely appreciative to the authors.

We gratefully thank the reviewer for their review and comments, which we now address.

W1. We recall that, among the baselines we analyze, both KernelDMD, DMD, DynamicalAE and ConsistentAE learn the transfer operator in a single step. Indeed, the first two do not require to pre-learn a representation, while the latter two optimize a sum of losses encompassing both representation (i.e. the reconstruction error of the AE) and dynamics. As explained in Sect. 2, and specifically in the paragraph “Learning transfer operators”, off-the-shelf representations often used alongside DMD-alike methods are detrimental in that they induce metric distortion as well as representation biases. In contrast, DPNets are designed to mitigate these effects. Further, the discussion around Eq. (3) entails that for all compact transfer operators, the cumulative error of DPNets plus consistent operator regression is controlled; This is also what we empirically show in every experiment. Does the reviewer agree on this point? We’d be happy to clarify any aspect of our response if needed.

Regarding establishing baselines using FNO/Deep-O-net, we would like to highlight that these techniques are solving a different problem. Although the concept of “neural-operator” is quite broad, the methods you mentioned have been specifically designed to learn surrogate maps for the solution operators of PDEs; see (Kovachki et al., 2023). Moreover, these methods often require knowledge of the underlying physical equations to generate the training data, which is something not at hand for general stochastic processes. Apart from the fluid flow example, for which we already included a physics-informed baseline, the applicability of FNO/Deep-O-net to solve the problem of transfer operator learning remains unclear. We are open to further discussion on this aspect if the reviewer deems it necessary.

W2 and W3. Ergodic theory is usually studied using the geometric approach or the operator theoretic approach. While the first is related to the study of Lyapunov exponents and fractal dimension, the latter is related to the Transfer/Koopman operators of Markovian processes. Connecting these two complementary approaches to ergodicity is an ongoing research, see e.g. (Tantet et al. 2018, Brunton et al  2017; 2022, Mezić and Avila 2023). Due to the nature of our work, and the central role played by stochastic dynamical systems, we have focused on the spectrum of the transfer operator as relevant information about the dynamics. On the other hand, the Lyapunov spectrum is much more interesting in the context of deterministic, and especially chaotic systems.  We can elaborate on both of these approaches in Sec. 1 also mentioning (Platt et al. 2023) as an example of deep learning methods for deterministic systems based on the Lyapunov spectrum approach. Would the reviewer consider this discussion appropriate and to be included in the revision? Finally, concerning the synthetic experiment of Noisy Logistic map we can report that, in the noiseless case the Lyapunov exponent is log(2), while introducing the noise in the model yields the Lyapunov exponent (in the stochastic sense) equal to 0, see (Baxendale, 1989). 

W4. Since representation learning is a broad field, in the “Previous work” paragraph we only covered representation learning in the context of transfer operator (Azencot et al., 2020; Lusch et al., 2018; Morton et al., 2018; Otto and Rowley, 2019). We are not aware of contrastive learning approaches in the context of dynamical systems. However, we welcome any suggestion on additional references to be included in the discussion. 

Q1. In the revised version of the manuscript (Appendix F.1) we added an in-depth discussion on the role of the feature dimension in the Logistic map experiment. It is also interesting to note that DPNet showed to be a quite “parsimonious” representation, needing only a small dimension to well approximate the dynamics. This is particularly relevant also on the ordered MNIST example, where a feature map of just 5 floats achieves > 98% accuracy on validation. 

Q2.  Our DPNets method learns a representation of dimension $r$ and subsequently, an OLS operator estimator can be trained using this representation. To evaluate the quality of the learned $r$-dimensional representation it is not meaningful to look at the OLS error (or risk) but rather to the forecasting error or the spectral error, which we have reported in Q1 for the logistic map. Indeed, although the OLS error seems the natural object to look at, it does not reflect the quality of the representation. For instance, a trivial representation mapping every state to zero would have a zero error. We’d be happy to further clarify any aspect of our response if needed.

Q3. We have added visualizations for both experiments (see Figures 8 and 9).

I want to thank the authors first for their detailed reply. I acknowledge the authors' replies regarding my W1 & W2. And I appreciate the authors provided additional experiments and visualizations for my Q1 & Q3. I still hold hesitations towards related works and notice Reviewer 1QLL's question regarding the comparison of the objective function against VAMPNets and DeepCCA (for which I have not found the authors' response).

An update: After careful rethinking and as we are approaching to the closing of the discussion period, I decided to change my score to 6. I feel in general impressed by the authors’ hard work. And I hope further clarifications regarding the novelty in terms of the objective function comparing against previous works, as well as the concerns regarding the computational cost would be added into the revision.

We gratefully thank the reviewer for their review and comments, which we now address. 

While we address the specific comments on the empirical evaluation of the methods below, we respectfully disagree on the point that the paper is “mainly empirical”. Our aim is to provide a method firmly grounded in state-of-the-art statistical learning theory, as evidenced by our detailed discussions in Sec. 2 and 3, as well as App. A–E, outlining the theoretical foundations of DPNets. In reply to your question, the metric distortion loss $\mathcal{R}$, as well as the relaxed score in eq. (9) haven’t been proposed before and are novel contributions. We will be happy to discuss any theoretical question/comment, should it arise during the rebuttal period.

The point on the high cost of evaluating covariances is important, and we are grateful to the reviewer for raising it. In this respect, we refer to the general answer for a detailed account on why we expect that in many scenarios a low dimension to be sufficient and what one can do if, for any reason, the representation dimension needs to be large. 

We also refer to point 1 of our answer to reviewer hdvQ for further comments on the computational complexity.

On the specific experiments:

Ordered MNIST: this example was actually introduced in Kostic et al. 2022. In the revised version of the manuscript we now acknowledge it, thanks for pointing it out. In our DPNets methods as well as in VAMPNets, oracle-features, and in the AE baselines we use the same CNN (the specific architecture is reported in App. F.4), and as you said non-NN methods fail because representing images is easier with convolutions. VAMPNet, which is the closest to DPNet in terms of methodology, fails to converge (see our discussion on its instability in Sec 3). The DynamicalAE actually learns a reasonable representation, making it the fourth best model behind DPNets and the oracle-features (which, however, are trained using also the true labels). Upon visual inspection, indeed, DynamicAE returns readable images (see the Figure on the very last page of the Appendix). Finally, we speculate that ConsistentAE does bad in this setting as they also enforce a “consistency” constraint for the time-reversed system, which is not satisfied by the true dynamics of this example. Finally, even though oracle-features have been trained using the true labels, they are oblivious of the dynamics of this system. This is why, even attaining very good performances, they eventually fail at long term prediction.

Logistic Map: We do not have any indication that DPNets-relaxed provides a better estimation in general. Indeed, for larger feature dimensions, the unrelaxed DPNets is on par with DPNets-relaxed. To clarify the takeaway of this experiment, we significantly expanded the App. F.1 with an in-depth discussion of the inherent difficulty when estimating the spectra (high sensitivity due to nonnormality of the operator) and the role of the dimension of the feature map, which is particularly important in this context.

Continuous dynamics: For improved readability we have changed the scale to logarithmic, and plotted the implied transition rate, that is the absolute value of the non-null eigenvalues. We are only aware of the alternative methods presented in (Klus et al. 2019) and (Hou et al 2023), which correspond to DMD and kernel-DMD, respectively, adapted for learning the generator. Our example corresponds to using the DMD method coupled with a DPNet representation. 

Fluid flow: We completely agree with your comment, and in the updated draft we plot the results in terms of ratios as you suggest. Working with RMSE, we find it sensible to use the standard deviation of the data (which is approximately ~ 0.1675) as a proxy for the general magnitude of the quantities involved. 

Chignolin: Indeed, quantities such as the free energy are very specific to the context of physical sciences. In the revised version of the draft we added an explanation of what free energy is, and how it relates to the metastable states. Concerning your question about the estimated quantities being far from the true ones, it basically boils down to the amount of data at our disposal. Molecular dynamics simulations are great in that they allow to simultaneously probe the dynamics of a physical system, as well as their equilibrium properties if the simulation is long enough. For big systems, however, to accurately estimate equilibrium quantities such as the transition time or the enthalpy one needs extremely long simulations, often out of reach. The trajectory that we used, while large scale, is simply not long enough to entirely capture the equilibrium properties with high accuracy. We point out that speeding up molecular dynamics simulations is an area of active and intense research, and algorithms to estimate the equilibrium properties from short, imperfect simulations are highly valuable.

We gratefully thank the reviewer for their review and comments, which we now address. 

1. This point was raised by different reviewers and we refer to the general response for a detailed discussion. Concerning the more specific questions about the computational complexity:
   1. DPNets: the computational complexity for both the forward and backward pass is exactly the same of VAMPNets, as the addition of the metric distortion loss $\mathcal{R}$ (see also Eq. 8) is actually of $\mathcal{O}(r^2)$, $r$ being the dimension of the representation. Thus, when $r > n$, the computational complexity is dominated by the (pseudo-)inversion which scales as $\mathcal{O}(r^3)$. Kernel-based operator regression algorithms such as DMD also share the same (forward) computational complexity, requiring the evaluation and inversion of a covariance ($\mathcal{O}(nr^2 + r^3)$) or kernel matrix ($\mathcal{O}(n^3)$). These algorithms do not-require backpropagation, but often fall short in representation power — as shown in the experiments.
   2. DPNets-relaxed: The computational complexity of the relaxed score is actually lower than that of DPNets and VAMPNets, since it does not require the full inversion of a matrix, but only its leading singular value (i.e. its operator norm), which can be computed via standard Arnoldi/Lanczos methods, yielding a total complexity of just $\mathcal{O}(nr^2)$. 
2. The optimization of the score in the Langevin dynamics experiment did not show any signs of instability. We speculate that this is due to both the low-dimensionality of the state space and the infinitely-smooth ($C^{\infty}$) potential function. Indeed, since the score for continuous dynamical systems is the sum of the (finite) eigenvalues of the symmetric eigenvalue problem $C^w_{X\partial} - \lambda C^w_X$, in non-pathological cases, its value and it derivatives can be computed efficiently in a numerically stable way, see e.g. (Andrew, A. L. and Tan, R.C.E., Computation of Derivatives of Repeated Eigenvalues and the Corresponding Eigenvectors of Symmetric Matrix Pencils, SIMAX 20/1, 1998), ). We have included this remark in the revised manuscript. Notwithstanding these favorable conditions, we already acknowledged the potential instability of the score for continuous dynamics in the conclusions. While beyond the scope of the paper, an intriguing avenue for future research involves an in-depth analysis of continuous dynamics.