Mechanistic PDE Networks for Discovery of Governing Equations
Authors: Adeel Pervez
Efstratios Gavves
Francesco Locatello
Venue: ICML 2025 poster (2025)
Keywords: Differential equations
discovery and inverse problems
AI for science
Primary area: deep_learning->other_representation_learning

Abstract: We present Mechanistic PDE Networks -- a model for discovery of governing *partial differential equations* from data.
Mechanistic PDE Networks represent spatiotemporal data as space-time dependent *linear* partial differential equations in neural network hidden representations. The represented PDEs are then solved and decoded for specific tasks. The learned PDE representations naturally express the spatiotemporal dynamics in data in neural network hidden space, enabling increased modeling power. Solving the PDE representations in a compute and memory-efficient way, however, is a significant challenge. We develop a native, GPU-capable, parallel, sparse and differentiable multigrid solver specialized for linear partial differential equations that acts as a module in Mechanistic PDE Networks. Leveraging the PDE solver we propose a discovery architecture that can discovers nonlinear PDEs in complex settings, while being robust to noise. We validate PDE discovery on a number of PDEs including reaction-diffusion and Navier-Stokes equations.

Mechanistic PDE Networks for Discovery of Governing Equations
Adeel Pervez' Efstratios Gavves? 2 Francesco Locatello
Abstract
tions (Brunton et al., 2020). In spite of progress (Brunton et al., 2016; Raissi et al., 2019; d' Ascoli et al., 2023), the development of machine learning discovery models that are flexible, handle complex nonlinear data and account for missing physics, while being robust noise and missing data, is very much an open problem.
We present Mechanistic PDE Networks model for discovery of governing partial differential equations from data. Mechanistic PDE Net- works represent spatiotemporal data as space- time dependent linear partial differential equa- tions in neural network hidden representations. The represented PDEs are then solved and de- coded for specific tasks. The learned PDE rep- resentations naturally express the spatiotempo- ral dynamics in data in neural network hidden space, enabling increased power for dynamical modeling. Solving the PDE representations in a compute and memory-effiiiet way, however, a significant challenge. We develop a native, GPU- capable, parallel, sparse, and differentiable multi- grid solver specialized for linear partial differen- tial equations that acts as module in Mechanistic PDE Networks. Leveraging the PDE solver, we propose a discovery architecture that can discover nonlinear PDEs complex settings while also be- ing robust noise. We validate PDE discovery on number of PDEs, including rractionddiffsiion and Navier-Stokes equations. Source code will be made available at: https: /ttph//ggttu...///// com/ alpz/mech-anee www.eehh--n-iiiooerr----e
A recent discovery architecture for ordinary differential equations (ODEs) was demonstrated by Mechanistic Neural Networks (MechNN) (Pervez al., 2024), enabling a flexi- ble and expressive network design for modeling explicit, in- terpretable nonlinear ODEs. MechNNs embed specialized, parallel ODE solver in the network, allowing explicit differ- ential equation representations in neural networks, which can be employed for ODE discovery and prediction, both relevant in applications including general time series and numerical simulations (Pervez et al., 2024) and in climate science (Chen et al., 2025; Yao et al., 2024).
In this paper focus on governing equation discovery and demonstrate Mechanistic Neural Networks for discovery of partial differential equations from data. We develop PDE solver, NeuRL_-PDDE, for parallel native differentiable solution of linear partial differential equations over multiple spatial dimensions (in addition to the time dimension). In contrast to ODEs, scaling is significantly more difficult in PDE solving, with memory requirements increasing expo- nentially with dimension and quickly becoming excessive even for mid-sized grids. With NeuRLP-PDE. sparse com- putation in all phases of the PDE solver ensures that the memory requirement of the solver over batches of PDEs remains limited for sizeable grids. We build the solver using constrained optimization and a sparse multigrid precondi- tioned iterative linear solver (Saad, 2003) for solving the large and sparse linear systems that appear in the constrained optimization. The NeuRL.-PPDE solver is then incorporated into the PDE discovery architecture enabling the learning of governing PDEs directly from data.
1. Introduction
Partial differential equations (PDEs) enjoy widespread use as interpretable, analytical models of spatio-temporal dy- namies in all areas of science (Temam, 2024; Turing, 1952). The development of PDE models still, however, is largely manual task that is performed by domain experts examining experimental data. The modern explosion of experimental and simulation data has underscored the need for data-driven machine learning models for discovery of governing equa-
The MechNN-PDE discovery architecture couples the NeuRLP-PDE solver with deep neural networks that pa- rameterize partial differential equations. Two natural con- straints on discovered differential equations are that they should be 1) simple: i.e., built from elementary functions and 2) concise: containing as few terms as possible. The MechNN-PDE architecture enables simplicity and concise-
'Institute of Science and Technology, Klosterneuburg, Aus- tria *Informaties Institute, University of Amsterdam, Amster- dam, The Netherlands. Correspondence to: Adeel Pervez < Redbze.Pereeziittac..t
Proceedings of the 4rnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Mechanistic PDE Networks for Discovery of Governing Equations
ness by building PDE expressions using family of elemen- tary basis functions, each with learnable parameter. Since elementary functions are restrictive, for expressive model- ing the basis functions inputs are parameterized by neural networks. The built PDE is then solved by NeuRLP-PDDE to obtain the solution u. Loss terms ensure consistency so that the parameterized basis functions equal the true basis functions at convergence. A sparsity loss ensures that the discovered expressions are concise and the entire pipeline is optimized by gradient descent.
set B of boundary coordinates we complete the PDE spec-
ification by specifying the initial and boundary conditions with a function w defined over B as
u(i) = (i). Vi B,
In general mechanistic PDE networks allow flexibility in the specification of initial and boundary conditions which may be of the Dirichlet (i.e., boundary conditions on u), Neu- mann (i.e., boundary conditions on partial derivatives) or mixed types and may be arbitrary linear constraints (subject to well-po well-posedness). The time and space varying coefficients allow the PDE representations of mechanistic PDE networks to be very expressive and allow for modeling of complex dynamical phenomena. In this paper we assume that PDE representations are well well-posed and that the PDEs together with the boundary conditions have a unique solution.
The proposed synthesis of neural networks and PDEs for discovery has the following advantages.
Derivatives never have to be directly on data and are only computed internally by the solver.
Discretization. To implement the PDE in neural network we discretize the rectangular domain D over time and space. We assume one time and one spatial dimension. A given time interval [0,4y] is discretized into n, intervals (4i, 4++1 with step size St, = f4+1 f, Similarly a spatial interval [0, xi.5l is discretized into nar, intervals [ Fi,4 1] with steps Sp,,t.
PDEs can contain arbitrary differentiable functions, going beyond generalized linear functions.
Numerical solving together with neural networks pro- vide robustness to noise and missing data.
Memory-efficient batch-parallel PDE solving on GPU with a sparse multigrid-precoondttinee iterative solver.
All functions Cma b, w produced by the mechanistic encoder are now assumed to be defined over the discretized spatio- temporal grid. With mechanistic networks the steps sizes can also learned and are then produced by the mechanistic encoder. All together the mechanistic encoder parameterizes the PDE by producing the coefficients, boundary parameters and the step sizes defined over the spatio-temporal grid for given multi-index set M.
2. Mechanistic PDE Networks
A Mechanistic PDE Network consists of a mechanis- tic encoder, fFne. differentiable solver, NeuRLP-PDE, specialized for linear PDEs, and an optional decoder, Dee. The input to the encoder are spatio-temporal data over time and one or more spatial dimensions, denoted Usann 12,,1,1122..... ,2,...,,TT.... Ina). Given input Udata the mechanistic encoder produces a set, of predefined size neg. of linear par- tial differential equations which we denote (Pu, hie[n...] Each Pu, is a linear PDE for a spatio-temporal function (t, 11, ....
Hn } mEM: b, w, (st. ), (82.3] = (Udat)
Next the generated PDE is fed to the NeuRLP-PDE solver which produces the solution and all partial derivatives spec- ified in M as { The output of the solver is fed to a task-specific decoder or directly used in a loss. All aspects of the PDE model in equation 1 including coeffi- cients, initial and boundary conditions and step sizes are differentiable.
(using the notation [n] :=
All functions u, are defined over a domain D which we as- sume to be Cartesian grid, i.e., D = T x X1 x X2... where T, X; are intervals. For conciseness we collect the partial derivative indices appearing in PDE Pu; into multi-index set denoted M. As an example M ft, 21, Z2, 21 {um IEE denotes the set of partial derivatives (uz, May Mayg ) We also use the empty index ó for u, i.c, Ug == u. For simplicity of exposition we assume a set with a single PDE denoted Pa, for a func- tion u(t, .r) over a spatial dimensions x := (x1 (#1,..., Ind) with E [0, ty] and 2; E [0,
Nonlinear PDEs. The PDE representations generated by MechNN-PDE are linear, time-space varying. However, the model is not restricted to learning linear PDEs. Nonlinear PDEs can be represented by using nonlinear basis functions, either directly over the input data udan or by learning non- linear functions of u. We illustrate this further in the section on discovery.
With these simplifications Pu has the following general form:
3. Linear PDE Solving With Differentiable Optimization
Pu: :
Σ Cm (t,, x;; udda at; Udd MEM
(1)
The workhorse of mechanistic PDE networks is NeuRLP- PDE a specialized, parallel and differentiable solver for
where M is a multi-index set, uhn are partial derivatives and Cm, are time-space varying coefficients over D. Given a
2

---

Mechanistic PDE Networks for Discovery of Governing Equations
linear PDEs. NeuRLP-PDE solves linear PDEs, with par- tial derivatives of arbitrary order, over discretized spatio- temporal grid by reducing PDE solution to a relaxed differ- entiable constrained optimization (Young, 1961).
grids point for that dimension. The constraint specifies that the Taylor approximation and the solution variables the next point should be close. We illustrate with second-order Taylor approximation over the t dimension.
Briefly, the constrained problem specifies that 1) the dis- cretized PDE holds over all points of the grid, 2) that the initial or boundary conditions hold and 3) that the computed solution is a smooth function over the grid, and the objec- tive is to minimize approximation error. We then solve relaxation of the optimization problem so that we can dif- ferentiate through the PDE solution to the PDE data using techniques from differentiable optimization.
Up, + St +
2 Umz.. = (i)
(6)
U - SyWmm, + $7 Um2.d Uprev,
(7)
2
where m1, m2 are multi-indices corresponding to the first and second partial derivatives of t.
We collect the variables um,ii m M, E into a vector z. We can then write the constraints in matrix form as
A= d,
(8)
To fully specify the constrained optimization, we specify the optimization variables and constraints. We consider the dis- cretized Cartesian grid with and x discretization (eje))
with A E R being the coefficient matrix with ne constraints, and d is the vector of constant coefficients.
{23 he(e_) Let I and B denote the set of grid and boundary indices, respectively. We also use the shorthand next.(i) (i) and prev,(i) where i E I and is a dimension for the next and previous adjacent grid points to point for coordinate c.
We can solve eq (8) in a least squares sense by solving the corresponding normal equations.
AT = ATd.
(9)
Solving these equations gives a solution of the optimiza- tion problem and thus of the PDE. However, we also require the dual variables for the differentiating through the least squares optimization. For this we consider the following saddle-point formulation of least squares (Saad, 2003).
Variables. We are given a multi-index set M with & E M. We create variables ug E I for the solution u at each grid point. We also create variables for each partial derivative for each i E I. obtaining variables denoted ve, and Mzz,i for partial derivatives i, and Uaz and similarly for all other partial derivatives from M. We denote the number of optimization variables by Tl.-
A
x*
[
(10)
AT
0
2
0
where À are the dual variables. This formulation can also be seen as the KKT constraints of the convex quadratic pro- gramming formulation of least-squares. Given the saddle- point formulation we can apply techniques from differen- tiable optimization for differentiating through the solution of constrained optimization problems (Amos & Kolter, 2017).
The solver has three types of constraints: equation, smooth- ness and initial and boundary constraints.
Equation constraints specify that the left and right hand sides of the given PDE are equal over all grid points.
Cm, = by Vi E
Forward pass. In the forward pass we solve the normal equations and compute the dual variables from 10 to obtain =* and X
MEM
Initial and boundary constraints are specified for points in B, where the form of the constraint can depend on the problem. Here we give Dirichlet-tyre typee conditions which specify the values of u at points in B.
Backward pass. Give scalar loss / we consider the gradient, 9: == 0.-l, relative to 2". We require the gradient of I relative to the matrix A and vector d which contain the optimization data, Following Amos & Kolter (2017), we can achieve this by solving the following saddle-point system:
Ug.i Wi, Vi E B
(4)
We use central difference constraints to compute partial derivatives. We illustrate for 3 point central difference approximation for a second order partial derivative u, with derivative multi-index m.
T
A
d2
0
(11)
AT
0
d
-92
The gradient of the loss 1 relative to A and d is 0nl dxæ*T + X dd and 0,l = -dx (Amos & Kolter, 2017).
of Um,m (1rev(i) - 2u, Wnext (1) Vi E I B
(5)
In practice we use 5 point central differences and use for- ward, backward finite differences at the edges of the grid.
Scaling PDE Solving. The linear systems in equations (10), (11), (9) can become very large with increasing dimension of the PDEs resulting in very large grids over which the PDE is defined. As an example, for a PDE defined over 32x32 grids we have 5120 variables and 9212 constraints, requiring 377 MB of double precision dense storage. How- ever for 32x32x32 (220k variables and 420k constraints)
Forward and backward smoothness constraints specify that the solution u is smooth over the grid in each dimension in both directions. This is achieved by computing a ID Taylor expansion at each grid point for each dimension and approximating the u value at the next and previous adjacent
3

---

Mechanistic PDE Networks for Discovery of Governing Equations
would need 780GB and for 64x64x64 we would need 50k GB. For such large grids the linear systems cannot be solved by dense methods. For the experiments we solve ID PDEs with dense solver since the memory requirement is reasonably bounded for small grids and batch sizes. How- ever, for 2D or higher PDEs with midsize grids (say, 32 to 64 per dimension), the memory required for dense solvers too great for GPU storage. The matrices. however, are highly sparse with 99. 99.9 zero values for 32x32 matrix and 99, 99999 zeros for a 64x64x64 matrix. For such problems we develop multigrid preconditioner for use with sparse iterative solver.
U
Ux
PDE Expression
E u
solve Ut = power
U
Figure 1. PDE expression building and discovery architecture. The expression is u, = u, with parameter e as the exponent.
3.1. An Iterative Sparse Multigrid Solver
In the second half of the V-cycle the error at grid k is transferred to the next finer grid k - 1, starting from the coarsest grid, by a prolongation operator, i.e., éx-1 = prolong(ew). The prolonged error is added the er- ror at the finer grid, to obtain the new error at grid k 1, i.e, ék-1 6k-1 + 6k-1- The error is smoothed by re- laxation ék-1 == re relaa velax(MM,
To make PDE solving feasible for more than two dimensions (including time) and for large grids, we build an iterative sparse solver for solving the normal equations. We use the FGMRES linear solver, the flexible variant of the Krylov- subspace linear solver GMRES (Saad & Schultz, 1986), as the base linear solver. As is well-known from numerical linear algebra, standalone iterative solvers can lead to poor solutions for complex problems and require precondition- ing for accurate solution. General purpose preconditioning methods such as incomplete LU or Cholesky factorization are not feasible for the large systems we encounter due to their high memory use and lack of GPU implementations. Instead, we build a multigrid V-cycle preconditioner for FGMRES. We briefly describe the multigrid V-cycle here and in Algorithm in the appendix. More details on the multigrid methods and sparse iterative solvers can be found in Saad (2003) and Briggs et al. (2000).
fk- rk-) The V-cycle ends when we return to the finest grid and we improve the solution zn by adding the estimated error, i.e, 2o == Zy Hotto G1-
In our solver we use the multigrid V-cycle as a precondi- tioner for FGMRES, which we found led to more accurate solutions compared with conjugate gradient methods and standalone V-cycle. We use Gauss-Seidel as the relaxation method and linear interpolation as the restriction and pro- longation methods. To construct coarse grid matrices we use grid sizes with powers of and successively halve each dimension whlist simultaneously doubling the step size in each dimension. We use size 8 for each dimension for the coarsest grid. The PDE coefficients and boundary condi- tions are coarsened by lincar interpolation and the multigrid preconditioned FGMRES is applied to the matrix AT A from the normal equations.
Multigrid V-cycle. The multigrid V-cycle approach solves linear system for a PDE on a succession of fine to coarse grids. Given square linear system Mæ with solution xx and current solution zo can compute the error c= x=="-xp Zo and residual as ro = b b Mro. We note that we can rewrite the linear system as Me = To, where we wish to solve for the error 6o- Given an estimate for 6n we can add it to Zo to obtain an improved solution. Next we assume that we have a sequence of ny successive boarse-grained versions of M (with the dimension of M1+1, being, say, half the dimension of M,) given as M1 M, M2, Mn,-
The resulting iterative solver is entirely sparse and batch- parallel on GPU for a batch of PDEs and is useful for higher dimensional PDEs where GPU memory cannot hold the full dense constraint matrices.
4. Mechanistic Networks for PDE Discovery
In the first half of the V-cycle, at the kth grid the linear system Ma = ты is partially solved by taking a few steps of a simple iterative relaxation method such as the Jacobi or Gauss-Seidel method, i.e., E := re velaax( rs). Next the new residual r/g := тк MrCa in the linear system at grid is is computed and transferred the next coarser grid k++1 by restriction operator designed for grid transfer, i.e., Te+1 := restrict (r.)) The linear system at the coarsest grid n, is small and we solve for the exact error En, for the given residual .ng.
The goal of discovery is to find the equations govern- ing spatio-temporal data representing some physical phe- nomenon. The result of governing equation discovery is a set of partial differential equations that, ideally, completely describe the given data. Further restrictions may be added to the form of the discovered equations so that the equations may be interpreted by human scientists. One such restric- tion is that the resulting equations should be composed only of simple functions. Another natural restriction is that the equations should concise, or sparse, with as few terms as

---

Mechanistic PDE Networks for Discovery of Governing Equations
possible (Brunton et al., 2016).
expression input is to make it robust to noise in the input
data. To ensure that the final learned equation is correct for u, we also add a loss term (oss(u, ū) so that ū is close to at convergence.
In this section we present a Mechanistic PDE Network dis- covery architecture that allows discovery of simple and concise governing partial differential equations from spatio- temporal data. With Mechanistic PDE Networks the PDE representations can contain arbitrary differentiable expres- sions over inputs parameterized by neural networks as in Figure 1. This allows the design to model complex nonlin- ear PDE expressions, going beyond linear combinations of fixed basis functions.
Discovery Example. As an informative example, we il- lustrate the discovery method for Burger's equation: ut = uuy + 0.1 where the diffusion term, Uga. has the fixed coefficient 0.1. Given Udata- we first compute a transfor- mation ū NN(udau). Choosing degree two polynomials as the expression functions we obtain the coefficients for Wy and Usz respectively as pı (ū) = Oo + 01 ü + 0zü² and P2 (z) = Φ + Φ ü + /2ū², with 0, and o; as learnable parameters. Choosing second order spatial and first order temporal derivatives we get a PDE expression of the form
4.1. Discovery Model
Given spatio-temporal data (t), where assume two dimensions for simplicity, the discovery method gener- ates a parameterized PDE expression along with initial and boundary conditions. P. == Pul(u (u, ue, Mgy Uet, Utr, Mee: +). In general the form of the PDE may include arbitrary differ- entiable functions of any complexity, such as deep neural networks. However satisfy the twin goals of simplicity and conciseness restrict the allowed forms of the expres- sions appearing in PDEs to simple tree structures where each node is an elementary function with some weighting param- eter. One common example of a tree expression structure in PDE modeling is a finite set of simple basis functions such as polynomials of a fixed maximum degree. Mechanistic net- works, however, are not constrained linear combinations of fixed basis funetions and can represent more complex expressions by combining the data and its derivatives with parameters in differentiable way (Figure 1).
U₁ - Pi (ū; 0)u,- + p2(ū; d)uzz.
(12)
The true equation corresponds to the case where 0₁ = 1, Φo = 0.1 and the remaining parameters are O. The ini- tial and boundary conditions are obtained from ã. In each iteration we solve Equation 12 using the NeuRIP- PDE solver obtaining solution u. We then compute the loss as a sum of the losses (either Ln or I2) between Udata+ u and u. ü. We also include an L1 loss over the parameters for sparsity. The total loss is computed as (Udaaa u. ū, 0, Φ) loss(udara, u) + loss(u, ū) + l (0, e).
The loss is iteratively minimized with gradient descent. Fi- nally a concise PDE is generated by thresholding the param- eters, with the ones below the threshold set to 0.
5. Related Work
Once the form of the expression is specified the aim of the discovery procedure is to find the appropriate parameters of the PDE expression given data. One could also imagine learning the structure of the expression itself; however, that is not direction we take, and in this paper, use fixed forms of PDE expressions.
The synthesis of machine learning and differential equations has been treated in a few different ways in the field. A major line of work approaches the problem from a machine learn- ing perspective with supervised or unsupervised learning for partial differential equations (Li et al., 2020; Brandstetter et al., 2022). Another line of work in machine learning for scientific application aims synthesis of differential equation models with a network-ing network-inn-solver approach by em- bedding neural networks in classical solvers (Chen et al., 2018; Rackauckas et al., 2020) enabling access to high qual- ity numerical solvers in the context of machine learning. Physics-informed networks use a physics-aware loss to su- pervise network training (Raissi, 2018: Raissi et al., 2019). Mechanistic neural networks conversely are solver-in- network approach that attempts to balance classical solvers and neural networks by combining restricted ODEs (Pervez et al., 2024: Chen et al., 2025) (and now PDEs), that allow- ing faster solving, with neural network learning together with techniques for handling nonlinear equations.
Parameterized PDE Expressions. Given the restriction of conciseness, there must only be a small number of expres- sion parameters that describe complex data, Furthermore, the initial values of these parameters might be far from the true parameters. If the PDE solution corresponding to the initial parameters is far from the true Udaaa the iterative learning can become brittle and non-smooth due the small number of parameters available to the optimization process.
To make the learning smoother and more flexible, we pa- rameterize the input to the PDE expression with a neural network. We do this by transforming the data Udata to á (N( with a neural network and compute the expression the transformed data î instead of Udata- This allows the learning to adapt both parameters and basis input to get to the target solution, leading to a more flexible leam- ing procedure. A further advantage of parameterizing the
As an application of the synthesis of ML and differential equations various data driven techniques for discovering governing equations have been explored for ODEs such as
5

---

Mechanistic PDE Networks for Discovery of Governing Equations
SINDy (Brunton al., 2016), SINDy-PI (Kaheman al., 2020) and PDEs including PDEFIND (Rudy al., 2017) and weak forms (Reinbold et al., 2020; Messenger & Bortz, 2021) extending the SINDy framework. Physics informed networks (Raissi et al., 2019; Raissi, 2018) have also been used for inverse problems but unlike UDEs (Rackauckas et al., 2020) and SINDy do not handle unknown physics. PDE-LEARN (Stephany& & Earls, 2024) combines features of physics informed networks and basis libraries for PDE discovery. MechNNs (Pervez et al., 2024) build discovery model for ODEs which we enhance and develop PDE mod- els that handle nonlinear equations and unknown dynamics.
Table 1. PDEs for Discovery
Name
Expression
Diffusion
Puze Us + uu, = Vuse
Viscous Burger's
Inviscid Burger's
Wt + uuk 0
Porous Medium
V (nm)
Ginzhurg-L.ame Landuuu
a
A+ V2 A (1 + A 8.U V
Navier-Stokes
6. Experiments
We demonstrate PDE discovery on number of PDEs in two and three dimensions (including time), including complex dynamical data from the 2D Navier Stokes and reaction diffusion PDEs. We examine the robustness of the method to noise. We also consider PDE parameter discovery in the case where the PDE cannot be expressed as a linear combina- tion of fixed basis functions. Table lists the PDEs that we use demonstrate our method, including their general form.
Figure 3. Easy case, showing A' (left) and harder, showing the norm |4|2 (right) at single time-step for reaction-diffusion data.
on a number of ID and 2D (in space) PDEs chosen to cover range of complexity. For the 1D equation, we demonstrate discovery with the diffusion equation and Burger's equation (both viscous and inviscid (i.e., without viscosity)). The in- viscid Burger's 's equation is known to develop shocks, which appear as a discontinuity in wave propagation. We show that our proposed method can perform discovery in the presence of this form of discontinuity. The porous medium equation is a nonlinear equation that cannot be written in the form of linear combinations of fixed basis functions because of real exponent. We show how the mechanistic network framework easily extends to such cases while many other methods cannot.
Robustness to Noisy Data. We test discovery in noisy settings by adding Gaussian noise of variance o² to the data. The noise standard deviation 0 is chosen to be a ratio R a of the data root mean squared value i.e., o == ONN ||RMs where U denotes the dataset and N E [0, 1] is the noise ratio (Messenger & Bortz, 2021). In the experiments, we report the noise level percentage as NN 100% noise.
Evaluation Metrics. We use two metrics for evaluating the quality of the discovered equation. The true positiv- ity ratio is defined as TPR- = TP-IN+FP (Lagergren et al., 2020) and measures how many terms have been correctly identified. We measure the maximum relative error of the K3 with E true coefficients as (E) max v [3, being the discovered and true coefficients (Messenger & Bortz, 2021). Exo (<) serves as the measure of accuracy of discovered coefficients.
TP
For 2D PDEs, we work with reaction-diffusion and incompressible Navier-Stokes equations. Ginzburg-Landau reaction-diffusion equations are frequently employed for modeling biological pattern development. We demonstrate discovery in two settings: 1) a relatively easy spiral dataset with a single spiral and 2) a harder setting with significantly more complex patterns where the baseline methods fail to recover the equations even in the absence of noise. We also demonstrate discovery for an incompressible Navier-Stokes example with a small viscosity term (Harpole et al., 2019). We set the viscosity term to 0.001 and show that our method is able to recover the equation while the baselines cannot. For baselines, et al., 2017) and WeakSINDy (Reinbold et al., 2020), with implementations by de Silva et al. (2020), which are known to be robust to noise in many settings. We especially focus on comparing the robustness under noise of our method with the baselines.
Datasets, experiments, baselines. We perform experiments
compare with PDEFIND (Rudy
Figure 2. Inviscid Burger's Equation (discontinuous shocks are visible) (left) and the vorticity field data V x U from an incom- pressible Navier-Stokk equation.
6

---

Mechanistic PDE Networks for Discovery of Governing Equations
Thue Positive Ratio
True Positive Ratio
550-
WEALSEND
Thee Poeitive Ratio
ASthed Meritto POr
88
MastNR-PUE Wa 64
Eree
WcAASIND
Nieaed MedaNN.PDD
1300
WeskSth
50 60
Percent Noise ONN
Percent. Noise di
Percent Naise
Percent Noise dNR
Figure 4. TPR for easy (left) and harder (right) reaction-diffusic experiments with varying amounts of noise.
Figure 6. TPR and Ess error for the Navier-Stokes dataset.
Easy Case. For the reaction-diffustio equations, we first test easy dataset with a single spiral pattern for which the baselines also perform well (Reinbold et al., 2020). For this example, the spatial domain has length of 20 with equal spatial step sizes and spatial resolution of 64x64 in the æ and y dimensions. For time we use 128 time steps with a step size of 0.05, An example of the data at a single time step is shown in Figure (left). The data was generated by spectral method following the implementation in de Silva et al. (2020), setting D1 D2 0.1 and B 1.
WeakSEND,
Erroz Ese
PDEPIND
Log Error
1.D
Wea PDEPIND
0.0
Percent Noise ONR
Percent Noise ONR
Figure 5. Log Ess errors for the casy (left) and harder (right)
daction-diffusion experiment with varying amounts of noise.
6.1. Solver Unit Test
A Harder Case. For reaction diffusion equations we also test on a more difficult dataset with more complex patterns, for which the baselines are unable to discover the correct equations even on noise-free data. The dataset has a spatial resolution of 256x256 over a Cartesian domain of size 10 with equal step size. The data was generated by using the Basilisk solver (Kenneally et al., 2020) starting from ran- dom initial condition and running with maximum time step size of 0.05. We set D1 D2 = 1 and B = 1.5. We take only 128 time steps from the data for training. An example of the data at a single time step is shown in Figure (right).
First, we verify that our solver is capable of solving some given linear PDEs. We test on the 2D Laplace PDE, v²u=0 on 32x32, 64x64, and 128x128 grids using the multigrid solver. We compare against a reference finite- difference implementation from the py-pde (PyPde) pack- age with sinusoidal boundary conditions obtaining I-2 errors of orders 10-1,5 5 x 10-2 and 10-2 2 respectively, showing the errors decreasing with increasing resolution. This simple evaluation shows that our solver is indeed capable of solving these equations to an accuracy comparable with traditional libraries while also being able to parallelize solving multiple equations on GPUs.
Training Data and Parameterization. We use the same training setup for both kinds of reaction-diffusion examples. with
For our model, the training data is mini-batch size of 8. Each data example in the minibatch is of size 32x32x32, where the first dimension is time. The data is parameterized by 10-layer 2D ResNets, where we consider the time dimension as batch dimension. The data consists of two scalar fields A", A', which are separately parameterized by neural networks before being used to build the PDE expression
6.2. PDE Discovery
In this section, we demonstrate PDE discovery for equations in one and two spatial dimensions in addition to time.
2D PDEs. For two dimensional PDEs, we test on Ginzburg- Landau reaction diffuuion equations and the incompressible huvier-Stokes equations. For 2D PDEs the PDEs are a set of coupled equations. For ease of training we demonstrate discovery on one equation from each of the coupled PDEs.
NN(A") and ē NN(A')
PDE Expression. The PDE expression model uses third degree polynomial functions built from the parameterized inputs u, 0 and second order derivatives in x and y. We use four separate polynomial expressions in the PDE: one serving as the coefficients of the term, another two for the Uga and Uyy terms, and one as the right hand side b. The polynomial basis parameters for discovery are computed by two-layer MLP with learnable input, which we find to improve training over a flat parameter vector. The boundary conditions are learned Dirichlet boundary conditions and are simply chosen to be the boundaries of ü. Given the PDE expression and the boundary conditions, the PDE expression
2D Reaction- Diffuusion Equations. The Ginzburg-Lannda reaction-diffusse equations are shown in Table in com- plex form in the variable A. For ease, we write the complex equations as a system of two equations, where we take the first equation in our experiments.
V2 A D, 2 + A [A + 3A
(13)
A: = D2V² + A'(1 -
(14)
where A and A are the real and imaginary parts of A and |4|2 = Ar² + 2 and Dj, D2, 3 are parameters.
7

---

Mechanistic PDE Networks for Discovery of Governing Equations
is fed to the solver to get solution u. Loss. The final loss is an L1 loss, which we find improves discovery. The loss is composed of four terms. The first is a loss (u, A"), minimizing the loss between solution and data. The other two terms are losses for the parameteriza- tions 7 and e as (2(ū, u) and (ū, A'). Since we only solve time, we do not have learned solution to the second coupled equation and minimizes the loss directly with the data A'. The final term in the loss is a spar- sity term, also L1 for the basis polynomial parameters. We use a weight of 10 4 for the sparsity term. In particular, we note that we do not threshold the parameters or repeat the op- timization in our experiments with thresholded parameters, using only a single training cycle. Repeating the optimiza- tion can and does improve the accuracy of the learned co- efficients, but also increases training time, which is why we do not repeat the optimization for the results in this paper.
(Figure 2). A comparison of the results with Wya is shown in 6, where WeakSINDy is unable to discover the system. For MechNN-PDE, the TPR is still lower than I with low noise. On inspection of the coefficients, it appears that this is due to the presence of zero terms with small coefficients that are of the order of the viscosity which appears as false positives, even though the viscosity and non-zeros terms are accurately recovered.
one equation
1D PDEs. We demonstrate discovery the ID diffusion equation and inviscid and viscous Burger's equations. We use degree 4 polynomials for the coefficients of the uy and Uaa terms. For all cases in the noise-free case the method discovers the true equation. For Burger's s equation we exper- iment with noisy data and compare with WeakSINDY and find both methods to be on-par with each other. The results for the are shown in Figure 7 and Appendix A.
True Positive Ratio
MediN--DDE
MecakN-PPDE
The model is trained with Adam with a learning rate of 10 which we fix for all experiments. We do not use learning rate scheduling for the experiments for simplicity and use a fixed learning rate throughout. We repeat the training with up to 80% noise added to A" and A'.
WeaESINIly
WeakStNit
Error
u22
00
00-
40
Results. The TPR results are shown in Figure 4 and the Es errors are shown in Figure 5 for the various quantities of noise for both the easy and hard experiments for MechNN- PDE, WeakSINDy and PDEFIND. For clean data in the easy example, we see that all methods recover the non-zeros terms with a TPR of with small error in the coefficients.
Percent Noise yy
Percent Noise dyr
Figure 7. TPR and error for viscous Burger's 's equation. Comparing with WeakSINDY
6.2.1. DISCOVERY BEYOND GENERALIZED LINEAR
Upon increasing noise levels, we find that PDEFIND wors- ens the most, and the TPR drops significantly to 0-0.2 for the easy case while the error becomes very large. WeakSINDy performs much better with noise on the easy dataset, and MechNN-PDE and WeakSINDy have similar TPR and error with noise except for 80% noise, where WeakSINDy has better error for similar TPR as MechNN-PDE.
MechNN-PDE can represent complex differentiable expres- sion in PDEs, going beyond expressions representable as linear combinations of fixed hasis functions. We demon- strate with the task of parameter discovering for the Porous Medium equation for which the expression cannot be writ- ten as a linear combination of fixed basis functions. The porous medium equation is given 4 (w")), m > O, and we want to recover the parameter m given data. We set m = 2.675 in the data and do not use any other basis functions. We note that m is a positive real value which precludes fixed polynomial basis functions. The task is to discover the exponent of u. We build the PDE expression with an exponent for the parameterized data ü. For this example we also parameterize the gradient with a neural net- work. With clean data recover the coefficient 2.64+0.03 close the true value of 2.675.
For the harder case, neither PDEFIND nor WeakSINDy is able to find any significant non-zeros terms even with clean data. The TPR seems to increase later on for PDEFIND, however, this is spurious due to the large error as seen in the plots. WeakSINDY also performs no better than random in this setting. MechNN-PDD PDE, however, able to maintain high TPR with noise for up to 40% noise, with the TPR drop- ping to 0.6 80% noise, showing consistent performance in TPR and error with noise.
Incompressible Navier-Stokes. We use the first component of the velocity U, from the coupled equations for discovery given as U1+ +U- VU1 + Pu = ~V2 where the pressure gradient Vp is given. We use similar setup for the reaction-diffusion equations, except that we use degree polynomials. We include interaction terms with the first derívatives. We use small viscosity 0.001 for the data
6.3. Multigrid Hyperparameters, Time, Memory
We show the effect of varying hyperparameters of the multi- grid algorithm on solving the Laplace equation (Figures 8, 9, 10 in Appendix C). The hyperparameters varied are the number of FGMRES iterations for two resolutions, the number of V-cycles and the number of Gauss Seidel steps.
We measure relative error and show decreasing error with
8

---

Mechanistic PDE Networks for Discovery of Governing Equations
the number of iterations.
References
In Figure 11 (Appendix C) we show the time and memory used in solving the Laplace equation with varying resolu- tions.
Amos, B. and Kolter, J. Z. Opinet:
Differentiable opti- mization as layer in neural networks. In International Conference on Machine Learning, pp. 136-145. PMLR, 2017.
7. Conclusion
Brandstetter, J., Welling, M., and Worrall, D. E. Lie point symmetry data augmentation for neural pde solvers. In International Conference on Machine Learning, pp. 2241- 2256. PMLR, 2022.
This paper presents a new method for discovering governing PDEs from data. We build specialized multigrid parallel and differentiable solver for solving PDEs in batch which is then incorporated in a discovery architecture. Experiments show that the method can discover complex equations from data while being robust to noise and expands the class of PDEs that can be modeled by machine learning methods.
Briggs, L.,, Henson, V. E., and McCormick, S. E. multigrid tutorial. SIAM, 2000.
Brunton, S. L., Proctor, J. L., and Kutz, J. N. Discovering governing equations from data by sparse identification of nonlinear dynamical systems. Proceedings of the national academy of sciences, 113(15):3932-3931 2016.
Limitations. Scalability is a main limitation of grid-based PDE methods and way forward might be to explore combination of grid-based and mesh-free techniques. An- other shortcoming is that the method sometimes sacrifices accuracy for speed. This could be resolved by better multi- grid solvers and faster GPU implementations. Finally, we do not consider the identifiability of the PDEs being discov- ered and the conditions under which discovery is possible. Extending the analysis of Yao et al. (2024) to PDEs would be an interesting starting point.
Brunton, S. L., Noack, B. R., and Koumoutsakos, P. Ma- chine learning for fluid mechanics. Annual review of fluid mechanics, 52(1):477-508, 2020.
Chen, J., Yao, D., Pervez, A., Alistarh, D., and Locatello, F. Scalable mechanistic neural networks. International Conference on Learning Representations, 2025.
Chen, R. T., Rubanova, Y., Bettencourt, J., and Duvenaud, D. K. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
Acknowledgements
AP. This project has received funding from the European Union's Horizon 2020 re- search and innovation programme under the Marie 3ktodowska-Curic Grant Agreement No. 101034413.
d' Ascoli, S., Becker, S., Schwaller, P., Mathis, A., and Kil- bertus, N. Odeformer: Symbolic regression of dynamical systems with transformers. In The Twelfth International Conference on Learning Representations, 2023.
FL. This research was funded in whole in part by the Austrian Science Fund (FWF) 10.55776/COE12. For open access purposes, the author has applied a CC BY public copyright license to any author accepted manuscript version arising from this submission.
de Silva, B., Champion, K., Quade, M., Loiseau, J.-C., Kutz, J. N., and Brunton, S. Pysindy: A python package for the sparse identification of nonlinear dynamical systems from data. Journal of Open Source Software, 5(49):1-4, 2020.
Impact Statement
Harpole. A., Zingale, M., Hawke, I., and Chegini, pyro: framework for hydrodynamies explorations and prototyping. Journal of Open Source Software, 4(34): 1265, 2019. doi: 10.21105/jjss.01255. URL https: / //doi. 10 10.00.211000 20101 /ww.1/////oss.... 01265.
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
Kaheman, K., Kutz, J. N., and Brunton, S. L. Sindy-pi: a robust algorithm for parallel implicit sparse identification of nonlinear dynamics. Proceedings of the Royal Society A, (44662222)200200277) 2020.
Kenneally, P. W., Piggott, S., and Schaub, H. Basilisk: A flexible, scalable and modular astrodynamics simulation framework. Journal of aerospace information systems, 17(9):96-557) 2020.

---

Mechanistic PDE Networks for Discovery of Governing Equations
Lagergren, H., Nardini, J. T., Michael Lavigne, G., Rut- ter, F.. M., and Flores, K. B. Learning partial differen- tial equations for biological transport models from noisy spatio-temporal data. Proceedings of the Royal Society A, 476(2234)22019088) 90800, 2020.
Temam, R. Navier-Stokes equations: theory and numerical analysis, volume 343. American Mathematical Society, 2024.
Turing, A. The chemical basis of morphogenesis. Philo- sophical Transactions of the Royal Society of London. Series B, Biological Sciences, 237(641):37 1952.
Li, Z., Kovachki, N. B., Azizzadenesheli. K.,, Bhattacharya, K., Stuart, A., Anandkumar, A., al. Fourier neural operator for parametric partial differential equations. In International Conference on Learning Representations, 2020.
Yao, D., Muller, C., and Locatello, E. Marrying causal representation learning with dynamical systems for sci- ence. Advances in Neural Information Processing Sys- lems (NeurIPS), 2024.
Messenger, and Bortz, D. M. Weak sindy for partial differential equations. Journal of Computational Physics, 443:110525, 2021.
Young, D. Linear programming applied to linear differen- tial equations. 1961.
Pervez, A., Locatello, F., and Gavves, S. Mechanistic neural networks for scientific machine learning. Forty-first International Conference on Machine Learning, 2024.
PyPde.
sttps://githbb.co com/ Swicker-growe Patppppy-pp.....pp pde.
py-pde.
Rackauckas, C., Ma, Y., Martensen, J., Warner, C., Zubov, K., Supekar, R., Skinner, D., Ramadhan, A., and Edelman, A. Universal differential equations for scientific machine
learning. arXiv preprint arXiv:2001.0443 2020.
Raissi, M. Deep hidden physics models: Deep learning of nonlinear partial differential equations. The Journal of Machine Learning Research, (0(1):932-955. 2018.
Raissi, M., Perdikaris, P., and Karniadakis, G. E. Physics- informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational physics, 378:686-707, 2019.
Reinbold, P. A., Gurevich, D. R., and Grigoriev, R. O. Us- ing noisy or incomplete data to discover models of spa- tiotemporal dynamics. Physical Review E, 101(1):010203, 2020.
Rudy, S. H., Brunton, S. L., Proctor, J. L., and Kutz, J. N. Data-driven discovery of partial differential equations. Science advances, 344)ee66 2017.
Saad, Y. Iterative methods for sparse linear systems. SIAM, 2003.
Saad, Y. and Schultz, M. H. Gmres: A generalized mini- mal residual algorithm for solving nonsymmetric linear systems. SIAM Journal on scientific and statistical com- puting, ((3):856-861) 1986.
Stephany, R. and Earls, C. Pde-learn: Using deep learning to discover partial differential equations from noisy, limited data. Neural Networks, 174:106242, 2024.
10

---

Mechanistic PDE Networks for Discovery of Governing Equations
A. 1D PDE Experiments
We test discovery on clean data on the diffusion equation with diffusion coefficient 0.01. We obtain a TPR and Eso value of 0.224. We show results for the viscous and inviscid Burger's equations in Figure and Table 2, comparing with WeakSINDy. We find both methods perform similarly on this problem.
Table 2. Inviscid Burger's Equation for clean and 10% noisy data.
MechNN-PDE
WeakSINDY
Noise
TRP
Eoo
TRP
E.c
1
0.043
0.001
0.0020
0.012
B. Multigrid V-cycle Algorithm
Algorithm shows the multigrid V-cycle (see Saad (2003)) that we use to precondition the FGMRES algorithm. We use Gauss-Seidel the relaxation method and linear interpolation the restriction and prolongation operators.
Algorithm 1 V-cycle
Require: relax, restrict, prolong operators
Input: M,, x,, by E [m]]
zy relax(A,, b,)
r b, A,
Ti+1 restrict(r))
if m = i 1 then
Solve Am 6m = I'
else
8, = 1, 0, ri+1)
end if
8, prolong(o++))
T; 2; + 5;
x; re relax( A4, 14, bi)

---

Mechanistic PDE Networks for Discovery of Governing Equations
C. Multigrid Hyperparameters and Memory
show effect varying hyperparameters of the multigrid algorithm on solving the Laplace equation (Figures 8, 9, 10). The hyperparameters varied the number of FGMRES iterations, the number of V-cycles and the number of Gauss Seidel steps.
In Figure 1l we show the time and memory used in solving the Laplace equation with varying resolutions.
x10-1
x10
Laplace 32x32
2.0
Relative Error
Laplace 64x64
Laplace 32x32
Relative Error
3:
2
1.0
0.5
40
60
80
100
120
0.0
3 Number of V-cycles
5
FGMRES Iterations
Figure 8. Relative error in linear system solving versus number FGMRES iterations. The number of FGMRES restarts is fixed to 20. The number preconditioning V-cycles is Showing results for solving Laplace equation on 32x32 and 64x64 grids.
Figure Relative error linear system solving versus number of preconditioning V-cycles with FGMRES. The number of FGMRES restarts is fixed 20 with iterations. Showing results for solving Laplace equation on 32x32 grid.
x10-3
Laplace 32x32
Relative Error
4
3
Figure 11. Time and GPU Memory usage for solving 2D Laplace equation for batch size of 32 with multigrid pre- conditioned FGMRES.
Grid
Memory (GB)
Time (seconds)
32x32
1.6
6.16
2
8
10
64x64
4.63
13.9
Gauss Seidel Steps
128x 128
16.7
33,9
Figure 10. Relative error linear system solving versus num- ber of Gauss-Seider cycles per preconditioning V-cyele with FGMRES. The number of FGMRES restarts is fixed to 20 with 40 iterations with V-eycle per preconditioning V-cycle. Showing results for solving Laplace equation a 32x32 grid.
12


Reviews:
[{"review_id": "aE9hF66mNX", "reviewer": "Reviewer_2aLy", "summary": "The paper introduces a new methodology for learning PDEs on data. The key contribution is an optimization of how partial derivatives are handled in the network. Firstly, the theoretical formulation includes a dual formulation that elides a way to backpropagate through a linear solve effectively. Secondly, they introduce an optimized linear solving using a multigrid algorithm implemented on GPU. The paper then demonstrates that their method learns a few test PDEs more accurately than baseline PDE discovery methods with various levels of noise.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the detailed review.\n\n**Clarification of our contribution.**\n\nWe wish to clarify the nature of our contribution. We do not claim to replace or improve existing PDE solvers. Our contribution is to enhance the MechNN model (Pervez et al, ICML 24) that works with ODEs to support PDE representations, together with an efficient way of solving the representations and applying the model for PDE discovery. The PDE solver we develop is specialized for this architecture, where the PDE terms are produced by a (non-smooth) neural network and solves a relaxation. The sparse solver is intended to make the MechNNs *feasible* for PDEs with multiple dimensions.\nThat said, we do not discount the possibility that the method could be useful for other applications.\n\nFurthermore, we make no claims of either speed or memory improvement over the baseline sparse regression methods which require much fewer resources. We do claim that our model is more expressive (handles complex expressions) and can handle more complex data as shown in the discovery experiments.\n\n**Missing Evidence and Failure.**\n\nWe provide quantitative evidence of discovery performance using the TPR and Error metrics from the literature. A TPR of 1 means exact discovery of terms. This is shown in Figures 4,5 and 6 for our method and the baselines. The plots also show partial failure to discovery when noise is added. For instance with 80% noise (Figure 4) we only see a TPR of 0.6 which indicates extra or missing terms are present.  From Figure 6 we see that the Navier-Stokes equation is also not exactly discovered (TPR ~0.8). We present quantitative metrics since they allow more objective comparison than qualitative comparison between equations. We will include examples of discovered equations.\n\n**Solver Eval**\n\nWe have performed further evaluation of the multigrid linear solver including time and memory usage, convergence of relative error with the number of V-cycles, GS smoothing iterations, and FGMRES iterations for the 2D Laplace equation. See https://filebin.net/wqv4t3oki732nuq0 for plots.\n\n**GPU Multigrid**\n\nGPU multigrid methods are frequently used in PDE solvers. However, our solver is specialized to work with (possibly non-smooth) neural network outputs over the grid that parameterize the PDEs. The non-smoothness is handled by adding extra smoothing constraints and solving relaxations as saddle-point problems. We are not aware of available solutions for the special case that we consider.\n\nWe plan to release the code for our experiments.\n\n**Comments**\n\nBy grid size we mean increasing the grid resolution for the same problem and running until convergence. We will make the correction.\n\nThe line about derivatives is with reference to the SINDy baselines which estimate numerical derivatives directly on data which makes them susceptible to noise. There is a typo and the word ‘computed’ is missing. \n\nThe statement about the parameter $m$ was meant to indicate that it is not a fixed value which implies that fixed basis functions cannot be used to represent the PDE. \n\nWe will make the corrections in the next draft.\n\n**Questions**\n\nThe V-cycle algorithm itself is fairly standard, however, it is applied in a non-standard setting to solve saddle-point problems in eq 10 and 11. For these problems the standalone V-cycle has subpar performance in our setting. However, we were able to obtain better results using the V-cycle as a preconditioner for FGMRES.\n\nYes, we also solve the backward pass with the V-cycle (with FGMRES).\n\nWe will include more details about the solver implementation in the next draft. The code is in Pytorch together with CuPy for features not available in Pytorch (such as sparse triangular solves). We did not write CUDA kernels, however, that is an option for further improvement. For FGMRES we start with the CPU implementation of GMRES from SciPy. We adapted the GMRES algorithm to FGMRES using the algorithm described in Saad (2003). We ported the code to Pytorch adding batching and GPU support and V-cycle preconditioning, written from scratch.\n\nFor simplicity we only threshold once at the end of training (unlike SINDy). The threshold is a hyperparameter per dataset (fixed across noise). We used 0.06, 0.3 and 0.0002 for the easy and hard reaction diffusion, and Navier-Stokes equations. Fine-tuning after thresholding can be used to obtain more accurate coefficients.\n\nThe parameterization with a ResNet refers to computing \\tilde{u} (line 267 or Figure 1). We parameterize the PDE terms (Figure 1) by a neural network (using \\tilde{u}) instead of feeding data (u_data) (which may be noisy) directly. This also helps with training convergence. \n\nLine 411 discusses discovery in the presence of noise. Results are shown in Figure 4 (left). WeakSINDy and MechNN have similar performance with at least TPR 0.6 whereas PDEFIND fails entirely.\n\nWe thank the reviewer again.", "claims_and_evidence": "There are three distinct dimensions that the paper improves:\n- Flexibility by allowing nonlinear backpropagation with PDE terms embedded.\n- A very fast and efficient GPU implementation of a multigrid V-cycle algorithm for linear solving.\n- Improved accuracy  on the PDE learning task.\nThe paper only demonstrates evidence of the improved accuracy.\n\nA key improvement purported by this paper is supposedly the scaling and optimization of the new solver. However, no performance metrics or comparisons are reported. There is only theoretical calculations of memory utilization. It would greatly strengthen the paper to show speed and memory usage comparison with the baseline methods, as a function of the domain size.", "methods_and_evaluation_criteria": "As mentioned, evaluation of the purported claims of a fast solver are missing.\n\nThe results are also missing evidence of the discovered PDE forms or failure to discover the forms; see line 418 on page 8.", "theoretical_claims": "I did not see any issues in the theoretical formulations in the paper.", "experimental_designs_or_analyses": "I did not see any issues in the comparison between PDE learners sound.\n\nAs mentioned, there could be more elaboration on the solver efficiency. There is a brief discussion of solver convergence in Section 6.1, but this does not have any data nor figure references. It is great to discuss testing methodology at the forefront, but this leaves some questions: What is meant by errors decreasing with increasing grid size? Shouldn’t it go the other way around? What is the convergence rate? Is this the extent of the testing of the code? What is the convergence rate in terms of iterations of the V-cycle solver?", "relation_to_broader_scientific_literature": "As mentioned, the novel way of representing PDE solving in nonlinear networks coupled with a very efficient solver will be a contribution to the literature. Such an algorithm could be dropped in to many other PDE learning methods beyond the proposed MechNN-PDE.", "essential_references_not_discussed": "Are there more papers about multigrid linear solvers implemented on GPU, even looking beyond the ML literature into scientifc computing? Is this the first time a GPU implementation of the multigrid method was published?", "other_strengths_and_weaknesses": "I am supportive of the paper, but it is missing more discussion about the advantages of the solver technique. The paper is not sufficiently demonstrating the scaling and speed improvements. It mentions the theoretical scaling, but does not show the application to problems that the other PDE learning techniques cannot handle. As mentioned above, this would strengthen the contribution of the paper. That aspect would have the potential to be applied to other PDE discovery methods as a drop-in module, and even be adapted to other ML problems requiring a similar operation.\n\nI also think that the method also has fewer constraints than the baseline methods: whereas e.g. WeakSINDy does a sparse linear regression, MechNN-PDE is doing backpropagation over the nonlinear expressions.", "other_comments_or_suggestions": "- Page 2: “Derivatives never have to be directly on data and are” Firstly, is there a mistake in this sentence? Secondly, what does it mean? Which derivatives? Do you mean parameter tangents, or partial derivatives?\n- Page 2: Is the discretization only done for 1D? The paper has 2D problems: could you write the math in a higher dimension instead?\n- Page 8: “We note that m is a positive real value which precludes..” What is that supposed to mean?\n- Define FGMRES and GMRES for the audience.\n- I don’t understand the paragraph headings that are bold versus italic on Page 5 and 7,8. It is not obvious which is higher or lower.", "questions_for_authors": "- What aspects of the multigrid V-cycle section are new in your paper, on page 4? Did you propose any changes to the algorithm, or is this just the standard V-cycle?\n- Do you have to solve the backward pass in Equation 11 with the multigrid solver as well?\n- Is the FGMRES solver also your implementation, or do you integrate with another library? Which library, or how did you implement it?\n- A few more details on the GPU implementation would be good for the paper, too. E.g., did you hand write CUDA kernels? For which library?\n- Page 5: “Finally a concise PDE is generated by thresholding the parameters...” How is that done? Is the thresholding applied iteratively in a loop with parameter optimization, as in SINDy?  What is the thresholding parameter used?\n- Page 7: “The data is parameterized by 10-layer 2D ResNets”: What does this mean? Which data, which parameterization? How does that fit in? What are the details of the architecture? Are they CNNs?\n- Page 8, paragraph on line 411: Are the methods able to discover the equation forms here?\n- The usage of the neural network for $\\tilde{u}$ in the MechNN-PDE architecture is unclear. Would it be possible to include it in Figure 1?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "6Xw2zkhTVN", "reviewer": "Reviewer_Ct17", "summary": "The paper presents mechanistic PDE networks, a method for discovering PDEs from spatiotemporal data. The proposed method integrates a differentiable PDE solver into the neural network and uses the neural network predictions to model the PDE instead of the raw data. The method discovers PDEs using the spatiotemporal data by expressing the PDE as a learnable polynomial kind of ansatz. The paper applies the proposed method to discover popular PDEs in the neural PDE solver domain and compares them with WeakSINDy and PDEFIND methods.\n\n## Update after rebuttal\n\nThe paper presents advancements in discovering equations from data. The paper also presents the limitations of baselines and mitigates the challenges therein for 2D problems. Hence, I have raised my score considering the author's responses.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the comments. We attempt to address the concerns raised below.\n\n**Non-autonomous systems**\n\nYes, in this paper we focus on autonomous systems. However, non-autonomous systems can be similarly handled since the method allows spatiotemporally varying terms. Any known external forces can be represented as time-space dependent terms similarly to what is achieved with basis functions. Any unknown forces that cannot be simply represented can be represented by a neural network. This is possible because the method allows arbitrary differentiable expressions in PDEs.\n\n**Identifiability**\n\nWhether the method finds the unique PDE requires further investigation into identifiability of the PDE learning problem. Identifiability is something we do not deal with in this paper and assume that the problems are identifiable and that there is a unique governing equation. \n\nWe have not considered stochastic or fractional differential equations thus far and it would be interesting to consider whether such types of equations can be modeled in this way. \n\n\n**PINN style methods**\n\nThe PDE-LEARN approach extends the PINN approach to PDE discovery. It reduces a PDE residual, constructed using basis functions, at random collocation points together with approximate L0 norm regularization. One difference of this approach with MechNN is that MechNN has a stronger physics informed bias due to the explicit modeling of PDEs as constraints.\n\nUnfortunately, the PDE-LEARN release does not have support for coupled equations, so it cannot be used with our 2D experiments in its current form. The main equations that we focused on were the 2D reaction diffusion and Navier-Stokes PDEs which are coupled equations.\n\nNevertheless, we trained PDE-LEARN on our 1D Burger’s equation dataset with a 1/cosh(x) initial condition. Despite repeated attempts with multiple settings of hyperparameters, we found it to always over or under prune the equation (even on clean data), performing worse than the SINDy baselines from our paper. The maximum TPR we obtain for this experiment is 0.5, compared with a TPR of 1 for the SINDy baselines and MechNN.\n\nIn particular one advantage of our method is that the modeled PDEs can contain arbitrary differentiable expressions and are not limited to linear combinations of fixed basis functions. Many other methods including PDE-LEARN model PDEs as linear combinations of fixed basis functions. A concrete example is the porous medium equation for which the equation cannot be modeled in this way, whereas we show that we can recover the true parameters.\n\n**Limitations**\n\nThe following are some limitations of the method.\n\n* The method is slower than the baseline sparse regression methods.\n*  The method is currently limited to Cartesian grids. \n* Although the method can represent arbitrary differentiable expressions, there remains the need for theory that dictates the conditions under which complex expressions can be exactly identified.\n* There is a tradeoff between accuracy and speed in the multigrid solver which needs to be tuned.\n* Application to very high dimensional data is not feasible in the current form.\n\nWe thank the reviewer again for the review and hope that we were able to meet any concerns.", "claims_and_evidence": "The claims in the paper regarding applicability of the proposed method for noisy data is validated through experiments.", "methods_and_evaluation_criteria": "The experimental setting regarding dataset and noisy data draw similarity with conventional methods employed for PDE discovery.", "theoretical_claims": "Not applicable", "experimental_designs_or_analyses": "Yes, I checked for all provided numerical experiments", "relation_to_broader_scientific_literature": "Discovering of governing equations from data is an important topic in many field of science and engineering. The proposed AI method is in the direction to model these problems motivated through canonical problems.", "essential_references_not_discussed": "References are adequate.", "other_strengths_and_weaknesses": "Strengths:\n\nThe paper presents a novel method for discovering PDEs using spatiotemporal data, where the derivative computations do not need to be performed on the data but are instead performed on the neural network outputs.\n\nThe experiments show that the proposed method is robust to noise and outperforms WeakSINDy and PDEFIND.\n\nThe proposed method is scalable and the polynomial kind of ansatz enables generalized PDE operators.\n\nWeaknesses:\n\nThe comparison with baseline methods is restricted to specific methods, and the paper does not compare it with recent methods.\n\nThe memory requirements increase with an increase in grid size. This is problematic for high-dimensional PDE discovery.\n\nThe computational cost of the method has a trade-off with accuracy, as also mentioned by the authors in the paper.", "other_comments_or_suggestions": "Typo:\n\nIn line 95, should the product be taken for two different spatial coordinates?\n\nSuggestions:\n\nThe authors might consider including a discussion on how the proposed method is advantageous over using rational neural networks for PDE discovery.", "questions_for_authors": "The paper does not consider PDEs under external acting forces and only showcases their applicability for autonomous systems. How would the proposed method be applied to discover the PDE in case of unknown and known external acting forces?\n\nHow does the method guarantee the uniqueness of the identified PDE? How will the method behave when multiple PDEs govern the data equally well?\n\nExtending the proposed method to discover stochastic PDEs or fractional differential equations seems complicated. How does one consider in priori the correct dynamical system framework modeling the dynamics?\n\nIt would be informative if the authors compare the proposed method with the rational neural networks approach of discovering governing PDEs, for instance, as presented in PDE-LEARN.\n\nWhat will be the limitations and failure modes of the proposed method?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "c5O3XzAZ2w", "reviewer": "Reviewer_coX1", "summary": "This paper presents a model that learns spatial-temporal PDEs from data samples. The model selects a set of basis functions with spatially and temporally varying coefficients over the problem domain. Next, it implements a multigrid solver to solve the proposed PDE. Finally, the model learns a PDE by backpropagating the gradients and optimizing a loss function on data terms. The paper demonstrates the method’s efficacy on classic 1D and 2D PDE instances.\n\n## update after rebuttal\n\nThank you for your rebuttal. I didn't have too many questions to ask in my initial review, and the rebuttal didn't greatly affect my overall evaluation of this work. I will maintain my current score.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the comments.\n\nThe reviewer’s point regarding use of the term ‘discovery’ is well received. However, the term is now endemic in the literature and we chose to employ the same term.\n\nWe would like to appreciate the reviewer’s recognition of the effort that went into building the solver. We will certainly release the source code for our models with the camera ready version of the paper.\n\n**Multigrid**\n\nIt is quite correct that multigrid solvers are also used for higher resolutions. Our use case is different in a few respects from the usual use of multigrid.\n*  One difference is that we solve for the saddle point systems in equations 10,11 and we have extra constraints for ensuring smoothness which are usually not present in solvers.  \n* Secondly, we work in a learning context with a nested iterative loop. The inner loop corresponds to the multgrid iteration and the outer loop is the learning loop that can run for hundreds of epochs. \n* Similarly we work in a batch setting where in each inner iteration we solve a minibatch of PDEs, all with their separate multgrid steps. \n* Furthermore the backward pass also uses the batched multigrid solver which doubles the memory use. \n* It is known that multigrid performance deteriorates with discontinuous coefficients. The data produced by neural networks over the grid (especially early in training) can be non-smooth which can deteriorate multigrid performance. \n* Our current implementation is in Pytorch which is not ideal for iterative processing. A native CUDA implementation can potentially produce significant computational improvements, but requires further work.\n\nVery high multigrid resolutions would have an adverse impact on the walltime, making the method very slow. Nevertheless, different from standard multigrid approaches, since we also have a learning component which can infer the underlying PDE given only a (batched) portion of the full resolution data. This implies that the solver is not required to solve for the full resolution data in each step. Rather it only has to solve for the given smaller resolution patch of data making it more efficient. \n\n**Encoder NN**\n\nThe neural network transform (line 267) allows parameter capacity to the model to make it flexible. Without the transformation the only learnable parameters would be the PDE parameters, parameterizing the expression, which are few in number. This can lead to a brittle model making it harder to recover from bad initialization or noise. We will consider adding an ablation in the next draft.\n\n $x_1 x_1$ should be $x_1 x_2$. Yes, that is a typo in this case.\n\nWe would like to thank the reviewer again for the review.", "claims_and_evidence": "Most of them look OK to me.\n\nI have one comment on the claim of “discovering” PDEs from data. I notice that the experiments generated data from known PDEs. I understand that having a ground-truth PDE is good for calibrating the efficacy of the “discovery” of PDEs, and many prior works used this setting as well. However, I feel it would be more proper to use the term “discovering” PDEs if the data are from a (real-life) experiment without a known PDE. Perhaps “recovering” or “reconstructing” a PDE is a more appropriate term for what is going on in this paper. I understand that this might be an unpopular opinion, and I won’t hold this against the paper.", "methods_and_evaluation_criteria": "Most of them look good to me. In particular, I want to commend the effort to build a multigrid solver. The authors definitely deserve some credit for this.\n\nI don’t quite get the neural network mapping (line 267) and would like to suggest an ablation study on it.", "theoretical_claims": "N/A.", "experimental_designs_or_analyses": "I read the experiments, and I think the results look reasonable.\n\nI have less experience with the baselines in the paper, so I am not sure whether they are the SOTAs for comparisons and will let the other reviewers decide. I also wonder whether PINNs and their variants are baselines for this work and would like to hear the authors’ thoughts on them.", "relation_to_broader_scientific_literature": "I think the proposed method overall is new and interesting. I can think of some related works that tackle similar problems or use similar techniques, none of which can challenge the novelty of this paper.", "essential_references_not_discussed": "Looks good to me. I don’t have extra references to suggest.", "other_strengths_and_weaknesses": "I want to second the multigrid aspect of this work. Quite a few previous works evaluate their methods on small-sized toy examples only. I highly appreciate that this paper is willing to implement a multigrid solver in this setting, which has the potential to scale up the problem size.\n\nOn the negative side, the problem size (256 x 256?) in this work is still relatively small compared with what modern multigrid solvers can achieve (i.e., solving linear systems with > millions of variables). I feel the current method hasn’t unleashed the full power of multigrid solvers, and I am curious to know which part in the method is now the bottleneck that prevents it from being applied to large-scale problems.", "other_comments_or_suggestions": "A minor comment: Line 95: is the notation x1x1 a typo? Based on the context, x1x2 seems to make more sense.", "questions_for_authors": "None for now. I look forward to a constructive reviewer-author discussion.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xtEExtInbb", "reviewer": "Reviewer_Zn4e", "summary": "This work extends mechanistic neural networks to discover PDEs. The approach relies on a numerical solver for PDEs, specialized for linear PDEs; however, the approach may still be applied to nonlinear PDEs by using nonlinear basis functions. Solutions are constructed as simple trees, which are restricted to be sparse and concise by Lasso regularization. \n\nThe method is as follows. Given data from some (perhaps set of) PDE(s) and a skeleton of the PDEs, an encoder produces a prediction of the solution $\\tilde{u}$. Next, coefficients of a polynomial combination of $\\tilde{u}$ are predicted and incorporated into the PDE structure. The constructed PDE is solved using the NeuRLP-PDE solver to compute $u$. Finally, a loss is calculated between the data and $u$, as well as $u$ and $tilde{u}$, and backpropagated through the solver and encoder.  The authors introduce a sparse solver with efficient implementation on GPUs, which serves as a backbone for the proposed approach.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer’s comments.\n\n**Main Contribution and Paper Structure**\n\nWe would like to strongly emphasize that we do not claim to replace or improve existing PDE solvers. Our contribution is to enhance the MechNN model (Pervez et al, ICML 24) that works with ODEs to support PDE representations, together with an efficient way of solving the representations and applying the model for PDE discovery. The PDE solver we develop is specialized for this architecture, where the PDE terms are produced by a (non-smooth) neural network output and solves a relaxation. The sparse solver is intended *to make MechNN feasible for PDEs* with multiple dimensions.\n\nFurthermore, we make no claims of either speed or memory improvement over the baseline sparse regression methods which are not neural networks and require much fewer resources. We do claim that our model is more expressive (handles complex expressions) and can handle more complex data as shown in the discovery experiments.\n\nWe are open to restructuring the paper for clarity as long as it correctly emphasizes our main contributions and will certainly consider the reviewer’s suggestions.\n\n**Plots look similar**\n\nWe disagree with the reviewer’s assessment of the plots. We claim significant improvement on the more complex reaction diffusion system (Figures 4,5 right) over the baselines which entirely fail for this dataset.\n\nWe note that the TPR and max coefficient error plots should be considered together, since the correct terms have to be discovered for the coefficient error to be meaningful. The baselines fail in all cases to discover the correct terms (with low error) including the case with no noise. For our model both TPR and coefficient error show gradual degradation with noise, with near perfect recovery in low noise settings even though we do not fine-tune after thresholding.\n\nOn the easier reaction diffusion system the WeakSINDy baseline performs well (although PDEFIND fails with noise) and our method matches its performance, with only some decrease in coefficient accuracy at the 80% noise level.\n\n**Evaluation/Simulation**\n\nThe metrics of TPR and max coefficient error are chosen from the literature. In our view they provide a more quantitative and objective basis for comparison of discovery methods. Raw equations allow limited insight, especially when exact discovery is not possible in noisy settings. However, we will consider adding examples of discovered equations in the updated draft.\n\nSimilarly, simulating a discovered equation does not serve as an accurate determinant of how well an equation’s form has been learned, especially in noisy settings. Nevertheless we link a video example of simulation for the harder reaction diffusion example simulating the ground truth and discovered equation on clean data (uploaded here: https://filebin.net/wqv4t3oki732nuq0). We will consider adding more examples in the next draft.\n\n\n**Burger’s equation**\n\nWith the inviscid Burger’s equation there is only one term to be discovered with the coefficient of 1. However, since we do not finetune after thresholding there can be occasional variation in the discovered coefficients. The learned coefficient in the noise-free case for this experiment is 0.957 which gives a relative error (|1-0.957|/|0.957|) of 0.04. For the noisy case the coefficient is 1.01 with a relative error of 0.001. Running the experiment with fine-tuning after thresholding gives a small relative error (~0.001) for the first case as well.\n\n**Recent Methods**\n\nDue to lack of space, please see our response to reviewer Ct17 for details of the comparison with PDE-LEARN.\n\n\n**Hyperparameters**\n\nWe clarify that the baselines (PDEFIND, WeakSINDy)  are sparse regression methods which do not employ any neural networks and therefore the learning rates etc, do not apply to the baselines. For our method, for simplicity, we fixed the learning rate across all our experiments and do not use any scheduling.\n\n**Retraining**\n\nYes, the models have to be retrained for new data.\nIn the discovery setting we consider, the PDE parameters to be discovered are global parameters and are not functions of the data. This is also true for the baselines. This setting does not allow direct evaluation of parameters on a new dataset. However, developing a setting where this is sensible is an interesting question for further work.\n\n**Threshold**\n\nWe chose the threshold per PDE qualitatively, for simplicity. The choice of threshold can affect small parameters which is also the reason why the TPR in the Navier-Stokes example is less than 1.\n\n**Analytical solutions** \n\nMost PDEs don’t have analytical solutions, but perhaps the question is with regard to learning the expression structure? \nIf so, then yes, we think that symbolic regression methods could be combined with this method to learn expressions together with parameters and we are considering approaches for this.\n\nWe would like to thank the reviewer again for the review.", "claims_and_evidence": "The authors claim that the PDE solver used in this work is fast and efficient, as well as scalable. However, there are no ablation studies which show to effectiveness of this solver in the proposed approach as compared to other differentiable solvers.\n\nFurthermore, the performance of the approach is not clear. The authors provide plots of TPR, but do not provide any details on the obtained solutions, for example, a table which shows the obtained parameters and solutions for the different test cases. This makes it difficult to judge the results.\n\nThe authors also claim that the combination of the machine learning approach with the numerical solver improve robustness, while the plots in Figures 4 and 5 seem to show similar performance with strictly machine learning based approaches. This claim seems difficult to accept given the limited evidence.", "methods_and_evaluation_criteria": "The experiments make sense, but the baseline methods seem limited. A recent approach, PDE-Learn, is referenced in this work, but not selected as a baseline. Is there a reason for this? \n\nThe evaluation of the approach for the experiments is also relatively limited. TPR and a maximum relative error over the coefficients are provided, but the relative performance of the solution to the discovered PDEs as compared with a numerical solution to the true PDEs is not given. Overall, the evaluation metrics make it difficult to determine exactly how well each approach is doing. Additional figures and organization of the results would be helpful to illustrate the ability to discover the PDEs. Since efficiency is also a central claim, the computation time and required computational resources of all approaches (including baselines) are also a critical point for discussion.", "theoretical_claims": "No theoretical claims are made here.", "experimental_designs_or_analyses": "Adam optimizer is used with a small learning rate and no scheduler is used. Is this the case across all baseline models as well? If so, this is a bit questionable, as learning rate scheduling and tuning can be critical to many applications.", "relation_to_broader_scientific_literature": "Differentiable solvers are useful for a variety of tasks, particularly in solving inverse problems and engineering design. Fast, scalable solvers of this type are crucial, and this paper aims to address that, as well as show an interesting application in inverse problems (PDE discovery from data). \n\nAdditionally, many areas of science rely heavily on modeling by PDEs, for example weather dynamics. In weather prediction, these PDEs are often not known exactly, and come from heuristics and experimental testing. Likewise, weather data is often noisy. An approach such as this one will aid understanding in these fields, and allow scientists to construct more accurate models of physical phenomena.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Strengths\n1. The systems considered in this work are relatively challenging\n2. The authors discuss limitations of their approach\n3. The proposed numerical solver seems to be robust, with the ability to solve many systems of PDEs\n\nWeaknesses\n1. As previously mentioned, the baseline methods seem to be relatively weak\n2. It is relatively difficult to draw comparisons between the performance of all approaches explored in this work and the evaluations are limited\n3. The contribution of this work to the literature seems minimal in the paper's current form, and the claims made by the authors are not strongly supported by experimental evidence\n4. In general, I find the structure of the paper to be hard to follow", "other_comments_or_suggestions": "I would suggest the follow major changes for the work to improve the clarity:\n- The paper focuses on PDE discovery, when the sparse solver is ultimately the main contribution in this work. I believe by first restructuring the paper to discuss the numerical solver and its contributions and advantages over existing works, the impact of this work would be more clear. \n- Following this, MechNNs can be introduced. Currently, both section 2 and 4 discuss MechNNs, and so this information is split. It should be unified so that the reader can more easily follow the goals of the experiments section. \n- This will naturally lead into the results section. Results for both the NeuRLP solver and MechNN-PDE can be shown here, and some details on the baseline methods can also be included. \n- Additionally, the results section should be expanded to include more details. I think the previous sections could be shortened to cover the fundamental aspects of the approach, but technical details can be moved to the appendix. \n\nI would like to ask the authors to review the work for typos as well. \n\nAlso, the NeuRLP acronym is not defined.\n\nLine 342, right column, references Figure 6, when it should reference Figure 3.", "questions_for_authors": "1. Are all the models retrained for each different noise setting? If so, how do the results compare with a model trained on a noiseless setting and evaluated on the noisy cases?\n\n2. How is the threshold chosen? What if a PDE has many small parameters?\n\n3. Can the approach be extended to the non-Cartesian domain? \n\n4. The structure of the PDE is defined a priori, as this problem is usually very ill-posed. Therefore, it makes sense to only consider solving for the parameters of the PDE. Recently, several papers on finding analytical solutions to PDEs from data have been proposed. Would it be possible to extend this approach to find analytical solutions and their parameters?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper presents Mechanistic PDE Networks (MechNN-PDE), an extension of the MechNN framework to partial differential equations (PDEs), combining neural representations with a custom GPU-capable sparse multigrid solver. The system enables learning of space-time-dependent linear PDEs from data, allowing for expressive and differentiable modeling of spatiotemporal dynamics, even in noisy regimes. The work targets the discovery (or recovery) of governing PDEs in challenging scenarios, including reaction-diffusion systems and the Navier-Stokes equations.

There was a good discussion around this submission and the authors did a good job in the rebuttal which lead to adjusted scores. The overall consensus is acceptance of the paper and I concur. However, I urge the authors to address the reviewers' point in their revision/camera ready.
Decision: Accept (poster)