Scalable Mechanistic Neural Networks

## Abstract

Abstract We propose Scalable Mechanistic Neural Network (S-MNN), an enhanced neural network framework designed for scientific machine learning applications involving long temporal sequences. By reformulating the original Mechanistic Neural Network (MNN) (Pervez et al., 2024 ) , we reduce the computational time and space complexities from cubic and quadratic with respect to the sequence length, respectively, to linear . This significant improvement enables efficient modeling of long-term dynamics without sacrificing accuracy or interpretability. Extensive experiments demonstrate that S-MNN matches the original MNN in precision while substantially reducing computational resources. Consequently, S-MNN can drop-in replace the original MNN in applications, providing a practical and efficient tool for integrating mechanistic bottlenecks into neural network models of complex dynamical systems.

## 1 Introduction

The Mechanistic Neural Network (MNN) (Pervez et al., ) has recently emerged as a promising approach in scientific machine learning. Unlike traditional black-box approaches for dynamical systems (Chen et al., , ; Kidger et al., ; Norcliffe et al., ) that primarily focus on forecasting, MNN additionally learns an explicit internal ordinary differential equation (ODE) representation from the noisy observational data that enables various downstream scientific analysis such as parameter identification and causal effect estimation (Yao et al., ).
Despite their advantages, the original formulation of MNN faces significant scalability challenges. Specifically, the computational time and memory usage severely limit the practical applicability of MNN to problems involving long time horizons or high-resolution temporal data, such as climate recordings (Verma et al., )
as the required computations become prohibitive even for the most advanced hardware.

Figure: Figure 1: 4-year sea surface temperature (SST) forecasting using the Mechanistic Identifier (Yao et al., ) and our *Scalable* Mechanistic Neural Network (S-MNN).
Refer to caption: x1.png

The inefficiency stems from the matrix operations required to solve the linear systems associated with the MNN. In the original framework, two solvers are provided: a dense solver and a sparse solver. The dense solver operates on dense matrices and employs standard methods for solving linear systems, resulting in cubic time and quadratic space complexities with respect to the sequence length. This computational inefficiency makes it unsuitable for long sequences. The MNN sparse solver, on the other hand, constructs sparse matrices and uses iterative methods such as the conjugate gradient algorithm to solve the linear systems. While this reduces memory usage by exploiting sparsity, the unstructured sparsity patterns of the matrices prevent the solver from fully leveraging the GPU’s parallelism potential. Additionally, iterative methods can suffer from slow convergence and numerical inaccuracies, particularly for large-scale problems.

This work proposes a *scalable* variant of MNN (S-MNN) that reduces the computational time and space complexities from cubic and quadratic with respect to the sequence length, respectively, to *linear*, while maintaining on-par accuracy. As a result, we successfully demonstrate the real-world applicability (Figure 1) of S-MNN on long-term climate data that the original MNN failed to handle. Our main contributions are as follows:

- •
Complexity Reduction. We reformulate the original MNN’s underlying linear system by eliminating the slack variables and central difference constraints, and reducing the quadratic programming problem to least squares regression (Section 3.1). This results in the left-hand side square matrix having a banded structure, allowing us to employ efficient algorithms (Section 3.2). The time and space complexities are reduced to *linear* with respect to the sequence length, making it suitable for long-sequence modeling.
- •
Efficient Solver Design. We develop an efficient solver that leverages the inherent sparsity and banded structure of the reformulated linear system (Section 3.2). The solver is optimized for GPU execution, fully exploiting parallelism to achieve significant speed-ups.
- •
Long-Term Sequence Modeling for Science. We validate the effectiveness of our S-MNN through experiments on various benchmarks, including governing equation discovery with the Lorenz system (Section 5.2), solving the Korteweg-de Vries (KdV) partial derivative equation (Section 5.3), and sea surface temperature (SST) prediction for modeling long real-world temporal sequences (Section 5.4). Our results demonstrate that S-MNN matches the precision of the original MNN while significantly reducing computational time and memory usage across the board.

## 2 Overview of Mechanistic Neural Networks

The Mechanistic Neural Network (MNN) by Pervez et al. () consists of three components: a mechanistic encoder, a specialized differentiable ordinary differential equation (ODE) solver based on constrained optimization, and a mechanistic decoder. The mechanistic encoder, realized as a neural network, generates a semi-symbolic representation of the underlying ODE from an input time series, effectively learning the dynamics from the data. The solver then constructs and solves a linear system that equivalently describes the original system. The mechanistic decoder processes the solver solutions to the final outputs.

For tasks such as forecasting future sequences from past data, the encoder processes the input data to produce the trajectory-specific semi-symbolic representation. Additionally, the encoder can be designed to overcome the limitations of the linear ODE solver by learning parameters of nonlinear basis functions (Brunton et al., ), enabling the MNN to model nonlinear ODEs effectively. The decoder is optional and their necessity depends on the specific task used for training.

Formally, given a multidimensional discretized time sequence $\bm{x}_{1},\bm{x}_{2},\dots,\bm{x}_{T}$ as the input, the mechanistic encoder maps $\bm{x}_{1:T}$ into semi-symbolic representations of a set of linear ODEs

$$ $\sum_{r=0}^{R}\bm{c}_{r}^{\top}\left(t,\bm{x}_{1:T}\right)\frac{\mathrm{d}^{r} \bm{y}}{\mathrm{d}t^{r}}=d\left(t,\bm{x}_{1:T}\right)$ (1) $$

of the variable $\bm{y}\in\mathbb{R}^{V}$ with the initial conditions

$$ $\frac{\mathrm{d}^{r}\bm{y}}{\mathrm{d}t^{r}}=\bm{u}_{r}\left(\bm{x}_{1:T}\right)$ (2) $$

where $R$ is the highest derivative order, $\bm{c}_{r}\left(t,\bm{x}_{1:T}\right)\in\mathbb{R}^{V}$ and $d\left(t,\bm{x}_{1:T}\right)\in\mathbb{R}$ are the coefficients, and $\bm{u}_{r}\left(\bm{x}_{1:T}\right)\in\mathbb{R}^{V}$ represents the initial conditions. $\bm{c}_{r}\left(t,\bm{x}_{1:T}\right)$, $d\left(t,\bm{x}_{1:T}\right)$, $\bm{u}_{r}\left(\bm{x}_{1:T}\right)$, as well as the step sizes ${s}\left(\bm{x}_{1:T}\right)$ for the time discretization are learned by the encoder. The representation $\theta$ is compactly denoted as

$$ $\theta\left(\bm{x}_{1:T}\right)=\left\{c\left(\bm{x}_{1:T}\right),d\left(\bm{x }_{1:T}\right),u\left(\bm{x}_{1:T}\right),s\left(\bm{x}_{1:T}\right)\right\}.$ (3) $$

The MNN solver then solves the linear ODE using this representation $\theta$, producing the discretized output time sequence $\bm{y}_{1},\bm{y}_{2},\dots,\bm{y}_{T}$. The decoder takes $\bm{y}_{1:T}$ as input and generates the final output $\bm{z}$. A loss $\ell$ can be computed based on $\bm{z}$ and the target data, and the encoder and decoder networks are updated using gradient descent methods.

An illustrative example is the task of discovering coefficients $\bm{\xi}\in\mathbb{R}^{k}$ for a time-independent one-dimensional first-order ODE in the form of

$$ $\mathrm{d}y/\mathrm{d}t=g\left(\bm{\xi}^{\top}\phi\left(y\right)\right)$ (4) $$

where $\phi:\mathbb{R}\rightarrow\mathbb{R}^{k}$ is a $k$ dimensional non-linear basis function, $g:\mathbb{R}\rightarrow\mathbb{R}$ is a differentiable function.
In this task, we first generate an initial $\bm{\xi}$ and specify the initial condition $u_{r=0}=x_{1}$. We fix $c$ to ones where they are multiplied with the first-order derivatives, and zeros otherwise. The step sizes $s$ are determined by the dataset’s time increments. During each gradient descent iteration, we set $d=g\left(\bm{\xi}^{\top}\phi\left(x\right)\right)$, and then the MNN solver solves for a discretized solution $y_{1:T}$ for Eq. 4. We update $\bm{\xi}$ by minimizing the loss $\ell=\sum_{t=1}^{T}\|y_{t}-x_{t}\|^{2}$.

## 3 Scalable Mechanistic Neural Networks

In this section, we define the linear system for the *Scalable* Mechanistic Neural Network (S-MNN) in Subsection 3.1, and present the solver’s implementation and complexity analysis in Subsection 3.2.

### 3.1 Linear System Formulation

A linear ordinary differential equation (ODE) system can be characterized by a set of linear equations involving the successive derivatives of an unknown time-dependent function $y$. Formally, in a system with $V$ variables (output dimensions), derivative orders up to $R$, and $Q$ governing equations, the state at $T$ discrete time points can be described by a series of clauses

$$ $\sum_{v=1}^{V}\sum_{r=0}^{R}c_{t,q,v,r}y_{t,v,r}=d_{t,q},\forall t\in\left\{1, \dots,T\right\},q\in\left\{1,\dots,Q\right\},$ (5) $$

where $y_{t,v,r}$ is the $r$-th derivative of the $v$-th variable at $t$-th time point, $c_{t,q,v,r}$ is the corresponding coefficient for the $q$-th governing equation, and $d_{t,q}$ is a constant term.

Initial values $u_{t,v,r}$ are specified for each $y_{t,v,r}$ up to time point $T_{\mathrm{init}}$ ($1\leq T_{\mathrm{init}}\leq T$) and derivative order $R_{\mathrm{init}}$ ($0\leq R_{\mathrm{init}}\leq R$):

$$ $y_{t,v,r}=u_{t,v,r},\forall t\in\left\{1,\dots,T_{\mathrm{init}}\right\},v\in \left\{1,\dots,V\right\},r\in\left\{0,\dots,R_{\mathrm{init}}\right\}.$ (6) $$

To ensure the smoothness of the trajectory, i.e., that the computed higher-order derivatives are consistent with the derivatives of lower-order terms, we introduce smoothness constraints. We define the forward and backward smoothness constraints using the Taylor expansion of the function $y$ at time points $t$ and $t+1$ respectively:

$$ $y_{t+1,v,r}=\sum_{r^{\prime}=r}^{R}\frac{s_{t}^{r^{\prime}-r}}{\left(r^{\prime }-r\right)!}y_{t,v,r^{\prime}},\forall t\in\left\{1,\dots,T-1\right\},v\in \left\{1,\dots,V\right\},r\in\left\{0,\dots,R\right\},$ (7) $$

$$ $y_{t,v,r}=\sum_{r^{\prime}=r}^{R}\frac{\left(-s_{t}\right)^{r^{\prime}-r}}{ \left(r^{\prime}-r\right)!}y_{t+1,v,r^{\prime}},\forall t\in\left\{1,\dots,T-1 \right\},v\in\left\{1,\dots,V\right\},r\in\left\{0,\dots,R\right\}.$ (8) $$

where $s_{t}$ is the time span between time points $t$ and $t+1$.

Combining the constraints Eqs. 5, 6, 7, and 8 yields a linear system. In total there are $m$ constraints and $n$ unknown variables where

$$ $m=TQ+T_{\mathrm{init}}V\left(R_{\mathrm{init}}+1\right)+2\left(T-1\right)V \left(R+1\right)\textnormal{ and }n=TV\left(R+1\right).$ (9) $$

We arrange the unknown variables $y_{t,v,r}$ into a vector $\bm{y}\in\mathbb{R}^{n}$, the left-hand side coefficients into a matrix $\bm{A}\in\mathbb{R}^{m\times n}$, and the right-hand side into a vector $\bm{b}\in\mathbb{R}^{m}$. The linear system can then be compactly represented as $\bm{A}\bm{y}=\bm{b}$.

This system is over-determined ($m>n$) under typical conditions ($T>1$) and can be solved for $\bm{y}$ using least squares regression. To balance the contributions of different constraints, we weight the smoothness constraints (Eqs. 7 and 8) by $s_{t}^{r}$. Additionally, we introduce optional importance weights $w_{\mathrm{gov}},w_{\mathrm{init}},w_{\mathrm{smooth}}\in\mathbb{R}$, applied to the governing equations (Eq. 5), initial conditions (Eq. 6), and smoothness constraints (Eqs. 7 and 8), respectively. This flexibility allows for emphasizing specific aspects of the model during optimization. These weights are encoded into a diagonal matrix $\bm{W}$, and the solution for $\bm{y}$ is then given by:

$$ $\bm{y}\left(c,d,u,s\right)=\left(\bm{A}^{\top}\bm{W}\bm{A}\right)^{-1}\bm{A}^{ \top}\bm{W}\bm{b}.$ (10) $$

Note that $\bm{y}$ is differentiable with respect to $c$, $d$, $u$, and $s$.

Our S-MNN formulation has three key differences from the original MNN formulation in Pervez et al. (): (1) the slack variables in the smoothness constraints are removed; (2) the forward and backward smoothness constraints are extended to the highest order ($r=R$), replacing the central difference constraints; (3) the quadratic programming is replaced by a least squares regression.

Also note that, unlike the finite difference method which approximates derivatives by discretizing differential equations, our approach formulates a linear system to directly involve the derivative terms.

### 3.2 Solver Design, Architecture, and Analysis

The primary motivation for our improvement is our observation that, if we remove the slack variables, the matrix $\bm{A}$ will exhibit a specific sparsity pattern that can be exploited for computational and memory gains. A direct implementation based on the dense matrix $\bm{A}$ using Eq. 10 incurs cubic time complexity and quadratic space complexity due to matrix multiplication and inversion. However, by analyzing the sparsity pattern of $\bm{A}$, we find that most (e.g., 95% in a typical application) values in the intermediate steps are zero and do not need to be computed or stored. In more detail, Eq. 10 can be decomposed into two steps: (1) a matrix-matrix multiplication $\bm{M}=\left(\bm{A}^{\top}\bm{W}\right)\bm{A}$ and a matrix-vector multiplication $\bm{\beta}=\left(\bm{A}^{\top}\bm{W}\right)\bm{b}$; (2) solving for $\bm{y}$ via the linear system $\bm{M}\bm{y}=\bm{\beta}$. The key idea is to structure $\bm{M}$ as a banded matrix. Observing the constraints, we note that the coefficients at time point $t$ are directly related only to those at time points $t-1$, $t$, and $t+1$. By ordering the variables by $t$, we can arrange $\bm{M}$ into a banded matrix. Specifically, the variable $y_{t,v,r}$ is placed at the position $\left(\left(t-1\right)V+v-1\right)\left(R+1\right)+r+1$ in the vector $\bm{y}$, and the columns of matrix $\bm{A}$ are ordered accordingly. The resulting $\bm{M}$ is a symmetric matrix in a block-banded form:

$$ $\bm{M}=\left[\begin{matrix}\bm{M}_{1}&\bm{N}_{1}^{\top}&&{\color[rgb]{1,1,1} \ddots}\\ \bm{N}_{1}&\bm{M}_{2}&\ddots&\\ &\ddots&\ddots&\bm{N}_{T-1}^{\top}\\ {\color[rgb]{1,1,1}\ddots}&&\bm{N}_{T-1}&\bm{M}_{T}\\ \end{matrix}\right]$ (11) $$

where each block is a square matrix of size $V\left(R+1\right)$. It is important to note that such a banded form is only possible after removing the slack variables from the original MNN formulation, as the slack variables introduce direct relationships between components at all time points.

For the matrix $\bm{M}$, only the non-zero blocks $\bm{M}_{t}$ and $\bm{N}_{t}$ need to be computed and stored. This can be achieved using efficient dense matrix operations with appropriately formatted dependencies $c$, $d$, $u$, and $s$. The matrix $\bm{A}$ does not need to be explicitly constructed. We present the detailed forward pass calculations for computing $\bm{M}$ and $\bm{\beta}$ in Appendix A.1 and omit them in this subsection. The backward pass is supported by automatic differentiation.

For solving the linear system $\bm{M}\bm{y}=\bm{\beta}$, we propose an efficient GPU-friendly algorithm (Algorithm 1). Specifically, we can factorize the positive-definite matrix $\bm{M}$ to a banded lower triangular matrix $\bm{P}$ and a block diagonal lower triangular matrix $\bm{L}$ using blocked versions of LDL and Cholesky decompositions such that

$$ $\bm{P}\bm{L}\bm{L}^{\top}\bm{P}^{\top}=\bm{M}$ (12) $$

where $\bm{P}$ and $\bm{L}$ are partitioned into block matrices in the same form as $\bm{M}$:

$$ $\bm{P}=\left[\begin{matrix}\bm{I}&&&{\color[rgb]{1,1,1}\ddots}\\ \bm{P}_{1}&\bm{I}&{\color[rgb]{1,1,1}\ddots}&\\ &\ddots&\ddots&\\ {\color[rgb]{1,1,1}\ddots}&&\bm{P}_{T-1}&\bm{I}\\ \end{matrix}\right]\ \mathrm{and}\ \bm{L}=\left[\begin{matrix}\bm{L}_{1}&&{ \color[rgb]{1,1,1}\ddots}&\\ &\bm{L}_{2}&{\color[rgb]{1,1,1}\ddots}&\\ &&\ddots&\\ &&{\color[rgb]{1,1,1}\ddots}&\bm{L}_{T}\\ \end{matrix}\right].$ (13) $$

with each block as a square matrix of size $V\left(R+1\right)$ and the diagonal blocks of $\bm{P}$ as an identity matrix $\bm{I}$.
We also partition $\bm{\beta}$ and $\bm{y}$ into sub-vectors, $\bm{\beta}=\left[\bm{\beta}_{1}^{\top},\bm{\beta}_{2}^{\top},\dots,\bm{\beta}_
{T}^{\top}\right]^{\top}$ and $\bm{y}=\left[\bm{y}_{1}^{\top},\bm{y}_{2}^{\top},\dots,\bm{y}_{T}^{\top}\right
]^{\top}$.
We then use forward and backward substitution to solve the following sequence of equations:

$$ $\bm{P}\bm{\beta}^{\prime}=\bm{\beta},\ \bm{L}\bm{\beta}^{\prime\prime}=\bm{ \beta}^{\prime},\ \bm{L}^{\top}\bm{\beta}^{\prime\prime\prime}=\bm{\beta}^{ \prime\prime},\ \bm{P}^{\top}\bm{y}=\bm{\beta}^{\prime\prime\prime}.$ (14) $$

The main advantage of this algorithm is that only the blocks $\bm{P}_{t}$ and $\bm{L}_{t}$ in the matrices $\bm{P}$ and $\bm{L}$ need to be computed and stored, which limits the computational complexities of the LDL and Cholesky decompositions and the subsequent forward and backward substitutions to be linear in $T$. In contrast, explicitly computing $\bm{M}^{-1}$ or $\bm{L}^{-1}$ would involve dense or triangular-dense full-size matrices and should be avoided.

Interestingly, for the backward pass of $\bm{M}\bm{y}=\bm{\beta}$, the back-propagated gradients of $\bm{M}$ and $\bm{\beta}$ have elegant analytic solutions that enable further (constant factor) speed-ups compared to automatic differentiation. Assuming that the final loss is $\ell$, and given $\partial\ell/\partial\bm{y}$, the gradients are

$$ $\frac{\partial\ell}{\partial\bm{\beta}}=\bm{M}^{-1}\frac{\partial\ell}{ \partial\bm{y}}\ \ \mathrm{and}\ \ \frac{\partial\ell}{\partial\bm{M}}=-\frac{ \partial\ell}{\partial\bm{\beta}}\bm{y}^{\top}.$ (15) $$

The proof can be found in Appendix A.2. Algorithm 3 details the backward pass. Computing $\partial\ell/\partial\bm{\beta}$ involves solving a similar linear system using the banded matrix $\bm{M}$, and $\partial\ell/\partial\bm{M}$ is a vector outer product that can be efficiently computed for the non-zero blocks $\bm{M}_{t}$ and $\bm{N}_{t}$. The matrices $\bm{P}$, $\bm{L}$, and the solution $\bm{y}$ can be stored during the forward pass and reused for the backward pass.

Figure: Algorithm 1 Solver Forward Pass

Numerical Stability Considerations.
An important aspect of our solver design is the numerical stability offered by the direct method of Cholesky decomposition compared to iterative methods like the conjugate gradient (CG) algorithm. Both direct and iterative methods have errors influenced by the condition number $\kappa$ of the matrix. However, direct methods tend to be more stable in practice because they compute an exact solution up to machine precision. In contrast, CG is iterative and can suffer from error accumulation across iterations, especially when the matrix is ill-conditioned or when the number of iterations is limited for computational reasons. This is crucial in our applications, where accurate solutions are necessary for modeling chaotic systems.

Complexity Analysis.
The original MNN dense (exact) solver described in Pervez et al. () has time complexity $O\left(T^{3}V^{3}R^{3}\right)$ and space complexity $O\left(T^{2}V^{2}R^{2}\right)$, for time $T$, $V$ variables, and $R$ derivative orders. The original sparse (approximate) solver working via conjugate gradient has both time and space complexities $O\left(T^{2}V^{2}R^{2}\right)$. In our proposed exact S-MNN, the time complexity is reduced by a $\Theta(T^{2})$ factor to $O\left(TV^{3}R^{3}\right)$ from the MNN dense solver, and the space complexity is reduced to $O\left(TV^{2}R^{2}\right)$. Thus, both time and memory requirements now depend *linearly* on the number of time points $T$, making the new method more scalable for longer trajectories.

## 4 Related Work in Scientific Machine Learning

Scientific machine learning has emerged as a transformative field that combines data-driven approaches with domain-specific knowledge to model complex dynamical systems. Various specialized methodologies have been developed to tackle different aspects of this challenge, particularly in solving differential equations using neural networks and, to a lesser extent, inverse problems.

Models for Prediction.
Neural Ordinary Differential Equations (Neural ODEs) (Chen et al., , ; Kidger et al., ; Norcliffe et al., ) model continuous-time dynamics by parameterizing the derivative of the hidden state with a neural network.
Neural ODEs, however, are constrained by the structure and sequential nature of ODE solvers and can be inefficient to train.
Neural Operators (Li et al., , , , ; Azizzadenesheli et al., ; Boullé & Townsend, ) are designed to learn mappings between infinite-dimensional function spaces, enabling the modeling of PDEs and complex spatial-temporal patterns.
However, their focus on lower frequencies in the Fourier spectrum can lead to poor prediction over longer roll-outs (Lippe et al., ).

Models for Discovery.
The line of work on Sparse Identification of Nonlinear Dynamical Systems (SINDy) (Kaheman et al., ; Brunton et al., , ; Kaptanoglu et al., ; Course & Nair, ; Lu et al., ; Rudy et al., ) aims to discover governing equations by identifying sparse representations within a predefined library of candidate functions. However, SINDy is only a generalized linear model that does not use neural networks and can struggle with highly complex, noisy, or strongly nonlinear systems.
Physics-informed networks and universal differential equations (Raissi et al., ; Rackauckas et al., ) also work as discovery methods for inferring unknown terms in PDEs.
Symbolic regression methods (Udrescu & Tegmark, ; d’Ascoli et al., ) constitute another line of work that aims to discover purely symbolic expressions from data.

Discussion.
Mechanistic Neural Networks (MNNs) (Pervez et al., ) have been proposed as a single framework for prediction and discovery.
MNNs compute ODE representations from data which provide a strong inductive bias for scientific ML tasks.
However, MNN training introduces significant challenges that require solving large linear systems during both the forward and backward passes and demands substantial computational resources. Our method addresses this scalability problem by reducing the computational complexity and enables applications on long sequences.

## 5 Experiments

To demonstrate the effectiveness and scalability of our proposed *Scalable* Mechanistic Neural Network (S-MNN), we conduct experiments across multiple settings in scientific machine learning applications for dynamical systems including governing equation discovery for the Lorenz system (Section 5.2), solving the Korteweg-de Vries (KdV) partial derivative equation (PDE) (Section 5.3), and sea surface temperature (SST) prediction for modeling long real-world temporal sequences (Section 5.4). We show that S-MNN matches the precisions and convergence rates of the original MNN (Pervez et al., ) while significantly reducing computational time and GPU memory usage.

### 5.1 Standalone Validation

To assess the correctness of our solver in solving linear ordinary differential equations (ODEs), we conducted a standalone validation. Our solver is designed to solve linear ODEs directly (Section 3.1) without incorporating additional neural network layers or trainable parameters.

Experiment Settings.
We selected five linear ODE problems from ODEBench (d’Ascoli et al., )—RC Circuit, Population Growth, Language Death Model, Harmonic Oscillator, and Harmonic Oscillator with Damping—that are commonly used in various scientific fields such as physics and biology, along with an additional third-order ODE. Mathematical details about these ODEs are provided in the Appendix B.1. For each problem, we discretized the time axis into 1,000 steps with a uniform step size of 0.01 and applied our S-MNN solver.

Results and Discussion.
We compared the numerical solutions obtained by our solver against the corresponding closed-form solutions. Figure 2 presents the results, where we plot the solutions $y(t)$ along with their first and second derivatives $y^{\prime}(t)$ and $y^{\prime\prime}(t)$. The numerical results from our solver closely match the analytical solutions, exhibiting negligible differences (mean squared error $<10^{-6}$ in all cases, $<10^{-8}$ in most cases). These results confirm that our solver is capable of correctly solving linear ODEs.

Figure: Figure 2: Standalone S-MNN solver validation results compared with the closed-form solutions.
Refer to caption: x2.png

### 5.2 Comparative Analysis: Discovery of Governing Equations

In this experiment, we evaluate the capability of our S-MNN in discovering the coefficients of the governing equations for the Lorenz system following Section 5.1 in the origin MNN paper (Pervez et al., ).
The Lorenz system is a set of nonlinear ODES known for its chaotic behavior, making it a standard benchmark for testing equation discovery methods in dynamical systems.
The governing equations are given by

$$ $\begin{cases}\mathrm{d}x/\mathrm{d}t=\sigma\left(y-x\right)=a_{1}x+a_{2}y\\ \mathrm{d}y/\mathrm{d}t=x\left(\rho-z\right)-y=a_{3}x+a_{4}y+a_{5}xz\\ \mathrm{d}z/\mathrm{d}t=xy-\beta z=a_{6}z+a_{7}xy\end{cases}$ (16) $$

where $a_{1},\dots,a_{7}\in\mathbb{R}$ are the coefficients to be discovered.

Figure: Figure 3: Lorenz discovery loss over first 1,000 optimization steps (exponential moving average factor = 0.9) using S-MNN (ours) compared with the original MNN dense and sparse solvers (Pervez et al., ).
Refer to caption: x3.png

Dataset.
The dataset was generated by numerically integrating the Lorenz system equations using the standard parameters $\sigma=10$, $\rho=28$, and $\beta=8/3$. We used the initial condition $x=y=z=1$ and integrated over 10,000 time steps with a step size of 0.01 using the scipy.integrate.odeint function.

Experiment Settings.
We applied our solver to the same network architecture and dataset in Pervez et al. (). We trained the model with the default settings: sequence length of 50 and batch size of 512. The training objective was to minimize the difference between the predicted and true trajectories by optimizing the coefficients $a_{1},\dots,a_{7}$. Then, to assess the scalability of our method, we measured runtime and GPU memory consumption across different sequence lengths and batch sizes using an NVIDIA H100 GPU with 80 GB of memory.

Results and Discussion.
Figure 3 shows the convergence of the loss over the optimization steps using our S-MNN solver compared to the original MNN dense and sparse solvers. The results indicate that our S-MNN solver achieves similar convergence rates, confirming that the removal of slack variables does not adversely affect accuracy. The final results of the discovered coefficients $a_{0},...,a_{7}$ are recorded in Appendix B.2.

Table 1 summarizes the performance and GPU memory usage. Our S-MNN solver not only maintains high accuracy but also offers substantial efficiency improvements. Specifically, compared to the MNN dense solver, our method achieves a 4.9$\times$ speedup and reduces GPU memory usage by 50% for the default setting (batch size 512, sequence length 50).
The more significant improvement in runtime compared to memory usage is expected, as our approach reduces runtime from $O\left(T^{3}\right)$ to $O\left(T\right)$, and memory from $O\left(T^{2}\right)$ to $O\left(T\right)$, with $T$ denoting the sequence length.
Our S-MNN solver maintains high performance even with larger batch sizes and sequence lengths where the MNN solvers run out of memory or become computationally infeasible.

**Table 1: Performance and GPU memory usage comparison for the Lorenz discovery experiment.**
| Batch Size | 512 | 512 | 512 | 64 | 4096 |  |
| --- | --- | --- | --- | --- | --- | --- |
| Sequence Length | 50 | 5 | 500 | 50 | 50 |  |
| Time per Optimization Step [ms] | MNN Dense | 36.4 | 9.5 | N/A(^1^1footnotemark: 1) | 10.0 | 208.1 |
| MNN Sparse | 104.4 | 76.5 | >589.7(^2^2footnotemark: 2) | 80.5 | 406.8 |  |
| S-MNN | 7.4 | 5.5 | 32.2 | 5.5 | 18.3 |  |
| GPU Memory Usage [GiB] | MNN Dense | 2.77 | 1.18 | >80.00(^1^1footnotemark: 1) | 1.33 | 14.85 |
| MNN Sparse | 1.69 | 0.93 | 9.83(^2^2footnotemark: 2) | 0.96 | 7.96 |  |
| S-MNN | 1.38 | 1.32 | 1.96 | 1.33 | 1.81 |  |
| (^1^1footnotemark: 1)Out of memory error. (^2^2footnotemark: 2)Loss does not converge after a large number (200) of conjugate gradient iterations. |  |  |  |  |  |  |

### 5.3 Comparative Analysis: Solving Partial Differential Equations (PDEs)

Next, we evaluate the capability of our S-MNN in solving partial differential equations (PDE), specifically focusing on the Korteweg-De Vries (KdV) equation, which is a third-order nonlinear PDE that describes the propagation of waves in shallow water and is expressed as

$$ $\frac{\partial y}{\partial t}+\frac{\partial^{3}y}{\partial x^{3}}-6y\frac{ \partial y}{\partial x}=0,$ (17) $$

where $y\left(x,t\right)$ represents the wave amplitude as a function of spatial coordinate $x$ and time $t$. Solving the KdV equation is challenging due to its nonlinearity and the involvement of higher-order spatial derivatives, making it a popular benchmark for PDEs.

Dataset.
We consider the KdV dataset provided by Brandstetter et al. (). The dataset consists of 512 samples each for training, validation, and testing. Each sample has a spatial domain of 256 meters and a temporal domain of 140 seconds, discretized into 256 spatial points and 140 temporal steps.

Experiment Settings.
We model the temporal evolution at each spatial point as an independent ODE. A ResNet-1D architecture (Brandstetter et al., ) is employed to encode the temporal and spatial dependencies in the input sequences and feed them into the mechanistic solver. The sequence length is set to 10 seconds, and the model is trained to predict the wave profile over the next 9 seconds. The model is trained for 800 epochs. We repeat the same experiment for the original MNN dense and sparse solvers as well as our S-MNN solver.

Results and Discussion.
Figure 4 illustrates a comparison between the ground truth solution, the prediction obtained from the original MNN dense solver, and that from our S-MNN. Both models produce predictions that closely align with the ground truth, demonstrating that our S-MNN effectively captures the intricate dynamics of the KdV equation. The MNN sparse solver cannot converge in this experiment and its result is not shown.

Figure: Figure 4: Visual comparison between the ground truth, the MNN dense solver (Pervez et al., ) prediction, and our S-MNN prediction for the KdV equation.
Refer to caption: x4.png

Table 2 presents a summary of the testing error, end-to-end training time, and GPU memory usage for each method. Our S-MNN achieves an RMSE (relative mean squared error) of 0.0042, slightly improving upon the 0.0046 of the MNN dense solver. In terms of efficiency, the training time for S-MNN is significantly reduced to 10.1 hours compared to 38.0 hours for the MNN dense solver, indicating a substantial speedup. Additionally, our method consumes less GPU memory, using 2.19 GiB versus 3.40 GiB for the original solver. The MNN sparse solver did not converge within a reasonable time frame, taking 82.4 hours without achieving satisfactory results.

**Table 2: Testing error, end-to-end training time, and GPU training memory usage for KdV.**
| Method | RMSE | Training Time [hour] | GPU Memory [GiB] |
| --- | --- | --- | --- |
| MNN Dense | 0.0046 | 38.0 | 3.40 |
| MNN Sparse | Did Not Converge | 82.4 | 3.07 |
| S-MNN | 0.0042 | 10.1 | 2.19 |

### 5.4 Real-world Application: Long-term Sea Surface Temperature Forecasting

Figure: Figure 5: Error visualization for the S-MNN 4-year sea surface temperature (SST) prediction.
Refer to caption: x5.png

The ability to handle longer sequences and larger batch sizes without sacrificing performance positions our S-MNN as a powerful tool for modeling complex dynamical systems. In this section, we demonstrate a real-world example use case: sea surface temperature (SST) prediction. SST exhibits long periodic features that can only be effectively captured with long sequences.

Dataset.
We use the SST-V2 dataset (Huang et al., ), which provides weekly mean sea surface temperatures for 1,727 weeks from December 31, 1989, to January 29, 2023, over a 1° latitude $\times$ 1° longitude global grid (180 $\times$ 360).

Experiment Settings.
We employ the Mechanistic Identifier proposed by Yao et al. () to predict SST data. The model leverages mechanistic layers to capture the underlying dynamics of SST. We set the default batch size to 12,960 (corresponding to 6,480 pairs of grid points and their randomly selected neighboring points) and the sequence length (chunk length) to 208 weeks. The dataset is split so that the latest chunk of measurements is reserved for testing while the remaining data is used for training. The model is trained for 1,000 epochs. To evaluate the scalability and stability, we benchmark the model with different sequence lengths.

Results and Discussion.
Figure 1 visualizes the 4-year (208 weeks) prediction made by our S-MNN using the Mechanistic Identifier. Figure 5 visualizes its prediction error over the ground truth. The S-MNN effectively captures the spatial patterns of SST, demonstrating high predictive precision.

To quantitatively assess the performance and scalability of our S-MNN, we measure the accuracy in terms of the mean squared error (MSE) over the standardized data, as well as the runtime and GPU memory usage across different sequence lengths. For comparison, we also perform the experiment using the MNN dense solver for smaller sequences. Figure 6 summarizes these results.

Figure: Figure 6: Testing error, training runtime, and GPU training memory usage comparisons between S-MNN (ours) and MNN *dense* solver (Pervez et al., ) for SST forecasting. Note that the x-axis is in log scale, so both runtime and memory consumption increase linearly as expected.
Refer to caption: x6.png

We observe that the MSE (averaged over both sequence length and batch size) remains low for both the MNN dense and S-MNN solvers, with our S-MNN maintaining low MSE even for longer sequences where the MNN dense solver cannot run due to resource limitations.
A slight increase in MSE for longer sequences is expected, as modeling longer-term dependencies introduces increased complexity, and the error accumulates over extended prediction horizons.
Note that for the sequence length of 1,727, we used the entire dataset for training, leaving no separate testing dataset for evaluation, which is why no MSE results can be provided in Figure 6.
Given the log-scale in the x-axis, we observe that S-MNN demonstrates a linear increase in both runtime and memory consumption with respect to sequence length, aligning with our theoretical complexities for time and space.
In contrast, the MNN dense solver exhibits a much steeper increase in memory usage due to its quadratic space complexity, quickly exceeding the 80 GiB limit of our GPU for longer sequences.
We were unable to run it for sequence lengths beyond 104 weeks.
We also tested the MNN sparse solver, which was too slow to complete the experiment in any reasonable time.

## 6 Conclusion

This paper introduced the *Scalable* Mechanistic Neural Network (S-MNN), addressing the scalability limitations of the original Mechanistic Neural Networks (Pervez et al., ) (MNN) by reducing computational complexities to linear in sequence length for both time and space. This was achieved by eliminating slack variables and central difference constraints, transitioning from quadratic programming to least squares regression, and exploiting banded matrix structures within the solver. Our experiments demonstrated that S-MNN retains the precision of the original MNN while significantly enhancing computational efficiency. Given these substantial advantages, S-MNN can drop-in replace the original MNN. We believe this advancement can provide a practical and efficient method for embedding mechanistic knowledge into neural network models for complex dynamical systems.

Limitations and Future Work.
While our S-MNN significantly enhances scalability, certain components, such as the sequential for-loops in the Cholesky decomposition (Algorithm 2) and the forward/backward substitution steps (Algorithm 4), still limit parallelism along the time dimension due to inherent data dependencies. This sequential execution can become a bottleneck when the batch and block sizes are small compared to the number of time steps, leading to underutilization of GPU resources and increased CPU overhead from launching small GPU kernels. Although we have employed CUDA Graphs to reduce this overhead, the fundamental sequential nature of the algorithms remains unaddressed. For future work, we aim to develop algorithms that retain linear time and space complexities but introduce full parallelism also across the time dimension.

## Appendix A Theoretical Derivations

### A.1 Compute 𝑴 𝑴 \bm{M} bold_italic_M and 𝜷 𝜷 \bm{\beta} bold_italic_β

Define matrix $\bm{C}_{t}\in\mathbb{R}^{Q\times V\left(R+1\right)}$ such that $\left[\bm{C}_{t}\right]_{q,i}=c_{t,q,v,r}$ where $v=\left\lfloor\left(i-1\right)/\left(R+1\right)\right\rfloor+1$ and $r=\left(i-1\right)\ \mathrm{mod}\ \left(R+1\right)$.

Define vector $\bm{d}_{t}=\left[d_{t,1},d_{t,2},\dots,d_{t,Q}\right]^{\top}\in\mathbb{R}^{Q}$.

Define constant matrix $\bm{U}_{t}\in\left\{0,1\right\}^{V\left(R+1\right)\times V\left(R+1\right)}$ such that

$$ $\left[\bm{U}_{t}\right]_{i,j}=\begin{cases}1&\mathrm{if}\ t\leq T_{\mathrm{ init}}\ \mathrm{and}\ \left(\left(i-1\right)\ \mathrm{mod}\ \left(R+1\right) \right)\leq R_{\mathrm{init}}\ \mathrm{and}\ i=j,\\ 0&\mathrm{otherwise}.\end{cases}$ (18) $$

Define vector $\bm{u}_{t}\in\mathbb{R}^{V\left(R+1\right)}$ such that

$$ $\left[\bm{u}_{t}\right]_{i}=\begin{cases}u_{t,v,r}&\mathrm{if}\ t\leq T_{ \mathrm{init}}\ \mathrm{and}\ r\leq R_{\mathrm{init}},\\ 0&\mathrm{otherwise},\end{cases}$ (19) $$

where $v=\left\lfloor\left(i-1\right)/\left(R+1\right)\right\rfloor+1$ and $r=\left(i-1\right)\ \mathrm{mod}\ \left(R+1\right)$.

Define constant matrix $\bm{F}\in\mathbb{R}^{\left(R+1\right)\times\left(R+1\right)}$ such that

$$ $\left[\bm{F}\right]_{i,j}=\begin{cases}0&\mathrm{if}\ i>j,\\ 1/\left(j-i\right)!&\mathrm{otherwise}.\end{cases}$ (20) $$

Define matrix $\bm{S}_{t}^{+}=\mathrm{diag}\left(s_{t}^{0},s_{t}^{1},\dots,s_{t}^{R}\right)
\in\mathbb{R}^{\left(R+1\right)\times\left(R+1\right)}$.

Define matrix $\bm{S}_{t}^{-}=\mathrm{diag}\left(\left(-s_{t}\right)^{0},\left(-s_{t}\right)^
{1},\dots,\left(-s_{t}\right)^{R}\right)\in\mathbb{R}^{\left(R+1\right)\times
\left(R+1\right)}$.

Define matrix $\bm{S}_{t}^{2}=\mathrm{diag}\left(s_{t}^{0},s_{t}^{2},\dots,s_{t}^{2R}\right)
\in\mathbb{R}^{\left(R+1\right)\times\left(R+1\right)}$.

Define matrix $\bm{S}_{t}^{*}\in\mathbb{R}^{\left(R+1\right)\times\left(R+1\right)}$ such that

$$ $\bm{S}_{t}^{*}=\begin{cases}\left(\bm{S}_{1}^{+}\right)^{\top}\bm{F}^{\top}\bm {F}\bm{S}_{1}^{+}+\bm{S}_{1}^{2}&\mathrm{if}\ t=1,\\ \left(\bm{S}_{T-1}^{-}\right)^{\top}\bm{F}^{\top}\bm{F}\bm{S}_{T-1}^{-}+\bm{S} _{T-1}^{2}&\mathrm{if}\ t=T,\\ \left(\bm{S}_{t}^{+}\right)^{\top}\bm{F}^{\top}\bm{F}\bm{S}_{t}^{+}+\left(\bm{ S}_{t-1}^{-}\right)^{\top}\bm{F}^{\top}\bm{F}\bm{S}_{t-1}^{-}+\bm{S}_{t}^{2}+ \bm{S}_{t-1}^{2}&\mathrm{otherwise}.\end{cases}$ (21) $$

Define matrix $\bm{S}_{t}^{**}=-\left(\bm{S}_{t}^{+}\right)^{\top}\bm{F}\bm{S}_{t}^{+}-\left(
\bm{S}_{t}^{-}\right)^{\top}\bm{F}^{\top}\bm{S}_{t}^{-}\in\mathbb{R}^{\left(R+
1\right)\times\left(R+1\right)}$ .

Define block diagonal matrix $\bm{S}_{t}=\mathrm{Diag}(\underbrace{\bm{S}_{t}^{*},\dots,\bm{S}_{t}^{*}}_{V\
\mathrm{times}})\in\mathbb{R}^{V\left(R+1\right)\times V\left(R+1\right)}$.

Then, $\bm{M}_{t,t}$, $\bm{M}_{t+1,t}$, and $\bm{\beta}_{t}$ can be computed as

$$ $\bm{M}_{t,t}=w_{\mathrm{gov}}^{2}\bm{C}_{t}^{\top}\bm{C}_{t}+w_{\mathrm{init}} ^{2}\bm{U}_{t}+w_{\mathrm{smooth}}^{2}\bm{S}_{t},\forall t\in\left\{1,\dots T \right\},$ (22) $$

$$ $\bm{M}_{t+1,t}=w_{\mathrm{smooth}}^{2}\mathrm{Diag}(\underbrace{\bm{S}_{t}^{** },\dots,\bm{S}_{t}^{**}}_{V\ \mathrm{times}}),\forall t\in\left\{1,\dots T-1 \right\},$ (23) $$

$$ $\bm{\beta}_{t}=w_{\mathrm{gov}}^{2}\bm{C}_{t}^{\top}\bm{d}_{t}+w_{\mathrm{init }}^{2}\bm{u}_{t},\forall t\in\left\{1,\dots T\right\}.$ (24) $$

### A.2 Gradients of 𝑴 𝑴 \bm{M} bold_italic_M and 𝜷 𝜷 \bm{\beta} bold_italic_β

In this proof, we use the subscripts $h,i,j,k$ to denote the element index, e.g., $y_{k}$ is the $k$-th component of vector $\bm{y}$, $M_{i,j}$ is the element in the $i$-th row and $j$-th column of matrix $\bm{M}$, and $\bm{M}_{i,:}$ is the row vector in the $i$-th row of $\bm{M}$.

Rewrite $\bm{y}=\bm{M}^{-1}\bm{\beta}$ as

$$ $y_{k}=\sum_{h}\left[M^{-1}\right]_{k,h}\beta_{h}.$ (25) $$

Differentiate $y_{k}$ with respect to $\beta_{i}$,

$$ $\frac{\partial y_{k}}{\partial\beta_{i}}=\left[M^{-1}\right]_{k,i}.$ (26) $$

Differentiate $l$ with respect to $\beta_{i}$ using the chain rule,

$$ $\frac{\partial\ell}{\partial\beta_{i}}=\sum_{k}\frac{\partial\ell}{\partial y_ {k}}\frac{\partial y_{k}}{\partial\beta_{i}}=\sum_{k}\left[M^{-1}\right]_{k,i} \frac{\partial\ell}{\partial y_{k}}=\left[\bm{M}^{-\top}\right]_{i,:}\frac{ \partial\ell}{\partial\bm{y}}.$ (27) $$

Express the derivative in vector form,

$$ $\frac{\partial\ell}{\partial\bm{\beta}}=\bm{M}^{-\top}\frac{\partial\ell}{ \partial\bm{y}}=\bm{M}^{-1}\frac{\partial\ell}{\partial\bm{y}}.$ (28) $$

Differentiate $y_{k}$ with respect to $M_{i,j}$,

$$ $\frac{\partial y_{k}}{\partial M_{i,j}}=\frac{\partial\left(\sum_{h}\left[M^{- 1}\right]_{k,h}\beta_{h}\right)}{\partial M_{i,j}}=\sum_{h}\beta_{h}\frac{ \partial\left[M^{-1}\right]_{k,h}}{\partial M_{i,j}}.$ (29) $$

Differentiate of $\left[M^{-1}\right]_{k,h}$ with respect to $M_{i,j}$,

$$ $\frac{\partial\left[M^{-1}\right]_{k,h}}{\partial M_{i,j}}=-\left[M^{-1}\right ]_{k,i}\left[M^{-1}\right]_{j,h}.$ (30) $$

Substitute the derivative back into the expression,

$$ $\frac{\partial y_{k}}{\partial M_{i,j}}=-\left[M^{-1}\right]_{k,i}\sum_{h} \left[M^{-1}\right]_{j,h}\beta_{h}=-\frac{\partial y_{k}}{\partial\beta_{i}}y_ {j}.$ (31) $$

Differentiate $l$ with respect to $M_{i,j}$ using the chain rule,

$$ $\frac{\partial\ell}{\partial M_{i,j}}=\sum_{k}\frac{\partial\ell}{\partial y_{ k}}\frac{\partial y_{k}}{\partial M_{i,j}}=-y_{j}\sum_{k}\frac{\partial\ell}{ \partial y_{k}}\frac{\partial y_{k}}{\partial\beta_{i}}=-y_{j}\frac{\partial \ell}{\partial\beta_{i}}.$ (32) $$

Express the derivative in matrix form,

$$ $\frac{\partial\ell}{\partial\bm{M}}=-\frac{\partial\ell}{\partial\bm{\beta}} \bm{y}^{\top}.$ (33) $$

## Appendix B Further Experimental Details

### B.1 Standalone Validation

We list the linear ODEs used for the validation experiment and their closed-form solutions. There are five ODEs from ODEBench (d’Ascoli et al., ) and an additional third-order ODE. $c_{0}$, $c_{1}$, $c_{2}$ are constant numbers. $u_{0}=y\left(0\right)$, $u_{0}=y\left(0\right)$, $u_{1}=y^{\prime}\left(0\right)$, $u_{2}=y^{\prime\prime}\left(0\right)$ are initial values.

RC-circuit (charging capacitor), $\left(c_{0},c_{1},c_{2}\right)=\left(0.7,1.2,2.31\right)$, $\left(u_{0}\right)=\left(10\right)$,

$$ $\frac{y}{c_{1}}+c_{2}\frac{\mathrm{d}y}{\mathrm{d}t}=c_{0},$ (34) $$

$$ $y=c_{0}c_{1}+\left(u_{0}-c_{0}c_{1}\right)\exp\left(-\frac{t}{c_{1}c_{2}} \right).$ (35) $$

Population growth (naive), $\left(c_{0}\right)=\left(0.23\right)$, $\left(u_{0}\right)=\left(4.78\right)$,

$$ $c_{0}y-\frac{\mathrm{d}y}{\mathrm{d}t}=0,$ (36) $$

$$ $y=u_{0}\exp\left(c_{0}t\right).$ (37) $$

Language death model for two languages, $\left(c_{0},c_{1}\right)=\left(0.32,0.28\right)$, $\left(u_{0}\right)=\left(0.14\right)$,

$$ $\left(c_{0}+c_{1}\right)y+\frac{\mathrm{d}y}{\mathrm{d}t}=c_{0},$ (38) $$

$$ $y=\frac{c_{0}}{c_{0}+c_{1}}-\left(\frac{c_{0}}{c_{0}+c_{1}}-u_{0}\right)\exp \left(-\left(c_{0}+c_{1}\right)t\right).$ (39) $$

Harmonic oscillator without damping, $\left(c_{0}\right)=\left(2.1\right)$, $\left(u_{0},u_{1}\right)=\left(0.4,-0.03\right)$,

$$ $c_{0}y+\frac{\mathrm{d}^{2}y}{\mathrm{d}t^{2}}=0,$ (40) $$

$$ $y=u_{0}\cos\left(t\sqrt{c_{0}}\right)+\frac{u_{1}}{\sqrt{c_{0}}}\sin\left(t \sqrt{c_{0}}\right).$ (41) $$

Harmonic oscillator with damping, $\left(c_{0},c_{1}\right)=\left(4.5,0.43\right)$, $\left(u_{0},u_{1}\right)=\left(0.12,0.043\right)$,

$$ $c_{0}y+c_{1}\frac{\mathrm{d}y}{\mathrm{d}t}+\frac{\mathrm{d}^{2}y}{\mathrm{d}t ^{2}}=0,$ (42) $$

$$ $y=\exp\left(-\frac{c_{1}}{2}t\right)\left(u_{0}\cos\left(\frac{t}{2}\sqrt{4c_{ 0}-c_{1}^{2}}\right)+\frac{c_{1}u_{0}+2u_{1}}{\sqrt{4c_{0}-c_{1}^{2}}}\sin \left(\frac{t}{2}\sqrt{4c_{0}-c_{1}^{2}}\right)\right).$ (43) $$

Additional third-order ODE, $\left(u_{0},u_{1},u_{2}\right)=\left(0,-1,1\right)$,

$$ $\frac{\mathrm{d}y}{\mathrm{d}t}+\frac{\mathrm{d}^{2}y}{\mathrm{d}t^{2}}+\frac{ \mathrm{d}^{3}y}{\mathrm{d}t^{3}}=0,$ (44) $$

$$ $y=u_{0}+u_{1}+u_{2}+\exp\left(-\frac{t}{2}\right)\left(-\left(u_{1}+u_{2} \right)\cos\left(\frac{\sqrt{3}}{2}t\right)+\frac{\sqrt{3}}{3}\left(u_{1}-u_{2 }\right)\sin\left(\frac{\sqrt{3}}{2}t\right)\right).$ (45) $$

### B.2 Lorenz Discovery

Table 3 lists the discovered coefficients after
final fine-tuning,
along with the results from the state-of-the-art method SINDy (Brunton et al., ). Our S-MNN closely matches the ground truth, SINDy, and both the original MNN solvers, with only minor differences observed.

**Table 3: Discovered coefficients for the Lorenz system using different methods.**
| Method | $a_{1}$ | $a_{2}$ | $a_{3}$ | $a_{4}$ | $a_{5}$ | $a_{6}$ | $a_{7}$ |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Ground Truth | -10 | 10 | 28 | -1 | -1 | -8/3 | 1 |
| SINDy | -10.000 | 10.000 | 27.998 | -1.000 | -1.000 | -2.667 | 1.000 |
| MNN Dense | -10.0003 | 10.0003 | 27.9760 | -0.9934 | -0.9996 | -2.6660 | 0.9995 |
| MNN Sparse | -10.0055 | 10.0061 | 27.7165 | -0.9304 | -0.9937 | -2.6641 | 0.9990 |
| S-MNN | -10.0003 | 10.0004 | 27.9915 | -0.9968 | -0.9997 | -2.6664 | 1.0000 |

## Peer Reviews

confidence: 2
contribution: 3
initial_rating: 6
presentation: 2
questions: - Are there any cases where intuitively MNN should outperform S-MNN?
soundness: 3
strengths: - The problem is well-motvated 
- S-MNN seems to achieve good performance while improving the computational cost 
- Detailed experimental results are presented
summary: The paper presents Scalable Mechanistic Neural Network (S-MNN) designed for long temporal sequences.
S-MNN is more efficient compared to a previously proposed solution MNN as it reduces the computational time and space to linear in the sequnce length.
The paper describes the algorithmic improvements to reduce the computational cost and provide experimental results
weaknesses: As discussed by the authors in the limitation section, parallelism along the time dimension can be challenging.

---

confidence: 2
contribution: 2
initial_rating: 5
presentation: 3
questions: I list below my two main suggestions. As previously mentioned, addressing the first point is what is most likely to change my current rating of the paper.

- [**Overall Relevance To Scientific Computing**] Although the authors do list related methods in scientific machine learning in section 4, it is not clear to me why improving on traditional MNNs is desirable over utilizing these other methods. Why choose to focus on MNNs? In order to answer this question, for example, the authors could run their experiments using some of the other methods listed in section 4, listing their accuracy and computational cost when compared to MNNs and S-MNNs. For example, some suggestions for comparisons are: 1) in section 5.1 and 5.2, they could run comparisons against Neural-ODEs and perhaps a traditional neural operator such as FNO or DeepONet; 2) in section 5.3 they could compare against a traditional MLP using both data and a Physics Informed loss (as detailed in line 300); 3) in section 5.4 the could compare against other autoregressive models including a traditional transformer or a neural operator. Doing so would help contextualize their work in a quantifiable way, and could help increase my score of their submission, even if S-MNNs do not beat other all other methods.

- [**Use of Relative L2 Errors**] As mentioned in the previous section, reporting errors using relative L2 error instead of MSE/RMSE provides a more reliable way of comparing methods across different benchmarks, since solutions can have drastically different magnitudes. Is there a particular reason why reporting absolute MSE/RMSE makes more sense for your benchmarks?
soundness: 4
strengths: ### Originality
This work innovates by proposing a more computationally efficient way of implementing Mechanistic Neural Networks (MNNs), making it more scalable for longer roll-outs and numerically stable. Although the paper largely basis itself on the recent work by Pervez et all., 2024, they improve upon this method in an original way, and present practical ways of making the MNN strategy less computationally expensive.


### Quality
The submission appears to be well thought-out, making precise contributions to the previous work of MNNs. The simplification of the underlying ODE solver implied by the architecture appears to yield similar performance at a computational cost 2-5 times lower, with greater grains when the sequence length $T$ is bigger. The author's computational improvements are substantial, and the theoretical arguments used to formulate and motivate S-MNNs are well crafted.


### Clarity
I was not originally familiar with MNNs, so there are parts of the paper that seemed unclear to me, but I think the authors overall do a decent job of introducing the important concepts and explaining how their alterations lead to computational efficiency.
summary: In their paper, the authors propose an alteration to the recently published Mechanistic Neural Networks (MNN), making the architecture less computationally demanding, while maintaining its predictive power. Their method, called Scalable MNNs (S-MNN), decreases computational cost primarily by removing certain slack variables present in the original method, as well as replacing a quadratic programing problem with a simple least squares regression, which has known analytical solution and can be more efficiently computed.

The authors begin by contextualizing their work among the scientific machine learning literature, then providing a brief overview of MNNs, an architecture aimed at modeling time-series data by learning an ODE in latent space, where the inputs are first encoded to represent the dynamics of the desired ODE. The solution of this ODE is then optionally passed through a decoder to map the solution to the desired output space. The authors then go on to describe their method, S-MNN, which simplifies certain aspects of the original MNN implementation. The main difference between the two methods is the linearization of the underlying ODE, which makes it solvable using a traditional least squares regression, with known analytical solution. The authors also propose a sparse method of computation for obtaining this solution, which is easily parallelizable using GPUs. Overall, their alterations reduce the time and memory complexity of the architecture so that it scales linearly with the sequence length $T$.

Finally, the authors briefly describe other related work in the scientific machine learning literature, and carry out experiments to empirically validate their method. The authors employ S-MNNs on 1) a simple set of ODEs, as a sanity check for the architecture; 2) the discovery of underlying parameters of a Lorenz system; 3) solving the Korteweg-De Vries (KdV) PDE from data; and 4) forecasting sea surface temperatures from the SST-V2 dataset. Overall, the experiments indicate that S-MNNs perform comparable or slightly better than traditional MNNs, but at considerably lower computational cost.

The authors conclude by summarizing their work and outlining limitations and possible future directions of work, including parallelizing their implementation over the time dimension, which currently need to be solved sequentially.
weaknesses: ### Contribution Is Specific To A Niche Method
Although the ideas in the paper seem interesting, my main criticism of this submission is that their contribution is very narrow, in the sense that it improves on a method that is generally not very employed in practice. This makes their work less relevant to the scientific machine learning community. While the authors carry out several well-designed experiments comparing MNNs and S-MNNs, it is unclear to me what the advantages of S-MNNs are over the other methods described in section 4 of the paper. Potentially including some of these other methods in the experiments carried out, and comparing their accuracy/computational cost would help contextualize their work and make a stronger case for their approach. I detail some suggestions for comparisons in the "Questions" section below.


### Errors Reported Using Absolute Error Instead of Relative Error (minor weakness)
Errors should ideally be reported as relative L2 error, instead of absolute error squared. Using relative L2 error helps contextualize the performance of the method for target functions of different magnitudes. For Example, in section 5.1 the harmonic damping problem has a solution that oscillates from $\approx -0.1$ to $\approx. 0.1$, while the population problem oscillates from 0 to $\approx 40$. Using relative L2 errors would also make their results more easily comparable against what is reported in other papers for comparable problems.


### Other comments that did not affect my score:
- [line 333] While it is great to see very close agreement between true and learned solutions, it would still be valuable to report exact errors for each benchmark, likely in the appendix.
- [figure 4] Since all predictions look very close to the ground truth solution when eyeballing the plots, it would be great to plot absolute error as well, which could help differentiate the performance of the two methods.

---

confidence: 3
contribution: 4
initial_rating: 8
presentation: 3
questions: - I don’t fully understand what determines the sparsity level of $A$. You say 95% in typical applications, which kinds of applications yield denser $A$? IIUC, the sparsity comes from the fact that none of the Eqs 5,6,7,8 entangle $y$ that are more than a timestep apart.

### Nits:
- $V$ is undefined when introducing Eq 1, only defined later
- maybe you can write out the definition of A, y and b once for clarity


### To the other reviewers
I am not an expert in this area, so take my review with a grain of salt.
soundness: 4
strengths: - straightforward and good idea, fuelled by several computational tricks
- strong complexity reduction!
- clear writing
- clear motivation
- experiments on toys to illustrate the point and on real world data to show feasibility
summary: The authors provide an alternative method to implement ODE solving in mechanistic neural networks (MNNs) to achieve better time and space complexity without sacrificing performance.
weaknesses: - while the description of MNNs is nice and concise, i believe a reader could benefit from a small additional description of the aspects of the network that are improved upon, like the slack variables.
- I am missing some ablations studies, for instance on the importance weights. do they do anything in practice?

---

confidence: 4
contribution: 2
initial_rating: 6
presentation: 4
questions: 1) The paper mentions the batch size of 512 as a default setting. For the Lorentz system with initial conditions $x=y=z=1$, what makes up an instance in the "batch", especially since small batch and block sizes seems to be a limitation of S-MNN (line 536).
3) The Cholesky decomposition in Eq. 12 hinges on the matrix $M$ being positive definite, which is claimed in line 201. But how is that guaranteed ? The same remark goes for Eq. 28 where the symmetric nature of $M^{-1}$ is implied but not (re)stated.
soundness: 4
strengths: Overall, the exposition is crisp, much clearer than the work it builds upon, of which the authors have given a great overview. The more direct formulation allows efficient forward pass, and a backward pass now directly supported by automatic differentiation. That said, the authors propose an analytical form that enables quicker numerical integration. The range of experiments is broad and the results are compelling, showing significant improvements over ordinary dense and sparse MNN.
summary: Building on mechanistic neural networks (MNN), this work proposes a strategy to efficiently solve long sequences using S-MNN. It addresses the scalability issues and reduces complexities to linear with 3 key improvements: elimination of slack variables, higher-order smoothness constraints, and least square regression. The result is a framework that can tackle long sequences with up to 4.9x speedup over MNN, and 50% GPU memory reduction, something the original MNN could not do.
weaknesses: 1) The work does not address the trade-offs from the 3 main changes it brings, notably the consequence of dropping the slack variables. For instance, MNN notably allows **nonlinear** ODEs to be integrated with relatively ease. Is there any reason S-MNN doesn't do the same ?
2) Limited evaluation. The original MNN showed improved performance over ResNets and FNOs, but the current method has only been compared to MNN. Particularly, in this same setting, it would be interesting to know how well models outside the MNN family perform on the KdV and long-term sea surface temperature forecasting experiments, at the very least to gauge how difficult the task is.
3) Training and testing. During training, a number of quantities are learned. But during testing, it is not clear whether only the $y$s are solved for, and whether other matrix quantities from Algorithm 1-4 are reused. Appropriate training and testing algorithms should help clarify this.

Meta-review: The paper extends the "mechanistic neural network" (MNN) framework, proposed this year at ICML. The MNN encodes the dynamics of a system in an interpretable ODE, which is then solved by an ODE solver and passed to a neural decoder. The extended version proposed here has significantly better computational complexity, which extends its possible applicability.

The paper had 4 four reviews, all positive after the rebuttal. Reviewers are generally happy about the the presentation, the technical novelty, and the possible impact. All the reviews are discussed better below. The only remaining concern is the scope of the paper, which extends a very recent method which, in the words of one reviewer, "has just recently been published and [is] not necessarily widely adopted (at least as of now)".

In light of (a) the technical validity of the paper, (b) the enthusiasm of some reviewers, and (c) the importance of AI for science, I am leaning towards acceptance.

Decision: Accept (Poster)