AROMA: Preserving Spatial Structure for Latent PDE Modeling with Local Neural Fields
Authors: Louis Serrano
Thomas X Wang
Etienne Le Naour
Jean-Noël Vittaut
Patrick Gallinari
Venue: NeurIPS 2024 poster (2024)
Keywords: PDE
Neural Operator
Neural Fields
Transformer
Diffusion
Primary area: machine_learning_for_physical_sciences
TL;DR: We introduce AROMA, a versatile framework for enhanced PDE modeling using local neural fields, which achieves stable, efficient processing of diverse spatial data and superior performance in simulating 1D and 2D equations.

Abstract: We present AROMA (Attentive Reduced Order Model with Attention), a framework designed to enhance the modeling of partial differential equations (PDEs) using local neural fields. Our flexible encoder-decoder architecture can obtain smooth latent representations of spatial physical fields from a variety of data types, including irregular-grid inputs and point clouds. This versatility eliminates the need for patching and allows efficient processing of diverse geometries. The sequential nature of our latent representation can be interpreted spatially and permits the use of a conditional transformer for modeling the temporal dynamics of PDEs. By employing a diffusion-based formulation, we achieve greater stability and enable longer rollouts compared to conventional MSE training. AROMA's superior performance in simulating 1D and 2D equations underscores the efficacy of our approach in capturing complex dynamical behaviors.

AROMA:
Preserving Spatial Structure for Latent PDE Modeling with Local Neural Fields
Louis Serranol *
Thomas X Wang'
Etienne Le Naour 1.2
Jean-Noël Vittaut 3
Patrick Gallinari
Sorbonne Université, CNRS, ISIR, 75005 Paris, France
EDF R&D, Palaiseau, France Sorbonne Université, CNRS, LIP6, 75005 Paris, France "Criteo AI Lab, Paris, France
Abstract
We present AROMA (Attentive Reduced Order Model with Attention), frame- work designed enhance the modeling of partial differential equations (PDEs) using local neural fields. Our flexible encoder-decoder architecture smooth latent representations spatial physical fields from a variety of data types, including irregular-giid inputs and point clouds. This versatility eliminates the need for patching and allows efficient processing of diverse geometries. The se- quential nature of our latent representation be interpreted spatially and per- mits the use of a conditional transformer for modeling the temporal dynamics achieve greater stability
obtain
PDEs. By employing diffusion-based formulation, and enable longer rollouts compared to conventional MSE training. AROMA'S superior performance in simulating 1D and 2D equations underscores the effi- cacy of our approach in capturing complex dynamical behaviors. Github page: https: htt///gtttuu wow./oooiseerrrae..... arrma
Introduction
In recent years, many deep learning (DL) surrogate models have been introduced to approximate solutions to partial differential equations (PDEs) (Lu al., 2021; Li et al., 2021; Brandstetter al., 2022; Stachenfeld et al., 2022). Among these, the family neural operators has been extensively adopted and tested across various scientific domains, demonstrating the potential of data-centric DI. models in science (Pathak et al., 2022; Vinuesa & Brunton, 2022).
Neural Operators were initially constrained by discretization and domain geometry limitations. Re- cent advancements, such as neural fields (Yin et al., 2022; Serrano et al., 2023) and transformer architectures (Li et al., 2023; Hao et al., 2023), have partially addressed these issues, improving both dynamic modeling and steady-state settings. However, Neural Fields struggle to model spatial information and local dynamics effectively, and existing transformer architectures, while being flex- ible, are computationally expensive due to their operation the original physical space and require large training datasets.
Our hypothesis that considering spatiality is essential modeling spatio-temporal phenomena, yet applying attention mechanisms directly is computationally expensive. We propose a new framework that models the dynamics in a reduced latent space, encoding spatial information compactly, by one two orders of magnitude relative to the original space. This approach addresses both the complexity issues of transformer architectures and the spatiality challenges of Neural Fields.
"Corresponding author: louis.serrange @ @isir.upmc.fr
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Our novel framework leverages attention blocks and neural fields, resulting in model that is easy train and achieves state-of-the-are results on most datasets, particularly for complex geometries, without requiring prior feature engineering. 'oo the best of our knowledge, we are the first to pro- pose a fully attention-based architecture for processing domain geometries and unrolling dynamics. Compared to existing transformer architectures for PDEs, our framework first encapsulates the do- main geometry and observation values in a compact latent representation, efficiently forecasting the dynamics at lower computational cost. Transformer-basee methods such (Li et al., 2023; Hao al., 2023) unroll the dynamics in the original space, leading to high complexity.
Our contributions are summarized as follows:
A principled and versatile encode incode-process. process- decode framework for solving PDEs that operate general imput geometries, including point sets, grids, or meshes, and can be queried at any location within the spatial domain.
A new spatial encode process / decode approach: Variable-size inputs are mapped onto a fixed-size compact latent token space that encodes local spatial information. This latent representation is further processed by a transformer architecture that models the dynamics while exploiting spatial relations both at the local token level and globally across tokens. The decoding exploits conditional neural field, allowing us to query forecast values at any point in the spatial domain of the equation.
We include stochastic components at the encoding and processing levels to enhance stabil- ity and forecasting accuracy.
Experiments performed on representative spatio-temporal forecasting problems demon- strate that AROMA on par with or outperforms state-of-the-ar baselines in terms of both accuracy and complexity.
2 Problem setting
In this paper, we focus on time-dependent PDEs defined over a spatial domain Q (with boundary 0l2 and temporal domain [0, 7'. In the general form, their solutions u(x, t) satisfy the following constraints
du F 1, u, du 8au ):): Or 0.
Vo E Q 2. (0,
(1)
Øt
B(u)(L, = 0 Ya E an, 14 E (0,7]
(2)
u(0,z) z) = u Vz E Ω
(3)
where represents a set of PDE coefficients, Equations (2) and (3) represent the constraints with respect to the boundary and initial conditions. We aim to learn, using solutions data obtained with classical solvers, the evolution operator G that predicts the state of the system the next time step: u + = ((u'). We have access to training trajectories obtained with different initial conditions, and want to generate accurate trajectory rollouts for new initial conditions at test time. A rollout is obtained by the iterative application of the evolution operator u mAr = gm (u²).
Model Description
3.1
Model overview
We provide below an overview of the global framework and each component is described in a sub- sequent section. The model comprises three key components, detailed in Figure 1.
Encoder Ey u + Z'. The encoder takes input values u'x sampled over the domain Ω at time 1, where X' denotes the discrete sample space and could be a grid, an irregular mesh or a point set. u'x is observed at locations æ (#1, I..,,), with values u (u' (a)))),,.. uu( ))) N is the number observations and can vary across samples. u is projected through a cross attention
mechanism onto set of M tokens z (z),... z/1)) with M fixed parameter. This allows

---

coordinales 2 E R- lunction values v'é gx
predicted value
Astent fokens with compressed dimension
predicted tokens
K
T1
T3
Encode geometry
Aggregate values
Diffusion Transformer
T3
Latent Refiner
TM
Decoder
Encoder
T E Mxd
T E Mxd
Ra
x M
geometry-aware lokens
coordinate query
Figure I: AROMA inference:
The discretization-Fe encoder compresses the information of set of N input values sequence of M latent tokens, where M < N. The conditional diffusion transformer is used to model the dynamics, acting as a latent refiner. The continuous decoder lever- ages self-attentions (SA), cross- coss-attentionn (CA) and local INR to map back to the physical space. Learnable tokens are shared and encode spatial relations. Latent token Z' represents u and + is the prediction corresponding U++Ai-
mapping any discretized input u onto a fixed dimensional latent representation encoding implicit local spatial information from the input domain. The encoder is trained VAE and Z is sampled from multivariate normal statistics as detailed in Section 3.2.
Latent time- time-marching refiner Rg Z 2+b We model the dynamics in the latent space latent space for any
through a transformer. The dynamics can be unrolled auto-regressively time horizon without requiring project back in the original domain Ω. Self-attention operates the latent tokens, which allows modeling global spatial relations between the local token represen- tations. The transformer is enriched with a conditional diffusion mechanism operating between two successive time steps of the transformer. We experimentally observed that this probabilistic model more robust than baseline deterministic transformer for temporal extrapolation.
Decoder D Z++Ar + A The decoder uses the latent tokens t++ to approximate the function value ((+A((a) (x) (xc z(+2) for any query coordinate & E Ω. We therefore denote a1+Ar = D (Z++t)) the predicted function.
Inference We encode the initial condition and unroll the dynamics in the latent space by succes- sive denoisings: mAr DØ o RT (u"). We then decode along the trajectory to get the reconstructions. We outline the full inference pipeline in Figure 1 and detail its complexity analysis in Appendix C.1.
Training We perform two-stage training: first train the encoder and decoder, secondly train the refiner. This is more stable than end-to-end training.
3.2 Encoder-decoder description
The encoder-decoder components are jointly trained using a VAE setting. The encoder is specifically designed to capture local input observation from any sampled point set in the spatial domain and encodes this information into fixed number of tokens. The decoder can be queried any position the spatial domain, irrespective of the input sample.
Encoder The encoder maps an arbitrary number N of observations (z, u(x)) ((x), ((11,u(2))) u(x1))............ (IN, u(xx)) onto a latent representation of fixed size M through the following series of transformations:

---

(i)
(æ, u(x))
(positional, value) embeddings
(((x), E B
(ii)
(T, (x))
geometry encoding
To E R Mxd
(iii)
TT v(x))
observation spatial encoding
Tobs E R Mxd
(iv) Tobs
dimension reduction
Z E RMxx
where v(z)),((x) = (1v(211, ((7(=)) (M(2)))))...,))))))))))))))))))) (()..))) (1(x)),((E))))
and h < d.
(i) Embed positions and observations: Given input sequence of coordinate-value pairs (I1, u(x)))....... (xN, u(Ix)), we construct sequences of positional embeddings 7 (r( (xN)) and value embeddings FourierFeatures (z; w) and u(z) = (nnerru(x)), with w a fixed set of frequencies. These embeddings are aggregated onto smaller set of learnable query tokens T (Ti..... TM) and then T' .T",..,,, T with M fixed, to compress the information and encode the geometry and spatial latent representations.
(v(z1),.... v(#N)), where >(x)
(ii) Encode geometry: Geometry-aware tokens T are obtained with multihead cross-attention layer and a feedforward network (FFN), expressed as T T + (Cr i ion (Q W = WX,, V = Wr))). This step does not include information on the observations, ensuring that similar geometries yield similar query tokens T irrespective of the values.
(iii) Encode observations: The T.0 tokens are then used to aggregate the observation values via a cross-attention mechanism: T = To + FFN(CrossAt tetttteentt tentiinn
Te0 W T K = Wk., V = w)). Here, the values contain information on the observation values, and the keys contain information on the observation locations.
(iv) Reduce channel dimension and sample Z: The information in the channel dimension of T is compressed using a bottleneck linear layer. To avoid exploding variance in this compressed latent space, we regularize it with penalty on the L2 norm of the latent code Z Introducing stochasticity through a variational formulation further helps to regularize the auto-encoding and obtain smoother representations for the forecasting step. For this, we learn the components of Gaussian multivariate distribution μ = Linear(T Tobs)) and log(or) Linear (Tohs) from which the final token embedding Z is sampled.
Decoder The decoder's role is reconstruct ++AA from ZI+ see Figure 1. Since training is performed in two steps ("encode-decode" first and then "process"), the decoder is trained to reconstruct û' for input u'. One proceeds as follows. (i) Increase channel dimensions and apply self-attention: The decoder first lifts the latent tokens Z to a higher channel dimension (this is the reverse operation of the one performed by the encoder) and then apply several layers of self-attention
to get tokens Z (ii) Cross-attend: The decoder applies cross-attention to obtain feature vectors CrossAttent in We(ng (x)), K
that depend on the query coordinate J, (f:)
W Z V = Wy Z'), where "y is a Fourier features embedding of bandwidth wig (iiil Decode small MLP to decode this feature vector and obtain the reconstruction (x) = MIP((f (z)). In contrast with existing neural field methods for dynamics modeling, the feature vector here is local. In practice, uses multiple cross attentions get feature vectors with different frequencies (see Appendix Figures and for further details).
with MLP: Finally,
Training The encoder and decoder are jointly optimized as a variational autoencoder (VAE) (Kingma & Welling, 2013) to minimize the following objective L Lreson + 3 LKLA where Lreson = MSE(u, atx) is the reconstruction loss between the input and the reconstruction De [Z', X)) on the grid X, with Z' ~ ~N("', loss LRI. DRi (N'" (I helps regularize the network and prevents overfitting. We found that using a variational formulation was essential obtain smooth latent representations while training the encoder-decoder.
and n²,o' (u'x). The KL divergence

---

T
T's
T goo
Ty
Figure 2: Spatial interpretation of the tokens through cross attention between T300 and y(x) for each x the domain. Here we visualize the cross-attention three different tokens for given head. The cross attentions can have varying receptive fields depending the geometries.
3.3 Transformer-based diffusion
Modeling the dynamies is performed in the latent Z space. This space encodes spatial information present in the original space while being a condensed, smaller-sized representation, allowing for re- duced complexity dynamics modeling. As indicated, the dynamics can be unrolled auto-regressively in this space for any time horizon without the need to map back to the original space. We use ab- solute positional embeddings Eps and a linear layer to project onto a higher dimensional space: Z Linear(Z) Epos- The backbone then applies several oeff-attention blocks, which process
tokens as follows:
Z[+1 Zμ Artemttion(cayeeNoom(222
Z[+1 Z[+1j FFN(LayerNorm (Z(+1))
(5)
We found out that adding diffusion component to the transformer helped enhance the stability and allowed longer forecasts. Diffusion steps are inserted between
time steps and ++ At of the time-marching process transformer. The diffusion steps are denoted by k and are different from the ones of the time-marching process (several diffusion steps It are performed between time- marching steps t and f At).
We then use a conditional diffusion transformer architecture close Peebles & Xie (2023) for Rø. where we detail the main block in Appendix B. At diffusion step k, the input to the network is sequence stacking the tokens at time t and the current noisy targets estimate (Z', ( See Appendix B. Figure 4 and Figure 5 for more details. To train the diffusion transformer RR, we freeze the encoder and decoder, and use the encoder to sample pairs of successive latent tokens (Z', Zt) We employ the "v-predict" formulation of DDPM (Salimans & Ho, 2022) for training and sampling.
Experiments
In this section, we systematically evaluate the performance of our proposed model across various experimental settings, focusing on its ability to handle dynamics on both regular and irregular grids. First, we investigate the dynamics regular grids, where we benchmark of-the-art neural operators, including Fourier Neural Operators (FNO), ResNet, Neural Fields, and Transformers. This comparison highlights the efficacy of our approach in capturing complex spatio- temporal patterns on structured domains. Second, we extend our analysis to dynamics on irregular grids and shared geometries, emphasizing the model's extrapolation capabilities in data-constrained regimes. Here, we compare our results with Neural Fields and Transformers, demonstrating the robustness of our model in handling less structured and more complex spatial configurations. Lastly, we assess the model's 's capacity to process diverse geometries and underlying spatial representations by comparing its performance on irregular grids and different geometries. This evaluation highlights the flexibility and generalization ability of our model in encoding and learning from varied spatial domains, showcasing its potential in accurately representing and predicting dynamics across wide range of geometric settings. We include additional results from ablation studies in Appendix C.6.
model against state-

---

4.1 Dynamics on regular grids
We begin our analysis with dynamics modeling on regular grid settings. Though our model is targeted for complex geometries, we believe this scenario remains an important benchmark to assess the efficiency of surrogate models.
Datasets 1D Burgers' Equation (Burgers): Models shock waves, using a dataset with periodic
initial conditions and forcing term in Brandstetter al. (2022). It includes 2048 training and 128 test trajectories, at resolutions of (250, 100). We create sub-trajectories of 50 timestamps and treat them independently. 2D Navier Stokes Equation: for a viscous and incompressible fluid. We use the data from Li et al. (2021). The equation is expressed with the vorticity form on the unit torus: the + u Vw = ~Aw + S. Vu = 0 for x E 22 1 > 0, where v is the viscosity coefficient. We x 10-4) and v = 10-5 (Navier-Stokes
ot
consider two different versions V = 10-4 (Navier-Stokee of 1000 and 200 trajectories with base spatial resolution of
1 x 10-9), 9), and use train and test size 64 x 64. We consider horizon of T = 30 for v = 10 4 and T 20 for v 10 since the phenomenon is more turbulent. At test time, we use the vorticity 4 10 the initial condition.
Setting We train all the models with supervision on the next state prediction to learn to ap- proximate the time-stepping operator u u'tår ((u'). At test time, unroll the dynamics auto-regressively with each model and evaluate the prediction with a relative L2 error defined as |a macoor 1okenee l|2
Lisat
New Etetest
Baselines We use diverse panel of baselines including state of the art regular-grid methods such as FNO (Li al., 2021) and ResNet (He et al., 2016; Lippe et al., 2023), flexible transformer architectures such as OFormer (Li et al., 2023), and GNOT (Hao et al., 2023), and finally neural- field based methods with DINO (Yin et al., 2022) and CORAL (Serrano et al., 2023).
Results Table 1 presents comparison of model and
performance on the Burgers, Navier-Stokesle-5 datasets, with metrics reported in
Table 1: Model Performance Comparison Test results. Metrics in Relative Lz
Relative 22- Our method, AROMA, demonstrates excellent performance across the board, highlighting its ability to capture the dynamics of turbulent phe- nomena, reflected in the Navier-Stokes datasets.
Model
Burgers
Navier-Stokes
Navier-Stokes
1 10
1 10
FNO ResNet
5.00 x 10
153 x 10
1.24 x 10
x 10
3.77
2.56 10
In contrast, DINO and CORAL, both global neu- ral field models, perform poorly in capturing turbu- lent phenomena, exhibiting significantly higher er- rors compared to other models. This indicates their limitations in handling complex fluid dynamics. On the other hand, AROMA outperforms GNOT on all datasets, though it performs reasonably well compared to the neural field based method.
DINO
4.57 x 10
7.25 x
3.72 10
CORAL
3.77
3.11 10
GNOT
10
1.85 10-1
1.65 10
OFormer
4.92 x 10-2
1.36 10-1
2.40 10-1
AROMA 3.65 x 10
1.05 10 1.24 10-
Regarding the regular-gris -rid methods, ResNet shows suboptimal performance in the pure teacher forc- ing setting, rapidly accumulating errors over time during inference. FNO stands out as the best base- line. demonstrating competitive performance on all datasets. We hypothesize that FNO's robustness error accumulation during the rollout can be attributed its Fourier block, which effectively cuts off -eeg-frryenncy components. Overall, the results underscore AROMA's effectiveness and high- light the challenges Neural Field-based models face in accurately modeling complex phenomena.
4.2 Dynamics irregular grids with shared geometries
We continue our experimental analysis with dynamics on unstructured grids, where we observe trajectories only through sparse spatial observations over time. We adopt data-constrained regime and show that our model can still be competitive with existing Neural Fields in this scenario.
Datasets To evaluate our framework, we utilize two fluid dynamics datasets commonly used as benchmark for this task (Yin et al., 2022; Serrano et al., 2023) with unique initial conditions for each trajectory: 2D Navier-Stokes Equation (Navier-Stokes x 10- 10-3): We use the same equation as

---

in Section 4.1 but with a higher viscosity coefficient v = le 3. We have 256 trajectories of size 40 for training and 32 for testing. We used a standard resolution of 64x64. 3D Shallow-Water Equation (Shallow-Water)
This equation approximates fluid flow on the Earth's surface. The data includes the vorticity w and height of the fluid. The training set comprises 64 trajectories of size 40, and the test set comprises 8 trajectories with 40 timestamps. We use standard spatial resolution of 64 x 128.
Setting Temporal Extrapolation: For both datasets, we split trajectories into two equal parts of 20 timestamps each. The first half is denoted as In-t and the second half as Out-t. The training set consists of In-t. During training, we supervise with the next state only. During testing, the model unrolls the dynamics from a new initial condition (IC) up to the end of Out-t, i.c. for 39 steps. Evaluation within the In- horizon assesses the model's ability forecast within the training regime. The Out-r evaluation tests the model's extrapolation capabilities beyond the training horizon. Sparse observations: For the train and test set we randomly select и percent the available regular mesh to create a unique grid for each trajectory, both in the train and in the test. The grid is kept fixed along a given trajectory. While each grid different, they maintain the same level sparsity across trajectories. In our case, x 100% amounts to the fully observable case, while in = 25% each grid contains around 1020 points for Navier-Stokes
10-3 and 2040 points for Shallow- Saallow-Water.
Baselines We compare our model to OFormer (Li al., 2023), GNOT (Hao et al., 2023), and choose DINO (Yin et al., 2022) and CORAL (Serrano et al., 2023) as the neural field baselines.
Training and evaluation During training, we only use the data from the training horizon (In-t). At test time, we evaluate the models to unroll the dynamics for new initial conditions in the training horizon (In-r) and for temporal extrapolation (Out-r).
Results Table demonstrates that AROMA consistently achieves low MSE across all levels of observation sparsity and evaluation horizons for both datasets. Overall, our method performs best with some exceptions. On Shallow- Watter our model is slightly outperformed by CORAL in the fully observed regime, potentially because of lack of data. Similarly, on Navier-Stokes 1 x 10-1 CORAL has slightly better scores in the very sparse regime म = 5%. Overall, this is not surprising as meta-learning models excel in data-constrained regimes. We believe our geometry-encodie block crucial for obtaining good representations of the observed values in the sparse regimes, potentially explaining the performance gap with GNOT and OFormer.
Table 2: Temporal Extrapolation Test results. Metrics in MSE.
datasct
Navier-Stokes x 10"3
Shallow- Water
No. J. Nim
In-1
Out-t
In-t
Oud-
DINO
2.51 x 10
9.91 X 10
4.15
10
3.55 x 10
π 100%
CORAL
5.76 x 10
3.00 x 10-3
2.12
10
6.00 x 10
OFormer
7.76 x 10
6.39 x 10 2
1.00 10
2.23 x 10
GNOT
3.21
10
2.33 10
2.48 10
2.17 x 10-3
AROMA
1.32
10
2.23 10
3.10 10
8.75
10
DINO
3.27 x 10
1.40 x 10
4.12 x 10
3.26 x 10
25%
CORAL
1.54 x 10
1.07 x 10-2
3.77 10-4
1.44 x 10
irregular grid
OFormer
3.73 x 10
1.60 x 10
6.19 10
1.40 x 10
GNOT
2.07 x 10
6.24 x 10
8.91 10
4.66 x 10-3
AROMA
7.02 10
6.31 10
1.49 10
1.02 10
DINO
3.63 x 10
1.35 x 10
4.47 x 10
9.88 x 10 10
= 5%
CORAL
2.87 x 10
1.48 x 10
2.72 x 10
6.58 x 10
irregular grid
OFormer
3.23 x 10
1.12 x 10
8.67 10
1.72 x 10
GNOT
7.43 x 10
1.89 x 10
5.05 10
1.49 10
AROMA
4.73 10
2.01 10
1.93 10
3.14 10
4.3 Dynamics on different geometries
Finally, we extend our analysis to learning dynamics over varying geometries.
7

---

Datasets We evaluate our model on two problems involving non-convex domains, described by Pfaff et al. (2021). Both scenarios involve fluid dynamics in a domain with an obstacle, where the area near the boundary conditions (BC) is more finely discretized. The boundary conditions are specified by the mesh, and the models are trained with various obstacles and tested on differ- ent, yet similar, obstacles. Cylinder (CylinderFlow). This dataset simulates water flow around a cylinder using a fixed 2D Eulerian mesh, representing incompressible fluids. For each node j in the mesh X, we have data on the node position x(3), momentum w(x("), and pressure p(æ ()) Our task is to learn the mapping from (w)( (x)) to (W++>x(z), Pe+ae ((++)(()))ee) (x)) for fixed AL. Airfoil (Airfoilllow): This dataset simulates the aerodynamics around an airfoil, relevant for compressible fluids. In addition to the data available in the Cylinder dataset, we also have the fluid density p(x(0) for each node j. Our goal is to learn the mapping from (wr (z), Pt (2)),, Pe(x))>E>'' to unique mesh.
(e)+se()) pt+ (x)), Pt+ Pt+ (x)) Each example the dataset corresponds On average, there are 5233 nodes per mesh for AirfoilFlow and 1885 for CylinderFlow. We tempo- rally subsample the original trajectories by taking one timestamp out of 10, forming trajectories of 60 timestamps. We use the first 40 timestamps for training (In-f) and keep the last 20 timestamps for evaluation (Out-I).
Setting We train all the models with supervision on the next state prediction. At test time, we unroll the dynamics auto-regressively with each model and evaluate the prediction with mean squared error (MSE) both in the training horizon (In-t) and beyond the training horizon (Out-t).
The results in Table 3 show that AROMA outperforms other models in predicting flow dy- namics on both CylinderFlow and AirfoilFlow geometries, achieving the lowest MSE values across all tests. This indicates AROMA's superior ability to encode geometric features accurately. Addi- tionally, AROMA maintains stability over extended prediction horizons, as evidenced by its consis- tently low Out-t MSE values.
Results
Table 3: Dynamics on different geometries Test results. MSE on normalized data.
Model
CylinderFlow
AirfoilFlow
I-tt
Out-t
In-t
Out-f
CORAL
4.458 10-2
8.695 10-2
1.690 x 10-1
3.420 10
DINO
1.349 x 10
1.576 x 10
3.770 x 10-T
4.740 10
OFormer
5.020 x 10-1
1.080 10"
5.620 x 10
7.620 x 10
AROMA
1.480 10
2.780 x 10
5.720 x 10
1.940 10
4.4 Long rollouts and uncertainty quantification
After training different models Burgers, we compare them on long trajectory rollouts. We start from to = 50 (i.e. use numerical solver for 50 steps), and unroll our dynamics auto-regressively for 200 steps. Note that all the models were only
Correlation
trained to predict the next state. We plot the cor-
0.8
relation over rollout steps of different methods, in-
cluding our model without the diffusion process,
in Figure 3. We can clearly the gain in stabil- ity in using the diffusion for long rollouts. Still,
8.6
FND
ResNet
the predictions will eventually become uncorre- lated over time as the solver accumulates errors compared with the numerical solution. As em- ploy a generative model, we can generate several rollouts and estimate the uncertainty of the solver with standard deviations. We can see in Appendix Figure that this uncertainty increases over time. This uncertainty is not a guarantee that the solution lies within the bounds, but is an indication that the model is not confident in its predictions.
8.
AROMA
AROMA No Difusion
25
50
120
125
150
175
200
Rollout Steps
Figure 3: Correlation over time for long roll- outs with different methods on Burgers
8

---

5 Related Work
Our model differs from existing models in the field of operator learning and more broadly from existing neural field architectures. The works most related to ours are the following.
Neural Fields for PDE Neural Fields have recently emerged as powerful tools model dynamical systems. DINO (Yin et al., 2022) is a space-time continuous architecture hased on a modulated multiplicative filter network (Fathony et al., 2021) and a NeuralODE (Chen & Zhang. 2019) for modeling the dynamics. DINO is capable of encoding and decoding physical states on irregular grids thanks to the spatial continuity of the INR and through auto-decoding (Park et al., 2019). CORAL is another neural-fie -ierld based architecture, which tackles the broader scope of operator learning, also builds on neta-learnngg (Zintgraf al., 2019; Dupont et al., 2022) to freely process irregular grids. CORAL and DINO are the most similar works to ours, as they are both auto-regressive and capable of processing irregular grids. On the other hand Chen et al. (2022) and Hagnberger et al. (2024) make use of spatio-temporal Neural Fields, for obtaining smooth and compact latent representations in the first or to directly predict trajectory solutions within temporal horizon in the latter. Moreover, they either use CNN or rely on patches for encoding the observations and are therefore not equipped for the type of tasks AROMA is designed for.
Transformers for PDE Several PDE solvers leverage transformers and cross-attention as back- bone for modeling PDEs. Transformers, which operate on token sequences, provide a natural solu- tion for handling irregular meshes and point sets. Li al. (2023) and Hao et al. (2023) introduced transformer architectures tailored for operator learning. Hao et al. (2023) incorporated an attention mechanism and employed mixture of experts strategy to address multi-scale challenges. However, their architecture relies on linear attention without reducing spatial dimensions, resulting in linear complexity in sequence size, but quadratic in the hidden dimensions, which can be prohibitive for deep networks and large networks. Similarly, Li et al. (2023) utilized Pross-attention to embed both regular and irregular meshes into a latent space and applied a recurrent network for time-marching in this latent space. Nonetheless, like GNOT, their method operates point-wise on the latent space. Transolver (Wu et al., 2024) decomposes a discrete input function into a mixture of "slices," each corresponding to a prototype in mixture model, with attention operating in this latent space. This approach, akin to our model, reduces complexity. However, has not been designed for temporal problems. (Alkin et al., 2024) recently proposed versatile model capable of operating on Eule- rian and Lagrangian (particles) representations. They reduce input dimensionality by aggregating information from input values onto "supernodes" selected from the input mesh via message pass- ing while decoding is performed with a Perceiver-like architecture. In contrast, AROMA performs implicit spatial encoding with cross-attention to encode the geometry and aggregate obsevation val- ues. Finally, their training involves complex end-to-end optimization, whereas we favor two simple training steps that are easier implement.
Conclusion and Limitations
AROMA offers a novel and flexible neural operator approach for modeling the lution of physical processes. It is able deal with general geometries and to forecast at any position of the spatial domain. It incorporates in an encode-process-dccod framework attention mecha- nisms, latent diffusion transformer for spatio-temporal dynamics and neural fields for decoding. Thanks to a very compact spatial encoding, its complexity is lower than most SOTA models. Exper- iments with small-size datasets demonstrate its effectiveness. Its reduced complexity holds potential for effective scaling to larger datasets. As for the limitations, the performance of AROMA are still to be demonstrated on larger and real world examples. Moreover, like all dynamical models that op- erate over a latent space, the reconstruction capabilities of the decoder is a bottleneck for the rollout accuracy. Since the encoder and decoder are learning spatial relationships from scratch, conferring the framework high flexibility, the training efficiency does not match that of CNN-based auto- encoders on regular grids. We therefore believe there could be further improvements to be made to achieve a similar performance while keeping the same level of flexibility. Finally, even though our model has some potential for uncertainty modeling, this aspect has still to be further explored and analyzed.
evo-

---

Acknowledgements
We acknowledge the financial support provided by DL4CLIM (ANR-19-CHIA-OOT NUM PEIA-0008, ANR, FRANCE 2030). This work was granted access to the HPC resources of IDRIS under the allocations 2023-A0161015133 made by GENCI.
DEEP-
PHLUSIM (ANR-23-CE23-0023, (20)2 and PEPR Sharp -3-R-23-
022-AD011011332R1 2023-AD011015133RI
References
Benedikt Alkin, Andreas Fürst, Simon Schmid, Lukas Gruber, Markus Holzleitner, and Johannes Brandstetter. Universal Physics Transformers. 2024. URL htt.http///axx...rg//aaa///// 12365.
Matthias Bauer, Emilien Dupont, Andy Brock, Dan Rosenbaum, Jonathan Schwarz, and Hyun- jik Kim. Spatial functa: Scaling fumcta to imagenet classification and generation. CoRR, abs/2302.03 1300, 2023.
Johannes Brandstetter, Daniel E Worrall, and Max Welling. Message passing neural pde solvers. International Conference on Learning Representations, 2022.
Peter Yichen Chen, Jinxu Xiang, Dong Heon Cho, Yue Chang, G A Pershing, Henrique Teles Maia, Maurizio Chiaramonte, Kevin Carlberg, and Eitan Grinspun. Crom: Continuous reduced-order modeling of pdes using implicit neural representations. International Conference on Learning Representation, 2022. URL http: /awxiviivvvv/////v////a////aaaaa wwwg/absee....ooooo.. 2202 02607.
Zhiqin Chen and Hao Zhang. Learning implicit fields for generative shape modeling. Proceedings of the IEEE Computer Society Conference on Computer Vision and Pattern Recognition, 2019-June, 2019. ISSN 10636919. doi: 10.1 N..110///PP..019.000
Emilien Dupont, Hyunjik Kim, S. M. Ali Eslami, Danilo Rezende, and Dan Rosenbaum. From data to functa: Your data point is a function and you can treat it like one. Proceedings of the 39 th International Conference Machine Learning, 1 2022. URL http://arxiv.co/ /bs/ 22C1. 1220..
Rizal Fathony, Anit Kumar Sahu, Devin Willmott, and J Zico Kolter. Multiplicative filter networks. International Conference on Learning Representations. 2021.
Jan Hagnberger, Marimuthu Kalimuthu, Daniel Musekamp, and Mathias Niepert. Vectorized Con- ditional Neural Fields: A Framework for Solving Time-dependent Parametric Partial Differential Equations. In Proceedings of the 41st International Conference on Machine Learning (ICML 2024), 2024.
Zhongkai Hao, Zhengyi Wang. Hang Su, Chengyang Ying, Yinpeng Dong, Songming Liu, Ze Cheng, Jian Song, and Jun Zhu. Gnot: A general neural operator transformer for operator learning. International Conference on Machine Learning, pp. 12556-125666 PMLR, 2023.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp- 770-778, 2016.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-66 6851, 2020.
Diederik P Kingma and Max Welling. Auto-encoding variational bayes. 131iv: 2013.
arXiv preprint
Doyup Lee, Chiheon Kim, Minsu Cho, and Wook-Shin Han. Locality-aware generalizable implicit neural representation. Advances in Neural Information Processing Systems, 2023.
Zijie Li, Kazem Meidani, and Amir Barati Farimani. Transformer for partial differential equations' operator learning. Transactions on Machine Learning Research (April/2023). 2023.
10

---

Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, An- drew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differ- ential equations. International Conference on Learning Representations. 10 2021. URI.
08895.
Phillip Lippe. Bastiaan S Veeling, Paris Perdikaris, Richard E Turner, and Johannes Brandstet- ter. Pde-refiner: Achieving accurate long rollouts with neural pde solvers. arXiv preprint arXiv:2308. 2023.
Lu Lu, Pengzhan Jin, and George Em Karniadakis. Deeponet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators. Nat Mach Intel, 3:218-222 229,, 10 2021. doi: 10. 00.1088/44255--21-0000000 org/abs 1910. 031 :p: / dx. doi. org/ 10. 011////222266-2100000000
URL http::/arxiiv
Aaron van den Oord, Oriol Vinyals. and Koray Kavukcuoglu. Neural discrete representation learn- ing. arXiv preprint arXiv:1711.00937 2017.
Jeong Joon Park, Peter Florence, Julian Straub, Richard Newcombe, and Steven Lovegrove. Deepsdf: Learning continuous signed distance functions for shape representation. Proceedings of the IEEE Computer Society Conference on Computer Vision and Pattern Recognition, 2019-June, 2019. ISSN 10636919. doi: 10. 10.1109/CVPR.20 w0w.P.....09000000 1.00205
Jaideep Pathak, Shashank Subramanian, Peter Harrington, Sanjeev Raja, Ashesh Chattopadhyay, Morteza Mardani, Thorsten Kurth, David Hall, Zongyi Li, Kamyar Azizzadenesheli, et Four- castnet: A global data-driven high-resolution weather model using adaptive fourier neural opera- tors. arXiv preprint arXiv:2202. 12144, 2022.
William Peebles and Saining Xie. Scalable diffusson models with transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 4195-4205, 2023.
Tobias Pfaff, Meire Fortunato, Alvaro Sanchez-Gonzalez and Peter W. Battaglia. Learning mesh- based simulation with graph networks. International Conference Learning Representations. 10 2021. URL http: / arxiv. org/ 2010. 03409.
Michal Rolinek, Dominik Zietlow, and Georg Martius. Variational autoencoders pursue pea direc- tions (by accident). In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 12406-12415, 2019.
Salva Rühling Cachay, Bo Zhao, Hailey Joren, and Rose Yu. DYffusion: dynamics-informed diffusion model for spatiotemporal forecasting. In Advances in Neural Information Processing Systems (NeurIPS). 2023.
Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. arXiv preprint arXiv:2202.0051/2 2022.
Louis Serrano, Lise Le Boudec, Armand Kassaï Koupaï, Thomas X Wang, Yuan Yin, Jean-Noël Vittaut, and Patrick Gallinari. Operator learning with neural fields: Tackling pdes on general geometries. Advances in Neural Information Processing Systems, 2023.
Kimberly Stachenfeld, Drummond B. Fielding, Dmitrii Kochkov, Miles Cranmer, Tobias Pfaff, Jonathan Godwin, Can Cui, Shirley Ho, Peter Battaglia, and Alvaro -Godzalzz. Learned coarse models for efficient turbulence simulation. International Conference Learning Repre- sentation, 2022.
Ricardo Vinuesa and Steven L Brunton. Enhancing computational fluid dynamics with machine learning. Nature Computational Science, 2(6): 366,, 2022.
Haixu Wu, Huakun Luo, Haowen Wang, Jianmin Wang, and Mingsheng Long. Transolver: A Fast Transformer Solver for PDEs on General Geometries. 2024. 2402. 02366.
Yuan Yin, Matthieu Kirchmeyer, Jean- Franceschi, Alain Rakotomamonjy, and Patrick Galli- nari. Continuous pde dynamics forecasting with implicit neural representations. International Conference on Learning Representations, 9 2022. URL 18855.
org/bbs...b...... 2209.

---

Luisa Zintgraf, Kyriacos Shiarlis, Vitaly Kurin, Katja Hofmann, and Shimon Whiteson. Fast context adaptation via meta-learning. 36th International Conference Machine Learning, ICML 2019, 2019-June, 2019.
A Extended Related Work
Diffusion models for PDE Recently, diffusion models have experienced significant growth and success in generative tasks, such as image or video generation (Ho et al., 2020). This success has motivated their application physics prediction. Rühling Cachay et al. (2023) propose DYffusion, framework that adapts the diffusion process to spatio-temporal data for forecasting on long- ting-tmme roll- outs, by performing diffusion-like timesteps in the physical time dimension. PDE-Refiner (Lippe et al., 2023) is CNN-based method that uses diffusion to stabilize prediction rollouts over long trajectories. Compared to these methods, we perform diffusion latent space, reducing the com- putational cost, and leverage the advanced modeling capabilities of transformers.
Local Neural Fields We are not the first work that proposes leverage locality to improve the de- sign of neural fields. In a different approach, Bauer et al. (2023) proposed a grid-based latent space where the modulation function the application of architectures with spatial inductive biases for generation on the latent represen- tations, such as a U-Net Denoiser for diffusion processes. Similarly, Lee al. (2023) developed locality-aware, generalizable Implicit Neural Representation (INR) with demonstrated capabilities generative modeling. Both of these architectures assume regular input structures, be through patching methods grid-based layouts.
dependent on the query coordinate æ. This concept enables
B Implementation details
We illustrate how our diffusion transformer is trained and used at inference in Figure 4 and Figure 5. We provide the diffusion step which acts conditioning input for the diffusion model. We use an exponential decrease for the noise level as in Lippe et al. (2023) i.e. k/K ag min We use the same diffusion transformer block in Peebles & Xie (2023), which relies on amplitude and shift modulations from the diffusion timestamp k:
Diffusion transformer
(0)) 10) + MLP1 (k)
(6)
(2 30) 7(2) MLP2(k)
Z1+11= + Zμy + a1) Attention((()) LayerNorm (1n) + g(0)
(8)
Z Z1t11 + (2) -FFN() LayerNorm (Z -((2)
(9)
12

---

Vāze v1 āzZ'
target tokens
MSE
oulput tokens
6
Linear
DIT Block
Linear
conditionngg embedding
z
k
previous latent token
nolsy estimate
diffusion step
Figure 4: During training, we noise the next-step latent tokens Z and train the transformer to predict the "velocity" of the noise. Each DIT block is implemented as in Peebles & Xie (2023).
DENOISE
oulpur
Latent Refiner
- A4
25 noisy estimate
k
previous latent token
diffusion
Figure 5: At inference, we start from 7 At
(O, I) and reverse the diffusion process to denoise
prediction. We set prediction 7+ 2.4
Encoder-Decoder We provide more detailed description of the encoder-decoder pipeline in Fig-
Local INR We show the implementation of our local INR, both with single-band frequency and muldi-band frequency, in Figure and Figure 8. The cross-attention mechanism enables retrieve a local feature vector f, (z) for each query position x. We then use an MLP to decode this feature vector to retrieve the output value. In practice, we retrieve several feature vectors corresponding each to separate frequency bandwidths. In this case, we concatenate the feature vectors before decoding them with the MLP.
13

---

coordinatos
Sinztion
oupput
f(F
festure wedtr
parstnelers
Encode geometry
Apgregatic values
Sampping
K
A
Attention
Attention
T
2i00
Iatont tolions Z a compaaasse otmenssions
2
+ (2())
hequency enbesiinng
lnverablle tokuens encoohnng spatial
geoottelyyy annane spatial reprssentation query dependling geomolyy
ceovdinaao quary
signal
Figure 6: Architecture of our encoder and decoder. We regularize the architecture as a variational auto-encoder. Pross-attention layers are used aggregate the N observations into M latent tokens, and to expand the M processed tokens to the queried values. We use a bottleneck layer to reduce the channel dimension of the latent space.
x
output value
MLP
Eg (z)
feature vector
K V
Z
Cross-Attention
Q
(z(
frequency embedding
Figure 7: Single-band local INR decoder
14

---

0 22
output value
MLP
(x)
fg
multi band
feature vectors
K V
Cross-Attention
Q
Q
multi band frequency embeddings
71 (x:)
72 (æz)
a' (x)
Figure 8: Multi-band local INR decoder
15

---

B.1 Hyperparameters
We detail the values the hyperparameters used on each dataset: Table 5 presents the hyperparam- eters of the Encoder-Decoder while Table presents the hyperparameters of the Diffusion Trans- former. We use cosine scheduler for the tuning learning rate for both trainings, with an initial max- imum learning rate of 10-3 annealing 10-5 All experiments were performed with NVIDIA TITAN RTX.
For the diffusion transformer, we use K 3 diffusion steps for all experiments and only vary the minimum noise min-
Table 4: Diffusion Transformer Hyperparameters for Different Datasets
Hyperparameters
Burgers
NSle-3
NSle-4
NSIe-5
Shallow-water
Cylinder-Flowe
Airfoil-Flon
hidden-size
128
128
128
128
128
128
128
depth
4
4
4
4
4
4
num heads
4
4
4
4
4
4
mlp ratio
4.0
4.0
4.0
4.0
4.0
4.0
4.0
min_noise
le-2
le-2
le-3
le-3
le-3
le-3
le-3
denoising_stepe
3
3
3
3
3
epochs
2000
2000
2000
2000
2000
2000
2000
For the encoder-decoder, we have the following hyperparameters:
hidden_dim:
The number h of neurons at each hidden layer.
num self attentions: The number of Self Attention layers used for the decoder.
num_laten latenttnt:: The number latent tokens used to spatially project the objervations and geometries.
latent.dim: The dimension e of each latent token.
latent heads: The number of heads use for the Self Attention layers.
latent_dim_head The dimension of each head in Self Attention layer.
cross heads: The number of heads use for the Cross Attention layers.
cross_dim_heate The dimension of each head in a Cross Attention layer.
dim: The number of neurons used in the MLP decoder.
depth The number of layers in the MLP decoder.
frequencies: The different frequencies used for the local INR. We use base 2 for experiments and select 16 frequencies in logarithmic scale per level. For ex- ample, [3, 4, 5] means that we construct 3 frequency embedding vectors, the (cos(2" (T)) sinn22nT)), ..., 72 (cos(22 (2) ..... cos cos((2 sin (cos(2* mx), s2n TT)..... cos (0s())xx) sin Tx))
first 71
cos))........("........ (2s(2"TT2) rin)) for the second
and for the third "3
dropout sequence: The ratio of points that are ignored by the encoder.
feature_dim: The dimension of the feature vector.
encode_ged If we use a cross-attention block to encode the geometry.
max encoding freq: The maximum frequency used for the frequency embedding y of the encoder.
kl_weihh:: The weight 3 used for the VAE training.
epochs: Number of training epochs.
The most important hyperparameter of the encoder-decoder is the number of tokens M that are used to aggregate the observations and geometries. We show the impact it has on the quality of reconstructions in Table 6.
16

---

Table Hyperparameters of the Decoder for Different Datasets
Hyperparameters
Burgers
NS1e-3
NS1e-4
NS1e-5
Shallow-wate water
Cylinder-Fai Flow
Airfoil-Flov
hidden dim
128
128
128
128
128
128
128
num_self self.attentions
2
2
3
2
2
3
num latents
32
32
256
256
32
64
64
latent_dim
8
16
16
16
16
16
16
latent_heade _eads
4
4
4
4
4
4
latent_dim. diim_eead
32
32
32
32
32
32
32
cross _head
4
4
4
4
4
4
4
cross gross_dim_head
32
32
32
32
32
32
32
dim depth_ing inr frequencies
128
128
128
128
64
128
128
3
3
3
3
3
3
[3.4, 3,4,55 5]
[2, 3]
[3,4,5, 5]
[3, 4, 5]
3]
[3.4, 4,
[3, 4, 5]
dropout seopout_sequence sequeece
0.1
0.1
0.1
0.1
0.1
0.1
0.1
feature dim encode geo max encoding freq
16
16
16
16
16
16
16
False
True
False
False
True
True
True
4
4
5
4
5
wliigt
lc-4
lc-4
lc-4
le-5
le-5
1c-5
lc-5
epochs
5000
5000
5000
5000
5000
5000
5000
C Additional results
C.1 Time complexity analysis
We denote as the number of observations u, M as the number of tokens used to compress the information, T as the number of autoregressive calls in the rollout, K as the number of refinement steps, and as the number of channels used in the attention mechanism. The most computation- ally expensive operations in our architecture the cross-attention and self-attention blocks. For simplification, we omit the geometry encoding block this study.
The cost of the cross-attention in the encoder is O(N and similarly, the cost of the cross- attention in the decoder is Md). Let L1 and L2 represent the number of layers the de- coder and diffusion transformer, respectively. The cost of the self-attention layers in the decoder is d22, while the diffusion transformer, it is (4L
To unroll the dynamics, we encode the initial condition, obtain the predictions in the latent space, and then decode in parallel, yielding a total cost of O((2N + 4KT L, Md) As expected, our architecture has linear complexity in the number of observations through the cross-attention layers. In contrast, GNOT relies on linear attention, resulting prediction, where J. is the depth of the network. At inference, the cost per step along trajectory is L.d for GNOT, compared to 4K T- M2d for AROMA.
time complexity of ((LNN)d) for each
For instance, using K = 3, M = 64, N = 4096, and d = 128, GNOT's cost is approximately 10 times that of AROMA for each prediction throughout the rollout. Therefore AROMA more efficient when M & N.
C.2 Encoding interpretation
We provide Figure 9 a qualitative analysis through cross-attention visualizations how the geom- etry encoding block helps to capture the geometry of the domain. In the first cross-attention block, the query tokens are not aware of the geometry and therefore attend to large regions of the do- mains. This lets the model understand, where the boundaries of the domain and therefore where the cylinder is. Once the query tokens have aggregated the mesh information, the cross attention between 8000 and the positions sharper and depend the geometry.
C.3 Example rollouts
We show examples of rollout predictions using AROMA on Burgers dataset in Figure 10, on Navier- Stokes 10-3 dataset in Figure 12 and on CylinderFlow in Figure 13. AROMA returns predictions that remain stable and accurate, even outside the training time horizon.
17

---

Cross-Attention T'
Cross-Attention T Tgeo
geo aggregates information of the different positions domain and therefore reffects geometry. Its cross-attention maps sharper and adapted the geometry.
Can be interpreted as prior spatial representation of the different domalns. The query tokens are learnt to attend to the - important regions of the domains.
Geometry encoding
Tolen 22
Geometry encoding
Token 23
Figure 9: Evolution of the cross-attention maps between the geometry encoding stage and the ob- servation encoding stage. Blue means the cross-attention value close to zero while yellow means the cross-attention score is close to one.
rollout prodiction
Figure 10: Test example long rollout trajectory with AROMA on Burgers. Left is the predicted trajectory and right is the ground truth.
203rchout
Figure 11: Uncertainty of AROMA over rollout steps. The blue line is the mean prediction while blue shade represents mean + 3 x standard deviation.
18

---

For Navier-Stokes, we show an example of test trajectory in the training horizon (Figure 12a) and in extrapolation (Figure 12b).
(a) In-t
21
126
D24
UM
(b) Out-
Figure 12: Test example rollout trajectories with AROMA on Navier-Stokes 10-3 Top: pre- dicted trajectory on In-t. Bottom: trajectory Out-t. First row in each subfigure shows the predic- tion, the second row shows the ground truth.
Predictions
Gronnd Truth
Delta
Figure 13: Visualization of AROMA's predictions on Cylinder for (Out-t). The left panel shows the prediction, the middle panel displays the ground truth, and the right panel is the point-wise error.
19

---

C.4 Scaling experiments
In Figure 14, we compare the reconstruction and prediction capabilities of CORAL and AROMA on Navier-Stokes 1 10 given the number of training trajectories. As evidenced, our architecture outperforms CORAL significantly when the number of trajectories is greater than 10 highlighting its efficacy in handling large amounts of data..
Thst reconstrucion eror
TMt prediction errer
CORAL
CORAL
AR
ANOMA
Dalatto fror
Rubdvv Eror
08
08
10
30
10
15
Number of Training Samples
Number of Training Samples
(a) Step Autoencoding
(b) Step Rollout prediction
Figure 14: Scaling comparison of AROMA & CORAL: relative L2 error with respect to the number of training trajectories
C.5 Spatial tokens perturbation analysis
To validate the spatial interpretation of our latent tokens, establish baseline code Z and introduce perturbations by sequentially replacing the j-th token, 2). with subsequent tokens along the trajectory, denoted as z]: ...,, 2/. Thus, the perturbed tokens mirror 2 in all aspects except for the j-th token, which evolves according to the true token dynamics. We show reconstruction visualizations of the perturbed tokens in figs. 15 to 22. On the right side, we show the groundtruth of the trajectory. On the left side, is the change in AROMA'S prediction in response to the token perturbation. These figures show that the perturbation of a token only impacts the reconstructed field locally. which validates the spatial structure our tokens. Additionally, we can notice some interesting effects of the token perturbations near the boundaries in figs. 17 and 22: our encoder- decoder has discovered from data and without explicit supervision that the solutions had periodic boundary conditions by leveraging the encoded geometry and the function values. This validates the architecture of our cross-attention module between the function values, the spatial coordinates and the geometry-awe tokens.
02
Figure 15: Perturbation analysis on Burgers. Token 0.
20

---

12-
02
-1
Figure 16: Perturbation analysis on Burgers. Token 1.
08
500
Figure 17: Perturbation analysis on Burgers. Token 2.
18
08
Figure 18: Perturbation analysis on Burgers. Token 3.
02
300
Figure 19; Perturbation analysis on Burgers. Token 5.
21

---

12-
48
Figure 20: Perturbation analysis on Burgers. Token 6.
15-
02
415
Figure 21: Perturbation analysis on Burgers. Token 7.
10
08
12-
Figure 22: Perturbation analysis Burgers. Token 8.
22

---

C.6 Ablation studies
Number of tokens M We show the impact of the number of latent tokens on the Navier-Stokes
1 x 10 dataset in Table 6. We train our auto-encoder with 10000 trajectories. We can see that the performance increases with the number of tokens.
#Latent Tokens
Test Reconstruction error
64
0.02664
128
0.0123
256
0.01049
Table 6: Influence of the number of latent tokens on the test reconstruction capabilities on Navier- Stokes x 10-a, Performance in Relative L2 Error.
Auto-encoding vs VAE Our framework can also used without the KL regularization, and could potentially be employed with other forms of regularization, such as L2 regularizaton vector- quantization (Oord et al., 2017). We investigated in Table the impact the KL regularization had on the overall rollout performance, and selected autoencoder with L2 regularization (weight decay) as baseline. Our conclusion that using an autoencoder with L2 regularization a viable alterna- tive the VAE in some cases for achieving a smooth latent space. The autoencoder demonstrated superior performance on two datasets (Burgers and Navier-Stole reconstruction errors, which translate into better rollout performance. However, for the more chal- lenging Navier-Stokes x 10-S case, the autoencoder's latent space exhibited high variance, which may explain the observed performance difference with the VAE.
x 10- explained by its lower
No-diffusion vs diffusion As an ablation, we also measured the influence of the diffusion formu- lation the rollout accuracy by comparing to the same transformer architecture trained directly with an MSE on the mean tokens. The deterministic version of AROMA shows consistently robust performance and even surpasses the diffusion version on the Navier-Stokes 1 10-4 case (Table 7). This demonstrates that the latent tokens obtained with AROMA contain meaningful information for dynamics modeling. On the other hand, the deterministic version yields less accurate long rollouts on Burgers or KS in Figure 23 and Figure 3. Note that using diffusion allows us to model the trajec- tory distribution, which opens the way to infer statistics on this distribution. This is key, for example, when modeling uncertainty, which is critical problem for these models.
Latent MILP Latent Transformer Modeling interactions the local and global levels is key to learn the dynamics faithfully. Experiments using MLPs (Table 7) as time steppers which do not con- sider interactions between tokens lead to significantly lower performance compared transformers.
Table 7: Ablation Study. Metrics in Relative L2 on the test set.
Model
Burgers
Navier-Stoker
Navier-Stokes
1 10 10-4
10
AROMA + auto-encoding
3.43 x 10-2
5.02 x 10-2 2
2.10 x 10-1
AROMA w/o diffusion
4.31 x 10-2
7.50 x 10-2
1.28 x 10
AROMA + mlp
1.11 x 10-1
1.00 10"
8.25 10-1
AROMA
3.65 x 10-2
1.05 10-1
1.24 10
C.7 Kuramoto-Sivashinsky failure case
We conducted additional experiments on a chaotic ID PDE, the Curamoto-SSvashing (KS) equa- tion. We found that AROMA currently struggles with dynamics that exhibit chaotic phenomena and non-decaying spectra, as shown in Figure 23. The primary limitation appears to be the reconstruction capabilities of the encoder-decoder For the KS equation, we found that obtaining reconstructions with an MSE in the range of le-10 to le- was necessary for accurate spectrum reconstruction. Like all models leveraging reduced latent representation space, AROMA inherently loses some of
23

---

(i.e. a high variance) the fine-grained details necessary for accurately capturing chaotic behavior. The diffusion frame- work slightly improves the high correlation time compared the deterministic version, however the main bottleneck comes from the decoder (Table 8). In conclusion, while AROMA performs very well on simpler dynamics, dealing with chaotic phenomena requires more involved modeling that explicitly targets the chaotic component. Note that using dedicated modules for this purpose is current practice fluid dynamics c.g. LES (Large eddy simulation).
Table 8:
Test results on the KS equation. The evaluated metrics include: 1-step prediction MSE, MSE over the entire rollout (160 timestamps), and the duration for which the correlation between the generated samples and the ground truth remains above 0.8.
Baseline
1 step. MSE
Rollout MSE
Cor.. 2 0.8
AROMA w/o diffusion
1.25 10
2.20 x 10"
32.8s
AROMA
1.81 x 10-5
2.07 x 10
36.0s
Spectrum the step predictions
301
Grsund buth
101
AROMA dffuaion
AROMA
10
107
Ampbuue
10-4
AROMA No Diffusion
10-4
100
120
20
60
160
120
Staven.mber
Time
(a) Spectrum of AROMA's 1-step predictions vs ground truth.
(b) Comparison of the MSE loss (1) with and without diffusion.
Figure 23: Qualitative results on KS equation.
C.8 Latent space dynamies
For Navier-Stokes, we show how the mean (Figure 24) and standard deviation tokens (Figure 25) evolve time for given test trajectory. We show the predicted trajectory of the latent tokens Z in the latent space in Figure 26. In practice, the tokens where the logvar on Figure 25 do not impact the prediction (Rolinek et al., 2019). We can therefore see, that ouf of the 16 tokens used, the most influential ones are Token 6, 7, 8, 9, 15, 16, as they clearly exhibit non-noisy patterns.
24

---

Token
Token
0.0
0.0
-0.5
-0.5
-1.0
-1.0
-1.0
40
20
40
20
Token
Token 6
Token
Token 8
4
-1
-2
-2
-2
20
Token
2
0.0
-0.5
-1.0
-1.0
-2
40
20
40
20
40
20
40
Taken 13
Taken 14
Taken 15
Taken 16
0.0
-0.5
-1.0
40
20
Figure 24: Latent space dynamics on Kavier-Stokes le-3 Mean tokens over time. Each color line is a different token channel.
25

---

Token
Token
Token
0.0
0.0
-05
-0.5
-0.5
-10
-1.0
11.0
-1.0
-1.5
-1.5
-1.5
-20
-2.0
-2.0
20
40
40
Token 5
Taken 6
Token
Taken
0
0-
0.0
-1
-2
-.5
-2-
-4-
-
-7.5
-6
1000
-8
20
Token
0.0
-0.5
00.5
-1
-1.0
-1.5
20
40
20
40
40
20
Token 13
Token 14
Token 15
Token 16
0.0
0.0
-2
-2
-4
-1.0
-6
-15
-1.5
-8-
-6-
-2.0
-8
20
40
20
Figure 25: Latent space dynamics on Kavier-Stokes Logvar tokens over time. Each color line is different token channel.
26

---

Token
Token
0.0
0.0
0.0
-0.5
-0.5
-1.0
-1.0
-1.0
40
20
40
20
Token
Token 6
Token
Token 8
-1
-2
-2
20
Token
0.0
0
-0.5
-1.0
-1.0
-2
40
20
40
20
40
20
40
Taken 13
Taken 14
Taken 15
Taken 16
0.0
-0.5
-1.0
40
20
20
Figure 26: Latent space dynamics on Navier-Stokes Predicted tokens over time. Each color line a different token channel.
27

---

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 conducted extensive experiments multiple settings to obtain robust results across several datasets.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made 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 perceived 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.
fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: There is limitation section.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are 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-specificattin, asymptotic approximations only holding locally). The au- thors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, c.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the ap- proach. For example, facial recognition algorithm may perform poorly when image resolution is low images taken in low lighting. Or a 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 applicable, the authors should discuss possible limitations of their approach to ad- dress 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 be that reviewers discover limitations that aren't it acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
28

---

Justification:
No theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theo- rems.
The proof' can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be comple- mented by formal proof's provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclu- sions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: Detailed hyperparameters and code.
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, a No answer this question will not be perceived well by the reviewers: Making the paper reproducible is 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 or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary either make it possible for others replicate the model with the same dataset, or provide access to the model. In general. releasing code and data often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all sub- missions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(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 (e.g., a large language model), then there should either be way to access this model for reproducing the results or a way to re- produce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case au- thors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers 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 and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
29

---

Answer: [Yes]
Justification:
Everything will be made available upon acceptation.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please the NeurlPS code and data submission guidelines (https: /wtps/ipsc.o. cc/
pub / gu des Coldeuumme i ss cisioooollioo cy) for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: //nips. cc pec/pubbii. ptppppccluuueeeeeeeepeepeppppeee /gw.ieuu///e///ee/ /Code SideeSubar Submis s ssisinnn sissPooll ooyyyy for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of 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 anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
All the settings are disclosed explicitly.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropri- ate information about the statistical significance of the experiments?
Answer: [No]
Justification:
No error bars yet.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, at least for the experiments that support main claims of the paper.
The factors of variability that the error bars are capturing should be clearly 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 a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
30

---

is OK to report I-sigma error bars, but one should state it. The authors should prefer- ably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative emror rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproducc the experiments?
Question:
Answer: [Yes]
Justification: Type of gpu detailed.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU 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 well 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 't make it into the paper).
9. Code Of Ethies
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https: /ttpa://neuuue /neurips.....
Answer: [Yes]
Justification:
Yes we conformed the code of ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is a special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [No]
Justification:
No social impacts mentioned.
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 no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact spe- cific 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 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 to
31

---

generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people 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 as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitiga- tion strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe 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, or scraped datasets)?
Question:
Answer: [No]
Justification:
This paper causes such risk.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by quiring that users adhere to usage guidelines or 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 challenging, and many papers do require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: All datasets mentioned, and code also credited when appropriate.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, if possible, include URL.
The name the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswitheee wwwoe..de.....oooooooooooooo.... com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's is creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and the documenta- tion provided alongside the assets?
32

---

Answer:
[No]
Justification:
There is no such dataset.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the d as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the pa- per include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: This paper is not related research this topic.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contri- bution of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, cura- tion, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: No such risk.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country which research conducted, IRB approval (or equiva-
lent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
33


Reviews:
[{"review_id": "xxPl9dLU12", "reviewer": "Reviewer_rnRB", "summary": "An innovative approach for improving the modeling of partial differential equations (PDEs) using local neural fields is presented in the paper \"AROMA: Preserving Spatial Structure for Latent PDE Modeling with Local Neural Fields\" (Attentive Reduced Order Model with Attention). Through the provision of a versatile encoder-decoder architecture that can handle a variety of data sources, including regular grids and point clouds, AROMA seeks to alleviate the shortcomings of current neural operator models.", "questions": "1. What are the main drawbacks that AROMA seeks to solve with respect to current neural operator models for PDEs?\n3. What function does the AROMA architecture's conditional transformer serve?\n5. How can AROMA guarantee the accuracy and stability of its forecasting models?\n5. How is the model's performance improved by the encoder and decoder's two-stage training process?\n6. In the experiments, what kinds of datasets were employed, and how did AROMA fare in comparison to other models?", "limitations": "na", "rating": 9, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "na", "weaknesses": "na", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "* **Q1** | What are the main drawbacks that AROMA seeks to solve with respect to current neural operator models for PDEs?\n\nAROMA addresses key limitations of existing neural operator models for PDEs. Traditional transformer-based methods, such as those by Li et al. (2023) and Hao et al. (2023), unroll dynamics directly in the original space, leading to high complexity and inefficiency. Neural-field methods such as DINO or CORAL unroll the dynamics in a latent space without spatial prior. AROMA mitigates these issues by encapsulating domain geometry and observation values into a compact latent representation, allowing for efficient forecasting with reduced computational cost. This approach simplifies the training process and avoids the need for prior feature engineering, making it particularly effective for complex geometries.\n\n\n* **Q2** | What function does the AROMA architecture's conditional transformer serve?\n\nIn AROMA, the conditional transformer plays a crucial role in modeling and forecasting dynamics efficiently. It operates on a fixed-size compact latent token space that encodes local spatial information. This transformer processes the encoded spatial data and models the dynamics while capturing spatial relations both locally and globally across tokens. Additionally, the conditional neural field utilized in the decoding stage enables querying forecast values at any point within the spatial domain, enhancing the model's flexibility and accuracy. Said otherwise, AROMA is a mesh free solution that accepts any geometry as input and can be queried at any point in the spatial domain.\n\n\n* **Q3** | How can AROMA guarantee the accuracy and stability of its forecasting models?\n\nAs AROMA is entirely data-driven, there is no absolute guarantee that the model will always approximate the true solution perfectly at inference. However, extensive experiments have been conducted throughout the paper to assess the accuracy and stability of the method under various initial conditions. Experiments have been performed on long rollouts (see appendix C.2 Fig. 9 for example), including forecasting beyond the training horizon domain. The results consistently align with state-of-the-art methods, suggesting that AROMA is stable for all the equations studied. Additionally, by incorporating a diffusion transformer to unroll the dynamics, it is possible to generate multiple solutions and analyze their variance or the average cross-correlation between these solutions. Said otherwise it allows to generate samples from the predictive distribution of the trajectories. This approach allows computing different statistics on this distribution, and helps in identifying when the model may have diverged from the expected ground truth. Preliminary results illustrating this behavior are presented in Figure 10 of the paper.\n\n\n* **Q4** | How is the model's performance improved by the encoder and decoder's two-stage training process?\n\nThe two-stage training process in AROMA is very stable and aligns with previous works in the literature. This separation allows for specialized training of each component, ensuring that the encoder-decoder effectively encodes spatial features and the processor accurately predicts dynamics. Experimentally, it is easier to train than the end-to-end training alternative.\n\n\n* **Q5** | In the experiments, what kinds of datasets were employed, and how did AROMA fare in comparison to other models?\n\nThe experiments were conducted on representative spatio-temporal forecasting problems, including 1D and 2D dynamics with periodic boundary conditions, as well as domains with complex geometries and very diverse input types, such as point sets, grids, and meshes. AROMA exhibited outstanding performance on these datasets, achieving state-of-the-art results when compared to existing neural field and transformer-based methods.", "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": "s3yCc2A5NQ", "reviewer": "Reviewer_gPBs", "summary": "The paper introduces a novel framework designed to enhance the modeling of partial differential equations (PDEs) using local neural fields. It proposes a flexible encoder-decoder architecture that achieves smooth latent representations of spatial physical fields from various data types and employs a diffusion-based formulation to achieve greater stability and enable longer rollouts compared to conventional MSE training. The authors show the superior performance of their framework on 1D and 2D PDEs.", "questions": "It is unclear what the authors mean by patching for data input types.\nThe visual results for the cylinder and aerofoil cases are not provided. Is there any specific reason these were not included?\nWhat would be the training cost of extending this to 3D PDEs.", "limitations": "While AROMA demonstrates effective performance on small-size datasets, its scalability to larger datasets and more complex real-world scenarios is not fully established. Scaling up neural field-based methods to handle larger volumes of data without compromising performance remains a challenge.", "rating": 5, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The authors claim that the proposed method can handle a variety of data types including regular-grid inputs and point clouds, eliminating the need for patching and allowing efficient processing of diverse geometries.\n\nThe diffusion-based formulation enhances stability and enables longer rollouts compared to traditional methods.\n\nAROMA demonstrates superior performance in simulating 1D and 2D equations, capturing complex dynamical behaviors effectively.\n\nThe framework leverages attention blocks and neural fields, resulting in a model that is easy to train and achieves state-of-the-art results without requiring prior feature engineering.", "weaknesses": "Although the proposed method reduces computational cost compared to existing transformer architectures, it may still be computationally expensive for very large datasets or extremely high-resolution simulations.\n\nThe experiments are performed on benchmark datasets, and the performance on larger and real-world examples remains to be demonstrated.\n\nIt is also unclear if the approach always produces stable results.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "* **W1** | Although the proposed method reduces computational cost,  [...] .\n\nOur architecture employs an encoder-decoder structure that includes cross-attention blocks. The computational complexity of these cross-attention operations is $\\mathcal{O}(NMd)$, where $N$ represents the number of observations, $M$ is the number of latent tokens, and $d$ is the hidden dimension. This gives our model a linear time complexity of $\\mathcal{O}(N \\times K)$ with respect to the number of observations, where $K$ is a constant factor. Note that $M$ and $d$ ($K=Md$) are parameters of the architecture that could be set at different values, allowing for a compromise between accuracy and complexity.\n\nWe acknowledge that, despite this linear complexity, the constant $K$ can still be substantial when dealing with large datasets or high-resolution simulations. This highlights a trade-off between computational efficiency and the scale of the data being processed. While our approach significantly reduces the computational cost compared to traditional transformer architectures, we recognize that further optimization or alternative strategies may be needed to handle extremely large datasets or high-resolution scenarios effectively.\n\n* **W2** | The experiments are performed on benchmark datasets, [...].\n\nThe current experiments utilize benchmark datasets to establish baseline performance and validate the effectiveness of the proposed method in a controlled setting. These datasets, introduced in various studies (Pfaff et al. 2020, Li et al. 2020, Yin et al. 2022, Brandstetter et al. 2022), were generated using different solvers and for different objectives. AROMA has demonstrated consistent performance across all these datasets, highlighting its robustness. Nevertheless, it is crucial to evaluate its applicability and performance in larger, real-world scenarios. Future work will focus on scaling the method to address larger datasets and more complex, real-world applications. In order to better assess the capabilities of AROMA, we have performed additional experiments on the more challenging and chaotic KS-equation (see Table 2 and Figure 1 of the pdf page) - see also the comments in the \"global response\" to the reviewers. In conclusion, when AROMA performs extremely well on a large range of dynamics and problems, like all the models using a reduced latent representation, it faces difficulties for more complex problems involving chaotic phenomena that require more specific approaches.\n\n\n* **W3** | It is also unclear if the approach always produces stable results.\n\nOur results provide strong evidence of the method's stability. Notably, the time-correlation plot in Figure 3 indicates that AROMA maintains stability during extended rollouts on Burgers' equation. Additionally, Table 2 demonstrates that the method can effectively extrapolate beyond the training horizon and remains stable under perturbations of the observations.\n\n* **Q1** | It is unclear what the authors mean by patching for data input types.\n\nThe term comes from visual transformers operating on regular grids, and in this case refers to a predefined regular spatial partition of the point grids. In a more general context, the term \"patching\" refers to a process that can be somewhat arbitrary, particularly when dealing with irregular meshes. In such cases, patching should ideally be derived from a systematic partitioning of the data. For example, [1] employs this idea by partitioning the domain on different resolutions to obtain frame patches. [2] considers graph representations of input points and agglomerates points into \"supernodes\" that are then embedded in a latent space. In contrast, AROMA eliminates the need for transforming data into patches. Instead, it directly learns a spatial representation of the data, making it more flexible and effective across various data structures without relying on predefined patches.\n\n[1] Solving High-Dimensional PDEs with Latent Spectral Models, Wu et al 2023.\n[2] Alkin, B., Fürst, A., Schmid, S., Gruber, L., Holzleitner, M., & Brandstetter, J. (2024). Universal Physics Transformers. 1–27. \n\n* **Q2** | The visual results for the cylinder and aerofoil cases are not provided. \n\nThank you for this remark. We have added visualizations of AROMA's predictions in the rebuttal pdf page for the dataset Cylinder in the *Out-t* regime. We will provide additional visualizations of the results for the camera ready. \n\n* **Q3** | What would be the training cost of extending this to 3D PDEs.\n\nAs detailed in our answer to **W1**, our model exhibits a linear complexity of $\\mathcal{O}(NMd)$, where $N$ is the number of observations, $M$ is the number of latent tokens, and $d$ is the hidden dimension. For a domain $\\Omega$ discretized on a regular 3D grid of size $L$, the total number of observations is $N = L^3$. We have not conducted 3D experiments with AROMA yet, and therefore do not know the number $M$ of tokens to allow for a faithful reconstruction and dynamics modeling. We then provide some indications below by considering two scenarios.\n\nIf we apply a downsampling factor of 4 to determine the number of tokens (as done in the N.S. experiments in the paper), then the number of tokens $M$ becomes $(L/4)^3$. Hence, the additional complexity due to scaling from 2D to 3D, with downsampling, is proportional to $L^2 / 4$. For example, with $L = 32$, the computational cost in 3D is multiplied by $256$ compared to 2D.\n\nIf we maintain the number of tokens fixed (i.e., not using a downsampling rule), the complexity simply scales with the increase in dimensions, resulting in a cost multiplication factor of $32$ when moving from 2D to 3D.\n\nIn future work, we plan to experimentally investigate how AROMA scales with these two strategies when moving from 2D to 3D.", "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": "cat9lHnGYG", "reviewer": "Reviewer_xV3P", "summary": "The paper proposes a framework using autoencoder and diffusion transformer for predicting the forward dynamics of time-dependent PDEs. Leveraging cross-attention and neural fields, the framework is able to handle different types of meshes and geometries. The authors demonstrate the effectiveness of their proposed framework on several 1D and 2D PDE problems with different geometries.", "questions": "1. For systems that feature a decaying spectrum (e.g. fluid flow with relatively high viscosity), the proposed method is a good fit, which is not surprising. Will the method still perform well on slightly more chaotic systems like KS-equation or 2D Komogorov flow and how does the spectrum of reconstructed snapshots from decoder look like in these cases?\n\n2. How many refinement steps does the model use?\n\n3. There are some closed related works that not have been discussed. For example, the similar idea of  learnable latent tokens have been explored in \"Solving high-dimensional pdes with latent spectral models.\" Using pretrained autoencoder to dervie a mesh-reduced space and then learn to forecast dynamics in the latent space is also studied in \"Predicting Physics in Mesh-reduced Space with Temporal Attention\" and \"Latent Neural PDE Solver: a reduced-order modelling framework for partial differential equations\".\n\n4. Typos: \n* line 228: 3D Shallow-Water Equation (Navier-Stokes1 × 10−3)\n* Table 4&5: \"opochs\"", "limitations": "The authors have discussed the limitation in the conclusion part.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The proposed framework adopted many existing techniques in a natural and reasonable way, such as using local neural fields and cross attention to pass information between nodes in different meshes, diffusion-based predictor/refiner. The motivation and benefits of various components in the framework are illustrated clearly.\n\n2. The empirical performance of the proposed model is strong across several benchmarks.\n\n3. The authors explain their method well and overall the paper is easy to follow.", "weaknesses": "1. Many of the techniques used in the paper are taken from existing models (authors have properly acknowledged them), so the technical advancement for the paper might not be that much. Nevertheless, I think the authors have done a decent job tweaking them to work well, evidenced by the experiments, and explained why these techniques can be helpful to PDE modeling.\n\n2. One issue in the experiment is that the authors have employed a diffusion predictor to predict the dynamics (which can be seen as a predictor-corrector scheme) whereas other baselines are deterministic single-step predictor. With that said, the proposed framework uses more NFE (proportional to the number of diffusion steps) than other baselines. The authors did provide an ablation on Burgers' which shows that AROMA without diffusion still outperforms other model, but it is unclear for other more chaotic and higher-dimensional cases.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "* **W1**  | Many of the techniques used in the paper are taken from existing models [...] .\n\nThank you for your feedback. It is true that our paper builds upon established techniques and models, which we have duly acknowledged. However, our contributions lie not only in leveraging these techniques but in integrating and adapting them in a novel way to advance PDE modeling. Specifically:\n\n1. *Novel Integration*: We present a unique framework that combines attention blocks and neural fields in a way that has not been done before, particularly for processing domain geometries and unrolling dynamics in PDEs. This integration allows for a more streamlined and effective approach to handling complex geometries and dynamics. Up to our knowledge, this is the first approach that allows for an automatic implicit latent encoding, leveraging local and global spatial relations, for a diversity of geometries, for modeling spatio-temporal dynamics. Concurrent approaches, [1,2] make use of pre-defined patching strategies that are less flexible for handling these diverse geometries.\n\n2. *Empirical Validation*: The experiments we conducted demonstrate that our approach achieves state-of-the-art performance across various datasets, confirming the effectiveness of our formulation. More importantly, these results validate our hypothesis that a latent space preserving spatial information is crucial for accurately modeling dynamics in a reduced representation. Additionally, our method systematically outperforms other transformer architectures, underscoring its robust performance and well-founded approach.\n\n3. *Insights and Rationale*: We experimentally validate through cross-attention and perturbations that our latent space has a spatial interpretation.\n\n[1] Solving High-Dimensional PDEs with Latent Spectral Models, Wu et al 2023.\n[2]  Universal Physics Transformers. Alkin et al. (2024).\n\n* **W2**. | One issue in the experiment is that the authors have employed a diffusion predictor [...] . \n\nIn the supplementary PDF, we present an ablation study showing that diffusion is not essential for accurate predictions, as AROMA without diffusion still outperforms Neural Field and Transformer baselines (Table 1). However, diffusion improves long rollouts, as seen in Figure 3 (Burgers) and Figure 1b (KS equation). \n\nThe diffusion approach opens the door to uncertainty modeling, and is state-of-the-art for e.g. video prediction. To fully leverage this for PDE, future work should consider metrics beyond MSE to better capture the advantages of diffusion models.\n\n* **W3** | With that said, the proposed framework uses more NFE [...].\n\nFirst, using a small number of diffusion steps $(K=3)$ strikes a good balance between accuracy and computational cost, as the processor works with fewer tokens ($M < N$).  Table 1 in the supplementary pdf compares three model versions: AROMA with $(K=3)$ steps, AROMA without diffusion, and a local MLP.  The MLP which ignores the relation between tokens does not perform well, while deterministic and diffusion transformers show similar results.  Finally, although using few diffusion steps helps with noise augmentation and for long rollouts, it does not fully utilize diffusion’s potential, which typically requires more steps (e.g., K=1000). For datasets characterized by chaotic dynamics, future work could explore training with a full diffusion process ($K=1000$).\n\n* **Q1** | For systems that feature a decaying spectrum [...].\n\nThis is an excellent question. As you correctly point out, the encoder-decoder architecture may indeed become a bottleneck in scenarios where high-frequency amplitudes are challenging to model and represent accurately.\n\nTo investigate this point, we conducted an experiment with the KS equation using the specific setup described by [3]. In this setup, the model predicts $u_{t+4\\Delta t}$ from  $u_t$ over training trajectories spanning $[0, 140 \\Delta t]$, and during testing, predictions are unrolled up to  $[0, 640 \\Delta t]$. Note that we do not employ the \"predict-difference\" trick here. In the dataset, each trajectory is generated with varying grid resolutions ($\\Delta x$) and time-stepping intervals ($\\Delta t$). To accommodate these variations, we provided additional tokens $T_{\\Delta x}$  and $T_{\\Delta t}$ as context for the transformer, while keeping the rest of the architecture unchanged. The results are detailed in Table 2 and Figure 1.\n\nInterestingly, while the deterministic version of the model achieves a lower 1-step prediction mean squared error (MSE) compared to the diffusion version, the diffusion model shows better performance in maintaining high correlation over longer time horizons. The primary identified limitation is the reconstruction capability of the encoder-decoder, which impacts the overall predictive accuracy of the transformer. Figure 1.a in the pdf page, represents the amplitude spectrum of AROMA's prediction vs the real spectrum. We can see in this figure that AROMA's decoder overestimates the weights of high frequencies and is faithful up to the 40th mode. In order to remedy to this problem, one would need to lower the reconstruction error corresponding to these frequencies. Specifically, AROMA achieved an MSE reconstruction of $1 \\times 10^{-6}$ (i.e. a Relative L2 error of $0.08$%) on this dataset, when a relative error around $0.001$% would be optimal. This is a non trivial issue requiring further investigations.\n\n[3] Lippe, et al. PDE-Refiner: Achieving Accurate Long Rollouts with Neural PDE Solvers. Neurips 2023.\n\n* **Q2** | How many refinement steps does the model use?\n\nWe use K=3 steps throughout the paper for all experiments. We chose this number based on the study of Lippe et al. 2023 and some preliminary trials. \n\n* **Q3** | There are some closed related works that not have been discussed. \n\nThank you for these references. We will add them together with a discussion to the related work section for the camera-ready.", "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": "uIjNU3vAUj", "reviewer": "Reviewer_qbYC", "summary": "This paper presents a new pipeline for modeling PDE systems, especially for learning local neural fields. It designs a new encoder-decoder for absorbing any type of input, which avoids the constraint on meshes and cloud points. It can handle diverse geometries. A diffusion-based transformer architecture is used to model the latent dynamics. The numerical experiments have shown the effectiveness of the proposed method.", "questions": "- In lines 90-91, for the second stage of training, do you fix the encoder-decoder or pretrain encoder-decoder and then fine-tune the entire network? \n\n- The authors claim that the encoder-decoder is principled. I didn’t see a principle or theoretical guarantee behind that.  \n\n- In Table 1, why does DNOT work worse than FNO on the 1D Burgers case? Also, as a standard baseline, FNO works pretty good among all baseline models and the proposed AROMA. \n\n- I am also curious why FNO is not considered in the experiments of temporal extrapolation. FNO can also be modified in an auto-regressive way.", "limitations": "Please see my concerns in **Weaknesses** and **Questions**.", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- This paper proposes an interesting encoder-decoder for handling arbitrary input geometries. \n\n- This paper has done multiple experiments on different PDEs to compare the model performance.\n\n- The paper is well-written and well-organized. \n\n- This research topic is of interest to the general scientific machine-learning community.", "weaknesses": "- The motivation of using a VAE-type encoder-decoder is not well justified. The authors might add more details about the benefits of using a VAE setup for an encoder-decoder instead of a deterministic auto-encoder. Also, my concern is that using Gaussian distribution might constrain the representation capability of the latent features. The real-world dynamics or complex turbulence dynamics sometimes present heavy-tailed characteristics. \n\n- This paper considers a diffusion-based transformer for learning latent dynamics. The training efficiency and computational memory might be an issue. Moreover, I think the clarification of its benefits is not sufficient. This module can be replaced by many other methods for modeling latent dynamics, such as NeuralODEs [1] or neural spectral methods [2]. It would be good to have an ablation study on the latent dynamics part. \n\n**Refs:**\n\n[1] Krishnapriyan, Aditi S., et al. \"Learning continuous models for continuous physics.\" Communications Physics 6.1 (2023): 319.\n\n[2] Wu, Haixu, et al. \"Solving high-dimensional pdes with latent spectral models.\" arXiv preprint arXiv:2301.12664 (2023).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "* **W1** | The motivation of using a VAE-type encoder-decoder is not well justified. \n\nThank you for your insightful comment. We chose to use a Variational Autoencoder (VAE) over a standard Autoencoder (AE) due to the VAE’s established reputation in the literature for producing compact latent representations while maintaining smoothness in the latent space. However, we acknowledge that there are various autoencoder types that employ different forms of regularization, and alternative methods could also be effective.\n\nWe agree that enforcing regularization is crucial for obtaining a smooth latent space, whether through a variational approach, L2 regularization, or spectral regularization. Without such regularization, we observed that the latent space can exhibit increasing variance, leading to less stable results.\n\nIn the supplementary pdf, Table 1 presents an ablation study comparing the VAE and AE across three datasets. This study demonstrates that an autoencoder with L2 regularization can be a viable alternative to the VAE for achieving a smooth latent space in some cases. The AE showed superior performance in terms of lower reconstruction errors for Burgers and Navier Stokes 1e-4, which translated into better rollout performance. However, for the more challenging Navier-Stokes 1e-5 dataset, the AE's latent space had high variance, which may account for the observed performance difference compared to the VAE.\n\nWe appreciate your feedback, which enriches our analysis, and hope this explanation clarifies our approach and findings.\n\n* **W2** |  This paper considers a diffusion-based transformer for learning latent dynamics. \n\nThank you for your feedback. In Table 1 of the supplementary material, we compare three versions of our model: AROMA with $K=3$ diffusion steps (as detailed in the paper), AROMA without diffusion, and a local MLP that disregards token relationships. As expected, the MLP, which ignores token interactions, performs poorly, while both the deterministic and diffusion transformers produce comparable results.\nIt's important to note that the diffusion-based approach has the added advantage of generating a distribution, which opens the door to uncertainty modeling. This is particularly relevant given that the state-of-the-art in video prediction often relies on probabilistic models. To fully capture the benefits of diffusion models, it may be necessary to evaluate them using metrics beyond MSE, as these could better reflect their ability to model uncertainty. This is an area that warrants further exploration in our future work.\n\nThank you for suggesting alternatives like NeuralODEs and neural spectral methods. We appreciate the opportunity to discuss our approach in this context.  We would like to emphasize that our chosen latent processor is particularly effective at capturing both local spatial information within tokens and global spatial information through token interactions, which is essential for accurately modeling global dynamics in our framework. In our study, the MLP is implemented with residual connections as described in [1], and while it has a structure similar to NeuralODEs, it lacks the ability to account for interactions between tokens, which is a key limitation.\n\nSimilarly, the model proposed in [2] processes tokens independently at a given resolution using neural spectral blocks. Given this design, we expect that these models would exhibit comparable performance when applied to AROMA's latent tokens. However, the integration of token interactions in our approach is a critical factor that we believe enhances the modeling of complex dynamics.\n\n[1] Serrano et al. Operator Learning with Neural Fields: Tackling PDEs on General Geometries.\n\n[2] Krishnapriyan, Aditi S., et al. \"Learning continuous models for continuous physics.\" Communications Physics 6.1 (2023): 319.\n\n* **Q1** |  In lines 90-91, for the second stage of training, do you fix the encoder-decoder ?\n\nYes, we keep the encoder and decoder frozen during the second stage. \n\n* **Q2** |  The authors claim that the encoder-decoder is principled.\n\nWe appreciate the feedback and agree that the wording may have been misleading. Rather than suggesting a theoretical guarantee, our intent was to emphasize that the encode-process-decode framework is a well-established and widely adopted approach in the research community. This method has proven effective in a variety of applications, and our use of it follows this standard practice.\n\n* **Q3** |  In Table 1, why does DNOT work worse than FNO on the 1D Burgers case? \n\nWe took the recommended default parameters proposed in the publications and observed that GNOT was indeed less robust in dynamics modeling. In the 1D case, FNO has $\\sim$ 4 million parameters while GNOT has about 0.8 million, this could explain the difference. As for FNO, we agree that it is a strong and robust baseline, however it is limited to regular grids and it cannot process point clouds as done here. Enhanced FNO versions have been developed for irregular meshes, but they are much less flexible than the method proposed here.\n\n* **Q4** | FNO can also be modified in an auto-regressive way.\n\nIndeed, prior research has explored and highlighted the temporal extrapolation capabilities of the Fourier Neural Operator (FNO). Given FNO's established performance, we anticipate that it will demonstrate comparable extrapolation effectiveness on our datasets as the AROMA model for regular grids only. \n\nIn the first experimental section, we concentrated on dynamics over regular grids, where we implemented and trained FNO in an auto-regressive manner. In the second section, our focus shifted to irregular grids, where we aimed to compare our model against the most relevant baselines, specifically neural-field-based methods and transformers. To provide a more comprehensive evaluation, we plan to include the performance results of FNO and UNet on the regular-grid case $\\pi = 100$ % in the camera-ready.", "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:
Dear Reviewers,

Thank you for your insightful feedback.

In response to your comments, we have addressed the key concerns by providing additional experimental results. Please refer to the supplementary PDF, specifically:

* An ablation study in Table 1 that compares different processing blocks  (Diffusion (K=3 steps) vs deterministic Transformer vs MLP) for AROMA on Burgers, Navier-Stokes 1e-4, and Navier Stokes 1e-5. The deterministic transformer and diffusion transformer share the same architecture, while the MLP is implemented with residual connections as in [1].
* An ablation study, also in Table 1, on the impact of the choice of encoder and decoder framework (Autoencoder with L2 regularization vs VAE) for AROMA on Burgers, Navier-Stokes 1e-4, and Navier Stokes 1e-5.
* Additional results on KS equation in Table 2 and Figure 1 of the pdf page, with a comparison between the diffusion and deterministic versions of AROMA.
* Visual comparisons between AROMA’s predictions and the ground truth on the Cylinder Flow dataset in Figure 2 of the pdf page.

From these new results, we can draw the following conclusions:

* Processing Blocks: Modeling  interactions at the local and global levels is key to learn the dynamics faithfully. Experiments using MLPs (table 1) as time steppers which do not consider interactions between tokens lead to significantly lower performance compared to transformers.

* Deterministic vs. Diffusion: The deterministic version of AROMA shows consistently robust performance and even surpasses the diffusion version on the Navier-Stokes 1e-4 case. This demonstrates that the latent tokens obtained with AROMA contain meaningful information for dynamics modeling. On the other hand, the deterministic version yields less accurate long rollouts on Burgers or KS. Note that using diffusion allows us to model the trajectory distribution, which opens the way to infer statistics on this distribution. This is key, for example, when modeling uncertainty, which is a critical problem for these models.

* Encoder-Decoder Frameworks: Using an autoencoder with L2 regularization is a viable alternative to the VAE for achieving a smooth latent space. The autoencoder demonstrated superior performance on two datasets, explained by its lower reconstruction errors, which translate into better rollout performance. However, for the more challenging Navier-Stokes 1e-5 case, the autoencoder's latent space exhibited high variance, which may explain the observed performance difference with the VAE.
* KS Equation Limitations: AROMA currently struggles with dynamics that exhibit chaotic phenomena and non-decaying spectra, as shown by the KS equation results. The primary limitation appears to be the reconstruction capabilities of the encoder-decoder. For the KS equation, we found that obtaining reconstructions with an MSE in the range of 1e-10 to 1e-12 was necessary for accurate spectrum reconstruction. Like all models leveraging a reduced latent representation space, AROMA inherently lose some of the fine-grained details necessary for accurately capturing chaotic behavior. In conclusion, while AROMA performs very well on simpler dynamics, dealing with chaotic phenomena requires more involved modeling that explicitly targets the chaotic component. Note that using dedicated modules for this purpose is a current practice in fluid dynamics - e.g. LES (Large eddy simulation).

We hope these additional results and insights address your concerns and enhance the understanding of AROMA’s capabilities and limitations.

[1] Serrano et al. Operator Learning with Neural Fields: Tackling PDEs on General Geometries.

Meta-review: This paper received initially somewhat mixed reviewed.  The authors submitted a rebuttal with some additional experimental results and there was an active discussion period with many of the issues noted in the initial reviews being addressed.  After discussion the reviewers are now consistently positive.

While there was some discussion around novelty and scalability, all reviewers believed the work warranted publication and the AC agrees.
Decision: Accept (poster)