Adaptive higher order reversible integrators for memory efficient deep learning

Dear reviewer, we are indebted for your careful reading and helpful comments. We would like to address your questions and concerns. 1. *Where the savings of the computation time come from? Is it because adaptive $Y^4$ in general uses larger step sizes? Is there a limit on $k$ for adaptive $Y^{2k}$ to outperform adaptive ALF on computation time?* In the considered examples, the 4th order integrator $Y^4$ and the (effectively) second order integrator ALF with adaptive time-stepping are applied to differential equations that resemble the planar Kepler problem and a nonlinear coupled Duffing harmonic oscillator with 10 masses. These are integrated forward (with adaptive step size control) and the adjoint equations are integrated backwards with the step sizes from the forward integration repeatedly. Indeed $Y^4$ can meet the local error tolerance with fewer steps than ALF as it is of higher order. This causes computational savings for the forward and the backward pass. This is because in this regime of local error tolerance and on the differential equation at hand, a fourth order method turns out to be more efficient. Indeed, higher order methods can be considered. Eventually, we expect a limit in efficiency gain just as in classical integration methods where there exists an optimal order. This is because if $c_p h^{p+1}$ is a bound for the local error of a numerical integrator that holds for all step sizes $0<h<h_{\mathrm {max}}$ then the error constant $c_p$ tends to rise with the order $p$ of a method, as is known in classical numerical integration theory. Indeed, these error constants grow when higher order methods $Y^{2k}$ are constructed as Yoshida compositions. 2. *Why the final loss of Y4 is lower? What if we increase the number of epochs? Will it help adaptive ALF to achieve lower loss?* We see that the prominent display of Figure 3 might be misleading as it is irrelevant for the hypothesis of the article. It has, therefore, been moved to the appendix (see Figure 7). Indeed, we do not see any reason to expect Y4 to be better trainable than ALF-based networks *per epoch*. The article's hypothesis is that computational advantages arise because an epoch can be computed faster thanks to higher order integration. Moreover, in our newly introduced example (see C.3 Discretized wave equation) the training accuracy after a sufficiently large training time levels out at roughly the same value. 3. *The author may consider adding more contexts for the "adaptive stepping" presented in Hairer et al. (2006) in the Background section.* We have added a paragraph in the introduction, which explains the concept of adaptive step size selection and provides references. 4. *The authors may consider adding known reversible methods other than asynchronous leapfrog and Heun* We have added this classical reference, on which the ALF method is based. 5. *for composition of mixed-order methods, although it is less developed than the theory for composing same-order integrators, there are some frameworks addressing this topic, especially in contexts like symplectic integrators and splitting methods* We have noticed that our terminology *mixed-order* could be mistaken for compositions of integrators of mixed orders, while in the context of the article *mixed-order* refers to the fact that a single step of ALF in its augmented phase space yields a prediction of the state variable to second order but a prediction of a velocity variable only to first order in the step-size. We have adapted our formulations and avoid the term *mixed-order*. Moreover, we have included the suggested excellent references. 6. *some other works on deep learning of dynamical systems* The suggested references are now included in the introduction. 7. *I don't quite understand the purpose of section 4.2. "LEARNING OF DYNAMICAL SYSTEMS PARAMETERIZED BY NEURAL NETWORK". Could the authors comment on the novel insights or advantages gained from this section compared to 4.1, or to clarify how this section extends or complements the results from 4.1?* Indeed, the additional experiment supports the same messages as the experiments of section 4.1. It is included to demonstrate that our method can be applied to problems with many parameters. Moreover, an additional experiment was included (discretized wave equation) to show the applicability to high-dimensional dynamical systems.

Dear reviewer, we highly appreciate your helpful comments and would like to address your questions and concerns. 1. *The validation falls short. The experiments are limited to simple datasets, which are insufficient to demonstrate the general applicability of the proposed architecture. In what settings is this architecture genuinely effective? For instance, while reversible dynamical systems, like the Kepler problem used in the study, are a clear application, this alone is somewhat narrow, as one could achieve similar outcomes by applying reversible integrators within Neural ODE frameworks. On the other hand, could applying this architecture to general tasks, like image processing, compromise performance due to the reversibility constraint? The study would benefit from a clarified scope of applicability and a more comprehensive set of experiments.* We have clarified the scope of the method in the introduction, which we mainly see in the identification of high-dimensional dynamical systems from irregular time-series data, for which gradient-based optimization methods can cause high memory requirements. This is now reflected better in the conclusion as well. To show the applicability of our method to such systems, we have included the example of learning a dynamical system which is governed by a spatially discretized partial differential equation. This is a highly active research area and references to alternative approaches based on model-order reduction are provided. Indeed, an application of high order reversible Neural ODE architectures to image processing tasks and normalizing flows is an exciting field for exploration in future works. 2. *While the theoretical memory efficiency is convincing, there is no experimental validation. Although ALF is presumably the baseline, only computation time is compared. An experimental comparison of memory consumption is also necessary, along with theoretical orders for computational and memory costs.* The memory requirements of the methods discussed in the article are now compared in Table 1. Experimental validation of lower memory costs for neural ODEs that are based on reversible methods in comparison to neural ODEs based on non-reversible methods that require storage of computational graphs has been provided in [1] 3. *Considering ICLR's focus on machine learning, this paper falls short of expected standards due to the ambiguity of its contributions and insufficient empirical support.* In the revised version, we state the contribution to this field more clearly. Moreover, the new example of learning a discretized partial differential equation constitute additional empirical support. The scope of the paper, learning of models of dynamical systems from data, is a highly active topic in machine learning and contributions are regularly presented at machine learning conferences. Examples include HNN [2] (NeurIPs 2019), LNN [3] (ICLR), MALI [1] (ICLR2021), or [4,5]({AAAI} 2020/2022 Symposium on Physics Guided AI). 4. *References are cited without parentheses, which hinders readability. Proper citation using citep tags with parentheses is recommended.* The citation has been improved. 5. *The composition method could construct more complex architectures, and such an exploration could add value to the study.* Indeed, it is an exciting avenue to explore higher order compositions or compositions which have at a given order an optimal error constant. Additionally, symmetry properties can be explored in combination with our reversible, time-adaptive approach. In order not to diffuse the main messages of the contribution and align with the conference's focus, we have decided not to include content of this context. **References** [1] Juntang Zhuang, Nicha C. Dvornek, Sekhar Tatikonda, and James S. Duncan. Mali: A memory efficient and reverse accurate integrator for neural odes, 2021 [2] Samuel Greydanus, Misko Dzamba, and Jason Yosinski. Hamiltonian Neural Networks. Advances in Neural Information Processing Systems, 2019. [3] Miles Cranmer, Sam Greydanus, Stephan Hoyer, Peter Battaglia, David Spergel, and Shirley Ho. Lagrangian neural networks, 2020. [4] Christine Allen-Blanchette, Sushant Veer, Anirudha Majumdar, and Naomi Ehrich Leonard. LagNetViP: A Lagrangian neural network for video prediction (AAAI 2020 symposium on physics guided ai), 2020. [5] Justice Mason, Christine Allen-Blanchette, Nicholas Zolman, Elizabeth Davison, and Naomi Leonard. Learning interpretable dynamics from images of a freely rotating 3d rigid body, 2022.

Dear reviewer, we would like to thank you for the helpful comments and to address your concerns. 1. **Paper doesn't provide significant theoretical results, the experimental results are weak.** We have developed the first, in its kind, class of higher order time reversible integration methods allowing adaptive time stepping. The adaptive stepping in turn permits to considerably accelerate computations. Notice that such methods in combination with the neural ODE approach are especially important in learning models of high dimensional systems exhibiting complex behaviour, such as molecular systems. Using checkpointing would lead to memory explosion in this case, whereas higher order reversible methods would allow to solve the learning problem with low computational costs. With the numerical examples we show that higher order reversible method outperform the known lower order method in terms of the computational time. Notice that the nonlinear oscillations example can be seen as a simplified model of molecular dynamics. 2. **Literature review and citations are inadequate.** We have extended the literature review. 3. **Never compared to alternative methods such as a non-reversible integrator** We have included in Appendix C the comparison with the Runge-Kutta 4(5) (RK45) method. On the Kepler example (Fig2 and Fig3), Y4 outperforms RK45 . On the other examples (Fig5 and Fig9) RK45 is faster but has a higher memory cost. This highlights the applicability domain of the proposed class of compositional methods. They permit to make computations faster, when the higher order of the integration is required and at the same time they have constant memory costs, which is not possible with the standard methods such as RK45. 4. **For parameter estimation tasks, consider accuracy with regard to the reference parameters** We provided additional results showing that the higher order method converges to parameters of high order of accuracy faster than the lower order method. This is displayed in Fig2, Fig3, Fig5 5. **A fixed step size was used in table 2** Table 2 shows the results for the adaptive time stepping. 6. **Fig2 exaggerates the benefits of Y4 over ALF** The figure was moved to the appendix. It illustrates the need for smaller time-steps in the integration for highly accurate parameter estimation. Additional discussion and clarification was added and can now be found in the Appendix C1 7. **The problems are fairly easy.** We have included a new example of discretized PDE. The obtained results displayed in Table 4 show that the higher order method is faster in getting to the predefined training accuracy. In additional, Fig9-10 support the observation. 8. **Questions on computation of $d\phi_{j+1}/d\phi_{j}$ in backprop** In Appendix B we have included an alternative approach for the integration of the adjoint, which does not require computation of $\frac{\partial \phi_{k+1}}{ \partial \phi_k}$ for the backpropagation. The method is based on the composition of rescaled steps backward of ALF method. As a result, computation of $\frac{\partial \phi_{k+1}}{ \partial \phi_k}$ is replaced by the composition of the terms depending of $\frac{\partial f(z,t)}{ \partial z}$. Notice that the Jacobian of $f$ is sparse in the case of discretized PDEs and the memory can be allocated more efficiently. 9. **Checkpointing.** We have included a paragraph on the checkpoitning in the end of the description of the gradient computations in the end of Section 2.2. In this paragraph we explain the difference between the approach by checkpointing and by the reversible methods. Checkpointing is highly efficient in combination with high order methods for the tasks, which require high accuracy in learning and use short trajectories in training. However, in case of systems with complex behaviour it is important to consider long time integration which leads to the linear growth in memory with respect to the time of integration. This is the context in which the reversible methods are required. 10. **Simulation-based inference.** We have added a discussion on statistical inference to Section 4.1 and related simulation-based inference to the proposed method. 11. **The notation in 2.1 is confused.** The section was shortened and the notations were improved. 12. **Proofreading** We have corrected mistakes. 13. **Are step sizes treated as fixed on the backward pass?** The time steps are computed in the propagation forward and saved in the memory. In the propagation backward the steps are treated as fixed, see the paragraph on adaptive stepping in Section 3 and Algorithm 2. 14. **Line 505** The neural network approximates the gradient of the potential, so to obtain the potential we need to integrate the function given by the neural network and we need to do it numerically, that is, to use a numerical integration methods. 15. **Heun's method** We removed the sentence on the initialization and only left the description of the important properties.