Accelerating Neural Differential Equations for Irregularly-Sampled Dynamical Systems Using Variational Formulation

Authors: hongjue  zhao
yuchen  wang
hairong  qi
jiajia  li
sha  lui
han  zhao
huajie  shao
Venue: ICLR_cc_2024_Workshop_AI4DiffEqtnsInSci
Published: 2024-02-10 14:53:01

Abstract: Neural ODEs have exhibited remarkable capabilities in continuously modeling dynamical systems from observational data. However, existing training methods, often based on adaptive-step-size numerical ODE solvers, are time-consuming and may introduce additional errors. Despite recent attempts to address these issues, existing methods rely heavily on numerical ODE solvers and lack efficient solutions. In this work, we propose the Fast-VF Neural ODE, a novel approach based on variational formulation (VF) to accelerate the training of Neural ODEs for dynamical systems. To further mitigate the influence of oscillatory terms in the VF loss, we incorporate the Filon's method into our design. Extensive experimental results show that our method can accelerate the training of Neural ODEs by 10 $\times$ to 100 $\times$ compared to the baselines while achieving comparable accuracy in irregularly-sampled dynamical systems.

Main paper: 
 introduction: 
 Neural ordinary differential equations (Neural ODEs) (Chen et al., 2018) represent a family of continuous-depth machine learning models. Inspired by ResNets (He et al., 2016) , Neural ODEs aim to parametrize vector fields of ODEs using neural networks,EQUATIONwhere f θ : [0, T ]×R d → R d is a neural network and θ are model parameters. The continuous nature and special inductive bias of Neural ODEs make them pretty suitable to model dynamical systems from irregularly-sampled time series data. (Chen et al., 2018; Rubanova et al., 2019; Kidger et al., 2020) . In existing training frameworks, numerical ODE solvers play a vital role. The forward pass results are directly calculated using numerical ODE solvers. In the backward pass, there are two major methods commonly employed to backpropagate through ODE solvers (Kidger, 2022; Onken & Ruthotto, 2020) : (1) discretize-then-optimize, which involves directly backpropagating through operations of ODE solvers, and (2) optimize-then-discretize, also known as the adjoint sensitivity method, as utilized in (Chen et al., 2018) , which introduces additional adjoint ODEs. Details of these training methods are presented in Appendix A.However, ODE solver-based training methods for Neural ODEs face two main limitations. First, they are inherently time-consuming. The internal mechanisms within numerical ODEs can incur significant computational costs in solving Neural ODEs. (Lehtimäki et al., 2022) . This is attributed to the numerous evaluations of neural networks beyond desired time points. For the optimizethen-discretize approach, the case may worsen due to the additionally introduced adjoint ODEs. Secondly, existing approaches may suffer from low accuracy for two primary reasons. The autoregressive nature of most numerical ODE solvers can lead to error accumulation. On the other hand, the optimize-then-discretize approach incurs additional numerical discretization error, resulting in inaccurate gradients, and potentially causing the training process to fail entirely (Gholami et al., 2019; Kidger, 2022) .Thus far, various approaches have been proposed to address the aforementioned limitations of Neural ODEs. However, these methods have not fully resolved these problems. To alleviate the computational bottleneck of Neural ODEs, some works attempt to constrain the complexity of learned dynamics, simplifying tasks for ODE solvers (Finlay et al., 2020; Kelly et al., 2020; Pal et al., 2021; Lehtimäki et al., 2022) . Nevertheless, these approaches are unsuitable for time series tasks, where capturing the underlying dynamics is crucial. Moreover, they only accelerate the inference process of trained Neural ODEs, but the training process still remains time-consuming. Some studies also try to modify existing training methods directly (Daulbaev et al., 2020; Kidger et al., 2021; Djeumou et al., 2022; Norcliffe & Deisenroth, 2023) . However, these approaches are still heavily dependent on the ODE solvers and the optimize-then-discretize technique, thus cannot effectively address the computational bottleneck. In certain cases, they may be slower than the discretize-then-optimize approach. Additionally, simulation-free training methods have emerged for continuous normalizing flows (Lipman et al., 2022; Ben-Hamu et al., 2022; Rozen et al., 2021) . Despite their promise, these methods are not applicable to dynamical systems. In the domain of irregularly-sampled time series tasks, several models built upon neural differential equations have been proposed (Chen et al., 2018; Rubanova et al., 2019; Kidger et al., 2020) . However, in essence, numerical ODE solvers are still the key components of these models, which potentially result in inaccuracies for long-term time series prediction.To address these challenges, we strive to mitigate the usage of numerical ODE solvers during the training of Neural ODEs. To this end, we propose the Fast-VF Neural ODE, a novel approach for speeding up the training of Neural ODEs using the variational formulation (VF) (Hackbusch, 2017) . Specifically, a loss function based on the VF (Qian et al., 2022) is employed in our proposed method. This VF loss only requires numerical integrations, thus neural networks are only needed to be evaluated on desired time points. Furthermore, we incorporate Filon's method (Deaño et al., 2017) into the loss function to address potential oscillatory integral issues effectively. We conduct extensive experiments to evaluate our method on various irregularly-sampled dynamical systems. Evaluation results show that our approach significantly outperforms existing baselines, achieving a 10× to 100× speed increase. In addition, it maintains higher or comparable accuracy levels.In summary, our contributions include: (1) we developed Fast-VF Neural ODE, a novel training method using VF for accelerating Neural ODEs for dynamical systems;(2) we incorporate Filon's method to alleviate the influence of oscillatory terms in the VF loss for further enhancing model performance, (3) our approach achieves a significant acceleration, outperforming other baselines by 10 to 100 times in irregularly-sampled dynamical systems while achieving competitive accuracy. 
method: 
 We first reviewed the basic concepts of VF and the Filon's method. Then, we elaborate on our proposed training framework designed to expedite Neural ODEs using these two techniques.First of all, we review the VF of ODEs and Filon's method that will be used in our work.Variational Formulation (VF) (Qian et al., 2022; Hackbusch, 2017) . The variational formulation presented in Theorem 1 establishes a direct connection between the trajectory x and the vector field f . This method enables Neural ODEs to learn parameters without relying on ODE solvers.Theorem 1. Consider d ∈ N + , T ∈ R + , continuous functions x : [0, T ] → R d , f : [0, T ] × R d → R d , and g ∈ C 1 [0, T ],where C 1 is the set of continuously differentiable functions. Here we define the functionalsEQUATION)Then x is the solution to the system ODEs ẋ = f (t, x) if and only ifEQUATIONThe proof of Theorem 1 can be found in (Qian et al., 2022 ).Filon's method for oscillatory integration. Consider the integral b a h(t) sin(ωt)dt and suppose that only tabular data of h(t) is available. As ω increases, the integrand will become oscillatory. In such case, it fails to perform Lagrange polynomial interpolation (Press, 2007) on the whole integrand using the general numerical integration techniques. To address this issue, Filon's method only performs Lagrange interpolation on h(t), and then we can approximate h(t) using a polynomial p(t) with closed-form. Subsequently, the integral b a p(t) sin(ωt)dt can be directly computed to approximate the original integral. In this context, the polynomial is not required to approximate the oscillatory part in the integrand, thus effectively addressing the main problem that general numerical integration schemes encounter in oscillatory integration tasks. In real-world applications, Filon's method can be enhanced by dividing the integration interval into several segments, commonly known as the composite version of Filon's method.We detail our proposed Fast-VF Neural ODE. Given trajectories {x [i] (t)} N i=1 of the given dynamical system ẋ = f (t, x). For the Neural ODE in Eq. (1), we aim to solve the following optimization problem to identify the optimal model parameters. According to (Qian et al., 2022) , we have:EQUATIONwhere C j is defined in Equation (2). Since Theorem 1 specifies an infinite number of constraints, rendering it impractical for implementation, we employ {g s } S s=1 to compute the loss function following (Qian et al., 2022) . This set represents a collection of Hilbert orthonormal basis for L 2 [0, T ], satisfying ∀s ∈ {1, . . . , S}, g s (0) = g s (T ) = 0 and g s ∈ C 1 [0, T ]. In this work, we set g s (t) = 2/T sin(sπt/T ), which can be viewed as the Fourier basis. However, the Fourier basis introduces oscillatory terms into the VF loss. As s increases, general numerical integration techniques will not work well, but this issue has not been addressed in (Qian et al., 2022) . In our approach, we calculate C j numerically using the composite Filon's method, as described in Section 2.1, in which a second-degree polynomial is interpolated on each segment.Next, we try to solve the optimization problem in Eq. (4). But the issue is that the functional C j in the VF loss is an integral from 0 to T . However, in irregularly-sampled dynamical systems tasks, sampling time points may not begin at 0, for example, [t 0 , t N ]. In this scenario, for autonomous systems, we can map the time interval [t 0 , t N ] to [0, T ] by defining T = t N -t 0 . Regarding nonautonomous systems, we can always transform them into autonomous systems by including t as a new variable (Brunton et al., 2016) . These adaptations allow our method to be applied effectively in irregularly-sampled dynamical systems. Now we try to demonstrate how we achieve speed enhancement in our method, which comes from two aspects: (1) a reduction in the number of function evaluations (NFEs), and (2) improved parallelizability. Consider a trajectory sampled K points. In our approach, the VF loss is calculated by evaluating the vector field exactly K times in parallel to compute numerical integrals. In contrast, traditional ODE solver-based training methods often require evaluating the vector field more than K times step by step, as they typically need to assess the vector field at additional time points due to step size settings in numerical ODE solvers to maintain accuracy. Moreover, to perform one-step forward predictions in ODE solvers, the vector field must be evaluated multiple times in high-order ODE solvers such as Dopri5, a common choice in practice.In summary, our method enables the learning of parameters in Neural ODEs without relying on numerical ODE solvers. Regarding computational efficiency, our approach requires only numerical integration for each trajectory. In these numerical integrals, vector fields of Neural ODEs are only evaluated at given time points in the dataset in parallel, significantly reducing the number of function evaluations and achieving better parallelizability compared with existing training methods. Additionally, our loss function effectively mitigates error accumulation from auto-regression. 
experiments: 
 We conduct extensive experiments to evaluate the proposed Fast-VF Neural ODE using irregularlysampled data collected from various dynamical systems.Datasets. We choose three dynamical systems pertinent to various fields such as biology, biochemistry, and genetics. These include the Gompertz model (Gompertz, 1825) ; the glycolytic oscillator (Sel'Kov, 1968) ; and the genetic toggle switch (Gardner et al., 2000) . For each system, we first generate 125 trajectories with randomly sampled initial points. Among these trajectories, 100 trajectories are allocated for training, while the remaining 25 trajectories are used for validation. For these trajectories, we first sample them regularly at 0, ∆t, 2∆t, . . . , T . Subsequently, we randomly select points from the sampled points based on a specific ratio r. Additionally, 25 trajectories sampled at 0, . . . , 2T are generated for testing. Notably, the points at 0, ∆t, . . . , T are utilized for interpolation tasks, while the points at T, T +∆t, . . . , 2T are reserved for extrapolation tasks. In our experiments, T is set to 10, and ∆t is set to 0.1. In addition, we choose r = 0.8 for our experiments. Details for datasets can be found in Appendix B.2.Baselines. To evaluate the acceleration performance, we compare our method with the following training approaches: (1) the discretize-then-optimize approach (Dis-Opt), (2) the optimize-thendiscretize approach (Opt-Dis) (Chen et al., 2018) , and (3) the seminorm approach (Kidger et al., 2021) . Additionally, we also compare our proposed method with following irregularly-sampled dynamical system tasks: (1) Vanilla Neural ODE (Chen et al., 2018) , (2) Latent ODE with RNN encoder (Chen et al., 2018) , (3) ODE-RNN (Rubanova et al., 2019) , (4) Latent ODE with ODE-RNN encoder (Rubanova et al., 2019) , (5) Neural CDEs (Kidger et al., 2020) . Detailed training settings are presented in Appendix B.Evaluation Results. We evaluate the performance of our method on interpolation and extrapolation tasks respectively. We evaluate the performance of different methods using the Mean Absolute Percentage Error MAPE = 1N N i=1 xi-xi xi. To evaluate the acceleration performance of our method, we compute the training time according to prior works (Djeumou et al., 2022; Norcliffe & Deisenroth, 2023) .First, we compare the acceleration performance of our method against the baselines by considering an irregularly-sampling ratio of r = 0.8. As shown in Fig. 1 , we can observe that our approach can speed up the training of Neural ODE by 10 × to 100 × compared to the baselines. In addition, we assess the mean absolute percent error across three dynamical systems for different methods, as illustrated in Table 1 . The results indicate that the proposed approach can achieve competitive performance yet faster training speed than the best baseline ODE-RNN. Note that the good performance of ODE-RNN is attributed to its intricate network structure. In conclusion, our method can markedly accelerate model training while simultaneously maintaining high accuracy.Figure 1 : Training time (minute) for each method. We can see that our method can achieve 10× to 100× faster than the baselines. Note that Neural ODE denotes the Vanilla Neural ODE, Latent ODE 1 denotes the Latent ODE with RNN ecnoder, and Latent ODE 2 denotes the Latent ODE with ODE-RNN encoder. Limitations of our method. As our model relies on Filon's method, which uses polynomial interpolation for sampled trajectories, there are instances where polynomials may not effectively approximate complex trajectories. In such cases, the method's ability to capture intricate system dynamics might be compromised. In addition, extreme measurement settings for dynamical systems still remain as open challenges, e.g., extremely sparse noisy observations.In our experiments, each neural differential equation model is trained for 5000 epochs using the Dopri5 solver. We employ the Adam optimizer (Kingma & Ba, 2014) with an initial learning rate of 0.001. The learning rate is scheduled using cosine onecycle schedule. All training data is loaded in one epoch. Data generation involves using the Dopri5 solver in Diffrax for numerical ODE solutions. In the VF loss, we set the S in Equation (4) to 100. For all RNN models throughout experiments, we use the Gated Recurrent Unit (GRU) (Cho et al., 2014) . To ensure a fair comparison, the number of parameters across all models is set to be approximately equal. Results in Table 1 are obtained using the optimize-then-discretize approach, following (Kidger, 2022) . With the same training settings, Vanilla Neural ODEs failed to learn the dynamics of the toggle switch based on diffrax proposed by (Kidger, 2022) . 
 
 Appendix:  


Decision: Dear Authors,

Thank you for submitting the draft.

Both reviewers agree that the presented work presents some interesting strengths. However, both reviewers do also raise some major points of concern, especially regarding the clarity of the presentation and some of the claims made. It is expected that authors will be addressing comments by the reviewers in the final draft. 

regards

AC

Reviews: review: title: title: Review of "Accelerating Neural Differential Equations for Irregularly-Sampled Dynamical Systems Using Variational Formulation"
main_review: main_review: # Summary
The paper proposes a novel training method for neural differential equations based on the variational formulation (VF) of ODEs.
The benefit of the proposed method is that it bypasses the need for numerical ODE solvers during training, with the result that training may be accelerated by multiple orders of magnitude.
Additionally, the paper addresses the challenges of oscillatory integral terms in the proposed VF loss using Filon's method.
In a number of experiments, the proposed method is 1-2 orders of magnitude faster than the baseline architectures.

# Strengths
1. (**significance**)
The proposed method is well motivated and has the potential to solve a real problem with high impact across multiple scientific disciplines.
2. (**quality**)
The use of Filon's method to address the oscillatory integral problem is a key component of the paper's methodological contribution.
3. (**quality**, **clarity**)
The paper is well written and enjoyable to read.

# Weaknesses
1. (**clarity**)
The paper builds upon the work of Qian et al. [1], but I think a few more sentences would help to rigorously connect the two settings, which are slightly different (neural differential equations versus closed-form ODE discovery).
In particular, Qian et al. [1] go to great lengths to justify their training objective, an optimization problem over a function space.
The training objective of this paper (Eq. 4) is presented without comment as equivalent to Qian et al.'s.
After some consideration, I think the equivalence is probably valid.
However, I think the clarity of this paper would be improved if this step were justified more explicitly.
1. (**quality**, **completeness**)
The main issue I have with the paper in its current form is the choice of experiments.
To the best of my understanding, the "Fast-VF Neural ODE" is a standard, vanilla neural ODE architecture trained using the proposed VF method.
This is then compared to a number of more exotic neural ODE architectures, each of which is trained using three standard solver-based approaches (discretize-then-optimize, optimize-then-discretize, seminorm).
The problem is that each of these baselines is effectively "two steps removed" from Fast-VF Neural ODE, that is to say, both the architecture *and* the training method are different, meaning I cannot properly evaluate the effect of the proposed training method in isolation.
What I would like to see, at a minimum, is a comparison of Fast-VF Neural ODE with another vanilla neural ODE that was trained using each of the three solver-based approaches.
Conversely, if possible, it would also be interesting to see the more exotic architectures trained using the proposed VF method, which I think is possible in principle.

# Other Comments
1. A number of methods already exist for simulation-free training of neural ODEs in the context of continuous normalizing flows, where they parametrize probability paths; see, for example, [2-4].
I'm not sure whether these methods could be adapted to the present setting, but I think this paper would benefit from a discussion of them.
1. It seems to me that applying Filon's method is essential to making this method work.
With that in mind, I think future versions of this paper would benefit from a discussion of its limitations.
For example, are there situations where Filon's method might fail?

# Conclusion
Overall, I like this paper a lot.
Unfortunately, the lack of more meaningful and direct comparisons, as discussed in the second weakness above, prevents me from giving the paper a higher score at the moment.

# Citations
[1] Zhaozhi Qian, Krzysztof Kacprzyk, & Mihaela van der Schaar. D-CODE: Discovering Closed-form ODEs from Observed Trajectories. In *The Tenth International Conference on Learning Representations*, 2022.

[2] Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow Matching for Generative Modeling. In *The Eleventh International Conference on Learning Representations*, 2023.

[3] Heli Ben-Hamu, Samuel Cohen, Joey Bose, Brandon Amos, Maximillian Nickel, Aditya Grover, Ricky T. Q. Chen, Yaron Lipman. Matching Normalizing Flows and Probability Paths on Manifolds. In *Proceedings of the 39th International Conference on Machine Learning*, PMLR 162:1749-1763, 2022. 

[4] Noam Rozen, Aditya Grover, Maximilian Nickel, Yaron Lipman. Moser Flow: Divergence-based Generative Modeling on Manifolds. In *Advances in Neural Information Processing Systems*, 2021.
title: title: Review of "Accelerating Neural Differential Equations for Irregularly-Sampled Dynamical Systems Using Variational Formulation"
main_review: main_review: The paper proposes a new training algorithm for neural ODEs based on a variational formulation. This effectively replaces the numerical ODE solver with numerical quadrature, with the goal of thereby accelerating the training procedure. 

Overall, the proposed method seems interesting and promising, and I think it is a valuable addition to the neural ODE literature. Unfortunately, I have concerns regarding the presentation and in particular regarding some claims made by the authors, that seem overly strong and insufficiently supported. I will elaborate on individual points in the following.
- *Numerical ODE solvers as the bottleneck for neural ODEs:* This claim is made throughout the paper, in particular in the introduction and in the abstract, but I find this statement quite unclear. 
  - "existing training method, often based on adaptive-step-size numerical ODE solvers, are time-consuming and may introduce additional errors": If anything, does step size adaptation not help with the efficiency by only evaluating the model when required? Also, in which sense does this introduce additional errors? Paragraph two states again that "ODE solver-based training methods [...] are inherently time-consuming", which again does not seem like a true statement.  
  - "the optimize-then-discretize approach incurs additional numerical discretization error, resulting in inaccurate gradients, and potentially causng the training process to fail entirely": To the best of my knowledge, this is due to stability issues of the adjoint ODE, and can be mitigated e.g. with checkpointing or computing and storing an interpolant; see e.g. [1].
  - Paragraph 3 again states multiple times that ODE solvers are the computational bottleneck and leads to inaccuracies, without substantial support.  
In summary, from the perspective of the initial Neural ODE paper by Chen et al, are ODE solvers not exactly the object that were added to ResNets that brought all of their nice benefits? An are adaptive steps not exactly the mechanism that saves computation? Of course, if there is a way to speed up the process that would be great! I just struggle to see the evidence for the claim that the paper tries to repeatedly make.  
- *The computational cost of the proposed method:* The introduction mentions that "this VF loss requires one numerical integration" (similar statement at the end of section 2.2). But does equation 4 not show that $C_j$ needs to be computed for $S$ different $g_s$ functions, implying that we need to solve $S$ integrals? In any case, with the strong claims made in the paper this statement could and should be much more rigorous: How many evaluations of the ODE vector field, i.e. the neural network, do existing methods need, and how many are needed by the proposed method? These numbers are directly comparable, as the main assumption here seems to be that exactly thos evaluations are the expensive part I believe? Or if the speed-up comes from a different aspect, e.g. the better parallelizability of numerical quadrature vs numerical ODE solvers, then this should be elaborated on. 
- *Insufficiently clear numerical evaluation:* The experiments compare many different methods, some of which are much more complicated than the simple neural ODEs that were required in the paper. I believe the evaluation would be strengthened by being much simpler. From my understanding, the main contribution of the paper is a novel loss function for training neural ODEs, as well as thereby also a novel way to compute parameter gradients. Just comparing this to existing methods would already provide helpful evidence. For example, benchmarking the cost of evaluating the loss function vs a standard loss function, and benchmarking the cost of evaluating gradients vs Opt-Dis and Dis-Opt and Seminorm; as well as also comparing the quality of both loss and gradients to some reference, e.g. computed by using an existing established ODE-solver based method with very low accuracy settings. A possible result of this would then hopefully be that all methods achieve similar quality of loss and gradients, but the proposed method presents a significant speed-up over the others. If then it was previously also discussed where exactly this comes from (number of NN evals or parallelisation?) then this would already be all the evidence that this paper needs. Of course, looking at multiple ODEs and even at ODE-RNNs etc is good additional evidence.

So overall, I do believe that the proposed method is valuable and should be presented to the neural ODE community, but in the current form the paper makes too many claims that are insufficiently supported. I would strongly recommend to the authors to revise the manuscript, soften many claims, and provide a simpler and clearer evaluation of the method, to better present this interesting method. 
  
  
[1] Ma et al, "A Comparison of Automatic Differentiation and Continuous
Sensitivity Analysis for Derivatives of Differential Equation Solutions"