Paper: Learning invariant representations of time-homogeneous stochastic dynamical systems (ICLR-2024)

Reviewer comment (Presentation):
The organization and the explanation of the method could be more straightforward. For example, Algorithm 1 is rather terse, and a more expanded version would be instructive to the reader. It's hard to get a birds-eye view of the method because the explanation is spread throughout the text. Perhaps a separate section in the supplementary materials that gives a more direct explanation of the method would help the reader.

Author response:
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.

Reviewer comment (Reproducibility):
The method requires computing covariance matrices of neural network outputs, which could affect the computation if the output dimension is high.

Author response:
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

Reviewer comment (Reproducibility):
In the beginning paragraph of section "5 Experiments", the authors state they only optimize over the learning rate to keep things fair, but I can't seem to find the set of learning rates that were optimized over.

Author response:
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.

Reviewer comment (Experiments):
The authors state at the top of page 7 that the mini-batch size should be sufficiently large to mitigate the impact of biased stochastic estimators. But, some neural network models benefit from the mini-batch stochasticity of SGD, as opposed to full-batch normal gradient descent. In this case, I would have liked to have seen optimizing over the batch sizes as well, and not just the learning rate.

Author response:
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.

Reviewer comment (Experiments):
I'm also not sure if optimizing over just the learning rate is completely fair. I would have liked to see optimizing over more hyperparameters, or at least in addition to the current experiments, to also provide the performance evaluation of the original model implemented in the paper, if possible. For example, in the fluid flow example, Azencot et al. (2022) has a similar fluid flow example, but uses a different architecture and I would have liked to see this one in comparison to the author's method.

Author response:
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).

Reviewer comment (Experiments):
In the molecular dynamics example, I see Ghorbani et al. (2022) have done similar experiments, but at a smaller scale. Even though the authors performed their experiment at a larger scale, I would have liked to have seen a comparison with Ghorbani et al. (2022) at the smaller scale, or even at the larger scale. This would allow comparing accuracy and training times.

Author response:
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.

Reviewer comment (Experiments):
It seems in most experiments, the output dimension of the neural networks, $r$, is low. Would this be the case for most practical applications?

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

Reviewer comment (Experiments):
On page 7 in "Downstream tasks", in order to make predictions of the next step, the authors need to compute: $(\hat{\Psi}_{w}^\top)^\dagger$, i.e. the pseudo-inverse, is that right? If so, it seems this would slow down the computation of the predicted next steps, so how do the computation times compare with other methods for predicting next steps? Any numerical instabilities?

Author response:
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).

Reviewer comment (Reproducibility):
In Algorithm 1, the authors use $\hat{\mathcal{S}}^\gamma_m$, implying this is the relaxed version, is that right? If so, this isn't stated anywhere in Algorithm 1. Perhaps the authors should have two Algorithms that give the regular DPNets and also DPNets-relaxed.

Author response:
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.

Reviewer comment (Reproducibility):
When training the unrelaxed DPNets, how was $\mathcal{P}^\gamma$ evaluated? It is stated in page 5 that in general the covariances are non-invertible and thus $\mathcal{P}^\gamma$ is non-differentiable during the optimization. Also, as the authors state, the use of the pseudo-inverse can introduce numerical instabilities. How much of a problem is this, in your experience? I see that some experiments were unable to use unrelaxed DPNets because of this.

Author response:
$\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.

Reviewer comment (Presentation):
For Figure 3, in the upper panel - "Ordered MNIST accuracy", I'd imagine that the "Reference" (dashed-line) is just random guessing, i.e. $1/5 = 0.2$? I don't see this stated anywhere, and would like it explicitly stated.

Author response:
This is right, we are now stating it explicitly, thank you.

Reviewer comment (Presentation):
For Figure 3 in the lower panel, "Fluid Flow RMSE", how many prediction steps were taken? Namely, what are the values of the $x$-axis?

Author response:
40 prediction steps, that is the full test set.

Reviewer comment (Presentation):
On page 5, on the line above equation (8), the authors say "we show in Lem. 3 in App. B", but I think you mean Section C in the appendix/supplementary materials.

Author response:
That’s right. We have now fixed it, thank you.

Reviewer comment (Writing):
On page 6, in "Learning the representation", might also need to state that we are also given $\psi'$, i.e. "In practice, given a NN architecture $\psi$ and $\psi'$."

Author response:
Absolutely. Fixed it as well in the newly uploaded draft.

Reviewer comment (Evaluation):
This is mainly an empirical paper, as it proposes to model a certain space using neural networks with a gradient descent optimised objective. However, the empirical evaluation of the method is not completely convincing, and some points are not clear.

Author response:
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.

Reviewer comment (Experiments):
Even the relaxed objective introduced in the paper is computationally heavy. It involves computing the covariance matrices as _differentiable functions_. This means the computation graph involves summing \( n \) matrices of size \( r \times r \), each coordinate of which is a differentiable function. As noted in the paper, minibatches could be used but would introduce bias. Even with minibatches, the representation dimension \( r \) would have to stay low, and appears to be low in all experiments (in which it was specified).

Author response:
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.

Reviewer comment (Evaluation):
**Ordered MNIST:** It is unclear why other methods fail on this example, which is apparently new in this paper. Is there an explanation or intuition for this? Are convolutional networks used to model the representation? It is assumed the non-NN methods fail simply because representing images is easier with convolutions. Why do other methods fail? Explaining this difference could make the paper considerably stronger.

Author response:
**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.

Reviewer comment (Evaluation):
**Logistic Map:** The takeaway from this experiment is not clear. DPNets has the worst spectral error among all methods. In addition, there appears to be no theoretical reason why DPNets-relaxed should perform better on this metric, as it was introduced as a computational relaxation. Do the authors suggest it has properties allowing it to be a better estimator in general?

Author response:
**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.

Reviewer comment (Presentation):
**Continuous dynamics:** The scale of Fig. 2 lower pane is that of 200, which makes it very hard to understand whether the smaller eigenvalues are estimated accurately and what is their order of magnitude. Displaying ratios of estimated/true values would be more informative. Also, are there other ways to estimate the eigenvalues?

Author response:
**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.

Reviewer comment (Evaluation):
**FlowRMSE:** The general magnitude of the quantities in this experiment is not clear. If the quantities are of order 10, then the difference between errors of order 0.1 and 0.001 is less impactful, and all methods perform well. What are the typical values of the quantities for which RMSE is computed? Percentage ratios would be much clearer than absolute values.

Author response:
**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.

Reviewer comment (Evaluation):
**Chignolin Experiment:** In Table 2, describing the Chignolin experiment, while estimated quantities are somewhat closer to the reference values than the true quantities, they still look quite far. It is thus not clear what is the point of this part of the experiment. The second part of that experiment, involving Fig. 4, unfortunately cannot be understood from the main text (e.g., what is the free energy surface, what is known about its proper minima, why matching the results of Bonati et al. is of interest).

Author response:
**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.

Reviewer comment (Novelty):
**Question:** Is the relaxation objective (9) a new contribution of this paper, or were similar relaxations employed in previous work (for instance VAMPSNets, DeepCCA)?

Author response:
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.

Reviewer comment (Experiments):
The pipeline design suggests that there might be a possibility of error accumulation from the first phase of learning representations to the second phase of learning dynamics. It would be beneficial if the authors could provide empirical justification to show that their design is comparable to or better than learning the neural operator in a single step. For instance, establishing baselines using Frameworks like FNO or Deeponet would be ideal.

Author response:
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.

Reviewer comment (Experiments):
Tying in with Point 1, the efficacy of the proposed method, particularly the operator regression in the second part, remains uncertain when applied to chaotic systems where acquiring accurate representations in the initial phase may prove challenging.

Author response:
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).

Reviewer comment (Evaluation):
It would be great if the authors could evaluate the performance on insightful metrics like Lyapunov spectrum. And including discussion on recent related works (e.g., https://arxiv.org/pdf/2304.12865.pdf) should be necessary.

Author response:
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).

Reviewer comment (Experiments):
If the primary goal is representation learning, I feel that discussions and comparisons related to popular representation learning methods (e.g., contrastive learning) should be discussed and included.

Author response:
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.

Reviewer comment (Experiments):
I am curious about the empirical performance of the results concerning different hyperparameters, particularly the value of $r$.

Author response:
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.

Reviewer comment (Experiments):
Additionally, what would be the loss for the OLS estimator of the transfer operator across different datasets? It would be concerning if the performance is highly sensitive to the value of $r$.

Author response:
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.

Reviewer comment (Presentation):
For the dynamical data, i.e., the 2D fluid flow and the molecular dynamics, it would be beneficial if the authors could provide the visualizations of the dynamics.

Author response:
We have added visualizations for both experiments (see Figures 8 and 9).

Reviewer comment (Experiments):
The computation for covariances for $\psi_{w_j}$ and $\psi'_{w_j}$ might be costly. How does this computation cost compare to the previous methods? Is there a way to speed up?

Author response:
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)$.

Reviewer comment (Experiments):
For the continuous-time dynamic, the score function might be unstable. Have you observed this in experiments?

Author response:
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.