Mechanistic Neural Networks

We thank the reviewer for the comments.

I also found eqs 8 and 9 to not be clear, could you please further explain their role and how they are obtained?

Smoothness Constraint Elaboration. The purpose of the constraints in eqs 8 and 9 is to enforce the solution of the ODE to be smooth.

The QP defines ODE solution as variables $u\_t, u'\_t, u''\_t$ over the time grid. However, these would not be smooth without additional constraints.

We use Taylor approximations to enforce smoothness between solution at time $t$ and $t+1$.

Given $u\_t, u'\_t, u''\_t$ we build the second order Taylor polynomial to approximate the function at time $t+1$ with step size $h$ as $\tilde{u}\_{t+1} = u\_t + hu'\_t + \frac{1}{2}h^2 u''\_t.$ Next we say that the difference between the approximation and the solution at $t+1$ should be $\epsilon$. $\tilde{u}\_{t+1} - u\_{t+1} = \epsilon,$ which is the first equation in 8 which enforces the function to be smooth by minimizing $\epsilon^2$ which is the optimization objective.

Similarly, the next equation in 8 enforces the first derivative to be smooth using the Taylor approximation.

The constraints in 8 say that we should be able to approximate the next time step $t+1$ using values at $t$. The constraints in 9 go backward using values at $t+1$ to approximate values at $t$. We find both constraints to be necessary to get proper solutions.

The paper emphasizes throughout how it is fundamentally different than neural ODEs, making it sound like neural ODEs cannot be used to solve tasks that mechanistic neural networks can be used for. Yet, most of the empirical comparisons are carried out against neural ODEs.

Neural ODE Comparison. We emphasize the difference of our approach in combining Neural networks and ODEs which is very different from that of Neural ODEs. Nevertheless, there is an overlap in the types of tasks for which Neural ODEs and MNN can be used, such as modeling dynamical systems and processes. We compare against Neural ODE for such tasks (Classification, Vibration prediction, N-body prediction) in Sections 5.2,5.3,5.5. MNNs have better performance on these tasks because the underlying solver performs better at learning dynamical processes than the solver type used in Neural ODE (see Section C.7).

On the other hand, MNNs are also designed for discovery. In discovery experiments we compare against a SINDy (Section 5.1) which is a popular sparse discovery method for ODEs.

For our PDE solving architecture we compare against FNO-based and similar methods.

Other than table 1, which compares exclusively solvers, the paper presents no quantitative comparisons against its baselines, only plots qualitatively comparing trajectories.

Quantitative Results. We point out that there are several other quantitative comparisons in the paper.

In Figure 4, we plot loss curves (left) and prediction error (second from left) which are quantitative and compare against Augmented Neural ODE and Second Order Neural ODEs.

In Table 2 in the appendix we compare the training and evaluation results on the 2-body and 10-body problem against Neural ODE.

In Table 3 in the appendix we compare the QP solver against the RK4 solver from torchdiffeq, where we look at the ability of the two solvers to fit long sinusoidal waves (300 and 1000 steps) and their respective training times.

Finally in Table 1 (Section 5.5) we give quantitative comparisons for a PDE solving architecture against other well-known deep learning based PDE solving architectures. These methods are not solvers in the traditional sense of differential equations solving but are learning-based approaches.

The above quantitative results are in addition to the host of results that we have shown for discovering governing equations (Sections 5.1, 5.4, C.1.2) and prediction (Sections 5.2,5.3,5.5).

Notation.

The notation is quite confusing, and I think the authors use the same symbols to denote different objects in different sections ...

Unfortunately, describing our approach has to combine methods from differential equations, optimization and deep learning where each has standard notation and some overlap has been difficult to avoid. We have made many of the changes suggested and appreciate the reviewer's suggestions for improvement.

What is $\hat{r}$ in 5.4?

$\hat{r}$ is the radius unit vector that gives the direction of force in Newton's law of gravitation. We have added a clarification.

Finally we appreciate the typographical corrections pointed out by the reviewer and are grateful for the review in general.

The methodology for handling nonlinear ODEs, as presented on page 4, lacks clarity and could benefit from a more detailed exposition.

Non-linear ODEs. We have added further clarification in the appendix A.2 with an example non-linear ODE.

Briefly, to solve an equation like

$

c_2(t) u'' + c_1(t) u' + d(t) u^3 = b,

$

where $u^3$ is non-linear functions of $u$, we solve the following two equations.

$

c_2(t) u'' + c_1(t) u' + d(t) \nu= b\\ u^3 = \nu

$

We solve the first equation by the QP solver by creating another set of variables $\nu\_t$ in addition to $u\_t$. We add equation constraints that include both types of variables and include smoothness constraints for $\nu\_t$.

Solving the program gives a solution for $u\_t$ and $\nu\_t$. We enforce the non-linearity in the neural network loss function by add an MSE term $\sum_t (u\_t^3 - \nu\_t)^2$ which imposes that constraint that $u^3 = \nu$. This shows how non-linear relationships in the ODE terms are learned.

The literature review in section 4 seems outdated, with the most recent references dating back to 2020. A comprehensive literature survey, including more recent and relevant baselines, would strengthen the paper's context. For example for the ODE modelling, [1][2] may be included.

References. Thank you for the suggestion about more recent references; we will consider this for the final draft. This is partly due to the fact that there hasn't been much significant development in Neural ODE style modeling over the past few years. The reviewer's suggested references are also from 2020 and 2021.

We demonstrate that our method overcomes significant limitations of the Neural ODE framework and adds directions such as explicit equation discovery which were not possible in the Neural ODE method, in addition to faster parallel solving and better learning.

The term $\Theta \xi$ on page 6 is introduced without clear definition or context. Could the authors provide clarity on this?

SINDy Notation. In Section 5.1, we assume and follow the notation used by Brunton et al (2016) that introduces the SINDy method. For clarity and completeness we have included further details in the appendix A.3 with a reference in the main text.

Briefly $\Theta$ is a data matrix that applies polynomial basis functions to time-series input $x\_0, x\_1 \ldots x\_t$.

$

\Theta = \begin{bmatrix} 1 & x_0 & x_0^2 & \ldots\\ 1 & x_1 & x_1^2 &\ldots\\ \vdots & \vdots & \vdots \\ 1 & x_t & x_t^2 &\ldots \end{bmatrix}

$

$\xi$ is a sparse vector that selects a small number of basis functions. So ODE like $x' = x + x^2$ may be represented with the appropriate choice of $\xi$ as $x' = \Theta\xi$.

The paper delineates two primary components: the learning of the ODE and its subsequent solving. However, the experimental section seems to lack comprehensive ablation studies that convincingly demonstrate the efficacy of each individual component.

Our original submission included the following ablations in the appendix (referenced in Section 5.7) due to lack of space.

RK4 Solver Comparison. In Section B.2 we validate the QP solver and compare the solution with the RK4 solver solution with the SciPy Python package (Figure 7).

Non-linear ODEs. We show learning a non-linear ODE also in B.2 and Figure 9.

Step Learning. In Section B.2 and Figure 8 we show that the method learns step sizes during training.

Learning with Noise. In Section B.3 we show that the method is able to learn ODEs from noisy data both with time-dependent and time-independent coefficients.

RK4 Learning and Timing Comparison. In Section C.7 we show a comparison with the RK4 solver from torchdiffeq for learning longer noisy sine waves with 300 and 1000 steps and compare the times. The results (Table 3. Figures 14,15) show that the QP solver is able to learn significantly better fits and is up to 200x times faster in this experiment.

However, if there are suggestions for any additional ablation experiments, please let us know.

We thank and appreciate the reviewer's comments and suggestions.

The paper's clarity is wanting. For example, around eq(2), there's an inconsistency in notation with both $c\_i(t;x)$ and $c\_i(t)$ being used. Which of these is the intended notation? Additionally, there's no explicit description or definition of $b(t;x)$ .

The notation $c\_i(t;x)$ was intended to make dependence of the ODE coefficients on neural network input $x$ explicit. For a fixed datum $x$, we use $c\_i(t)$. To avoid confusion, we have made the notation consistent in the updated draft and have given a definition of $b(x;t)$, defined similarly to the other coefficients.

On page 2, the statement "In general, one ODE is generated for a single data example and a different example would lead to a different ODE" is made. Could the authors elucidate why this is a characteristic of the modeling approach presented in eq(2)?

Per-instance ODEs. The ODE in equation 2 depending on the input datum $x$ means that we have one ODE per input datum in the dataset in the most general case. For instance, if wanted the MNN can derive a different governing ODE for each new input sequence trajectory, initial condition, etc. Of course, if wanted the MNN can also be defined to derive a single governing ODE from all the data in the dataset, e.g when there is a single governing equation in the discovery experiments in Sections 5.1 and 5.4.

This draws a contrast with Neural ODEs, which are only able to model a global ODE for an entire dataset.

If a new instance $x'$ necessitates retraining the model, wouldn't it be more streamlined to directly learn a neural ODE through parameter optimization, bypassing the need for coefficients as functions of $x$ ?

Retraining. It is not correct that a new instance $x'$ requires retraining the model. In our model a neural network produces the parameters $c\_i, b, \phi\_k$ (equation 2) of an ODE for a given input $x$. A dataset can then be modeled by a family of ODEs by training over the dataset. As we mention above, MNN is more general than neural ODE and is able to model per-instance ODEs without retraining, simply returning the datum-specific coefficients by conditioning on the datum. This is indicated by the dependence of the produced ODE on the input datum $x$.

The paper's approach to solving any nonlinear ODEs using equality-constrained quadratic programming must have gaps. These gaps aren't clearly addressed. Relying on such an algorithm to solve any ODE without theoretical guarantees is precarious. A more transparent discussion on potential limitations is needed.

Theoretical Analysis. We have added an error analysis section in the updated draft in appendix A.1 for linear second order ODEs with the QP solver $c\_2 u'' + c\_1 u' + c\_0 u = b.$

Sketch. The analysis proceeds by estimating total error introduced by the error in the second order Taylor approximation $O(h^3)$ and the error from the constraint $\epsilon$. In the first step the function and derivative values are known from the initial value constraints leading to an error of $O(\epsilon + h^3).$ For the next step we use the approximate function and derivative values, plugged into the ODE, and show that the total error over $N$ steps remains bounded by $O(N(\epsilon + h^3))$, under the assumption that $\frac{c\_1}{c\_2}$ and $\frac{c\_0}{c\_2}$ are bounded by $O(\frac{1}{h^2})$ and $O(\frac{1}{h})$, respectively.

The analysis shows that the method has a local error of $O(\epsilon + h^3)$, where $\epsilon$ is the QP error variable and $h$ is a fixed step size. For an $N$ step grid, the analysis shows that this error becomes $O(N(\epsilon + h^3))$.

In particular for $\epsilon = h^3$, which is reasonable in practice, and $N=1/h$ the global error simplifies to $O(h^2)$. This error depends on the number of terms of the Taylor expansion and could be further reduced by using central difference approximations for all derivatives, investigation of which we leave to future work.

The analysis implies that the cumulative error becomes potentially large only when $\frac{c0}{c2}$, $\frac{c1}{c2}$ are very large: In little omega notation, when $\frac{c0}{c2}$ is $\omega(\frac{1}{h^2})$ and $\frac{c1}{c2}$ is $\omega(\frac{1}{h})$.

We thank the reviewer for reviewing and appreciating our work. We address some of the points and questions raised below.

I wonder how it scales with the number of dimensions in the input data, and for more complex systems like Navier-Stokes.

The results on PDEs like Darcy Flow show that the network is not as much better (at least in performance) as compared to FNO baseline.

PDE. In this paper we present a general framework for learning the governing equation from data and not a specific solution to a specialized class of PDEs. Naturally, when narrowing down the focus to PDE forecasting, a specialized solver like FNO works very well.

MNN are a general and encompassing framework compared with the overarching state-of-the-art of ODE-related ML methods where

1. Neural ODEs model sequences but cannot discover governing equations,
2. SINDy discovers governing equations but is not with neural networks so it cannot scale and it cannot discover non-linearly entangled governing equations,
3. FNO can forecast future trajectories but is purely data driven, has no notion of governing equations.

MNNs cut across these types of methods and in many cases perform better.

Scaling. Scaling to large dimensions especially for PDEs depends on how effectively sparse matrix methods can be used for solving linear equations, since QP constraint matrices are highly sparse. Currently we use dense matrix factorizations which are fast and have manageable memory usage for the problems we have considered. We plan to investigate sparse matrix methods for PDEs in follow-up work.

I am a bit unclear about how the authors are able to parallelize in terms of the sequences, I understand the parallel batchwise training aspect of the training. Is is that since we have ground truth training data, and the we can write down the forward and backward Taylor approximations, we get different sets of equations that are solved in parallel?

Parallelization. Yes, since we can write the discretized ODE over an entire time sequence in terms of the constraints, we only have to solve the constrained optimization to get the entire solution over the grid without stepping through time.

One minor point is that the parallel solution does not depend on the ground-truth data and is a feature of the forward pass.

It seems that the authors are solving a relatively complex system of equations stemming from the discretization of the ODEs. Are there any set of equations that may not be solvable due to the given methodology.

Solving the equality-constrained quadratic program only requires solving linear equations so as long as the constraints are not inconsistent (such as contradictory initial conditions) we do not encounter any problems.

However, as we mention in Section 3.1 regarding non-linear ODEs, the QP method on its own does not solve non-linear ODEs (since the constraints are linear). For this we introduce new variables representing non-linear terms (say a squared term) in the QP and solve the QP. The the solution corresponding to the non-linear variables is related to the original in the MSE loss function using the required non-linear function (say the square function).

We thank the reviewer again for the review and appreciation and hope that we were able to answer the reviewer's questions.