Space-Time Continuous PDE Forecasting using Equivariant Neural Fields
Authors: David M Knigge
David Wessels
Riccardo Valperga
Samuele Papa
Jan-Jakob Sonke
Erik J Bekkers
Stratis Gavves
Venue: NeurIPS 2024 poster (2024)
Keywords: pde solving
neural fields
equivariance
attention
Primary area: machine_learning_for_physical_sciences
TL;DR: We introduce an equivariant continuous PDE solving method based on Equivariant Neural Fields that preserves boundary conditions and known symmetries of the PDE.

Abstract: Recently, Conditional Neural Fields (NeFs) have emerged as a powerful modelling paradigm for PDEs, by learning solutions as flows in the latent space of the Conditional NeF. Although benefiting from favourable properties of NeFs such as grid-agnosticity and space-time-continuous dynamics modelling, this approach limits the ability to impose known constraints of the PDE on the solutions -- such as symmetries or boundary conditions -- in favour of modelling flexibility. Instead, we propose a  space-time continuous NeF-based solving framework that - by preserving geometric information in the latent space of the Conditional NeF - preserves known symmetries of the PDE. We show that modelling solutions as flows of pointclouds over the group of interest $G$ improves generalization and data-efficiency. Furthermore, we validate that our framework readily generalizes to unseen spatial and temporal locations, as well as geometric transformations of the initial conditions - where other NeF-based PDE forecasting methods fail -, and improve over baselines in a number of challenging geometries.

Space-Time Continuous PDE Forecasting using Equivariant Neural Fields
David Knigge"1, 1, David R. Wessels" Riccardo Valperga', Samuele Papa', Jan-Jakob Sonke?,
Efstratios Gavves f.1, Erik J. Bekkers
I University of Amsterdam Netherlands Cancer Institute
smiggeeuva... nl, d.r. wesselsduva.n.
Abstract
Recently, Conditional Neural Fields (NeFs) have emerged as powerful modelling paradigm for PDEs, by learning solutions flows in the latent space of the Conditional NeF. Although benefiting from favourable properties of NeFs such as grid-agnosticity and space-time-continuen dynamics modelling, this approach limits the ability to impose known constraints of the PDE on the solutions e.g. symmetries boundary conditions in favour of modelling flexibility, Instead, we propose space-time continuous NeF-based solving framework that by preserving geometric information the latent space respects known symmetries of the PDE. We show that modelling solutions as flows of pointclouds over the group of interest improves generalization and data-efficiency. We validated that our framework readily generalizes to unseen spatial and temporal locations, as well as geometric transformations of the initial conditions where other NeF-bbaed PDE forecasting methods fail and improve over baselines in a number of challenging geometries.
1 Introduction
Partial Differential Equations (PDEs) are foundational tool in modelling and understanding spatio- temporal dynamics across diverse scientific domains. Classically, PDEs are solved using numerical methods such as finite elements, finite volumes, or spectral methods. In recent years, Deep Learning (DL) methods have emerged promising alternatives due to abundance of observed and simulated data as well as the accessibility to computational resources, with applications ranging from fluid simulations and weather modelling [51] to biology (33]
* shared first author. shared lead advising
38th Conference on Neural Information Processing Systems (NeurIPS 2024).
[]
Lmon
fo
p:
D
fo(z:
t
T
Solan
Figure 1: We propose to solve equivariant PDE in function space by solving equivariant ODE in latent space. Through our proposed framework, which leverages Equivariant Neural Fields for a field is represented by a set of latents 2 {P) c {") J consisting of pose p; and context vector c,, Using meta-learning, the initial latent = is fit only SGD steps, after which an equivariant neural ODE F3 models the solution as a latent flow.

---

The systems modelled by PDEs often have underlying symmetries. For example, heat diffusion or fluid dynamics can be modeled with differential operators which are rotation equivariant, e.g.. given a solution to the system of PDEs, its rotation is also a valid solution [ sensible, and even desirable, to design neural networks that incorporate and preserve such symmetries to improve generalization and data-efficiency [12 48
In such scenarios it is
Crucially, DL-based approaches often rely on data sampled on regular grid, without the inherent ability to generalize outside of it, which is restrictive in many scenarios 40]. To this end, [51] propose to use Neural Fields (NeFs) for modelling and forecasting PDE dynamics. This is done by fitting a neural ODE [II] to the conditioning variables of a conditional Neural Field trained to reconstruct states of the PDE [3]. However, this approach fails leverage aforementioned known symmetries of the system. Furthermore, using neural fields as representations has proved difficult due to the non-linear nature of neural networks [3 3 35,, limiting performance in more challenging settings. We posit that NeF-based modelling of PDE dynamics benefits from representations that account for the symmetries of the system as this allows for introducing inductive biases into the model that ought to be reflected in solutions. Furthermore, we show that through meta-learning [28] 45 the NeF backbone improves performance for complex PDEs by further structuring the NeF's latent space, simplifying the task of the neural ODE.
We introduce a framework for space-time continuous equivariant PDE solving. by adapting a class of Eqqiivaraat Neural Fields (ENFs) to PDE-specific symmetries. We leverage the ENF as representation for modelling spatiotemporal dynamics. We solve PDEs by learning flow in the latent space of the ENF starting at point 2g corresponding to the initial state of the PDE with an equivariant graph-based neural ODE [1] we develop from previous work [5]. We extend the ENF to equivariances beyond SE(n), by extending its weight-sharing scheme to equivalance classes for specific symmetries relevant to our setting. Furthermore, we show how meta-learning [1 28 45 [3), can not only significantly reduce inference time of the proposed framework, but also substantially simplify the structure of the latent space of the ENF, thereby simplifying the learning process of the latent dynamics for the neural ODE model. We present the following contributions:
We introduce a framework for spatio-temporally continuous PDE solving that respects known symmetries of the PDE through equivariance constraints.
We show that correctly chosen equivariance constraints as inductive bias improves per- formance of the solver in terms of MSE in spatio-temporall continuous settings, i.c. evaluated off the training grid and beyond the training horizon.
We show how meta-learning improves the structure of the latent space of the ENE, simplify- ing the learning process, leading to better performance in solving PDEs.
We structure the paper as follows: in Sec. 2 we provide an overview of the mathematical preliminarics and describe the problem setting. Our proposed framework is introduced in Sec. 3 We validate
our framework on different PDEs defined over a variety of geometries in Sec. H with differing equivariance constraints, showing competitive performance over other neural PDE solvers. We provide
an in-depth positioning of our approach in relation to other work in Appx. A
2 Mathematical background and problem setting
Continuous spatiotemporal dynamics forecasting. The setting considered is data-driven learning of the dynamics of system described by continuous observables. In particular, we consider flows of fields, denoted with D ddd x [0, + RR We use i as a shorthand for D(.,t). .). We assume the flow governed by a PDE, and consider the Initial Value Problem (IVP) of predicting ù, from a given vo- The dataset consists of field snapshots v X' x ['] + RS in which [2]] == 1, 1.2...., T denotes the set of time points on which the flow is sampled and X C Rd is a set of coordinate values. For each time point we are given a set of input-output pairs (X,((X) where v(X) C R are the values of the field at those coordinates. Importantly, the location at which the field is sampled need not be regular, i.e., we do not require the training data to be coordinate values be identical for train and test sets. Following (51), we distinguish between in- referring values within the training time horizon [0, T and four analogously to values beyond T.
grid or be regularly spaced in time, nor need
Asssuinng boundary conditions symmetric, ie. they transform according to the relevant group action.

---

Neural Fields in dynamics modelling. Conditional Neural fields (NeFs) are class of coordinate- based neural networks, often trained to reconstruct discretely-sampled input continuously. More specifically, conditional neural field fo R - IRd is field -arameterized by neural network with parameters 0- that maps input coordinates a: E RM in the data domain alongside conditioning latents to d-dimensional signal values v(xx) E IR... By associating conditioning latent z" E TR
to each signal single conditional NeF fe IE x IE + Rd can learn to represent families D of continuous signals such that Vv E D: f()) = (2; 2"). [5 propose to use conditional NeFs for PDE modelling by learning continuous flow in the latent space of a conditional neural field. In particular, set of latents {= are obtained by fitting conditional neural field to a given set of observations {=))²= at timesteps L..., T; simultaneously, a neural ODE [ F& is trained to map pairs of temporally contiguous latents s.t. solutions correspond to the trajectories traced by the learned latents. Though this approach yields impressive results for sparse and irregular data in planar PDEs, we show it breaks down on complex geometries. We hypothesize that this is due to lack of a latent space that preserves relevant geometric transformations that characterize the symmetries of the systems we are modelling, and as such propose an extension of this framework where such symmetries are preserved.
Symmetries and weight sharing. Given a group G with identity element e E G. and a set X, group action is a map T: G + X. For simplicity, we denote the action of g E G on X as gz := T(g. z), and call G-space smooth manifold equipped with a G action. A group action is homomorphic to G with its group product, namely it is such that ez = x and (gh)æ g(hx). As an example, we are interested in the Special Euclidean group SE(n)=R" x SO(n); group elements of SE(n) are identified by a translation L E Ie and rotations R E SO(n) with group operation g9' = (1. R&) (L', Rgy) (Rx' +x RR ; We denote by Cg the left action of G on function spaces defined as (x' = f
f(R," 'x' x)). Many PDEs are defined by equivariant differential operators such that for a given state v: LgN(v] =M[A,"). If the boundary conditions do not break the symmetry, namely if the boundary symmetric with respect to the same group action, then a G-transformed solution to the IVP for some vo corresponds to the solution for the G-transformed initial value. For example, laws of physies do not depend on the choice of coordinate system, this implies that many PDEs are defined by operators. The geometric deep learning literature shows that models can benefit from leveraging the inherent symmetries or invariances present in the data by constraining the searchable function space through weight sharing g 25 5). Recall that in our framework we model flows of fields, solutions to PDES detined by equivariant differential operators, with ordinary differential equations in the latent space of conditional neural fields. We leverage the symmetries of the system for two key aspects of the proposed method: first by making the relation between signals and corresponding latents equivariant; second, by using equivariant ODEs, namely ODEs defined by equivariant vector fields: dz
differential
dr F(2) is such that F (g=) = g (z), then solutions are mapped solutions by the group action.
3 Method
We adapt the work of [51, and consider the following optimization problem? x
min BuED, || [o(;;= where x = z + L (o( 0,0.2
(1)
with fo(x; 2f') a decoder tasked with reconstructing state from latent =/" and Fg. a neural ODE that Fx modelling the solution as flow latent space starting at the initial latent 2/ see Fig. for visual intuition.
maps a latent to its temporal derivative: dzz
Equivariant space-time continuous dynamics forecasting. A PDE defined by a G-equivariant differential operator for which Lv,N[v = NN are such that solutions are mapped other solutions by the group action if the boundary conditions are symmetric. We would like to leverage this property, and constrain the neural ODE F.. such that the solutions it finds in latent space be mapped onto each other by the group action. Our motivation for this is twofold: (1) it is natural for
We highlight that [5] optimize latents t.. neural field fa. and ODE F. using two separate objectives. We instead found that our framework is stable under single-objective optimization.
3

---

our model to have, by construction, the geometric properties that the modelled system is known to posses (2) to get more structured latent representations and facilitate the job of the neural ODE. To achieve this we first need the latent space Z to be equipped with well-defined group action with respect to which Vg E G. E Z Fa(gz)= Fo (z), and, most importantly, we need the relation between the reconstructed field and the corresponding latent to be equivariant, i.e.,
(2) Note that, somewhat imprecisely. we call this condition equivariance to convey the idea even though it is not, strictly speaking, the commonly used definition of equivariance for general operators. If we consider the decoder as a mapping from latents to fields, we can make the notion of equivariance of this mapping more precise. Namely
Vgg E G, x E X Lyfo(x; 2 = - 2/) fo(:; g=(").
f(x) = Do(z),) Dz(z) z + fo(:= f(gg/(()))))) Da(gz), Do(gz) ) g + f
:2() (3)
In Sec. 3.1 we describe the Equivariant Neu- ral Field E -baased decoder, which satisfies equation 2) Second, in Sec. 3.2 we outline the graph-based equivariant neural ODE. Sec. 3.3
solve F/Fo
explains the motivation for- and use of- meta- learning for obtaining the ENF backbone param- eters. We show how the combination of equiv- ariance and meta-learning produce much more structured latent representations of continuous signals (Fig. 3)
transform
wransoom
fa 30)
fola; Ep
solve F/Fe
9
gof
fo(1:0=))
3.1 Representing
Figure 2: The proposed framework respects pre- defined symmetries of the PDE: a rotated solution Ly may be obtained either by solving from la- tent a5 (top-left) and transforming the solution 24 (top-right) to gay (bottom-right) or transforming z to ga (bottom-left) and solving this.
PDE states with Equivariant Neural Fields
We briefly recap ENFs here, referring the reader to [49] for more detail. We extend ENFs to symmetries for PDEs over varying geometries.
Attention -based conditional neural fields [52) This class of conditional neural fields through a cross-attention operation between the latent set 2" and input coordinates ENFs 49 extend this approach by imposing equivariance constraints w.r.t a group G C SE(n) on the relation between the neural field and the latents such that transformations the signal y correspond to transformation of the latent z" (Eq. 2 For this condition to hold, we need a well-defined action on the latent space Z of Sg. To this end, ENFs define elements of the latent set z" as tuples of pose p: E G and context c; E R2d 2" == {pr. { The latent space is then equipped with group action defined g2 ((p: ) To achieve equivariance over transformations ENFs follow S where equivariance is achieved with convolutional weight-sharing over equivalence classes of points pairs x. x'. ENFs instead extend weight-sharing to cross-attention over bi-invariant attributes of 2, x pairs.
ENFs as cross-attention over bi-invariant attributes. represent signal D with corresponding latent set 2"
obtain signal-specific reconstructions (xx) Rs fo (x: 2²)
Weight-sharing over bi-invariant attributes of 2, is motivated by Eq. 2 by which we have:
fo(:; 2) fo(gx;g))
Intuitively, the above equation says that a transformation on the domain of fe. i.c. 'r, can be undone by also acting with g on 2. In other words, the output of the neural field fo should be bi ri-invariant to g transformmations of the pair bi-invariant attribute a,m describes function a (24, F.m) + a(2;, Fmm such that a(2;, Rm) =
For specific pair (2s, E Z x X, the term
a(gas, grm). Thorughout the paper we use aim shorthand for ((z)
To parameterize Je, we can accordingly choose any function that bi-invariant Grrransformation of 2, x. In particular, for an input coordinate xm, ENFs choose to make fo a cross-attention operation between attributes a,m and the invariant context vectors C,;
fo(Em; z) cross_ates pooss_att((,,,. C., c.)
(5)
As an example, for SE(n) -quivvraui,,,,, we can define the bi-invariant simply using the group action:
SE(n) = p, Zm = R. (am Iz),, which is bi-invariant by:
m.
Vg E SE(n): (p),)) + (g)..g))
+ p (gp.) 10'' gx p, -p---'' P, x (6)
gz

---

Bi-invariant attributes for PDE solving, As explained above, ENF is equivariant to SE(n)- transformations by defining fo as a function of an Siini-bi-nnvarian attribute a a Although many physical processes adhere to oooo-rranslational symmetries, we are also interested in solving PDEs that due to the geometry of the domain, their specific formulation, and/or their boundary conditions are not fully
As such, we are interested in extending ENFs to equivariances that are not strictly (subsets of) SE(n), which we show we can achieve by finding bi-invariants that respect these particular transformations. Below, we provide two examples, the other invariants we use in the experiments including a "bi-invariant" a that is not actually bi-invariant to any geometric transformations, which we use to ablate over equívariance constraints are in Appx. D
The flat 2-torus. When the physical domain of interest is continuous and extends indefinitely, periodic boundary conditions are often used, i.e. the PDE is defined over a space topologically equivalent to that of the 2-torus. Such boundary conditions break S0(2) symmetries; assuming the domain has periodicity т and none of the terms of this PDE depend on the choice of coordinate frame, these boundary conditions imply that the PDE is equivariant periodic translations: the group of translations modulo x: T R²/2². In this case, periodic functions over with periods would
work as bi-invariant, i.e. using poses p E T a cos(2r Po)) cos(2 (#1 pı)) which т happens to be bi-invariant to rotations by 2 as well. Instead, since we do not assume any rotational symmetries to exist on the torus, we opt for non-rotationally symmetric function:
ag m. - coss(2 p))) @ cos(2 - p.)),
T
where @ denotes concatenation. This bi-invariant is used in experiments on Navier-Stokes over the flat 2-Torus.
The 2-sphere. In some settings a PDE may be symmetric only to rotations along a certain axes. An example is that of the global shallow-water equations on the two-sphere- used model geophysical processes such as atmospheric flow [16), which are characterised by rotational symmetry only along the earth's 's axis of rotation due inclusion of term for Coriolis acceleration that breaks full SO(3) equivariance. We use poses p E SO(3) parametrised by Euler angles o, 0, 7, and spherical coordinates 6, for x E S We make the first two Euler angles coincide with the spherical coordinates and define a bi-invariant for rotations around the axis 0 π
SW a i,m
Appr.am ① 6p. + Tp: + 6.m:
(8)
where A =Op. mm to adjust for periodicity.
2T if dp. - -oam > and Dop. =0p. - -mm + + if pp. -mam -
In summary, to parameterize an ENF equivariant with respect to a specific group we are simply required to find attributes that are bi-invariant with respect to the same group. In general we achieve this by using group-valued poses and their action on the PDE domain.
3.2 PDE solution as latent space flow
Let = be a latent set that faithfully reconstruets the initial state Po. We want to define a neural dzy d= =F that is equivariant with dr ===f)==((==()) To this end, use message passing neural
ODE Fg that map latents 2 to their temporal derivatives respect to the group action: network (MPNN) to learn a flow of poses and contexts c; over time. We base our architecture on PONITA (5), which employs convolutional weight-sharing over bi-invariants for SE(n). For an in-depth recap of pesssae-passsig frameworks, we refer the reader to Appx. A Since F6 is required to be equivariant w.r.t. the group action, any updates the poses P; should also be equivariant, 4I propose to parameterize an equivariant node position update by using a basis spanned by relative node positions 2 Fg. In our setting, poses p; are points on a manifold M equipped with a group action. As such, we analogously propose parameterizing pose updates by weighted combination of logarithmic maps log,, (p,), which intuitively describe the relative position between pi; Pj in the tangent space Tp, M, or the displacement from p: to Pj. We integrate the resulting pose update over the manifold through the exponential map expp,- In the euclidean case logp, (P))= - Ir and we get back node position updates per [1]. In short, the message passing layers we use consist of the
5

---

(a)
(b)
(c)
Figure 3: We show the impact of meta-learning and equivariance the latent space of the ENF when
representing trajectories of PDE states. Fig. Shows a T-SNE plot of the latent space of fa when of optimized with autodecoding, and no weight sharing over bi-invariants enforced. Fig. 3b shows
the latent space when meta-learning is used, but no weight sharing is enforced. Fig. 3c shows the latent space when z are obtained using meta-learning and fo shares weights over a SE
following update functions:
ct+
Σ
context (a.j)dds,
Ps I+1 = expy!
Σ
hmne (al ( )e, logg! (p/)
(p).e))Ea"d
(p),c)))am
with Acontex Apose message functions weighting the incoming context and pose updates, parameterized by two-layer MLP function of the respective bi-invariant.
3.3 Obtaining the initial latent
Until now we' not discussed how obtain latents corresponding the initial condition z An approach often used in conditional neural field literature is that of autodecoding [36), where latents =" are optimized for reconstruction of the input signal with SGD, Optimizing a NeF for reconstruction does not necessarily lead to good quality representations 35 i.e. using MSE- -SEed autodecoding to obtain latents 2 as is proposed by [51] may complicate the latent space, impeding optimization of the neural ODE Fg. Moreover, autodecoding requires many optimization steps at inference (for reference, [ST] use 300-500 steps). [33 propose meta-learmnng as way overcome long inference times, as it allows for fitting latents in a few steps typically three four. We hypothesize that meta-learning may also structure the latent space similar to the impact of equivariance constraints, since the very limited number of optimization steps requires efficient organization of latents z/¹ around the (shared) initialization, forcing together the latent representation of contiguous states. To this end, we propose the neural ode Fg to find solutions z/.
meta-learning for obtaining the initial latent zi which is then unrolled by
3.4 Equivariance and meta-learning structure the latent space
As first validation of the hypotheses that both equivariance constraints and meta-learning introduce structure to the latent space of for we visualize latent spaces of different variants of the ENF. We fit ENFs to a dataset consisting of solutions to the heat equation for various initial conditions (details in Appx. Eb. For each sample Vt, we obtain a
latents 2/,, which we average over the invariant context vectors E IR to obtain single vector in Ie invariant to group action according to the chosen bi-invariant. Next, we apply T-SNE [47 to resulting vectors in R We use three setups: (a) meta-learning, 0 and latents = optimized for every separately using autodecoding (36), and no equivariance imposed (per Eq. [5]. shown Fig. [3a (b) meta-learning is used to obtain 0, but no equivariance imposed, shown Fig. 3b and (c) meta-learning is used to obtain 0.2/ and SEE() -quivaancue is imposed by weight-sharing over a SE(n) bi-invariants, shown in Fig. 3c The results confirm our intuition that both meta-learning and equivariance improve latent-space structure.
13 to obtain the initial latent 2 under supervision coordinate-val pairs (z, v(x)))rex from un- This latent is unrolled for frain timesteps using Fg. The obtained latents are used to reconstruct states along the trajectory
Recap: optimization objective. We use meta-learning inner-loop
28

---

of v. and parameters of fe. F6 are optimised for reconstruction MSE, as shown in the left-hand side of Eq. See Appx. for detailed pseudocode of this process.
4 Experiments
We intend to show the impact of symmeirrypreseeraaiio preservation in continuous PDE solving. To this end we perform range of experiments assessing different qualities of our model on tasks with different symmetries. First, we investigate the equivariance properties of our framework by evaluating it against unseen geometric transformations of the initial conditions. Next, we assess generalization and extrapolation capabilities w.r.t. unseen spatial locations and time horizons inside and outside the time ranges seen during training respectively, robustness to partial test-time observations, and data- efficiency. As the continuous nature of NeF-based PDE solving allows, we verify these properties for PDEs defined over challenging geometries: the plane 2 2-torus T and the sphere S and the 3D ball [B', Architectural details and hyperparameters are in Appx. E Code is available on Gitllub
Additionally, we validate our model on a benchmark of PDEs that exhibits no transformation symmetries: the CFDBench (32) benchmark. We include details on parameter counts, memory usage and runtimes of our model compared to baselines in Appx. F
4.1 Datasets and evaluation
All datasets are obtained by randomly sampling disjoint sets of initial conditions for train and test and solving them using numerical meth- ods. Dataset- speciffc details generation can be found in Appx E *Heat equation on IR² and S.. The heat equation describes diffusion
Table 1: MSE & for heat equation 1R3
tw TRAIN
furr TRAIN
TEST
forr TEST
DING BJT Ours
5.920 04
2,40B 0s
3,85E 33
5:128 05
623
450:30B-06
219
5.08anzD-00
1.18
E..00
2.50
surface: de de DV²c, where is scalar field, and is the diffusivity coefficient. We solve on the 2D plane where V a's
+ 8% and on the 2-sphere S where in spherical coordi- nates: V sin σ6
dzz
&he a (sin 0 00)) + si Dez
Although a relatively simple PDE, we find that defining it over non-trivial geometry such as the sphere proves hard for non- eouivvuiva methods. *Navier-Stokes
Jq train
test
T We solve 2D Navier Stokes [43] for an incompressible fluid de
Figure 4: A train and test sample from the planar diffusion dataset. Initial conditions for train and test are spikes in disjoint subsets of IR..
with dynamics dt Ve + vāyu f,u
V x Vu = 0, where is the velocity field, the vorticity, a the viscosity and f a forcing term (see Appx. Ep. We create a dataset of solutions for the vorticity using Gaussian random fields as initial conditions. Due to the incompressibility condition, it is natural to solve this PDE with periodic boundary conditions corresponding to the periodic translation. -haellow-water The global shallow- -wtee equations model large-scale oceanic and atmospheric flow on the globe, derived from Navier-Stokes under assumption of shallow fluid depth. The global shallow-water equations (see Appx. E include terms for Coriolis accelleration, which makes this problem equivari- rotation along the globe's axis of rotation. We follow the IVP specified by [16), and create a dataset of paired vorticity-fluid height solutions. +Internally- heated convection in a 3D ball. We solve the Boussinesq equation for internally heated convection in ball, model relevant for example in the context of the Earth's 's mantle convection. It involves continuity equations for mass conservation, momentum equations for fluid flow under pressure, viscous forces and buoyancy, and term modelling heat transfer. We generate initial conditions varying the internal temperature using (0, 1) noise and obtain solutions for the temperature defined over a regular spherical Φ, 8, grid. CFDBennc 322 consists of a set of one-step solutions (pairs of input and output states Vt, 1/+1) for classic computational fluid dynamics (CFD) problems, with varying fluid properties, boundary conditions and geometries. The goal of this dataset is to assess generalizability of DL-based PDE solvers over problem parameters. This dataset does not exhibit transformation symmetries because of the absolute position of obstacles in the geometry and orientation of the flow.
topology of 2-Torus T implying equivariance
S?.
7

---

groun
predictio
DIoo 256 ralecovies
Igi peod
e0 eq.ieeeent
hn Lext
1=0 1=10 2=
10
1-40
t=00
train
Figure 5: A Navier-Stokes test sample (top) and corre- sponding predictions from our model (bottom). We visu- alize predictions in the train horizon in- [O,..., 9] fout [10...., 20] and beyond. The model remains stable well be- yond the train horizon, but due accumulated errors fails to capture dynamics beyond t > 40.
Figure 6: Test MSE in for increasing training set sizes for the heat equa- tion over the sphere. Equivariant im- proves over non-equivariant. For ref- erence we show performance of DINo [51] trained on 256 trajectories.
Evaluation. All reported MSE values are for predictions obtained given only the initial condition to. with std over 3 runs. We evaluate two settings for train and test sets both: generalization setting with time evolution happening within the seen horizon during training (fmn); and, extrapolation setting with the time evolution happening outside the seen horizon during training (lom). For both cases we measure the mean-squared error (MSE). To position work relative to competitive data-driven PDE solvers, on the 2D-NNaver--Stokes experiment we provide comparisons with a range of baselines. In most other settings these models cannot straightforwardly be applied, and we only compare to [51]. to our knowledge the only other fully continuous PDE solving method in literature. For the Navier-Stokes and Internally- Convection experiments, we compare with Transolver (50). which has shown SOTA results as DL-based PDE solving method for general geometries.
Equivariance properties - heat equation on the plane. To verify our framework respects the posed equivariance constraints, create a dataset of solutions to the heat equation that requires a neural solver to respect equivariance constraints to achieve good performance. Specifically, for initial conditions we randomly insert pulse of variable intensity in
x2) E IR² s.t.
-I<; <1, 0<r=<< for the training data and -1<x; <1, -11,,- <0 for the test data. Intuitively, train and test sets contain spikes under different
Table 2: MSE for
on T
disjoint sets of noto-translations (see Fig. 4 We
train variants of our framework with (a (ET)) Eq.
fn TRAIN
furr THAIN
fu TEST
Tapp TEST
6) and without (a Eq. 15) equivariance con- straints. In this dataset, we set im = [0, 9). and evaluation horizon tout = [10, 20). Re- sults in Tab. [ show that the non-vauivariaatt model, as well as the baseline [51] are un- able to successfully solve test initial conditions, whereas the equivariant model performs well.
100% OF
DBSERVED
CNODDE P
6:026 02
5.350 01
5.48E 02
5.17E 01
FND
9:43--05
2:115-05
3.44E-05
1 1 03
G-PNO
3.13E-05
3499-04
3.156-05
3.12-04
Transoveer 50
1,50-02
4.85 01
1.85--2
4 01
DINo 51
8,20E 00
6358 02
1,11E 02
9.8E 02
Ours 57
560:1ol
6.75
4.00saat-00
1.1
167-12E-00
2.60
2.14:
Ours a 2i/e
1.45
9.34 03
1.57 05
L.16 02
50% or ve OBSERVED
CNODD 2
L 01
6.338 01
1,53E 01
6.76 01
FND
3.31E-02
1.396-01
3.21F-002
1477-01
G-PND
2.73E-02
1.176-01
2,536-02
1 101E-01
Transceer 50
4.69E-01
1199
4.76E-01
0.9E.01
Robustness to subsampling & time-horizons Navier-Stokes on the 2-Torus. We perform an experiment assessing the impact of equivari- robust-
DINo 30
3.57E.02 02
2 01
3,74E 02
2133
Ours 42
6.31 02
3.98
2 01
7.0 02
4.011SE 01
105:10EB02
1.4
2600
2.14
Ous 2/s/
1,50 A3:
5.5 13
ance constraints and meta-lcarning ness to sparse test-time observations of the ini- tial condition. To this end, we train model with (a Eq. 7) without (a". Eq. 15) equivari- ance constraints, and one with equivariance con- straints and without meta-learning (AD a T", Eq. 77.. on a fully-observed train set. The training horizon fm = [0........ .... 9), and evaluation horizon tou = [10,... .... 20). Subsequently, we apply the trained model to the problem of solving from sparse initial conditions Uo. with observation rates where 50% and 5% of the initial condition is observed (Tab. 2 Approaches operating on
CNODF P
1.214400
2:14440
OBSERVED
1.F0400
4335-00
FND
4.13F-01
7-706-01
1.84F-01
7075-01
G-FND
3.56E-01
7096-01
3,00-01
6478-01
Trunsdven
9.48E-01
1.235400
8.28E-01
1.24E400
DINo 30
3,67E-02
2818:01
3.4E-12
2918:01
Ouns now
3.05 PP01
4.01 PP-001
Oun 4/1
7.31
2.77
7.6
n.-01
319
1.33
344:108-02
1.61:400-01
Table 3: Zero-shot temporal superresolution on Navier-Stokes.
Ti5T
lerr T2ST
d:-1.0 tuiso
di -.5
de A11.itt
L1SE-IC
3835 02
3.6E
4.822-02
2.9E-01
01
Ous
1.6E40
1962.63
191E-83
L42E-03
L49E-42
LS22-12
discrete (CNODE 2 and regular grids (FNO [29], G-FNO 20) perform very well when evaluated
on ously-observed regular grids, outperforming continuous approaches (ours, (5l). However, we note
8

---

that all discrete/regular models greatl deteriorate in performance when observation rates decrease. Equivariance constraints and meta-learning clearly improve performance overall, achieving best perfomance in all sparse settings. Our proposed framework performs competitively to discrete base- lines and other NeF based PDE solving methods [51] in the fully observed setting. To qualitatively assess long-term stability well-heyond the train horizon, we visualizate test trajectory and the solution
found by our model for in = [0. ..., 9] feut [10,.... 20] and beyond Fig. We show error accumulation for solutions 100% observed and 50% observed initial states in Appx. F To validate the time-continuous nature of our model, we train a model with supervision at a resolution of dr=1.0, and subsequently evaluate resolutions dr=0.5 and d7=0.25 (Tab. B), showing that our framework does not accumulate significant error when deploying on higher temporal resolution.
Data-efficiency- Diffusion on the sphere. To assess the impact of equivariance on data efficiency, we vary the size of the training set of heat equation solutions from 16 to 64 trajectories and apply model with (a$0(3), Eq. 13) and without (a Eq. 15) equivariance constraints. In this dataset, we set lim = [o, 9). and evaluation horizon for = 10. 20]. We visualize in test- and train MSE in Fig. 6 These results show the non-equívariant model overfitting the training set for smaller numbers of trajectories while unable to solve the PDE satisfactorily, whereas the equivariant model generalizes
well even with only 16 training trajectories.
Super-resolution Shallow-Water on the sphere. Due to their continuous nature, NeF- based approaches inherently support zero-shot super-resolution. In this setting, we generate set of solutions for the global shallow-water equations over S at 2x resolution, and apply mean-pooling with a kernel size of 2 to obtain a low-resolution dataset. We train a model that respects rotational symmetries along the rota- tion axis of the globe (aSW, Eq. 8 at train resolution, and evaluate the model by solving initial conditions at 2x resolution (Tab. 4 Fig.
Table 4: MSE & on Shallow Water equations on the sphere.
to TRAIN
Eus TRAIN
to TEST
fuur TEST
TRAIN RESOLUTION
DINO 5J
177E 01
L.GE-03
2001 04
1.7E-43 1.87
Ours
9.44
1.9
109
DING 3T
3.03E:04
2,09--0
3.37E-04
203E-03
2
ground truth
7) In this dataset, we set lin [O..... 9.. and evaluation horizon lout = [10, 14]. First, we note that our model has difficulty capturing the dynamics near four and beyond the training horizon, i.c. => we suspect because of accumulation of reconstruction errors impacting the ability of Fg. model the relatively volatile dynamics of these equations. This points to drawback of NeF-bbaeed solvers: error accumu- lation starts with the reconstruction error on the initial condition. Ranging over our experiments, we found that this error can be reduced by in- creasing model capacity, at steep cost of com- putational complexity attributable to the global attention operator in the ENF backbone. Regard- ing super-resolution; the model is able solve the high-resolution initial conditions without in- ducing significantly increased MSE it does not produce significant artefacts in the process.
ground truth high res
zero-shot prediction
f-3
Figure 7: Test samples at train resolution (top), 2x train resolution (middle) and corresponding predictions from our equivariant model (a Eq. 8 (bottom). The model does not produce significant upsampling artefacts, but fails to capture dynamics outside the training horizon.
Table 5: MSE 4 on Internally in the ball.
-Heanedly-Heated Heated Convection
Ey TRAIN
Cour TRAI
Isc TEST
Sun TIST
100% mp OBSTEVED
Transdver DINo 511
3.99-04
4.13E-04
2.095-02
2998-05
7.56L-02
3,06E-6
7.788-02
Ourow"
79 2P-44
7,72 2000E-03
1.99g119-004
7.97
Challenging geometries Internally heated convection in 3D ball. We show the value of inductive biases modelling over a challenging geometry. We apply an equivariant model (a Eq. 14 to a set of solutions to Boussinesq internally heated convection in a ball defined over a regular d. 0, r-grid, where we set 4n = [0, ..., 9), and
Tranedsez SD DINo 51
4 4.99F-01
4.99--01
50% OF OBSERVED
4.38F-01
5000-001
501E 05
LOGE
3,02
111E
Cues o
6.7E-44
7.76E-03
6638-04
8.21
evaluation horizon laut = [10,... .... 14]. Results (Tab. [6] Fig. 8 for our equivariant model show good
generalization compared to a non-equivariant baseline [51]. The significantly larger Transolver (50)
model (see Appx. F) obtains better performance within the train time horizon on train and test

---

sets, but overfits to this time horizon, generalizing poorly beyond. We additionally show results for sparsely observed input states, showing the limitations of non NeF-based solvers to generalize over irregular changing observation grids. We interpret these results as an indication of a marked reduction solving-complexity and improved generalization when correctly accounting for PDE's symmetries.
Non-symmetric PDEs Lastly, we evaluate our model in the setting when solving PDEs that do not exhibit any global symmetries to assess whether the preservation of symmetries in the latent space of our model precludes application to non-symmetrie PDEs. We train translation- equivariant (a**) model, i.e. one with equivari- ance properties identical to a CNN, on the Cav- ity. Dam and Cylinder flows from CFDBench 32]. Comparing to the baselines set by the dataset authors, we see our method improves sig- nificantly over a number of classical baselines, despite no global transformational symmetries being present in the problems being solved.
grund buth
prediction
Figure 8: Test samples (top) and corresponding predictions from our model equivariant to $². rotations in the ball. (Eq. 14p
Table 6:
Test MSE + for CFDBench auto-
regressive one forward propagation considering all properties for the cavity, dam and cylinder flows.
5 Limitations & Future work
Baselines taken from 32
CAVITY
DAM
CYLINDER
We're interested in exploring the application of our based PDE solving framework to larger-scale, more complex problems. Through- experiments with ENFs noticed that more complex signals, e.g. higher resolution PDEs, may be fit easily either by increasing the ENF hidden size or by increasing the number of latents used. However, either of these changes significantly impacts computational complexity, due to the calculation of attention coefficients in the latent space the ENF for every latent-input coordinate pair. possible way of addressing this
Identity
6.42F-02
1.69F-03
7.54F-02
Auto-FFN
6.42E-02
1668-003
7.54E-02
Auto- Meto-DeepONet
6.39E-02
1.64E-03
7.53E-02
Deee-EDee
6.45F-02
1.49E-03
7.43F-02
Neto-DeePONeeCNN
6.33F-02
1.68F-03
7.52F-02
FNO
2.61E-02
8.75B-05
1,15E-03
U-Net
1,58E-02
1.70E-03
5,49E-05
Ours
aa
1.10E:02
8.19E-05
1.42E-05
is detailed in [49]; we can approximate the output of the attention operation through limiting the number of latents attended (using k-nearest neighbours). This may open the door to modelling more complex, larger-scale dynamics than learned in present experiments.
6 Conclusion
We introduce a novel equivariant space-time continuous framework for solving partial differential equations (PDEs). Uniquely the initial state while respecting symmetry consttainss and boundary conditions of the underlying PDE. We clearly show the benefit of commetry--reeervaatee over range of challenging tasks, where existing methods fail to capture the underlying dynamics.
method handles sparse irregularly sampled observations of
References
[1] Ilze Amanda Auzina, Çağatay Yildiz, Sara Magliacane, Matthias Bethge, and Efstratios Gavves. Modulated neural odes. Advances in Neural Information Processing Systems, 36, 2024.
[2] Ibrahim Ayed, Emmanuel De Bezenac, Arthur Pajot, and Patrick Gallinari. Learning the spatio- temporal dynamics of physical processes from partial observations. In ICASSP 2020-2020 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pages 3232-3236. IEEE, 2020.
[3] Mathias Bauer, Emilien Dupont, Andy Brock, Dan Rosenbaum, Jonathan Richard Schwarz, and Hyunjik Kim. Spatial functa: Scaling functa to imagenet classification and generation. arXiv preprint arXiv:2302.0313 2023.
10

---

[4] Erik Bekkers. B-spline cnns on lie groups. In International Conference on Learning Repre- sentations, 2019.
[5] Erik Bekkers, Sharvaree Vadgama, Rob D Hesselink, Putri A van der Linden, and David W Romero. Fast, expressive (n) equivariant networks through weight-sharing in position- orientation space. arXiv preprint arXiv:2310.02970, 2023.
[6] Johannes Brandstetter, Rob Hesselink, Elise van der Pol, Erik J Bekkers, and Max Welling. Geometric and physical quantities improve e (3) equivariant message passing. arXiv preprint arXiv:2110.02905, 2021.
[7] Johannes Brandstetter, Rianne van den Berg, Max Welling, and Jayesh K Gupta. Clifford neural layers for pde modeling. arXiv preprint arXiv:2209.0499 2022.
[8] Johannes Brandstetter, Daniel Worrall, and Max Welling. Message passing neural pde solvers. arXiv preprint arXiv:2202.03376 2022.
[9] Michael M Bronstein, Joan Bruna, Taco Cohen, and Petar Veličković. Geometric deep learning: Grids, groups, graphs, geodesics, and gauges. arXiv preprint arXiv:2104.13478, 2021.
[10] Keaton J. Burns, Geoffrey M. Vasil, Jeffrey S. Oishi, Daniel Lecoanet, and Benjamin P. Brown. Dedalus: A flexible framework for numerical simulations with spectral methods. Physical Review Research, 2(2):023068, April 2020. doi: 10.1
[11] Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
[12] Taco Cohen and Max Welling. Group equivariant convolutional networks. In International conference machine learning, pages 2990-2999. PMLR, 2016.
[13] Emilien Dupont, Hyunjik Kim, SM Eslami, Danilo Rezende, and Dan Rosenbaum. From data to functa: Your data point is a function and you can treat it like one. arXiv preprint arXiv:2201. 12204, 2022.
[14] Chelsea Finn, Pieter Abbeel, and Sergey Levine. aggostig meta-learning for fast adap- tation of deep networks. In International conference on machine learning, pages 1126-1135. PMIRR, 2017.
[15] Mare Finzi, Samuel Stanton, Pavel Izmailov, and Andrew Gordon Wilson. Generalizing convolutional neural networks for equivariance to lie groups arbitrary continuous data. In International Conference on Machine Learning, pages 3165-3176. PMLR, 2020.
[16] Joseph Galewsky, Richard Scott, and Lorenzo M Polvani. An initial-value problem for testing numerical models of the global shallow-water equations. Tellus A: Dynamic Meteorology and Oceanogruphy, 56(5):429-440, 2004.
[17] Justin Gilmer, Samuel S Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. International conference on machine learning. pages 1272. PMLR, 2017.
[18] Samuel Greydanus, Misko Dzamba, and Jason Yosinski. Hamiltonian neural networks, Advances in neural information processing systems, 32, 2019.
[19] Xiaoxiao Guo, Wei Li, and Francesco lorio. Convolutional neural networks for steady flow ap- proximation. In Proceedings of the 22nd ACM SIGKDD international conference knowledge discovery and data mining, pages 481 190, 2016.
[20] Jacob Helwig, Xuan Zhang. Cong Fu, Jerry Kurtin, Stephan Wojtowytsch, and Shuiwang Ji. Group equivariant fourier neural operators for partial differential equations. Proceedings of the 40 th International Conference on Machine Learning, Honolulu, Hawaii, USA. PMLR 202, 2023., 2023.
[21] Quercus Hernández, Alberto Badías, David González, Francisco Chinesta, and Elfas Cueto. Structure-preserving neural networks. Journal of Computational Physics, 426: 109950, 2021.
11

---

[22] Pengzhan Jin, Zhen Zhang, Aiqing Zhu, Yifa Tang. and George Em Karniadakis. Sympnets: Intrinsic structure- preseervigg symplectic networks for identifying hamiltonian systems. Neural Networks, (e421,6,17,,20000
[23] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
[24] Thomas N Kipf and Max Welling. Semi-supervised classification with graph convolutional networks. arXiv preprint M0Kiv:66000290000 2016.
[25] David M Knigge, David W Romero, and Erik J Bekkers. Exploiting redundancy: Separable group convolutional networks on lie groups. In International Conference on Machine Learning, pages 11359-1133 11559-13866 11386. PMLR, 2022.
[26] Miltiadis Miltos Kofinas, Erik Bekkers, Naveen Nagaraja, and Efstratios Gavves. Latent field discovery in interacting dynamical systems with neural fields. Advances in Neural Information Processing Systems, 36, 2023.
[27] Nikola Kovachki, Zongyi Li, Burigede Liu, Kamyar Azizzadenesheli, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Neural operator: Learning maps between function spaces. arXiv preprint 1eXiv:2208.00 08481, 2021.
[28] Zhenguo Li, Fengwei Zhou, Fei Chen, and Hang Li. Meta-sgd: Learning to learn quickly for few-shot learning. arXiv preprint 17X1v:170..09835, 2017.
[29] Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differen- tial equations. arXiv preprint arXiv:2010.0885 10.0895, 2020.
[30] Yongtuo Liu, Sara Magliacane, Miltiadis Kofinas, and Efstratios Gavves. Graph switching dynamical systems. In International Conference Machine Learning, pages 21867-21883. PMLR, 2023.
[31] Yongtuo Liu, Sara Magliacane, Miltiadis Kofinas, and Efstratios Gavves, Amortized equation discovery in hybrid dynamical systems, 2024.
[32] Yining Luo, Yingfa Chen, and Zhen Zhang. Cfdbench: comprehensive benchmark for machine learning methods in fluid dynamics. arXiv preprint arXiv:2310.0596. 2023.
[33] Philipp Moser, Wolfgang Fenz, Stefan Thumfart, Isabell Ganitzer, and Michael Giretzlehner. Modeling of 3d blood flows with physics-infomme neural networks: Comparison of network architectures. Fluids, 8(2):46, 2023.
[34] Alex Nichol, Joshua Achiam, and John Schulman. On first-order meta-learning algorithms. arXiv preprini /00::v/800009999 2018
[35] Samuele Papa, David M Knigge, Riccardo Valperga, Nikita Moriakov, Miltos Kofinas, Jan- Jakob Sonke, and Efstratios Gavves. Neural modulation fields for conditional cone beam neural tomography. arXiv preprint 2rXivv2307.08351 2023.
[36] Jeong Joon Park, Peter Florence, Julian Straub, Richard Newcombe, and Steven Lovegrove. Deepsdf: Learning continuous signed distance functions for shape representation. In Proceed- ings of the IEEE/CVF conference computer vision and pattern recognition, pages 165-174, 2019.
[37] Ethan Perez, Florian Strub, Harm De Vries, Vincent Dumoulin, and Aaron Courville. Film: Visual reasoning with a general conditioning layer. In Proceedings of the AAAI conference artificial intelligence, volume 32, 2018.
[38] Adeel Pervez, Francesco Locatello, and Efstratios Gavves. Mechanistic neural networks for scientific machine learning. arXiv preprint arXiv:2402. 2024.
[39] Tobias Pfaff, Meire Fortunato, Alvaro Sanchez-Gonzalez, and Peter W Battaglia. Learning besh-based simulation with graph networks. arXiv preprint arXiv:2010.03409, 2020.
12

---

[40] Michael Prasthofer, Tim De Ryck, and Siddhartha Mishra. Variable-input deep operator networks. arXiv preprint arXiv:2205.1140k 2022.
[41] Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling, E (n) equivariant graph neural networks. In International conference machine learning, pages 9323-9332 PMLR, 2021.
[42] Vincent Sitzmann, Eric Chan, Richard Tucker, Noah Snavely, and Gordon Wetzstein. Metasdf: Meta-learning signed distance functions. Advances in Neural Information Processing Systems, 2020.
[43] George Gabriel Stokes al. On the effect of the internal friction of fluids on the motion of pendulums. 1851.
[44] Matthew Tancik, Pratul Srinivasan, Ben Mildenhall, Sara Fridovich-Keil Nithin Raghavan, Utkarsh Singhal, Ravi Ramamoorthi, Jonathan Barron, and Ren Ng. Fourier features let networks learn high frequency functions in low dimensional domains. Advances in neural information processing systems, 33:7537-7547, 2020.
[45] Matthew Tancik, Ben Mildenhall, Terrance Wang. Divi Schmidt, Pratul Srinivasan, Jonathan Barron, and Ren Ng. Learned initializations for optimizing coordinate-basee neural repre- sentations. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 2846-2855, 2021.
[46] Riccardo Valperga, Kevin Webster, Dmitry Turaev, Victoria Klein, and Jeroen Lamb. Learning reversible symplectic dynamics. In Proceedings of The 4th Annual Learning for Dynamics and Control Conference, volume 168 of Proceedings of Machine Learning Research, pages 906-916. PMLR, 23-24 Jun 2022.
[47] Laurens Van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of machine learning research, (11)) 2008
[48] Maurice Weiler and Gabriele Cesa. General e (2)-equivariant steerable cnns, Advances in neural information processing systems, 32, 2019.
[49] David R Wessels, David M Knigge, Samuele Papa, Riccardo Valperga, Efstratios Gavves, and Erik Bekkers. Grounding continuous representations in geometry: Equivariant neural fields. ArXiv Preprint arXiv: 2024.
[50] Haixu Wu, Huakun Luo, Haowen Wang, Jianmin Wang, and Mingsheng Long. Transolver: fast transformer solver for pdes on general geometries. arXiv preprint arXiv:2402.02366. 2024.
[51] Yuan Yin, Matthieu Kirchmeyer, Jean-Yves Franceschi, Alain Rakotomamonjy, and Patrick Gallinari. Continuous pde dynamics forecasting with implicit neural representations. arXiv preprint arXiv:2209. 14855, 2022.
[52] Biao Zhang, Jiapeng Tang, Matthias Niessner, and Peter Wonka. 3dshape2vecset: A 3d shape representation for neural fields and generative diffusion models. ACM Transactions on Graphics (TOG), 42(4):1-16, 2023.
[53] Maksim Zhdanov, David Ruhe, Maurice Weiler, Ana Lucic, Johannes Brandstetter, and Patrick Forré. Clifford-steerable convolutional neural networks. arXiv preprint arXiv:2402. 14730, 2024.
[54] David Zwicker. py-pde: A python package for solving partial differential equations. Journal of Open Source Software, 5(48):2158, 2020. doi: 10.2 00.21105/joss.02111 02158. URL httph://///::
org/10. 10. 21105/joss. 02158
13

---

A Related work
DL approaches to dynamics modelling In recent years, the learning of spatiotemporal dynamics has been receiving significant attention, either for modelling interacting systems B1 30 scientific Machine Learning
(88 26 [53] or even videos M. Most DL methods for solving PDEs attempt to directly replace solvers with mappings between dineesiiinniinal Euclidean spaces, i.c. through the use of CNNs [19] 2 or GNNs
often applied autoregressively to an observed (discretized) PDE state. Instead, the Neural Operator (NO) paradigm attempts to learn fhiiite-dimensionll operators, i.c. mappings between function spaces, with limited success. Fourier Neural Operator (FNO) [2]] extends this method by performing convolutions in the spectral domain, FNO obtains much improved performance, but due to its reliance on FFT limited to data on
39
27
regular grids.
Inductive biases in DL and dynamics modelling Geometric Deep Learning aims to improve model generalization and performance by constraining/designin
model's space of learnable functions based on geometric principles. Prominent examples include Group Equivariant Convolutional Networks and Stecrable CNNs [2 4. generalizations CNNS that respect symmetries of the data such as dilations and continuous
15 25) Analogously, Graph Neural Networks (GNNs) (4 or Message Passing Neural Networks (MPNNS) 7 are a variant of neural network that respects set-permutations naturally found in graph data. They are typically formulated for graphs 9 (V, E). with nodes E V and edges E. Typically nodes are embedded into node vector ff, which is subsequently updated over multiple layers of message passing. Message passing consists (1) computing messages mas, over edges i. from node to with the message function (taking into account edge attributes as.: mi.j = 2m (fi f fj. (4.5) (2) aggegating incoming messages: m, EeeA(()) Mi.j. (fl, m.).
rotations 48
(3) computing updated node features f!²
Recently, such methods have also been adapted for sparse physical data, c.g. for molecular property prediction where the GNN is additionally required to respect transformation symmetrics. (5] unifics these approaches to equivariance under the guise of weight sharing over equivalence classes defined by bi-invariant attributes of pairs of nodes j, a viewpoint we leverage in constructing the equivariant conditioning latent zt" corresponding a PDE state vt. In the context of dynamies modelling, equivariant architectures have been employed to incorporate various properties of physical systems in the modelling process, examples of such properties are the symplectic structure [2]] discrete symmetries such as reversing symmetries [46] and energy conservation [8 21
41 6
Neural Fields in dynamics modelling Conditional Neural fields (NcFs) are a class of coordinate based neural networks, often trained to reconstruct bisscretely-sample input continuously. More specifically, conditional neural field fo R -> TRd a field -prrmmterrred by neural network with parameters 6- that maps input coordinates I E TR" in the data domain alongside conditioning latents z to d-dimensional signal values f(xx) E RR... By associating a conditioning latent 2" E IR to each signal v, a single conditional NeF fs R x _2 + _ed can learn to represent families of continuous signals such that V E D: f(x) = (x;:"") 13 showed the viability of using the latents z' as representations for downstream tasks (e.g. classification, generation) proposing framework for learning on neural fields. This framework inherits desirable properties of neural fields, such as inherent support for sparsely and/or irregularly sampled data, and independence to signal resolution.
51 propose to use conditional NeFs for PDE modelling by learning a continuous flow in the latent space of a conditional neural field. In particular, a set of latents {zī Vi= are obtained by fitting a conditional neural field a given set observations (v, T at timesteps 1, ..., T`: simultaneously, a neural ODE [I] Fs is trained to map pairs of temporally continuous latents s.t. solutions correspond to the trajectories traced by the learned latents. Though this approach yields impressive results for sparse and irregular data in planar PDEs, we show it breaks down on more challenging geometries. We hypothesize that this is due a lack of latent space that preserves relevant geometric transformation with respect to which systems we are modelling are symmetric, and as such propose an extension of this framework where such symmetries are preserved.
Obtaining Neural Fields representations Most NeF-based approach to representation or reconstruction use SGD to optimize (a subset of) the parameters of the NeF, inevitably leading to significant overhead in inference: conditional NeFs require optimizing (set of) latents from initialization to reconstruct for novel sample. Accordingly, research has explored ways of addressing this limitation. 22 45 propose using Meta- Learning [14 34 optimize for an initialization for the NeF from which is possible to reconstruct for a novel sample in as few as 3 gradient descent steps. [3] propose to meta-learn the NeF backbone, but fix the initialization for the latent z and instead optimize the learning rate used in its optimization using Meta-SGD
of a NeF and its value as downstream representation; 335 show that (1) using a shared NeF initialization and (2) limiting the number of gradient updates to the NeF improves performance
(28). Recently, work has also explored the relation between inttialization/opimiiza
downstream tasks, as this simplifies the complex relation between a NeFs parameter space and its output function space. We combine these insights and make Meta-Learning Learning part of our equivariant PDE solving pipeline, as it enables fast inference and we show it to simplify the latent space of the ENF, improving performance of the neural ODE solver.
14

---

B Pseudocode for optimization objective
See Alg. for pseudocode of the training loop that use, written for single datasample for simplicity of notation. For simplicity, we further assume we're 'e using an euler stepper to solve the neural ODE, but this can be replaced by any solver. For inference, this stratagem is identical, except we do not perform gradient updates to 0, 4b.
Algorithm 1 Optimization objective
Randomly initialize neural field fe Randomly initialize neural ode Fe while not done do
Sample initial states and coordinates Vo- Initialize latents 2/1 -((): c:)
for all step E iinitill state ope = 3 do
z - = - VVay Lmc end for
for all t E [1, tinl do
2/ 2 + So d2)
end for Update 0, per: per: 4-0 0 1 V/ v0 t- v - 1/ Cmse with Lime = -(f((==)))
end while
C Equivariant Neural Fields
ENF to reconstruct PDE states For ease of notation we denote P and C the matrices containing poses and corresponding appearances stacked row-wise, i.. P..: pf and C..: cf Furthermore, we denote A as the matrix containing bi-invariants asm. stacked row-wise, i.e. Ai, a.
Q(A) (C) + G(A) V(C:A). V(C.A) vdz
fx(x;2") = softmax
(10)
where the softmax is applied over the latent set and with de the hidden dimensionality of the ENF. The query matrix Q is constructed Q=W, 70 (A), 74 Gaussian RFF embedding (44) followed by linear layer W., i.c. Q consists the RFF embedded bi-invariants of the input coordinate xm and each of the latent poses stacked row-wise. The key matrix is given by a learnable linear transformation W; of the context vectors & C". The attention coefficients which result from the inner product of Q, K are weighted by Gaussian window G whose magnitude is conditioned
c: K=
distance measure the relative distance between latent poses and input coordinates as: G. = (lpp x(l²). with du hyperparameter which determines the locality of cach of the latents. Finally the value matrix is calculated as a Icarnable linear transformation W. of the appearances A, conditioned through FiLM modulation
(37 by second RFF embedding of the relative poses split into scale- and shift modulations: V=W,, A & W na Yon (A) + W Yoa (A). The latents zi are optimized for single state vs whereas the parameters 0 of the ENF backbone which consist of all the
learnable parameters of the linear layers W.. WR. W., W.. W used to construct Q, K, are shared over all states.
The overall architecture consists of linear layer WR - Idd applied to c, E IR followed by a layemnrmm After this, the cross attention listed above is applied, followed by three d-dim linear layers, the final one mapping the output dimension nu
Equivariance follows from sharing Q K, V over equivalence classes Note that the latent space of the ENF is equipped with group action as: yaī [[gps, a.) As an example, (E(2))qquivarian)) of the ENF follows from bi-invariance of the quantity used construct Q under the group action:
Vg SF(n): (pt,
(gp..gx)
- (gpx) p. gx pi g. (11)
And so, constructing the matrix containing the relative poses of bi-transformed poses and coordinates (gP)-1 as ((gP) gx);;; = p. r gx = p. x, we trivially have:
Vg SE(n) (pi, x) (g
Q(A) Q(AA) Q(A).
(12)
15

---

D Defining additional bi-invariant attributes
Other examples of the bi-invariants attributes that are used in the experiments section are listed here.
Full iotatio symmetries on the 2-sphere For the global shallow water equations defined a SW attrihute that is bi-invariant only to rotations over globe's axis, ie. rotations over Φ. In experiments also solve diffusion over the sphere. which is fully SO(3) rotationally symmetric. To achieve equivariance to 3d rotations, we take poses p E (3) parameterized by culer angles which act on points æ E S parameterized by 3D unit vectors x through 3D-rotation matrices, allowing us to calculate the bi-invariant p 2::
80(3)
R,Xm
(13)
This bi-invariant is used in our experiments for diffusion on the 2-sphere.
The 3D ball B². We experiment with Boussinesq equation for internally heated convection in a ball. The PDE fully rotationally symmetrie, but since the source K is at fixed point (the center of the ball resp.), is not symmetric to translations of the initial conditions within the ball. As such, we let p E SO(3) x with Φ, 0. s.t. 0 < r < 1. The PDE is defined over spherical coordinates (ø, 0, r), which we map to vectors in x E R We then use the following bi-invariant, which is only symmetric rotations in SO(3):
Rm Tpp Tam-
(14)
No transformation symmetries. A simple "bi- "bi-invariant" this setting that preserves all geometric information is given by simply concatenating coordinates p with coordinates x:
p: Im
(15)
Parameterizing the cross-attention operation in Eq. 5 as function of this bi-invariant results in a framework without any equivariance constraints. We use this in experiments to ablate over equivariance constraints and its impact performance.
E
Experimental Details
E.1
Dataset creation
For creating the dataset of PDE solutions we used py-pde 5 for Navier-Stokes and the diffusion equation on the plane. For the shallow-water equation and the diffusion equation on the sphere, well as the internally heated convection in 3D ball we used Dedalus [T0]
Diffusion on the plane. For the diffusion equation on the plane we use as initial conditions narrow spikes centred at random locations in the left half of the domain for the train set, and in the right half of the domain for the test set. States are defined on 64 x 64 grid ranging from -3 to 3. Initial conditions are randomly sampled uniformly between -2 and 2 for and 0 and 2 for y in the training for y. A random value uniformly sampled between 5.0 and 5.5 is inserted at the randomly sampled location. We solve the equation with an Euler solver for 27 steps, discarding the first 7, with a timestep di = 0.01. We generate 1024 training and 128 test trajectories.
and between -2 and 2 for and -2 and 0
Navier-Stokes on the flat 2-torus. For Navier-Stokes on the flat 2-torus we Gaussian random fields initial conditions and solve the PDE using a Cranck-Nicholson method with timestep dt 1.0 for 20 steps. u, Vu O, where u the velocity field, e the vorticity, μ the
The PDE du - Vvv + vap + f.v J7 viscosity and f forcing term
du
= a VVv + vAp +
di
V
Vu O,
where u is the velocity field, the vorticity, the viscosity and f forcing term. We set viscosity to 1E 3. resulting with our setup in Reynolds number of ~ 1 1000. States are defined on 64 x 64 grid. We generate 8192 training and 512 test trajectories.
Diffussin on the 2-sphere. For diffusion dataset on sphere, states are defined over 128 64 s 0 grid. Initial conditions are generated as a gaussian peak inserted at a random point on the sphere with = 0.25. The equation is solved for 20 timesteps with RK4 and dt = 1.0. We generate 256 training and 64 test trajectories.
16

---

Spherical whallow-water equations [16].
The global shallow-water equations are
du fk x u VVh + vAu
di
dh AV PAh,
di
where d. is the material derivative, is the unit vector orthogonal to the surface of the sphere, u is the velocity field that is tangent to the spherical surface and and is the thickness of the fluid layer. The rest are constant parameters of the Earth (see [16] for details). As initial conditions follow (16) and use basic zonal flow, representing a mid-latitude tropospheric jet, with correspondingly balanced height field.
for S do
((6)=
Mmax exp
for do < < Φ1
(0) (1
0
for 2 11
Where Mmux 80ms''. 6o x x/7,, +1 7/2 Φ1, and cn expl 1(01 }). With this initial zonal flow, numerically integrate the balance equation
f au((') s+ tan(o") u(dd) 3 de',
(() g/u
to obtain the height A. We then randomly generate small un-balanced perturbations Jt' to the height field
h' 1' ()) is cos())e-(02--1/0) -(((2=-0)//))
by uniformly sampling a. B. h. 02, and d2 within neighbourhood of the values use in 16. States are defined
192 x 96 grid for the high- resolutioo dataset, which is subsequently downsampled by 2 2 mean pooling a 96 x grid. We generate 512 training trajectories and 64 test trajectories.
Internally- heated convection in the ball. The equations for the internally- heated convection system listed here, they include thermal diffusivity (x) and kinematic viscosity (1)) given by:
(Ra Pr)
1/2
Ra
1/2
Pr
We set Ra = and Pr
Incompressibility condition (continuity equation):
u Tp =0
2. Momentum equation (Navier-Stokees equation):
du v Vp lift((T.) - (V x xu)
Temperature equation:
ØT 0t
2 lift((T) VT + KT source
4. Shear stress boundary condition (stress-free condition):
Shear Stress = 0 on the boundary
5. No penetration boundary condition (radial component of velocity at r 1):
(u(a = 1))
Thermal boundary condition (radial gradient of temperature 1):
radial(VT(r 1)) -2
7. Pressure gauge condition:
f pdv
The boundary conditions imposed stress-free and no-penetration for the velocity field and constant thermal flux at the outer boundary. These conditions are enforced using penalty terms (r) that are lifted into the domain using higher-one order basis functions.
x 24 Φb 0, grid. We use SBDF2 solver which we constrain by dtmin le 4 and dtmux = 2e 2. We evolve the PDE for 26 timesteps, discarding the first 6. We generate 512 training trajectories and 64 test trajectories.
States are defined over a 64 x 24x
17

---

E.2 Training details
We provide hyperparameters per experiment. We optimize the weights of the neural field fe, and neural ODE F& with Adam [33 with learning rate of IE-4 and IE-3 respectively. We initialize the inner learning rate that we use in Meta-SGD (28) for learning x" at 1.0 for p and 5.0 for c. For the neural ODE Fs. we use of our message passing layers in the architecture specified in [5]. with hidden dimensionality of 128, The std parameter of the RFF embedding functions "e- "wo : YYo (see Appx. C is chosen per experiment. We run all experiments single A100. All experiments are ran 3 times.
Diffusion on the plane. We use 4 latents with E R We set the hidden dim of the ENF to 64 and use 2 attention heads. We train the model for 1000 epochs. We set Y4 0.05, Yeu 0.01, Yog 0.01. We use batch size of 8. The model takes approximately 8 hours train.
Navier-Stokes on the flat 2-torus. We use 4 latents with e E R We set the hidden dim of the ENF to 64 and use 2 attention heads. We train the model for 2000 epochs. We set % 0.05, Yo.. 0.2, Yug 0.2. We use a batch size of 4. The model takes approximately 48 hours to train.
Diffusion on the 2-sphere. We use 18 latents with c E R We set the hidden dim of the ENF to 16 and use 2 attention heads. We train the model for 1500 epochs. We set To = 0.01, Te. 0.01, Yog 0.01. We use batch size of The model takes approximately 12 hours to train.
Spherical whallow-water equations [16]. We use 8 latents with e E 2. We set the hidden dim of the ENF to 128, and use 2 attention heads. We train the model for 1500 epochs. We set Ti 0.05, Tea 0.2, Tua- Tvp = 0.2. We use batch size of 2. The model takes approximately 24 hours to train.
Internally-heated convection in the ball We use 8 latents with c E R We set the hidden dim of the ENF to 128, and use 2 attention heads. We train the model for 1500 epochs. We set Yy = 0.05, Vea 0.2, Tvg = 0.2. We use a batch size of 2. The model takes approximately 24 hours to train.
We use 25 latents with E RR16. We set the hidden dim of the ENF to 128, and use I attention head. We set 7/4 = 0.05, Ph. = 0.1, 0.1. We train the model for 100 epochs on a single A 100 GPU, taking about hours.
CFDBench 32
Baselines As baseline models on Navier-Stokes we train FNO and GFNO 22 with modes and 32 channels for 700 epochs (until convergence). We train CNODE 2 with 4 layers of size 64 for 300 epochs (until convergence). We train DINo on experiments for 2000 epochs with an architecture as specified in [5T]. For the IHC and shallow water experiments, we increase the latent dim from 100 to 200, the number of layers for the neural ODE from 3 to 5, and the latent dim of the neural field decoder from 64 to 256, as per [51]
For the Navier-Stoli Stokes and Internally- -nternaly-HHeated Convection experiments we additionally train the Transolver [50] model. We adapt the Transolver model in PyTorch from the official github repository, and use the same hyperparameter settings as in the original paper (resulting in 7.1M parameter model)- - modifying the training objective to be identical to the autoregressive one we use in our experiments (i.c. mapping from frame to frame). During training in the Navier-Stokes task we observed noticeable instabilities. with the method producing feeuueryy artefacts in rollouts in this autoregressive setting. For sparsely subsampled initial conditions. performance deteriorates further- highlighting the benefit of NeF-based continuous PDE solving.
For experiments internally-heat convection, due to the large size of the input frames. we were required to scale down the Transolver model size in order to able fit the model on our A100 GPU, from 256 64 hidden units, resulting in a 1.2M parameter model somewhat comparable in size to the model we use in our experiments (889K). We train the model for 2000 epochs on an A100 GPU, taking approximately 30 hours.
F Additional results
Parameter counts, memory and time efficiency In order to compare parameter, memory and time efficiency of our method to other baselines, we provide details of our models compared to [1]] in Tab. 7 and memory usage/inference time compared to the models used as baselines in the Navier-Stokes experiments i We note that although our method is the most memory efficient- meta-leaming requires a significant computational overhead it the fastest in inference settings when unrolling an unseen state. This is attributable to (1) the fact that we're re using a latent-space space solver which operates on a drastically compressed representation of the PDE state, and use (2) meta-learning over decoddngg which means we only require 3 gradient descent steps to obtain the initial state compared to 300-500 typically used in auto- decoding [51].
18

---

Table 7: Parameter count for model used in Table 8: Parameter count and runtimes for for models Navier-
Navier-Stokes, Stokes, Shallow- Water and IHC exper- used in Navier-Stokee Stoker experiments. Note that training runtimes are measured per epoch, inference runtimes measured for unrolling single st-s trajectory. GPU memory allocations are measured per trajectory.
iments.
PARAMS /s
FARAMS
NAVIER-STOKES
DINo 51
355K
502K
835K 885K
RUNIMEILY TRAINING
NINe//Ta GPU INTERENCE
354k
531k
S SHALLOW- WATTE
INO
5415
00.75
0.3502b
5t.0ss
0
DINo 3T
639K
902K
1.5M
CNODE
GENO
1.8M
3,551
1450.0s
200.56
1.450h
2,0802
423 Ims
1704ms
0.7C
1,2K
356k
S33K
859
DINo
835L
25.8%
0.230h
61.7us
O.14GB
Translvvrr 50
7.1M
619.25
9,42Gb
43.6ms
5,19GB
DINo S1
639X
702K
L3M
356k
533K
889K
1.35
gound
0.25
FMO
0.25
pind
0.00
Dit
040
TND
Figure 9: Error accumulation Figure 10: Error accumulation
Figure 11: Visualization of an 80-step rollout for Navier-
over long rollouts for Navier- Stokes test set with vo 100% observed, train horizon lin is marked in green.
over long rollouts for Navier- Stokes test set with up 50% observed, train horizon lin is marked in green.
Stokes test sample.
Error accumulation on long rollouts for Navier-Stokes In order to better empirically assess accumulation of method for long-term extrapolation, provide experimental results for unrolling of Navier-Stokes for 80 timesteps. We apply our model both setting with fully observed initial conditions sparsely observed (50%) initial conditions, and provide plots of accumulated MSE along these trajectories in Figs. and visualization of the solution and error in Fig. [I] We compare against DINo [1] and FNO (29), and show that consistently improve reduced error accumulation the fully observed setting. but very rapidly deteriorates
DINo in terms long-tern -erm extrapolation. FNO achieves better
extrapolation limits even within train horizon with sparsely observed initial conditions, whereas our proposed approach loses very little in terms of extrapolation performance in this sparse setting.
19

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We feel the abstract and introduction accurately reflect contributions made paper, the experiments in Sec. | clearly show that our framework achieves improved performance exactly because of the equivariance constraints.
Guidelines:
The answer NA mcans that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be pereeived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine include aspirational goals motivation long it is clear that these goals attained by the paper.
2.
Limitations
Question: Does paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We discuss some limitations of the model's performance, namely the inability to generalize well on fou for the global shallow-wate wattrr equitions, and our hypothesis for the cause of this; namely limited reconstruction accuracy of the decoder on complex states v.
Guidelines:
The answer NA mcans that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated practice and what implications would
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance the approach. example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicaalle, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their hest judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Question:
Answer: [Yes]
Justification: The main claim of the paper; that we obtain equivarance by shahitg over bi- invariances is proven in Eq.
Guidelines:
20

---

The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should he numbered and cross-referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in core paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: results should reproducible based only on the contents of this paper and references made in it. Code will be made available for a camera-ready version, which also includes code for generating the datasets of solutions and experiments.
Guidelines:
answer means that the paper does not include experiments.
the paper includes experiments, answer to this question will not be perceived well by the reviewers: Making the paper reproducible important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible verifiable.
Depending on the contribution. reproducibility can be accomplished in various ways. For example, the contribution novel architecture, describing the architecture fully might suffice, or the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access to a hosted model (e-g., the case of a large language model), releasing model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily reproduce that algorithm.
algorithm, the paper should make clear how
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (c.g., a large language model), then there should either be way to access this model for reproducing the results way to reproduce the model (c.g.. with open-source dataset instructions for how construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-soure source models, may be that access the model is limited in some way (e.g., to registered users), but it should be possible for other rescarchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data code, with sufficient instructions to faithfully reproduce the main experimental results, as described supplemental material?
Answer: [Yes]
Justification:
Code added to the submission.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines ttttps://nips.cc/puttt
quuddes/CooeeStttmiesssooooooode for more details.
21

---

While we encourage the release of code and data, we understand that this might not be possible. "No" is acceptable answer. Papers cannot be rejected simply for not including code, unless this central to the contribution (e.g.. for a new open-soure benchmark).
The instructions should contain the exact command and environment needed the results. See the NeurlPS code and data submission guidelines (https: http:s////iiiccpeeeeeeeeeeees cc/pbiiee/
to reproduce
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for new proposed method and baselines. If only a subset experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release amonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g.. data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
Justification: These details are listed in Appx.
E
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, supplemental material.
Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate informa- tion about the statistical significance of the experiments?
Answer: [Yes]
Justification: std is given for experimental results over 3 runs.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error hars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be elearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to library function, bootstrap, etc.)
The assumptions made should given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK report I-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error than state that they have 96% CI, if the hypothesis of Normality of errors not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
error bars are reported tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
Question: For each experiment. does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Answer: [Yes]
Justification: We provide details on the hardware on which we run the experiments, the approximate runtime per This should be enough information to reproduce.
22

---

Guidelines:
The answer NA means that the paper does not include experiments.
paper should indicate the type compute workers CPU or GPU, internal cluster. or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't It make into the paper).
9. Code Of Ethics
Question:
Does the research conducted the paper conform. every respect, with NeurIPS Code Ethics
Answer:
[Yes]
Justification:
We respect the code of ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Codc of Ethics.
If' the authors answer No. they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure preserve anonymity (e.g.. if there is special consideration duc laws regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
The paper provides way for improving PDE forecasting. This can have many positive societal impacts, e.g. improving weather forecasting for early warning, material sciences. However, the framework could provide erroneous predictions, without any indication confidence, which could have major negative consequences. Future work should look into associating confidence intervals to predictions, to get indications of the viahility of results.
Justification:
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g.. disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deploy- ment technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used generate deepfakes for disinformation. On the other hand, it not needed point out a generic algorithm optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the tcchnology.
there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release models, providing defenses in addition to attacks, mechanisms monitor- ing misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of MI.).
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g.. pretrained language models, image generators, scraped datasets)?
Question:
Answer: [NA]
23

---

Justification:
We want the paper to reproducibel and as such we provide code. There are no major risks for misuse.
Guidelines:
answer means that the paper poses no such risks.
Released models that have high risk misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners assets (c.g., code, data, models), used the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: Authors original owners, credit relevant owners required.
Guidelines:
The answer means that the paper does not use existing assets.
authors should cite the original paper that produced the code package or dataset.
authors should state which version
asset is and, possible, include a URL.
name of license (e.g.. CC-BY 4.0) should be included each asset. scraped data from particular source (e.g., website), the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms of use the package should be provided. For popular datasets, papppaaii www.rewithhccd.... com/ dataseta has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information not available online, the authors encouraged reach out the asset's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: Paper does not provide assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Rescarchers should communicate details of the dataset/code/moodl as part of their sub- missions via structured templates. This includes details about training, license, limitations,
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: Not applicable.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
24

---

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then much detail possible should be included in the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least minimum wage country data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Does paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Question:
Answer: [NA]
Justification: Not applicable.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained TRB approval, you should clearly state this the paper.
recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere to the NeurlPS Code of Ethics the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applica- ble), such as the institution conducting the review,
25


Reviews:
[{"review_id": "6wdFB0qNb4", "reviewer": "Reviewer_NkrS", "summary": "The paper introduces a novel framework that leverages Equivariant Neural Fields (ENFs) to solve Partial Differential Equations (PDEs). By preserving geometric information in the latent space, the proposed method respects the known symmetries of the PDE, enhancing generalization and data efficiency. The framework demonstrates improved performance in various challenging geometries, validated through experiments against other neural PDE forecasting methods.", "questions": "None", "limitations": "The framework assumes that the boundary conditions are symmetric and that the PDEs exhibit certain symmetries. In real-world applications, these assumptions might not always hold, potentially limiting the applicability of the proposed method to PDEs with different or more complex boundary conditions.\n\nThe use of a global attention operator in the Equivariant Neural Field (ENF) backbone increases the computational complexity. This can lead to high computational costs, especially when scaling the model for larger datasets or more complex PDEs.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper presents an innovative approach by integrating equivariant neural fields, which respect the symmetries of PDEs, thereby enhancing model performance.\nThe methodology addresses significant limitations of existing NeF-based PDE solvers, particularly in generalization to unseen spatial and temporal locations and geometric transformations.\nExtensive experimental validation across various geometries (e.g., plane, torus, sphere) demonstrates the robustness of the proposed framework over existing methods.", "weaknesses": "The framework's performance decreases when extrapolating beyond the training horizon for complex PDEs.\nWhile the approach shows competitive performance, the computational complexity due to the global attention operator in the ENF backbone can be high.\nError accumulation in long-term predictions could be mitigated with increased model capacity, but this comes at the cost of computational resources.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their thorough assessment of our work, and we’re happy to see that the reviewer deems our approach innovative and appreciates the experimental validation we provide. We thank the reviewer for highlighting a number of important considerations with regards to our proposed approach, and discuss limitations mentioned by the reviewer in detail here. We would love to hear your thoughts.\n\n**Equivariance constraints in real-world applications** The reviewer mentions that our method assumes symmetric boundary conditions and/or symmetric PDE formulations, and remarks that in real-world applications, such information might not always be known - limiting applicability of our approach in these cases.\n\n*Note: Since R. i6aX, hu7A also ask after usability with more complex boundary conditions and applicability when PDE symmetries are unknown, we provide a joint response here.*\n\nIn general, not knowing the boundary conditions in advance when solving a PDE would result in an ill-posed problem. Boundary conditions help ensure the uniqueness and stability of the solution. From given boundary conditions it is immediate to verify if they are preserved by a given group action - we show examples of this in the manuscript Sec. 3.1 with Eqs 7, 8.\n\nThe concern raised about not always knowing the symmetry in advance is a very valid point, but it must be placed within the broader context of inductive biases in deep learning. First, our method extends existing approaches by introducing equivariance - an inductive bias that has shown to be beneficial in many DL-based scientific research applications [3, 4] - to improve forecasting accuracy and consistency for PDEs with known symmetries. Furthermore, our proposed attention architecture with Gaussian windows results in local latent variables that can be beneficial regardless of whether or not symmetries are present. Second, it is often the case - also in real-world applications - that some a priori knowledge about the system is available. It is reasonable to assume, for example, that a spherical domain, such as in the shallow water system, is equivariant with respect to rotations around the axis of rotation. On the contrary, full SO(3) equivariance in real-world climate data defined over the globe can be ruled out a priori by noting that Coriolis (fictitious) forces, due to the rotation of the sphere, would break this symmetry. \n\nNote that our method can indeed be applied when symmetries are unknown or inexact, as we show with an additional experiment applied on CFDBench - as suggested by R. kxDM and shown in Tab. 6. We hypothesise that often the differential operators that define PDEs themselves are symmetric (i.e. derivatives, Laplacians), so introducing these same symmetries into the building blocks of the neural solver itself might still be beneficial to reduce the problem complexity, but an extensive investigation would be needed to properly verify this intuition. Another possibility would be to allow for soft-equivariance constraints to be imposed on the modelled solutions. A simple approach to this end - gleaned from application of GNNs used in latent force field discovery in particle physics [8] - would be to attach one or more symmetry-breaking dataset-wide shared “reference frame nodes” to the latent ODE graph, that could encode for any symmetry-breaking biases that the data contains, to e.g. account for absolute positioning of hindrances or objects within the domain.\n\nRecent works in equivariant deep learning have explored symmetry-discovery / soft symmetry constraints in model design, highlighting another interesting research direction. E.g. [5] initialise their model as being fully equivariant, but allow relaxation of this constraint through optimization, learning an interpolation factor between a set of equivariant and non-equivariant kernels. [6] instead proposes to learn generators for symmetry groups from data directly, overcoming the need for explicit specification of the specific equivariance constraints for novel PDEs - and provides a very extensive list of examples of symmetries common to a variety of PDEs (moreover noting that “simplifying systems of PDEs using their symmetries and consequent coordinate transformations was, in fact, the primary reason Sophus Lie discovered Lie groups”). Though it falls outside of the scope of our current work, similar adaptations might help the proposed framework transfer effectively to settings with inexact or unknown symmetries.\n\n **Scalability and model complexity** The reviewer remarks that the global attention operation used in the decoder may limit scalability of our method. Since our latents are localised we can approximate the attention operator by restricting the number of latents that are attended to based on their relative distance to the sampled coordinate. We provide a little further detail in our response to R. hu7A.\n\n[1] Li, Z., et al. Fourier Neural Operator for Parametric Partial Differential Equations. In ICLR.\n\n[2] Yin, Y., et al. Continuous PDE Dynamics Forecasting with Implicit Neural Representations. In The Eleventh International Conference on Learning Representations.\n\n[3]  Helwig, J., et al (2023). Group equivariant fourier neural operators for partial differential equations. arXiv preprint arXiv:2306.05697.\n\n[4] Bogatskiy, A., et al. (2022). Symmetry group equivariant architectures for physics. arXiv preprint arXiv:2203.06153.\n\n[5] Wang, R., et al. (2022, June). Approximately equivariant networks for imperfectly symmetric dynamics. In International Conference on Machine Learning (pp. 23078-23091). PMLR.\n\n[6] Gabel, A, et al. (2024). Data-driven Lie point symmetry detection for continuous dynamical systems. Machine Learning: Science and Technology, 5(1), 015037.\n\n[7] Kofinas, M, et al. (2024). Latent field discovery in interacting dynamical systems with neural fields. Advances in Neural Information Processing Systems, 36.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "cGtuhQ8m20", "reviewer": "Reviewer_kxDM", "summary": "The work proposes a novel framework combining equivariant neural fields and neural ODEs, providing a continuous space-time solution for PDEs while respecting associated equivariance constraints. The author uses PDE-specific bi-invariant attributes for equivariant neural fields and a meta-learning approach for learning the initial latent state. The proposed method achieves better performance on the chosen PDE problems.", "questions": "1. What is the training and inference time of the proposed method compared to existing methods like FNOs and Deeponets?\n\n2. what is the Reynolds number of the Navier-Stokes equation problem?", "limitations": "yes", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The work addresses an important and complex issue of equivariance in the context of solving partial differential equations (PDEs). The proposed architecture is not only space-time continuous but also respects the equivariance constraint. This characteristic makes it particularly valuable and effective for various types of scientific research and applications.\n\n\n2. The proposed method is well-motivated and clearly explained in the paper.", "weaknesses": "I have found the empirical study to be the weak point of the work. In order to argue the effectiveness of the proposed solution over existing approaches, the authors need to consider established benchmarks, large-scale datasets, PDEbench, and CFDBench, especially with irregular domains (domains with holes or solid objects as hindrances).\n\nI also find that the choice of baselines is not extensive. For example, SFNO [a] is used for the shallow water equation. Also, baselines like [b,c] are not considered.\n\nAlso, as the proposed solution is claimed to be time continuous, zero-shot super-resolution along the time domain should be demonstrated (analogous to Table 3). \n\n\na. Spherical Fourier Neural Operators: Learning Stable Dynamics on the Sphere\n\nb. GNOT: A General Neural Operator Transformer for Operator Learning\n\nc. Geometry-Informed Neural Operator for Large-Scale 3D PDEs", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank reviewer kxDM for their effort in reviewing our work. We’re glad to see that the reviewer agrees on the value of adding equivariance constraints to neural PDE solving.\n\n**Improving comparative analysis** The reviewer raised concerns about comparison with a wider set of baselines. To this end, we include a general purpose transformer PDE solving architecture (Transolver [1]) to the list of baseline by applying it both to 2D Navier Stokes and 3D internally heated convection, since this model has shown SOTA results on PDE solving over different geometries.\n\nWe adapt the Transolver model in PyTorch from the official github repository, and use the same hyperparameter settings as in the original paper (resulting in a 7.1M parameter model) - modifying the training objective to be identical to the autoregressive one we use in our experiments (i.e. mapping from frame to frame). Tab. 5 shows the results for the Transolver model on Navier-Stokes in 2D. Notably, Transolver achieves 1.80E-02 and 1.85E-02 train and test mse respectively within the training horizon, and 4.85E-01 and 4.90E-01 train and test mse outside the training horizon. During training in the Navier-Stokes task we observed noticeable instabilities, with the method producing high-frequency artefacts in rollouts in this autoregressive setting. For sparsely subsampled initial conditions, performance deteriorates further - highlighting the benefit of NeF-based continuous PDE solving.\n\nFor experiments on internally-heated convection, due to the large size of the input frames, we were required to scale down the Transolver model size in order to be able to fit the model on our A100 GPU, from 256 to 64 hidden units, resulting in a 1.2M parameter model - somewhat comparable in size to the model we use in our experiments (889K). We train the model for 2000 epochs on an A100 GPU, taking approximately 30 hours. Results in Tab. 4 show the performance of this model within the training time horizon - achieving 4.13E-04 test MSE where our framework achieves 5.99E-04, indicating that it is indeed a strong baseline for solving PDEs over complicated geometries. However, we also note that outside of the training horizon, error accumulates more quickly for the Transolver model, indicating it has somewhat overfit the training horizon dynamics. Here, Transolver achieves 2.09E-02 test MSE, where our framework achieves 8.21E-03. We hypothesise that - due to the equivariance constraints placed on our model reflecting inductive biases about this PDE, it extrapolates more reliably beyond the train horizon.\n\nWe also provide results in Tab. 4 for sparsely observed initial conditions. These results clearly show the advantage of NeF-based continuous solvers in this sparsely observed setting, both DINo and our framework significantly outperform Transolver, which is unable to provide accurate solutions either within the training horizon or outside of it.\n\n**Application to irregular domain** Secondly, the reviewer raised the concern that the method is only tested on regular domains and asked whether we could extend the empirical study to irregular domains as well. We apply our model to the recommended CFDBench dataset [2], which contains solutions to initial conditions for computational fluid dynamics problems with varying boundary conditions, fluid properties and domain shapes - and applied our method on three different problems (cavity, dam and cylinder). We fit these problems with translational bi-invariants (which results in similar translation equivariance constraints to CNN-based neural PDE solvers) - keeping the same model architecture we used in all our experiments, with 25 latents and context vectors $\\mathbf{c_i} \\in \\mathbb{R}^{16}$ and without any specific finetuning. The results can be seen in Tab. 6 in the PDF and show that our approach also handles these more complicated and varied geometries and boundary conditions well - even without the presence of the global symmetries. We hypothesise that the weight-sharing that the equivariance constraints result in might have a regularising effect beneficial to PDE solving; PDEs are built from differential operators that are themselves generally equivariant.\n\n**Continuous-time solving** We provide results for a zero-shot temporal superresolution experiment to show the continuous-time property of our solving framework. We generate a dataset of Navier-Stokes solutions at a higher time-resolution, with timestep size $d\\tau=0.25$. We train our model on temporally subsampled frames resulting in a training time resolution of $d\\tau=1.0$. We then evaluate the model on step-sizes $d\\tau=0.5, d\\tau=0.25$. As shown in Tab. 3 of the attached PDF, both inside and outside the train horizon the higher sampling resolution introduces very little error, even with $4\\times$ as many unrolling steps, showcasing reliable continuous-time performance of the proposed approach.\n\n**Experimental details / computational requirements** The reviewer asked what the Reynolds number was used to generate the Navier Stokes data. In our experiment we use viscosity $\\nu=$ 1e-3, which in our setup results in a Reynolds number of $\\frac{1}{\\nu} \\sim 1000$.\n\nWe provide training and inference times for our method compared to the baseline models in Tab. 2 of the PDF. We note that although memory-wise our method has somewhat significant requirements (due to memory-intensive meta-learning), during inference our method outperforms the baselines.\n\n[1] Wu, H., Luo, H., Wang, H., Wang, J., & Long, M. (2024). Transolver: A fast transformer solver for pdes on general geometries. arXiv preprint arXiv:2402.02366.\n\n[2] Luo, Y., Chen, Y., & Zhang, Z. (2023). Cfdbench: A comprehensive benchmark for machine learning methods in fluid dynamics. arXiv preprint arXiv:2310.05963.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "Vp195XuHrB", "reviewer": "Reviewer_i6aX", "summary": "The paper attempts to learn the dynamics of certain PDEs from time series data using implicit neural representations, while encoding symmetry information of the domain. In fact, constructing a neural model that is aware of Euclidean transformations is the primary focus of this paper. To this end, the authors design two equivariant neural networks. Given an initial condition, a latent state is obtained by meta-learning. This latent state is then integrated in time by a neural ODE (first network) to obtain a final latent state. The second network then takes this final latent state as an input, and maps any given point coordinate (in the domain) to the solution at the final time. Examples are presented on periodic domains in $ \\mathbb{R}^2 $, 2-torus, 2-sphere and the 3D-ball. The paper builds on a 2022 ICLR paper [1] which attempts the same, but without any symmetry assumption.\n\n[1] Yuan Yin, Matthieu Kirchmeyer, Jean-Yves Franceschi, Alain Rakotomamonjy, and Patrick Gallinari. Continuous pde dynamics forecasting with implicit neural representations. 2022.", "questions": "How will this method work with other boundary conditions?\n\nIn the examples, the invariance groups are chosen according to the equation at hand. But how does one choose the invariance groups when the underlying operators and functions (RHS) are unknown?\n\nThe extrapolation results are compared with ground truth data, and it is seen that the accuracy deteriorates as the inference horizon goes farther from the  training horizon. In a new application, how to determine a limit for accurate extrapolation, i.e., the temporal horizon where extrapolation is always successful? Does this limit exist?\n\nThe forward problem triggers an ODE solve. What is the typical DOFs associated with this ODE solve?\n\nWhat is the boundary condition applied on the heat equation example?", "limitations": "See weaknesses.", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The paper is well written.\n\nPDEs are often posed on domains that have symmetric properties. This is in addition to the fact that the operators appearing in the PDE have their own symmetry / directional properties. While learning from data, most of the existing methods attempting to learn PDE dynamics ignore the symmetry information. Therefore, this is a welcome idea.\n\nThe method exhibits impressive extrapolation results.", "weaknesses": "Only infinite domain (or periodic boundary conditions) are considered.\n\nIn the examples, the transformation groups are chosen by carefully considering the nature of the domain and the operators appearing in the equations. But in a real application, this information, especially the operator information, is not known a priori.\n\nExtrapolation results are shown where results outside the training time horizon are predicted. The problem with such prediction is that they look good until they do not. And there is no logical or analytical bound on the time horizon where the extrapolation is supposed to work. The time horizon is always chosen so as to exhibit the effectiveness of the method. But no analysis is presented in that regard. Therefore such extrapolation results, even though impressive in some respects, do not add to either the understanding or the applicability of this method to a new application.\n\nMemory and execution (training) times are not compared (only the training times of the proposed method are included). Error comparisons are made with other methods. Sometimes this method outperforms the other methods (e.g., Table 2), but in some cases, it is marginally better than the others (e.g., Table 3, 4). Providing memory and training times would make these comparisons more well rounded.\n\n(Minor) Line 276: should be $ \\nabla \\cdot u = 0 $.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their very thorough treatment of our manuscript, and appreciate the recognition of the importance of encoding symmetry information in neural PDE solvers. We also acknowledge the reviewer's constructive feedback and address their concerns in detail below.\n\n**How will this method work with other boundary conditions?** In the experiments, we consider both domains like the flat torus (square with periodic boundaries) or the 2-sphere, where no boundary conditions are needed because they are enforced by construction of the bi-invariant, and domains such as the 3-sphere where we use explicit boundary conditions, in particular, zero radial velocity and fixed temperature gradient at the boundary of the sphere. The method proved to work well in both cases. \n\nOther boundary conditions may be incorporated likewise, through appropriate construction of a bi-invariant that respects these boundary conditions. From given boundary conditions it is immediate to verify if they are preserved by a given group action - we show examples of this in the manuscript Sec. 3.1 with Eqs 7, 8. As a further example, for instance for many cylinder flow problems there would exist vertical flip symmetries (the $Z_2$ symmetry group), which could be incorporated by choice of appropriate bi-invariant.\n\nReviewers R. i6aX, R. NkrS and R. hu7A all ask about the usability in settings when the PDE has no underlying symmetries or when the symmetries are unknown - as such we give an answer jointly. Please see our response to R. NkrS. - under _Equivariance constraints in real-world applications_.\n\n**Limits for successful extrapolation** The reviewer touches upon a very valid drawback of deep-learning based surrogates, namely the fact that to a large extent these methods do not provide any guarantees on extrapolation error bounds as they are purely learned from data with neural networks, and so this would require being able to analytically determine generalisation bounds. We feel this would be a very valuable future research direction.\n\nIn order to better empirically assess error accumulation of our method for long-term extrapolation, we provide experimental results for unrolling of Navier-Stokes for 80 i.o. 50 timesteps (see our response to R. zC2e for details). We apply our model both in a setting with fully observed initial conditions and sparsely observed (50%) initial conditions, and provide plots of accumulated MSE along these trajectories in Figs 1, 2, 3. We compare against DINo and FNO, and show that we consistently improve over DINo in terms of long-term extrapolation. FNO achieves better extrapolation limits due to reduced error accumulation in the fully observed setting, but very rapidly deteriorates even within the train horizon with sparsely observed initial conditions, whereas our proposed approach loses very little in terms of extrapolation performance in this sparse setting.\n\nAlthough we can not provide guarantees on extrapolation limits, results in all experiments show that train and test extrapolation MSEs are relatively consistent. As such, a possible way to obtain empirical extrapolation limits for a new dataset would be to analyse train extrapolation MSE and determine empirical soft bounds for test extrapolation limits by assessing the window in which train extrapolation MSE remains bounded. \n\n**Memory and execution times** We understand that information on the execution / inference times and memory complexity would allow for a better comparison of our approach with baselines. We provide this information in Table 1, 2, which will be included in the appendix of our manuscript! Notably, although our framework has increased GPU memory consumption during training and inference (compared to DINo), we show a marked reduction in inference time attributable to the use of meta-learning to obtain the latents for the initial state.\n\n**What are the typical DOF associated with the ODE solve?** The neural ODE solver operates on the latent space of the ENF, as such, the number of DOF of this ODE is equal to the total size of a single set of latents. We vary this as a hyperparameter over the different datasets:\n- Navier-Stokes: 4 latents, each with a 2-D pose $p_i$ and a 16-D context vector $\\mathbf{c}_i$, 72 DOF.\n- Planar diffusion: 4 latents, each with a 2-D pose $p_i$ and a 16-D context vector $\\mathbf{c}_i$, 72 DOF.\n- Spherical diffusion: 18 latents, each with a 2-D pose $p_i$ and a 4-D context vector $\\mathbf{c}_i$, 108 DOF.\n- Shallow-water: 8 latents, each with a 2-D pose $p_i$ and a 32-D context vector $\\mathbf{c}_i$, 272 DOF.\n- Internally heated convection: 25 latents, each with a 3-D pose $p_i$ and a 32-D context vector $\\mathbf{c}_i$, 875 DOF.\n\nWe typically obtain these settings through running a small hyperparameter search over a range of different settings. We will add these numbers to the appendix for the camera-ready version.\n\n**What is the boundary condition applied in the heat equation example** We use Dirichlet boundary conditions with value 0.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "sAUnps5IGx", "reviewer": "Reviewer_hu7A", "summary": "This work proposes a space-time continuous method for solving PDEs that respects the inherent symmetries of the PDE via equivariance constraints. Building upon prior work which (a) fits a conditional neural field to output latent vectors and (b) evolves the latent state through time via a Neural ODE, the contribution of this work is to additionally enforce equivariance constraints in the latent space itself. Secondly, the work employs meta-learning to obtain the initial latent representation, which improves the structure of the latent space representation and accelerates the inference time. The authors show improved performance of the method for linear and nonlinear PDEs on complex geometries such as the 2d torus, 3d sphere and 3d ball.", "questions": "- In figure 2, why is the bottom right image not the rotated version of the solution field on the top right?\n - Complex geometries can be a strong use-case for symmetry-preserving PDEs. However, complex geometries can often have non-trivial boundary conditions as well. Is there any way the method can be extended to handle non-equivariant boundary conditions?\n - Due to the global attention mechanism in the ENF there could be scalability concerns. Could you comment on how the method could be applied to larger scale problems in this case?", "limitations": "N/A", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The proposed method significantly reduces overfitting compared to non-equivariant baselines. \n - The method shows good stability at 3–5x the length of the training regime (even though the error accumulates slowly).", "weaknesses": "- The computational cost of the method v/s baselines is not shown. Relatedly, do the DINo baselines have similar parameter counts as your method?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their assessment of our work, and appreciate their recognition of the benefits of incorporating the symmetry constraints in PDE solving. We elaborate on the raised concerns below.\n\n**Computational cost**. We provide comparison of the computational cost of our method to all other baselines - as well as their parameter counts in tables 1 and 2 of the provided PDF. When looking at the results there can be observed that the training time per epoch is low in comparison to the best performing models (transolver and GFNO), although a bit higher than the less performing models (FNO and DINO). However, we want to stress that looking at inference time, our method is the most time-efficient, mostly attributable to the fact that we’re using a combination of meta-learning with latent-space ODE solving. For new trajectories, initial states only need to be fitted for 3 SGD steps compared to DINO’s 300-500 SGD steps of plane auto-decoding, which provides a clear speed-up during inference. \n\n**Application to non-equivariant boundary conditions** The reviewer asks if/how our method could be extended to non-equivariant boundary conditions. As R. nkRS, i6aX ask the same question, we provide a joint answer in our response to R. nkRS.\n\n **Scalability and model complexity** The reviewer raises valid concerns regarding the computational cost and scalability of our approach, especially in the context of solving complex/ high-dimensional PDEs. To better assess the computational cost of our approach and contrast it with existing methods, we added information about parameter count, runtimes and GPU memory usage during training and inference (see Tab. 2 of the PDF) for the experiments on Navier-Stokes. These results show that our method excels in inference compared to the baseline methods, although it is not as memory efficient as  FNO[1] or DINo [2]. Although we feel it is outside the scope of our current work, we do agree with the reviewer that aiming for improved scalability should be a focus of future research into this approach, as this would enable application to higher-resolution / higher-dimensional problem settings. We outline an idea for future research into ENF-based PDE solving here, which revolves around a more efficient implementation of the cross-attention operation.\n\nMuch of the computational cost of our method is attributable to the calculation of the query and key and values $q_{i,j},k_i, v_{i,j} \\in \\mathbb{R}^\\text{hidden}$ between a point $x_j$ and each of the latents $z_i$ in a latent state $Z$. To sample solution values with a set of $16$ latents for a $64x64$ output grid - for example - this requires the calculation of query and key values for each of $16 \\times 64^2$ combinations of latents and points. Note however that, within the ENF architecture’s main operation (described in the manuscript in appx C. Eq 10.) we enforce a latent $ ( p_i, \\mathbf{a_{i}} ) $ to be local in the domain of the PDE by weighting the attention value calculated between a latent $i$ and a point $x_j$ through applying a Gaussian window weighted by the relative distance between this point $x_j$ and the latents pose $p_i$ on the domain of the PDE. For latents far removed from a sampled point $x_j$, this Gaussian window forces the attention coefficient from $j$ to $i$ to be negligibly small, in turn nullifying contributions for such latents $i$ to the output for $f_\\theta(x_j, Z)$. For an approximation of Eq 10, it could thus be economical to forego calculating these attention coefficients altogether, by first applying an efficient knn algorithm to sort the relative distances from each coordinate $x_j$ to each latent $z_i$, and only calculating query key and value vectors for the nearest $N$ latents. Preliminary results on image regression tasks show that performance loss is negligible, and this method allows for sampling from much larger numbers of latents, making it an interesting possible future direction for scaling the proposed framework to much larger and more complex PDEs.\n\n**In figure 2, why is the bottom right image not the rotated version of the solution field on the top right?** Indeed the function on the bottom right is a rotated version of the solution, but for this particular solution that is quite hard to tell. We will change this figure to be more easily legible in the camera ready version.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "4rwtcsEs03", "reviewer": "Reviewer_zC2e", "summary": "The paper presents a novel framework for solving Partial Differential Equations (PDEs) by leveraging the power of Equivariant Neural Fields (ENFs). The authors propose a space-time continuous approach utilizing the symmetry of the PDEs, which is crucial for improving generalization and data-efficiency. The framework is tested on various geometries and PDEs, showing its effectiveness in handling complex dynamics.", "questions": "Using ODESolver often incurs higher computational costs. Did the authors implement any acceleration methods to speed up the integration process?", "limitations": "1. More experiments should be conducted to prove the model's efficiency in longer timespan.\n2. More baselines should be listed in the paper so that model's efficiency can be thoroughly tested and confirmed.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. **Data efficiency**: By designing a system that preserves the symmetries of PDEs, the proposed framework enhances the model's ability to generalize from limited data.\n2. **Novel initialization method**: The use of meta-learning to structure the latent space of the ENF simplifies the learning process and leads to better performance than autodecoding.", "weaknesses": "1. **Error Accumulation:** The usage of ODESolver might pose a challenge with error accumulation over time, particularly for dynamics occurring beyond the training horizon, which could affect the model's long-term predictive accuracy. So it would be helpful if the model is tested in a longer timespan.\n2. **Lack of Comparative Analysis:** While the paper compares its approach to a baseline method, a more comprehensive comparison with existing state-of-the-art methods in PDE solving would strengthen the paper's claims, such as Geo-FNO[1], GNOT[2], Transolver[3].\n\n[1] Li, Z., Huang, D. Z., Liu, B., & Anandkumar, A. (2023). Fourier neural operator with learned deformations for pdes on general geometries. *Journal of Machine Learning Research*, *24*(388), 1-26.\n\n[2] Hao, Z., Wang, Z., Su, H., Ying, C., Dong, Y., Liu, S., ... & Zhu, J. (2023, July). Gnot: A general neural operator transformer for operator learning. In *International Conference on Machine Learning* (pp. 12556-12569). PMLR.\n\n[3] Wu, H., Luo, H., Wang, H., Wang, J., & Long, M. (2024). Transolver: A fast transformer solver for pdes on general geometries. *arXiv preprint arXiv:2402.02366*.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their thorough assessment, and for the valuable criticism of our work. We address each of the reviewers' concerns in detail here, and hope to continue the discussion if any details remain unclear.\n\n**Error accumulation** The reviewer is concerned about error accumulation in long-term rollouts as a result of requiring a neural ODE forward propagation. We agree, and as such provide an additional experiment with a more in-depth analysis of long-term rollouts. We extend the dataset of Navier-Stokes solutions we use in the paper to 80 timesteps, and provide test-set results in Figs 1, 2, 3 of the PDF. In Fig. 1 we provide average test MSE per timestep for our model in comparison with FNO and DINo on the setting of having a fully observed initial condition $\\nu_0$. Here, we can see that FNO clearly outperforms both NeF-based neural solvers in terms of long-term error accumulation - in agreement with the results shown in our paper in Tab. 1 - although our proposed method retains relatively low MSE up to 20-30 steps after the train horizon (marked in green), where DINo starts error accumulation inside the training time horizon. In Fig. 2 we show results where the initial condition $\\nu_0$ for each test trajectory is only sparsely observed, i.e. where only 50% of the initial state is given as input to the model. Here, FNO quickly deteriorates in performance and accumulates error at a higher rate than the NeF-based framework we propose in this work. To us these results indicate that (1) our equivariant NeF-based continuous space-time solving framework provides better generalisation and less error accumulation than its non-equivariant counterpart and (2) retains this relatively limited rate of long-term error accumulation in settings with sparsely observed initial conditions.\n\n**Improving comparative analysis** The reviewer indicates that they would like to see experimental comparison with a wider range of baseline methods. We provide results for a baseline proposed by the reviewer; Transolver, which has shown very promising results for PDE modelling on general geometries.\n\nWe took the Transolver model in PyTorch from their github repository, and use the same hyperparameter settings (resulting in a 7.1M parameter model) - modifying the training objective to be identical to the autoregressive one we use in our experiments (i.e. mapping from frame to frame i.o. taking 10 frames as input). Tab. 5 shows the results for the Transolver model on Navier-Stokes in 2D after training on a single A100 to convergence over 700 epochs in approximately 49 hours. Notably, Transolver achieves 1.80E-02 and 1.85E-02 train and test mse respectively within the training horizon, and 4.85E-01 and 4.90E-01 train and test mse outside the training horizon. During training in the Navier-Stokes task we observed noticeable instabilities, with the method producing high-frequency artefacts in rollouts in this autoregressive setting. For sparsely subsampled initial conditions, performance deteriorates further - highlighting the benefit of NeF-based continuous PDE solving.\n\nFor experiments on internally-heated convection, due to the large size of the input frames, we were required to scale down the Transolver model size in order to be able to fit the model on our A100 GPU, from 256 to 64 hidden units, resulting in a 1.2M parameter model - somewhat comparable in size to the model we use in our experiments (889K). We train the model for 2000 epochs on an A100 GPU, taking approximately 30 hours. Results in Tab. 4 show the strong performance of this model in the training time horizon - achieving 4.13E-04 test MSE where our framework achieves 5.99E-04, indicating that it is indeed a strong baseline for solving PDEs over complicated geometries. However, we also note that outside of the training horizon, error accumulates more quickly for the Transolver model, indicating it has somewhat overfit the training horizon dynamics. Here, Transolver achieves 2.09E-02 test MSE, where our framework achieves 8.21E-03. We hypothesise that - due to the equivariance constraints placed on our model reflecting inductive biases about this PDE, it extrapolates more reliably beyond the train horizon and over the validation set.\n\nFurthermore, we provide results in Tab. 4 for sparsely observed initial conditions. These results clearly show the advantage of NeF-based continuous solvers in this sparsely observed setting, both DINo and our framework significantly outperform Transolver, which is unable to provide accurate solutions either within the training horizon or outside of it.\n\nIn the camera ready version, we will further include the results for the transolver baseline in our other experiments.\n\n**Computational cost of use of ODESolver** The reviewer raises concerns about computational efficiency of our method. We did not implement any specific acceleration methods for solving the forward ODE in the latent space. The main reason we did not feel we need to is because the neural ODE solver operates on a drastically compressed representation of the PDE state, i.e. in the Navier-Stokes experiment we operate on a set of 4 latents with a context vector living in $\\mathbb{R}^{16}$. Since the neural ODE solver operates on such a small representation, its overhead is relatively limited. In order to provide more insight into the computational requirements of our method we provide training and inference runtimes and memory usage of our method compared with the different baselines we show in our paper in Tab. 1, 2 of the PDF. We will add these details to the appendix of our manuscript.\n\nWe also provide a direction for how to significantly reduce the memory overhead of our decoder by approximating the integral listed in appx C. Eq. 10 in our response to R. hu7A. We will be exploring this in future work.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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:
We thank the reviewers for their thorough investigation of our work, and for investing the time to write out valuable criticism. We’re happy to see that reviewers regard our equivariant space-time continuous PDE solving method as valuable and effective for scientific research applications. Additionally, we appreciate the recognition of our efforts to address significant limitations of existing NeF-based solvers, and are happy to see that the reviewers found the work well-motivated and clearly explained.

Though the reviewers do not indicate any critical flaws, they raise - and agree on - a number of concerns, primarily aimed at (1) strengthening the motivation of our method for a wider family of PDEs and boundary conditions (Rev. NkrS, kxDM, hu7A) and (2) expanding the experimental validation of our method by including additional baselines and datasets (Rev. hu7A, zC2e, kxDM, i6aX). We address the most pressing of these concerns in this general response, referring to the attached single-page pdf, and provide additional detailed comments to each reviewer individually.

### **Questions regarding motivation/usability for non-symmetric PDEs / PDEs with unknown symmetries**

R. NkrS  and R. i6aX point out that in our analysis “boundary conditions are symmetric and the chosen PDEs exhibit certain symmetries. In real-world applications, these assumptions might not always hold, potentially limiting the applicability of the proposed method” and that “the invariance groups are chosen according to the equation at hand. But how does one choose the invariance groups when the underlying operators and functions (RHS) are unknown”. This is a valid point, however, we highlight that our method *can* be employed with relaxed or without equivariance constraints both through use of a non-symmetric “bi-invariant” (see appx. D $\mathbf{a}^\emptyset_{i,m}$) and through a procedure we outline in our response to R. NkrS under **Equivariance constraints in real-world applications**, which means that it would retain its utility in scenarios where the symmetries are not (fully) known or are deemed irrelevant. 

We argue that the attention-based conditional neural field trained with meta-learning can enhance forecasting accuracy due to its ability to capture local dependencies effectively, also in these settings. To underscore this, we provide additional experimental results on CFDBench (see response to R. kxDM), which doesn’t exhibit symmetric boundary conditions or global symmetries and contains varying geometries with hindrances and varying boundary conditions - achieving performance competitive to the baselines set by the authors.

On the other hand, in science and engineering it is not rare that known symmetries are used to develop models even when the complete details of the system are not fully known (e.g., in atmospheric and oceanic sciences rotational symmetry is often assumed). For these reasons, we posit that an equivariant model for spatio-temporally continuous PDE dynamics forecasting constitutes a valuable contribution, also in real-world settings.

### **Improving comparative analysis (R. hu7A, zC2e, kxDM, i6aX)**

**Comparison with additional baselines** Rev. zC2e and kxDM point out that on a number of experiments we compare with Yin. (2023), and argue for including additional baseline results for neural PDE solvers that can similarly handle the varying geometries we’re operating on, as this would strengthen the paper’s claims. To this end, we implemented an additional baseline proposed by the reviewers, which we detail below.

We adopt the Transolver model (2024), and train it on the Navier-Stokes and Internally-Heated Convection tasks. Find a detailed description of the experimental setup in our response to Rev. kxDM, results are listed in Tab 4, 5. The Transolver proved unstable in the autoregressive Navier-Stokes experiment, obtaining comparable performance to DINo, due to high-frequency artefacts arising on rollout - and so it underperforms FNO and our method. Moreover, it deteriorates heavily with sparsely observed initial conditions. Transolver does achieve competitive error in the internally-heated convection experiment in the fully observed setting, i.e. when 100% of the initial state is observed, indicating that it is a strong baseline in settings with complicated geometry and boundary conditions. However, the Transolver model deteriorates considerably in performance with subsampled initial conditions in both experiments - indicating that indeed it is not a model suitable for space-time continuous solving on complex geometries. In these settings, the proposed model architecture outperforms Transolver. We feel these results further strengthen the case for our proposed model architecture, and we thank the reviewers for this suggestion. In the camera ready version, we will further include results for this baseline in our other experiments.

**Computational cost / parameter count** (R. i6aX, hu7A) We provide the parameter counts of our model versus each of the baselines, as well as their respective runtime and GPU memory usage during training and inference in tables 1 and 2 of the pdf. We note that although our method is not the most memory efficient - meta-learning requires a significant computational overhead - it is the fastest in inference settings when unrolling an unseen state. This is attributable to (1) the fact that we’re using a latent-space solver which operates on a drastically compressed representation of the PDE state, and use (2) meta-learning over auto-decoding, which means we only require 3 gradient descent steps to obtain the initial state compared to 300-500 typically used in auto-decoding (Yin 2023).

This information, as well as all extra provided experiments, will be included in the appendix of the manuscript. We again thank the reviewers for their diligence, and hope to discuss further if any questions or concerns remain.

Meta-review: The paper introduces an encode-process-decode architecture for modeling physical spatio-temporal dynamics using implicit neural representations (INRs), also known as neural fields, that respect the symmetries of physical problems. The primary focus and novelty of the paper lie in the preservation of these symmetry properties within the problem's geometry. This is achieved by incorporating equivariance constraints at the level of the latent representation via equivariant neural fields and within the "process" component. Specifically, the neural ODE solver, implemented using a graph-based message-passing architecture, is constrained to be equivariant to symmetries. The experiments demonstrate the advantages of these constraints across diverse benchmarks.

Building on previous INR solvers, the paper introduces symmetry constraints both in the signal representation and in the solver. Its main contribution is the introduction of equivariance properties within this framework. In particular, the authors adapt the design of equivariant neural fields to the context of partial differential equations (PDEs). The results highlight the importance of such prior physics-informed characteristics for generalizing in time and space for physical dynamical processes. Following the reviewers' comments, the authors included additional results with baselines and datasets during the rebuttal.

Although the paper presents a somewhat incremental contribution, the reviewers regard it as a solid and significant advancement that improves the state of the art. It demonstrates the importance of imposing physics constraints to help neural physics models generalize, particularly from limited data and across diverse geometries. The adaptation of equivariant INR to this domain is novel and contributes to the development of the emerging trend of INR-based PDE solvers.
Decision: Accept (poster)