Paper: Mechanistic Neural Networks (ICLR-2024)

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

Author response:
**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.

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

Author response:
**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).

Reviewer comment (Writing):
What is $\hat{r}$ in 5.4?

Author response:
>> 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.

Reviewer comment (Writing):
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)$.

Author response:
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.

Reviewer comment (Theory):
On page 2, the statement "In general, one ODE is generated for a single data example $x$ and a different example $x'$ 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)?

Author response:
**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.

Reviewer comment (Novelty):
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$?

Author response:
**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$.

Reviewer comment (Writing):
The methodology for handling nonlinear ODEs, as presented on page 4, lacks clarity and could benefit from a more detailed exposition.

Author response:
**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 
\begin{equation}
c\_2(t) u'' + c\_1(t) u' + d(t) u^3 = b,
\end{equation}
where $u^3$ is non-linear functions of $u$, we solve the following two equations.
\begin{align}
c\_2(t) u'' + c\_1(t) u' + d(t) \nu= b\\\\
u^3 = \nu
\end{align}

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.

Reviewer comment (Novelty):
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.

Author response:
**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.

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

Author response:
**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$.
\begin{equation}
\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}
\end{equation}

$\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$.

Reviewer comment (Experiments):
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.

Author response:
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.

Reviewer comment (Experiments):
While the methodolgy is pretty interesting, I wonder how it scales with the number of dimensions in the input data, and for more complex systems like Navier-Stokes. Some discussion related to it would be useful!

Author response:
**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.

Reviewer comment (Experiments):
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.

Author response:
**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.

Reviewer comment (Theory):
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?

Author response:
**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.

Reviewer comment (Theory):
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.

Author response:
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).