Universal Physics Transformers: A Framework For Efficiently Scaling Neural Operators

Authors: benedikt  alkin
andreas  fürst
simon  schmid
lukas  gruber
markus  holzleitner
johannes  brandstetter
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-13 12:15:04

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.

Main paper: 
 introduction: 
 In scientific pursuits, extensive efforts have produced highly intricate mathematical models of physical phenomena, many of which are naturally expressed as 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 [51, 61, 47] , have emerged as a computationally efficient alternative [99, 119] , and impact e.g., weather forecasting [48, 8, 1] , molecular modeling [7, 4] , or computational fluid dynamics [107, 30, 51, 43, 31, 19] . Additional to computational efficiency, neural surrogates offer potential to introduce generalization capabilities across phenomena, as well as generalization across characteristics such as boundary conditions or PDE coefficients [66, 14] . Consequently, the nature of neural operators inherently complements handcrafted numerical solvers which are characterized by a substantial set of solver requirements, and mostly due to these requirements tend to differ among sub-problems [3] . However, similar to their numerical counterparts, different neural network techniques are prevalent across applications. For example, when contrasting particle-and grid-based dynamics in computational fluid dynamics (CFD), i.e., Lagrangian and Eulerian discretization schemes. This is in contrast to other areas of deep learning where the flexibility of transformers [106] 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 [10] that are pretrained on huge passive datasets [25, 36] .We introduce Universal Physics Transformers (UPTs), an efficient and unified neural operator learning paradigm with strong focus on scalibility over a wide range of spatio-temporal problems. UPTs flexibly encode different grids, and/or different number of particles into a compressed latent space representation which facilitates 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 at any point in space-time. UPTs operate without grid-or particle-based latent structures, and demonstrate the beneficial scaling-behavior of transformer backbone architectures. Figure 1 sketches the UPT modeling paradigm.We summarize our contributions as follows: (i) we introduce the UPT framework for efficiently scaling neural operators; (ii) we formulate encoding and decoding schemes such that dynamics can be propagated efficiently in a compressed and fixed-size latent space; (iii) we demonstrate applicability on diverse applications, putting a strong research focus on the scalability of UPTs.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 on a wide range of physical simulations. Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
background: 
 Partial differential equations. We focus on experiments on (systems of) PDEs, that evolve a signal u(t, x) = u t (x) ∈ R d in a single temporal dimension t ∈ [0, T ] and m spatial dimensions x ∈ U ⊂ R m , for an open set U . With 1 ≤ l ∈ N, systems of PDEs of order l can be written as F (D l u t (x), . . . , D 1 u t (x), ∂ l ∂t l u t (x), . . . , ∂ ∂t u t (x), u t (x), x, t) = 0, for x ∈ U, t ∈ [0, T ] ,(EQUATIONwhere F is a mapping to R d , and for i = 1, ..., l, D i denotes the differential operator mapping to all i-th order partial derivatives of u with respect to the spacial variable x, whereas ∂ i ∂t i outputs the corresponding time derivative of order i. Any l-times continuously differentiable u : [0, T ]×U → R d fulfilling the relation Eq. ( 1) is called a classical solution of Eq. (1) . Also other notions of solvability (e.g. in the sense of weak derivatives/distributions) are possible, for sake of simplicity we do not go into details here. Additionally, initial conditions specify u t (x) at time t = 0 and boundary conditions B[u t ](x) at the boundary of the spatial domain.We work mostly with the incompressible Navier-Stokes equations [98] , which e.g., in two spatial dimensions conserve the velocity flow field u(t, x, y) : [0, T ] × R 2 → R 2 via:EQUATIONwhere u • ∇u is the convection, i.e., the rate of change of u along u, µ is the viscosity parameter, µ∇ 2 u the viscosity, i.e., the diffusion or net movement of u, ∇p the internal pressure gradient, and f an external force. The constraint ∇ • u = 0 yields mass conservation of the Navier-Stokes equations.A detailed depiction of the involved differential operators is given in Appendix B.1.Operator learning. Operator learning [60, 61, 52, 51, 47] 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 U, V to be Banach spaces of functions on compact domains X ⊂ R dx or Y ⊂ R dy , mapping into R du or R dv , respectively. The goal of operator learning is to learn a ground truth operator G : U → V via an approximation Ĝ : U → V. This is usually done in the vein of supervised learning by i.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 a given data set consisting of N function pairs (u i , v i ) = (u i , G(u i )) ⊂ U × V, i = 1, ...N , we aim to learn Ĝ : U → V, so that G can be approximated in a suitably chosen norm.In the context of PDEs, G can e.g. be the mapping from an initial condition u(0, x) = u 0 (x) to the solutions u(t, x) = u t (x) of Eq. ( 1) at all times. In the case of classical solutions, if U is bounded, U can then be chosen as a subspace of C( Ū , R d ), the set of continuous functions from domain Ū (the closure of U ) mapping to R d , whereas V ⊂ C([0, T ] × Ū , R d ), so that U or V consist of all l-times continuosly differentiable functions on the respective spaces. In case of weak solutions, the associated spaces U and V can be chosen as Sobolev spaces.1 2 4 8 16 32 [52] radius ✓ ✓ ✗ ✗ FNO [51] global ✗ ✓ ✗ ✗ GINO [54] global✓ ✓ ✓ ✗ UPT global ✓ ✓ ✓ ✓Figure 2 : Qualitative exploration of scaling limits. Starting from 32K input points (scale 1), we train a 68M parameter model for a 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 a 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. n nodes × h, where usually n nodes = n points . For large number of inputs, this becomes infeasible as every layer has to process a large number of tokens. Contrary, UPTs compress the inputs into a low-dimensional latent space, which drastically decreases computational requirements. Different architectures and their scaling limits are compared in Fig. 2 .Question: Does the paper discuss the limitations of 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 in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
universal physics transformers: 
 Problem formulation. Our goal is to learn a mapping between the solutions u t and u t ′ of Eq. ( 1) at timesteps t and t ′ , respectively. Our dataset should consist of N function pairs(u t i , u t ′ i ), i = 1, .., N , where each u t i is sampled at k spatial locations {x 1 i , . . . , x k i } ∈ U . Similarly, we query each output signal ût ′ i at k ′ spatial locations {y 1 i , . . . , y k ′ i } ∈ U . Then each input signal can be represented by u t i,k = (u t i (x 1 i ), . . . , u t i (x k i ))T ∈ R k×d as a tensor of shape k × d, similar for the output. For particle-or mesh-based inputs, it is often simpler to represent the input as graph G = (V, E) with k nodes {x 1 i , . . . x k i } ∈ V , edges E and node features {u t i (x 1 i ), . . . , u t i (x k i )}. Architecture desiderata. We want Universal Physics Transformers (UPTs) to fulfill the following set of desiderata: (i) an encoder E, which flexibly encodes different grids, and/or different number of particles into a unified latent representation of shape n latent × h, where n latent is the chosen number of tokens in the latent space and h is the hidden dimension; (ii) an approximator A and a training procedure, which allows us to forward propagate dynamics purely within the latent space without mapping back to 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 E is to compress the input signal u t i , which is represented by a point cloud u t i,k . Importantly, the encoder should learn 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 E first embeds k points into hidden dimension h, adding position encoding [106] to the different nodes, i.e., u t i,k ∈ R k×d → R k×h . In the first hierarchy, information is exchanged between local points and a selected set of n s supernode points. For Eulerian 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 s supernodes on the mesh, choosing n s such that the mesh characteristic is preserved. Similarly, in the Lagrangian discretization scheme, choosing supernodes based on particle positions provides the same advantages as selecting them based on the mesh.Information is aggregated at the selected n s supernodes via a message passing layer [28] using a radius graph between points. Importantly, messages only flow towards the n s supernodes, and thus the compute complexity of the first hierarchy scales linearly with n s . The second hierarchy consists of transformer blocks [106] followed by a perceiver block [40, 39] with n latent learned queries of dimension h. To summarize, the encoder E maps u t i ∈ U to a latent space viaE : u t i ∈ U evaluate ----→ u t i,k ∈ R k×d embed ---→ R k×h MP --→ R ns×h transformer ------→ R ns×h perceiver ----→ z t i ∈ R nlatent×h, where tyically n latent ≪ n s < 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 an implicit "importance sampling" of the underlying simulation. Additional implementation details are provided in Appendix F.1. This encoder design projects into a fixed size latent space as it is an 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 is equal to the number of latent tokens and complex problems could be tackled by a larger supernode count.Approximator. The approximator propagates the compressed representation forward in time. As n latent is small, forward propagation in time is fast. We employ a transformer as approximator.A : z t i ∈ R nlatent×h → z t ′ i ∈ R nlatent×h .Notably, the approximator can be applied multiple times, propagating the signal forward in time by ∆t each time. If ∆t is small enough, the input signal can be approximated at arbitrary future times t ′ .Decoder. The task of decoder D is to query the latent representation at k ′ arbitrary locations to construct the prediction of the output signal u t ′ i at time t ′ . More formally, given the output positions {y 1 i , . . . , y k ′ i } ∈ U at k ′ spatial locations and the latent representation z t ′ i , the decoder predicts the output signalu t ′ i,k ′ = (u t ′ i (y 1 i ), . . . , u t ′ i (y k ′ i )) T at these spatial locations at timestep t ′ , D : (z t ′ i , {y 1 i , . . . , y k ′ i }) → ût ′ i,k ′ ∈ R k ′ ×d .The decoder is implemented via a perceiver-like cross attention layer using a positional embedding of the output positions as query and the latent representation z t ′ i as keys and values. Since there is no interaction between queries, the latent representation can be queried at arbitrarily many positions 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 t 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 a dimension-wise scale, shift and gate operation that are 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, x) = u 0 (x) is encoded into a latent representation z 0 . From there on, instead of autoregressively feeding the decoder's prediction into the encoder, UPTs propagate z 0 forward in time to z t ′ through iteratively applying the approximator A in the latent space. We call this procedure latent rollout. Especially for large meshes or many particles, the benefits of latent space rollouts, i.e. fast inference, pays off.To enable latent rollouts, the responsibilities of encoder E, approximator A and decoder D need to be isolated. Therefore, we invert the encoding and decoding by means of two reconstruction losses during training as visualized in Fig. 3 . First, an inverse encoding is performed, wherein the input u For the modeling of temporal PDEs, a common scheme is to map the input solution at time t to the solution at next time step t ′ [51, 13, 95] . Especially for systems that are modeled by graphbased representations, predicted accelerations at nodes are numerically integrated to model the time evolution of the system [87, 81] . Recently, equivariant graph neural operators [116] were introduced which model time evolution via temporal convolutions in Fourier space. More related to our work are methods that propagate dynamics in the latent space [50, 111] . Once the system is encoded, time evolution is modeled via LSTMs [111] , or even linear propagators [63, 73] . In Li et al. [53] , attention-based 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. [16] use recurrent MLPs for temporal updates within the latent space, while utilizing a 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 models [85] where a pre-trained compression model is used to compress the input into a latent space from which a 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.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? Answer: [NA]• Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
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 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 baseline regular grid methods on regular grid datasets D.3. 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 in 3D space. We randomly create a train/test split containing 700/189 samples. We regress the pressure at each surface point with a mean-squared error (MSE) loss and sweep hyperparameters per model. Due to the small scale of this dataset, we train the largest possible model that is able to generalize the best.Training even larger models resulted in a 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, 31] , FNO [51] ) are significantly smaller and range from 15M to 100M. Additional details are listed in Appendix D.4.ShapeNet-Car is a small-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 Model-wise, UPT uses the hierarchical encoder setup with all optional components depicted in Fig. 3 .A message passing layer aggregates local information into n s = 2048 randomly selected supernodes, a transformer processes the supernodes and a perceiver pools the supernodes into n latent = 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-Spectral" 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. 5 , shows that UPTs outperform compared methods by a 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 A100 hours, respectively. Figure 4 shows a rollout and Appendix D.5.1 presents additional ones. We also study out-of-distribution generalization (e.g. more obstacles) in Appendix D.5.6.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. 5 where we vary the number of input/output points and the number of supernodes. We use the 68M models without any retraining, i.e., we test models on "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 or 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.Finally, we evaluate training with inverse encoding and decoding techniques, see Fig. 3 . We investigate the impact of the latent rollout by training our largest model -a 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. 1. However, in its current implementation the latent rollout requires a non-negligible overhead during training. We discuss this limitation in Appendix A.Scaling particle-based methods such as discrete element methods or smoothed particle hydrodynamics to 10 million or more particles presents a significant challenge [117, 9] , yet it also opens a distinctive opportunity for neural surrogates. Such systems are far beyond the scope of this work. We however present a framing of how to model such systems via UPTs such that the studied scaling properties of UPTs could be exploited. In order to do so, we demonstrate how UPTs capture inherent field characteristics when applied to Lagrangian SPH simulations, as provided in LagrangeBench [101] .Here, GNNs, such as Graph Network-based Simulators (GNS) [87] and Steerable E(3) Equivariant Graph Neural Networks (SEGNNs) [14] are strong baselines, where predicted accelerations at the nodes are numerically integrated to model the time evolution of the particles. In contrast, UPTs learn underlying dynamics without dedicated particle-structures, and propagate dynamics forward without the guidance of numerical time integration schemes. An overview of the conceptual differences between GNS/SEGNN and UPTs is shown in Fig. 6 . We use the Taylor-Green vortex dataset in three dimensions (TGV3D). The TGV system was introduced as a test scenario for turbulence modeling [97] . It is an unsteady flow of a decaying vortex, displaying an exact closed form solution of the incompressible Navier-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 to be queried at 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 in 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 D.6.x t x t ′ ∆t Model predicts ât v t ′ = v t + ât ∆t x t ′ = x t +v t ′ ∆t x t x t ′ v t v t ′For UPT training, we input two consecutive velocities of the particles in the dataset at timesteps t and t -1, and the respective particle positions. We regress two consecutive velocities at a later timesteps {t ′ -1, t ′ } = {t + ∆T -1, t + ∆T } with mean-squared error (MSE) objective. For all experiments we use ∆T = 10∆t. We query the decoder to output velocities at 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 7 compares the rollout performance of GNS, SEGNN and UPT and shows the speedup of both methods compared to the SPH solver. The results demonstrate that UPTs effectively learn the underlying field dynamics. 
discussion: 
 Potential for extreme scale. UPTs consist of mainly transformer [106] layers which allows application of the same scaling and parallelism principles as are commonly used for training extremescale language models. For example, the recent work Llama 3 [27] trained on up to 16K H100 GPUs. Additionally, 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. 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 7 . 
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. UPTs operate without grid-or particle-based latent structures, enabling flexibility across meshes and number of particles. The UPT training procedure separates responsibilities between components, allowing a 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. 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA] Justification: -Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. 12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected? Answer: [Yes] Justification: Authors of code and data are cited accordingly. Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. 
 
 Appendix: Latent rollout. We show that the latent rollout can be enabled via a simple end-to-end training procedure, but ultimately think that it can be improved by more delicate training procedures such as a 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 UPTs to other domains for future work. We additionally show that UPTs are neural operators and since neural operators have been shown to generalize well across domains, this should also hold for UPTs.Large-scale Lagrangian simulations. A particularly intriguing direction is to apply UPTs to largescale Lagrangian simulations. However, there do not exist readily available large-scale Lagrangian datasets. Therefore, we 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 on the transient flow dataset, which contains up to 59K mesh points. In this setting, we train with a distributed setup of up to 32 A100 GPUs.Unifying Lagrangian and Eulerian simulations. UPTs can encode both Lagrangian and Eulerian simulations. It is therefore a natural future direction to exploit those different modalities. Especially, since particle-and grid-based simulations are used to describe different phenomena but model similar underlying dynamics. It is however to note that our method is designed primarily with the purpose of efficient scaling, and multi-modality training follows as side-concept thereof.This appendix discusses the Navier-Stokes equations and selected numerical integration schemes, which are related to the experiments presented in this paper. This is not 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-Stokes equations, relate compressible and incompressible formulations, and discuss why Navier-Stokes equations are difficult to solve. Second, we discuss turbulence as one of the fundamental challenges of computational fluid dynamics (CFD), explicitly working out the difference for two dimensional and three dimensional turbulence modeling. Third, we introduce Reynolds-averaged Navier-Stokes equations (RANS) as an numerical approach for steady-state turbulence modeling, and the SIMPLE and PISO algorithms as steady-state and transient solvers, respectively. Finally, we discuss Lagrangian discretization schemes, focusing on smoothed particle hydrodynamics (SPH).In our experiments, we mostly work with the incompressible Navier-Stokes equations [98] . In two spatial dimensions, the evolution equation of the velocity flow field u :[0, T ] × U ⊂ R 3 → R 2 , u(t, x 1 , x 2 ) = (u 1 (t, x 1 , x 2 ), u 2 (t, x 1 , x 2 )) with internal pressure p(x 1 , x 2 ) = (p 1 (x 1 , x 2 ), p 2 (x 1 , x 2 )) and external force field f (x 1 , x 2 ) = (f 1 (x 1 , x 2 ), f 2 (x 1 , x 2 )) is given via:EQUATIONThis system is usually written in a shorter and more convenient form:EQUATIONwhereu • ∇u = ( u 1 u 2 ) •     ∂u 1 ∂x 1 ∂u 2 ∂x 1 ∂u 1 ∂x 2 ∂u 2 ∂x 2     (5)is called the convection, i.e., the rate of change of u along u, andEQUATIONthe viscosity, i.e., the diffusion or net movement of u with viscosity parameter µ. The constraintEQUATIONyields mass conservation of the Navier-Stokes 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 of 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 ∇ • u = 0 is an implicit constraint on the velocity fields. Additionally, the nonlinear nature of the NS equations makes them notoriously harder to solve than parabolic or hyperbolic PDEs such as the heat or wave equation, respectively. Also, the occurrence of turbulence as discussed in the next subsection which is due to the fact that small and large scales are coupled such that small errors, can have a large effect on the computed solution.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.Turbulence [82, 64] is one of the key aspects of CFD and refers to 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 a 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 numbersEQUATION)where ρ is the density of the fluid, ν is the velocity of the fluid, L is the linear dimension, e.g., the width of a pipe, and µ 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 vs 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 a 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. In three dimensions, this energy cascade is responsible for the continuous formation of vortex structures at various scales, and thus for the emergence of high-frequency 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 ω is the curl of the flow field, i.e., ω = ∇ × u, with × here denoting the cross product. This derivation can be found in many standard texts on Navier-Stokes equation , e.g. Tsai [104, Section 1.4.], and for the reader's convenience we briefly repeat the derivation here. We start using Eq. ( 4), setting f = 0 for simplicity:EQUATIONwhere we applied the dot product rule for derivatives of vector fields:EQUATIONNext, we take the curl of the right and the left hand side:EQUATION)EQUATION)using the property that the curl of a gradient is zero. Lastly, using that the divergence of the curl is again zero (i.e. ∇•ω = 0), via the identity∇×(u × ω) = (ω • ∇) u-(u • ∇) ω+u ∇ • ω =0 -ω ∇ • u =0 ,we obtain:EQUATIONwith the material derivativeDω Dt = ∂ω ∂t + (u • ∇) ω.For three dimensional flows the material derivative of each component ω i can be written asDω i Dt = 3 j=1      ω j ∂u i ∂x jvortex turning and stretchingEQUATIONfor example picking i = 2:EQUATIONFor two-dimensional flows, as represented byEQUATIONthe terms for vortex turning and vortex stretching vanish, as is evident in the material derivative of ω 3 , given byEQUATIONConsequently, the length and the angle of a vortex do not change in two 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 [81, 51, 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, 56] .Various approaches are available for numerically addressing turbulence, with different schemes suited to 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 not necessitate additional closure equations, a 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-scales 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 to 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 timeaveraging. The foundational concept is Reynolds decomposition, attributed to Osborne Reynolds, where an instantaneous quantity is decomposed into its time-averaged and fluctuating components.Turbulent scenarios where the fluid conditions change over time are an example of transient flows.These are typically harder to model/solve than steady state flows where the fluid properties exhibit only negligible changes over time.In the following, we discuss Reynolds-averaged Navier-Stokes equations (RANS) as a numercial approach for steady-state turbulence modeling, and the SIMPLE and PISO algorithms as steady-state and transient numerical solvers, respectively.Reynolds-averaged Navier-Stokes equations (RANS) are used to model time-averaged fluid properties such as velocity or pressure that result in a steady state which does not change over time. Writing (4) in Einstein notation we getEQUATION)Taking the Reynolds decomposition [84] g(t, x 1 , x 2 ) := ḡ(x 1 , x 2 ) + g ′ (t, x 1 , x 2 ) with ḡ(x 1 , x 2 ) := lim T →∞ 1/T T 0 g(x 1 , x 2 , t)dt being the time-average on [0, T ] of a scalar valued function g, and splitting each term of both equations into it's time-averaged and fluctuating part we get∂(ū i + u ′ i ) ∂t = -(ū j + u ′ j ) ∂(ū i + u ′ i ) ∂x j + µ ∂ 2 (ū i + u ′ i ) ∂x j ∂x j - ∂(p + p ′ ) ∂x i + ( fi + f ′ i ) 0 = ∂(ū i + u ′ i ) ∂x i .Time-averaging these equations together with the property that the time-average of the fluctuating parts equals zero results inūj ∂ ūi ∂x j + u ′ j ∂u ′ i ∂x j =µ ∂ 2 ūi ∂x j ∂x j - ∂ p ∂x i + fi (21) 0 = ∂ ūi ∂x i .Using the the mass conserving equation ( 20) the momentum equation ( 21) can be rewritten asūj ∂ ūi ∂x j = fi + ∂ ∂x j µ ∂ ūi ∂x j + ∂ ūj ∂x i -pδ ij -u ′ i u ′ jwhere -u ′ i u ′ j is called the Reynolds stress which needs further modeling to solve the above equations. More specifically, this is referred to as the Closure Problem which led to many turbulence models such as k-ϵ [32] or k-ω [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 be closed by modeling the Reynolds stress as a 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 [11] for 3-d incompressible Navier Stokes:-u ′ i u ′ j = ν t ∂ ūi ∂x j + ∂ ūj ∂x i - 2 3 kδ ijwith the turbulence eddy viscosity ν t > 0 and the turbulence kinetic energy k = 1/2 u ′ i u ′ i based on Boussinesq's hypothesis that turbulent shear stresses act in the same direction as shear stresses due to the averaged flow. The k-ϵ 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 k and dissipation ϵ∂(ρk) ∂t + ∂(ρku i ) ∂x i = ∂ ∂x j ν t σ k ∂k ∂x j + 2ν t E ij E ij -ρϵ ∂(ρϵ) ∂t + ∂(ρϵu i ) ∂x i = ∂ ∂x j ν t σ ϵ ∂ϵ ∂x j + C 1ϵ ϵ k 2ν t E ij E ij -C 2ϵ ρ ϵ 2 kwith the rate of deformation E ij , eddy viscosity ν t = ρC ν k 2 /ϵ, and adjustable constants C ν , C 1ϵ , C 2ϵ , σ ϵ , σ k . In our experiment section 4.1 solutions of simulations based on the RANS k-ϵ turbulence model are used as ground truth where the quantity of interest is the pressure field given the shape of a vehicle.Next, we introduce two popular algorithms that try to solve Eq. ( 4) numerically, where the main idea is to couple pressure and velocity computations. Specifically, we will discuss the SIMPLE (Semi-Implicit Method for Pressure-Linked Equations) [77] and PISO (Pressure Implicit with Splitting of Operators) [38] algorithm, which are implemented in OpenFOAM as well. The essence of these algorithms are the following four main steps:1. In a first step, the momentum equation is discretized by using suitable finite volume discretizations of the involved derivatives. This leads to a linear system for u for a given pressure gradient ∇p:EQUATIONHowever, this computed velocity field u does not yet satisfy the continuity equations ∇ • u = 0. 2. In a next step, let us denote by A the diagonal part of M and introduce H, so thatEQUATION)Taking into account Mu = -∇u allows to easily rearrange Eq. ( 23) for u, since A is easily invertible. This gives us:EQUATION)3. Now we substitute Eq. ( 24) into the continuity equation yieldingEQUATIONwhich is a Poisson equation for the pressure p that can be 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-corrector stage and this u now satisfies the continuity equation. Now, however, the pressure equation is 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 = -∇p (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 time-dependent flows, one usually considers the PISO algorithm (since in practice, for each timestep several thousands of iterations are usually required until convergence, which is computationally very expensive).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, 12] , and smoothed particle hydrodynamics (SPH) [29, 62, 69, 70] .The core idea behind SPH is to divide a fluid into a set of discrete moving elements referred to as particles. These particles interact through using truncated radial interpolation kernel functions with characteristic radius known as 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 i its position x i is given by its velocity u i such that ∂xi ∂t = u i . 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 ρ i of a particle can be expressed as ρ(x i ) = j ρ j V j W (||x i -x j ||, h) with W being a kernel, h its smoothing length, and V j the volumes of the respective particles. Rewriting the weakly-compressible Navier-Stokes equations with quantities as in (4) and additional Reynolds number Re and density ρ ∂u ∂t= 1 Re ∇ 2 u - 1 ρ ∇p + 1 ρ f ∂ρ ∂t = -ρ(∇ • u)in terms of these kernel interpolations leads to a 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 Eulerian 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 a mesh significantly simplifies the model implementation and opens up more possibilities for parallelization compared to Eulerian schemes [35, 23] . However, for accurately resolving particle dynamics on small scales a 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.A seemingly great fit for particle-based dynamics are graph neural networks (GNNs) [88, 42] with graph-based latent space representations. In many cases, predicted accelerations at the nodes are numerically integrated to model the time evolution of the many-particle systems [87, 65, 100, 102] .We provide a brief sketch of how universality can be established for our architecture. For transformerbased 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 nonlocality are crucial for universality. By demonstrating that the attention mechanism can, under 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.All experiments are conducted mostly on A100 GPUs. For large-scale experiments we use two research clusters equipped with either 8xA100-40GB nodes or 4xA100-64GB nodes. For smaller experiments and evaluations, we use a mix of internal servers equipped with varying numbers of A100-40GB or A100-80GB cards.We estimate the total number of GPU-hours (mostly 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 a standard pre-norm architecture as used in ViT [26] .We apply feature modulation to condition models to 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-Spectral" conditioning of [31] is 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.We use two datasets from [31] , the "Navier-Stokes 2D" and "ShallowWater-2D" to compare against baseline transformer architectures. Baseline results, data preprocessing and benchmark setups are taken from DPOT [34] and CViT [108] .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 as (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 [108] .We use a peak learning rate of 1e-4 with 10% warmup and a 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.We run comparisons against different Transformer baselines on a regular gridded Navier-Stokes equations dataset [31] in Table 2 (baseline results and evaluation protocol taken from DPOT [34] ). UPT outperforms all compared methods, some of which are specifically designed for regularly gridded data.As baseline transformers often train small models, we compare on a small scale, where UPT significantly outperforms other models. We also compare on larger scales, where UPT again outperforms competitors.Table 2 : Comparison on a regular grid Navier-Stokes dataset [31] . (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 Model # Params Rel. L2 Error FNO [51] 0.5M 9.12% FFNO [103] 1.3M 8.39% GK-T [18] 1.6M 9.52% GNOT [33] 1.8M 17.20% OFormer [53] 1.9M We run comparisons against UNet, FNO, Dilated ResNet variants on the regular gridded ShallowWater-2D climate modeling dataset [31] 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.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 of 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 . 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 is necessary. We use the transformer positional encoding [106] , 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 a learning rate of 1e-3/1e-3/5e-4 for grid resolutions 32 3 /48 3 /64 3 , respectively. We also tried using 32 Fourier modes for a grid resolution of 64 3 , 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 [115] →Conv3d→ReLU [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, we considered a higher initial hidden dimension for grid resolution 64 3 which performed worse.For GINO, we create a positional embedding of mesh-and grid-positions which are then used as input to a 3 layer MLP with GELU [37] non-linearities to create messages. Messages from mesh-positions within a radius of each grid-position are accumulated which serves as input to the FNO part of GINO. The FNO uses 64 hidden dimensions with 16/24/32 modes per dimension for grid resolutions 32 3 /48 3 /64 3 , respectively. The GINO decoder encodes the query positions and again uses a 3 layer MLP with GELU non-linearities to create messages. Each query position aggregates the messages from grid-points within a radius. The radius for message aggregation is set to 10. When using SDF features, the SDF features are encoded via a 2 layer MLP to the same dimension as the hidden dimension, and added onto the positional encoding of each grid point.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 [113] 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 we 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 a constant value of 1 and linearly interpolate onto a grid.Extended results. Table 4 gives a detailed overview of all results from ShapeNet-Car with varying resolutions and comparison to baseline models. 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 a short benchmark run with the largest possible batchsize take and extrapolate the time to 1 epoch. All benchmarks are conducted on a A100 80GB SXM GPU.Dataset. The dataset consists of 10K simulations which we split into 8K training simulations, 1K validation simulations and 1K test simulations. Each simulations has a length of 100 seconds where a ∆t of 1s is used to train neural surrogate models. Note that this ∆t is different from the ∆t used by the numerical solver, which requires a much lower ∆t to remain stable. Concretely, we set the temporal update of the numerical solver initially to ∆t = 0.05s. If instabilities occur, the trajectory is rerun with smaller ∆t. Overall, each trajectory comprises 2K timesteps at the coarsest ∆t setting and 200K at the finest (∆t = 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 to 0.06 m/s. Each simulation contains between 29K and 59K mesh points where each point has three features: pressure, and the xand y-component of the velocity. In total, the dataset is converted to float16 precision to save storage and amounts to roughly 235GB.From all training samples, data statistics are extracted to normalize inputs to approximately mean 0 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 is not supported to handle such a large amount of values, we assume that values follow a Gaussian distribution with mean and variance from the data, which allows us to infer the inter-quartile range from the Gaussian CDF.Additionally, we convert the normalized values ũt i,k to log scale to avoid extreme outliers which can lead to training instabilities:EQUATIONwhere all functions/operations are applied point-wise to the individual vector components. We apply the same log-scale conversion to the decoder output ût i,k . Model scaling. We scale models by scaling the dimension of the model. For grid-based 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 are as follows: 128 for UPT-8M, 192 for UPT-17M and 384 for UPT-68M.For UPT, we use 12 transformer blocks in total, which we distribute evenly across encoder, approximator 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 a 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 a 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 Unet mod 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 [0, 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 a graph out of mesh points, we use a radius graph with r = 5 and limit the maximum number of graph edges per node to 32 to avoid memory fluctuations. We choose r = 5 as it covers the whole domain when using a grid resolution of 64 2 .Training. We train all models for 100 epochs using a batchsize of 1024 where we use gradient accumulation if the full batchsize does not fit into memory. Following common practices, we use AdamW [41, 59] for U-Net, FNO and GINO. As we 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 use a learning rate of 5 × 10 -5 and 1 × 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 bfloat16 precision.Other hyperparameters are chosen based on common practices when training ViTs [26] . We linearly increase the learning rate for the first 10 epochs [58] followed by a 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 MLP are copied from ViTs.We find float32 precision to be beneficial for GINO and UPTs, which is likely due to positional embeddings being too inaccurate with bfloat16 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 work similarly well with AdamW but we kept Lion as we already trained models and re-training all models is quite expensive. Interpolated models. We train U-Net and FNO models as baseline on this task. To project the mesh onto a regular grid, we use a k-NN interpolation which takes the k nearest neighbors of a grid position and interpolates between the values of the nerarest neighbors. Going from the grid back to the mesh is implemented via the grid_sample 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.We show more qualitative results of the UPT-68M model in Fig. 13 .While one would ideally use lots of supernodes also during training, it also requires more computational 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 8 shows that UPTs are discretization convergent and inference performance can be improved by using more supernodes for inference than during training. 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 a 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.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. 9 show that UPTs scale well with larger latent spaces, allowing a flexible compute-performance tradeoff.Figure 5 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 to UPT (GINO-68M is worse than UPT-8M). 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 10 right, we compare UPT 1M/2M/8M against GINO 8M/17M/68M. UPT shows similar scaling on that loss scale. 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 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. : UPT is much more expressive than GINO and shows good scaling when increasing model size. UPT-1M achieves comparable performance to GINO-17M using 17x less parameters.We 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 11 right show that UPT scales well with data and is data efficient, achieving comparable results to GINO-8M with 4x less data.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 of different sizes and contains between 1 and 4 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 a new dataset containing 100 simulations and evaluate the model that was trained on the in-distribution dataset. The results in Figure 12 show that UPTs behave similar to other models when evaluated on OOD tasks where UPT outperforms all other models in all tasks. Therefore, the grid-indpendent architecture of UPTs does not impair OOD generalization performance. Generalization to more objects. Adding additional obstacles to a 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 not encode the number of 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 v ∈ [0.01, 0.06] m/s. The velocity is used (together with the timestep) as input to 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 a 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 different geometries. The in-distribution dataset contains randomly placed circles of varying size where the mesh is generated via an adaptive meshing algorithm. To investigate robustness to the meshing algorithm, we generate a OOD dataset with a uniform meshing algorithm. In this dataset, mesh cells are approximately uniform, i.e ṫhe distance between two points is roughly the same. This is 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, we investigate generalization to different obstacle geometries by using polygons (with up to 9 edges) or triangles instead of circles. Note that even though polygons are a more "complicated" shape, they are 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 1 in 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. 12 shows that UPT achieves the best performances on OOD meshes, despite having to adjust to a different input position distribution.Datasets. We use the Taylor-Green vortex datasets from Toshev et al. [101] 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, a sequence length of 126, and is characterized by a 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-based Simulator (GNS) model [87] is a popular learned surrogate for physical particle-based simulations. The architecture is kept simple, adhering to the encoder- processor-decoder principle [6] , where the processor consists of multiple graph network blocks [5] .The Steerable E(3)-equivariant Graph Neural Network (SEGNN) architecture [15] is a 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 Clebsch-Gordan tensor products interleaved with gated non-linearities [109] . SEGNN layers are message-passing 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 et al. [101] 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 l = 1, for more details see [101] .Hyperparameters. We train using a batchsize of 128 for 50 epochs with a warmup phase of 10 epochs. We sample inputs between 50% to 100% of total particles. For optimization, we use AdamW [59] for all experiments with learning rate 10 -3 and weight decay of 0.05.UPT architecture. For the TGV2D and TGV3D experiments, we use n s = 256 and n s = 512 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.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 to 12.6M parameters. Experiments on Taylor Green Vortex 2D.We show results on the TGV2D dataset from LagrangeBench [101] and the time to generate a simulated trajectory in Figure 14 . 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. 20We qualitatively study the scaling limits of representative methods by evaluating the memory consumption when training a 68M parameter model with a batchsize of 1 while scaling the number of inputs in the left part of Fig. 2 .We 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 a GNO (i.e. a 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 is a heavy restriction and these models would most likely need a larger degree to perform well when actually trained.• Transformer: 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 to 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] .• GINO: Aggregating information into a lower dimensional latent space allows models to be scaled to a 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 a very fine-grained grid to achieve comparable results to GINO or UPTs, which makes a fair comparison hard. Note that GINO usea a 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.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 Sec 4.2. For GINO-3D, we use the same grid resolution for the 3rd dimension as for GINO-2D (i.e. GINO-2D uses 64 2 and GINO-3D 64 3 ). 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 a quantitative comparison between completely different architectures is challenging, we additionally list the theoretical model complexities in Tab. 5. This study is meant to give a qualitative impression of the practical implications of the theoretical asymptotic complexities. (M D) ✓ ✗ ✗ ✗ CNN local O(G) ✗ ✗ ✗ ✗ Transformer global O(M 2 ) ✓ ✓ ✗ ✗ Linear Transformer global O(M ) ✓ ✓ ✗ ✗ GNO [52] radius O(M D) ✓ ✓ ✗ ✗ FNO [51] global O(G log G) ✗ ✓ ✗ ✗ GINO [54] global O(GD + G log G) ✓ ✓ ✓ ✗ UPT global O(SD + S 2 ) ✓ ✓ ✓ ✓Neural PDE surrogates play an important role in modeling many natural phenomena, many of which are related to computational fluid dynamics. Examples are weather modeling, aerodynamics, biological engineering, or plasma physics. Given the widespread application of computational fluid dynamics, obtaining shortcuts or alternatives for computationally expensive simulations, is essential for advancing scientific research, and has direct or indirect implications for reducing carbon emissions.In contrast to traditional computational fluid dynamics simulations, our approach aims to eliminate the need for complex meshing, bypass the requirement for detailed knowledge of boundary conditions, and overcome bottlenecks caused by solver intricacies. However, it is important to note that relying on simulations always necessitates thorough cross-checks and monitoring, especially when employing a "learning to simulate" methodology.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 particle 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 do not use any information (such as cell connectivity/adjacency) from the original mesh in the Eulerian case. 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, our design via radius graph is agnostic to the simulation type and also to randomly dropping input nodes.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 input nodes and avoid "squashing" too much information into a single supernode. We choose this edge limit to be 32 in all cases.We embed input and query positions via the sine-cosine position embedding from transformers [106] . Each position dimension is embedded separately. We rescale all positions to be in the range [0, 200] 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 a hyperparameter where a broad range of values works. If input nodes have additional input features, these features are first linearly projected into a hidden dimension and then the positional embedding vector is added.We also experimented with random fourier features [96] , but did not find a significant difference.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 the LayerNorm [2] layers before the attention and MLP parts of each transformer block. Additionally, a third scalar for gating the residual connections is learned. This methodology follows DiT [78] . For predicting these scalars, a simple learnable linear projection is used.  


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

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 1/ Clarity and detail in methodology:

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

2/ Scalability with extremely large datasets:

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

3/ Insufficient analysis of generalization capabilities:

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

4/ Potential overfitting concerns:

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

5/ Analysis from particle and grid based methods:

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

A/ Lack of detailed comparison with established methods:

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

B/ Handling of boundary conditions:

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

C/ Adaptation to high-resolution grids and large particle systems:

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

D/ Numerical stability and accuracy:

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

E/ Computational efficiency in complex geometries:

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

F/ Interoperability with existing simulation tools:

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

G/  Handling of multiscale phenomena:

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


6/ Model conditioning standpoint:

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

A/ Insufficient detail on conditioning mechanisms:

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

B/ Limited analysis of conditioning performance:

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

C/ Scalability of conditioning methods:

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

D/ Generalization to unseen conditions:

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

E/  Comparison with other conditioning approaches:

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

F/  Impact of conditioning on training stability:

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

G/ Real-world application scenarios:

While 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.
questions: 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.

How does the UPT framework handle complex boundary conditions compared to traditional particle and grid-based methods?

2/  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.

What measures are in place to ensure the numerical stability and accuracy of UPTs over long simulation periods?

3/ 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.

 What are the scalability and generalization capabilities of UPTs when applied to extremely high-resolution grids, large particle systems, and unseen physical scenarios?
limitations: limitations: Yes.
paper_summary: paper_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.
strength_weakness: strength_weakness: -	The 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. 
-	The 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.
-	There 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. 
	Overall, 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. 
1.	Haixu Wu, Tengge Hu, Huakun Luo, Jianmin Wang, Mingsheng Long, Solving High-Dimensional PDEs with Latent Spectral Models. https://arxiv.org/abs/2301.12664
2.	Tailin Wu, Takashi Maruyama, Jure Leskovec, Learning to Accelerate Partial Differential Equations via Latent Global Evolution, https://arxiv.org/abs/2206.07681
3.	Zijie Li, Dule Shu, Amir Barati Farimani, Scalable Transformer for PDE Surrogate Modeling, https://arxiv.org/abs/2305.17560
4.	Zijie Li, Kazem Meidani, Amir Barati Farimani, Transformer for Partial Differential Equations' Operator Learning, https://arxiv.org/abs/2205.13671
5.	Zhongkai 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
6.	Michael 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
questions: questions: -	I’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?
-	What 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: limitations: -	For 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. 
-	Fixing 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: 1. A bit more analysis and information could be included to improve completeness. See Questions.
2. 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,
questions: 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)
2. 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.
3. For a more comprehensive display of results, could the authors include accuracy metrics (e.g., mean error) and memory benefits to  Table 1?
4. 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: Limitations are well-discussed in Appendix A.
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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.
- 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. 
- 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.

[1] GNOT: A General Neural Operator Transformer for Operator Learning

[2] Transolver: A Fast Transformer Solver for PDEs on General Geometries

[3] Transformer for Partial Differential Equations' Operator Learning
questions: 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?
- How dose DiT modulation influences UPTs' performance? Additionally, which features have you found to be most effective as conditions for DiT modulation? 
- 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?
- 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? 
- For each optimization step, are the supernodes fixed or do they change dynamically? 
- 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: limitations: Yes
paper_summary: paper_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.
strength_weakness: strength_weakness: - The framework is labeled as universal but is only tested on fluid dynamics, lacking experiments on other PDEs.
- 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)
- The scalability of UPT concerning training size remains unclear, which is an important property.
- The stability of the latent rollout technique for long rollouts is uncertain, raising concerns about its general reliability
questions: 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?
- Why do you believe that latent rollouts do not require stabilization?
limitations: limitations: The authors explained well the limitations.

Comments: title: Thank you for the follow-up
Pseudo-code for training UPT
Reply to follow-up questions
Missing pseudo-code and extra details on model conditioning
Follow-up questions
Suggestions
Summary of the revision and follow-up
Reply to rebuttal
comment: Thanks for your detailed reply, most of my concerns have been addressed and I've raised the score. Best wishes!
Please add the discussion in the revised version.
Due to space constraint in the rebuttal, we were not able to include it there (the rebuttal did also not allow to update the submitted paper). Please find below pseudocode for UPT in the setting of our transient flow experiments (2D positions with 3 features per node).

```
# input_embed: linear projection, projecting from the number of input features to a hidden dimension
# pos_embed: positional embedding (with sine and cosine waves of different frequencies) as common in transformers
# message_mlp: shallow MLP to create messages
# encoder_transformer: stack of transformer blocks of the encoder
# latent_queries: `n_latent_tokens` learnable query vectors
# encoder_perceiver: cross attention block
# approximator_transformer: stack of transformer blocks of the approximator
# query_mlp: shallow MLP in the decoder to encode query positions
# decoder: cross attention block

def encoder(input_features, input_pos, radius, n_supernodes):
    """
    encode arbitrary pointclouds into a fixed latent space
    inputs:
        input_features `Tensor(n_input_nodes, 3)`: features of input nodes 
        input_pos `Tensor(n_input_nodes, 2)`: positions of input nodes 
        n_supernodes `integer`: number of supernodes 
        radius `float`: radius for creating the radius_graph
    outputs:
        latent `Tensor(n_latent_tokens, hidden_dim)`: encoded latent space
    """
    # create radius graph (using all input nodes)
    # edges are uni-directional and are passed from nodes_from to nodes_to
    nodes_from, nodes_to = radius_graph(input_pos, radius)

    # select supernodes from input_nodes
    n_input_nodes = len(input_features)
    supernode_idx = randperm(n_input_nodes)[:n_supernodes]

    # filter out edges that do not involve supernodes
    is_supernode_edge = nodes_to in supernode_idx
    nodes_from = nodes_from[is_supernode_edge]
    nodes_to = nodes_to[is_supernode_edge]

    # encode inputs and positions
    encoded_nodes = input_embed(input_features) + pos_embed(input_pos)

    # create messages
    messages = message_mlp(encoded_nodes[nodes_from])
    # accumulate messages per supernode by averaging messages
    supernodes = accumulate_messages(messages, nodes_to, reduce="mean")

    # process supernodes with some transformer blocks
    supernodes = encoder_transformer(supernodes)
    # perceiver pooling from supernodes to latent tokens
    latent = encoder_perceiver(query=latent_queryies, key=supernodes, value=supernodes)
    return latent

def approximator(latent):
    """
    propagates latent forward in time
    inputs:
        latent_t `Tensor(n_latent_tokens, hidden_dim)`: encoded latent space at timestep t
    outputs:
        latent_t_plus_1 `Tensor(n_latent_tokens, hidden_dim)`: encoded latent space at timestep t + 1
    """
    return approximator_transformer(latent)

def decoder(latent, query_pos):
    """
    decode latent space pointwise at arbitrary positions
    inputs:
        latent `Tensor(n_latent_tokens, hidden_dim)`: 
        query_pos `Tensor(n_outputs, 2)`: positions for querying the latent space
    outputs:
        decoded `Tensor(n_outputs, 3)`: evaluation of the latent space at query positions
    """
    # encode query positions
    query_pos_embed = query_mlp(pos_embed(query_pos))
    # query latent space
    decoded = decoder(query=query_pos_embed, key=latent, value=latent)
    return decoded


def train_step(input_features, input_pos, n_supernodes, radius, query_pos):
    """
    encode arbitrary pointclouds into a fixed latent space
    inputs:
        input_features `Tensor(n_input_nodes, 3)`: features of input nodes at timestep t
        input_pos `Tensor(n_input_nodes, 2)`: positions of input nodes at timestep t
        n_supernodes `integer`: number of supernodes 
        radius `float`: radius for creating the radius_graph
        query_pos `Tensor(n_outputs, 2)`: positions for querying the latent space
        target_features `Tensor(n_input_nodes, 3)`: features of nodes at timestep t + 1
    outputs:
        loss `Tensor`: skalar loss value
    """
    # next-step prediction
    latent_t = encoder(input_features, input_pos, radius, n_supernodes)
    latent_tplus1 = approximator(latent_t)
    decoded_tplus1 = decoder(latent_tplus1, query_pos)
    next_step_loss = mse(decoded_tplus1, target_features)

    # inverse decoder (decode latent into inputs)
    decoded_t = decoder(latent_t, input_pos)
    inverse_decoding_loss = mse(decoded_t, input_features)

    # inverse encoder (encode predictions at t + 1 into latent of t + 1)
    inverse_encoded = encoder(decoded_tplus1, query_pos, radius, n_supernodes)
    inverse_encoding_loss = mse(inverse_encoded, latent_tplus1)

    return next_step_loss + inverse_decoding_loss + inverse_encoding_loss
```
We updated the paper with additional details concerning modulation which would be included in the camera-ready version, but the rebuttal unfortunately did not allow updating the paper.

**1. Details on modulation**

We use the DiT modulation as it was introduced in the original paper [1] and condition onto boundary conditions (timestep, inflow velocity) in the following way:

```
# embed: embedding layer (with sine and cosine waves of different frequencies) as also common to encode positions in transformers
# condition_mlp: shallow MLP to combine boundary conditions
def create_condition(timestep, velocity):
    timestep_embed = embed(timestep)
    velocity_embed = embed(velocity)
    embed = concat([timestep_embed, velocity_embed])
    condition = condition_mlp(embed)
    return condition
```

From this condition, DiT modulates its features by (i) shifting the features after layer normalization (ii) scaling the features after layer normalization or (iii) gating the features before adding the residual.
This is done in the same way as DiT does it. We provide a link to the original implementation [2] but are also happy to provide pseudocode if needed.


This form of conditioning allows us to encode the inflow velocity and the current timestep into the model. Therefore, the model can better adapt to different inflow velocities as this information can be explicitly encoded into the model. Without encoding such boundary conditions into the model, it would need to derive it from the mesh points or particles alone, which is sometimes impossible.
Note that we also apply modulation to UNets in the form of FiLM [3] (similar to DiT it scales and shifts intermediate features) and for FNO/GINO we apply the "spatial-spectral parameter conditioning" from [4] (see Appendix B.2.6) which applies modulation in the fourier space of FNO.


[1] Peebles et al., "Scalable Diffusion Models with Transformers", ICCV 2023, https://arxiv.org/abs/2212.09748

[2] https://github.com/facebookresearch/DiT/blob/main/models.py#L101

[3] Perez et al., "FiLM: Visual Reasoning with a General Conditioning Layer" AAAI 2018, https://arxiv.org/abs/1709.07871

[4] Gupta et al., "Towards Multi-spatiotemporal-scale Generalized PDE Modeling", arXiv 2022, https://arxiv.org/abs/2209.15616


**2. Boundary conditions in ShapeNet-Car**

In the ShapeNet-Car experiments, we showcase the flexible encoding of UPT by additionally encoding the signed distance function (SDF) of the geometry.
Given SDF evaluated on a 3D grid, we encode it with a shallow CNN with 3D convolutions. For example, we use a 64^3 grid of SDF values and encode it into a 8^3 latent representation with said CNN.
This grid is then flattened out into a 1D sequence of 8^3=512 tokens. Afterwards these 512 tokens are concatenated to the output of the encoder (the 1024 latent tokens from encoding the mesh of the car) resulting in 1536 tokens as input to the approximator.

**3. Numerical stability of UPT**

Traditional solvers resolve the physics in very small timesteps as this is required to preserve the physical properties of the simulation. For example, if two particles are about to collide but the time resolution is not small enough to accurately detect when the two particles collide (and therefore change direction/velocity), the two particles will overlap with each other in the next simulation timestep. If this overlap is large, the resulting force acting on these particles will be extremely high, which in turn results in numerical instability and a collapse of the simulation.
Contrary, UPTs do not learn to explicitly model the interactions between two particles, but rather an abstract representation thereof. Therefore, UPT does not need to resolve at a fine-grained time resolution as it can learn that "if two particles are about to collide, they are about to change direction". This makes it much more resilient against numerical stability as it focuses on abstract concepts instead of low-level interactions.


Other neural operators like GNNs or graph neural operators (GNOs) also resolve particle-particle interactions and therefore require smaller timesteps than UPT. We showcase this in Figure 6 of our paper, where the UPT modeling paradigm (via an abstract compressed latent space) can resolve much larger timesteps without problems. We show this in our experiments on TGV3D where UPT is trained on 10x coarser timesteps than GNS or SEGNN.


**4. Scaling properties of UPTs**

As UPTs consist of mainly transformer layers, it does not have any issues with scaling. The same techniques that are applied to train large language models (LLMs) can be applied to UPTs as well. As an example, Llama3-405B [1] was scaled to 16K H100 GPUs. 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.

[1] Dubey et al., "The Llama 3 Herd of Models", arXiv 2024, https://arxiv.org/abs/2407.21783
I am unable to find the pseudo-code and the additional details on model conditioning that are mentioned in your response.
1. In your response, you mentioned enhancing the model conditioning with techniques such as feature modulation, but I could not find the additional details in the paper. Could you elaborate on how the DiT modulation and any other conditioning mechanisms are implemented in the UPT framework, and how these techniques specifically improve model adaptability to varying boundary conditions and simulation inputs?
2. You highlighted the use of transformers to manage complex boundary conditions, particularly in experiments like the ShapeNet car. Could you provide more details on how the transformers are configured to encode these boundary conditions effectively? Are there specific architectural adjustments or tokenization strategies employed to ensure accurate representation and processing of boundary conditions?
3. In the context of your framework's numerical stability, you mentioned that UPT rollouts are stable without specific measures. Could you clarify how the UPT framework maintains stability and accuracy over extended simulations, especially when compared to traditional numerical solvers that require small timesteps for stability? Are there specific aspects of the transformer architecture that contribute to this stability?
4. While you provided insights into UPT's scalability with large datasets and distributed training, can you discuss any potential limitations or bottlenecks encountered when scaling the model further? How does the framework handle challenges such as increased computational demands or memory usage as the complexity of the simulations and datasets grows?
1. While real-world applications are beyond the current scope, discussing potential use cases or hypothetical scenarios could highlight the practical impact of UPTs.
2. Consider exploring how UPTs might be integrated with traditional particle or grid-based methods in practice. Discussing potential hybrid approaches could broaden the applicability of your work
3. Discussing specific measures or techniques to enhance long-term stability and robustness in simulations, particularly with noisy data, could strengthen the paper's claims.
4. As hardware capabilities evolve, exploring how UPTs can leverage advancements in GPU and TPU technologies could provide insights into the future potential of your framework.
Thank you for addressing the feedback provided in my initial review. Here is a summary of the revisions and responses you have made, along with some additional feedback and observations.

**Scalability to extremely large datasets / high-resolution grids**

I commend your efforts to demonstrate scalability within the constraints of your computational resources. The additional experiments on the Shallow Water and Navier-Stokes datasets effectively illustrate your framework's ability to handle large datasets. Your discussion on using distributed training across multiple GPUs highlights the potential of your approach in handling complex simulations.

**Generalization capabilities**

The expanded tests on generalization, including across different flow regimes and discretizations, strengthen the evidence for your framework's robustness. This comprehensive analysis is a significant enhancement and effectively addresses previous concerns about the adaptability of your model.

**Handling of boundary conditions**

The examples provided on handling complex boundary conditions through experiments like the ShapeNet car and transient flow simulations highlight the flexibility and adaptability of UPTs in various scenarios. This aspect of your work is well-explained and a notable strength of your framework.
Thank you for your answers. I increased the rate to 7. Good luck!
Thank you for the clarifications, and the additional experiments in general response which I believe strengthen the paper. I have raised my score. However, it would be better if the supernode pooling/radius graph construction could be included either in the maintext or supplementary for reference.
Questions have been addressed rigorously. Well done. Changing score from 6 -> 7.