Paper: How Temporal Unrolling Supports Neural Physics Simulators (ICLR-2024)

Reviewer comment (Writing):
The NOG + MIG could be explained more clearly, specifically giving a longer explanation of how the gradients are computed and defining what the authors mean by “loss landscape changes” and “data distribution shifts.”

Author response:
Thank you for the comment. We have now added a theoretical section in the general response that formally introduces the data shift concept, and should help clarify your concerns. Furthermore, appendix A has a detailed explanation of gradient calculation for NOG and WIG setups. The theoretical analysis then expands on these and studies the loss-gradient mismatch in NOG unrolling.

The revised paper will also motivate the NOG setup more clearly. Crucially, no-gradient unrolling could be achieved by interfacing existing numerical code with a machine-learning library. It is a viable option and first step for scientific domains that are dominated by known, complex, and validated code bases, like computational fluid dynamics.

Reviewer comment (Evaluation):
Based on the provided experimental results would the authors conclude that the scale of the model is more important in determining good performance than the specific error scheme?

Author response:
Network size does indeed significantly influence accuracy. Our results actually provide insights beyond this simple fact: We provide convergence rates of the inference error with respect to the parameter count for machine-learned correction and prediction methods. For the physical systems, numerical solvers, and network architectures studied, these convergence rates are small in comparison to purely numerical approaches. A grid refinement of the underlying solver might thus yield better results than an extension of the neural network parameter space for the same additional computational burden. Based on this observation, we conclude that slimmer architectures are of greater interest to the scientific computing community.
We believe that this is an important and novel consideration that goes beyond previous studies.

Reviewer comment (Evaluation):
How does the compute of the best performing deep learning models compare to direct solvers?

Author response:
Thank you for this comment, it touches upon one central insight of our work. We have established convergence rates for network architectures and seen how their scaling motivates slimmer sizes over overparametrized variants. This is also true in our tests, where large models require substantially more operations than numerical solvers, which translates to the computing time. We also include convergence studies for our datasets and network architectures. A rate of $n^{⅓}$ can be fit to the network's convergence rate, where n represents the number of network parameters. In practice, this means that a doubling in the computational complexity only yields an accuracy improvement factor of 1.26. This convergence rate is poor in comparison to numerical approaches, where rates larger than $n^2$ are the standard (i.e. improvement of 4x for a doubling in numerical gridpoints). This observation is crucial to the scientific machine-learning community and has often been overlooked in previous studies.

At the same time, the goal of our work is not to propose a specific architecture that outperforms classical approaches. Rather, we compare strategies that ultimately help train such models as well as extract and study the relevant quantities in these training strategies.

Reviewer comment (Novelty):
Many of the claims in the manuscript are fairly trivial or shown in previous works. (1) Large networks yield better results. (2) Learning on trajectory is better (a large family of works on Lagrangian and Hamiltonian Neural Networks have already implemented this and demonstrated it. For example, see https://proceedings.neurips.cc/paper/2020/hash/9f655cc8884fda7ad6d8a6fb15cc001e-Abstract.html ) where they have also used a horizon of 4 timesteps while training the systems, (3) differentiable solvers are better for training ML models of dynamical systems (there is a family of works on JAX and specifically JAX-MD which demonstrates this extensively).

Author response:
The observation that larger networks yield better results might seem trivial. However, our results allow insight beyond that fact: We provide convergence rates of the inference error with respect to the parameter count for machine-learned correction and prediction methods. For the physical systems, numerical solvers, and network architectures studied, these convergence rates are small in comparison to numerical approaches. A grid refinement of the underlying solver might thus yield better results than an extension of the neural network parameter space for the same additional computational burden. Based on this observation, we conclude that slimmer architectures are of greater interest to the scientific computing community.
We believe that this is an important and novel consideration that goes beyond previous studies. 

We’d also like to point out that, unfortunately, many current submissions in the field (and also in this year’s ICLR) do not respect the seemingly obvious fact above: they compare models with strongly differing parameter counts in a single graph or table. We hope that stating and illustrating this effect clearly raises awareness in the community.

Differentiable trajectory unrolling in itself is of course not new, and various papers have demonstrated the benefits of this approach. However, differentiability is generally not given in codebases used for scientific computing. These codebases are the product of long development and validation processes, where establishing differentiability would be a tremendous effort. The no-gradient setup provides a very attractive strategy that would only require an interface between a numerical simulator and the machine learning augmentation. Reliable statistical evaluations are necessary to estimate the efficacy of these training modalities and to steer future developments in scientific computing.

Reviewer comment (Experiments):
The experiments are performed on an extremely limited number of datasets and architectures. For a benchmarking and dataset paper, more extensive studies are expected, especially in a dense area where a large number of models and datasets are available. The architectures used are not state of the art (see questions for details).

Author response:
In contrast to many previous studies, our work focuses on correction tasks and deploys a comparison of differentiable and non-differentiable training approaches. As such, a differentiable solver is at the base of each of our 3 main datasets. This stands in stark contrast to other studies (e.g. Stachenfeld et al. 2022) which are limited to predictive setups. For these, mostly existing or commercial solvers were used. In contrast, our work combines popular dataset choices in machine learning on differentiable solvers (e.g. Um et al. 2019, Kochkov et al. 2021), and thus constitutes a valuable asset for training and testing architectures based on these differentiable solvers. 

We trained a broad range of architectures that cover more than 3 orders of magnitude in network size, resulting in over 2000 trained models for this study. As further elaborated upon in the paper and our answer to W4, this capacity dimension of neural networks is often overlooked, but of critical importance to scientific computing. The range of networks covers common choices in prominent previous work. For instance, the 1.0M parameter graph-network aligns with Branstetter et al. (2022), while the convolutional network with 250k parameters from Kochkov et al. (2021) is also contained in our capacity range. Similarly, the convolutional network sizes (60k and 250k) studied by Um et al. (2020) fall within the same range. This is also true for the convolutional ResNet architectures in Stachenfeld et al. (2022) with 200k parameters for 1D problems and 580k parameters in 2D.

Reviewer comment (Evaluation):
It is mentioned that across architectures and network sizes, WIG has the lowest error. From the results, it seems to not be clear statistically as the error bars are mostly overlapping. A paired t-test or so may be conducted to establish the statistical significance of the result.

Author response:
We thank the reviewer for the suggestion. To improve the interpretability of our results, a series of significance tests were performed on our trained models. We used Welch’s t-test with one-sided p-value calculation and compared NOG and WIG distributions to the ONE baseline. The resulting p-values are reported in the updated appendix of the paper. The differences in the studied training modalities are statistically significant for smaller network setups. As model sizes increase, the distribution of trained models becomes more similar. Our heavily overparameterized networks (e.g. 1.0M parameters for 48 degree of freedom KS system) all show highly accurate and stable predictions (see divergence time metric, Figure 20 Appendix C). Based on the theoretical considerations added in the general comment, this means that the attractor of the learned dynamics is closely aligned with the ground truth system. These setups thus display less data shift and the benefits of unrolling are reduced. When these conditions in the overparameterized regime are met, NOG models are at a disadvantage due to their mismatch between gradients and loss landscape. As discussed in our answer to W4, the heavily oversized architectures are not practically relevant for scientific computing due to their weak scaling compared to numerical approaches. This evaluation confirms that for relevant, small to medium sized networks our results are statistically significant and hence that conclusions can be drawn.

Reviewer comment (Experiments):
There are two approaches for predicting the long-term dynamics of systems known as Deep Koopman operators, and Fourier Neural Operators. Both of them have shown superior performance in predicting the long-term dynamics of physical and chaotic systems. Especially, FNO with vision transformers has shown promise in predicting extremely complex systems, such as weather forecast at a planetary scale. For benchmarking work, it is important to consider the SOTA models. Moreover, such models may raise additional questions, as outlined in the next question.

Author response:
Of course, there are exciting developments on the architecture side that could ultimately improve the accuracy of long-term correction and prediction networks. However, the underlying architecture of these networks is not at the core of our work, and we also do not present these architectures necessarily as the best-performing choice. The chosen architectures are, however, a fundamental building block of more recent developments and still serve as a basis of comparison in related work. The goal of our study is to establish a guide for designing training procedures that yield the best results for a given network. Given the scale of our empirical study across architectures, solvers, physical systems, and network sizes, as well as the added theoretical analysis, we are confident that our results translate to other architectures. To address your skepticism on this issue, we have trained an attention-based U-Net on the KS system. Similar to FNOs, the attention mechanism in this architecture is designed to make the network learn an effective mapping based on both global as well as local feature structures. The test loss values are added to Appendix D of our paper. The results support our conclusions and align well with previously tested architectures, from which we conclude that the attention mechanism does not fundamentally change the behavior of unrolling.

For Koopman operators, we are skeptical whether a separate Koopman embedding would really be beneficial in a corrective setup, which fundamentally requires communication between the network and a numerical solver at each timestep in real space. While the integration of this architecture in correction setups would surely be interesting, we'd prefer to leave this development for future work.

Reviewer comment (Experiments):
In the context of neural operators and transformer architectures (which employ all pair attention), do the conclusions made in the paper still hold true? This is important as the SOTA models rely on these architectures, and commenting on the applicability of these approaches toward the newer architectures is important.

Author response:
Of course, there are exciting developments on the architecture side that could ultimately improve the accuracy of long-term correction and prediction networks. However, the underlying architecture of these networks is not at the core of our work, and we also do not present these architectures necessarily as the best-performing choice. The chosen architectures are, however, a fundamental building block of more recent developments and still serve as a basis of comparison in related work. The goal of our study is to establish a guide for designing training procedures that yield the best results for a given network. Given the scale of our empirical study across architectures, solvers, physical systems, and network sizes, as well as the added theoretical analysis, we are confident that our results translate to other architectures. To address your skepticism on this issue, we have trained an attention-based U-Net on the KS system. Similar to FNOs, the attention mechanism in this architecture is designed to make the network learn an effective mapping based on both global as well as local feature structures. The test loss values are added to Appendix D of our paper. The results support our conclusions and align well with previously tested architectures, from which we conclude that the attention mechanism does not fundamentally change the behavior of unrolling.

For Koopman operators, we are skeptical whether a separate Koopman embedding would really be beneficial in a corrective setup, which fundamentally requires communication between the network and a numerical solver at each timestep in real space. While the integration of this architecture in correction setups would surely be interesting, we'd prefer to leave this development for future work.