RECURSIVE NEURAL ORDINARY DIFFERENTIAL EQUATIONS FOR PARTIALLY OBSERVED SYSTEM

On the reply on weakness 2:

Regarding SINDy's strategies, to our knowledge, they are not designed to deal with unobserved states. Furthermore, SINDy does not learn global models. Instead, SINDy learns a local model $\Xi_k$ at each time step $k$. This is why we did not compare with it as we are learning a global model.

SINDy is supposed to be learning a global model of the dynamics, not local ones. On your second point, it has been applied in conjunction with delay embedding to identify dynamical systems from partial observations [1].

[1] Bakarji, J., Champion, K., Kutz, J. N., & Brunton, S. L. (2022). Discovering governing equations from partial measurements with deep delay autoencoders. arXiv preprint arXiv:2201.05136.

We thank the reviewer for the detailed comments which definitely improved the clarity and quality of the revised manuscript. Below, we reply to all Weaknesses provided by the reviewer point by point. When referring to equations we refer to the numbers in the original manuscript unless explicitly stated.

• Weakness 1:

We agree that the presentation, especially the positioning of the method, was a bit confusing in the original manuscript. Besides, we would like to highlight that the proposed approach is not based on the adjoint method, and that giving too much emphasis to NODE and its possible problems may have added to the confusion the reviewer experienced. In the case of the proposed RNODE approach, we leveraged an alternating optimization to avoid vanishing gradients that would be caused by our employed methodology. Again this point is not necessarily related to the standard NODE.

In fact, after the reviewer's comments, we noticed that in the case of standard NODE, the initial condition of the adjoints does indeed vanish, but its integral does not. For this reason, and to improve the clarity of the paper regarding our contributions, we removed Section 3.3 of the original manuscript and revised our discussions regarding vanishing gradients in the manuscript. Furthermore, we included results with NODE in all experiments.

It becomes evident that NODE leads to poor performance for the more challenging scenarios since it was not designed to deal with unobserved latent states. We thank the reviewer for their critical review which certainly improved the clarity of our contributions and positioning of the RNODE approach.

The open problem is stated in the abstract, in the second paragraph of the introduction and related work section. The contribution of the proposed method is also discussed at the end of the introduction, and in the related work sections. To emphasize both the open problem and contribution we modified to third paragraph of the introduction to improve the clarity of the problem and contributions. We also introduced a paragraph "The partial observation problem" at the end of Section 3.1 to address this issue.

Regarding the usage of Hessians, they are always used. In Appendix A, we provide the full derivation of Theorem~1 where the connection of the different quantities P’s, Q’s, R’s, etc, and the Hessians are clearly shown.

Although in our experiments we consider a small-variance random walk, mainly for numerical reasons, for the parameters process the formulation is general and other processes could be used. For this reason, we kept the general parameter process formulation presented in Section 3.1.

• Weakness 2:

We would like to point out that we did not compare it with vanilla ODE in our original manuscript because existing NODE is not demonstrated to learn with partially observed state constraint. Nevertheless, we agree with the reviewer that standard neural ODE comparison with RNODE would be valuable and provide more context, that’s why we modified our experimental results section to include benchmarks of RNODE against the vanilla NODE approach.

Regarding SINDy's strategies, to our knowledge, they are not designed to deal with unobserved states. Furthermore, SINDy does not learn global models. Instead, SINDy learns a local model $\Xi_k$ at each time step $k$. This is why we did not compare with it as we are learning a global model.

Concerning the run-time analysis, we included the complexity of the proposed algorithm in Appendix D of the revised manuscript. We refer to it at the end of Section 4.1.

• Weakness 3:

Thank you for your suggestion, however, we would like to clarify that we did include a real-world experiment which is the Electro-Mechanical Positioning System. Moreover, even though we believe the Hodgkin-Huxley model is fairly complex, we added a complex biochemistry example, namely, yeast glycolysis to Appendix B.

Below we reply to the Questions point by point.

• We agree with the reviewer. We refer to $dx$ and $d\theta$ as stochastic differential equations in the revised manuscript as defined in [5]. We also modified the definitions in Section 3.1, and also the generative model in Section 3.4 of the original manuscript to cope with the stochasticity of the approach. We apologize for the confusion. Regarding the second part of the question, $x$ represents the dynamical system states, which we want to learn parts of its dynamics using neural networks parameterized by $\theta$. In the more general scenario of non-stationary processes, $\theta$ may have non-zero dynamics to cope with non-stationary scenarios such as [1], to speed up the learning procedure, or to avoid numerical issues.

• We understand the reviewer's point. Nevertheless, we opted for this simplistic and at the same time, very restrictive condition, since we require some type of distinguishability. We would like to point out that states following Definition 3.1 lead to unique observations. This is necessary for recovering unique states from the observations. Without this condition, there is no guarantee that states can be uniquely recovered from the observations. A simple example is $y = x^2$ where y represents the observations and $x$ is the latent variable. It is clear that in this scenario one cannot uniquely recover $x$ from $y$. We agree that such a condition is sufficient but may not be necessary. However, determining the necessary conditions for state estimation in dynamical models with neural networks is a very challenging open problem and is outside the scope of this paper.

• The work focuses on learning unknown dynamics from data where the state related to the unknown dynamics is unobserved. Since the dynamics is unknown, we assume parametric models such as NNs. We clarified this point in the third paragraph of Section 1. Regarding non-stationarity, although the RNODE approach could indeed be used to deal with such systems, in this paper we learn a constant dynamical model as shown in our experiments.

• The reviewer is correct. We revised this point and noticed that although the initial condition of the adjoint vanishes, this is not necessarily the case for it's integral as we pointed out in the reply for Weakness 1.

• To generate $x(t_i)$ in Equation 5 (Equation 4 in the revised manuscript), we need $x(t_{i-1})$ and $u(t_{i-1})$ and hence $x_{t_i}=f_o(x(t_{i-1}),u(t_{i-1}))$. Note that this parametrization of input variables is ubiquitous in the dynamical system literature [2]. Please, note the solution for $\theta$ in Equation 6 (Equation 5 in the revised manuscript).

• Note that in the text before Equation 1 we define $\theta$ as the system's parameters. In the scenario where part of the dynamics are modeled by NNs, then $\theta$ also refers to the parameters of the NNs. We refer to $x$ as dynamical system states which are modeled by $f$. In our work, we tested the proposed approach using hybrid models where partial knowledge of $f$ was assumed, and where the dynamics of non-observed states were considered unknown and modeled as a neural network parameterized by $\theta$. We emphasized this point in the third paragraph of the Introduction in the revised manuscript.

• The Hessians come from the application of Newton's method as shown in Appendix A. In the Appendix we show how the Hessians can be written as functions of other quantities $P$'s, $Q$'s, etc, leading to the update expressions shown in Theorem 4.1. We assume the observation function $h$ to be known for the known states. For instance, in a 2D system, if $h=I$, $y=[x_1,x_2]$, if we observe $x_1$, then $H=[1 \quad 0]$. This copes with many systems and applications in control, physiological systems modeling, etc, where the measurement functions for known states are known.

Considering the vanishing adjoints, see the reply to Weakness 1.

Regarding Equations 10 and 11 in the original manuscript, the reviewer's understanding is inaccurate since we do not just set $x$ and $\theta$ to the ODE solution. We split the problem in Equation 9 (in the original manuscript) into four sub-problems, see, Equations 10, 11, 13, and 15 of the original manuscript, to obtain alternating optimizations for $x$ and $\theta$ for prediction (Eqs. 10 and 11) and update steps (Eqs. 13 and 15). The benefit of this approach is that we can estimate $\theta$ in (15) using the solution for $x$ found in (13) which contains the estimated unobserved state of the hidden dynamics. Finally, $x(t_i)$ and $\theta(t_i)$ are updated at each time step according to (17) and (18) in the original paper. We modified Section 4.1 to clarify these points.

Dear reviewer,

we submitted the revised manuscript. We believe to have addressed most of the reviewer's points and hope the revised manuscript matches your expectations. Thank you for your valuable input.

There was indeed a confusion. We tried to address this point in the last few hours by including plots of the true and estimated vector fields for the Harmonic Oscillator example in Appendix B4 (see Fig. 9). The revised manuscript has been submitted. We thank the reviewer one more time for the valuable comments.

• These works focused on the filtering aspects where during training and testing parameters and states were continually updated. In RNODE, we aim at learning a global model. While we performed recursive joint estimation of states and parameters during the training phase, we only integrated the learned system during test. We also included a sentence in the third paragraph of Section~5 to highlight this point: "All results were obtained with learned mean vector field integrated over time."

• That is true, we thank you for pointing that out. We corrected our model description sections 3.1 and 3.3 in the revised version where we refer to the model as a stochastic differential equation.

• Since the optimization steps are performed according to Newton’s second-order method, they enjoy the same convergence properties as Newton’s method for every optimization step.

• We agree that a complexity analysis should be presented, for this reason, we added a complexity section in the appendix. The complexity of the overall method at each epoch is $\mathcal{O}(N(2d_{\theta}^2d_x + 2d_{\theta}d_x^2))$, which is time-expensive if compared to gradient methods such as SGD and Adam. We would like to reiterate that computation speed was not the core focus of the work, instead, we targeted learning under partial observation. We would also like to point also that the method can be improved by Hessian approximations and other second-order optimization approximations used in [1], [2], and [3]. Nevertheless, that was not the focus of our work and we leave it for future work.

• Thank you for pointing that out, we added the reference.

• RNODE is a sequential joint estimation procedure for learning systems dynamics. Our understating is that the proposed approach is not directly connected to Expectation-Maximization (EM) in the sense that we are not maximizing a lower-bound of the marginal likelihood.

• We emphasized in section 5 of the paper that PR-SSM does not account for the ODE structure and replaces all of the ODE dimensions with a neural network. Moreover modifying the PR-SSM method to account for ODE structure was out of the scope of the paper. We suspect that since all the dimensions are modeled as neural networks, the hidden state is not recoverable since the neural networks used are not state identifiable, which is a requirement to learn hidden state dynamics as we explained in the previous reply. Therefore, incorrect values of the hidden state are used in parameter inference.

• During testing, we did not perform parameter estimation. All results, in Section 5 and the appendix, were performed by simply integrating the learned vector fields which were compared with the ground-truth. The set of parameters found in the final step, $\hat{\theta}(t_N)$, are saved and then plugged into equation (4) of the original manuscript, which integrates the learned vector field and produces the plots and nRMSEs corresponding to RNODE. To clarify this in the manuscript we include the sentence "All results were obtained with learned mean vector field integrated over time." in the third paragraph of Section~5.

• We revised the manuscript for typos.

[1] Gupta, Vineet, Tomer Koren, and Yoram Singer. "Shampoo: Preconditioned stochastic tensor optimization." International Conference on Machine Learning. PMLR, 2018.

[2] Anil, Rohan, et al. "Scalable second order optimization for deep learning." arXiv preprint arXiv:2002.09018 (2020).

[3] Peirson, Abel, et al. "Fishy: Layerwise Fisher Approximation for Higher-order Neural Network Optimization." Has it Trained Yet? NeurIPS 2022 Workshop. 2022.

[4] Øksendal, B. and Øksendal, B., 2003. Stochastic differential equations (pp. 65-84). Springer Berlin Heidelberg.

Thank you very much for your deep understanding and valuable suggestions. We are pleased that you recognized the novelty and performance of the proposed approach. First, we would like to highlight to the reviewer that, during testing, we did not perform parameter estimation, we learned the parameters, and the results displayed in the paper show the learned vector field with respect to the ground truth. The set of parameters found in the final training step $\hat{\theta}(t_N)$ are saved, then plugged into equation (4) of the original manuscript, which integrates the learned vector field and produces the plots corresponding to RNODE. We thought it was important to clarify this point because we felt it was not properly conveyed on our part and not properly understood by the reviewer.

Bellow we reply to the Weaknesses point by point:

• That is true, we thank you for pointing that out. We added the stochastic terms to our model description in sections 3.1 and 3.3 of the revised manuscript. Added zero-mean noise, which makes our model equivalent to a stochastic differential equation in [4].

• We thank the reviewer for bringing up this important point and we would like to take this opportunity to clarify our contribution further. Our framework learns dynamics governing unobserved states, and the only way these states can be identified is through their relationship with observed states. For this reason, unobserved states should be distinguishable. For example, let's assume we have a system of neural ODEs described as follows:

$\dot{x}_1(t) =f_1(x_1(t),x_2(t),\theta_1)$

$\dot{x}_2(t) =f_2(x_1(t),x_2(t),\theta_2)$

$y(t)=h(x_1(t),x_2(t))$

where our task is to learn $f_2(x_1(t),x_2(t),\theta_2)$ without measuring $x_2(t)$. Therefore, to identify $x_2(t)$, the system has to be state identifiable, that is $x_2(t)$ has to be recoverable from observations $y(t)$.

A simple example of a undistinguishable scenario is $y(t) = h(x_1(t)=x_1^2(t)$. It is clear that in this scenario one cannot uniquely recover $x_1(t)$ from $y(t)$ since $x_1(t)$ is undistinguishable and has two possible values given one measurement $y(t)$.

Another example of a distinguishable scenario is $\dot{x}_1(t)=x_2(t)$ and $y(t)=x(t)$. Assuming we observed $y(t)=x_1(t)$ at $t_0$ and $t_1$, then $x_1(t_1)-x_1(t_0) = (t_1-t_0)x_2(t_1)$, therefore $x_2(t_1)=\frac{x_1(t_1)-x_1(t_0)}{t_1-t_0}$ is the unique solution. It is clear that in this scenario one can uniquely recover $x_2$ from $y$ with knowledge of $t_1$ and $t_0$ since $x$ is distinguishable.

We agree although, that distinguishability condition is sufficient to ensure identifiability but may not be necessary. In that sense, yes this could be quite limiting.

To our knowledge, and as demonstrated by our experiments, there is no machine learning approach capable of learning ODEs governing unmeasured processes without requiring identifiability even if not stated.

In control theory, necessary conditions for state identifiability are for classical systems is extensively covered under Observability of linear and nonlinear system. However, determining the necessary conditions for state identifiability in dynamical models with neural networks is still a very challenging open problem and is outside the scope of this paper. Therefore, we still do not have the necessary tools to identify under which scenario distinguishability condition could be eased.

Below is the reply to the questions point by point:

• In the proposed approach, the network parameters are included as part of the state space. Therefore, since the size of the neural network is already way larger than the system dimension, the proposed approach would not have a problem with moderate dimensional systems since it would not add many variables to the state space compared to the number of neural network parameters. However, our approach may scale badly to very higher-dimensional neural network architectures since it’s a second-order Newton method at its core. To calculate the complexity of our algorithm, first, we need to calculate the complexity of each Jacobian. Since we are using automatic differentiation to calculate them, the complexity of getting $F_{x_{i-1}}, G_{\theta_{i-1}}$, and $F_{\theta_{i-1}}$ is $\mathcal{O}(d_x^2), \mathcal{O}(d_{\theta}^2)$ and $\mathcal{O}(d_{\theta}d_x)$, respectively. However, the complexity of $G_{\theta_{i-1}}$ could be removed if $g(\theta(t_i))=0$ since $G_{\theta_{i-1}}=I$ in that case. Note that this is the case in our experiments. Assuming $n_x \approx n_y$, the complexity of calculating $P_{\theta_i}^{-}$, $P_{x_i}^{-}$, $P_{x_i}$ and $P_{\theta_i}$ are $\mathcal{O}(2d_{\theta}^2d_x + d_x+ (2d_x^2 + 1)d_{\theta})$, $\mathcal{O}(2d_x^3 + d_x)$, $\mathcal{O}(4d_x^3 + 2d_x)$ and $\mathcal{O}(d_\theta)$, respectively. Moreover, the complexity of computing $\hat{x}(t_i)$ is $\mathcal{O}(2d_x^3+d_x^2)$, and the complexity of $\hat{\theta}(t_i)$ is $\mathcal{O}(d_{\theta}(d_x+1))$. Since the Jacobians are differentiated once, then evaluated at each time step, the complexity of the whole algorithm with a dataset of size $N$ for each epoch becomes:

$

\mathcal{O}(N(d_x^3 + d_\theta^2 d_x + d_x^2 d_\theta)).

$

Finally, assuming $d_{\theta} \gg d_x$, the total cost of each training epoch can be simplified as:

$

\mathcal{O}(N(2d_{\theta}^2d_x + 2d_{\theta}d_x^2))

$

It is clear that the main source of complexity would be $d_{\theta}$. Thus, RNODE may scale badly to very higher-dimensional neural network architectures. However, fast approximations of Newton's method exist, as pointed out by the reviewer, such as Shampoo. Although merging Shampoo with our proposed approach could reduce the computational burden, we will analyze this hypothesis in future works.

• A practical example of RNODE estimating the wrong state and turn in finding the wrong parameters would be the following example. let's assume we have a system of neural ODEs described as follows: $\dot{\theta}(t)=g(\theta(t))$ $\dot{x}_1(t)=x_2^2(t)$ $\dot{x}_2(t)=f_2(x_1(t),x_2(t),\theta(t))$ $y(t)=x_1(t)$ where our goal is to learn $f_2(x_1(t),x_2(t),\theta(t))$ with $y=x_1(t)$. Assuming we observed $y=x_1(t)$ at $t_1$ and $t_2$, then $x_1(t_1)=\frac{x_1(t_1)-x_1(t_0)}{t_1-t_0}=x_2^2(t_1)$ therefore $x_2(t_1)=\pm\sqrt{\frac{x_1(t_1)-x_1(t_0)}{t_1-t_0}}$. It is clear that in this scenario one cannot uniquely recover $x_2(t)$ from $y=x_1(t)$ since $x_2(t)$ is not distinguishable and has two possible solutions.\ Therefore, $\theta(t)$ would be learned according to equation (18) in the paper using a wrong value of $x_2$ and would result in learning the wrong vector field $f_2(x_1(t),x_2(t),\theta(t))$.

• The computation complexity of sequential Newton optimization is $\mathcal{O}(N(2d_{\theta}^2d_x + 2d_{\theta}d_x^2))$ per each epoch as calculated above.

• Yes, we have validated the claim that gradients do not vanish with the proposed method RNODE since we were able to learn the hidden dynamics.

• In large-scale data assimilation problems $d_x$ is often very-large. In this context, the proposed approach will suffer during training since large neural networks and number of states are present. In our experiments, we only integrated the learned mean vector field during testing. In this scenario, the complexity at each time step is fairly small (just the application of the dynamical model). If data assimilation is performed with respect to the states but keeping the learned parameters $\theta$ constant, the cost will be as follows assuming $d_x \approx d_y$: $\mathcal{O}(d_{x}^2) + N[\mathcal{O}(2d_x^3 + d_x) + \mathcal{O}(4d_x^3 + 2d_x) + \mathcal{O}(2d_x^3+d_x^2) ]$ which is equivalent to $\mathcal{O}(Nd_x^3)$ for the whole dataset of size N. However, if the learning is also performed during assimilation periods, which can be beneficial in non-stationary scenarios, then the cost will be similar to the training, in which case the method will suffer due to large dimensional states and parameters.

Thank you very much for your detailed review. We are very pleased that the reviewer find our contributions clear and our mathematical derivation of high quality. We thank you also for highlighting our paper strength and weakness which made us improve the paper’s quality significantly.

Below is reply to the weaknesses point by point :

• We thank the reviewer for pointing out these valuable references which we included in the revised manuscript in our related work section. We would like to highlight that in our literature review we focused on articles that tackled the partially observed systems which poses additional challenges to the learning process, overlooking speeding up second order methods since that is not the main focus of our work. Nevertheless, we believe addressing this complexity issue is of very high importance to be addressed in future work. Moreover, NODE extensions such as Heavy ball NODE [4] deal with irregularly sampled time-series but assumes that the system is fully observed. For this reason we only included this reference in the Introduction.

• We agree with the reviewer that the positioning of the algorithm should be clarified. RNODE is a sequential (recursive) optimization approach that at each step solves an alternating optimization problem for learning system dynamics under partially observed states. The benefit of the sequential strategy is twofold: $(1)$ reduce the need for accurate initial conditions during training; $(2)$ avoids simultaneous estimation of all states, making second-order optimization methods feasible. Similarly, the alternating optimization approach improves the optimization of system parameters since it estimates unobserved hidden states and uses them in learning system parameters. In the case of RNODE, it also prevents vanishing gradients. To make these points clearer in the manuscript we modified the last paragraph of the introduction. Regarding the NODE, in the proposed method we do not use adjoints and having a dedicated section for this end made the positioning of the paper confusing. For this reason, we removed section 3.3 from the revised manuscript.

• We agree with the reviewer that the neuron model and the retinal circulation results are close together. We tried to highlight that RNODE shows better performance in terms of RMSE (Table 1). We understand that adding confidence intervals would be beneficial, but we highlight that in our experiments and results, we are only integrating the mean learned dynamics. That is why we did not include uncertainty measures in the plots. Regarding the complexity of examples, although we believe the neuron model is fairly complex. We also agree that adding a more difficult yeast glycolysis example would strengthen the paper. We included this new experiment in the appendix of the revised manuscript. Moreover, the reviewer pointed out that there is no comparison to existing NODE approaches, and we added the vanilla ODE to our experimental benchmarks as mentioned in the previous reply.

Dear reviewer,

we submitted the revised manuscript. We believe to have addressed most of the reviewer's points and hope the revised manuscript matches your expectations. Thank you for your valuable input.