Title: AROMA: Preserving Spatial Structure for Latent PDE Modeling with Local Neural Fields
Venue: NeurIPS 2024
Accepted: yes
Keywords: PDE

Neural Operator

Neural Fields

Transformer

Diffusion
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

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. Github page : https://github.com/LouisSerrano/aroma

1 Introduction

In recent years, many deep learning (DL) surrogate models have been introduced to approximate solutions to partial differential equations (PDEs) (Lu et al., 2021; Li et al., 2021; Brandstetter et al., 2022; Stachenfeld et al., 2022). Among these, the family of neural operators has been extensively adopted and tested across various scientific domains, demonstrating the potential of data-centric DL models in science (Pathak et al., 2022; Vinuesa & Brunton, 2022).

Neural Operators were initially constrained by discretization and domain geometry limitations. Recent 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 flexible, are computationally expensive due to their operation in the original physical space and require large training datasets.

Our hypothesis is that considering spatiality is essential in 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 or 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.

1,2

Our novel framework leverages attention blocks and neural fields, resulting in a model that is easy to train and achieves state-of-the-art results on most datasets, particularly for complex geometries, without requiring prior feature engineering. To the best of our knowledge, we are the first to propose a fully attention-based architecture for processing domain geometries and unrolling dynamics. Compared to existing transformer architectures for PDEs, our framework first encapsulates the domain geometry and observation values in a compact latent representation, efficiently forecasting the dynamics at a lower computational cost. Transformer-based methods such as (Li et al., 2023; Hao et al., 2023) unroll the dynamics in the original space, leading to high complexity.

Our contributions are summarized as follows:

Aprincipled and versatile encode-process-decode framework for solving PDEs that operate on general input 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 a 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 stability and forecasting accuracy.

Experiments performed on representative spatio-temporal forecasting problems demonstrate that AROMA is on par with or outperforms state-of-the-art 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 Ω (with boundary ∂ Ω ) and temporal domain [0 , T ] . In the general form, their solutions u ( x, t ) satisfy the following constraints :

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 at the next time step: u t +∆ t = G ( u t ) . We have access to training trajectories obtained with different initial conditions, and we 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 m ∆ t = G m ( u 0 ) .

3 Model Description

3.1 Model overview

We provide below an overview of the global framework and each component is described in a subsequent section. The model comprises three key components, as detailed in Figure 1.

Encoder E w : u t X → Z t . The encoder takes input values u t X sampled over the domain Ω at time t , where X denotes the discrete sample space and could be a grid, an irregular mesh or a point set. u t X is observed at locations x = ( x 1 , . . . x N ), with values u t = ( u t ( x 1 ) , · · · , u t ( x N )) . N is the number of observations and can vary across samples. u t X is projected through a cross attention mechanism onto a set of M tokens Z t = ( z t 1 , · · · , z t M ) with M a fixed parameter. This allows

coordinate query

predicted value

SA

CA

Decoder

MLP

Latent

Refiner

Diffusion

Transformer

latent tokens with a

compressed dimension

predicted tokens

...

Encode

geometry

Q

coordinates

function values

K

Aggregate

function

values

...

V

V

K

Q

Encoder

learnable tokens

geometry-aware tokens

mapping any discretized input u t X onto a fixed dimensional latent representation Z t encoding implicit local spatial information from the input domain. The encoder is trained as a V AE and Z t is sampled from a multivariate normal statistics as detailed in Section 3.2.

Latent time-marching refiner R θ : Z t → ˆ Z t +∆ t . We model the dynamics in the latent space through a transformer. The dynamics can be unrolled auto-regressively in the latent space for any time horizon without requiring to project back in the original domain Ω . Self-attention operates on the latent tokens, which allows modeling global spatial relations between the local token representations. 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 was more robust than a baseline deterministic transformer for temporal extrapolation.

Decoder D ψ : ˆ Z t +∆ t → ˆ u t +∆ t . The decoder uses the latent tokens ˆ Z t +∆ t to approximate the function value ˆ u t +∆ t ( x ) = D ψ ( x, ˆ Z t +∆ t ) for any query coordinate x ∈ Ω . We therefore denote ˆ u t +∆ t = D ψ ( Z t +∆ t ) the predicted function.

Inference We encode the initial condition and unroll the dynamics in the latent space by successive denoisings: ˆ u m ∆ t = D ψ ◦ R m θ ◦ E w ( u 0 ) . 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 a two-stage training: we 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 a fixed number of tokens. The decoder can be queried at any position in the spatial domain, irrespective of the input sample.

Encoder The encoder maps an arbitrary number N of observations ( x , u ( x )) := (( x 1 , u ( x 1 )) , . . . , ( x N , u ( x N )) onto a latent representation Z of fixed size M through the following series of transformations:

where ( γ ( x ) , v ( x )) = (( γ ( x 1 ) , v ( x 1 )) , . . . , ( γ ( x N ) , v ( x N ))) , and h ≪ d .

(i) Embed positions and observations : Given an input sequence of coordinate-value pairs ( x 1 , u ( x 1 )) , . . . , ( x N , u ( x N )) , we construct sequences of positional embeddings γ = ( γ ( x 1 ) , . . . , γ ( x N )) and value embeddings v = ( v ( x 1 ) , . . . , v ( x N )) , where γ ( x ) = FourierFeatures ( x ; ω ) and v ( x ) = Linear ( u ( x )) , with ω a fixed set of frequencies. These embeddings are aggregated onto a smaller set of learnable query tokens T = ( T 1 , . . . , T M ) and then T ′ = ( T ′ 1 , . . . , T ′ M ) with M fixed, to compress the information and encode the geometry and spatial latent representations.

(ii) Encode geometry : Geometry-aware tokens T are obtained with a multihead cross-attention layer and a feedforward network (FFN), expressed as T geo = T + FFN ( CrossAttention ( Q = W Q T , K = W K γ , V = W V γ )) . This step does not include information on the observations, ensuring that similar geometries yield similar query tokens T geo irrespective of the u values.

(iii) Encode observations : The T geo tokens are then used to aggregate the observation values via a cross-attention mechanism: T obs = T geo + FFN ( CrossAttention ( Q = W ′ Q T geo , K = W ′ K γ , V = W ′ V v )) . 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 a penalty on the L 2 norm of the latent code ∥ Z ∥ 2 . 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 a Gaussian multivariate distribution µ = Linear ( T obs ) and log( σ ) = Linear ( T obs ) from which the final token embedding Z is sampled.

Decoder The decoder's role is to reconstruct ˆ u t +∆ t from ˆ Z t +∆ t , see Figure 1. Since training is performed in two steps ('encode-decode' first and then 'process'), the decoder is trained to reconstruct ˆ u t for input u t . 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 that depend on the query coordinate x , ( f u q ( x )) = CrossAttention ( Q = W Q ( γ q ( x )) , K =

W K Z ′ , V = W V Z ′ ) , where γ q is a Fourier features embedding of bandwidth ω q . (iii) Decode with MLP : Finally, we use a small MLP to decode this feature vector and obtain the reconstruction ˆ u ( x ) = MLP ( f u q ( x )) . In contrast with existing neural field methods for dynamics modeling, the feature vector here is local. In practice, one uses multiple cross attentions to get feature vectors with different frequencies (see Appendix Figures 7 and 8 for further details).

Training The encoder and decoder are jointly optimized as a variational autoencoder (VAE) (Kingma & Welling, 2013) to minimize the following objective : L = L recon + β · L KL ; where L recon = MSE ( u t X , ˆ u t X ) is the reconstruction loss between the input and the reconstruction D ψ ( Z t , X ) on the grid X , with Z t ∼ N ( µ t , ( σ t ) 2 ) and µ t , σ t = E w ( u t X ) . The KL divergence loss L KL = D KL ( N ( µ t , ( σ t ) 2 ) || N (0 , I )) helps regularize the network and prevents overfitting. We found that using a variational formulation was essential to obtain smooth latent representations while training the encoder-decoder.

...

T1

T1

T3

T2

T3

TM

M

3.3 Transformer-based diffusion

Modeling the dynamics 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 reduced 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 absolute positional embeddings E pos and a linear layer to project onto a higher dimensional space: Z [0] = Linear ( Z ) + E pos. The backbone then applies several self-attention blocks, which process tokens as follows:

We found out that adding a diffusion component to the transformer helped enhance the stability and allowed longer forecasts. Diffusion steps are inserted between two time steps t and t + ∆ t 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 k are performed between two timemarching steps t and t +∆ t ).

We then use a conditional diffusion transformer architecture close to Peebles & Xie (2023) for R θ , where we detail the main block in Appendix B. At diffusion step k , the input to the network is a sequence stacking the tokens at time t and the current noisy targets estimate ( Z t , ˜ Z t +∆ t k ) . See Appendix B, Figure 4 and Figure 5 for more details. To train the diffusion transformer R θ , we freeze the encoder and decoder, and use the encoder to sample pairs of successive latent tokens ( Z t , Z t +∆ t ) . We employ the 'v-predict' formulation of DDPM (Salimans & Ho, 2022) for training and sampling.

4 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 on regular grids, where we benchmark our model against stateof-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 spatiotemporal 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 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 a wide range of geometric settings. We include additional results from ablation studies in Appendix C.6.

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 as in Brandstetter et 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: ∂w ∂t + u · ∇ w = ν ∆ w + f , ∇ u = 0 for x ∈ Ω , t > 0 , where ν is the viscosity coefficient. We consider two different versions ν = 10 -4 ( Navier-Stokes 1 × 10 -4 ) and ν = 10 -5 ( Navier-Stokes 1 × 10 -5 ), and use train and test sets of 1000 and 200 trajectories with a base spatial resolution of size 64 × 64 . We consider a horizon of T = 30 for ν = 10 -4 and T = 20 for ν = 10 -5 since the phenomenon is more turbulent. At test time, we use the vorticity at t 0 = 10 as the initial condition.

Setting We train all the models with supervision on the next state prediction to learn to approximate the time-stepping operator u t +∆ t = G ( u t ) . At test time, we unroll the dynamics auto-regressively with each model and evaluate the prediction with a relative L 2 error defined as

Baselines We use a diverse panel of baselines including state of the art regular-grid methods such as FNO (Li et 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 neuralfield based methods with DINO (Yin et al., 2022) and CORAL (Serrano et al., 2023).

Results Table 1 presents a comparison of model performance on the Burgers , Navier-Stokes1e-4 , and Navier-Stokes1e-5 datasets, with metrics reported in Relative L 2 . Our method, AROMA, demonstrates excellent performance across the board, highlighting its ability to capture the dynamics of turbulent phenomena, as reflected in the Navier-Stokes datasets.

In contrast, DINO and CORAL, both global neural field models, perform poorly in capturing turbulent phenomena, exhibiting significantly higher errors 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.

Regarding the regular-grid methods, ResNet shows suboptimal performance in the pure teacher forcing setting, rapidly accumulating errors over time during inference. FNO stands out as the best baseline, demonstrating competitive performance on all datasets. We hypothesize that FNO's robustness to error accumulation during the rollout can be attributed to its Fourier block, which effectively cuts off high-frequency components. Overall, the results underscore AROMA's effectiveness and highlight the challenges Neural Field-based models face in accurately modeling complex phenomena.

4.2 Dynamics on 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 a 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 a 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 1 × 10 -3 ): We use the same equation as

in Section 4.1 but with a higher viscosity coefficient ν = 1 e -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 h of the fluid. The training set comprises 64 trajectories of size 40, and the test set comprises 8 trajectories with 40 timestamps. We use a standard spatial resolution of 64 × 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.e. for 39 steps. Evaluation within the In-t horizon assesses the model's ability to forecast within the training regime. The Out-t evaluation tests the model's extrapolation capabilities beyond the training horizon. · Sparse observations : For the train and test set we randomly select π percent of 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 is different, they maintain the same level of sparsity across trajectories. In our case, π = 100% amounts to the fully observable case, while in π = 25% each grid contains around 1020 points for Navier-Stokes 1 × 10 -3 and 2040 points for Shallow-Water .

Baselines We compare our model to OFormer (Li et 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-t ) and for temporal extrapolation ( Out-t ).

Results Table 2 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-Water our model is slightly outperformed by CORAL in the fully observed regime, potentially because of a lack of data. Similarly, on Navier-Stokes 1 × 10 -3 CORALhasslightly 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-encoding block is crucial for obtaining good representations of the observed values in the sparse regimes, potentially explaining the performance gap with GNOT and OFormer.

4.3 Dynamics on different geometries

Finally, we extend our analysis to learning dynamics over varying geometries.

Datasets We evaluate our model on two problems involving non-convex domains, as 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 different, 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 ( j ) , momentum w ( x ( j ) ) , and pressure p ( x ( j ) ) . Our task is to learn the mapping from ( w t ( x ) , p t ( x )) x ∈X to ( w t +∆ t ( x ) , p t +∆ t ( x )) x ∈X for a fixed ∆ t . · Airfoil ( AirfoilFlow ): 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 ρ ( x ( j ) ) for each node j . Our goal is to learn the mapping from ( w t ( x ) , p t ( x ) , ρ t ( x )) x ∈X to ( w t +∆ t ( x ) , p t +∆ t ( x ) , ρ t +∆ t ( x )) x ∈X . Each example in the dataset corresponds to a unique mesh. On average, there are 5233 nodes per mesh for AirfoilFlow and 1885 for CylinderFlow . We temporally 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-t ) and keep the last 20 timestamps for evaluation ( Out-t ).

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 a mean squared error (MSE) both in the training horizon (In-t) and beyond the training horizon (Out-t) .

Results The results in Table 3 show that AROMA outperforms other models in predicting flow dynamics 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. Additionally, AROMA maintains stability over extended prediction horizons, as evidenced by its consistently low Out-t MSE values.

4.4 Long rollouts and uncertainty quantification

After training different models on Burgers , we compare them on long trajectory rollouts. We start from t 0 = 50 (i.e. use a numerical solver for 50 steps), and unroll our dynamics auto-regressively for 200 steps. Note that all the models were only trained to predict the next state. We plot the correlation over rollout steps of different methods, including our model without the diffusion process, in Figure 3. We can clearly see the gain in stability in using the diffusion for long rollouts. Still, the predictions will eventually become uncorrelated over time as the solver accumulates errors compared with the numerical solution. As we employ a generative model, we can generate several rollouts and estimate the uncertainty of the solver with standard deviations. We can see in Appendix Figure 11 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.

0

25

50

75

100

125

150

175

200

Rollout Steps

0.5

0.6

0.7

0.8

0.9

1.0

Correlation

FNO

ResNet

AROMA

AROMA No Diffusion

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 to model dynamical systems. DINO (Yin et al., 2022) is a space-time continuous architecture based 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-field based architecture, which tackles the broader scope of operator learning, also builds on meta-learning (Zintgraf et 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 a temporal horizon in the latter. Moreover, they either use a 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 a backbone for modeling PDEs. Transformers, which operate on token sequences, provide a natural solution for handling irregular meshes and point sets. Li et al. (2023) and Hao et al. (2023) introduced transformer architectures tailored for operator learning. Hao et al. (2023) incorporated an attention mechanism and employed a 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 cross-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 a mixture model, with attention operating in this latent space. This approach, akin to our model, reduces complexity. However, it has not been designed for temporal problems. (Alkin et al., 2024) recently proposed a versatile model capable of operating on Eulerian and Lagrangian (particles) representations. They reduce input dimensionality by aggregating information from input values onto 'supernodes' selected from the input mesh via message passing 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 values. Finally, their training involves complex end-to-end optimization, whereas we favor two simple training steps that are easier to implement.

6 Conclusion and Limitations

AROMAoffers a novel and flexible neural operator approach for modeling the spatio-temporal evolution of physical processes. It is able to deal with general geometries and to forecast at any position of the spatial domain. It incorporates in an encode-process-decode framework attention mechanisms, a 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. Experiments 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 operate 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 a high flexibility, the training efficiency does not match that of CNN-based autoencoders 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.

References

Benedikt Alkin, Andreas F¨ urst, Simon Schmid, Lukas Gruber, Markus Holzleitner, and Johannes Brandstetter. Universal Physics Transformers. 2024. URL http://arxiv.org/abs/2402. 12365 .

Matthias Bauer, Emilien Dupont, Andy Brock, Dan Rosenbaum, Jonathan Schwarz, and Hyunjik Kim. Spatial functa: Scaling functa to imagenet classification and generation. CoRR , abs/2302.03130, 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 , 6 2022. URL http://arxiv.org/abs/2206.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.1109/CVPR.2019.00609.

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 on Machine Learning , 1 2022. URL http://arxiv.org/abs/ 2201.12204 .

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 Conditional 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. In International Conference on Machine Learning , pp. 12556-12569. PMLR, 2023.

Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. 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-6851, 2020.

Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114 , 2013.

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.

Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. International Conference on Learning Representations. , 10 2021. URL http://arxiv.org/abs/2010.08895 .

Phillip Lippe, Bastiaan S Veeling, Paris Perdikaris, Richard E Turner, and Johannes Brandstetter. Pde-refiner: Achieving accurate long rollouts with neural pde solvers. arXiv preprint arXiv:2308.05732 , 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 Intell , 3:218-229, 10 2021. doi: 10.1038/s42256-021-00302-5. URL http://arxiv. org/abs/1910.03193http://dx.doi.org/10.1038/s42256-021-00302-5 .

Aaron van den Oord, Oriol Vinyals, and Koray Kavukcuoglu. Neural discrete representation learning. 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.1109/CVPR.2019.00025.

Jaideep Pathak, Shashank Subramanian, Peter Harrington, Sanjeev Raja, Ashesh Chattopadhyay, Morteza Mardani, Thorsten Kurth, David Hall, Zongyi Li, Kamyar Azizzadenesheli, et al. Fourcastnet: A global data-driven high-resolution weather model using adaptive fourier neural operators. arXiv preprint arXiv:2202.11214 , 2022.

William Peebles and Saining Xie. Scalable diffusion 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 meshbased simulation with graph networks. International Conference on Learning Representations. , 10 2021. URL http://arxiv.org/abs/2010.03409 .

Michal Rolinek, Dominik Zietlow, and Georg Martius. Variational autoencoders pursue pca directions (by accident). In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pp. 12406-12415, 2019.

Salva R¨ uhling Cachay, Bo Zhao, Hailey Joren, and Rose Yu. DYffusion: a 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.00512 , 2022.

Louis Serrano, Lise Le Boudec, Armand Kassa¨ ı Koupa¨ ı, Thomas X Wang, Yuan Yin, Jean-No¨ el 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 Sanchez-Gonzalez. Learned coarse models for efficient turbulence simulation. International Conference on Learning Representation , 2022.

Ricardo Vinuesa and Steven L Brunton. Enhancing computational fluid dynamics with machine learning. Nature Computational Science , 2(6):358-366, 2022.

Haixu Wu, Huakun Luo, Haowen Wang, Jianmin Wang, and Mingsheng Long. Transolver: A Fast Transformer Solver for PDEs on General Geometries. 2024. URL http://arxiv.org/abs/ 2402.02366 .

Yuan Yin, Matthieu Kirchmeyer, Jean-Yves Franceschi, Alain Rakotomamonjy, and Patrick Gallinari. Continuous pde dynamics forecasting with implicit neural representations. International Conference on Learning Representations , 9 2022. URL http://arxiv.org/abs/2209. 14855 .

Luisa Zintgraf, Kyriacos Shiarlis, Vitaly Kurin, Katja Hofmann, and Shimon Whiteson. Fast context adaptation via meta-learning. 36th International Conference on 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 to physics prediction. R¨ uhling Cachay et al. (2023) propose DYffusion, a framework that adapts the diffusion process to spatio-temporal data for forecasting on long-time rollouts, by performing diffusion-like timesteps in the physical time dimension. PDE-Refiner (Lippe et al., 2023) is a CNN-based method that uses diffusion to stabilize prediction rollouts over long trajectories. Compared to these methods, we perform diffusion in a latent space, reducing the computational cost; and leverage the advanced modeling capabilities of transformers.

Local Neural Fields We are not the first work that proposes to leverage locality to improve the design of neural fields. In a different approach, Bauer et al. (2023) proposed a grid-based latent space where the modulation function ϕ is dependent on the query coordinate x . This concept enables the application of architectures with spatial inductive biases for generation on the latent representations, such as a U-Net Denoiser for diffusion processes. Similarly, Lee et al. (2023) developed a locality-aware, generalizable Implicit Neural Representation (INR) with demonstrated capabilities in generative modeling. Both of these architectures assume regular input structures, be it through patching methods or grid-based layouts.

B Implementation details

Diffusion transformer Weillustrate how our diffusion transformer is trained and used at inference in Figure 4 and Figure 5. We provide the diffusion step k which acts as a conditioning input for the diffusion model. We use an exponential decrease for the noise level as in Lippe et al. (2023) i.e. α k = 1 -σ k/K min . We use the same diffusion transformer block as in Peebles & Xie (2023), which relies on amplitude and shift modulations from the diffusion timestamp k :

...

...

DiT Block

x H

conditioning embedding

...

...

Linear

Linear

output tokens

...

MSE

target tokens

noisy estimate

previous latent token

diffusion step

...

...

...

output tokens

DENOISE

noisy estimate

previous latent token

diffusion step

Encoder-Decoder We provide a more detailed description of the encoder-decoder pipeline in Figure 6.

Local INR We show the implementation of our local INR, both with single-band frequency and multi-band frequency, in Figure 7 and Figure 8. The cross-attention mechanism enables to retrieve a local feature vector f q ( x ) 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.

...

Encode

geometry

Q

...

...

Linear

Linear

...

multivariate normal parameters

sampling

Self-

Attention

latent tokens

x L

Cross-

Attention

K V

Q

coordinate query

modulation vector

...

...

...

MLP

output value

feature vector

frequency embedding

coordinates

function values

K

Aggregate

function

values

...

V

V

K

Q

learnable tokens

encoding prior spatial

information

geometry-aware tokens to

query depending

on the geometry

spatial representation

of the signal

with a compressed

channel dimensions

...

Cross-Attention

MLP

output value

feature vector

frequency embedding

K V

Q

...

Cross-Attention

MLP

output value

multi band

feature vectors

multi band

frequency embeddings

K V

Q

Q

B.1 Hyperparameters

We detail the values of the hyperparameters used on each dataset: Table 5 presents the hyperparameters of the Encoder-Decoder, while Table 4 presents the hyperparameters of the Diffusion Transformer. We use a cosine scheduler for the tuning learning rate for both trainings, with an initial maximum learning rate of 10 -3 annealing to 10 -5 . All experiments were performed with an NVIDIA TITAN RTX.

For the diffusion transformer, we use K = 3 diffusion steps for all experiments and only vary the minimum noise σ min.

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 latents: The number M of latent tokens used to spatially project the objervations and geometries.

latent dim: The dimension c 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 a Self Attention layer.

cross heads: The number of heads use for the Cross Attention layers.

cross dim head: The dimension of each head in a Cross Attention layer.

dim: The number of neurons used in the MLP decoder.

depth inr: The number of layers in the MLP decoder.

frequencies: The different frequencies used for the local INR. We use base 2 for all experiments and select 16 frequencies in logarithmic scale per level. For example, [3, 4, 5] means that we construct 3 frequency embedding vectors, the first γ 1 ( x ) = (cos(2 0 πx ) , sin(2 0 πx ) , . . . , cos(2 3 πx ) , sin(2 3 πx )) , for the second γ 2 = (cos(2 3 πx ) , sin(2 3 πx ) , . . . , cos(2 4 πx ) , sin(2 4 πx )) , and for the third γ 3 = (cos(2 4 πx ) , sin(2 4 πx ) , . . . , cos(2 5 πx ) , sin(2 5 πx ))

dropout sequence: The ratio of points that are ignored by the encoder.

feature dim: The dimension of the feature vector.

encode geo: If we use a cross-attention block to encode the geometry.

max encoding freq: The maximum frequency used for the frequency embedding γ of the encoder.

kl weight: The weight β used for the V AE 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.

C Additional results

C.1 Time complexity analysis

We denote N as the number of observations of 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 d as the number of channels used in the attention mechanism. The most computationally expensive operations in our architecture are the cross-attention and self-attention blocks. For simplification, we omit the geometry encoding block in this study.

The cost of the cross-attention in the encoder is O ( NMd ) , and similarly, the cost of the crossattention in the decoder is O ( NMd ) . Let L 1 and L 2 represent the number of layers in the decoder and diffusion transformer, respectively. The cost of the self-attention layers in the decoder is O ( L 1 M 2 d ) , while in the diffusion transformer, it is O (4 L 2 M 2 d ) .

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 ((2 N +4 KTL 2 M + L 1 M ) 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 in a time complexity of O (( LN ) d 2 ) for each prediction, where L is the depth of the network. At inference, the cost per step along a trajectory is LNd 2 for GNOT, compared to 4 KL 2 M 2 d for AROMA.

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 is more efficient when M ≪ N .

C.2 Encoding interpretation

We provide in Figure 9 a qualitative analysis through cross-attention visualizations how the geometry encoding block helps to capture the geometry of the domain. In the first cross-attention block, the query tokens T are not aware of the geometry and therefore attend to large regions of the domains. This lets the model understand, where the boundaries of the domain are and therefore where the cylinder is. Once the query tokens have aggregated the mesh information, the cross attention between T geo and the positions are sharper and depend on the geometry.

C.3 Example rollouts

Weshow examples of rollout predictions using AROMA on Burgers dataset in Figure 10, on NavierStokes 1 × 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.

Cross-Attention on

T

Cross-Attention on

T

geo

Can be interpreted as a prior spatial

representation of the different domains. The query

tokens are learnt to attend to the « important

»

regions of the domains.

Token 22

Token 23

Geometry

encoding

Geometry

encoding

aggregates information of the different

positions in the domain and therefore reflects its

geometry. Its cross-attention maps are sharper

and adapted to the geometry.

T

geo

·

0

20

40

60

80

100

2

1

0

1

2

Long rollout prediction

0

20

40

60

80

100

gt

0

25

50

75

100

125

150

175

200

Timestamp

0

25

50

75

100

125

150

175

200

Timestamp

0

20

40

60

80

100

Index

3

2

1

0

1

2

3

Value

100 rollouts

0

20

40

60

80

100

Index

200 rollouts

0

20

40

60

80

100

Index

400 rollouts

Mean

For Navier-Stokes , we show an example of test trajectory in the training horizon (Figure 12a) and in extrapolation (Figure 12b).

pred

t1

t2

t3

t5

t6

t7

t8

t9

t10

t11

t12

t13

t14

t15

t16

t17

t18

t19

t20

gt

(a) In-t

pred

t20

t21

t22

t23

t24

t25

t26

t27

t28

t29

t30

t31

t32

t33

t34

t35

t36

t37

t38

t39

gt

(b) Out-t

t=40

Predictions

Ground Truth

Delta

t=41

t=42

t=44

C.4 Scaling experiments

In Figure 14, we compare the reconstruction and prediction capabilities of CORAL and AROMA on Navier-Stokes 1 × 10 -4 given the number of training trajectories. As evidenced, our architecture outperforms CORAL significantly when the number of trajectories is greater than 10 3 , highlighting its efficacy in handling large amounts of data..

10 1

10 2

10 3

10 4

Number of Training Samples

0.00

0.05

0.10

0.15

0.20

0.25

0.30

Relative Error

Test reconstruction error

CORAL

AROMA

(a) Step 1: Autoencoding

10 1

10 2

10 3

10 4

Number of Training Samples

0.1

0.2

0.3

0.4

0.5

0.6

Relative Error

Test prediction error

CORAL

AROMA

(b) Step 2: Rollout prediction

C.5 Spatial tokens perturbation analysis

To validate the spatial interpretation of our latent tokens, we establish a baseline code Z 0 , and introduce perturbations by sequentially replacing the j -th token, z 0 j , with subsequent tokens along the trajectory, denoted as z 1 j , z 2 j , . . . , z t j . Thus, the perturbed tokens mirror Z 0 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 of our tokens. Additionally, we can notice some interesting effects of the token perturbations near the boundaries in figs. 17 and 22: our encoderdecoder 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-aware tokens.

pred

0

20

40

60

80

100

2.0

1.5

1.0

0.5

0.0

0.5

1.0

1.5

0

20

40

60

80

100

0

25

50

75

100

125

150

175

200

Timestamp

0

25

50

75

100

125

150

175

200

Timestamp

gt

0

20

40

60

80

100

2.0

1.5

1.0

0.5

0.0

0.5

1.0

1.5

pred

0

20

40

60

80

100

gt

0

25

50

75

100

125

150

175

200

Timestamp

0

25

50

75

100

125

150

175

200

Timestamp

0

20

40

60

80

100

2.0

1.5

1.0

0.5

0.0

0.5

1.0

1.5

pred

0

20

40

60

80

100

gt

0

25

50

75

100

125

150

175

200

Timestamp

0

25

50

75

100

125

150

175

200

Timestamp

0

20

40

60

80

100

2.0

1.5

1.0

0.5

0.0

0.5

1.0

1.5

pred

0

20

40

60

80

100

gt

0

25

50

75

100

125

150

175

200

Timestamp

0

25

50

75

100

125

150

175

200

Timestamp

0

20

40

60

80

100

2.0

1.5

1.0

0.5

0.0

0.5

1.0

1.5

pred

0

20

40

60

80

100

gt

0

25

50

75

100

125

150

175

200

Timestamp

0

25

50

75

100

125

150

175

200

Timestamp

0

20

40

60

80

100

2.0

1.5

1.0

0.5

0.0

0.5

1.0

1.5

pred

0

20

40

60

80

100

gt

0

25

50

75

100

125

150

175

200

Timestamp

0

25

50

75

100

125

150

175

200

Timestamp

0

20

40

60

80

100

2.0

1.5

1.0

0.5

0.0

0.5

1.0

1.5

pred

0

20

40

60

80

100

gt

0

25

50

75

100

125

150

175

200

Timestamp

0

25

50

75

100

125

150

175

200

Timestamp

0

20

40

60

80

100

2.0

1.5

1.0

0.5

0.0

0.5

1.0

1.5

pred

0

20

40

60

80

100

gt

0

25

50

75

100

125

150

175

200

Timestamp

0

25

50

75

100

125

150

175

200

Timestamp

C.6 Ablation studies

Number of tokens M We show the impact of the number of latent tokens on the Navier-Stokes 1 × 10 -4 dataset in Table 6. We train our auto-encoder with 10000 trajectories. We can see that the performance increases with the number of tokens.

Auto-encoding vs VAE Our framework can also be used without the KL regularization, and could potentially be employed with other forms of regularization, such as L2 regularizaton or vectorquantization (Oord et al., 2017). We investigated in Table 7 the impact the KL regularization had on the overall rollout performance, and selected an autoencoder with L2 regularization (weight decay) as baseline. Our conclusion is that using an autoencoder with L2 regularization is a viable alternative to the V AE in some cases for achieving a smooth latent space. The autoencoder demonstrated superior performance on two datasets ( Burgers and Navier-Stokes 1 × 10 -4 , explained by its lower reconstruction errors, which translate into better rollout performance. However, for the more challenging Navier-Stokes 1 × 10 -5 case, the autoencoder's latent space exhibited high variance, which may explain the observed performance difference with the VAE.

No-diffusion vs diffusion As an ablation, we also measured the influence of the diffusion formulation on 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 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.

Latent MLP vs Latent Transformer Modeling interactions at the local and global levels is key to learn the dynamics faithfully. Experiments using MLPs (Table 7) as time steppers which do not consider interactions between tokens lead to significantly lower performance compared to transformers.

C.7 Kuramoto-Sivashinsky : a failure case

We conducted additional experiments on a chaotic 1D PDE, the Kuramoto-Sivashinsky ( KS ) equation. 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 1e-10 to 1e-12 was necessary for accurate spectrum reconstruction. Like all models leveraging a reduced latent representation space, AROMA inherently loses some of

the fine-grained details necessary for accurately capturing chaotic behavior. The diffusion framework slightly improves the high correlation time compared to 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 a current practice in fluid dynamics - e.g. LES (Large eddy simulation).

0

20

40

60

80

100

120

Wavenumber

10

6

10

5

10

4

10

3

10

2

10

1

10 0

10 1

10 2

Amplitude

Spectrum of the one-step predictions

Ground truth

AROMA No diffusion

AROMA

0

20

40

60

80

100

120

Time

10

5

10

4

10

3

10

2

10

1

10 0

MSE loss

AROMA

AROMA No Diffusion

C.8 Latent space dynamics

For Navier-Stokes , we show how the mean (Figure 24) and standard deviation tokens (Figure 25) evolve over time for a 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 is 0 (i.e. a high variance) 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.

0

20

40

1.0

0.5

0.0

Token 1

0

20

40

1.0

0.5

0.0

Token 2

0

20

40

1.0

0.5

0.0

Token 3

0

20

40

1.0

0.5

0.0

Token 4

0

20

40

2

1

0

1

Token 5

0

20

40

2

0

2

Token 6

0

20

40

4

2

0

2

4

Token 7

0

20

40

2

1

0

1

2

Token 8

0

20

40

2

0

2

Token 9

0

20

40

1.0

0.5

0.0

Token 10

0

20

40

1.0

0.5

0.0

Token 11

0

20

40

2

1

0

1

Token 12

0

20

40

1.0

0.5

0.0

Token 13

0

20

40

1.0

0.5

0.0

Token 14

0

20

40

2

0

2

Token 15

0

20

40

1

0

1

2

Token 16

c0

c1

c2

c3

c4

c5

c6

c7

c8

c9

c10

c11

c12

c13

c14

c15

0

20

40

2.0

1.5

1.0

0.5

0.0

Token 1

0

20

40

2.0

1.5

1.0

0.5

0.0

Token 2

0

20

40

2.0

1.5

1.0

0.5

0.0

Token 3

0

20

40

2.0

1.5

1.0

0.5

0.0

Token 4

0

20

40

4

3

2

1

0

Token 5

0

20

40

10

8

6

4

2

0

Token 6

0

20

40

10.0

7.5

5.0

2.5

0.0

Token 7

0

20

40

8

6

4

2

0

Token 8

0

20

40

8

6

4

2

0

Token 9

0

20

40

2.0

1.5

1.0

0.5

0.0

Token 10

0

20

40

1.5

1.0

0.5

0.0

Token 11

0

20

40

4

3

2

1

0

Token 12

0

20

40

2.0

1.5

1.0

0.5

0.0

Token 13

0

20

40

2.0

1.5

1.0

0.5

0.0

Token 14

0

20

40

10

8

6

4

2

0

Token 15

0

20

40

8

6

4

2

0

Token 16

c0

c1

c2

c3

c4

c5

c6

c7

c8

c9

c10

c11

c12

c13

c14

c15

0

20

40

1.0

0.5

0.0

Token 1

0

20

40

1.0

0.5

0.0

Token 2

0

20

40

1.0

0.5

0.0

Token 3

0

20

40

1.0

0.5

0.0

Token 4

0

20

40

2

1

0

1

Token 5

0

20

40

2

0

2

Token 6

0

20

40

4

2

0

2

4

Token 7

0

20

40

2

1

0

1

2

Token 8

0

20

40

2

0

2

Token 9

0

20

40

1.0

0.5

0.0

Token 10

0

20

40

1.0

0.5

0.0

Token 11

0

20

40

2

1

0

1

Token 12

0

20

40

1.0

0.5

0.0

Token 13

0

20

40

1.0

0.5

0.0

Token 14

0

20

40

2

0

2

Token 15

0

20

40

1

0

1

2

3

Token 16

c0

c1

c2

c3

c4

c5

c6

c7

c8

c9

c10

c11

c12

c13

c14

c15

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 on 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 in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be 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.

It is 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 of the work performed by the authors?

Answer: [Yes]

Justification: There is a 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 not discussed in the paper.

The authors are encouraged to create a separate 'Limitations' section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a 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 approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it 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 address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [NA]

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 crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

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 conclusions 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 to 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 a 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 a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how 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 submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a 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 a way to access this model for reproducing the results or a way to reproduce 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 authors 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., to 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 instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: Everything will be made available upon acceptation.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a 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/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a 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 as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary 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 to a 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 appropriate 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, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

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

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: 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 or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: Yes we conformed to the code of ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [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 specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to

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 to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used 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 mitigation strategies (e.g., gated release of models, providing defenses in addition to 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

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [No]

Justification: This paper causes no such risk.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines 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 is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or 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?

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 or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of 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, paperswithcode.com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [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 of the dataset/code/model as part of 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 is 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 paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: This paper is not related to research on 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 contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, 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 an equivalent approval/review based on 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 in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to 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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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).\n\nQUESTIONS:\n- 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.\n\nLIMITATIONS:\nPlease see my concerns in **Weaknesses** and **Questions**.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-08"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\n1. 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\"\n\nLIMITATIONS:\nThe authors have discussed the limitation in the conclusion part.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nAlthough 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.\n\nQUESTIONS:\nIt 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.\n\nLIMITATIONS:\nWhile 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-13"}

comments: 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.

STRENGTHS:
na

WEAKNESSES:
na

QUESTIONS:
1. What are the main drawbacks that AROMA seeks to solve with respect to current neural operator models for PDEs?
3. What function does the AROMA architecture's conditional transformer serve?
5. How can AROMA guarantee the accuracy and stability of its forecasting models?
5. How is the model's performance improved by the encoder and decoder's two-stage training process?
6. In the experiments, what kinds of datasets were employed, and how did AROMA fare in comparison to other models?

LIMITATIONS:
na
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2024-07-16