Universal Physics Transformers: A Framework For Efficiently Scaling Neural Operators

We appreciate your review and respond to the raised concerns below.

Experiments with other PDEs

UPTs are neural operators known to be applicable across various PDE types. Due to the nonlinear nature of the NS equations, they are notoriously more challenging to solve than parabolic or hyperbolic PDEs, like the heat or wave equations. Therefore, we hypothesize that UPTs should also generalize well to different PDEs.

Additionally, as requested by reviewer JKEE, we added comparisons of UPT with other methods on small-scale regular grid datasets. In total, UPT outperforms competitors on 5 diverse datasets that span over different dataset sizes (900 to 1M frames), resolutions (2K to 60K inputs), spatial dimensions (2D, 3D), simulation types (steady state, transient), different boundary conditions and specifications (Eulerian, Lagrangian).

Scalability with model size

We see now that Figure 5 could suggest that UPT does not scale well with parameter size. However, the scaling of GINO only looks good because GINO underperforms in contrast to UPT (GINO-68M is worse than UPT-8M). As you correctly identified, for well trained (UPT) models it gets increasingly difficult to improve the loss further. Ideally, one would show this by training larger GINO models, however this is prohibitively expensive (68M parameter models already take 450 A100 hours per model). We therefore go in the other directions and train even smaller UPT models that achieve a similar loss to the GINO models and compare scaling there. In Figure 2 of the supplemental pdf, we compare UPT 1M/2M/8M against GINO 8M/17M/68M. UPT shows similar scaling on that loss scale.

We see how this could be easily misinterpreted from Figure 5 and adjust the paper accordingly to remove this misunderstanding.

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.

Scalability with dataset size

We added experiments to show the scalability and data efficiency of UPTs by training UPT-8M on subsets of the data used for the transient flow experiments (we train on 2K and 4K out of the 8K train simulations). The results in Figure 3 of the supplementary rebuttal pdf show that UPT scales well with data and is data efficient, achieving comparable results to GINO-8M with 4x less data.

We also show that UPTs can handle various dataset sizes. ShapeNet-Car is a small-scale dataset consisting of 889 car shapes with 3.6K mesh points each. TGV3D is a bit bigger with 8K particles per simulation and 200 simulations of length 61 (12K frames). The dataset for our transient flow experiments contains around 50K mesh cells per simulation with 10K simulations of length 100 (1M frames). UPT shows strong performances across all considered dataset sizes.

Additionally, UPT consists of mostly transformer blocks, which have demonstrated outstanding scalability in other domains such as language modeling [1] and computer vision [2].

[1] Kaplan et al., "Scaling laws for neural language models", arXiv 2020, https://arxiv.org/abs/2001.08361

[2] Zhai et al., "Scaling vision transformers", CVPR 2022, https://arxiv.org/abs/2106.04560

Stability of latent rollout

We found UPT to be fairly stable without any special techniques to stabilize the rollout (e.g. [3]). However, such methods could further improve UPTs performance but these methods are not specific to UPT and typically require additional computations during training. Therefore, we leave exploration of this combination to future work.

Additionally, the latent rollout opens new avenues to potentially apply existing stabilization tricks with less compute. For example, one could do the forward propagation for the stabilization technique from [4] in the latent space, which would greatly reduce training costs thereof. However, as these tricks can be tricky to train (e.g. due to requiring a precise trade-off between training the next-step prediction vs n-step prediction), we leave this direction to future work.

[3] Lippe et al., "Pde-refiner: Achieving accurate long rollouts with neural pde solvers", NeurIPS 2023, https://arxiv.org/abs/2308.05732

[4] Brandstetter et al., "Message passing neural pde solvers", ICLR 2022, https://arxiv.org/abs/2209.15616

Thank you for your review and helpful comments which were very useful to improve the paper. We address your points individually.

Comparison to other transformer methods

The fundamental building principles of UPT are (i) an encoding that is designed to handle irregular grids of various sizes, (ii) a compressed and unified latent space representation that allows for efficient forward propagation in the latent space, and (iii) a neural field-type decoder. We consider all of these points as fundamental for scaling neural operators.

GNOT processes each mesh point as a token and therefore quadratically scales with the number of mesh points. Furthermore the output can only be evaluated at the positions of the input.

Transolver's concept of "Physics-aware Tokens" is somewhat similar to UPTs supernodes. However, only the substitutes attention part operates on a reduced space, whereas the FFN part of the transformer operates on the uncompressed space. Therefore, it is a type of linear complexity transformer, which quickly becomes infeasible for larger input sizes (see Figure 2). Similarly, OFormer also operates on the uncompressed input.

In contrast, UPT heavily compresses the input, leading to sub-linear complexity w.r.t. input size in all transformer layers.

Additionally, we added comparisons to transformer baselines on regular grid datasets (see general response) where UPT outperforms e.g. OFormer by quite a big margin.

Benefits of the latent rollout

While the latent rollout does not provide a significant performance improvement, it is almost an order of magnitude faster. The UPT framework allows to trade-off training compute vs inference compute. If inference time is crucial for a given application, one can train UPT with the inverse encoding and decoding objectives, requiring more training compute but greatly speeding up inference. If inference time is not important, one can simply train UPT without the reconstruction objectives to reduce training costs.

Additionally, the latent rollout enables applicability to Lagrangian simulations. As UPT models the underlying field instead of tracking individual particle positions it does not have access to particle locations at inference time. Therefore, autoregressive rollouts are impossible since the encoder requires particle positions as input. Using the latent rollout, it is sufficient to know the initial particle positions as dynamics can be propagated without any spatial positions. After propagating the latent space forward in time, one can simply query the latent space at arbitrary positions to evaluate the underlying field at given positions. We showcase this in Figure 7 where the latent space is queried with regular grid coordinates (white arrows).

We discuss a potential improvement for the latent rollout to make it more efficient in Appendix A. While we show in the paper that a latent rollout can be enabled via a simple end-to-end training, we think that it can be improved, e.g. via a two stage procedure of first training encoder/decoder in an autoencoder setting, followed by freezing encoder/decoder and training the approximator on the fixed pre-trained latent space akin to [1]. Such an approach would not require inverse encoding/decoding objectives as the separation of components is enforced through the multi-stage training.

[1] Rombach et al., "High-resolution image synthesis with latent diffusion models", CVPR 2022, https://arxiv.org/abs/2112.10752

DiT modulation

Conditioning onto external features is crucial to encode this external information into the model. We condition onto the current timestep and also onto the inflow velocity in the transient flow experiments. Note that this is not specific to UPT and we also apply conditioning to compared models.

Supernodes and radius graph creation

We realize that our description in the paper is a bit misleading. Supernodes are, in the case of Eulerian data, sampled according to the underlying mesh and, in the case of Lagrangian data, sampled according to the particle density. A random sampling procedure which follows the mesh or particle density, respectively, allows us to put domain knowledge into the architecture. (We change "randomly sampled" to "sampled according to the underlying mesh / underlying particle density".) 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 thus having a regularization effect.

The radius graph encodes a fixed region around each supernode. While one could use the original edge connections for Eulerian data, Lagrangian data does not have edges. Additionally, we employ randomly dropping input nodes as a form of data augmentation which makes using the original edge connections more complicated from an implementation standpoint. In contrast, a radius graph is agnostic to randomly dropping input nodes.

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 16, i.e. on average each supernode represents 16 inputs. We found our model to be fairly robust to the radius choice. Also, the circles of different supernodes can overlap. Therefore, the encoding of dense regions can be distributed among multiple supernodes.

We impose the edge limit of the radius graph by randomly dropping connections to preserve the average distance from supernode to input nodes.

Additionally, the radius graph facilitates Lagrangian settings where there is no predefined connectivity between particles.

We extended discussion in the corresponding sections and added a guide on how to choose these hyperparameters to the paper.

Thank you for your comments which helped to improve the paper. We addressed all your comments and followed all your suggestions.

Memory and time complexity of different architectures

We provide theoretical complexities in Appendix C.6 in Table 3. Note that when scaling only the input size towards infinity to calculate an asymptotic runtime/memory complexity, the runtime complexity is mostly the same as the memory complexity as storing intermediate activation dominates the other factors. For example, transformers scale quadratically with the number of inputs M. For large M, a transformer needs to calculate the MxM attention matrix and also store it in memory. Therefore, it has $O(M^2)$ runtime and memory complexity. For simplicity, we only consider the theoretical case and do not take optimizations such as hardware-aware implementations (e.g. [1]) into account that can trade-off additional computations for a reduced memory footprint.

As the theoretical complexities introduce many variables due to different architectures processing the input in vastly different ways, we believe that Figure 2 presents the practical complexities in a cleaner way and provide the theoretical complexities later in the appendix.

[1] Dao et al., "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness", NeurIPS 2022, https://arxiv.org/abs/2205.14135

Positional encoding type

We employ the same positional embedding approach used in the original transformer paper [2], applying it separately to each dimension, as is also common in vision transformers [3]. Namely, we employ a combination of sine and cosine functions with different frequencies to encode each position. The revised version of the paper offers a clearer explanation of this positional embedding.

We also experimented with different types of positional embeddings such as the Fourier feature mapping from [4], which uses randomly sampled frequencies from a Gaussian distribution. In an experiment on TGV2D of our Lagrangian experiments it shows a slight improvement (see Figure 1 of the supplemental rebuttal pdf). The results show a slight improvement for longer rollouts, which we hypothesize is because this method doesn't overly emphasize features aligned with the axes. However, as improvements of other positional embeddings were minor, we stuck to the transformer positional embedding for simplicity.

[2] Vaswani et al., "Attention is all you need", NeurIPS 2017, https://arxiv.org/abs/1706.03762

[3] Dosovitskiy et al., "An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale", ICLR 2021, https://arxiv.org/abs/2010.11929

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

Additional accuracy metrics and memory benefits of Table 1

We included memory benefits in Table 3 of Appendix C.6. However, we find it difficult to assign a representative accuracy metric to each architecture as most methods are specialized for a certain type of task. For example, the accuracy of a CNN on a task with irregular data will naturally be worse than the accuracy of a GNN on that problem. UPT shows strong performances across various input domains and input scales, but as other methods are not as flexible we find it difficult to quantify their performance with a single metric.

Application in other domains

Our method could also be applied to other fields, such as video modeling. In that context, it would be possible to train on a dataset with varying resolutions. Since the decoder's output is a point-wise neural field, it could generate videos at arbitrarily high resolutions. However, in the context of language modeling, this approach is challenging to apply because language data is finite-dimensional.

Thank you for your thorough review and helpful comments which were very useful to improve the paper. We address your points individually.

Why UPT is a better scalable latent space model

Latent Spectral Model uses geo-FNO (which could be consider as predecessor of GINO) to handle irregular grids. Therefore, it is somewhat similar to GINO as it has to map its input into a fixed regular grid representation. UPT doesn't map to a regular grid which makes it more flexible/expressive.

LE-PDE only considers small-scale regular grid data and their design requires train and test data to have the exact same resolution (as mentioned in Appendix C of their paper). Therefore, its application to irregular grid data is impractical as a lot of data has different number of particles/mesh cells during test time. UPT can handle arbitrary input resolutions during training and generalizes well to higher input resolutions (as shown in Figure 5).

FactFormer does not compress its input into a smaller latent space and therefore doesn't scale to large input sizes. Their method proposes an efficient attention mechanism, but their compute requirement would be similar to linear transformers and therefore become infeasible on large systems (as shown in Figure 2).

Comparison to transformer based models on regular grid

We compare UPT on two regular grid benchmarks as shown in the general response and in the supplemental pdf Tables 1 to 3. UPT outperforms all compared methods (such as OFormer or DPOT) -- often by quite a margin -- without being specifically designed for regular grid datasets.

Justification why UPTS are neural operators

Thank you for bringing up this important point. To avoid overloading the notation, we will provide a brief sketch of how universality can be established for our architecture. For transformer-based neural operators, universal approximation has been recently demonstrated in [1], Section 5. The arguments in this work are heavily based on [2], which establishes that nonlinearity and nonlocality are crucial for universality. By demonstrating that the attention mechanism can, under appropriate weight choices, function as an averaging operator, the results from [2] are directly applicable. For detailed proof, refer to Theorem 22 in [1]. Our algorithm fits within this framework as well: we employ nonlinear, attention-based encoders and decoders (as allowed by the results of [2], Section 2) and utilize attention layers in the latent space.

Please let us know if you require further explanations. We will provide detailed information in an updated version of the manuscript.

[1] Calvello et al., "Continuum Attention for Neural Operators", arXiv 2024, https://arxiv.org/abs/2406.06486

[2] Lanthaler et al., "The nonlocal neural operator: Universal approximation" arXiv 2024, https://arxiv.org/abs/2304.13221

Rollout stability

Our training procedure doesn't use any special methods to stabilize the rollout because this typically comes at increased training costs. However, one could easily apply such techniques to UPT. As we found our models to produce fairly stable rollouts even without additional techniques, we leave this direction for future work.

Thoughts on latent space vs physics space

We see the latent rollout as a promising way to speedup inference, particularly for large-scale systems where encoding and decoding takes up most of the runtime. UPT offers the flexibility to invest resources into training it as a latent rollout model (by using the inverse encoding/decoding objectives) or to save resources during training (by omitting the reconstruction objectives) at the cost of increased inference time. We consider this a valuable tool to have for neural operators.

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. Using the latent rollout, it is sufficient to encode the initial particle positions into the latent space, which can then be propagated without the knowledge of any spatial positions. After propagating the latent space forward in time, one can simply query the latent space at arbitrary positions to evaluate the underlying field at given positions. We showcase this in Figure 7 where the latent space is queried with regular grid coordinates (white arrows).

Limitations of non-regular gridded data

We are interested in complex physics simulations which are often simulated using e.g., finite element meshes. Several phenomena are modeled by particle-based simulations such as smoothed particle hydrodynamics, material point methods, or discrete element methods. Many phenomena even require a coupling of different aforementioned simulation types. This is very much driven by daily life engineering applications, and as such was a main motivation for developing UPT.

Limitations on fixed latent space

In the current architecture of UPT, we consider a fixed latent space as it has proven to be an efficient way to compress the input into a fixed size representation to enable scaling to large-scale systems while remaining compute efficient. However, if an application requires a variable sized latent space, one could also remove the perceiver pooling layer in the encoder. With this change the number of supernodes is equal to the number of latent tokens and complex problems could be tackled by a larger supernode count. While we currently do not consider a setting where this is necessary, we show that the performance of UPT steadily increases with the number of supernodes during training (Figure 9 in Appendix C.4.3). Additionally, UPT's performance improves when more supernodes are used during evaluation than were utilized during training (Figure 8 in Appendix C.4.2).

Thank you for your profound review and suggestions that helped us to improve our paper a lot. We addressed all your comments and followed all your suggestions. Please let us expand a bit on your comments.

Clarity and detail in methodology

In order to make the paper more accessible we -- as suggested -- add pseudocode of a UPT forward pass with adaptive sampling. This together with sketch Figure 3a, sketch Figure 3b, and the encoder/approximator/decoder description of Section 3 should make the paper understandable wrt. all necessary details.

Scalability to extremely large datasets / high-resolution grids

We agree that scalability to extremely large datasets is a desired property. In the paper, we have done our best to test this. Our models are trained using distributed training up to 32 A100 GPUs. We have self-simulated a dedicated dataset which is larger (> 50000 mesh points), and more complex (different flow regimes, different number of differently placed obstacles) as standard grid based datasets. As requested by several reviewers we have added experiments on the Shallow Water dataset and the Navier-Stokes dataset to the paper, which demonstrate the diverse applicability of UPT. However, even larger simulations than our transient flow simulations are beyond the scope of this paper and the available compute budget. Nevertheless, we have tested GPU usage of several model classes (Figure 2), which demonstrates that our framework is able to process meshes / particle clouds with up to 4 million nodes on a single GPU.

Insufficient analysis of generalization capabilities

We agree that generalization capabilities are one of the most important aspects to test for neural operators. Therefore, in the transient flow experiment -- our largest experiment -- we test generalization across number of input / output points (discretization convergence), discretization across different flow regimes (input velocity), and generalization across different scenarios (differently placed obstacles).

Potential overfitting concerns

We observe that the variable selection of the supernodes and the variable selection of input/output nodes allows for a strong regularization and data augmentation. The general training pipeline of UPT is therefore very robust against overfitting. This is further strengthened by the additional Navier-Stokes and Shallow water experiments.

Lack of detailed comparison with established numerical methods

We are not claiming to be better than numerical methods. We compare with neural methods and especially focus on scaling. We follow the general procedures of the community, and report MSE, correlation time measures, and runtime comparisons. The latter two give an idea of the potential of UPT when compared with numerical methods. Similar to e.g. weather modeling (see Aurora [2], Pangu [3], ...) we follow the belief that benefits of neural operators will become more pronounced at scale.

Handling of boundary conditions

The ShapeNet car experiment allows for testing neural operators on complicated geometries. UPT performs favorably, even when using a strongly reduced latent space representation. Additionally, in the transient flow example we have varying numbers of differently placed obstacles which can also be seen as complicated in-domain boundaries. In fact, UPT is specifically designed for large data on non-regular domains including different boundary conditions.

Model conditioning

We have added a more detailed paragraph on the model conditioning choices in this paper. In our paper we follow closely the detailed studies described in [1].

Real world applications

Real world applications are beyond the scope of this work, but a very strong motivation for developing our framework. We note that current weather modeling such as Aurora [2] or Pangu [3] follow a similar design principle, but for regular gridded data, and without discretization convergence properties.

Question on complex boundary conditions

Transformers offer a flexible way to encode various boundary conditions. Skalar boundary conditions, such as inflow velocity or the current timestep, can be encoded via feature modulations. We use DiT modulation as it has shown strong performances in transformers.

Additionally, transformers offer a flexible mechanism to encode additional information by encoding the information into tokens and concatenating them to the supernodes or latent tokens. We make use of this flexible encoding in the ShapeNet-Car experiments where we additionally encode the signed distance function evaluated on a 3D grid via a CNN and feed the resulting tokens to the approximator.

Question on numerical stabilitiy and accuracy

Traditional methods require extremely small timesteps in order to preserve numerical stability when resolving the physics. As neural operators only approximate the solution of the traditional solver, they can operate on much larger timesteps as their powerful modeling capabilities enables them to learn dynamics also from a coarse time resolution.

Furthermore, UPT rollouts are fairly stable without any specific measures to ensure rollout stability. Sophisticated techniques to improve rollout stability (e.g., [4]) could be easily applied to UPTs to further enhance rollout stability. However, as these techniques typically impose a runtime overhead, we leave exploration thereof to future work.

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

[2] Bodnar et al., "Aurora: A foundation model of the atmosphere", arXiv 2024, https://arxiv.org/abs/2405.13063

[3] Bi et al., "Accurate medium-range global weather forecasting with 3D neural networks", Nature 2023, https://www.nature.com/articles/s41586-023-06185-3

[4] Brandstetter et al., "Message passing neural pde solvers", ICLR 2022, https://arxiv.org/abs/2209.15616

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