Thank you for the review and feedback. We want to address your remaining concerns in the following:

Stronger Baselines While there can certainly be more baselines, we politely disagree that the baselines chosen are "relatively weak". For example, we compare extensively against scalable operator transformer scOT [1a] and UDiTs [2b]. They are both very strong and recent baselines (both published NeurIPS) for scientific machine learning and a SOTA diffusion transformer architectures.

Thank you for mentioning the paper [3c], FactFormer, which is also a possible baseline, so we include it in our comparison of transformer architectures for a more comprehensive evaluation. See the following table for a comparison between FactFormer and PDE-Transformer (extending table 1 in the main paper):

PDE-S clearly outperforms FactFormer. Note that FactFormer has fewer parameters, since we were using the implementation by the authors of [3c] (from github). As a weight-computation tradeoff, we preserved the original architecture with fewer weights, but gave FactFormer a significantly larger computational budget. Hence, despite fewer parameters than PDE-S, FactFormer trains much longer and requires more than three times as many floating point operations. We believe this is fair comparison for FactFormer. The resulting performance of FactFormer is significantly lower despite the additional operations.

For the PDEformer model [4d], we will include it in the related work. PDEformer only targets 1D PDEs and an extension to 2D PDEs seems nontrivial. Moreover, it constructs a graph using the target PDE, which is not always available in the general setup we are targeting. Similarly, Unisolver [5e] is mostly orthogonal to our work: it primarily focuses on conditioning on PDE specific information (PDE equation, boundaries, etc.) using language embeddings from LLMs. Note that both [4d] and [5e] are only available as preprints so far.

[1a] Poseidon: Efficient Foundation Models for PDEs, https://arxiv.org/pdf/2405.19101

[2b] U-DiTs: Downsample Tokens in U-Shaped Diffusion Transformers, https://arxiv.org/pdf/2405.02730v1

[3c] Scalable Transformer for PDE Surrogate Modeling, https://arxiv.org/pdf/2305.17560

[4d] PDEformer: Towards a Foundation Model for One-Dimensional Partial Differential Equations, https://arxiv.org/pdf/2402.12652

[5e] Unisolver: Unisolver: PDE-Conditional Transformers Are Universal PDE Solvers, https://arxiv.org/pdf/2405.17527

Novelty Even though PDE-Transformer combines improvements from different more established architectures in computer vision, the final architecture is novel and follows the paradigm “use what works best”. It gives SOTA performance on learning PDEs with significantly improved scalability. The modifications are carefully evaluated against the newest, similar SOTA transformer architecture. Additionally, the separate channel (SC) variant is not used in any previous work, and fundamentally improves the downstream performance. For building scientific foundation models, this is one of the most critical aspect: effective pretraining on large datasets so that finetuning on difficult new PDEs works. We have shown that finetuning a pretrained network works much better for the separate channel SC version than when mixing channels (MC). We believe this is an important empirical finding.

We therefore kindly ask you to reconsider your overall recommendation, and we’d be happy to discuss any remaining open aspects.

Thank you for the positive review and feedback.

Pretraining on Downstream Tasks Baseline models are not pretrained. We have a version of our own model that is trained from scratch and one that is pretrained. In all cases, our model trained from scratch is better than the baselines and performance is further improved if it is pretrained, so we think that's fair for a comparison. We also compare with Poseidon (scOT-B, see figure 9 in the appendix) which is pretrained by the corresponding authors on their own large pretraining dataset.

Superresolution experiments In principle, PDE-Transformer can be trained on a coarse resolution and then applied to high-resolution data later. This is tied to both how the attention operation can be viewed as an integral kernel operation (similar to many neural operator papers). Within each attention window, a 2D grid with relative positions is used as an input to a MLP that outputs position-dependent attention scores. We can modify the relative position grid and window size to obtain a model that operates at an increased resolution.

Because of two central modifications that improve the scalability of the architectures, super-resolution without finetuning is difficult to achieve. Those two modifications are patching and the down- and upsampling of tokens within the multi-scale architecture, which both assume data to be at a fixed resolution.

Based on your suggestions, we have evaluated PDE-Transformer trained at resolution 256^2, modified the window size and evaluated the modified model at resolution 512^2. While results still showed the correct dynamics, the accuracy was affected. When finetuning PDE-Transformers to new resolutions we expect the proposed architecture to generalize well and achieve a high performance.

Graph-based models We agree that graph-based models are an important baseline for unstructured meshes. For regular grids, previous work has shown that graph networks have no advantages, but only require a slightly larger weight count compared to networks that leverage the inductive biases from the grid. See e.g. [2] for a comparison.

Nonetheless, we have started to train Message Passing Neural PDE Solvers (MPNN) on the pretraining dataset and plan to add a comparison in the updated manuscript. In the MPNN paper, the 2D experiments only involved grids of size 32^2. For the 256^2 resolution, training has started, but takes several days to finish, because of much bigger memory and compute requirements. Analyzing the architecture, a clear conclusion to draw is that graph networks like MPNN are not designed to be scaled up to large-scale tasks. Due to the required training time, we can only report very preliminary results (epoch 2/100) for the rebuttal.

As mentioned above, we plan to target unstructured grids and datasets as future work, and will make sure to compare to graph networks more extensively for these cases.

[2] Differentiability in Unrolled Training of Neural Physics Simulators on Transient Dynamics, https://arxiv.org/pdf/2402.12971

Different models for pretraining/downstream tasks: We considered different models for pretraining and the downstream tasks. First, we wanted to focus on transformer architectures that are more similar to our architecture for pretraining. For the downstream tasks, we have more "standard" scientific ML baselines. It would be possible to train every model on all tasks, giving an even more extensive comparison, but we believe that grouping models (transformer/scientific ML) is a fair approach that demonstrates relative performance differences. There is one architecture (scOT) that fits in both categories and it was included in both pretraining/finetuning tasks. We included an additional SOTA transformer model, FactFormer. See our response to reviewer yvZY.

UDiT-S better than PDE-S? In Table 1, UDiT-S is slightly better than PDE-S for nRMSE(1). Still, in almost any case it is better to use PDE-S. The different configs S,B,L are not always directly comparable between model architectures, e.g. UDiT-S has more parameters than PDE-S. The performance of transformer models is known to increase when scaling the model size, so a "larger" UDiT-S can beat a "smaller" PDE-S. The B config PDE-B trains much faster than UDiT-S (10h40m vs. 18h 30m) but now clearly beats UDiT-S in terms of nRMSE (0.038 vs. 0.042). Inference speed is correlated to training time. UDiT does not scale well for higher resolutions as demonstrated in figure 3.

We will improve the paper based on your many other helpful suggestions. We had to keep our answers short due to the strict rebuttal character limit.

We are happy to continue in more detail during the discussion phase.

Thank you for the positive review and feedback.

MSE vs. Diffusion This is a very important and interesting issue. We believe that scientific machine learning models should in general be capable of learning the full posterior; however how useful this is still depends on the specific task and the underlying data distribution.

When training a neural solver for example, then the mapping that the neural solver learns is deterministic when given simulation hyperparameters and initial conditions. Another way to phrase this is that given a point in the input space (representing the simulation hyperparameters and initial conditions), then this is mapped to a single point in the output space (the solution field at a certain time). In this situation, supervised training is a suitable choice, especially if we use a metric based on the MSE/L2 distance for evaluation and training.

This situation changes, when the mapping becomes probabilistic, for example when the exact simulation hyperparameters are unknown. This case more closely resembles the setup in the paper. Then, instead of learning a deterministic mapping, we aim to learn a mapping that transports a single point in the input space to a distribution of possible solutions in the output space. In this case, we can learn the mapping via a diffusion model. What happens when we train a diffusion model, but use the common MSE/L2 distance for evaluation? For simplicity, let's say that this distribution is approximately Gaussian. We only have samples from the target distribution that we can compare to. At the same time, our diffusion model will draw a sample from the learned target distribution, which will match the target distribution if our model is trained well.

Now, we evaluate the MSE/L2 distance between the sample from the target distribution and the sample from the "learned" distribution. Let's denote this value by the random variable X. One factor plays a key role now: if we use a metric based on MSE/L2 distance, then it is not optimal when the diffusion model learns the full posterior. In fact, it would be better, if the model just learned to predict the mean of the posterior, because that will decrease the mean and variance of X. If we use supervised training, the network will learn exactly the mean of the posterior.

To summarize, if we use metrics based on MSE/L2 distance for the evaluation, then in theory they will favor supervised models over diffusion models, because it is just not optimal and incentivized to learn the full posterior for these metrics.

Of course that does not mean that the supervised models are better than diffusion models. In fact, we advocate for using diffusion models. However, if the main objective is based on MSE, then it is very difficult to beat a supervised model in terms of the accuracy-compute tradeoff. For more complicated data domains and tasks, we should always use expert knowledge and metrics that make sure the full posterior is considered. Then we can reasonably show the advantages of diffusion models over supervised training. In this paper, we use nRMSE, because of the large number of different PDEs and because there is no single other metric that would work well for all PDEs. For nRMSE, we see that the diffusion models improve significantly when increasing the number of parameters from config S to L and when increasing the number of steps used for inference, closely approaching the supervised baseline.

Generalization to Unstructured Meshes

It's possible to generalize the the architecture to irregular grids, which however this is not the focus of the current paper.

• One straightforward approach is to couple PDE-Transformer with GNO layers as encoder/decoders (instead of the patchification) that map from a given geometry to a latent regular grid as used in e.g. [1]

• Alternatively, it is also possible to generalize the notion of attention window to be defined on local neighbourhoods. Correspondingly, token up- and downsampling are replaced by respective graph pooling operations. This requires more fundamental changes to the architecture making it a type of graph neural network.

Both approaches are planned as directions of future work.

[1] Geometry-Informed Neural Operator for Large-Scale 3D PDEs, https://arxiv.org/pdf/2309.00583

Please let us know if this could clear your open questions. We are happy to discuss more.

Thank you for the positive review and feedback.

BF16 mixed precision That's a good point. With BF16 mixed precision we lose precision compared to FP32, but can train a lot faster. In our experiments, we did not see a difference in the evaluation metrics and training loss when switching between BF16 mixed precision and FP32. Note that in practice we can always train first with the faster BF16 mixed precision until the training loss converges and then finetune with full FP32 precision.

Diffusion models In our opinion, scientific foundation models should be designed with the possibility in mind to produce probabilistic samples. Uncertainty estimates are especially useful when dealing with turbulence. Direct numerical simulation is very expensive, so more simplified turbulence models such as Reynolds-averaged Navier-Stokes simulation and large eddy simulations are still prevalent in the engineering community. However, these models produce posterior distributions, because of unknown latent parameters of the models. As an example of an application that requires the full posterior, see [1] which learns the supersonic flow around aircraft airfoils using diffusion models. Even more computationally challenging is turbulence in 3D, which can be phrased as a generative modeling task to learn all possible turbulent flow states [2]. Scaling 2D models to 3D requires very computationally efficient 2D models to begin with, which we focus on in this paper. We are working on our follow-up work, which addresses efficient scaling of our transformer architectures to high resolutions in 3D.

[1] Uncertainty-aware Surrogate Models for Airfoil Flow Simulations with Denoising Diffusion Probabilistic Models, https://arxiv.org/pdf/2312.05320v1

[2] From Zero To Turbulence: Generative Modeling for 3D Flow Simulation, https://arxiv.org/pdf/2306.01776

Token Upsampling/Pixel Shuffle Token downsampling refers to an operation that merges multiple (4) tokens into a single token. Token downsampling can be implemented for example via PixelUnshuffle, which is part of the Pytorch library (see torch.nn.PixelUnshuffle of the official Pytorch documentation). Respectively, there is token upsampling, which splits a single token into multiple (4) tokens, which can be implemented with PixelShuffle.

Patchification Our implementation uses a single convolutional layer for the patchification where the kernel size and stride are set to the patch size, which is an efficient standard implementation of the patchification. We have performed experiments with extending this single convolutional layer for patchification to small convolutional encoders/decoders. However we didn't see any improvements here and using a pure transformer backbone with a single patchification layer was optimal for us. Our transformer also works in pixel space directly. In principle it can be easily coupled with additional encoders/decoders. However, doing this needs to be carefully engineered. In computer vision, transformer architectures such as the Diffusion Transformer, are mostly coupled with pretrained VAEs and work in a reduced latent space. For scientific machine learning, a pretrained VAE is often problematic, as it makes it difficult to achieve low MSE values.

Which is better: MC or SC? For the pretraining tasks, we are in a situation where there are very large amounts of simulation data to train on. In this case, the mixed channel (MC) and separate channel (SC) variants achieve the same accuracy, however SC requires more computation. In this case, MC still wins the accuracy-compute tradeoff. For finetuning, there are domain-specific applications, where data is more valuable and scarce. Here we see that when trained from scratch, MC generalizes better than SC. However, when we don't train from scratch and finetune pretrained networks, then (1) results always improve for both SC and MC and (2) SC now clearly beats MC and shows much bigger improvements from pretraining. This is because the SC version is better at transferring what it has learned to new types of simulations. Thus for typical use-cases of finetuning a pretrained network, SC is clearly preferable.

Comparison with Poseidon/MPP We did in fact compare our model extensively to Poseidon. The model from Poseidon is the scOT model we compare to in all experiments. We also compare to the pretrained model, Poseidon-B for the downstream tasks. Because of space limitations, we could not include a full comparison in the main text, but it is discussed and we refer to figure 9 in the appendix.

We saw Poseidon as the most recent and "toughest" competitor, so we wanted to make sure we have extensive comparisons with it and show that we can outperform it. MPP is already used as a baseline in the Poseidon paper and Poseidon achieves much better performance there, so our priority was to focus on Poseidon.

Please let us know if this could clear your open questions. We are happy to discuss more.