Universal Physics Transformers: A Framework For Efficiently Scaling Neural Operators
Authors: Benedikt Alkin
Andreas Fürst
Simon Lucas Schmid
Lukas Gruber
Markus Holzleitner
Johannes Brandstetter
Venue: NeurIPS 2024 poster (2024)
Keywords: neural operator
computational fluid dynamics
Lagrangian simulations
transformers
latent space modeling
Primary area: machine_learning_for_physical_sciences
TL;DR: We introduce Universal Physics Transformers, an efficiently scalable neural operator framework to model a wide range of spatio-temporal problems – for Lagrangian and Eulerian discretization schemes.

Abstract: Neural operators, serving as physics surrogate models, have recently gained increased interest. With ever increasing problem complexity, the natural question arises: what is an efficient way to scale neural operators to larger and more complex simulations - most importantly by taking into account different types of simulation datasets. This is of special interest since, akin to their numerical counterparts, different techniques are used across applications, even if the underlying dynamics of the systems are similar. Whereas the flexibility of transformers has enabled unified architectures across domains, neural operators mostly follow a problem specific design, where GNNs are commonly used for Lagrangian simulations and grid-based models predominate Eulerian simulations. 

We introduce Universal Physics Transformers (UPTs), an efficient and unified learning paradigm for a wide range of spatio-temporal problems. UPTs operate without grid- or particle-based latent structures, enabling flexibility and scalability across meshes and particles. UPTs efficiently propagate dynamics in the latent space, emphasized by inverse encoding and decoding techniques. Finally, UPTs allow for queries of the latent space representation at any point in space-time. We demonstrate diverse applicability and efficacy of UPTs in mesh-based fluid simulations, and steady-state Reynolds averaged Navier-Stokes simulations, and Lagrangian-based dynamics.

Universal Physics Transformers: A Framework For Efficiently Scaling Neural Operators
Benedikt Alkin 1.2 Andreas Fürst
Simon Schmid Lukas Gruber
Markus Holzleitner
Johannes Brandstetter 1,2
ELLIS Unit Linz, Institute for Machine Learning, JKU Linz, Austria
2 NXAI GmbH, Linz, Austria Software Competence Center Hagenberg GmbH, Hagenberg, Austria MaLGa Center, Department of Mathematics, University of Genoa, Italy {alkin, fuerst, brandstetter)(ml. jmk. at
Abstract
Neural operators, serving as physics surrogate models, have recently gained in- creased interest. With ever increasing problem complexity, the natural question arises: what is an efficient way to scale neural operators to larger and more complex simulations most importantly by taking into account different types of simulation datasets. This is of special interest since, akin to their numerical counterparts, used across applications, of the systems are similar. Whereas the flexibility of transformers has enabled unified architectures across domains, neural operators mostly follow a problem specific design, where GNNs are commonly used for Lagrangian simulations and grid-based models predominate Eulerian simulations. We introduce Universal Physics Transformers (UPTs), an efficient and unified learning paradigm for wide range spatio-tempora problems. UPTs operate without grid- or particle-bas based latent structures, enabling flexibility and scalability across meshes and particles. UPTs efficiently propagate dynamics in the latent space, emphasized by inverse encoding and decoding techniques. Finally, UPTs allow for queries of the latent space representation at any point in space-time. We demonstrate diverse appli- cability and efficacy of UPT's in mesh-based fluid simulations, and steady- Reynolds averaged Navier-Stokes simulations, and Lagrangian-bain based dynamics. Project page: attps://ml. ju. github. io/UPT
different techniques
if the underlying dynamics
Introduction
In scientific pursuits, extensive efforts have produced highly intricate mathematical models of physical phenomena, many of which are naturally expressed partial differential equations (PDEs) [76). Solving most PDEs is analytically intractable and necessitates falling back on compute-expensive numerical approximation schemes. In recent years, deep neural network based surrogates, most importantly neural operators [1]
47), have emerged as a computationally efficient alternative [99 119), and impact e.g., weather forecasting [48 [ molecular modeling [7 (4, or computational fluid dynamics [07 30
19). Additional to computational efficiency, neural surrogates offer potential to introduce generalization capabilities across phenomena, well generalization across characteristics such as boundary conditions or PDE coefficients [66] 14]. Consequently. the nature of neural operators inherently complements handcrafted numerical solvers which characterized by substantial set of solver requirements, and mostly due to these requirements tend to differ among sub-problems
5T
43
31
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Discretized Eulorion Lagrangian Data at
Discretized Eulorian/ / Lagrangian Data at
Features
Position
Latent Token:
Positions Quary
Ropeat until d
Encoder
Approximator
Decoder
Figure 1: Schemantic sketch of the UPT learning paradigm. UPTs flexible encode different grids, and/or different number of particles into unified latent space representation, and subsequently unroll dynamics in the latent space. The latent space is kept fixed size to ensure scalability to larger systems. UPTs decode the latent representation any query point.
However, similar to their numerical counterparts, different neural network techniques are prevalent across applications. For example, when contrasting particle- and grid-based dynamics in compu- tational fluid dynamics (CFD), i.e., Lagrangian and Eulerian discretization schemes. This is in contrast other areas of deep learning where the flexibility of transformers [06] has enabled unified architectures across domains, allowing advancements in one domain to also benefit all others. This has lead to an efficient scaling of architectures, paving the way for large "foundation" models [IO] that are pretrained on huge passive datasets 25 36]
We introduce Universal Physics Transformers (UPTs), efficient and unified neural operator learning paradigm with strong focus on scalibility over a wide range of spatio-temporal problems. UPT's flexibly encode different grids, and/or different number of particles into compressed latent space representation which faacilitates scaling to large-scale simulations. Latent space rollouts are enforced by inverse encoding and decoding surrogates, leading to fast simulated trajectories which is particularly important for large systems. For decoding, the latent representation can be evaluated any point in space-time UPTs operate without grid- or particle-baat based latent structures, and demonstrate the beneficial scaling-behavior of transformer backbone architectures. Figure I sketches the UPT modeling paradigm.
We summarize our contributions as follows: (i) we introduce the UPT framework for efficiently scaling neural operators; (ii) formulate encoding and decoding schemes such that dynamics be propagated efficiently in a compressed and fixed-size latent space; (iii) we demonstrate applicability on diverse applications, putting strong research focus the scalability of UPTs.
2 Background
Partial differential equations. We focus experiments on (systems of) PDEs, that evolve signal utt, a:) = u((x) E Rd in a single temporal dimension / E [o, T] and m spatial dimensions U C Th, for an open set U. With 1 < l E N, systems of PDEs of order I can be written as
O' (x)..... (x),
F(D'u'(z)).... .en
oe ue
u' (x), u (x), x, = 0, for a E E [0,T),
St
where F is a mapping to Rd, and for 1,..., D denotes the differential operator mapping to all i-th order partial derivatives of u with respect to the spacial variable æ, whereas g¹ ØF outputs the corresponding time derivative of order i. Any l-times continuously differentiable u: [o. 7] U - Rd fulfilling the relation Eq. (1) is called a classical solution of Eq. ① Also other notions of solvability (e.g. in the sense of weak derivatives/distriuttooo into details here. Additionally, initial conditions specify u (2x) at time = and boundary conditions B uu (x) at the boundary of the spatial domain.
are possible, for sake of simplicity we do not go

---

We work mostly with the incompressible Navier-Stokes equations 198] which e.g., in two spatial dimensions conserve the velocity flow field ut, x. y) [0. T x E ? via:
du Ot
u Vu uV Vp+ u.
(2)
where u. Vu is the convection, i.c., the rate of change of u along u, Il is the viscosity parameter, u the viscosity, i.e., the diffusion or net movement of u, Vp the internal pressure gradient, and
an external force. The constraint V u = 0 yields mass conservation of the Navier-Stokes equations. A detailed depiction of the involved differential operators is given in Appendix [B.1
learns a mapping between function spaces a concept which is often used to approximate solutions of PDEs. Similar to Kovachki et al. [47], we assume 24, V to be Banach spaces of functions on compact domains X C R or y C R mapping into Rd or R respectively. The goal of operator learning is to learn a ground truth operator 9 ::2. 1 V via an approximation 9 - V. This is usually done in the vein of supervised learning by i..d. sampling input-output pairs, with the notable difference, that in operator learning the spaces sampled from are not finite dimensional. More precisely, with given data set consisting of N function pairs (u, v;) (u,; G(u.)) C U V, i = I,.,,, aim to learn 9 U + V, so that G can be approximated in suitably chosen norm.
Operator learning. Operator learning [60
61 52 51 47
In the context of PDEs, G can c.g. be the mapping from initial condition u(0, x) = (x)) to the solutions (tt,, x) = u' ( of Eq. ① at all times. In the case of classical solutions, if U is bounded, can then be chosen as a subspace of C(U, Rd),, the set of continuous functions from domain Ü (the closure of U) mapping to Rd whereas c C(OO, TJ U, RR), so that i or V consist of all I-times continuosly differentiable functions on the respective spaces. In case of weak solutions, the associated spaces Ul and can be chosen Sobolev spaces.
We follow the popular approach approximate G via three maps [89] G R 9 DoA E. The encoder & 14 + IR takes an input function and maps it to a finite dimensional latent feature chosen hidden dimension
representation. For example, & could embed continuous function 1 for a collection of grid points. Next, A: : Ra + R approximates the action of the operator G. and D decodes the hidden representation, and thus creates the output functions via D IR + V. which in many cases is point-wise evaluated at the output grid or output mesh.
Particle vs. grid-based methods. Often, numerical simulation methods can be classified into two distinct families: particle and grid- based methods. This specification is notably prevalent, for instance, in the field of computational fluid dynamics (CFD), where Lagrangian and Eulerian discretization schemes offer different characteristics dependent on the PDEs. In simpler terms, Eulerian schemes essentially monitor velocities at specific fixed grid points. These points, represented by a spatially limited number of nodes, control volumes, or cells, serve to discretize the continuous space. This process leads grid-based or mesh-based representations. In contrast to such grid- and mesh-based representations, in Lagrangian schemes, the discretization is carried out using finitely many material points, often referred to as particles, which move with the local deformation of the continuum. Roughly speaking, there are three families of Lagrangian schemes: discrete element methods (24), material point methods 94 [2]. and smoothed particle hydrodynamics (SPH) 29 62
69 70). In
this work, we focus on SPH methods, which approximate the field properties using radial kernel interpolations over adjacent particles at the location of each particle. The strength of SPH lies in its ability to operate without being constrained by connectivity issues, such as meshes. This characteristic proves especially beneficial when simulating systems that undergo significant deformations.
Latent space representation of neural operators. For larger meshes or larger number of particles, memory consumption and inference speed become more and more important. Fourier Neural Operator (FNO) based methods work on regular grids, or learn a mapping to a regular latent grid, e.g., geometry- informed neural operators (GINO) [54]. In three dimensions, the stored Fourier modes have the shape h x n₂ X ny X nz where h is the hidden size and nz. Ny. n2 are the respective Fourier modes. Similarly, the latent space of CNN-based methods, e.g., Raonić et al. [83, Gupta Brandstetter [I]. is of shape h X Wr Wy X W,, where Wz, Wy wa are the respective grid points. In three dimension, the memory requirement in each layer increases cubically with increasing number of modes or grid points. In contrast, transformer based neural operators, e.g., Hao et al. [3 Cao [8] Li et al. [53
operate on a boken-based latent space of dimension Tookens x h, where usually Thokens XX Ppoints and GNN based neural operators, c.g., Li ct al. [52]. operate on node based latent space of dimension
3

---

80
Irregular Grid
Diser.
Learns Field
Latent
Model
Range
Conv.
Rollout
Memory [GB]
60
GNN
local
X
X
CNN
local
X
40
Transformer
global
GNO 52
radius
J
/
X
X
20
FNO S1
global
X
V
X
GINO
global
V
seco
UPT
global
2
4
16 32 64
128
Scale
Figure 2: Qualitative exploration of scaling limits. Starting from 32K input points (scale 1), we train 68M parameter model for few steps with batchsize 1 and measure the required GPU memory. Models without a compressed latent space (GNN, Transformer) quickly reach their limits while models with a compressed latent space (GINO, UPT) scale much better with the number of inputs. However, as GINO compresses the latent space onto regular grid, the scaling benefits are largely voided on 3D problems. The efficient latent space compression of UPTs can fit up to 4.2M points (scale 128). We use a linear attention transformer [90] for this study. "Disc. Conv." denotes "Discretization Convergent". Appendix D.7 outlines implementation details and complexities.
Input
Data
Postionn
Bakect
Model
Enooder
Decoder
Latent
Propagaie Timn
Aporoximator (Treraformme)
Perratiorr Pronling
Aop roxmmaaoo
A'"
Postams
Decader (Phrtakkerr
Daca
Leammalle
Decoder
Enooder
Piediction
Forilange
G
Figure 3: Left: UPT compresses information from various grids or differing particles with an encoder, propagates this information forward in time through the approximator and decodes arbitrary query positions. Right: Training procedure to enable latent rollouts via inverse encoding/decoding losses.
Onodes X h, where usually Piodes Ppoints- For large number of inputs, this becomes infeasible as every layer has to process large number of tokens. Contrary, UPTs compress the inputs into low-dimensional latent space, which drastically decreases computational requirements. Different architectures and their scaling limits are compared in Fig.
3 Universal Physies Transformers
Problem formulation. Our goal is to learn a mapping between the solutions w' and u" of Eq. ( at timesteps t and t', respectively. Our dataset should consist of N function pairs (uf, u. ) 1= 1,.., ......, N, where each uf is sampled
spatial locations {a: ...;.... at} E U. Similarly, we query each output signal u, at k spatial locations .y....., y E U. Then each input signal can be represented by uh,x = (a!())),... ui(x)) ( E BRxd as a tensor of shape x d, similar for the output. For particle- or mesh-based inputs, it is often simpler to represent the input as graph = (V, F)) with nodes (x].....f] E V. edges and node features (u{ ... ut (x))).
Architecture desiderata. We want Universal Physics Transformers (UPTs) to fulfill the following set of desiderata: an encoder ε, which flexibly encodes different grids, and/or different number of particles into a unified latent representation of shape Niatent X h, where Naaten is the chosen number of tokens the latent space and the hidden dimension: (ii) an approximator A and training procedure, which allows us to forward propagate dynamics purely within the latent space without mapping back the spatial domain at each operator step; and iii) a decoder D that queries the latent representation at different locations. The UPT architecture is schematically sketched in Fig. 3

---

Encoder. The goal of the encoder & is to compress the input signal uj, which is represented by a point cloud " Importantly. the encoder should lear to selectively focus on important parts of the input. This is a desirable property as, for example, in many computational fluid dynamics simulations large areas are characterized by laminar flows, whereas turbulent flows tend to occur especially around obstacles. If k is large, we employ a hierarchical encoder.
The encoder & first embeds k points into hidden dimension h, adding position encoding [006] to the different nodes, i.e., u1,x x E R fxd - R In the first hierarchy, information is exchanged between local points and a selected set of n, supernode points. For Fulerian discretization schemes those supernodes can either be uniformly sampled on a regular grid as in (54), or selected based on the given mesh. The latter has the advantage that mesh characteristics are automatically taken into account, e.g., dense or sparse mesh regions are represented by different numbers of nodes. Furthermore, adaptation to new meshes is straightforward. We implement the first hierarchy by randomly selecting n, supermodes on the mesh, choosing n, such that the mesh characteristic is preserved. Similarly, in the Lagrangian discretization scheme, choosing supernodes based particle positions provides the same advantages as selecting them based on the mesh.
Information is aggregated at the selected n, supernodes via message passing layer 28 using radius graph between points. Importantly, messages only flow towards the n, supernodes, and thus the compute complexity of the first hierarchy scales linearly with ng.- The second hierarchy consists of transformer blocks [106] followed by a perceiver block [40] [3] with Naatent learned queries of dimension h. To summarize, the encoder & maps u E U to latent space via
E : u E U evaluate, witk E 1 smbed, IR jakxh
MP, 1
transformer, R perceiver
2f RR Tinne
where tyically Valteen < nl < k. If the number of points is small, the first hierarchy can be omitted.
Note that randomly sampling mesh cells or particles implicitly encodes the underlying mesh or particle density and allocates more supernodes to highly resolved areas in the mesh or densely populated particle regions. Therefore, this can be seen as implicit "importance sampling" of the underlying simulation. Additional implementation details are provided in Appendix [F1
This encoder design projects into fixed size latent space as it is efficient way to compress the input into a fixed size representation to enable scaling to large-scale systems while remaining compute efficient. However, if an application requires a variable sized latent space, one could also remove the perceiver pooling layer. With this change the number of supernodes equal to the number of latent tokens and complex problems could be tackled by larger supernode count.
Approximator. The approximator propagates the compressed representation forward in time. As Mattet is small, forward propagation in time is fast. We employ transformer as approximator.
A: o E B xh + + E I Ilomn
Notably, the approximator can be applied multiple times, propagating the signal forward in time by At each time. If At is small enough, the input signal can be approximated at arbitrary future times l'.
Decoder. The task of decoder D to query the latent representation at k' arbitrary locations to construct the prediction of the output signal ut at time t'. More formally, given the output positions .... y } E U at spatial locations and the latent representation 2!" the decoder predicts the output signal uixx- (ut (u.))..., ,., (v! ))² at these spatial locations at timestep t',
D: (2f', .u!....."!)) of E F'd
The decoder is implemented via perceiver-like cross attention layer using a positional embedding of the output positions query and the latent representation z!' as keys and values. Since there is no interaction between queries, the latent representation can be queried at arbitrarily many posi- tions without large computational overhead. This decoding mechanism establishes a connection of conditioned neural fields to operator learning [79].
Model Conditioning. To condition the model to the current timestep and to boundary conditions such as the inflow velocity, we add feature modulation to all transformer and perceiver blocks. We use DiT modulation (78), which consists of dimension-wise scale, shift and gate operation that
5

---

applied to the attention and MLP module of the transformer. Scale, shift and gate are dependent on an embedding of the timestep and boundary conditions (e.g. velocity).
Training procedure. UPTs model the dynamics fully within a latent representation, such that during inference only the initial state of the system u(0, a²) = = is encoded into a latent representation z". From there on, instead of autoregressively feeding the decoder's prediction into the encoder, UPTs propagate z" forward in time to z" through iteratively applying the approximator A in the latent space. We call this procedure latent rollour. Especially for large meshes or many particles, the benefits of latent space rollouts, i.c. fast inference, pays off.
To enable latent rollouts, the responsibilities of encoder ε, approximator A and decoder 'D need to be isolated. Therefore, we invert the encoding and decoding by means of two reconstruction losses during training visualized in Fig. First, an inverse encoding is performed, wherein the input us! reconstructed from the encoded latent state 2 by querying it with the decoder at A input locations
(a!...., ). Second, we invert the decoding by reconstructing the latent state =!" from the output signal it at k' spatial locations {y! ..... y }- Using two reconstruction losses, the encoder is forced to focus on encoding a state u into a latent representation x², and similarly the decoder is forced to focus on making predictions out of a latent representation z²".
Related methods. The closest work to ours are transformer neural operators of Cao [8] Li et al. [3]. Hao et al. (33] which encode different query points into a tokenized latent space representation of dimension Poodes x ht where modes varies based the number of input points, i.e., Poodes O Apoints- Wu et al. [14] adds a learnable mapping into a fixed latent space of dimension nnodes x h to each transformer layer, and projects back to dimension Pppints h after self-attention. In contrast, UPTs use fixed Matent for the unified latent space representation Naatent X h.
For the modeling of temporal PDEs, common scheme is to map the input solution at time t to the solution at next time step [51 13 955. Especially for systems that are modeled by graph- based representations, predicted accelerations at nodes numerically integrated to model the time evolution of the system [87 8T.. Recently, equivariant graph neural operators [116] were introduced which model time evolution via temporal convolutions in Fourier space. More related our work are methods that propagate dynamics in the latent space (50, 111. Once the system is encoded, time evolution is modeled via LSTMs [1I1], even linear propagators [63 73.. In Li et al. (53), attention-basel layers are used for encoding the spatial information of the input and query points, while time updates in the latent space are performed using recurrent MLPs. Similarly, Bryutkin et al. [6] use recurrent MLPs for temporal updates within the latent space, while utilizing graph transformer for encoding the input observations.
Building universal models aligns with the contemporary trend of foundation models for science. Recent works comprise pretraining over multiple heterogeneous physical systems, mostly in the form of PDEs 66). foundation models for weather and climate (75), or material modeling 67 118 4
Methods similar to our latent space modeling have been proposed in the context of diffusion mod- els (85] where a pre-trained compression model used to compress the input into a latent space from which diffusion model can be trained at much lower costs. Similarly, our approach also compresses the high-dimensional input into a low-dimensional latent space, but without a two stage approach. Instead, we learn the compression end-to-end via inverse encoding and decoding techniques.
4 Experiments
We ran experiments across different settings, assessing three key aspects of UPTs: (i) Effectiveness of the latent space representation. We test on steady state flow simulations in 3D, comparing against methods that use regular grid representations, and thus considerably larger latent space representations. (ii) Scalability. We test on transient flow simulations on large meshes. Specifically, we test the effectiveness of latent space rollouts, and assess how well UPTs generalize across different flow regime, and different domains, i.e., different number of mesh points and obstacles. (iii) Lagrangian dynamics modeling. Finally, we assess how well UPTs model the underlying field characteristics when applied to particle-based based simulations. We outline the most important results in the following sections and provide implementation details and additional results in Appendix D Most notabily, UPT also compares favorably against bascline regular grid methods on regular grid datasets D.3

---

[= 30
t= 40
Ground Truth
Error Prediction
2
Figure 4: Example rollout trajectories of the UPT-68M model, visually demonstrating the efficacy of UPT physics modeling. The UPT model is trained across different obstacles, different flow regimes, and different mesh discretizations. Interestingly, the absolute error might suggest that UPT trajectories diverge, although physics are still simulated faithfully. This stems from subtle shifts in predictions throughout the rollout duration, likely attributed the point-wise decoding the latent field.
4.1 Steady state flows
For steady state prediction, we consider the ShapeNet-Car dataset generated by [105]. It consists of 889 car shapes from ShapeNet [20). where each car surface is represented by 3.6K mesh points 3D space. We randomly create train/test split containing 700/189 samples. We regress the pressure each surface point with a mean-squared error (MSE) loss and sweep hyperparameters per model. Due the small scale of this dataset, we train the largest possible model that able generalize the best, Training even larger models resulted in performance decrease due to overfitting. We optimize the model size for all methods where the best mesh based models (GINO, UPT) contain around 300M parameters. The best regular grid based models (U-Net [86] 31J FNO (5T)) are significantly smaller and range from 15M to 100M. Additional details are listed in Appendix D.4
ShapeNet-Car is a scale dataset. Consequently, methods that map the mesh onto a regular grid can employ grids of extremely high resolution, such that the number of grid points is orders of magnitude higher than the number of mesh points. For example, a grid resolution of 64 points per spatial dimension results in 262K grid points, which is 73x the number of mesh points. As UPT is designed to operate directly on the mesh, we compare at different grid resolutions.
UPT is able to outperform models on smaller resolutions, for example UPT with only 64 latent tokens achieves a test MSE of 2.31 whereas GINO with resolution 48 (110K tokens) achieves only 2.58 while taking significantly more runtime and memory. When additionally using feature engincering in the form of signed distance function and 64 grid resolution UPT achieves a competitive performance of 2.24 compared to 2.14 of GINO while remaining efficcintt. All test errors are multiplied by 100. UPT achieves a favorable cost-vs-pen cett-s-peerfommacce tradeoff where the best GINO models requires 900 seconds per epoch whereas only slightly worse (-0.17) UPT model takes only 4 seconds. We show a comprehensive table with all results in Appendix Tab,
4.2 Transient flows
We test the scalability of UPTs on large-scale transient flow simulations. For this purpose, we self-generate 10K Navier-Stokes simulations within a pipe flow using the pisoFoam solver from OpenFOAM [110], which we split into 8K training, 1K validation and 1K test trajectories. For each simulation, between one and four objects (circles of variable size) are placed randomly within the pipe flow, and the uni-directional inflow velocity varies between 0.01 to 0.06 m/s. The simulation is carried out for 100s resulting 100 timesteps that are used for training neural surrogates. Note that pisoFoom requires much smaller A/ to remain stable (between 2K and 200K timesteps for 100s of simulation time). We use an adaptive meshing algorithm which results in 29K to 59K mesh points. Further dataset details are outlined in Appendix D.5

---

1e-3
10
70
GINO
Test MSE (+)
UPT
FND
Correlation Time
60-
U-Net
Test MSE
50
UPT
40
8
17
68
17 eters [M]
68
4K
8K
16K
24K 32K
Parameters [M]
Number InputiOutput Points
Figure 5: Left and middle: MSE and correlation time the testset. UPTs outperform compared methods on all model scales by a large margin. Right: We study discrctization convergence by varying the number of input/output points or the number of gridpoints/supernodss of models that were trained on inputs between 8K and 24K points, 8K target points. UPT demonstrates a stable performance across different number of input/outputs even though it has never seen that number of input/output points during training. We study discretization convergence of supernodes in App. A.5. where UPT also shows a steady improvement when more supernodes are used during inference. Additionally, we study smaller UPT models and training on less data in Appendix D.5.4 and D.5.5
Model-wise, UPT uses the hierarchical encoder setup with optional components depicted in Fig. 3 A message passing layer aggregates local information into 2048 randomly selected supernodes, a transformer processes the supernodes and a perceiver pools the supernodes into Raatent = 512 latent tokens. Approximator and decoder are unchanged. We compare UPT against GINO, U-Net and FNO. For U-Net and FNO, we interpolate the mesh onto a regular grid. We condition the models onto the current timestep and inflow velocity by modulating features within the model. We employ FiLM conditioning for U-Net 80), the "Spatial-Spectrall conditioning method introduced in (31] for FNO and GINO, and DiT for UPT (78). Implementation details are provided in Appendix [D.5
We train all models for 100 epochs and evaluate test MSE as well as rollout performance for which we
use the number of timesteps until the Pearson correlation of the rollout drops below 0.8 as evaluation
metric 43 We do not employ any techniques to stabilize rollouts [57]. The left side of Fig. [
shows that UPTs outperform compared methods by large margin. Training even larger models becomes increasingly expensive and is infeasible for our current computational budget. UPT-68M
and GINO-68M training takes roughly 450 A100 hours. UPT-8M and UPT-17M take roughly 150 and 200 A 10 hours, respectively. Figure + shows a rollout and Appendix D.5.] presents additional ones. We also study dfistribuitiioooonnnnnn generalization (e.g. more obstacles) in Appendix [D.366]
While one would ideally use lots of supernodes and query the latent space with all positions during training, increasing those quantities increases training costs and the performance gains saturate. Therefore, we only use 2048 supernodes and 16K randomly selected query positions during training. We investigate discretization convergence in the right part of Fig. where we vary the number of input/output points and the number of supernodes. We use the 68M models without any retraining, i... we test models "discretization convergence" as, during training, the mesh was discretized into 2048 supernodes and 16K query positions. UPT generalizes across a wide range of different number of input output positions, with even slight performance increases when using more input points. Similarly, using more supernodes increases performance slightly. Additionally, we investigate training with more supernodes and/or more latent tokens in a reduced setting in Appendix D.5.3 D.5.3
Finally, we evaluate training with inverse encoding and decoding techniques, see Fig. B We investigate the impact of the latent rollout by training our largest model 68M UPT. The latent rollout achieves on par results to autoregressively unrolling via the physics domain, but speeds up the inference significantly as shown in Tab. [
However, in its current implementation the latent rollout requires non-negligible overhead during training. We discuss this limitation in Appendix A
4.3 Lagrangian fluid dynamics
Scaling particle-based methods such as discrete element methods or smoothed particle hydrodynamics to 10 million or more particles presents a significant challenge [17 99 yet it also opens a distinctive opportunity for neural surrogates. Such systems are far beyond the scope of this work. We however present framing of how to model such systems via UPTs such that the studied scaling properties
8

---

Table Required time to simulate full trajectory rollout. UPT and GINO are orders of magnitude faster than traditional finite volume solvers. The latent rollout is additionally more than 5x faster than an autoregressive rollout via the physics domain. Neural surrogate models are also faster on CPUs as traditional solvers require extremely small timescales to remain stable (ÁI ε 0.05 vs. A 1).
Model
Time on 16 CPUs
Time GPU
Speedup
pisoFoam
120s
1x
GINO-68M (autoreg.)
48s
1.2s
100x
UPT-68M (autoreg.)
46s
2.0s
60x
UPT-68M (latent)
3s
0.3s
400x
of UPTs could be exploited. In order to do so, demonstrate how UPTs capture inherent field characteristics when applied to Lagrangian SPH simulations, as provided in LagrangeBench [101]. Here, GNNs, such Graph Network-bascd Simulators (GNS) [87] and Steerable E(3) Equivariant Graph Neural Networks (SEGNNs) [14] are strong baselines, where predicted accelerations at the nodes numerically integrated to model the time evolution of the particles. In contrast, UPTs learn underlying dynamics without dedicated particle-structume
and propagate dynamics forward without the guidance of numerical time integration schemes. An overview of the conceptual differences between GNS/SEGNN and UPTs shown in Fig. 6
Model predicts a'
+a
Figure 6: Conceptual difference between GNS/SEGNN on the left and UPT on the right side. GNS/SEGNN predicts the acceleration of a particle which then integrated to calculate the next position. UPTs directly model the velocity field and allow for large timestep predictions.
We use the Taylor-Green vortex dataset in three dimensions (TGV3D). The TGV system was introduced as test scenario for turbulence modeling (97). It is unsteady flow of decaying vortex, displaying an exact closed form solution of the incompressible Ntvier-Stokes equations in Cartesian coordinates. We note that the TGV3D dataset models the same trajectories but does so by tracking particles using SPH to solve the equations. Formulating the TGV system as a UPT learning problem allows the same trajectories be queried different positions, enabling the recovery of the complete velocity field, whereas GNNs can only evaluate the velocities of the particles. Consequently, the evaluation against GNNs should be viewed as an illustration of the efficacy of UPTs learning field characteristics, rather than a comprehensive GNN versus UPT comparison. More details and experiments on the 2D version of the Taylor-Green vortex dataset (TGV2D) are in Appendix [.6
For UPT training, we input two consecutive velocities of the particles in the dataset at timesteps and t 1, and the respective particle positions. We regress two consecutive velocities at later timesteps {t' 1,t = {t AT 1, + AT) with mean-squared error (MSE) objective. For all experiments we use AT 10At. We query the decoder output velocities target positions. UPTs encode the first two velocities of a trajectory, and autoregressively propagate dynamics forward in the latent space. We report the Euclidean norm of velocity differences across all k particles. Figure F compares the rollout performance of GNS, SEGNN and UPT and shows the speedup of both methods compared the SPH solver. The results demonstrate that UPTs effectively learn the underlying field dynamics.
Discussion
Potential for extreme scale. UPTs consist of mainly transformer [066 layers which allows application of the same scaling and parallelism principles as are commonly used for training extreme- scale language models. For example, the recent work Llama 3 [27] trained on up 16K H100 GPUs.

---

1e-2
GNS
SECRIN
UPT
vabeity Erre
Timesttep
Figure 7: Left: Velocity error over all particles for different timesteps. UPTs effectively learn the underlying field dynamics, resulting in lower error as the trajectory evolves in time. Right: Visualization of the velocity field modeled by UPT (white) vs the ground truth particle velocities.
While we do not envision that we train UPT on such a massive scale in the foreseeable future, the used techniques can be readily applied to UPTs.
Benefits of the latent rollout. While the latent rollout does not provide a significant performance improvement, it is almost an order of magnitude faster. The UPT framework allows to trade-off' training compute vs inference compute. If inference time is crucial for a given application, one can train UPT with the inverse encoding and decoding objectives, requiring more training compute but greatly speeding up inference. inference time not important, one can simply train UPT without the reconstruction objectives to reduce training costs.
Additionally, the latent rollout enables applicability to Lagrangian simulations. As UPT models the underlying field instead of tracking individual particle positions does not have access to particle locations inference time. Therefore, autoregressive rollouts are impossible since the encoder requires particle positions as input. When using the latent rollout, it is sufficient to know the initial particle positions as dynamics are propagated without any spatial positions. After propagating the latent space forward in time, the latent space can be queried at arbitrary positions to evaluate the underlying field at given positions. We show this by querying with regular grid positions in Figure
Conclusion
We have introduced Universal Physics Transformers (UPTs) framework for efficiently scaling neural operators, demonstrating its applicability to a wide range of spatio-temporal problems. UPT' operate without grid- or particle-bas based latent structures, enabling flexibility across meshes and number of particles. The UPT training procedure separates responsibilities between components, allowing forward propagation in time purely within the latent space. Finally, UPTs allow for queries of the latent space representation at any point in space-time.
Acknowledgments
We would like sincerely thank Artur P. Toshev and Gianluca Galletti for ongoing help with and in-depth discussions about LagrangeBench. Johannes Brandstetter acknowledges Max Welling and Paris Perdikaris for numerous stimulating discussions.
We acknowledge EuroHPC Joint Undertaking for awarding us access to Karolina at IT4Innovations, Czech Republic and Leonardo at CINECA, Italy.
The ELLIS Unit Linz, the LIT AI Lab, the Institute for Machine Learning, are supported by the Federal State Upper Austria. We thank the projects Medical Cognitive Computing Center (MC3), INCONTROL-RL (FFG-88 (--8888164)) 1064). PRIMAL (FFG-873979), S3AI (FFG-872172). DL for Gran- ularFlow (FFG-871302), EPILEPSIA (FFG-892171), AIRI FG 9-N (FWF-36284 FWF-36235), AAGrrenHeattinGrriss (FFG- 899943), INTEGRATE BFFG-892418), ELISE H1202-IICT-2015 We thank Kudi.JKU Deep
ID: 951847), Stars4Waters HORIZON-CL6-202 WWW.ONNCLL--0I--CLMMMEEEEEEOOO -LL.AATE-----OI1111 Learning Center. TGW LOGISTICS GROUP GMBH, Silicon Austria Labs (SAL), FILL Gesellschaft mbH, Anyline GmbH, Google, ZF Friedrichshafen AG, Robert Bosch GmbH, UCB Biopharma SRL, Merck Healthcare KGaA, Verbund AG, Software Competence Center Hagenberg GmbH, Borealis AG, TÜV Austria, Frauscher Sensonic, TRUMPF and the NVIDIA Corporation.
10

---

References
[1] Andrychowicz, M., Espcholt, L.,, Li, D., Merchant, S., Merose, A., Zyda, F,, Agrawal, S., and Kalch- brenner, N. Deep learning for day forceasts from sparse observations. arXiv preprint arXiv:2306.0000 2023.
[2] Ba, I., J., Kiros, J. R., and Hinton, G. E. Layer normalization. CoRR, 0tte///607.06500, 2016.
[3] Bartels, S. Numerical Appreximation of Partial Differential Equations. Springer, 2016.
[4] Batatia, I., Kovacs, D. P.. Simm, G., Ortner, C., and Csányi, G. Mace: Higher order equivariant message passing neural networks for fast and accurate force fields. Advances in Neural Information Processing Systems, 35: 15:11423-111 11436, 2022.
[5] Battaglia, P., Pascanu, R., Lai, M., Jimenez Rezende, D., et al. Interaction networks for learning about objects, relations and physics. Advances in neural information processing systems, 29, 2016.
[6] Battaglia, P. W., Hamrick, J. B., Bapst, V., Sanchez-Gonzalez, A., Zambaldi, V., Malinowski, M., Tacchetti, A., Raposo, D., Santoro, A., Faulkner, R., el al. Relational inductive biases, deep learning, and graph networks. arXiv preprint 006.:v::8860/2666000 2018
[7] Batzner. S., Musaelian, A., Sun, L., Geiger, M., Mailoa, J. P., Kornbluth, M., Molinari, N., Smidt, T. E.. and Kozinsky, B. E (3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials. Nature communications. 13(1):2453, 2022.
[8] Bi, K., Xie, L., Zhang. H., Chen, X., X., and Tian, Q. Pangu-weather. A 3d high-resolution model for fast and accurate global weather forecast. arXiv preprint arXiv:2211.02556 2022.
[9] Blais, B., Vidal, D., Bertrand, F., Patience, G. S., and Chaouki, J. Experimental methods in chemical engineering: Discrete element methhod-dmm. The Canadian Journal of Chemical Engineering, 97(7): 1964-1973, 2019.
[10] Bommasani, R., Hudson, D. A., Adeli, E., Altman, R., Arora, S., von Arx, S., Bernstein, M. S., Bohg. J., Bosselut, A., Brunskill, E., Brynjolfsson, E.,, Buch, S., Card, D., Castellon, R., Chatterji. S., Chen, A. S., Creel, K. A., Davis, J., Demszky, D., Donahue, C., Doumbouya, M., Durmus, E., Ermon, S., Etchemendy, J., Ethayarajh, K., Fei-Fei, L., Finn, C., Gale, T., Gillespic, L. E., Goel, K., Goodman, N. D., Grossman, S., Guha, N., Hashimoto, T., Henderson, P., Hewitt, J.. Ho. D. E., Hong. J., Hsu, K., Huang, Ieard, F., Jain, Jurafsky, D., Kalluri, P., Karamcheti, S., Keeling, G., Khani, F,, Khattab, O., Koh, P. W., Krass, M. S., Krishna, R. Kuditipudi, R., Kumar, A., Ladhak, F., Lee, M., Lee, T., Leskovee, J., Levent, I., Li, X. L.., Li, X., Ma, T., Malik, A., Manning, C. D., Mirchandani, S. P., Mitchell, E., Munyikwa, 7., Nair, S., Narayan, A., Narayanan, D., Newman, B., Nie, A., Niebles, J. C., Nilforoshan, H.. Nyarko, J. F., Ogut, G., Orr L., Papadimitriou. I., Park, J. S., Piech, C., Portelancee E., Potts, C.. Raghunathan, A., Reich, R. Ren, H., Rong. F., Roohani, Y. H.. Ruiz, C., Ryan, J., R' C., Sadigh, D.. Sagawa, S., Santhanam, K., Shih, A., Srinivasan, K. P., Tamkin, A., Taori, R., Thomas, A. W., Tramèr, F., Wang, R. E., Wang, W., Wu, B., Wu, J., Wu, Y., Xie, S. M., Yasunaga, M., You, J., Zaharia, M. A., Zhang, M., Zhang, T., Zhang, X., Zhang, Y., Zheng, L., Zhou, K., and Liang, P. On the opportunities and risks of foundation models. ArXiv, 2021. URL htttps:/crrft
stanforddd..... http:////rf...ttfor.....oo edu/ assets/report...t.... pdf
[11] Boussinesq, J. V. Essai sur la théorie des caux courantes. Mémoires présentés par divers savants à l'Académie des Sciences, 23, 1877.
[12] Brackbill, J. U. and Ruppel, H. M. Flip: A method for adaptively zoned, particle-in-cell calculations of fluid flows in two dimensions. Journal of Computational physics, 65(2):314-35 343, 1986.
[13] Brandstetter, J., Berg. R. d., Welling, M., and Gupta, J. K. Clifford neural layers for pde modeling. arXiv preprint arXiv:2209.049.4 2022.
[14] Brandstetter, J., Hesselink, R., van der Pol, E., Bekkers, E. J., and Welling, M. Geometric and phys- ical quantities improve e(3) equivariant message passing. In International Conference on Learning Representations, 2022.
[15] Brandstetter, J., Worrall, D., and Welling, M. Message passing neural pde solvers. arXiv preprint DXXi::2002003366 2022.
[16] Bryutkin, A., Huang, J., Deng, Z., Yang, G.. Schönlieb, C.-B., and Aviles-Rivero, A. Hamlet: Graph transformer neural operator for partial differential equations. arXiv preprint arXiv:2402.03544 2024.
[17] Calvello, E., Kovachki, N. B., Levine, M. E., and Stuart, A. M. Continuum attention for neural operators. CoRR, abs/2406. 06486, 2024.

---

[18] Cao, S. Choose a transformer: Fourier or galerkin. Advances in neural information processing systems. 34:24924 24940, 2021.
[19] Carey, N., Zanisi, L.,, Pamela, S., Gopakumar, V., Omotani, J., Buchanan, J., and Brandstetter, J. Data efficiency and long term prediction capabilities neural operator surrogate models of core and edge plasma codes. arXiv preprint arXiv:2402.08550 2024.
[20] Chang, A. X., Funkhouser, T. A., Guibas, L. J., Hanrahan, P., Huang. Q. Li, Z.. Savarese, S., Savva, M., Song. S., Su, H., Xiao, J., Yi, L., and Yu, F. Shapenet: An information-rich 3d model repository. CoRR, abs/1512.000122 2015.
[21] Chen, X., Liang, C., Huang, D., Real, E.. Wang, K., Liu, Y., Pham, H., Dong. X., Luong, T.. Hsich, C.. Lu. Y., and Le. Q V. Symbolic discovery optimization algorithms. CoRR, abs/2302.06675. 2023.
[22] Colagrossi, A. and Maurizio, Numerical simulation of interfacial flows by smoothed particle hydrody- namics. Journal of Computational Physics, 191, 2003.
[23] Crespo, A., Dominguez, J. M., Barreiro, A., Gomez-Gasteira, M., and D, R. B. Gpus, a new tool of acceleration in cid: efficiency and reliability on smoothed particle hydrodynamics methods. PLOS ONE, 6, 2011.
[24] Cundall, P. A. and Strack, O. D. A discrete numerical model for granular assemblies. geotechnique, 29 (1):47-65, 1979.
[25] Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805 2018.
[26] Dosovitskiy, A.. Beyer, L., Kolesnikov, A., Weissenborn. D., Zhai, X., Unterthiner, T.,, Dehghani, M., Minderer, M., Heigold, G., Gelly, S., Uszkoreit, J. and Houlsby, N. An image is worth 16x16 words: transformers for image recognition scale. In ICLR. OpenReview.net, 2021.
[27] Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A., Yang, A., Fan, A.. Goyal, A., Hartshorn, A., Yang, A., Mitra, A., Sravankumar, A., Korenev, A., Hinsvark, A.. Rao, A., Zhang. A., Rodriguez, A., Gregerson, A., Spataru, A., Rozière, B., Biron, B., Tang. B., Chern, B. Caucheteux, C., Nayak. C., Bi. C., Marra, C., McConnell, C., Keller, C., Touret, C., Wu, C., Wong.
C., Ferrer, C. C., Nikolaidis, C., Allonsius, D., Song, D., Pintz, D., Livshits, D., Fsiobu, D., Choudhary, D., Mahajan, D., Garcia-Olano, D., Perino, D., Hupkes, D., Lakomkin, E.,, AlBadawy, E.,, Lobanova, E., Dinan, F., Smith, M., Radenovic, F., Zhang, F., Synnaeve, G., Lee, G., Anderson, G. L.., Nail, G., Mialon, G., Pang. G.. Cucurell, G., Nguyen, H., Korevaar, H., Xu, H., Touvron, H., Zarov, I., Ibarra, I. A., Kloumann, I. M.. Misra, I., Evtimov, I., Copet, J., Lee, J., Geffert, J., Vranes, J., Park, J., Mahadeokar, J.. Shah, J., van der Linde, J., Billock, J., Hong, Lee, Fu, J., Chi, J.. Huang, J., Liu, J., Wang, Yu, J., Bitton, J., Spisak, Park, Rocca, J., Johnstun, J., Saxe, J., Jia, J., Alwala, K. V., Upasani, K., Plawiak, K., Li, K., Hcafield, K., Stone, K., and ct al. The Ilama herd of models. CoRR, abs/2407. 21783, 2024.
[28] Gilmer, J., Schoenholz, S. S., Riley, P. F., Vinyals, O., and Dahl, G. E. Neural message passing for quantum chemistry. In International conference on machine learning, pp. 1263-1274 PMLR, 2017.
[29] Gingold, R. A. and Monaghan. J. J. Smoothed particle hydrodynamics: theory and application to non-spherical stars. Monthly notices of the royal astronomical society, 181 1977.
[30] Guo, X.. Li. W., and lorio, F. Convolutional neural networks for steady flow approximation. In Proceedings of the 22nd ACM SIGKDD international conference on knowledge discovery and data mining, pp. 481 191-490,70 2016.
[31] Gupta. J. K. and Brandstetter, J. Towards multi-spatiooemppoall-caal generalized pde modeling. arXiv preprint ar 2209 15616, 2022.
[32] Hanjalic, K. and Brian, L. A reynolds stress model of turbulence and its application to thin shear flows. Journal of Fluid Mechanics, 52(4):609 638, 1972.
[33] Hao, Z., Wang. Z., Su, H., Ying, C., Dong, Y., Liu. S., Cheng, Z., Song. J., and Zhu, Gnot: A general neural operator transformer operator learning. International Conference on Machine Learning, pp- 125661 PMLR, 2023.
[34] Hao, Z., Su. C., Liu, S., Berner, J., Ying. C., Su, H., Anandkumar, A., Song, J., and Zhu, J. DPOT: auto- regressive denoising operator transformer for acale PDE pre-training. In ICML. OpenReview.net. 2024.
12

---

[35] Harada, T. Koshizuka, S., and Kawaguchi, Y. Smoothed particle hydrodynamics gpus. Computer Gruphics International, pp. 63-70, 2007.
[36] He, K., Chen, X., Xie. S., Li, Y., Dollár, P., and Girshick. R. Masked autoencoders are scalable vision learners. In IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR). pp. (6000-16009, 2022.
[37] Hendrycks, D. and Gimpel, K. Gaussian error linear units (gelus). arXiv preprint arXiv:1600.08415. 2016.
[38] Issa, R. Solution of the implicitly discretized fluid flow equations by operator-spliting. Journal of Computational Physics, 62, 1986.
[39] Jaegle, A.., Borgeaud, S., Alayrac, J.-B., Doersch, C., Ionescu, C., Ding. D., Koppula, S., Zoran, D.,
Brock, A., Shelhamer, E., et al. Perceiver io: general architecture for structured inputs outputs. In International Conference on Learning Representations, 2021.
[40] Jaegle, A., Gimeno, F., Brock, A., Vinyals, O., Zisserman, A., and Carreira, J. Perceiver: General perception with iterative attention, In International conference on machine learning, pp. 4651-445666 4651-4664. PMLR, 2021.
[41] Kingma, D. P. and Ba, J. Adam: method for stochastic optimization. In Bengio, and LeCun, (eds.). 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015.
[42] Kipf, T. N. and Welling, M. Semi-supervised classification with graph convolutional networks. In International Conference on Learning Representations, 2017.
[43] Kochkov, D., Smith, J. A., Alieva, A., Wang, Q., Brenner, M. P., and Hoyer, S. Machine learning- accelerated computational fluid dynamics. Proceedings of the National Academy of Sciences, 18(21): e2101784118, 2021.
[44] Kohl, G., Chen, L.-W., and Thuerey, N. Turbulent flow simulation using autoregressive conditional diffuuson models. arXiv preprint arXiv:2309.01744 2023.
[45] Kolmogorov, A. N. A refinement of previous hypotheses concerning the local structure of turbulence in viscous incompressible fluid at high reynolds number. Journal of Fluid Mechanics, 13(1):82 85, 1962.
[46] Kolmogorov, A. Local structure of turbulence in an incompressible fluid at very high reynolds numbers. CR Ad. Sei, UUSR, 30:305, 1941.
[47] Kovachki, N., Li, Z., Liu. B., Azizzadenesheli, K., Bhattacharya, K., Stuart. A.. and Anandkumar, A. Neural operator: Learning maps between function spaces. arXiv preprint 100.::2008.08888 2021.
[48] Lam, R.,
A., Willson, M., Wimsberger, P., Fortunato, M., Pritzel, A., Ravuri, S., Z., et al. Grapheast: Learning skillful medium-range global weather
Ewalds, T., Alet, F., forecasting. arXiv preprint arXiv:22 12794, 2022.
[49] Lanthaler, S., Li, Z., and Stuart, A. M. The nonlocal neural operator: Universal approximation. CoRR, has/2304.13221) 2023.
[50] Lee, K. and Carlberg, T. Deep conservation: A latent-dynamies model for exact satisfaction of physical conservation laws. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 277-285, 2021.
[51] Li, Z.,, Kovachki, N., Arizzadenesheli, K., Liu, B., Bhattacharya, K., Stuart, A., and Anandkumar, A. Fourier neural operator for parametric partial differential equations. arXiv preprint arXiv:2010.08895, 2020.
[52] Li, Z., Kovachki, N., Aeizzadenesheli, K., Liu, B., Bhattacharya, K., Stuart, A., and Anandkumar, Neural operator: Graph kernel network for partial differential equations. arXiv preprint 20X1v:2003.03-485 2020.
[53] Li, Z. Meidani, K., and Farimani, A. B. Transformer for partial differential equations' operator learning. arXiv preprint arXiv:2205.13677 2022.
[54] Li, Z., Kovachki, N. B., Choy, C., Li, B., Kossaifi, J., Otta, S. P., Nabian, M. A., Stadler, M., Hundt, C., Azizzadenesheli, K., et al. Geometry informed neural operator for large-scale 3d pdes, arXiv preprint 000502309.00588. 2023.
13

---

[55] Li. Z.,, Patil, S., Ogoke. F., Shu, D., Zhen, W., Schneier, M. Jr.. J. R. B., and Farimani, A. B. Latent neural PDE solver: reduced-order modelling framework for partial differential equations. CoRR, abs/2402.17853, 2024.
[56] Lienen, M., Hansen Palmmn, J., Lüdke, D., and Günnemann, S. Generative diffusion turbulent flows. arXiv preprint arXiv:2306.0177 2023.
[57] Lippe, P., Veeling, B., Perdikaris, P., Turner, R., and Brandstetter, J. Pde-refiner, Achieving accurate long rollouts with neural pdc solvers. Advances in Neural Information Processing Systems, 36, 2024.
[58] Loshchilov, I. and Hutter, SGDR: stochastic gradient descent with warm restarts. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net, 2017.
[59] Loshehilov, L. and Hutter, F. Decoupled weight decay regularization. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019.
[60] Lu, L., Jin, P., and Karniadakis, Deeponet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators. arXiv preprint arXiv:1910003193 2019.
[61] Lu, L., Jin, P., Pang, G., Zhang, Z., and Karniadakis, G. E. Learning nonlinear operators via deeponet based on the universal approximation thcorem of operators. Nature machine intelligence, 3(3):218 2222118-229 2021.
[62] Lucy, L. B. numerical approach the testing of the fission hypothesis. Astronomical Journal, vol. 82, Dec. 1977, p. 1013-1024,, 12:1013-1024, 1977.
[63] Lusch, B., Kutz, N., Brunton, S. L. Deep learning universal linear embeddings nonlinear dynamics. Nature communications, 91):4950, 2018.
[64] Mathieu, J. and Scott, J. An introduction to turbulent flow. Cambridge University Press, 2000.
[65] Mayr, A., Lehner, S., Mayrhofer, A., Kloss, C., Hochreiter, S., and Brandstetter, J. Boundary graph neural networks for 3d simulations. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pp.
2023.
[66] McCabe, M., Blancard, B. R.-S,, Parker, L., H., Ohana, R., Cranmer, M., Bietti, A., Fickenherg, M., Golkar, S., Krawezik, G., Lanusse, F., al. Multiple physics pretraining for physical surrogate models. arXiv preprint 10.092310.009944 2023.
[67] Merchant, A., Batzner, S., Schoenholz, S. S., Aykol, M., Cheon, G., and Cubuk, D. Scaling deep learning for materials discovery. Natwre, pp. 1-6, 2023.
[68] Moin, P. and Mahesh, K. Direct numerical simulation: a tool in turbulence research. Annual review of fluid mechanics, 30(1))539-578 1998.
[69] Monaghan, J. J. Smoothed particle hydrodynamics. Annual review of astronomyy and astrophysics, 30(1): 543 574, 1992.
[70] Monaghan, J. J. Smoothed particle hydrodynamics. Reports progress in physics, 68(8):1708 2005.
[71] Monaghan, J. and Gingold, R. A. Shock simulation by the particle method sph. Journal of computational physics, 52(2):374 389, 1983.
[72] Morris, J., Fox, P., and Zhu, Modeling low reynolds numher incompressible flows using sph. Journal of Computational Physics, 136, 1997.
[73] Morton, J., Jameson, A., Kochenderfer, M. and Witherden, Deep dynamical modeling and control unsteady fluid flows. Advances Neural Information Processing Systems, 31, 2018.
[74] Nair, V. and Hinton, G. E. Rectified linear units improve restricted boltzmann machines. In ICML, pp. 807- 814. Omnipress. 2010.
[75] Nguyen, T., Brandstetter, J., Kapoor, A., Gupta, J. K., and Grover, A. Climax: A foundation model for weather and climate. arXiv preprint arXiv:2301.10344 2023.
[76] Olver, P. J. Introduction to partial differential equations. Springer, 2014.
14

---

[77] Patankar, S. V. and Spalding, B. D. A calculation procedure for heat. mass and momentum transfer in three-dimensional parabolic flows. Int. J. of Heat and Mass Transfer, 15, 1972.
[78] Peebles, W. and Xie, S. Scalable diffusion models with transformers. In ICCV, pp. 4172-4182. IEEE, 2023.
[79] Perdikaris, A unifying framework for operator learning via neural fields, Dec 2023.
[80] Perez, E., Strub, F., de Vries, H., Dumoulin, V., and Courville, A. C. Film: Visual reasoning with general conditioning layer. pp.
AAAI Press, 2018.
[81] Pfaff, T., Fortunato, M., Sanahez-Gonzllez, A., and Battaglia, P. W. Learning mesh-based simulation with graph networks. arXiv preprint arXi::2010.03409 2020.
[82] Pope, S. B. Turbulent flows. Measurement Science and Technology, 12(1 (2(11):220-2022 2001.
[83] Raonić, B., Molinaro, R., Rohner, T.. Mishra, S., and de Bezenac, E. Convolutional neural operators. arXiv preprint ar arXiv:2302.01178 2023.
[84] Reynolds, O. Iv. on the dynamical theory of incompressible viscous fluids and the determination of the criterion. Philosophical transactions of the royal sociery of london.(a.), (186):133-164, 1895.
[85] Rombach, R., Blattmann, A., Lorenz, D., Esser, P., and Ommer, B. High-resolution image synthesis with latent diffusion models. In CVPR, pp.
IEEE, 2022.
[86] Ronneberger, O., Fischer, P., and Brox, T. U-net: Convolutional networks for biomedical image seg- mentation. In MICCAI (3), volume 9351 of Lecture Notes Computer Science, pp. 234-241. Springer, 2015.
[87] Goncaez-Gonzalez, A., Godwin, J., Pfaff, T., Ying, R.,, Leskovee, and Battaglia, P. Learning to simulate complex physics with graph networks. International conference on machine learning, pp. $459-8468. PMLR, 2020.
[88] Scarselli, F., Gori, M., Tsoi, A. C.,, Hagenbuchner, M., and Monfardini, G. The graph neural network model. IEEE transactions on neural networks, 20(1):61-88) 2008
[89] Scidman, J., Kissas, G., Perdikaris, P., and Pappas, G. Nomad: Nonlinear manifold decoders for operator learning. Advances in Neural Information Processing Systems, 35:5601-5613, 2022.
[90] Shen, Z., Zhang. M., Zhao, H., Yi, S., and Li, H. Efficient attention: Attention with linear complexities. In WACV, pp. 3530-3538. IEEF, 2021.
[91] Smagorinsky, J. General circulation experiments with the primitive equations: I. the basic experiment. Monthly weather review, 91(3):99-188 164, 1963.
[92] Stachenfeld, K., Fielding, D. B., Kochkov, D., Cranmer, M., Pfaff, T., Godwin, J., Cui, C., Ho, S.. Battaglia, P., and Manchez-Gonzalee A. Learned simulators for turbulence. In International Conference Learning Representations, 2022.
[93] Stachenfeld, K. L.,, Fielding, D. B., Kochkov, D., Cranmer, M. D., Pfaff, T., Godwin, J., Cui, C., Ho, S., Battaglia, P. W., and Sanchez- Sanchez-Gonzalez Gonzalez, A. Learned coarse models for efficient turbulence simulation. CoRR, abs/21 15275, 2021.
[94] Sulsky, D., Chen, Z., and Schreyer, H. L. A particle method for history-dependent materials. Computer methods in applied mechanicy and engineering, 118(1-2):179-19, 1994.
[95] Takamoto, M., Praditia, T., Leiteritz, R., MacKinlay, D., Alesiani, F,, Pflüger, D., and Niepert, M. Pdehench: extensive benchmark for scientific machine learning. Advances Neural Information Processing Systems, 35:1596-1611, 2022.
[96] Tancik, M., Srinivasan, P. P., Mildenhall, B., Fridovich-Kaan S., Raghavan, N., Singhal, U., Ramamoorthi, R., Barron, J. T., and Ng, R. Fourier features let networks learn high frequency functions in low dimensional domains. In NeurIPS, 2020.
1971 Taylor, G. I. and Green, A. E. Mechanism of the production of small eddies from large ones. Proceedings of the Royal Society of London. Series A- A-Mathematical and Physical Sciences, 158(895):499-577 1937.
[98] Temam, R. Navier Stokes equations: theory and numerical analysis, volume 343, American Mathematical Soc., 2001.
15

---

[99] Thuerey, N., Holl. P., Mueller, M., Schnell, P., Trost, F., and Um. K. Physics-based deep learning, arXiv preprint urXiv:2109.05237, 2021.
[100] Toshev, A. P., Galletti, G., Brandstetter, J., Adami, S., and Adams, N. A. Learning lagrangian fluid mechanics with equiivarar graph neural networks, arXiv preprint arXiv:2305. 15603, 2023.
[101] Toshev, A. P.. Galletti, G., Fritz, F., Adami, S., and Adams, N. A. Lagrangebench: A lagrangian fluid mechanies benchmarking suite. In 37th Conference Neural Information Processing Systems (NeurIPS 2023) Track on Datasets and Benchmarks, 2023.
[102] Toshev, A. P.. Erbesdobler, J. A.. Adams, N. A., and Brandstetter. Neural sph: Improved neural modeling lagrangian fluid dynamies. arXiv preprint urXiv:2402.06275 2024.
[103] Tran, A., Mathews, A. Xie, L.,, and Ong, C. S. Factorized fourier neural operators. In ICLR. OpenReview.net, 2023.
[104] Tsai, T.-P. Lectures on Navier-Stokes equations, volume 192. American Mathematical Soc., 2018.
[105] Umetani, N. and Bickel, B. Learning three-dimensional flow for interactive acrodynamic design. ACM Trans. Graph., (7,4):89, 2018.
[106] Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, t., and Polosukhin, I. Attention is all you need. Advances neural information processing systems, 30, 2017.
[107] Vinuesa, R. and Brunton, S. L.. Enhancing computational fluid dynamics with machine learning. Nature Computational Science, 26))358-366, 2022.
[108] Wang. S., Scidman, J. H., Sankaran, S., Wang, H., Pappas, G. J., and Perdikaris, P. Bridging operator learning and conditioned neural fields: A unifying perspective. CoRR, abs/2405.13995 2024.
[109] Weiler, M., Geiger, M., Welling, M., Boomsma, W., and Cohen, T. S. 3d steerable enns: Learning rotationally equivariant fcatures volumetric data, Advances in Neural Information Processing Systems, 31, 2018.
[110] Weller, H. G., Tabor, G., Jasak, H., and Fureby, C. tensorial approach computational continuum mechanics using orject-oriented techniques. Computers in physics, 12(6))620-631, 1998.
[111] Wiewel, S., Becher, M., and Thuerey, Latent space physics: Towards learning the temporal evolution fluid flow. Computer graphics forum, volume 38, pp. 71-82. Wiley Online Library, 2019.
[112] Wilcox, D. Formulation of the k-omega turbulence model revisited. AIAA Journal, 46( (6(11):28232288) 2008.
[113] Woo, S., Debnath, S., Hu, R., Chen, X., Liu, Z., Kweon, and Xie, S. Convnext V2: co-designing and scaling convnets with masked autoencoders. In IEEF/CVF Conference on Computer Vision and Pottern Recognition, CVPR 2023, Vancouver; BC, Canada, June 17-24, 2023, pp- 16133-16142. IEFF, 2023.
[114] Wu, H., Luo, H., Wang, H., Wang. J., and Long, M. Transolver: A fast transformer solver for pdes on general geometries. arXiv preprint arXiv:2402.0243 02366, 2024.
[115] Wu, Y. and He, K. Group normalization. In ECCV (13), volume 11217 Lecture Notes in Computer Science, pp. 3-19. Springer, 2018.
[116] Xu, M., Han, J., Lou, A., Kossaifi, J., Ramanathan, A., Azizzadenesheli, K., Leskovee, J., Ermon, S., and Anandkumar, A. Equivariant graph neural operator for modeling 3d dynamics. arXiv preprint arXiv:2401.1 1037, 2024.
[117] Yang, E., Bui, H. H., De Sterek, H.. Nguyen, G. D., and Bouazza, A. A scalable parallel computing sph framework for predictions of geophysical granular flows, Computers and Geotechnics, 121:103474, 2020. ISSN 0266-352X.
[1 [118] Zeni, C., Pinsler, R., Zügner, D., Fowler, A., Horton, M., Fu, X., Shysheya, S., Crabbé, J., Sun, I.,, Smith, J., al. Mattergen: generative model for inorganic materials design. arXiv preprint arXiv:23/2.03667, 2023.
[119] Zhang, X., Wang, I.., Helwig, J., Luo, Y., Fu, C., Xie, Y., Liu, Min, Y., Xu, Z., Yan, K., et al. Artificial intelligence for science quantum, atomistic, and continuum systems. arXiv preprint arXiv:2307.08423 2023.
16

---

A Limitations and Future work
Latent rollout. We show that the latent rollout can be enabled via simple end-to-end training procedure, but ultimately think that can be improved by more delicate training procedures such as two stage procedure akin to diffusion models 85 55 As there are many possible avenues to apply or improve the latent rollout, it would exceed the scope of this paper and we therefore leave it for future work.
Generalization beyond fluid dynamics. We show the generalization capabilities of UPTs across different simulation types and leave the application of UPT' to other domains for future work. We additionally show that UPTs are neural operators and since neural operators have been shown generalize well across domains, this should also hold for UPTs.
Large-scale Lagrangian simulations. A particularly intriguing direction is apply UPTs large- scale Lagrangian simulations. However, there do not exist readily available large-scale Lagrangian datasets. Therefore, consider large-scale Lagrangian simulations beyond the scope of this paper as generating such a dataset requires extensive domain knowledge as well as engineering effort. Note that we show the favorable scaling properties of UPTs the transient flow dataset, which contains up to 59K mesh points. In this setting, we train with a distributed setup of up to 32 A 100 GPUs.
Unifying Lagrangian and Eulerian simulations. UPTs can encode both Lagrangian and Fulerian simulations. It is therefore a natural future direction to exploit those different modalities. Especially, since particle- and grid-based simulations underlying dynamics. It is however to note that our method is designed primarily with the purpose of efficient scaling, and multi- multi-modality training follows side-concept thereof.
used to describe different phenomena but model similar
B
Computational fluid dynamics
This appendix discusses the Navier-Stoker equations and selected numerical integration schemes, which are related the experiments presented in this paper. This is an complete introduction into the field, but rather encompasses the concepts which are important to follow the experiments discussed in the main paper.
First, we discuss the different operators of the Navier-Stoker equations, relate compressible and incompressible formulations, and discuss why Navier-Stokes equations difficult to solve. Second, we discuss turbulence as one of the fundamental challenges of computational fluid dynamics (CFD), explicitly working the difference for two dimensional and three dimensional turbulence modeling. Third, we introduce seyyolds-averageee for steady-state turbulence modeling, and the SIMPLE and PISO algorithms steady-state and transient solvers, respectively. Finally, we discuss Lagrangian discretization schemes, focusing on smoothed particle hydrodynamics (SPH).
Navier-Stokes equations (RANS) as an numerical approach
B.1 Navier-Stokes equations
In our experiments, we mostly work with the incompressible Navier-Stokes equations 988 In two spatial dimensions, the evolution equation the velocity flow field u [0, T x C
R3 - R², ull, 11, x22) (w1((e 11, I2), (2(1, with internal pressure p(11,T2) = (pr (#1, x2), P22222221p (#1, x2)) and external force field f x2) = (fi(21, fe(x,, is given via:
Oui
Oui
Oul
8ul m1
Juj
Op1 f1
U1
U2 drz
+ p
+
Øt
dx
dx3
dri
Ouz
Ou2
Oug
8
02
0p2 + f2
ui
42 0.22
+ P
+
(3)
01
0x
0x7
0u7
Orz
Own
Ou2
0a1
002
17

---

This system is usually written in shorter and more convenient form:
iu
uu Vut PutuVer - Vp +
0t
(4)
0 = V u
where
Oun
Ou2
dzi
dzi
Vu
U1
U2 )
(5)
du1
Ou2
Oxg
Orz
called the convection, i.e., the rate of change of a along u. and
4 u =
a
0
az
U1
02
0
u1
(6)
dx1
Ø
0x1
Orz
the viscosity, i... the diffusion or net movement of u with viscosity parameter p. The constraint
Ox1 (oe
(
Ou
Ouz
U1
+
0
dx²
U22
de
dxz
yields mass conservation of the Navier-Stoker equations.
The incompressible version of the NS equations above are commonly used in the study of water flow, low-speed air flow, and other scenarios where density changes are negligible. Its compressible counterpart additionally accounts for variations in fluid density and temperature which is necessary to accurately model gases at high speeds (starting in the magnitude speed of sound) or scenarios with significant changes in temperature. These additional considerations result in a more complex momentum and continuity equation.
Why are the Navier-Stokes equations difficult to solve? One major difficulty is that no equation explicitly models the unknown pressure field. Instead, the conservation of mass V implicit constraint on the velocity fields. Additionally, the nonlinear nature of the NS equations makes them notoriously harder solve than parabolic hyperbolic PDEs such as the heat wave equation, respectively. Also, the occurrence of turbulence as discussed in the next subsection which due to the fact that small and large scales are coupled such that small errors, can have large effect on the computed solution.
is
Computational fluid dynamics (CFD) uses numerical schemes to discretize and solve fluid flows. CFD simulates the free-stream flow of the fluid, and the interaction of the fluid (liquids and gases) with surfaces defined by boundary conditions. As such, CFD comprises many challenging phenomena, such as interactions with boundaries, mixing of different fluids, transonic, i.e, coincident emergence of subsonic and supersonic airflow, or turbulent flows.
B.2 Turbulence
Turbulence 82 64 is one of the key aspects of CFD and refers chaotic and irregular motion of fluid flows, such as air or water. It is characterized by unpredictable changes in velocity, pressure, and density within the fluid. Turbulent flow is distinguished from laminar flow, which is smooth and orderly. There are several factors that can contribute to the onset of turbulence, including high flow velocities, irregularities in the shape of surfaces over which the fluid flows, and changes in the fluid's viscosity. Turbulence plays significant role in many natural phenomena and engineering applications, influencing processes such as mixing, heat transfer, and the dispersion of particles in fluids. Understanding and predicting turbulence is crucial in fields like fluid dynamics, aerodynamics, and meteorology. Turbulent flows occur at high Reynolds numbers
pVL.
Re
(8)
1
18

---

where is the density of the fluid, v/ is the velocity of the fluid, L is the linear dimension, e.g., the width of a pipe, and y is the viscosity parameter. Turbulent flows are dominated by inertial forces, which tend to produce chaotic eddies, vortices and other flow instabilities.
Turbulence in 3D turbulence in 2D. For high Reynolds numbers, i.e., in the turbulent regime, solutions to incompressible Navier-Stokes equations differ in three dimensions compared to two dimensions by phenomenon called energy cascade. Energy cascade orchestrates the transfer of kinetic energy from large-scale vortices to progressively smaller scales until it is eventually converted into thermal energy. three dimensions, this energy cascade is responsible for the continuous formation of vortex structures at various scales, and thus for the emergence of high-frequenge features. In contrast to three dimensions, the energy cascade is inverted in two dimensions, i.e.. energy is transported from smaller to larger scales, resulting in more homogeneous, long-lived structures. Mathematically the difference between three dimensional and two dimensional turbulence modeling can be best seen by rewriting the Navier-Stokes equations in vorticity formulation, where the vorticity w is the curl of the flow field, i.e., w V u, with x here denoting the cross product. This derivation can be found in many standard texts Navier-Stokes equation e.g. Tsai [T04 Section 1.4.], and for the reader's convenience we briefly repeat the derivation here. We start using Eq. 4 setting = 0 for simplicity:
du
u Vu Vu+uvvu Vp
(9)
0L
1
2
u
+ V Vp.
(10)
where we applied the dot product rule for derivatives of vector fields:
u = u VutuxVxu
(11)
Next, we take the curl of the right and the left hand side:
① V x u)
v x v(u u
V
V Vp
(12)
Ot
dw
V (u w pV
(13)
Ot
using the property that the curl of gradient is zero. Lastly, using that the divergence of the curl is again zero (i.e. v-w = 0), via the identity Vx(u x w) = (w V) u uu V) w w V
we obtain:
dw
+ (u V)w (w VVu+pV2
(14)
Dw (w V)u+wVw
(15)
Dt
Dw Ow
with the material derivative
+ (u V.
Di
Ø1
For three dimensional flows the material derivative of each component wi can be written as
Dw
Ou: Wi Day
8o
Σ
+ /1 Dx, diffusion
(16)
DI.
vortex turning and stretching
for example picking i = 2:
Dw2
Ouz
Ou2 W2
Ouz Wi dzs
Puz @x10r1
Pus Oxzdzz
Pur Oxzdrs
W1 dx
+
(17)
dxy vortex stretching
vortex turning
vortex toming
dillusion
19

---

For two-dimensional flows, as represented by
U1
0
0
U2
and
(18)
auz 0m
Ous Orz
the terms for vortex turning and vortex stretching vanish, as is evident the material derivative of 4a,, given by
Dws
0us
dus
dug
8ws 8ws 0c, Ozi 2x2220
Pww
W1
w2
wB
+ lu
+
+
(19)
DI
0m
Orz
Oxa
difluusion
Consequently, the length and the angle of vortex do not change in dimensions, resulting in homogeneous, long-lived structures.
There's been several approaches to model turbulence with the help of Deep Learning. Predominantly
models for two dimensional scenarios have been suggested [8 15 43 44]. Due to the higher complexity (as discussed above) and memory- and compute costs comparably less work was done in the 3D case 92 566.
B.3 Numerical modeling based on Eulerian discretization schemes
Various approaches are available for numerically addressing turbulence, with different schemes suited different levels of computational intensity. Among these, Direct Numerical Simulation (DNS) [68] stands out as the most resource-intensive method, involving the direct solution of the unsteady Navier-Stokes equations.
DNS is renowned for its capability to resolve even the minutest eddies and time scales present in turbulent flows. While DNS does necessitate additional closure equations, significant drawback is its high computational demand. Achieving accurate solutions requires the utilization of very fine grids and extremely small time steps. This is based on Kolmoqorov 46] and Kolmogorov 45] which give the minimum spacial and temporal scales that need to be resolved for accurately simulating turbulence. The discretization- scallesssssclsseeeeeeeee for both time and space needed for 3D-problems become very small and simulations computationally extremely expensive.
This computational intensity of DNS resulted in the development of turbulence modeling, with Large Eddy Simulations (LES) (91 and Reynolds-averaged Navier-Stokes equations (RANS) [84] being two prominent examples. LES aim reduce computational cost by neglecting the computationally expensive smallest length scales in turbulent flows. This is achieved through a low-pass filtering of the Navier-Stokes equations, effectively removing fine-scale details via time- and spatial-averaging. In contrast, the Reynolds-averaged Navier-Stokes equations (RANS equations) are based on time- averaging. The foundational concept is Reynolds decomposition, attributed Osborne Reynolds, where an instantaneous quantity is decomposed into its time-avee averaged and fluctuating components.
Turbulent scenarios where the fluid conditions change time are example of transient flows. These are typically harder to model/solve than steady state flows where the fluid properties exhibit only negligible changes time.
In the following, we discuss Reynolds-averaged Kavier-Stokes equations (RANS) as numercial approach for steady-state turbulence modeling, and the SIMPLE and PISO algorithms steady-state and transient numerical solvers, respectively.
B.3.1 Reynolds-averaged Navier-Stokes equations
Reynolds-averaged Havier-Stoken equations (RANS) are used to model time-averaged fluid properties such as velocity or pressure that result steady state which does not change over time. Writing in Einstein notation we get
Ou: Øt
Ous
S
ap Dx;
Uij Dz,
+
+ ft
Oz "z,0x;
Jus
(20)
Dr.
20

---

Taking the Reynolds decomposition [4] g(t, := g(21, 22) + 21, 22) with (()1,22) := lim 1/T So g I2, t)dt being the timeaaeerage on T] of a scalar valued function g. and splitting each term of both equations into it's time-averaged and fluctuating part we get
(ū; + u/))
0(ū + u/)
i u/) + 0z,0z,
o(p+p) dr.
(ū, + u/)
+ (Ji - 11)
01
dxy
0 aüu + u!))
Oz,
Time-averaging these equations together with the property that the time-average of the fluctuating parts equals zero results in
θū.
+ u 0u
dp dr,
új 0z
+ fr
(21)
0x,
dz 0z,0x,
0ū:
Ox:
Using the the mass conserving equation 200 the momentum equation 21
can be rewritten
0 Oer
Ou;
Juj
=j
pb.
where -u is called the Reynolds stress which needs further modeling to solve the above equations. More specifically, this referred as the Closure Problem which led to many turbulence models such as k-c [32] or k-w [112]. Analogously the 3D RANS equations can be derived. This turbulence model consists of simplified equations that predict the statistical evolution of turbulent flows. Due to the Reynolds stress there still remain velocity fluctuations in the RANS equations. To get equations that contain only time-averaged quantities the RANS equations need to closed by modeling the Reynolds stress as function of the mean flow such that any reference to the fluctuating parts is removed. The first such approach led to the eddy viscosity model Stokes:
11
for 3-d incompressible Navier
0u0: Ozy
Oug Ox;
2
u/u =
k6. 3
with the turbulence eddy viscosity in O and the turbulence kinetic energy h: 1/2 uiuu, based Boussinesq's 's hypothesis that turbulent shear stresses act in the same direction as shear stresses due to the averaged flow. The k-e model employs Boussinesq's hypothesis by using comparably low-cost computations for the eddy viscosity by means of two additional transport equations for turbulence kinetic energy and dissipation e
0(pk) Ōt Oz;
(pkuu)
1/ Ok
+ 2v, Eij - pe
0x,
σh Ōx.,
OpE))
(peu.) θ
V/3 de
2
+ Cle i -2v 2vr Eij C2.P
+
at
Ox:
0x;
o. 0x,
with the rate of deformation Eij, eddy viscosity V4 = =CCkkp/ and adjustable constants Co, Cie, C2e, or, σk. our experiment section solutions of simulations based on the RANS k-e turbulence model are used as ground truth where the quantity of interest is the pressure field given the shape of vehicle.
B.3.2 SIMPLE and PISO algorithm
Next, we introduce two popular algorithms that try to solve Eq. numerically, where the main idea to couple pressure and velocity computations. Specifically, we will discuss the SIMPLE (Semi- Equations) [77] and PISO (Pressure Implicit with Splitting of
Implicit Method for Pressure-Linde Operators) 38 algorithm, which are implemented OpenFOAM as well. The essence of these algorithms are the following four main steps:
21

---

1. In first step, the momentum equation is discretized by using suitable finite volume dis- cretizations of the involved derivatives. This leads to linear system for for given pressure gradient Vp:
Mu Vp.
(22)
However, this computed velocity field u does not yet satisfy the continuity equations V u = 0.
2. In next step, us denote by A the diagonal part of M and introduce H. so that
H = Au Mu.
(23)
Taking into account Mu = - Vu allows to easily rearrange Eq. 23 for u, since A is easily invertible. This gives us:
u = - - A
(24)
3. Now we substitute Eq.
(24 into the continuity equation yielding
V (A--vp)) = V (A-
(25)
which is a Poisson equation for the pressure p that can solved again numerically.
4. This pressure field can now be used to correct the velocity field by again applying Eq. 24 This is the pressure- crossurecoorecctor stage and this u now satisfies the continuity equation. Now, however, the pressure equation no longer valid, since H. depends on u as well. A way out here is iterating these procedures, which is the main idea of the SIMPLE and PISO algorithm.
The SIMPLE algorithm just iterates steps 1-4 several times, which is then called an outer corrector loop. In contrast, the PISO algorithm solves the momentum predictor Mu -Vp (i.e. step 1) only once and then iterates steps 2-4, which then result in an inner loop. In both algorithms, in case the meshes are non orthogonal, step 3, i.e. solving the Poisson equation, can also be repeated several times before moving to step 4. For stability reasons, the SIMPLE algorithm is preferred for stationary equations (since it implicitly promotes under-relaxation, whereas in case of timeddeendeent flows, one usually considers the PISO algorithm (since in practice, for each timestep several thousands of iterations are usually required until convergence, which computationally very expensive).
B.4 Lagrangian discretization schemes
In contrast to such grid- and mesh-based representations, in Lagrangian schemes, the discretization is carried out using finitely many material points, often referred to particles, which move with the local deformation of the continuum. Roughly speaking, there are three families of Lagrangian schemes: discrete element methods (24]. material point methods 994 12]. and smoothed particle hydrodynamics (SPH)
70
B.4.1 Smoothed particle hydrodynamics
The core idea behind SPH is to divide fluid into set of discrete moving elements referred as particles. These particles interact through using truncated radial interpolation kernel functions with characteristic radius known the smoothing length. The truncation is justified by the assumption of locality of interactions between particles which allows to approximate the properties of the fluid at any arbitrary location. For particle é its position æ; is given by its velocity u4 such that da; u, The modeling of particle interaction by kernel functions implies that physical properties of any particle can be obtained by aggregating the relevant properties of all particles within the range of the kernel. For example, the density Px of particle can be expressed as p(æ:,) E, Pj V,V, ([] æzyl, h) with W being kernel, h its smoothing length, and V, the volumes of the respective partieles. Rewriting the seaaklycoomressibeeee Navier-Stokes equations with quantities as (4) and additional Reynolds number Re and density P
Ou
1 22
1 -
-
01
Re
P
0p
= PVV 4
Ot
22

---

in terms of these kernel interpolations leads to system of ordinary differential equations (ODEs) for the particle accelerations 72 where the respective velocities and positions can be computed by integration. One of the advantages of SPH compared to Fulerian discretization techniques is that SPH can handle large topological changes as no connectivity constraints between particles are required, and advection is treated exactly [100. 71]. The lack of mesh significantly simplifies the model implementation and opens up more possibilities for parallclization compared to Eulerian schemes [35 23] However, for accurately resolving particle dynamics small scales large number of particles is needed to achieve a resolution comparable to grid-based methods specifically when the metric of interest is not directly related to density which makes SPH more expensive [22]. Also, setting boundary conditions such as inlets, outlets, and walls is more difficult than for grid-based methods.
seemingly great fit for particle-based dynamics are graph neural networks (GNNs) graph-based latent space representations. In many cases, predicted accelerations at the nodes are numerically integrated model the time evolution of the many-particle systems 87
with
102
C Justification that UPTs are universal neural operators
We provide a brief sketch of how universality can be established for our architecture. For transformer- based neural operators, universal approximation has been recently demonstrated in (17), Section 5. The arguments in this work are heavily based on [49]. which establishes that nonlinearity and under nonlocality are crucial for universality. By demonstrating that the attention mechanism appropriate weight choices, function as an averaging operator, the results from 49] are directly applicable. For detailed proof, refer to Theorem 22 in [17.. Our algorithm fits within this framework as well: we employ nonlinear, attention-based encoders and decoders (as allowed by the results of 49). Section 2) and utilize attention layers in the latent space.
D Experimental details and extended results
D.1 General
All experiments are conducted mostly on A100 GPUs. For large-scale experiments we use two research clusters equipped with either 8xA 10000000G nodes or 4xA 100000-64GB nodes. For smaller experiments and evaluations, we use mix of internal servers equipped with varying numbers of A100-40GB or A100-80GB cards.
We estimate the total number of GPU-hours (mostly A A100-hours) used for this project to be 45K. This includes everything from model architecture design, exploratory training runs, investigating training instabilities to training/evaluating baseline models and UPTs.
All experiments linearly scale the learning rate with the batchsize, exclude normalization and bias parameters from the weight decay, follow a linear warmup > cosine decay learning rate schedule (58) and use the AdamW 41 59) optimizer. Transformer blocks follow standard pre-norm architecture
as used in ViT 26
D.2 Feature modulation
We apply feature modulation condition models time/velocity. We use different forms of feature
modulation in accordance with existing works depending on the model type. FiLM [80] is used for U-Net, the "Spatial-Spectraa" conditioning of [1 used for FNO based models and DiT [78] is used for transformer blocks. For perceiver blocks, we extend the modulation of DiT to seperately modulate queries from keys/values.
D.3 Experiments regular grid datasets
We use two datasets from 31 the "Navier-Stokes 2D" and "ShallowWater-22D" to compare against baseline transformer architectures. Baseline results, data preprocessing and benchmark setups are taken from DPOT (34] and CViT [108].
23

---

As these datasets contain only regular grid data, we make the following modifications to UPT: (i) we replace the supernode pooling with a patch embedding (ii) we remove the perceiver pooling in the encoder (iii) we decode patchwise. These modifications were done (i) patch embedding is more efficient than supernode pooling for regular grid data (ii) these datasets are small-scale datasets, so compressing the latent space to make the model more efficient is not necessary (iii) as the decoding positions are always the same, there is no reason to decode position-wise as the missig variability of positions in the dataset makes superresolution impossible and patch-wise decoding is more efficient than position-wise decoding. This setup is conceptually similar to CViT [08].
We use a peak learning rate of le-4 with 10% warmup and cosine decay schedule afterwards, patchsize 8, batchsize 256 and train for 1000 epochs. We use hidden dimension 96 and 12 blocks for 2M parameter models and hidden dimension 256 with 14 blocks for 13M parameter models to best match model sizes of previous models. Additionally, we mask out between 0% and 5% of the input patches via an attention mask in the transformer and perceiver blocks. Similar performances could be achieved with less epochs, but our flexible masking strategy allows long training without overfitting.
D.3.1 NavierStokes-2D dataset
We comparisons against different Transformer baselines regular gridded Navier-Stokes equations dataset [31 in Table 2 (baseline results and evaluation protocol taken from DPOT [4]). UPT outperforms all compared methods, some of which specifically designed for regularly gridded data.
As baseline transformers often train small models, we compare on a small scale, where UPT signifi- cantly outperforms other models. We also compare on larger scales, where UPT again outperforms competitors.
Table Comparison on a regular grid Navier-Stokes dataset [3 (a) UPT outperforms competitor methods that train only small models by a large margin. (b) UPT also performs well on larger model sizes, outperforming competitors even if they train much larger models or pre-train (PT) on more data followed by fine-tuning (FT) on the Navier-Stokes dataset.
(a) Comparison on small model sizes
(b) Comparison on larger model sizes
Model
# Params
Rel. L.2 Error
Model
# Params
Rel. L.2 Error
FNO [51
0.5M
9.12%
DPOT-Ti 34
7M
12.50%
FENO
103
1.3M
8.39%
DPOT-S 34
30M
9.91%
GK-T
18
1.6M
9.52%
DPOT-L (PT)
500M
7.98%
GNOT
33
1.8M
17.20%
DPOT-L (FT)
500M
2.78%
OFormer
(53
1.9M
13.50%
DPOT-H (PT)
1.03B
3.79%
UPT-T
1.8M
5.08%
CViT-S 108
13M
3.75%
CViT-B 108
30M
3.18%
UPT-S
13M
3.12%
UPT-B
30M
2.69%
D.3.2 Shallow Water-2D dataset
We run comparisons against UNet, FNO, Dilated ResNet variants on the regular gridded Shallow 2D climate modeling dataset [3]] in Table 3 (baseline results and evaluation protocol taken from CViT (108). UPT outperforms all compared methods, which are specifically designed for regularly gridded data.
D.4 ShapeNet-Car
Dataset. We test on the dataset generated by [105), which consists 889 car shapes from ShapeNet (20] where side mirrors, spoilers and tires were removed. We randomly split the 889 simulations into 700 train and 189 test samples. Each sample consists of 3682 mesh points, including a small amount of points that are not part the car mesh. We filter all points that do not belong to the car mesh resulting in 3586 points per sample. [105] simulated 10 seconds of air flow and averaged the results
24

---

Table Comparison on the regular gridded small-scale Shallow Water-2D climate modeling dataset. UPT outperforms models that are specifically designed for regular grid data.
Model
# Params
Rel. I.2 Error
NilResNet [93]
4.2M
13.20%
U-Net 866 31
148M
5.68%
FNO [51
268M
3.97%
CViT-S [108]
13M
4.47%
UPT-S
13M
3.96%
the last seconds. The inflow velocity fixed at 20 m/s with an estimated Reynolds Number of Re = 5 x 10". For each point, we predict the associated pressure value. As there is no notion of time and the inflow velocity is constant, no feature modulation necessary. We use the transformer positional encoding [066], and, therefore rescale all positions to the range [0, 200]. We normalize the pressure values to have zero mean and unit variance.
Baselines. As baselines we consider FNO, U-Net and GINO. For FNO, we use a hidden dimension of 64 and 16/24/24 Fourier modes per dimension, and train with learning rate of 1e-3/1e-3/5e-4 for grid resolutions 32 /480//8///// /643, respectively. We also tried using 32 Fourier modes for grid resolution of 64 which performed worse than using 24 modes.
The U-Net baseline follows the architecture used as baseline in GINO which consists of 4 downsampling and 4 upsampling blocks where each block consists of two Group- Norm [T5]] -oooondnnneeeel (74) subblocks. The initial hidden dimension is set to 64 then doubled in each downsampling block and halved in each upsampling block. Similar to FNO, considered higher initial hidden dimension for grid resolution 64 which performed worse.
For GINO, we create positional embedding of mesh- and grid- which are then used as input to 3 layer MLP with GELU (37] non-lincarities to create messages. Messages from mesh- possi-osstionns within radius of each grid-position GINO. The FNO uses 64 hidden dimensions with 16/24/32 modes per dimension for grid resolutions 32 /48 /643, respectively. The GINO decoder encodes the query positions and again uses a layer MLP with GELU non- non-linearities to create messages. Each query position aggregates the messages from grid-points within radius. The radius for message aggregation is set to 10. When using SDF features, the SDF features are encoded via a layer MLP to the same dimension as the hidden dimension, and added onto the positional encoding of each grid point.
accumulated which serves input to the FNO part of
Hyperparameters. We train for 1000 epochs with 50 warmup epochs, a batchsize of 32. We tune the learning rate and model size for each model and report the best test loss.
UPT architecture. Due to a relatively low number of mesh points (3586), we use only a single perceiver block with 64 learnable query tokens as encoder, followed by the standard UPT architecture (transformer blocks as approximator and another perceiver block as decoder).
When additionally using SDF features as input, the SDF features are encoded by a shallow ConvNeXt V2 [ that processes the SDF features into 512 (8x8x8) tokens, which are concatenated to the perceiver tokens. To distinguish between the two token types, a learnable vector per type is added to each of the tokens.
As UPT operates on the mesh directly, we can augment the data by randomly dropping mesh points of the input. Therefore randomly sample between 80% and 100% of the mesh points during training and evaluate with all mesh points when training without SDF features. The decoder still predicts the pressure for all mesh points by querying the latent space with each mesh position. We also tried randomly dropping mesh points for GINO where it degraded performance.
Interpolated models. We consider U-Net and FNO as baseline models. These models only take SDF features as input since both models are bound to a regular grid representation, which prevents using the mesh points directly. The output feature maps are linearly interpolated to each mesh point from which a shallow MLP predicts the pressure at each mesh point. In the setting where no SDF features are used, we assign each mesh point constant value of I and lincarly interpolate onto grid.
25

---

Extended results. Table gives detailed overview of all results from ShapeNet-Car Car with varying resolutions and comparison to baseline models.
Table 4: Normalized test MSE for ShapeNet-Car pressure prediction. Memory is the amount of memory required to train on a single sample. UPTs can model the dynamics with fraction of latent tokens compared to other models. SDF additionally uses the signed distance function from each gridpoint to the geometry as input features. To include SDF features into UPT, we encode the SDF features with resolution 32 48 64 into 8 tokens using a shallow ConvNeXt V2 [13 and concatenate these tokens to the tokens coming from the UPT encoder. To balance the number of SDF tokens with the number of latent tokens, we increase the number of latent tokens to 1024 in the settings where we use SDF features for UPT. Runtime is measured on an A100 GPU.
Model
SDF
#Tokens
MSE [1e-2]
Mem. [GB]
Runtime per Epoch [s]
U-Net
X
64
6.13
86
FNO
X
64
4.04
3,8
148
GINO
X
64
2.34
19.8
900
UPT (ours)
X
64
2.31
U-Net
V
32"
3.66
0.2
11
FNO
V
323
3,31
0.5
17
GINO
V
32
2.90
2.1
103
UPT (ours)
/
8*
1024
2.35
2.7
U-Net
48
3.33
0.5
35
FNO
48
3.29
64
GINO
48
2.57
7.9
360
UPT (ours)
83
1024
2.25
2.7
156
U-Net
V
64"
2,83
1,3
86
FNO
643
3.26
3.8
148
GINO
64
2.14
19.8
900
UPT (ours)
8
1021
2.24
2.7
156
Profiling memory and runtime. We evaluate the memory per sample and runtime per epoch (Table 4) by searching the largest possible batchsize via a binary search. To get the memory consumption per sample, we divide the peak memory consumption by the largest possible batchsize. For the runtime per epoch, we conduct short benchmark run with the largest possible batchsize take and extrapolate the time to epoch. All benchmarks are conducted on A100 80GB SXM GPU.
D.5 Transient flows
The dataset consists of 10K simulations which we split into 8K training simulations, IK validation simulations and 1K test simulations, Each simulations has length of 100 seconds where At Is is used to train neural surrogate models. Note that this At is different from the At used by the numerical solver, which requires much lower At to remain stable. Concretely, we set the temporal update of the numerical solver initially to Al 0.05s. If instabilities occur, the trajectory is rerun with smaller At. Overall, each trajectory comprises 2K timesteps at the coarsest Al setting and 200K at the finest (A = 0.0005s). The trajectories are randomly generated 2D windtunnel simulations with 1-4 objects (circles of varying size) placed randomly in the tunnel. The uni-directional inflow velocity varies between 0.01 0.06 m/s. Each simulation contains between 29K and 59K mesh points where each point has three features: pressure, and the z- and y comomentttt of the velocity. In total, the dataset is converted to 16 precision to save storage and amounts to roughly 235GB.
Dataset.
From all training samples, data statistics are extracted normalize inputs approximately mean and standard deviation 1. In order to get robust statistics, the normalization parameters are calculated only from values within the inter-quartile range. As the Pytorch function to calculate quartiles not supported to handle such large amount of values, we assume that values follow Gaussian distribution with mean and variance from the data, which allows us infer the inter-quartile range from the Gaussian CDF.
26

---

Additionally, we convert the normalized values ai & to log scale to avoid extreme outliers which can lead to training instabilities:
-i,A- - sign( (7,,)) log(1 + ((
(26)
where all functions/operations are applied point-wise to the individual vector components. We apply the same log-scale conversion to the decoder output a
Model scaling. We scale models by scaling the dimension of the model. For grid grid-based baased methods, we choose 64x64 as grid size which results in similar compute costs of GINO and UPT. As encoder and decoder operate on the mesh- - which is more expensive than operating in the latent space - we use half of the approximator hidden dimension as hidden dimension for encoder and decoder. For example, for UPT-68M, we use hidden dimension 384 for the approximator and 192 for encoder/decoder. The hidden dimensions of the approximator as follows: 128 for UPT-SM, 192 for UPT-17M and 384 for UPT-68M.
For UPT, we use 12 transformer blocks in total, which we distribute evenly across encoder, approxi- mator and decoder (i.e. 4 transformer blocks for each component). As the decoder directly follows the approximator, this implementation detail only matters when training with inverse encoding and decoding to enable latent rollout. In this setting, the decoder needs more expressive power than a single perceiver layer to decouple the latent propagation and the inverse encoding and decoding.
For GINO, we use 10/12/16 Fourier modes per dimension with hidden dimension of 40/50/76 for the 8M/17M/68M parameter models, respectively. The hidden dimension of the encoder/decoder are the same ones as used for the UPT models.
For U-Nets, we use the Unetmwd architecture from Gupta & Brandstetter 31 where we adjust the number of hidden dimensions to be approximately equal to the desired parameter counts. We use hidden dimensions 13/21/42 for the 8M/17M/68M parameter models, respectively.
Implementation. To efficiently use multi-GPU setups, we sample the same number of mesh points for each GPU. Without this modification, the FLOPS in the encoder would fluctuate drastically per GPU, which results in "busy waiting" time for GPUs with less FLOPS in the encoder. We also keep the number of mesh points constant between batches to avoid memory fluctuations, which allows using a larger batchsize. We rescale the x positions from the range [-0.5, 0.5] to [o, 200] and the y positions from [-0.5, 1] to [0, 300]. We do this to avoid complications with the positional embedding [106] which was designed for positive integer positions. To create graph out of mesh points, we use a radius graph with r = and limit the maximum number of graph edges per node to 32 to avoid memory fluctuations. We choose = 5 as it covers the whole domain when using a grid resolution of 642.
Training. We train all models for 100 epochs using batchsize of 1024 where we use gradient accumulation if the full batchsize does not fit into memory. Following common practices, we use AdamW [1 59 for U-Net, FNO and GINO. As encountered training instabilities when training UPTs on this dataset, we change the optimizer for UPTs to Lion (21). These instabilities manifest in sudden loss spikes from which the model can not recover. Using the Lion optimizer together with float32 precision training solves these instabilities in our experiments. When training with Lion, we learning rate of 5 x 10 and 1 x 10-4 otherwise. We also use float32 for GINO as we found that this improves performance. U-Net and FNO do not benefit from float32 and are therefore trained in 6 precision.
We linearly increase the learning rate for the first 10 epochs (58) followed by cosine decay. A weight decay of 0.05 is used when training with AdamW, which is increased to 0.5 when using Lion as suggested in (21). Architectural choices like the number of attention heads, the latent dimension for each attention head or the expansion dimension for the transformer's 's MLP are copied from ViTs.
Other hyperparameters are chosen based on common practices when training ViTs 26)
We find float32 precision to be beneficial for GINO and UPTs, which is likely due to positional embeddings being too inaccurate with 6 precision. Training with float16 instead of bfloat16 (to increase the precision of the positional embedding while speeding up training) resulted in NaN values due to overflows. Therefore, we train all GINO and UPT models with float32. We found this to be very important as otherwise UPTs become unstable over the course of training, leading to a large loss spike from which the model can not recover. During model development, we tried various ways to solve this instabilities where using the Lion optimizer was one of them. With fp32 training, UPTs
27

---

work similarly well with AdamW but we kept Lion as we already trained models and re- re-training all models quite expensive.
Interpolated models. We train U-Net and FNO models as baseline on this task. To project the mesh onto regular grid, position and interpolates between the values of the nerarest neighbors. Going from the grid back the mesh implemented via the grid_ __spapllleeeeee function of pytorch. For FNO, we use the same dimensions as in latent model of GINO and for U-Net, we scale the latent dimensions such that the model is in the same parameter count range as the other models.
k-NN interpolation which takes the nearest neighbors of grid
D.5.1
Visualization of rollouts
We show more qualitative results of the UPT-68M model in Fig.
D.5.2 Supernode discretization convergence
While one would ideally use lots of supernodes also during training, also requires more computa- tional power. Therefore, we investigate increasing the number of supernodes only during inference. We use the 68M parameter models from Sec. 4.2 and increase the number of supernodes during inference. Figure showwws that UPTs are discretization convergent and inference performance can be improved by using more supernodes for inference than during training.
1e-3
GINO
Test MSE (-)
UPT
4
1K 2K 4K 8K 16K 32K
Number of Gridpoints/Sppernodes
Figure 8: We study discretization convergence by varying the number the number of gridpoints/- supernodes of models that were trained 2K supernodes (UPT) or 4K grid points (GINO). UPT demonstrates stable performance across different number of supernodes even though it has never seen that number of supernodes during training. Increasing the number of supernodes even improves the performance of UPT slightly. In contrast, the performance of GINO plummets when the number of gridpoints is different during inference.
D.5.3 Impact of larger latent space
As training with larger latent spaces becomes expensive, we investigate it in a reduced setting where we train for only 10 epochs and fix the number of input points to 16K. The results in Fig. show that UPTs scale well with larger latent spaces, allowing flexible
tradeoff.
D.5.4 Performance of smaller models
FigureBccoubl could suggest that UPT does not scale well with parameter size due to the steeper decline of GINO. However, the scaling of GINO only looks good because GINO underperforms in contrast UPT (GINO-68M is worse than UPT-8M). For well trained (UPT) models gets increasingly difficult to improve the loss further. Ideally, one would show this by training larger GINO models, however this is prohibitively expensive (68M parameter models already take 450 A 1000 hours per model). We therefore go in the other directions and train even smaller UPT models that achieve similar loss to the GINO models and compare scaling there. In Figure TO right, we compare UPT IM/2M/8M against GINO 8M/17M/68M. UPT shows similar scaling on that loss scale.
28

---

5.5 1.022
105
5.5 10.2
5.5
10.2
165
Throughpus (-)]
(-)
4.5
Test MSE
3.5
Test MSE
3.5
Through
Test MSE
3.5
2.5
2.5
2.5
1.5 512
128 1.5
1.5 1/
Number of superoodes
Number of latent tokens
Compound scale
Figure 9: Latent space scaling investigations of I 17M parameter UPT model for 10 epochs. Com- pound scaling scales the number of supernodes and latent tokens simulataneously where a compund scale of uses Tssupemodes =2048 and P i.e. compound scale 2 uses nsspermodess 4009 and MMateo==1224 Throughput measured as number samples processed per GPU-hour. Models are trained reduced setting with 10 epochs and 16K input points.
We strongly hypothesis that the effect of larger UPT models would become apparent in even more challenging settings or even larger datasets. However, challenging large-scale datasets are hard to come by, which is why created one ourselves. Creating even larger and more complex ones beyond the scope of our work as it exceeds our current resource budget, but it is definitely an interesting direction for future rescarch/applicattins..
10
GINO
8M
UPT
17M
Test MSE (+)
1M
2M
68 M
8M
10
Expressitivity
Figure 10: UPT much more expressive than GINO and shows good scaling when increasing model size. UPT-IM achieves comparable performance to GINO-17M using 17x less parameters.
D.5.5 Scalability with dataset size
We show the scalability and data efficiency UPTs by training UPT-SM on subsets of the data used for the transient flow experiments (we train on 2K and 4K out of the 8K train simulations). The results in Figure right show that UPT scales well with data and is data efficient, achieving comparable results to GINO-8M with 4x less data.
D.5.6 Out-of-distribution generalization study
We study generalization to out-of-distribution datasets by evaluating the 68M parameter models that were trained on the transient flow dataset from Section 4.2 The in-distribution dataset uses an adaptive meshing algorithm, i.e. the mesh resolution around objects is increased, leading to mesh cells different sizes and contains between and circles of variable size at random locations.
We evaluate three different settings of variable difficulty: more objects, higher inflow velocity. For each setting, we generate new dataset containing 100 simulations and evaluate the model that was trained on the in-distribution dataset. The results in Figure 12 show that UPT's behave similar to other models when evaluated on OOD tasks where UPT outperforms all other models all tasks. Therefore, the grid-inde inddiidddendent architecture of UPTs does not impair OOD generalization performance.
29

---

10
Test MSE (-)
GINO
UPT
10 3
2
4
8
#Simulations [K]
Figure 11: UPT scales well with more data and data efficient, achieving comparable results to competitors with 4x less data.
Generalization to more objects. Adding additional obstacles to flow makes it more turbulant. The left side of Fig. 12 shows that also for in-distribution, more objects correspond to higher loss values due to the increased difficulty. Models show stable generalization to larger object counts. Note that we do
encode the number objects or the objects explicitly in the model. The model simply does not get any inputs in the locations where obstacles are. Therefore, only the dynamics of the flow change but not the distribution of input positions.
Generalization to higher velocities. The in-distribution dataset contains random velocities sampled from E 10.01, 0.06] m/s. The velocity is used (together with the timestep) as input the modulation layers (e.g. DiT (78 modulation) which modulate the features of all layers in the network. This leads to OOD velocities distorting all features within forward pass. The center plot of Fig. 12 shows that UPT has the best OOD generalization among all models. Notably, the performance of GINO drastically drops for higher velocities.
Generalization to difffernt geometries. The in-distribution dataset contains randomly placed circles of varying size where the mesh is generated via adaptive meshing algorithm. To investigate robustness to the meshing algorithm, we generate a OOD dataset with uniform meshing algorithm. In this dataset, mesh cells approximately uniform, i.ethe distance between two points is roughly the same. This in contrast to an adaptive mesh, where regions around an object have more mesh cells and therefore the distance between two points is smaller in these regions. Additionally, investigate generalization to different obstacle geometries by using polygons (with up to 9 edges) or triangles instead of circles. Note that even though polygons more reminiscent of a circle than triangles. The size of the obstacle is also varied here. For simplicity, the number of objects per simulation is set to this study. Note that U-Net and FNO interpolate the mesh onto a 2D grid and therefore their distribution of input positions does not change here, only the dynamics of the simulation. For GINO and UPT, also the distribution of input position changes. The right side of Fig. [2 shows that UPT achieves the best performances on OOD meshes, despite having to adjust to a different input position distribution.
a more "complicated" shape, they
D.6 Lagrangian fluid mechanics
Datasets. We use the Taylor-Green vortex datasets from Toshev et al. [00]] for our Lagrangian experiments, comprising two dimensions (TGV2D) and three dimensions (TGV3D). The Taylor- Green vortex problem is characterized by a distinctive initial velocity field without an external driving force, resulting in a static decrease in kinetic energy over time. The TGV2D dataset includes 2500 particles, sequence length of 126, and is characterized by Reynolds number (Re) of 100. The data is split into 100/50/50 for training, validation, and testing, respectively. The TGV3D dataset consists of 8000 particles, a sequence length of 61, and a Reynolds number of 50. The data is partitioned into 200/100/100 for training, validation, and testing, respectively.
Baselines. The Graph Network-base Simulator (GNS) model [8 is popular learned surrogate for physical particle-based simulations. The architecture is kept simple, adhering to the encoder-
30

---

10c
FNO GINO
10
FNO
10
GRO
UPT
UPT
Teat MSE
Test MSE
Test MSE
10-2
10
10
GINO
UPT
Polygon Triangle
*Objects
Velocity [m/s]
Geomeby Setting
Figure 12: OOD generalization study. Trained 68M parameter models evaluated OOD datasets with more objects, higher inflow velocities and different geometries. "ID" refers to an adaptive meshing algorithm with circles as obstacles. "UMesh" changes the meshing algorithm from adaptive to uniform. "Triangle" uses triangles instead of circles and "Polygon' uses polygon obstacles (with up edges) instead of circles. UPTs have similar OOD generalization capabilities other models, outperforming them in all evaluations. Grey indicates OOD settings.
processor-decoder principle [6]. where the processor consists of multiple graph network blocks [5] The Steerable -qquvvrauaat Graph Neural Network (SEGNN) architecture [3] is general implementation of an E(3) equivariant GNN, where layers are directly conditioned on steerable attributes for both nodes and edges. The main building blocks are steerable MLPs, i.e., stacks of
learnable linear flebsch-Gordan tensor products interleaved with gated non-linearities [109]. SEGNN layers are pessaae--assiing layers (28] where steerable MLPs replace the traditional non-equivariant MLPs for both message and node update functions.
We utilize the checkpoints provided by Toshev al. [0I for the comparisons in 4.3 Both GNS and SEGNN baseline models comprise 10 layers with latent dimensions of 128 and 64, respectively. The maximum irreps order of SEGNN are = 1, for more details [101).
Hyperparameters. We train using batchsize of 128 for 50 epochs with warmup phase of 10 epochs. We sample inputs between 50% 100% of total particles. For optimization, AdamW [59] for all experiments with learning rate -3 and weight decay of 0.05.
UPT architecture. For the TGV2D and TGV3D experiments, we use supernodes, respectively. For TGV3D more supernodes are required due to the increased number of particles. Message passing is done with a maximum of four input points per supernode. The flexibility of the UPT encoder allows us to randomly sample 50% up to 100% of the total particles, diversifying the training procedure and forcing the model to learn the underlying dynamics.
256 and = 512
The encoder comprises four transformer blocks with a latent dimension of 96 and two attention heads. The encoder perceiver expands the latent dimension to 192 and outputs 128 latent tokens, using 128 learned queries and three attention heads. The approximator comprises four transformer blocks with a latent dimension of 192 and three attention heads. The decoder perceiver comprises again three attention heads. Overall, the parameter count of UPT amounts 12.6M parameters.
31

---

t=
[=
[=
Ground Truth
Prediction
Abs. Error
t= 10
t= 20
t= 30
t= 40
Ground Truth
Prediction
Abs. Error
t= 10
t= 20
t= 30
t= 40
Ground Truth
Abs. Error Prediction
t=
Ground Truth
Prediction
Abs. Error
Figure 13: Exemplary visualizations for transient flow simulation rollouts. Best viewed zoomed in.
32

---

Experiments on Taylor Green Vortex 2D.
We show results the TGV2D dataset from LagrangeBench [101 and the time generate simulated trajectory in Figure [4 The velocity error behaves similar to the 3D version of the dataset in Figure 7 UPT is roughly 10 times faster than GNNs and 100 times faster than the classical SPH
solver.
TGV2D
GNS
SEGNN
Velocity Emro
Speedup
SPH
4.9s
1x
A6000
SEGNN
2.76s
2x
A100
4-
GNS
0.56s
A100
UPT
0.05s
98x
A100
20
40
60
80
100 120
Timestep
Figure 14: Left: Mean Euclidean norm of the velocity error over all particles for different timesteps. UPTs effectively learn the underlying field dynamics, resulting in lower velocity error as the trajectory evolves in time. Right: Comparison of simulation/rollout runtimes for a TGV2D trajectory with 125 timesteps and 2500 particles across SPH simulation.
D.7 Qualitative study of scaling limits details
We qualitatively study the scaling limits of representative methods by evaluating the memory con- sumption when training a 68M parameter model with batchsize while scaling the number of inputs in the left part of Fig. 2
study the following models:
GNN: Lagrangian based simulations are currently dominated by GNN architectures due to their strong ability to update individual particles based on representation of neighboring particles. In this study, we consider GNO (i.e. GNN with radius graph and mean aggregation) since edge information is typically absent in physical simulations, but this does not change the asymptotic complexity. Additionally, we limit the number of edges per node (degree of the graph) to 4 as otherwise the GNN could not fit 32K inputs on 80 GB of GPU memory. Note that this heavy restriction and these models would most likely need larger degree to perform well when actually trained.
The amount of operations required for the vanilla transformer self-attention is quadratic in the number of inputs and therefore becomes increasingly expensive when the number of input grows. Linear transformers change the attention mechanism to have linear complexity and therefore scale better larger input spaces. We use the linear attention implementation of [90), but note that the complexity is equivalent to recent transformer architectures proposed as neural surrogate for partial differential equations [18 53 33
Transformer:
GINO: Aggregating information into a lower dimensional latent space allows models to be scaled to large number of inputs. GINO is one of the prominent models that operates with this paradigm by projecting the input point cloud onto points on a regular grid.
UPTs: The efficient latent space compression of UPTs allows scaling to large number of input points.
We do not include FNO or CNN/U-Net into this study as they can only handle inputs from a regular grid. Therefore, they would need very fine- -rainnd grid to achieve comparable results to GINO or UPTs, which makes fair comparison hard. Note that GINO usea FNO after it projects inputs onto a regular grid. Therefore, the scalability of GINO is highly correlated with the scalability of FNOs or CNN/U-Nets (i.e. decent scalability on 2D problems but drastically higher compute requirements on 3D problems). We also exclude transformers with quadratic attention as they quickly become infeasible and their linear complexity counterpart scales linear with the number of inputs.
33

---

We also scale the latent space size for GINO (number of gridpoints) and UPTs (number of supernodes and number of latent tokens). As starting point, we use the latent space sizes used in S.c For GINO-3D, we use the same grid resolution for the 3rd dimension as for GINO-2D (i.e. GINO-2D uses 64 and GINO-3D 64 Our experiments in Sec. 4.1 and Sec. 4.2 show that GINO needs this high resolution for the grid to achieve good performances. For GINO-3D we set the number of fourier modes to 5 to keep the parameter count close to 68M.
As quantitative comparison between completely different architectures is challenging, we addi- tionally list the theoretical model complexities in Tab. [ This study meant to give a qualitative impression of the practical implications of the theoretical asymptotic complexities.
Table 5: Extending the right table of Fig. 2 with theoretical asymptotic complexities. Complexity includes number of mesh points M, and maximum degree of the graph D. Grid-based methods project the mesh to G' grid points. UPTs instead use small amount of supernodes S as discretization, where is typically much larger than S. The UPT training procedure separates responsibilities between components, allowing to forward propage dynamics purely within the latent space.
Irregular Grid
Discretization Convergent
Learns Field
Latent
Model
Range
Complexity
Rollout
GNN
local
D(MD)
X
X
X
CNN
local
O
X
Transformer
global
O(M)
X
Linear Transformer
global
O(M)
X
X
GNO
R22
radius
O(MD)
V
V
X
FNO
global
X
X
GINO
global
Glgg G)
UPT
global
O(SD+
V
E
Impact Statement
Neural PDE surrogates play an important role in modeling many natural phenomena, many of which are related computational fluid dynamics. Examples weather modeling, aerodynamics, biological engineering, or plasma physics. Given the widespread application of computational fluid dynamics, obtaining shorteuts or alternatives for computationally expensive simulations, is essential for ad advancing scientific research, and has direct or indirect implications for reducing carbon emissions. In contrast traditional computational fluid dynamics simulations, our approach aims eliminate the need for complex meshing, bypass the requirement for detailed knowledge of boundary conditions, and overcome bottlenecks caused by solver intricacies. However, is important to note that relying on simulations always necessitates thorough cross-checks and monitoring, especially when employing "learning to simulate" methodology.
F Implementation Details
F.1 Supernode pooling
Supernodes are, in the case of Eulerian data, sampled according to the underlying mesh and, in the case of Lagrangian data, sampled according to the partiele density. A random sampling procedure which follows the mesh or particle density, respectively. allows us to put domain knowledge into the architecture. Consequently, complex regions are accurately captured, as these regions will be assigned more supernodes than regions with a low-resolution mesh or few particles.
The sampling of supernodes is done for each optimization step which provides regularization.
The radius graph encodes a fixed region around each supernode. We not use any information (such from the original mesh in the Eulerian case. While one could use
as cell comnectivity/adjaeemn original edge connections for Fulerian data, Lagrangian data does not have edges. Additionally, we employ randomly dropping input nodes as form of data augmentation which makes using the original edge connections more complicated from an implementation standpoint. contrast, our design via radius graph is agnostic to the simulation type and also to randomly dropping input nodes.
34

---

We choose the radius depending on the dataset. For each dataset, we first analyze the average degree of the radius graph with different radius values. We then choose the radius such that the degree is around 24, i.e., on average each supernode represents 24 inputs. We found our model to be fairly robust to the radius choice. Also, the circles of different supernodes can overlap. Therefore, the encoding of dense regions can be distributed among multiple supernodes.
We impose the edge limit of the radius graph by randomly dropping connections to preserve the average distance from supernode to imput nodes and avoid "squashing" too much information into single superoode. We choose this edge limit be 32 in all cases.
F.2 Position embedding
We embed input and query positions via the sine-cosine position embedding from transformers [006] Each position dimension embedded separately. We rescale all positions to be in the range [0, 200J to avoid "unsmooth" transitions when going from positive positions to negative positions. This also makes sure that positions are in a suitable scale for the transformer positional embedding. For example, input positions could be arbitrarily small, which would lead all positions to be the same due the limited precision of float values. The upper bound 200 is hyperparameter where broad range of values works. If input nodes have additional input features, these features are first linearly projected into hidden dimension and then the positional embedding vector is added.
We also experimented with random fourier features 96 but did not find significant difference.
F.3 Conditioning pseudocode
We provide pseudocode how to create a conditioning vector to encode scalar properties into the model. This conditioning vector is then used to predict the parameters of LayerNorm [2] layers before the attention and MLP parts of each transformer block. Additionally, third scalar for gating the residual connections is learned. This methodology follows DiT (78). For predicting these scalars, simple learnable linear projection is used.
embed: sine cosine positional embedding condition_ shallow MLP to combine boundary conditions def create_ condition (timestep, velocity): temestep_embed = embed (timestep) velocity-embes
embed (velocity)
embed concat [[timestep_embed velocity_embed)) condition = condition_ulp (embed) return condition
7
Listing 1: PyTorch-style pseudocode for conditioning onto scalar variables.
E.4 Training pseudocode
provide pseudocode a training step in the setting of the transient flow experiments (2D positions with 3 features per node) in Listing 2 including inverse encoding/decoding objectives.
35

---

input_embed: linear projection from input features dim to hidden dim poa_embed: sime-cosine positional embedding
message_mip: shallow MLP to create messages
encoder..rrassform::: stack of transformer blocks of the encoder latent_queries: 'n_latent_tokens
learnable query vectors
encoder.perceeive:: cross attention block irpssimattr__ransfooroe
stack of transformer blocks
query_mlp: shallow MLP in the decoder to encode query positions
decoder: cross attention block
10
def encoder (input input_pos, radius, a_supernodes):
12
13
encode arbitrary pointclouds into fixed latent space inputs:
14
input_feature Tensor (n_input_nodes, 3); features of input nodes
16
input_pos Tensor (n_iaput_nodes, 2): positions of input nodes n_aupernodes integer: number of supernodes radius float: radius for creating the radius-graph
17
18
outputs:
20
latent Tensor (n_latentt_okees, hidden_dim): encoded latent ...
21
22
create radius graph (using all input nodes)
edges are uni-directional
messages passed from codes_from to nodes_to nodes_from, nodes_to
25
dadus-graphhinpttee
radius)
26
27
# select supernodes from input_nodes n_input_nodes (en(iutt [: supernodes]
29
supernode_idx randpers(n_.inut.nooo
30
31
filter out edges that do not involve supernodes
is _dg nodes_to in supernode_idx
nodes_from nodes_from sis.ruoee edge] nodes_to todes_
34
35
36
encode inputs and positions encoded_nodes = input_embed taaput_fettureeee + pos_embed (input_pos)
39
create messages
40
messages =
[nodes_from]]
# accumulate messages per supernode by averaging messages
supernodes accumulate (messages, nodes_to, reduce="mean")
43
44
process supernodes with some transformer blocks supernodes
45
perceiver pooling from supernodes to latent tokens
latent encoder..prreeiver query-latemt.querjiee, key=supernodes, value-supernodes,
48
49
50
52
return latent
53
54
def approximator (latent):
propagates latent forward in time inputs:
37
58
latent_t Tensor (n_latent tokens, hidden_dim): encoded latent space at timestep
59
outputs:
tatent_t_plus_1 Tensor (n_laatttttttttoeeee hidden_dim): encoded latent space at timestep
62
63
return approxinatto__tansffoooo (latent)
36

---

def decoder (latent, query-pos):
67
68
decode latent space pointwise at arbitrary positions
69
inputs:
latent Tensor (nnoorrn_lttent_toooe query-pos Tensor (n_ (n_outputa, 2): positions for querying the latent space
hidden_dim): latent space
72
73
outputa:
24
decoded Tensor outputs, 3);
evaluation of the latent space query positions
26
77
encode query positions
78
query-pos_embll embed = query_mlp (pos, embed (query-pos))
query latent space
decoded decoder (query embed, key=latent, value=latent)
81
return decoded
82
83
def step (tepiutt__featurees input_pos, n_supernoess, radius, query-pos):
encode arbitrary pointelouds into fixed latent space inputs:
86
87
input_ featuree Tensor (n_ iinptt_nodes 3): features of input nodes at timestep t input_pos Tensor (n_input_nodes positions
2): timestep n_supernodes integer: number of supernodes
90
input nodes
91
92
radius float: radius for creating the radius_graph query-pos Tensor (n_outputs, 2):
positions for querying the latent space target_features Tensor (n_ipput_odd 3): features of nodes at timestep t
95
96
outputs:
loss Tensor (1); (
skalar loss value
***
100
next-sted prediction
101
latent_t encoder(
input_features, input_pos,
104
radius, a_supernodes,
105
latent_ tplusi approxiiattorlaatet.... tplusi decoder (Iatenttpllt query-pos)
109
next next_step_loss mse (decooeee__tpllu target_featuret
# inverse decoder (decode latent into inputs)
decoded_t = decoder paccdettaatnt.t... input_pos)
inverse_decoding_.oo
mae (decooded, input_features)
113
inverse encoder (encode predictions at t+1 into latent t+1) inverse_
114
encoder(
decoded_tplust,
query_pos, radius,
118
n_supernodes,
119
Lnverse_encoding_loo = mse (inverse_encoded, latent_tplusi)
losses
122
return
123
next_step_los
+ inverse_decoding_Je
+ Inverse_encoding
Listing 2: PyTorch-style pseudocode for UPT with inverse encoding an decoding objecitves.
37

---

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 introduce a scalable neural operator architecture that performs well on a wide range of physical simulations.
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.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We discuss limitations in Appendix A
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 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 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 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 approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or 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 used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged 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 to not penalize honesty concerning limitations.
the paper. The authors should use their best
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]
38

---

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 the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide 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 supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibiliyy
Does the paper fully disclose all the information needed to reproduce the main ex- perimental 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)?
Question:
Answer: [Yes]
Justification: We provide implementation details for our experiments Appendix D Additionally, we provide the code that was used for this work in the supplemental materials.
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, 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 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 to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way 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 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 submis- sions 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 reproduce that algorithm. (b) If the contribution is primarily
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 way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe 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?
39

---

Answer: [Yes]
Justification: We provide the code that was used for this work in the supplemental materials which also includes data preparation scripts.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips... cc/
public/guidest httpe///eeooeeueebesssssoooooooo
for more details.
While we encourage the release of code and data, we understand that this might 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 new open-source benchmark).
The instructions should contain the exact command and environment needed to reproduce the results. See the NeurlPS code and data submission guidelines (https
wetgguddddeeeuuubbueeee ssinoolicyy 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 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 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, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We provide experimental details and hyperparameters in Appendix D The code in the supplemental materials also includes the full hyperparameter configurations used for each experiment.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: Repeating experiments with different random seeds is expensive, especially in the large-scale experiments.
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, 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 with given experimental conditions).
40

---

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).
It should be clear whether the error bar is the standard deviation 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 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful figures symmetric error bars that would yield results that are out of range (e.g. negative rates).
show in tables or
error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We provide details on compute infrastructure in Appendix D.I and provide the most important runtimes in the main paper.
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, 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 into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics htttp:://eurr/ http://pp.....bbcc/ccaaeeccceeea
Answer: [Yes]
Justification: Our paper conforms with the guidelines.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive socictal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
We provide impact statement in Appendix
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 why the paper does not address societal impact.
41

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess 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 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 generate deepfakes for disinformation. On the other hand, it is not needed to point out that 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 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 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 monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper deseribe 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)?
Question:
Answer: [NA]
Justification:
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 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 best faith effort.
12. Licenses for existing assets
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?
Question:
Answer: [Yes]
Justification: Authors of code and data are cited accordingly.
Guidelines:
The answer NA means that the paper does 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, possible, include 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 package should be provided. For popular datasets, paperswithcone com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
42

---

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 s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/mooel 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. Crowdsoureing 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 details about compensation (if any)?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According 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 equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending 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.
43


Reviews:
[{"review_id": "1nGuGhHL5W", "reviewer": "Reviewer_5AUz", "summary": "In this paper a framework for efficiently scaling neural operators is introduced under the name of Universal Physics Transformers (UPTs). It is a novel paradigm that scales neural operators across diverse spatio-temporal problems without relying on specific grid or particle-based structures. Leveraging transformer architectures, UPTs propagate dynamics within a compressed latent space, enabling efficient and flexible simulations.", "questions": "1/  The paper briefly mentions feature modulation for model conditioning but lacks detailed discussion on handling complex boundary conditions. Providing specific examples or additional experiments that illustrate UPTs' effectiveness in managing complex boundary conditions would strengthen the paper’s claims. Addressing this can clarify the practical applicability of UPTs in real-world scenarios where boundary conditions play a crucial role.\n\nHow does the UPT framework handle complex boundary conditions compared to traditional particle and grid-based methods?\n\n2/  While UPTs are presented as efficient and scalable, the paper does not delve into the specifics of their numerical stability and accuracy compared to traditional methods. Discussing strategies to maintain stability and accuracy over extended simulations and providing relevant experimental evidence would address potential concerns about the reliability of UPTs in long-term applications.\n\nWhat measures are in place to ensure the numerical stability and accuracy of UPTs over long simulation periods?\n\n3/ The paper demonstrates scalability and generalization within certain limits but does not extensively explore these aspects in more demanding scenarios. Providing additional experimental results or theoretical insights on UPTs’ performance in high-resolution and large-scale environments, as well as their ability to generalize to entirely new conditions, would significantly bolster the paper’s contribution and practical relevance.\n\n What are the scalability and generalization capabilities of UPTs when applied to extremely high-resolution grids, large particle systems, and unseen physical scenarios?", "limitations": "Yes.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1/ Originality:\n\nThe paper presents an original contribution to the field of neural operators by introducing Universal Physics Transformers (UPTs). This novel paradigm removes the traditional reliance on grid or particle-based latent structures, enabling greater flexibility and scalability across various simulation types. The innovative use of transformer architectures to propagate dynamics within a compressed latent space is a quite original combination of existing ideas applied to a new domain.\n\n2/ Quality:\n\nThe quality of the overall contribution is good. The authors address a critical challenge in scaling neural operators for complex simulations and provide a robust framework with practical applications. The UPT framework is well-formulated, and the encoding and decoding schemes are designed to ensure efficient simulation rollouts. The experiments are comprehensive, covering multiple types of simulations, and the results demonstrate the superiority of UPTs in terms of performance and scalability. \nMoreover, the paper's methodology is sound, and the technical claims are well-supported by evidence from the experiments.\n\n3/ Clarity:\n\nThe paper is well-written and clearly presented. The authors provide a thorough background. \n\n4/ Significance:\n\nThe significance of the paper lies in its potential impact on the broader scientific and engineering communities. By providing a unified and scalable framework for neural operators, UPTs can be applied to a wide range of spatio-temporal problems, including those in fluid dynamics. The ability to efficiently handle large-scale simulations can lead to significant advancements in these fields, offering valuable insights and solutions to complex physical phenomena. \nMoreover, the experimental design is robust, with appropriate datasets and baselines used for comparison. The methodology is clearly described, allowing for reproducibility of the results. The use of inverse encoding and decoding techniques to enable latent rollouts is particularly noteworthy, as it demonstrates a deep understanding of the underlying principles and challenges in neural operator learning.\n\n5/ Latente space and scalability:\n\nThe Universal Physics Transformers (UPTs) framework is efficient since it allows to get compact yet expressive latent space representation, capturing the essential dynamics of physical systems while significantly reducing memory usage and computational overhead. This compactness, combined with efficient encoding and decoding schemes, ensures that the transformation to and from the latent space is both efficient and accurate. By unifying the encoding of various grids and particles, UPTs simplify the modeling process and improve generalization across different spatio-temporal problems.\n\nThe empirical validation provided in the paper, through steady-state and transient flow simulations, demonstrates that UPTs achieve lower mean-squared error (MSE) and faster computation times compared to other models, highlighting the effectiveness of the latent space representation.\n\nIn terms of scalability, UPTs are designed to handle large-scale simulations effectively. The fixed-size latent space, regardless of input size, allows UPTs to manage large inputs without a proportional increase in computational cost. This scalability is evident in the framework's ability to maintain performance and efficiency with large meshes and numerous particles. The efficient latent space rollouts further enhance scalability by enabling quick updates and predictions, making UPTs suitable for time-dependent simulations. Additionally, the transformer architecture provides a robust foundation for scalability, ensuring that UPTs can handle complex spatio-temporal simulations efficiently. \n\n\n6/ Lagrangian dynamic modeling:\n\nThe Universal Physics Transformers (UPTs) framework presents several notable strengths when viewed from the perspective of Lagrangian dynamic modeling. One of the primary advantages is the framework's ability to model particle-based simulations effectively without relying on traditional particle-structures. In Lagrangian simulations, particles move with the local deformation of the continuum, and modeling these dynamics accurately requires handling a large number of particles and their interactions.\n\nUPTs use a unified latent space to encode particle information, enabling the framework to handle varying numbers of particles flexibly. This flexibility is particularly beneficial for Lagrangian methods, such as Smoothed Particle Hydrodynamics (SPH), where particle counts can vary significantly based on the simulation's complexity. By compressing this information into a fixed-size latent space, UPTs manage to reduce computational overhead while maintaining the ability to capture intricate particle interactions.\n\nThe efficient encoding and decoding processes in UPTs ensure that particle dynamics are propagated accurately and swiftly within the latent space. This efficiency is crucial for large-scale Lagrangian simulations, where the computational cost can be prohibitive with traditional methods. The ability to perform latent rollouts enables UPTs to predict future states quickly, making the framework suitable for real-time or near-real-time applications in Lagrangian dynamic modeling.\n\nThe paper provides strong empirical evidence of UPTs' effectiveness in Lagrangian dynamic modeling through experiments on datasets such as the Taylor-Green vortex in three dimensions (TGV3D). The results demonstrate that UPTs can effectively learn the underlying field dynamics and predict particle velocities with lower error compared to traditional Graph Neural Networks (GNNs) and other baselines. This empirical validation underscores the practical applicability of UPTs in complex Lagrangian simulations\n\nUPTs exhibit strong generalization capabilities, which are critical for Lagrangian dynamic modeling. The ability to query the latent representation at any point in space-time allows UPTs to adapt to different particle distributions and simulation conditions without extensive retraining. This robustness ensures that UPTs can handle a wide range of Lagrangian problems, from small-scale particle systems to large-scale simulations involving a representative number of particles.\n\nThe use of transformers to manage the latent space representation in UPTs is a significant technical innovation. Transformers are known for their efficiency in handling large datasets and sequences, and their application in UPTs leverages this strength to manage the complexities of Lagrangian dynamics. This choice of architecture allows UPTs to model particle interactions accurately while maintaining computational efficiency.", "weaknesses": "1/ Clarity and detail in methodology:\n\nThe methodology, particularly the detailed implementation of the encoding and decoding schemes, could be more clearly articulated. Some parts of the algorithms are difficult to follow. Providing additional diagrams, detailed explanations, or pseudo-code would help clarify these complex processes. Ensuring that each step of the process is well-explained and easily understandable would make the paper more accessible and replicable.\n\n2/ Scalability with extremely large datasets:\n\nWhile the paper demonstrates UPTs' scalability, there is limited discussion on how the framework performs with extremely large datasets or in distributed computing environments. Given the increasing size of datasets in fields like high-resolution climate modeling and genomics, a discussion or preliminary results on the scalability of UPTs in such contexts would be valuable. This could include potential challenges, proposed solutions, and the impact on computational resources.\n\n3/ Insufficient analysis of generalization capabilities:\n\nWhile UPTs are shown to generalize across different simulation types, the paper lacks a deep analysis of this capability. Neural operator networks are often praised for their ability to generalize across different boundary conditions and initial states. A more thorough examination of UPTs' generalization performance, especially in unseen or out-of-distribution scenarios, would strengthen the claims. Detailed experiments and discussions on how UPTs perform under varied conditions would be valuable.\n\n4/ Potential overfitting concerns:\n\nGiven the complexity and high capacity of transformer-based models, there is a risk of overfitting, especially on smaller datasets. The paper briefly mentions overfitting issues in some experiments but does not provide a detailed strategy for mitigating this. More discussion on regularization techniques, data augmentation strategies, or how UPTs handle overfitting would be beneficial. Insights into how the model can be generalized better and made more robust against overfitting are crucial for practical applications.\n\n5/ Analysis from particle and grid based methods:\n\nThe Universal Physics Transformers (UPTs) framework offers a novel approach, but there are specific areas where it could be improved, particularly when compared to traditional particle and grid-based methods.\n\nA/ Lack of detailed comparison with established methods:\n\nThe paper does not provide a thorough comparison with well-established particle-based methods (such as Smoothed Particle Hydrodynamics, SPH) and grid-based methods (such as Finite Volume Methods, FVM). Including detailed performance metrics, such as accuracy, computational cost, and scalability, would help highlight UPTs' advantages and areas needing improvement. Specifically, comparative studies on benchmark problems typically addressed by SPH (except figure 12) and FVM would provide more context.\n\nB/ Handling of boundary conditions:\n\nParticle and grid-based methods have well-established techniques for handling complex boundary conditions, which are often crucial in physical simulations. The paper does not sufficiently address how UPTs manage complex boundary conditions compared to these traditional methods. A more detailed discussion or experiments showcasing UPTs' effectiveness in dealing with various boundary conditions would strengthen the paper's claims.\n\nC/ Adaptation to high-resolution grids and large particle systems:\n\nWhile UPTs are shown to be scalable, the paper lacks detailed insights into their performance on extremely high-resolution grids or very large particle systems. Traditional methods often excel in these areas due to their specialized structures and optimizations. Providing more evidence on how UPTs handle such scenarios, including any potential bottlenecks and solutions, would be beneficial.\n\nD/ Numerical stability and accuracy:\n\nGrid-based methods, such as FVM, are known for their numerical stability and accuracy, especially in simulating fluid dynamics. The paper does not provide an in-depth analysis of UPTs' numerical stability and accuracy compared to these methods. Detailed experiments and discussions on how UPTs ensure stability and accuracy over long simulation times would enhance the paper’s credibility.\n\nE/ Computational efficiency in complex geometries:\n\nParticle and grid-based methods have specific strategies to efficiently handle complex geometries, such as adaptive meshing in grid-based methods or kernel adjustments in particle methods. The paper does not clearly explain how UPTs manage complex geometries and whether they can maintain computational efficiency in such scenarios. More detailed experiments or case studies involving complex geometrical domains would be informative.\n\nF/ Interoperability with existing simulation tools:\n\nTraditional methods are often integrated into comprehensive simulation tools (e.g., OpenFOAM for grid-based methods or LAMMPS for particle-based methods). The paper does not discuss how UPTs can be integrated or used alongside these existing tools, which is crucial for practical adoption. Providing insights into interoperability and how UPTs can complement or enhance traditional methods within established simulation frameworks would be useful.\n\nG/  Handling of multiscale phenomena:\n\nParticle and grid-based methods have developed sophisticated techniques to handle multiscale phenomena, such as adaptive mesh refinement (AMR) in grid-based methods. The paper lacks a discussion on how UPTs address multiscale phenomena, which are common in many physical simulations. Including experiments or theoretical discussions on UPTs' capabilities in multiscale modeling would be beneficial.\n\n\n6/ Model conditioning standpoint:\n\nModel conditioning is crucial for ensuring that neural networks adapt accurately to varying inputs and simulation conditions. The Universal Physics Transformers (UPTs) framework has room for improvement in this aspect. Here are specific areas where the paper could enhance its discussion and implementation of model conditioning.\n\nA/ Insufficient detail on conditioning mechanisms:\n\nThe paper briefly mentions the use of feature modulation (e.g., DiT modulation) for conditioning UPTs to various inputs, such as the current timestep and boundary conditions. However, it lacks a detailed explanation of how these conditioning mechanisms are implemented and their impact on model performance. Providing more comprehensive descriptions and theoretical justifications for the chosen conditioning methods would help readers understand their effectiveness and potential limitations.\n\nB/ Limited analysis of conditioning performance:\n\nWhile the paper demonstrates that UPTs can handle different flow regimes and domains, it does not provide a thorough analysis of how well the conditioning mechanisms work across a broader range of scenarios. Including detailed experiments that specifically evaluate the performance of UPTs under different conditioning inputs, such as varying boundary conditions, initial states, and external forces, would strengthen the paper's claims.\n\nC/ Scalability of conditioning methods:\n\nThe scalability of the conditioning mechanisms is not thoroughly discussed. As simulations grow in complexity, the effectiveness of conditioning methods can degrade if not properly scaled. The paper should address how the conditioning mechanisms scale with increasing input dimensions and simulation complexities, providing insights into any potential bottlenecks and how they are mitigated.\n\nD/ Generalization to unseen conditions:\n\nOne of the strengths of neural operator networks is their ability to generalize to unseen conditions. The paper does not sufficiently explore how well UPTs generalize to entirely new boundary conditions or physical scenarios that were not present in the training data. Including experiments that test the generalization capabilities of UPTs to unseen conditions would provide valuable insights into the robustness of the conditioning methods.\n\nE/  Comparison with other conditioning approaches:\n\nThe paper does not compare its conditioning mechanisms with other advanced conditioning techniques used in neural operator networks or related fields. A comparative analysis of different conditioning approaches, such as those used in Fourier Neural Operators and DeepONets, would highlight the strengths and weaknesses of the methods employed in UPTs.\n\nF/  Impact of conditioning on training stability:\n\nConditioning mechanisms can significantly impact the training stability of neural networks. The paper does not discuss how the chosen conditioning methods affect the stability of UPT training, especially in the presence of complex and noisy data. Analyzing and addressing potential stability issues arising from conditioning would be beneficial.\n\nG/ Real-world application scenarios:\n\nWhile the paper presents conditioning in the context of fluid dynamics simulations, it does not discuss its applicability to other real-world scenarios that require complex conditioning. Exploring and providing evidence of UPTs' effectiveness in diverse applications, such as climate modeling or structural analysis, where conditioning to various external factors is crucial, would enhance the practical relevance of the framework.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your profound review and suggestions that helped us to improve our paper a lot. We addressed all your comments and followed all your suggestions. Please let us expand a bit on your comments.\n\n\n**Clarity and detail in methodology**\n\n\nIn order to make the paper more accessible we -- as suggested -- add pseudocode of a UPT forward pass with adaptive sampling. This together with sketch Figure 3a, sketch Figure 3b, and the encoder/approximator/decoder description of Section 3 should make the paper understandable wrt. all necessary details.\n\n**Scalability to extremely large datasets / high-resolution grids**\n\nWe agree that scalability to extremely large datasets is a desired property. In the paper, we have done our best to test this. Our models are trained using distributed training up to 32 A100 GPUs. We have self-simulated a dedicated dataset which is larger (> 50000 mesh points), and more complex (different flow regimes, different number of differently placed obstacles) as standard grid based datasets. As requested by several reviewers we have added experiments on the Shallow Water dataset and the Navier-Stokes dataset to the paper, which demonstrate the diverse applicability of UPT. However, even larger simulations than our transient flow simulations are beyond the scope of this paper and the available compute budget. Nevertheless, we have tested GPU usage of several model classes (Figure 2), which demonstrates that our framework is able to process meshes / particle clouds with up to 4 million nodes on a single GPU.\n\n**Insufficient analysis of generalization capabilities**\n\nWe agree that generalization capabilities are one of the most important aspects to test for neural operators. Therefore, in the transient flow experiment -- our largest experiment -- we test generalization across number of input / output points (discretization convergence), discretization across different flow regimes (input velocity), and generalization across different scenarios (differently placed obstacles).\n\n\n**Potential overfitting concerns**\n\nWe observe that the variable selection of the supernodes and the variable selection of input/output nodes allows for a strong regularization and data augmentation. The general training pipeline of UPT is therefore very robust against overfitting. This is further strengthened by the additional Navier-Stokes and Shallow water experiments.\n\n**Lack of detailed comparison with established numerical methods**\n\nWe are not claiming to be better than numerical methods. We compare with neural methods and especially focus on scaling. We follow the general procedures of the community, and report MSE, correlation time measures, and runtime comparisons. The latter two give an idea of the potential of UPT when compared with numerical methods. Similar to e.g. weather modeling (see Aurora [2], Pangu [3], ...) we follow the belief that benefits of neural operators will become more pronounced at scale.\n\n**Handling of boundary conditions**\n\nThe ShapeNet car experiment allows for testing neural operators on complicated geometries. UPT performs favorably, even when using a strongly reduced latent space representation. Additionally, in the transient flow example we have varying numbers of differently placed obstacles which can also be seen as complicated in-domain boundaries. In fact, UPT is specifically designed for large data on non-regular domains including different boundary conditions.\n\n**Model conditioning**\n\nWe have added a more detailed paragraph on the model conditioning choices in this paper. In our paper we follow closely the detailed studies described in [1].\n\n**Real world applications**\n\nReal world applications are beyond the scope of this work, but a very strong motivation for developing our framework. We note that current weather modeling such as Aurora [2] or Pangu [3] follow a similar design principle, but for regular gridded data, and without discretization convergence properties.\n\n**Question on complex boundary conditions**\n\nTransformers offer a flexible way to encode various boundary conditions.\nSkalar boundary conditions, such as inflow velocity or the current timestep, can be encoded via feature modulations. We use DiT modulation as it has shown strong performances in transformers.\n\nAdditionally, transformers offer a flexible mechanism to encode additional information by encoding the information into tokens and concatenating them to the supernodes or latent tokens. We make use of this flexible encoding in the ShapeNet-Car experiments where we additionally encode the signed distance function evaluated on a 3D grid via a CNN and feed the resulting tokens to the approximator.\n\n**Question on numerical stabilitiy and accuracy**\n\nTraditional methods require extremely small timesteps in order to preserve numerical stability when resolving the physics.\nAs neural operators only approximate the solution of the traditional solver, they can operate on much larger timesteps as their powerful modeling capabilities enables them to learn dynamics also from a coarse time resolution.\n\n\nFurthermore, UPT rollouts are fairly stable without any specific measures to ensure rollout stability. Sophisticated techniques to improve rollout stability (e.g., [4]) could be easily applied to UPTs to further enhance rollout stability. However, as these techniques typically impose a runtime overhead, we leave exploration thereof to future work.\n\n\n[1] Gupta et al., \"Towards Multi-spatiotemporal-scale Generalized PDE Modeling\", arXiv 2022, https://arxiv.org/abs/2209.15616\n\n[2] Bodnar et al., \"Aurora: A foundation model of the atmosphere\", arXiv 2024, https://arxiv.org/abs/2405.13063\n\n[3] Bi et al., \"Accurate medium-range global weather forecasting with 3D neural networks\", Nature 2023, https://www.nature.com/articles/s41586-023-06185-3\n\n\n[4] Brandstetter et al., \"Message passing neural pde solvers\", ICLR 2022, https://arxiv.org/abs/2209.15616", "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": "edaVVa6T90", "reviewer": "Reviewer_JKEE", "summary": "The authors present a new framework for unifying PDE surrogate modeling across domains. The proposed encoder, approximator, and decoder structure can accommodate PDEs of different discretizations and simulation types. Solutions are approximated in latent space, which aids in scalability and reducing computational cost of high-dimensional PDE data.", "questions": "-\tI’m curious about the general behavior of stability of error accumulation of your model during auto-regressive rollout. Are there specific training details that you use to mitigate this? Does using a latent space help mitigate this effect?\n-\tWhat are your thoughts on using a latent space vs. directly working in physics space when it comes to the accuracy of the solutions? Latent unrolling is obviously faster during inference, but even in the physical domain many neural surrogates are still orders of magnitude faster than numerical solvers with the added benefit of being simpler and potentially more accurate than working in a latent space.", "limitations": "-\tFor common benchmark problems on a regular grid, I don’t see why this architecture would be preferred over a baseline or fine-tuning a pretrained model. Most physical problems use a mesh or irregular grid, but certain problems such as weather or isotropic turbulence come to mind. \n-\tFixing the latent space is presented as a method to standardize different physics to a common space. However, more complex physics problems (3D turbulence vs. 2D laminar flows) may conceivably require larger latent spaces and the model would need to be retrained for different latent sizes.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "-\tThe method for encoding/decoding to a common latent space is exciting and seems to provide a good balance between compute/accuracy. \n-\tThe experiments are done with challenging fluids problems, which provides a good evaluation of model capabilities.\n-\tThe authors make a good effort to run a variety of experiments and provide good justification for when certain experiments cannot be performed. Furthermore, it is evident that the authors spent a lot of time thoroughly evaluating different model choices/hyperparameters and their effects.\n-\tThe generalization study and latent space evaluation are convincing.\n-\tThe authors consider Lagrangian fluid simulation, which is not commonly done in PDE surrogate modeling, and being able to accommodate this underexplored domain is good.\n-\tThe authors provide a good background of the fluid data used and the clarity of the paper is good.", "weaknesses": "-\tThe Unet, FNO, and GINO benchmarks are good, however, they are not specifically tailored for efficient computation. There have been prior works that address latent space modeling or efficient PDE modeling of 3D data, namely Latent Spectral Models [1], Latent Evolution of PDEs (LE-PDE) [2], and FactFormer [3]. I would like to see some reasoning or experiments as to why your model is better. \n-\tThe proposed architecture seems similar to OFormer [4]. It follows a similar encoding/propagating/decoding scheme as well as uses the transformer architecture. There are good improvements to adjust to higher dimensional and diverse data (e.g, Perceiver, supernodes, decoder design), but the broad architectural claims are shared with this previously proposed model (e.g. using a latent space, training an encoder/decoder, relying on transformer models for scalability). There are also other transformer-based works that present scalable architectures (DPOT [5], MPP [6]). I understand that these models do not include mechanisms for high-dimensional or irregular data, but I am curious to see how your model compares to these transformer-based models on a regular domain.\n-\tThere is no justification given for why your model is a neural operator. I know that not every paper that proposes a neural operator justifies the operator approximation theorem, and there are good discretization invariance results, but I would like to see some reasoning for how your model satisfies operator approximation. \n\tOverall, I think the paper is a good contribution and is a valuable step in extending neural operators to large, complex PDE problems. There is a lack of comparable benchmarks, however this does not significantly detract from the experimental rigor and results. \n1.\tHaixu Wu, Tengge Hu, Huakun Luo, Jianmin Wang, Mingsheng Long, Solving High-Dimensional PDEs with Latent Spectral Models. https://arxiv.org/abs/2301.12664\n2.\tTailin Wu, Takashi Maruyama, Jure Leskovec, Learning to Accelerate Partial Differential Equations via Latent Global Evolution, https://arxiv.org/abs/2206.07681\n3.\tZijie Li, Dule Shu, Amir Barati Farimani, Scalable Transformer for PDE Surrogate Modeling, https://arxiv.org/abs/2305.17560\n4.\tZijie Li, Kazem Meidani, Amir Barati Farimani, Transformer for Partial Differential Equations' Operator Learning, https://arxiv.org/abs/2205.13671\n5.\tZhongkai Hao, Chang Su, Songming Liu, Julius Berner, Chengyang Ying, Hang Su, Anima Anandkumar, Jian Song, Jun Zhu, DPOT: Auto-Regressive Denoising Operator Transformer for Large-Scale PDE Pre-Training, https://arxiv.org/abs/2403.03542\n6.\tMichael McCabe, Bruno Régaldo-Saint Blancard, Liam Holden Parker, Ruben Ohana, Miles Cranmer, Alberto Bietti, Michael Eickenberg, Siavash Golkar, Geraud Krawezik, Francois Lanusse, Mariel Pettee, Tiberiu Tesileanu, Kyunghyun Cho, Shirley Ho, Multiple Physics Pretraining for Physical Surrogate Models, https://arxiv.org/abs/2310.02994", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thorough review and helpful comments which were very useful to improve the paper. We address your points individually.\n\n\n**Why UPT is a better scalable latent space model**\n\nLatent Spectral Model uses geo-FNO (which could be consider as predecessor of GINO) to handle irregular grids. Therefore, it is somewhat similar to GINO as it has to map its input into a fixed regular grid representation. UPT doesn't map to a regular grid which makes it more flexible/expressive.\n\nLE-PDE only considers small-scale regular grid data and their design requires train and test data to have the *exact* same resolution (as mentioned in Appendix C of their paper). Therefore, its application to irregular grid data is impractical as a lot of data has different number of particles/mesh cells during test time. UPT can handle arbitrary input resolutions during training and generalizes well to higher input resolutions (as shown in Figure 5).\n\nFactFormer does not compress its input into a smaller latent space and therefore doesn't scale to large input sizes. Their method proposes an efficient attention mechanism, but their compute requirement would be similar to linear transformers and therefore become infeasible on large systems (as shown in Figure 2).\n\n**Comparison to transformer based models on regular grid**\n\nWe compare UPT on two regular grid benchmarks as shown in the general response and in the supplemental pdf Tables 1 to 3. UPT outperforms all compared methods (such as OFormer or DPOT) -- often by quite a margin -- without being specifically designed for regular grid datasets.\n\n\n\n**Justification why UPTS are neural operators**\n\nThank you for bringing up this important point. To avoid overloading the notation, we will provide a brief sketch of how universality can be established for our architecture. For transformer-based neural operators, universal approximation has been recently demonstrated in [1], Section 5. The arguments in this work are heavily based on [2], which establishes that nonlinearity and nonlocality are crucial for universality. By demonstrating that the attention mechanism can, under appropriate weight choices, function as an averaging operator, the results from [2] are directly applicable. For detailed proof, refer to Theorem 22 in [1].\nOur algorithm fits within this framework as well: we employ nonlinear, attention-based encoders and decoders (as allowed by the results of [2], Section 2) and utilize attention layers in the latent space.\n\nPlease let us know if you require further explanations. We will provide detailed information in an updated version of the manuscript.\n\n\n\n[1] Calvello et al., \"Continuum Attention for Neural Operators\", arXiv 2024, https://arxiv.org/abs/2406.06486\n\n[2] Lanthaler et al., \"The nonlocal neural operator: Universal approximation\" arXiv 2024, https://arxiv.org/abs/2304.13221\n\n\n**Rollout stability**\n\nOur training procedure doesn't use any special methods to stabilize the rollout because this typically comes at increased training costs. However, one could easily apply such techniques to UPT.\nAs we found our models to produce fairly stable rollouts even without additional techniques, we leave this direction for future work.\n\n\n\n\n\n**Thoughts on latent space vs physics space**\n\nWe see the latent rollout as a promising way to speedup inference, particularly for large-scale systems where encoding and decoding takes up most of the runtime.\nUPT offers the flexibility to invest resources into training it as a latent rollout model (by using the inverse encoding/decoding objectives) or to save resources during training (by omitting the reconstruction objectives) at the cost of increased inference time. We consider this a valuable tool to have for neural operators.\n\nAdditionally, the latent rollout enables applicability to Lagrangian simulations. As UPT models the underlying field instead of tracking individual particle positions it does not have access to particle locations at inference time. Therefore, autoregressive rollouts are impossible since the encoder requires particle positions as input. Using the latent rollout, it is sufficient to encode the initial particle positions into the latent space, which can then be propagated without the knowledge of any spatial positions. After propagating the latent space forward in time, one can simply query the latent space at arbitrary positions to evaluate the underlying field at given positions. We showcase this in Figure 7 where the latent space is queried with regular grid coordinates (white arrows).\n\n\n\n**Limitations of non-regular gridded data**\n\n\nWe are interested in complex physics simulations which are often simulated using e.g., finite element meshes. Several phenomena are modeled by particle-based simulations such as smoothed particle hydrodynamics, material point methods, or discrete element methods. Many phenomena even require a coupling of different aforementioned simulation types. This is very much driven by daily life engineering applications, and as such was a main motivation for developing UPT.\n\n\n\n\n\n**Limitations on fixed latent space**\n\n\nIn the current architecture of UPT, we consider a fixed latent space as it has proven to be an efficient way to compress the input into a fixed size representation to enable scaling to large-scale systems while remaining compute efficient.\nHowever, if an application requires a variable sized latent space, one could also remove the perceiver pooling layer in the encoder. With this change the number of supernodes is equal to the number of latent tokens and complex problems could be tackled by a larger supernode count. While we currently do not consider a setting where this is necessary, we show that the performance of UPT steadily increases with the number of supernodes during training (Figure 9 in Appendix C.4.3).\nAdditionally, UPT's performance improves when more supernodes are used during evaluation than were utilized during training (Figure 8 in Appendix C.4.2).", "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": "tZk0Xyr9tW", "reviewer": "Reviewer_CUfq", "summary": "This paper introduces a transformer-based architecture to scale neural operators to larger and more complex conditions involving spatiotemporal modeling. The novelty of this architecture is the use of the so-called latent rollout, which performs autoregressive modeling on the latent space, as opposed to the decoded space in other approaches. Authors claim that this architecture can operate without spatial structures, which is an inherent benefit of transformer-based approaches, and inspection of the latent-space in space-time. The authors show the usefulness of this approach across different physics-based problems. This is a meaningful contribution to design of NNs for physics applications, with potential broader applications in other ML domains. However, there are some questions that this reviewer hopes that the authors can address.", "questions": "1. Authors should be commended on the discussion on memory complexity in Line 105 to 119. However, this could be summarized as a table for cleaner presentation. Can the authors also include time complexity information for a more transparent comparison of the trade-offs associated with the different architectures? (I know this has been somewhat benchmarked in Table 1, but it could be useful to have some more theoretical information)\n2. What kind of position encoding is used for the UPT? This is unclear in the paper. Does the type of position encoding matter for Lagrangian vs Eulerian.\n3. For a more comprehensive display of results, could the authors include accuracy metrics (e.g., mean error) and memory benefits to  Table 1?\n4. This has been briefly discussed in Appendix A. But for further clarification, is the present architecture approach limited to physics-related applications? Or are there opportunities for extending the present approach towards other autoregressive ML applications, such as video or language modeling?", "limitations": "Limitations are well-discussed in Appendix A.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "1. Paper is clearly written, and appendices help with self-completeness.\n2. Proposed method is shown to be better than existing popular architectures for physics-based modeling.\n3. Code is provided for reproducibility.", "weaknesses": "1. A bit more analysis and information could be included to improve completeness. See Questions.\n2. Only test MSE has been used as an accuracy metric. Physics applications typically care about conservation of mass and performance near boundaries. However, the lack of further metrics is somewhat justified by the speedup and memory gains shown,", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your comments which helped to improve the paper. We addressed all your comments and followed all your suggestions.\n\n**Memory and time complexity of different architectures**\n\n\nWe provide theoretical complexities in Appendix C.6 in Table 3. Note that when scaling only the input size towards infinity to calculate an asymptotic runtime/memory complexity, the runtime complexity is mostly the same as the memory complexity as storing intermediate activation dominates the other factors. For example, transformers scale quadratically with the number of inputs M. For large M, a transformer needs to calculate the MxM attention matrix and also store it in memory. Therefore, it has $O(M^2)$ runtime and memory complexity.\nFor simplicity, we only consider the theoretical case and do not take optimizations such as hardware-aware implementations (e.g. [1]) into account that can trade-off additional computations for a reduced memory footprint.\n\nAs the theoretical complexities introduce many variables due to different architectures processing the input in vastly different ways, we believe that Figure 2 presents the practical complexities in a cleaner way and provide the theoretical complexities later in the appendix.\n\n[1] Dao et al., \"FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness\", NeurIPS 2022, https://arxiv.org/abs/2205.14135\n\n\n\n**Positional encoding type**\n\n\n\nWe employ the same positional embedding approach used in the original transformer paper [2], applying it separately to each dimension, as is also common in vision transformers [3]. Namely, we employ a combination of sine and cosine functions with different frequencies to encode each position. The revised version of the paper offers a clearer explanation of this positional embedding.\n\nWe also experimented with different types of positional embeddings such as the Fourier feature mapping from [4], which uses randomly sampled frequencies from a Gaussian distribution. In an experiment on TGV2D of our Lagrangian experiments it shows a slight improvement (see Figure 1 of the supplemental rebuttal pdf). The results show a slight improvement for longer rollouts, which we hypothesize is because this method doesn't overly emphasize features aligned with the axes. However, as improvements of other positional embeddings were minor, we stuck to the transformer positional embedding for simplicity.\n\n\n[2] Vaswani et al., \"Attention is all you need\", NeurIPS 2017, https://arxiv.org/abs/1706.03762\n\n[3] Dosovitskiy et al., \"An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale\", ICLR 2021, https://arxiv.org/abs/2010.11929\n\n[4] Tancik et al., \"Fourier features let networks learn high frequency functions in low dimensional domains\", NeurIPS 2020, https://arxiv.org/abs/2006.10739\n\n\n\n**Additional accuracy metrics and memory benefits of Table 1**\n\nWe included memory benefits in Table 3 of Appendix C.6. However, we find it difficult to assign a representative accuracy metric to each architecture\nas most methods are specialized for a certain type of task. For example, the accuracy of a CNN on a task with irregular data will naturally be worse than the accuracy of a GNN on that problem. UPT shows strong performances across various input domains and input scales, but as other methods are not as flexible we find it difficult to quantify their performance with a single metric.\n\n\n**Application in other domains**\n\nOur method could also be applied to other fields, such as video modeling. In that context, it would be possible to train on a dataset with varying resolutions. Since the decoder's output is a point-wise neural field, it could generate videos at arbitrarily high resolutions.\nHowever, in the context of language modeling, this approach is challenging to apply because language data is finite-dimensional.", "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": "MtuoHROcMZ", "reviewer": "Reviewer_QzDc", "summary": "This paper introduces Universal Physics Transformers (UPTs) to provide a unified learning paradigm for grid- or particle-based structures, enabling scalability across meshes and particles. UPTs mainly follow Encode-Process-Decode paradigm and allow queries at any space-time point through perceiver-like cross attention. To separate the responsibilities of individual components, the authors introduce inverse encoding and decoding losses in addition to the next-step prediction loss. Extensive experiments are conducted in diverse applications, including mesh-based fluid simulations and Lagrangian dynamics.", "questions": "- Given that the latent rollout approach does not provide a significant improvement in experimental results compared to autoregressive unrolling, can you elaborate on the potential reasons behind this? Specifically, do you believe that the lack of clear separation between the encoder, approximator, and decoder in this training framework contributes to this issue? Additionally, how effective do you find the inverse encoding and decoding processes? Considering the substantial training overhead and marginal performance gains, what are the key motivations for choosing the latent rollout method over other approaches? Are there any foreseeable optimizations or modifications that could enhance its efficacy?\n- How dose DiT modulation influences UPTs' performance? Additionally, which features have you found to be most effective as conditions for DiT modulation? \n- How do you specifically determine the number of supernodes in each region of mesh-structured data? After sampling nodes on the mesh, the original edge connections are lost. What impact does using a radius graph to replace the original mesh connectivity have on the results? Can you provide details on how this transformation affects the model's performance and accuracy?\n- How is the radius for the radius graph determined when aggregating information at the selected supernodes via a message passing layer? Additionally, do the circles of different supernodes overlap? \n- For each optimization step, are the supernodes fixed or do they change dynamically? \n- It is mentioned that the edge limit is imposed on nodes in a radius graph to prevent memory fluctuations. Is the edge limit satisfied if the radius is smaller than the set value, or if the limit is exceeded, are edges randomly removed to comply with the constraint?", "limitations": "Yes", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "UPTs provide a unified framework that can handle both grid- and particle-based structures, enhancing flexibility across different simulation types. UPTs reduce computational overhead by selecting and aggregating features on supernodes within mesh- and particle-based structures. The paper is well-written with a clear structure.", "weaknesses": "- Regardless of the structure of unstructured mesh/structured mesh/point clouds, when processing with UPTs, each node is actually modeled as a token, and a transformer architecture is built on this basis. However, there have been several works [1, 2, 3] that utilize transformers or attention mechanisms to model PDE problems. The novelty of using UPTs is somewhat limited in this way.\n- The paper misses comparisons to existing work that employs transformer and attention mechanisms for similar purposes.  While [1, 2, 3] have designed methods to reduce the computation overhead of transformer or attention mechanisms, the UPTs approach uses a simple random sampling to aggregate supernode features to reduce computation overhead. \n- While the latent rollout approach significantly accelerates inference speed, the experimental results do not show a clear improvement over the autoregressive unrolling via the physics domain. This raises concerns about the effectiveness and practicality of the latent rollout method.\n\n[1] GNOT: A General Neural Operator Transformer for Operator Learning\n\n[2] Transolver: A Fast Transformer Solver for PDEs on General Geometries\n\n[3] Transformer for Partial Differential Equations' Operator Learning", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review and helpful comments which were very useful to improve the paper. We address your points individually.\n\n\n**Comparison to other transformer methods**\n\n\n\n\n\n\n\nThe fundamental building principles of UPT are (i) an encoding that is designed to handle irregular grids of various sizes, (ii) a compressed and unified latent space representation that allows for efficient forward propagation in the latent space, and (iii) a neural field-type decoder. We consider all of these points as fundamental for scaling neural operators.\n\nGNOT processes each mesh point as a token and therefore quadratically scales with the number of mesh points.\nFurthermore the output can only be evaluated at the positions of the input.\n\nTransolver's concept of \"Physics-aware Tokens\" is somewhat similar to UPTs supernodes.\nHowever, only the substitutes attention part operates on a reduced space, whereas the FFN part of the transformer operates on the uncompressed space. Therefore, it is a type of linear complexity transformer, which quickly becomes infeasible for larger input sizes (see Figure 2).\nSimilarly, OFormer also operates on the uncompressed input.\n\nIn contrast, UPT heavily compresses the input, leading to sub-linear complexity w.r.t. input size in all transformer layers.\n\nAdditionally, we added comparisons to transformer baselines on regular grid datasets (see general response) where UPT outperforms e.g. OFormer by quite a big margin.\n\n\n\n\n\n\n\n\n\n\n\n\n\n**Benefits of the latent rollout**\n\nWhile the latent rollout does not provide a significant performance improvement, it is almost an order of magnitude faster.\nThe UPT framework allows to trade-off training compute vs inference compute. If inference time is crucial for a given application, one can train UPT with the inverse encoding and decoding objectives, requiring more training compute but greatly speeding up inference. If inference time is not important, one can simply train UPT without the reconstruction objectives to reduce training costs.\n\nAdditionally, the latent rollout enables applicability to Lagrangian simulations. As UPT models the underlying field instead of tracking individual particle positions it does not have access to particle locations at inference time. Therefore, autoregressive rollouts are impossible since the encoder requires particle positions as input. Using the latent rollout, it is sufficient to know the initial particle positions as dynamics can be propagated without any spatial positions. After propagating the latent space forward in time, one can simply query the latent space at arbitrary positions to evaluate the underlying field at given positions. We showcase this in Figure 7 where the latent space is queried with regular grid coordinates (white arrows).\n\n\nWe discuss a potential improvement for the latent rollout to make it more efficient in Appendix A. While we show in the paper that a latent rollout can be enabled via a simple end-to-end training, we think that it can be improved, e.g. via a two stage procedure of first training encoder/decoder in an autoencoder setting, followed by freezing encoder/decoder and training the approximator on the fixed pre-trained latent space akin to [1]. Such an approach would not require inverse encoding/decoding objectives as the separation of components is enforced through the multi-stage training.\n\n[1] Rombach et al., \"High-resolution image synthesis with latent diffusion models\", CVPR 2022, https://arxiv.org/abs/2112.10752\n\n**DiT modulation**\n\nConditioning onto external features is crucial to encode this external information into the model. We condition onto the current timestep and also onto the inflow velocity in the transient flow experiments.\nNote that this is not specific to UPT and we also apply conditioning to compared models.\n\n\n\n\n\n\n**Supernodes and radius graph creation**\n\n\nWe realize that our description in the paper is a bit misleading. Supernodes are, in the case of Eulerian data, sampled according to the underlying mesh and, in the case of Lagrangian data, sampled according to the particle density. A random sampling procedure which follows the mesh or particle density, respectively, allows us to put domain knowledge into the architecture. (We change \"randomly sampled\" to \"sampled according to the underlying mesh / underlying particle density\".)\nConsequently, complex regions are accurately captured, as these regions will be assigned more supernodes than regions with a low-resolution mesh or few particles.\n\n\n\n\nThe sampling of supernodes is done for each optimization step thus having a regularization effect.\n\nThe radius graph encodes a fixed region around each supernode. While one could use the original edge connections for Eulerian data, Lagrangian data does not have edges. Additionally, we employ randomly dropping input nodes as a form of data augmentation which makes using the original edge connections more complicated from an implementation standpoint. In contrast, a radius graph is agnostic to randomly dropping input nodes.\n\nWe choose the radius depending on the dataset. For each dataset, we first analyze the average degree of the radius graph with different radius values. We then choose the radius such that the degree is around 16, i.e. on average each supernode represents 16 inputs. We found our model to be fairly robust to the radius choice. Also, the circles of different supernodes can overlap. Therefore, the encoding of dense regions can be distributed among multiple supernodes.\n\nWe impose the edge limit of the radius graph by randomly dropping connections to preserve the average distance from supernode to input nodes.\n\n\n\n\n\nAdditionally, the radius graph facilitates Lagrangian settings where there is no predefined connectivity between particles.\n\n\nWe extended discussion in the corresponding sections and added a guide on how to choose these hyperparameters to the paper.", "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": "99RTXtlMzw", "reviewer": "Reviewer_1stF", "summary": "The authors presented Universal Physics Transformers (UPTs), a framework for efficient learning and scaling neural operators. UPTs offer flexibility in handling various data types, whether grid-based or particle-based. UPTs compress data into a low-dimensional latent space and perform dynamics propagation within this space, resulting in fast simulations. Additionally, the latent representation allows evaluation at any point in space-time. UPTs demonstrated strong performance when tested on various fluid dynamics problems.", "questions": "- What do you think causes the performance to plateau with increasing model size as observed in Figure 5? Is it because the error is already sufficiently low, preventing further enhancement?\n- Why do you believe that latent rollouts do not require stabilization?", "limitations": "The authors explained well the limitations.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The authors developed a model based on linear attention that handles both mesh-based and particle-based data effectively.\n- It leverages a low-dimensional latent space for efficient temporal dynamics propagation, enabling rapid evaluations.\n- They conducted very interesting Navier-Stokes experiments.\n- They demonstrated that UPT serves as a robust and efficient baseline.\n- In scalability experiments, UPT showed excellent performance, scaling efficiently with input size and outperforming other baselines across all tested scales.\n- They examined the model's convergence concerning the number of I/O points, obtaining good outcomes.", "weaknesses": "- The framework is labeled as universal but is only tested on fluid dynamics, lacking experiments on other PDEs.\n- Although the model scales well with input size, its scalability with model size appears poor based on a conducted experiment (Figure 5). While it outperforms other baselines, its scaling rate is significantly lower than that of models like GINO. Additional experiments are necessary to address this issue (if the UPT is meant to be used in alrge scale applications)\n- The scalability of UPT concerning training size remains unclear, which is an important property.\n- The stability of the latent rollout technique for long rollouts is uncertain, raising concerns about its general reliability", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your review and respond to the raised concerns below.\n\n**Experiments with other PDEs**\n\nUPTs are neural operators known to be applicable across various PDE types. Due to the nonlinear nature of the NS equations, they are notoriously more challenging to solve than parabolic or hyperbolic PDEs, like the heat or wave equations. Therefore, we hypothesize that UPTs should also generalize well to different PDEs.\n\n\nAdditionally, as requested by reviewer JKEE, we added comparisons of UPT with other methods on small-scale regular grid datasets. In total, UPT outperforms competitors on 5 diverse datasets that span over different dataset sizes (900 to 1M frames), resolutions (2K to 60K inputs), spatial dimensions (2D, 3D), simulation types (steady state, transient), different boundary conditions and specifications (Eulerian, Lagrangian).\n\n\n\n**Scalability with model size**\n\n\nWe see now that Figure 5 could suggest that UPT does not scale well with parameter size. However, the scaling of GINO only looks good because GINO underperforms in contrast to UPT (GINO-68M is worse than UPT-8M). As you correctly identified, for well trained (UPT) models it gets increasingly difficult to improve the loss further. Ideally, one would show this by training larger GINO models, however this is *prohibitively* expensive (68M parameter models already take 450 A100 hours per model). We therefore go in the other directions and train even smaller UPT models that achieve a similar loss to the GINO models and compare scaling there. In Figure 2 of the supplemental pdf, we compare UPT 1M/2M/8M against GINO 8M/17M/68M. UPT shows similar scaling on that loss scale.\n\nWe see how this could be easily misinterpreted from Figure 5 and adjust the paper accordingly to remove this misunderstanding.\n\nWe strongly hypothesis that the effect of larger UPT models would become apparent in even more challenging settings or even\nlarger datasets. However, challenging large-scale datasets are hard to come by, which is why we created one ourselves. Creating even larger and more complex ones is beyond the scope of our work as it exceeds our current resource budget, but it is definitely an interesting direction for future research/applications.\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n**Scalability with dataset size**\n\nWe added experiments to show the scalability and data efficiency of UPTs by training UPT-8M on subsets of the data used for the transient flow experiments (we train on 2K and 4K out of the 8K train simulations). The results in Figure 3 of the supplementary rebuttal pdf show that UPT scales well with data and is data efficient, achieving comparable results to GINO-8M with 4x less data.\n\nWe also show that UPTs can handle various dataset sizes. ShapeNet-Car is a small-scale dataset consisting of 889 car shapes with 3.6K mesh points each. TGV3D is a bit bigger with 8K particles per simulation and 200 simulations of length 61 (12K frames). The dataset for our transient flow experiments contains around 50K mesh cells per simulation with 10K simulations of length 100 (1M frames).\nUPT shows strong performances across all considered dataset sizes.\n\n\nAdditionally, UPT consists of mostly transformer blocks, which have demonstrated outstanding scalability in other domains such as language modeling [1] and computer vision [2].\n\n\n[1] Kaplan et al., \"Scaling laws for neural language models\", arXiv 2020, https://arxiv.org/abs/2001.08361\n\n\n[2] Zhai et al., \"Scaling vision transformers\", CVPR 2022, https://arxiv.org/abs/2106.04560\n\n\n**Stability of latent rollout**\n\n\nWe found UPT to be fairly stable without any special techniques to stabilize the rollout (e.g. [3]). However, such methods could further improve UPTs performance but these methods are not specific to UPT and typically require additional computations during training. Therefore, we leave exploration of this combination to future work.\n\nAdditionally, the latent rollout opens new avenues to potentially apply existing stabilization tricks with less compute. For example, one could do the forward propagation for the stabilization technique from [4] in the latent space, which would greatly reduce training costs thereof. However, as these tricks can be tricky to train (e.g. due to requiring a precise trade-off between training the next-step prediction vs n-step prediction), we leave this direction to future work.\n\n\n\n[3] Lippe et al., \"Pde-refiner: Achieving accurate long rollouts with neural pde solvers\", NeurIPS 2023, https://arxiv.org/abs/2308.05732\n\n[4] Brandstetter et al., \"Message passing neural pde solvers\", ICLR 2022, https://arxiv.org/abs/2209.15616", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank all reviewers for their positive feedback and for their constructive comments and suggestions.
We are pleased to see that the reviewers highlighted the clarity of our paper and appreciated our detailed and thorough experiments. Several reviewers recognized the efficiency of our approach by encoding and decoding into a latent space as well as UPTs ability to handle both Lagrangian and Eulerian data.



To address the questions of the reviewers, we included the following new experiments. Please find visualizations thereof in the supplemental rebuttal pdf.

**Experiments on a regular grid Navier-Stokes equations dataset**

We run comparisons against different Transformer baselines on regular gridded Navier-Stokes equations data (data, baseline results and evaluation protocol taken from [1]). UPT outperforms all compared methods, some of which are specifically designed for regularly gridded data.

As baseline transformers often train small models, we first compare on a small scale, where UPT significantly outperforms other models.

| Model | # Params | Rel. L2 Error  |
|---|---|---|
| FNO | 0.5M | 9.12 \%  |
| FFNO | 1.3M | 8.39 \%  |
| GK-T |  1.6M | 9.52 \%  |
| GNOT |  1.8M | 17.20 \%  |
| OFormer | 1.9M |  13.50 \%  |
| UPT-T | 1.8M | **5.08** \%  |

We also compare on larger scales, where UPT again outperforms competitors, even if they train much larger models or pre-train (PT) on more data followed by fine-tuning (FT) on the Navier Stokes dataset.


| Model | # Params | Rel. L2 Error  |
|---|---|---|
| DPOT-Ti |  7M | 12.50 \%  |
| DPOT-S | 30M |  9.91 \%  |
| DPOT-L (PT) | 500M | 7.98 \%  |
| DPOT-L (FT) | 500M | 2.78 \%  |
| DPOT-H (PT) | 1.03B | 3.79 \%  |
| CViT-S | 13M | 3.75 \%  |
| CViT-B | 30M | 3.18 \%  |
| UPT-S | 13M | 3.12 \%  |
| UPT-B | 30M | **2.69** \%  |


[1] Hao et al., "DPOT: Auto-Regressive Denoising Operator Transformer for Large-Scale PDE Pre-Training", ICML 2024, https://arxiv.org/abs/2403.03542



**Experiments on a regular grid Shallow Water equations dataset**


We run comparisons against UNet, FNO, Dilated ResNet variants on regular gridded Shallow Water eqations data (data, baseline results and evaluation protocol taken from [2]). Similarly to the Navier-Stokes experiments, UPT can outperform methods which are specifically designed for regularly gridded data.


| Model | # Params | Rel. L2 Error  |
|---|---|---|
| DilResNet | 4.2M | 13.20 \%  |
| U-Net | 148M | 5.68 \%  |
| FNO | 268M | 3.97 \%  |
| CViT-S | 13M | 4.47 \%  |
| UPT-S | 13M | **3.96** \%  |

[2] Wang et al., "Bridging Operator Learning and Neural Fields:. A Unifying Perspective", arXiv 2024, https://arxiv.org/abs/2405.13998

**Impact of different positional encodings**

We added ablation results with different positional encoding.
Using a Fourier feature mapping from [3] resulted in an slightly better performance in the TGV2D experiments. We visualize this in Figure 1 of the supplemental rebuttal pdf.

[3] Tancik et al., "Fourier features let networks learn high frequency functions in low dimensional domains", NeurIPS 2020, https://arxiv.org/abs/2006.10739


**UPT scaling clarifications**

We added experiments to clarify the scaling of UPTs in the transient flow experiments by training small UPT models that perform similar to larger GINO models. We visualize the results in Figure 2 of the supplemental rebuttal pdf, which shows that UPT scales well when increasing parameter counts. This experiment should clarify that UPTs do scale well with parameter counts, but as the testloss of UPT (in Figure 5 of the paper) is significantly lower than for GINO, it is much harder to improve it further.


**Scaling with dataset size**

We investigate scaling behavior of UPT w.r.t. dataset size by training a UPT-8M model on smaller subsets in the setting of our transient flow experiments. Figure 3 in the supplementary rebuttal pdf visualizes the scaling curves. UPT scales well with data and achieves comparable results to GINO-8M with 4x less data.


**Further additions to the paper**

We followed the reviewers comments and suggestions, e.g., by including pseudocode of an UPT forward pass, detailed paragraphs on positional encoding and conditioning mechanisms, implementation details of the supernode pooling/radius graph and additional related works.


In addition, we followed the reviewers comments and suggestions, e.g. by discussing their raised points in the Discussion section or added clarifications, contextualization, and references.

Meta-review: This paper studies an encoder-decoder arch of Transformers as neural operators for several PDE tasks, mainly flow problems. The authors had nice discussions with the reviewers. However, after reading the paper, I felt quite unsatisfied with the method presented, especially with regard to the scaling properties. The scaling law of Transformer considers the increasing "complexity" of the dataset. Yet, here in this paper, the models are still trained separately for each dataset (4.1-4.3), and the study of scaling is with respect to the size of discretizations (# of super nodes). Note that this "fact", even though ad-hoc, is quite commonly acknowledged in the CFD community from the work of Temam, Shen, Tadmor, and many others: when isotropic turbulence formed (so that one can observe the energy cascade), the nonlinearity is quite "simple" and has a lower intrinsic dimension (the reason reduced-basis approach works for this regime) if the initial conditions are sampled from an energy distribution with certain decay. In this regard, unlike language, the "complexity" of the dataset does not increase with the increase of the tokens. Another note is that OpenFOAM's solver prioritizes speed over the conformity to the NSE in the continuous Hilbertian framework, e.g., RANS already makes things tamer than DNS. Nevertheless, the architecture is simple and may have long lasting impact on the operator learning community.
Decision: Accept (poster)