Conventional notions of generalization often fail to describe the ability of learned models to capture meaningful information from dynamical data. A neural network that learns complex dynamics with a small test error may still fail to reproduce its \emph{physical} behavior, including associated statistical moments and Lyapunov exponents. To address this gap, we propose an ergodic theoretic approach to generalization of complex dynamical models learned from time series data. Our main contribution is to define and analyze generalization of a broad suite of neural representations of classes of ergodic systems, including chaotic systems, in a way that captures emulating underlying invariant, physical measures. Our results provide theoretical justification for why regression methods for generators of dynamical systems (Neural ODEs) fail to generalize, and why their statistical accuracy improves upon adding Jacobian information during training. We verify our results on a number of ergodic chaotic systems and neural network parameterizations, including MLPs, ResNets, Fourier Neural layers, and RNNs.
Paper
Similar papers
Peer review
Summary
The present manuscript concerns the use of neural networks to fit physical dynamical systems of generic kind (including and focusing on chaotic maps). It is shown that adding the information of the Jacobian of the dynamical map in the supervised learning process leads to better performances.
Strengths
Originality: the work builds on previous ideas and problems. The original part (as claimed by the authors) concerns adding the information on the Jacobian of the dynamical map to the supervised learning process. The author show that this significantly improves the statistical properties of the points generated by the fitted dynamical systems. Quality: The paper is very well written, well structured and clear. Section 4 is very clear. Significance: I believe that the problem tackled in this work is a relevant one and the results show an interesting perspective on it.
Weaknesses
I believe that the most important weakness of the manuscript relies on the limitations of the results. The paper clearly shows that having information on the Jacobian of a dynamical map is essential for reproducing the statistical properties of the generated orbits. However such information is typically lacking.
Questions
1) I believe that it must be carfully stated that it is assumed that the map F does not depend on time. This is an important assumption which may be relaxed.
Rating
6
Confidence
3
Soundness
3
Presentation
4
Contribution
2
Limitations
The main and severe limitation is that the improvement training method needs the knowledge of the Jacobian of the ground truth dynamical system. It is unclear how this can be estimated from data contain just points in the orbits of the maps and what are realistic settings where this is known. I believe that this is a severe limitation of the work. This does not spoil the result in the sense that the manuscript clearly shows that the information coming from the Jacobian is essential, but on the other hand it seems to me that this information is hard to get.
Summary
The paper addresses the problem that classical ERM training of dynamical systems models often fails to capture invariant measures of the observed dynamics, even when test errors are low. The authors use ergodic theory to explain this failure from a theoretical viewpoint. They further demonstrate that incorporating Jacobian information during training enables much better reconstructions of invariant properties and validate their approach on various neural flow operators trained on several common benchmark chaotic systems.
Strengths
- I think the authors study an interesting problem, i.e. when classical ERM for dynamical models can still lead to a reconstruction of invariant properties of the observed dynamical system, as classical ERM (i.e. here simply “one-step-ahead” predictions) comes with benefits of decreased training time and less training difficulties - The authors use dynamical systems and ergodic theory to accompany and explain the practical behavior of several neural flow operators (N-ODEs, RNNs, MLPs, etc.) when training them using ERM on observed dynamical systems - The paper is generally well written
Weaknesses
- l. 53 - 56: Indeed there are also connections of training recurrent architectures to model chaotic DS and exploding loss gradients in gradient descent based training [1]. Moreover, recently methods for dealing with these training instabilities have been introduced and proven useful [1, 2]. - There is also a recent paper which discuss (out-of-domain) generalization in dynamical systems reconstruction [3] which might make sense to add to related work, as it also introduces measures to assess generalization with a focus on evaluating invariant measures of the underlying dynamics and addresses the problem that classical ERM framework is not enough to assess generalization in dynamical systems reconstruction - While it is good to see that the Jacobian matching loss improves the reconstruction of invariant measures of the underlying dynamical system, it’s application in practice is fairly limited as Jacobians have to be estimated from data if the ground truth vector field is not known. - I see that the authors want to study the case where training is only perform using “one-step-ahead” predictions (i.e. no unrolling of dynamics), however, I think considering the case of unrolling dynamics as a comparison to the Jacobian matching loss would be helpful as 1) unrolling dynamics is often enough to get a good reconstruction of the invariant measures of the observed attractor and 2) it is easy to apply in practice (no need for extra knowledge of the Jacobians). I.e. I think the more interesting question is whether knowledge of the Jacobians outperforms unrolling of dynamics during training significantly - I think section 4, i.e. dynamic generative models, could be written a tad clearer, i.e. more explicit in terms of what exactly is now assumed to be stochastic in the specific VAE architecture and which part is deterministic. Maybe a Figure demonstrating the architecture setup in favor of moving e.g. proof of Theorem 1 into the Appx. - To me the paper lacks a clear connection of the theoretical findings to the practical real-world setting, i.e. how these insights influence how we should train models to reconstruct real-world dynamics I really like the idea and insights of this paper, but given the Weaknesses above, I do think it lacks actual practical relevance and the novelty is also very much limited to the (seemingly impractical) theoretical findings. Miscellaneous comments: - Figure 2 readability could be improved by removing the legend from within the plots and moving it to the outside (e.g. right, or between the plots), as the legend is the same for both plots - Hyperlink of Table 3 jumps to section 3 - typo l. 192: “shdaowing” [1] [Mikhaeil et al. (NeurIPS, 2022), On the difficulty of learning chaotic dynamics with RNNs](https://proceedings.neurips.cc/paper_files/paper/2022/hash/495e55f361708bedbab5d81f92048dcd-Abstract-Conference.html) [2] [Hess et al. (ICML, 2023), Generalized teacher forcing for learning chaotic dynamics.](https://proceedings.mlr.press/v202/hess23a.html) [3] [Göring et al. (ICML, 2024), Out-of-domain generalization in dynamical systems reconstruction.](https://arxiv.org/abs/2402.18377)
Questions
- how exactly are the empirical measures $\mu$ and $\mu_{NN}$ and their $W_1$ distance computed (for Table 1 e.g.)? - Could the authors include comparisons of the Jacobian matching loss to the more common training method of unrolling the dynamics for at least some time $t$? - How can one draw practical consequences from the theoretical findings?
Rating
5
Confidence
4
Soundness
3
Presentation
2
Contribution
2
Limitations
The authors do mention limitations of their approach (e.g. that the Jacobian matching loss is hard to implement in practice when ground truth is not known), but do not address how their findings might be used in real-world data settings.
Summary
This paper extends generalization results to models trained on dynamical data, especially Neural ODEs. The paper shows and then attempts to explain why Neural ODE trained without a Jacobian matching term fail to capture physical behaviour even when they have low generalization error. Under a generalization assumption, the paper shows that the Jacobian-matching loss can lead to statistically accurate models. The paper shows a number of experiments that show the behaviour stipulated by the theoretical results.
Strengths
The paper is technical but well written and ideas are clearly explained. The mathematical formulation is quite clean and targets an important problem of generalization in ODE models for dynamical data. The explanation of generalization for Neural ODEs is I think novel.
Weaknesses
I think the main weakness is assumption 1. There should be some explanation of the difference between C1 and strong C1 generalization. Furthermore, there should at least be some justification of why this should hold and when we could not expect it to hold. Another possible problem I see is this: The paper explains why minimizing loss 3 (with the jacobian) implies statistical accuracy. I don’t think that it properly explains why minimizing loss 2 does not imply statistical accuracy. I would understand why a standard MLP might not match derivatives. But I have a harder time understanding why a neural ODE with low generalization error isn’t able to do that. Some more intuition about this would be useful.
Questions
Related to the above. In lines 264-265 the paper states that C0 generalization is insufficient for learning shadowing orbits. Is it an obvious fact? Further explanation would be useful. Also useful would be a comparison of a standard MLP minimizing loss 2 and a neural ODE minimizing the same loss.
Rating
6
Confidence
2
Soundness
3
Presentation
3
Contribution
3
Limitations
The paper does not have a section on limitations. I think a discussion of limitations would add to the value of the paper.
Summary
The authors focus on analyzing why MSE loss fails to capture the physical behavior of dynamical systems. Narrowing their analysis to invariant ergodic systems, they provide theoretical insights on when generalization implies statistical accuracy. They propose that for models to be statistically accurate, they must reproduce dynamical invariants, which is achieved by ensuring that the learned model closely follows the true system's orbits. Specifically, the authors provide theoretical justification for why Jacobian matching loss can better capture statistical properties than MSE loss. Empirically, they verify their analysis on the Lorenz 63 system using different architectures, including MLP, Neural ODE, and FNO.
Strengths
The authors provide a thorough theoretical analysis and propose some useful notions to characterize models' ability to reproduce the dynamics statistical measure. Overall, the paper is well-written and easy to follow.
Weaknesses
1. Jacobian loss (Eqn.3) considered in the paper is a special case of the Sobolev norm in [1]. The authors should consider extending their analysis to the Sobolev norm. 2. Recent related works have also discussed the problem of MSE and shown multiple ways to improve upon MSE loss, e.g., the Wasserstein loss used in [2,3] and the theoretical analysis provided in [4]. Although some of these works might focus more on the empirical side, my understanding is that this problem has raised more attention than the authors' claim, and some comparison to the existing work is necessary. [1] Learning Dissipative Dynamics in Chaotic Systems (https://arxiv.org/pdf/2106.06898) [2] DySLIM: Dynamics Stable Learning by Invariant Measure for Chaotic Systems (https://arxiv.org/abs/2402.04467) [3] Training neural operators to preserve invariant measures of chaotic attractors (https://arxiv.org/abs/2306.01187) [4] On the difficulty of learning chaotic dynamics with RNNs (https://arxiv.org/pdf/2110.07238)
Questions
Q1: Figure 1 is unclear, especially when the y-axis is not aligned in the 4&5 columns, and the message shown in Figure 1 seems to contradict the authors' theoretical claims, as the probability distribution of the modeled dynamics is still mismatched when the dynamics of the attract seems to be well learned. Q2: It's confusing when the authors interchangeably used notations $h$ and $F$ for the dynamics map. Is there a particular reason for switching between these notations, and could you elaborate on why? Q3: Line 183: It's not clear why the authors state that implementing Wasserstein distance is difficult when the map $h$ is chaotic. Q4: As Lyapunov exponents are calculated using local dynamics, could you show a comparison using some long-term evaluation metrics?
Rating
7
Confidence
4
Soundness
2
Presentation
3
Contribution
3
Limitations
N/A
Response to your rebuttal
Thank you for your response! Your response regarding the Sobolev norm was valuable, and your further clarifications cleared up my concerns. I appreciate your discussion regarding the related works, which helped me better assess the status of your work. I will increase my score to accept.
Thank you for the feedback
Dear Reviewers, We are extremely grateful to you for spending the time to understand our paper and raise many interesting questions. Due to lack of space, we have not quoted from your review in our responses, but we have answered all your comments to the best of our knowledge under easily identifiable headings. We outline the revisions/edits that we have already made to our paper, and as soon as the option becomes available, we will upload the revised version. For new experiments that we have conducted, we have shared **anonymous links to the figures** since we cannot yet share the revised paper. We request you to visit these links to see the plots for the new experiments on unrolled loss functions. We have indicated in our response wherever there are similarities in the questions/comments by more than one Reviewer. We request you to therefore read the *relevant sections of the rebuttal to other reviewers* as well. We thank you once again for your time and valuable feedback that has improved our paper substantially!
Re: Rebuttal
> Related work on invariant measures Happy to hear that the authors add the references to the related work and thanks for the additional clarification and connection to their own work! > Practicality of Jacobian loss and implications of theory Indeed, the complexity of computing the *model Jacobian* using AD/Finite differences is one part, however, did the authors also consider the added complexity by including this Jacobian in the loss function, which means one has to a backward-over-forward AD problem (-> second order differentiation)? Does runtime/complexity still only scale with problem dimensionality squared? My main concern with the practicality is the estimation of the *Jacobians from time series data*, where **no** ground truth knowledge is available (I appreciate the geophysics examples and I see that in these cases the Jacobian matching loss might improve things). I really miss any experiments on trying to make this loss more applicable in real-world scenarios (e.g. estimating Jacobians from time delay embeddings of real-world data -> do crude approximations already help learning dynamical invariants?). > Unrolling dynamics Thanks for conducting the additional experiments! That zero/negative LEs are still underesimated when unrolling dynamics is indeed a problem I encountered in practice, too. Hence it is great to see that knowledge of the Jacobians can eliminate this problem. However, I checked the paper again and found that the authors report a ground-truth max. LE for the Lorenz63 system with standard settings of ~0.86 (cf. Table 2). However, to my knowledge the Lyapunov spectrum is given by ~[0.905, 0, -14.57]. I checked this with minimal code using Julia (see below). Do the authors mind to share how they estimated the ground-truth Lyapunov spectrum or whether they used specific literature values? ``` using DynamicalSystems, ChaosTools # Lorenz63 ds = Systems.lorenz(ρ=28.0, σ=10.0, β=8 / 3) # Lyapunov spectrum lyap = lyapunovspectrum(ds, 100_000, Ttr=1000) # ~ [0.905, 6e-6, -14.57] ``` > W1 distance Thanks for the clarification! > Edits I see, thanks!
Jacobian and LEs
Thank you very much for reading through the rebuttal! We really appreciate your interest and excellent questions! > Does runtime/complexity still only scale with problem dimensionality squared? This is a very keen observation -- thank you. We meant that computing the Jacobian (using Finite Difference or AD) has quadratic complexity, but you are absolutely right that training time increases also because of differentiating the model at the training points. During training the Jacobian (with respect to the input state) is differentiated with respect to the parameters of the neural network. The complexity with respect to the dimension of the problem still remains quadratic but the training complexity increases just like a physics-informed neural training; the factor of increase therefore also depends on the complexity of the network. Inference time is the same as a vanilla neural ODE however since only the training is modified. Practically we found that KS system (128-dimensional system) was still training in about 10 hours on a single RTX30xx GPU. > estimation of the Jacobians from time series data, where no ground truth knowledge is available This is exactly the scenario we had in mind for the rebuttal. We can estimate the Jacobian from timeseries data using finite difference as long as the data sample the system frequently (there are classical results known for the impossibility of system identification for longer observation frequency and very noisy observations of chaotic systems). As you correctly point out, there is an error made by finite difference (which is on order of epsilon^2, where epsilon is the step size norm). However, even though our theoretical results only use exact Jacobians, the results with unrolling dynamics, e.g., where only indirect Jacobian information is available, suggest that even inexact Jacobians can help us learn the dynamical invariants better. For instance, even Lyapunov vectors and exponents can be estimated with high accuracy using just finite difference estimates of the Jacobian (this is easy to try by replacing the AD with inexact Jacobians in the code snippet below). But as you suggest, we should test this on real world high-dimensional examples (even ones where the attractor is only known through delay embedding), which we defer to a future work. Our present paper emphasizes on understanding why Jacobian information leads to statistical accuracy. Our main contribution is using shadowing theory to provide an explanation and suggest that there is hope of learning chaotic systems with elementary supervised learning methods, and that these can still learn the invariant measure. About the Lyapunov exponents, small differences (statistical error/noise) may arise depending on the integration time (len(traj_gpu) in the code below), time integration method (RK4 in the code below), and subsequently how we estimate Jacobians. Here is the code snippet we wrote for computing 3 LEs: ``` def lyap_exps(dyn_sys_info, traj, iters): model, dim, time_step = dyn_sys_info LE = torch.zeros(dim).to(device) traj_gpu = traj.to(device) f = lambda x: rk4(x, model, time_step) Jac = torch.vmap(torch.func.jacrev(f))(traj_gpu) Q = torch.rand(dim,dim).to(device) eye_cuda = torch.eye(dim).to(device) for i in range(iters): if i > 0 and i % 1000 == 0: print("Iteration: ", i, ", LE[0]: ", LE[0].detach().cpu().numpy()/i/time_step) Q = torch.matmul(Jac[i], Q) Q, R = torch.linalg.qr(Q) LE += torch.log(abs(torch.diag(R))) return LE/iters/time_step ``` Thank you very much for your questions and please feel free to let us know if anything is unclear! Thank you again for your time and help!
Re: Jacobian and LEs
Thanks for the clarification - makes sense that the runtime complexity increases similar to PINN training! Regarding the LEs; First of all, thanks for providing the code. I think this needs more careful evaluation. If you call something "true" LEs, you should indeed compute the right quantity - for the true spectrum, the statistical error or noise should be negligible. I.e. you need to let the trajectories converge into their limit set and draw a long enough trajectory to pin down the correct spectrum (you may plot time steps/integration time vs error in the spectrum to verify you draw long enough orbits to converge). But if the reported "true" LE is indeed ~0.905 instead of ~0.85, then this changes the interpretation of Table 2 (is it just that the models need to be evolved for longer or are they then indeed underestimating the max. LE even with Jac. loss?). To be clear, I really appreciate the effort the authors put into the paper and rebuttal and I think the paper is of great value, but to me it seems it could really benefit from another revision and submission at a later point. I will raise my score to 5 with the assumption that the authors will include all additional experiments in the final version (reevaluate LE spectrum calculation, making sure the true spectrum is reported, estimating Jacobians from data alone, addressing computational complexity).
Re: Jacobian and LEs
Thank you for reading through our comment and for your response! We are so grateful to you for responding promptly and thoughtfully! You are absolutely right. Since LEs are statistical quantities, and for a singular hyperbolic system like the Lorenz system, the Central limit theorem holds, we do see a 1/sqrt{N} convergence of ergodic averages and LEs computed along trajectories of N length. The "true" label does refer to long averaging times conducted with the ground truth system -- i.e., the true vector field. It is indeed a fair comparison to maintain an equally long time averaging but compute the LEs with the Neural ODE as a vector field instead. These are what are referred to as "JAC" or "MSE" values -- the LEs computed with these learned models. A 5% difference in LEs can happen even in between runs with different initial conditions and this should, for instance, not be considered as an indication of lack of ergodicity. Overall, we want to clarify that the LEs are expectations of random variables and therefore cannot be estimated like deterministic quantities. It would not be meaningful to report them (in a chaotic system) with many digits of precision, e.g. Please let us know if any further clarification is needed!
Request to read rebuttals
Dear Reviewers, Since the discussion period is set to close soon, we request you to please read our rebuttals and ask for clarification. If we have answered all your questions satisfactorily, we request you to please re-evaluate our paper.
Thank you for the clarifications. From the other reviews, I see there are some concerns regarding practicality of obtaining the Jacobian. However, I think the explanation of generalization also has value. I maintain my accept recommendation.
Thank you!
Thank you very much -- we especially appreciate your questions! The purpose of this paper, as you indeed correctly point out, is dynamics-aware generalization bounds. We show that learning the Jacobian leads to statistical accuracy and *explain why* by deriving a generalization bound via shadowing theory. Computing the Jacobian is actually not impractical compared to more sophisticated training approaches using recurrent architectures or generative modeling of the SRB measure (physical measure in the paper). We appreciate your time and effort! Please let us know if we can provide any further clarification!
Thank you and request for response!
Dear Reviewers, Thank you for reading our rebuttals, seeking further clarification and re-evaluating our paper! We are extremely grateful for all the time you have spent on improving our work. If you have not done so already, could you please read our rebuttal, and re-assess our work in light of the clarifications in the rebuttal? Sorry for the multiple reminders -- we are almost at the end of the discussion period. Thank you again!
Thank you for the comments and discussion on learning chaotic dynamics
Dear Reviewers and ACs, We have truly grateful for all the time you have spent in giving us a thorough review. The subsequent discussions (including the references) have been very fruitful and led to substantial improvements in the revision. We are excited about both directions our work contributes to and is influenced by: practical methods for learning physical measures (ensemble/long-term behavior) of chaotic dynamics using supervised learning and introducing dynamical systems/ergodic theory to analyze and understand generalization. Overall, the intersection of nonlinear dynamics with machine learning theory and practice is growing and we are very happy for the opportunity to be able to improve our work at this intersection through these discussions with you.
Decision
Accept (poster)