PDE-Transformer: Efficient and Versatile Transformers for Physics Simulations
Authors: Benjamin Holzschuh
Qiang Liu
Georg Kohl
Nils Thuerey
Venue: ICML 2025 poster (2025)
Keywords: physics simulations
PDEs
transformer architectures
scalability
Primary area: applications->chemistry_physics_and_earth_sciences
TL;DR: A novel transformer architecture for physics simulations with improved performance on complex downstream tasks.

Abstract: We introduce PDE-Transformer, an improved transformer-based architecture for surrogate modeling of physics simulations on regular grids. We combine recent architectural improvements of diffusion transformers with adjustments specific for large-scale simulations to yield a more scalable and versatile general-purpose transformer architecture, which can be used as the backbone for building large-scale foundation models in physical sciences. We demonstrate that our proposed architecture outperforms state-of-the-art transformer architectures for computer vision on a large dataset of 16 different types of PDEs. We propose to embed different physical channels individually as spatio-temporal tokens, which interact via channel-wise self-attention. This helps to maintain a consistent information density of tokens when learning multiple types of PDEs simultaneously. We demonstrate that our pre-trained models achieve improved performance on several challenging downstream tasks compared to training from scratch and also beat other foundation model architectures for physics simulations.
Our source code is available at https://github.com/tum-pbs/pde-transformer.

PDE-Transformer:
Efficient and Versatile Transformers for Physics Simulations
Benjamin Holzschuh
Qiang Liu
Georg Kohl'
Nils Thuerey'
Abstract
et al., 2024; Zhang et al., 2024a), where they can be used in variety of downstream tasks such as surrogate modeling (Kim et al., 2019; Sun et al., 2020) or inverse problems involving physics simulations (Ren et al., 2020; Holzschuh et al., 2023). They are especially promising in areas where only low amounts of training data are available.
We introduce FDE-Transformer an improved transformer-based architecture for surrogate mod- cling of physics simulations on regular grids. We combine recent architectural improvements of dif- fusion transformers with adjustments specific for large-scale simulations to yield a more scalable and versatile general-purpose transformer archi- tecture, which can be used as the backbone for building large-scale foundation models in physi- cal sciences. We demonstrate that our proposed architecture outperforms state-of-the-art trans- former architectures for computer vision on a large dataset of 16 different types of PDEs. We propose to embed different physical channels in- dividually as spatio-temporal tokens, which inter- act via channel-wise self-attention. This helps to maintain a consistent information density of to- kens when learning multiple types of PDEs simul- tancously. We demonstrate that our pre-trained models achieve improved performance on several challenging downstream tasks compared to train- ing from scratch and also beat other foundation model architectures for physics simulations. Our source code is available at /ttps://githhb. han//tu--bbs /talb///de-eeee trans former.
The adoption of machine learning for physies simulations faces several characteristic challenges. The underlying physics often exhibit an inherent multi-scale nature (Smith, 1985). The representation of data is tied to the numerical methods used for simulations, ranging from regular grids to meshes to pettcle-based representations (Anderson et al., 2020). While numerical simulations produce vast amounts of data, they oftentimes cannot readily be used as input to machine learning models further preprocessing re- quired. Additionally, machine learning models in this area directly compete with traditional numerical methods. As such, they need to either outperform such solvers in terms of accuracy or speed while exhibiting high reliability. Alter- natively, they need to provide solutions that solvers cannot easily obtain, for example, uncertainty estimates (Geneva Zabaras, 2019; Jacobsen et al., 2025; Liu & Thuerey, 2024; Kohl et al., 2024) working with partial inputs (Wu et al., 2024a).
We address these problems by introducing a multi-purpose transformer architecture 2DE-Transforme that is carefully tailored to scientific machine learning: the same architecture can be used to generalize between different types of PDEs, different resolutions, domain extents, boundary conditions. and it includes deep conditioning mechanisms for PDE- and task- information. Our model works on regular grids in two dimensions, and can be trained in supervised manner as an efficient surrogate model, or as a diffusion model for downstream tasks where solutions have wider posterior distributions. Specifically, our contributions are:
1. Introduction
Large multi-purpose networks that are trained on diverse, high-quality datasets have shown great achievements when adapted to specific downstream tasks. These so-called fown- dation models have received widespread recognition in fields such as computer vision (Yuan et al., 2021; Awais al., 2025), decision making (Yang et al.. 2023b). and time series prediction (Liang et al., 2024). Naturally, these models are likewise extremely interesting for the scientific machine learning community (Bodnar et al., 2024; Herde
We augment existing SOTA diffusion transformer ar- chitectures to tailor them to PDE and physics simula- tions, among others, by token down- and upsampling for efficient multi-scale modeling and shifted window attention for improved scaling to high-resolution data.
'School of Computation, Information and Technology, Techni- cal University of Munich, Germany. Correspondence to: Benjamin Holzschuh <benjamin.holeeee holzschuh @tum.de>.
We modify the attention operation to decouple token interactions between the spatio-tempore axes and the physical channel axis. This improves accuracy and
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

TranTronsfo:me:: Efficient and Versatile Transformers for Physics Simulations
Burgers
Decaying Turbulence
Kolmogorav
Gray (e)
Scott
Soott
Kuramol-SSrashinnk
Reference
Redietionn
3
Figure 1. 3oE--ransformer is transformer model tailored to scientific data, the images above show its autoregressive predictions after 20 time steps on large dataset comprising 16 different PDE dynamics, given initial conditions only. Additional simulation parameters (viscosity, domain extent, etc.) are unknown to the model and need to be inferred from the observed data. Transfornfrre is especially well suited to be pre-trained for out-of-distribution downstream tasks.
allows for better generalization to different PDEs.
class labels or text encodings can be used as an additional imput. Diffusion transformers use a significantly lower patch size than many previous transformer models for computer vision tasks. While diffusion models show an increasing per- formance with lower patch sizes, they can only be used for high-dimensional data by working in the latent space. This is caused by the quadratic scaling of the required compute with the number of tokens, due to the global self-attention mechanisms of the diffusion transformer architecture, mak- ing training and inference computationally expensive. In this paper, focus modeling the raw data directly with- out including any pre-trained autoencoders.
We perform a detailed ablation study on accuracy- compute tradeoffs when scaling and modifying PDE- Transformer.
We demonstrate that PDE-Transfommer generalizes well to challenging downstream tasks when pre-trained on set of generic PDEs.
2. Related Work
Translormers Transformers have become one of the dom- inant deep learning architectures. While first established for sequence-to-sequence models in natural language pro- cessing (Vaswani et al., 2017), transformers have been suc- cessfully applied to computer vision by splitting images into patches, which are embedded jointly with their position as tokens (Dosovitskiy al., 2021, VIT). VITs have been shown to scale better vision tasks than classical convo- lutional neural networks and benefit from high accelerator utilization (Maurício et al., 2023; Takahashi al., 2024: Rodrigo et al., 2024). Large transformer- based architectures have led to foundation models such as GPT (Radford et al., 2018; 2019; Brown et al., 2020) and BERT (Devlin et al., 2019).
A Attention mechanisms The global self-attention opera- tion of transformers is the main computational bottleneck of the transformer architecture. There have been two main directions addressing this. Windowed attention limits the computation of self-attention to non-overlapping local win- dows (Liu et al., 2021). For different layers, this window is shifted to prevent discontinuities at the window borders. Simiarly, in axial attention (Ho et al., 2019), tokens only interact with each other along a specific axis, i.e. when they belong to the same image row or column. When combined with down- and upsampling in hierarchical architecture, the resulting models show the same or improved perfor- mance while operating at lower compute cost. The cost of attention can also be lowered by algorithmic modifications to the attention mechanisms, which make the attention scale linearly (Wang et al., 2020: Cao, 2021). However, these models often report subpar performance compared to the quadratic self-attention. For physics simulations, the impor- tance of transformers and attention mechanism has likewise been recognized. Galerkin transformers (Cao, 2021) uti- lize a linearized attention variant that removes the softmax normalization and can be related to learnable layer-wise Petrov-Galerkin projection. Multiple Physics Pertaining (McCabe al., 2023, MPP) uses custom transformer
Diffusion Transformers While carly applications of transformers in computer vision focus on visual recogni- tion tasks, the success of latent diffusion models has also led the adoption of wansformer-basee diffusion models, the diffusion transformer (Peebles & Xie, 2023, DiT). In- stead of modeling the data space directly, latent diffusion models (Ho et al., 2020; Rombach et al., 2022) use a vari- ational autoencoder to embed the data and operate in the resulting latent space. Additionally, diffusion transform- ers have very powerful conditioning mechanisms based on adaptive layer normalization (Perez et al., 2018) in which
2

---

PrE-Transformr:: Efficient and Versatile Transformers for Physics Simulations
Mixed channel (MC)
Input
Separate channel (SC)
Channel A
Outpur
Projection
Scale
MLP
Transformer Stage
Scale, Shift
Transformer Stage
Layer
Mixed channels
Separate channels
Channel-wist wise MHSA
Transformer Stage
Scale
Attentton windows
Channel types
Type
Transformer Stage
Shifted Window
MHSA
Transformer Stage
Scale, Shift
"Embedding
Shifted Window MHSA
Shifted-window self- seff-attention
Channel-wise self-attention
MLLP
Tokemizer
Layer
Conditioning
Input
Figure 2. Architecture overview of PDE-Transtonne The multi-scalc architecture combines up- and downsampling of tokens with skip connections between transformer stages of the same resolution. The attention operation is restricted to local window of tokens. The window shifted between two adjacent transformer blocks. Conditionings are embedded and used to scale and shift the intermediate token representations. The mixed channel (MC) version embeds different physical channels within the same token. The separate channel (SC) version embeds different physical channels independently. Tokens of different physical channels only interact via axial self-attention over the channel dimension. The types of channel (velocity, density, etc.) are part of the conditioning, which is distinct for each channel.
backbone based on an axial ViT, where axial attention is computed along the temporal and spatial dimensions.
masked reconstruction (Devlin et al., 2019; He et al., 2022) and contrastive learning (Chen et al., 2020). For PDEs, pre- training via autoregressive next-step prediction can include additional simulation parameters (Gupta & Brandstetter, 2023; Takamoto et al., 2023; Subramanian et al., 2023) or can be based on previous snapshots only (Herde et al., 2024). Pre-training has also been explored for PDEs in the context of neural operators (Li et al., 2021; Goswami et al., 2022: Wang et al., 2022), and domain transfers (Xu et al., 2023). Recently, models have also been pre-trained on multiple PDE dynamics at the same time (Subramanian et al., 2023; Yang et al., 2023a; McCabe et al., 2023).
Learning for PDEs While physical residuals pose diffi- cult learning tasks (Raissi et al., 2019; Bruna et al., 2024), neural networks can also be combined with existing PDE solvers via learned corrections (Um et al., 2020; Dresdner et al., 2023; Thuerey et al., 2021), Iearned closures (Du- raisamy al., 2019; Sirignano & MacArt, 2023), or learned algorithmic components (Bar & Sochen, 2019; Kochkov et al., 2021), such computation stencils -Bar-Sinai al., 2019). Significant interest has been sparked in building neural operators (Lu et al.. 2021; Kovachki et al., 2023), which map between function spaces. The attention mecha- nisms can be generalized to arbitrary meshes by including the distance between nodes as weighting, converging to an integral kernel operator as the resolution increases (Li et al., 2023a;b). As result, transformers and attention can be broadly generalized, allowing them to work ar- bitrary input and output points (Wu et al., 2024a). While enjoying beneficial theoretical properties, the added flexibil- ity makes it difficult to scale to high resolutions and many points. Scalable operator transformer (Herde et al., 2024, scOT) is a hierarchical vision transformer with shifted win- dows for the attention computation. In contrast scOT, PDE-Transfommer is designed from the ground up improved diffusion transformer for PDEs, and significantly outperforms the former.
3. Transfooffrrmrr
We denote spatiotemporal system by S. This system encompasses n physical quantities u(x, f) Is x [0, 7] + Rn over spatial domain Ug C R We assume the data is discretized in time and space, i.e. the system is described by sequence [uğ _ w where each snap- shot uب is sampled on the points determined by the spatial discretization. We denote all additional information about such as the type of PDE or simulation hyperparameters by the vector c. The network, denoted Me, is parameterized by weights 0.
Notation
Autoregressive Prediction For autoregressive prediction tasks, the goal is to predict the snapshot ug given sequence
Pre-training Networks can be pre-trained using different strategies: autoregressive prediction (Radford al., 2018),
of Tр preceding snapshots u TT. .... u. Art Throughout
the paper, we abbreviate uin [uş_ TT-TT,A At ... S Ad and
3

---

PDE-Transformer: Efficient and Versatile Transformers for Physics Simulations
Mout [u]] for autoregressive prediction. The definitions of Unn and Hout can be easily adapted represent different tasks, e.g., Uin = 0 and Hout [ for unconditional generative model, or Um [u/-s a and Uout
our implementation.
Shifted Windows To prevent the quadratic blowup of the global self-attention in DiTs, we adopt the shifted win- dow multi-head self-attention (MHSA) operation used by Taannrranforrrrss (L.iu et al., 2021), which limits the self- attention between tokens to a local window. A window with window size w encompasses w spatio-temporal tokens at each block. To prevent discontinuities at the window bor- ders between two adjacent layers, the windows are shifted by w/2 tokens. We evaluate the impact of the window size on the accuracy-compute tradeoff in Section 4. In contrast to DiTs, no absolute positions are added to the token em- beddings. We use log-spaced relative positions of tokens within each window combined with a feed-forward neural network when computing attention scores (Liu et al., 2022). This improves translation-invariaayy for learning PDEs and increases generalization across different window resolutions.
[u:] for temporal interpolation.
3.1. Design Space
We augment the diffusion transformer backbone (Peebles & Xie, 2023, DiT) to obtain PDE- Trr-Trrnsformrr. Its architecture is summarized in Figure 2. We describe the DiT architecture and our modifications step by step in the following.
Patching DiT operates in a latent space with fixed di- mension, whereas PDE-Transforme operates on the raw data, For data consisting of single physical channel, given patch size p, an input of size T x H x W is partitioned into H/p W/p patches of size T p X p. These patches are embedded into tokens representing a spatio-temporal subdomain via a linear layer that maps each patch to a d-dimensional vector. The patch size is a critical hyperpa- rameter, as it directly controls the granularity of the token representation. By halving the patch size, the number of tokens quadruples, significantly increasing the required com- pute. We define an additional quantity, which we call the expansion rate E(p)= == d/ d/)"T) that specifies the rate at which the input data expands. Low expansion rates are bet- ter for scalability, as they coincide with fewer tokens. On the other hand, there is a clear correlation between patch size and the performance of the trained model, where lower patch sizes achieve better results. We explore this trade-off between accuracy and compute costs for physies simulations in Section 4.
Mixed and separate channel representations PDEs can involve multiple physical quantities and hence require differ- ent numbers of physical channels, fundamental difference to the fixed channel count of DiT architectures. The dynam- ics and scales of the physical channels can be fundamentally different. A naive strategy for this situation is to define a maximum number of channels Cmax as an additional di- mension of the input. Therefore, the embedding layer will map the spatio-temporal patch T x Cmax X p x p to the d-dimensional token embedding. If the data has fewer than Cmax channels, they are padded with zeros. This is the mixed channel (MC) variant of TogE-Tranformer. Effec- tively, this reduces the expansion rate E(p) by the factor 1/Cmax and mixes channels with different physical mean- ing. This leads to an overly compressed token representa- tion, which can deteriorate performance and lead to poor generalization capabilities. Instead, we propose to scale the compute with the number of channels and keep the expan- sion rate of each token the same independent of the number of channels. We achieve this by embedding each channel independently and by implementing interactions between tokens of different channels via an additional channel-wise axial MHSA operation in each block. The windowed self- attention is not computed between tokens of different chan- nels. This variant is called separate channel (SC) below, and we evaluate both variants in detail for the pre-training dataset and when finetuning in Section 4.
Multi-scale architecture While U-shaped architectures like the classic UNet (Ronneberger et al., 2015) have shown remarkable success, token down- and upsampling have not been used by DiT's in favor of a more streamlined archi- tecture by the original authors (Peebles & Xie, 2023). The hierarchical structure of U-shaped architectures resembles the multi- nature of features in nature and adds a strong inductive bias. Several works (Bao et al., 2023; Hoogeboom et al., 2023; Tian et al., 2024) have combined U-shaped architectures with a transformer backbone architecture. We introduce down- and upsample tokens at the end of each transformer stage via PixelShuffle and PixelUnsuuffe lay- ers. In contrast to Bao et al. (2023) and Hoogeboom et al. (2023), we rely on adaptive layer normalization for condi- tioning. Tian et al. (2024) also feature a U-shaped design, but it is combined with token downsampling on the query- key-value tuple of the self-attention operation. We found that this slightly improves performance but comes at the cost of increased training and inference time, due to subop- timal accelerator utilization. Thus, we did not downsample tokens within the self-attention operation for PDE data in
Conditioning mechanism Similar to DiTs, we use adaLN-Zero blocks (Goyal et al., 2017: Perez et al., 2018; Dhariwal & Nichol, 2021). The embedding vectors of all conditionings are added and vectors for scale and shift op- erations are regressed using a feed-forward network. In addition, the scale and shift vectors are initialized such that each residual block is equal to the identity function, which accelerates training. DiTs use the class label and diffusion
4

---

PDE-Transformer: Efficient and Versatile Transformers for Physics Simulations
time as conditioning. When training PDE-Transformer diffusion model, we also include the diffusion time condi- tioning. The class label corresponds to the PDE type in our setup. Moreover, for each physical channel in the SC ver- sion, we use a label embedding for the type of channel (e.g. density, vorticity). All labels use dropout with probability 10%, therefore our model can be used in both a conditional and unconditional setting. Extending this conditioning to additional simulation parameters is straightforward.
are mapped to samples x; from the posterior P1 via the or- dinary differential equation (ODE) dxx, o(xx,t) dt. The network Me learns the velocity by regressing a vector field that generates a probability path between Po and pr Samples along the probability path are generated via the forward process
= + [1 - (1-mmnnn) e
(2)
for E [0, 1] with c NN I) and a hyperparameter
min 10-1, We denote ct [c, and u [uin xt]
Boundary conditions We explicitly consider periodic and non-periodic boundary conditions for the simulation do- mains. When shifting the attention windows, the tokens are rolled along the x- and y-axis as if they were aligned on spatial grid matching their position. This mimics pe- riodic boundary conditions for both axes. If periodicity is not required, it can be explicitly disabled in the architecture by masking token interactions in the computation of the attention scores.
The velocity can be regressed by training with the loss
LFM - E [[] c')
Uont + (1 (1-a0 omin)e ) (3)
Once trained, samples from the posterior can be generated conditioned on Uin and с by sampling Xo ~ (0, and solving the ODE dx, = MM c') dt from t = 0 until 1. We use the explicit Euler method and experiment with choosing good step size At for PDEs in Section 4.
Algorithmic improvements We normalize Q and K of the self-attention operation using RMSNorm (Zhang& & Sen-
4. Experiments
nrich, 2019) to avoid instabilities due to an uncontrolled growth of the attention entropy (Dehghani et al., 2023). In addition, we find that the training configuration of diffusion transformers from Peebles Xie (2023) leads to training instabilities and suffers from spikes in the training loss for our datasets. Thus, we alter the learning rate from 1.0- 10-4 to 4.0 10-5 and employ the AdamW optimizer using small amount of weight decay with factor of 10-15 15 for bfle-mixed precision training as recommended by Esser et al. (2024). Moreover, we find that gradient clipping based on the exponential moving average (EMA) of gradients pre- vents any remaining spikes in the loss curve and ensures stable training. Further details of the training approach can be found in Appendix A.
We evaluate the performance of toregressive prediction with Tр = preceding snapshots. Our experiments are divided into two parts: first we compare PDE- Transformer to other SOTA transformer architectures on a large pre-training set of different PDEs, focusing on accuracy, training time and required compute. We moti- vate our design choices for PDE-Transformer in an abla- tion study. Second, we finetune the pre- trained network on three different challenging downstream tasks involving new boundary conditions, different resolutions, physical channels and domain sizes, showcasing its generalization capabilities to -fistt--iitiiuutiin data. Our models use three different configurations S, B and L. that have different token embedding dimensions d (96, 192 and 384 respectively). We use the name PDE-S to refer to Transfonsfree with configuration S. Unless otherwise noted and w are set to their default values = 4 and w = 8 and the mixed channel (MC) version is used.
for au-
3.2. Supervised and Diffusion Training
Supervised PDE-Transformer can be trained both in a supervised manner and as diffusion model. For tasks with deterministic solution, for example when training surrogate for a deterministic solver, supervised training of #DE-Transformer with the MSE can be used, which allows for a fast inference in one step. In this case, the network is directly trained with the MSE loss
4.1. Pre-training dataset
Training We train our models on a large set of 16 linear and non-linear PDEs, including Kolmogorov flow, Burgers' equation, different variants of the Gray Scott equation and many more. The datasets are based on APEBench (Kochler et al., 2024), and described in detail in Appendix C. each type of PDE, we consider 600 trajectories of 30 simulation steps each, which are randomly split into a fixed training, validation and test set. The data is generated with a spectral solver at resolution 2048 x 2048 and downsampled to 256 x 256 with PDEs having either or 2 physical channels. We apply gradient accumulation when training larger models to keep the batch size unchanged. For evaluation, use an EMA of the models weights with decay of 0.999.
Ss E [I|Me c) Houe
Diffuusion training If the solution is not deterministic, then diffusion training is preferable, as it enables sampling from the full posterior distribution, instead of learning an av- craged solution. We adopt the flow matching (Lipman et al., 2023; Liu et al., 2023) formulation of diffusion models (Ho et al., 2020) for training. Given the input um and condition- ing c, samples Xo from a noise distribution pa = .N(0 T)
5

---

PDE-Transformer: Efficient and Versatile Transformers for Physics Simulations
512
decreases, irrespective of learning rate and gradient clipping strategies. This behavior indicates stability issues of the DiT architecture for the characteristics of PDE datasets. For evaluating DiT-S, we therefore use the checkpoint with the lowest validation loss.
DiT-S
0.5 GB
400
U-DIT-S
1.5 GB
PDE-S
GFlops
300
200
256
100
128
Table I. Training S configurations for 100 epochs on 4x H100 GPUs.
0-
Domain size
Model
nRMSE;
nRMSE10
Time (h)
Params
DiT-S
0.066
0.78
13h 4m
39.8M
Figure 3. Scaling the domain size for patch size p 4. An input resolution 256 256 corresponds to x tokens. The area of each point corresponds to the required GPU memory for inference (batch size I).
UDiTS
0.042
0.39
18h 30m
58.9M
scOT-S
0.051
0.59
21h 11m
39.8M
FactFormer
0.069
0.65
12h 25m
3.8M
UNet
0.075
0.68
48h 00m
35.7M
PDE-S
0.044
0.36
7h 42m
33.2M
Evaluation metries For the evaluation, we use the nor- malized RMSE, defined as
Efficient scaling of domain size and model parameters We compare the compute for inference (GFlops) and GPU memory requirements (GB) of PDE-S, UDiT-S and DiT-S for different input sizes assuming a patch size of p 4. DiT- S is computationally very expensive for larger domains due to the global self-attention. UDiT-S scales better in general; however, GPU memory requirements scale badly due to convolutional layers within the attention operator for token up- and downsampling. PDE-S uses few convolutional layers possible for up- and downsampling within the U-shaped architecture, and achieves the best scaling for GFlops and GPU memory. We train the PDE-Transfore Transformer with confi ooffiuurations S, and L. The overall architecture is kept identical, but we increase the dimension of the token embeddings d with each configuration. see Table 2. Larger token embeddings lead to an improved performance. While the number of GFlops for inference increases by factor of ca. 4x when doubling the token dimension, the total training time increases at a lower pace. This is due to an efficient accelerator utilization of matrix-multiplicatoon in the self aelf-attention operator on modern GPUs. The supervised and flow matching losses during training for the different configurations are shown in Figure in the appendix.
M nRMSE = M Σ
MSE oout Mout)
(4)
MSE(O, 0, Uout)
where üout is the network prediction and M corresponds to the number of trajectories in the test dataset. We can autoregressively generate the entire trajectory for system S and define the rRMSE at time comparing the predicted state uf at time with the reference u.S. We average over all systems in the test datasets for each timestep.
Windowed attention and multi-scale architecture We train several distinct and representative transformer mod- els: DiT (Peebles Xie, 2023) with no token down- or upsampling, UDiT (Tian et al., 2024) featuring U-shape architecture, scOT (Herde et al., 2024), neural operator transformer that includes a hierarchical architecture and shifted window attention, FactFormer (Li et al., 2023b) a transformer model based on axial factorized kernel inte- grals, a modern UNet (Ronneberger et al., 2015; Ho et al., 2020), and our proposed cable, all models are trained with the S configuration of their architecture. We evaluate the trained models using the nRMSE and report values for the first step and after 10 steps of autoregressive prediction, see Table 1. We use patch size of = 4, resulting in 64 x 64 spatiotemporal tokens. UDiT-S and PDE-S achieve the best scores. Addi- tional important differences between the two are visible in terms of required training time: first, PDE-S can be trained much faster (7h 42m). DiT-S has much lower training time (13h 4m) compared UDiT-S (18h 30m), even though the number of GFlops is higher. We believe this is caused by better accelerator utilization of the DiT architecture on modern GPU hardware in comparison to UDiT. When training DiT-S, spikes in the loss curve occur that the model does not recover from. This happens repeatedly as the loss
When appli-
Axial attention for different physical channels We com- pare the mixed channel (MC) and separate channel (SC) versions of PDE- TDE-Transformer for all three configurations in Figure 4 (bottom). Most importantly, the flexibility of the proposed SC embeddings does not lead to any deterioration of inference accuracy. However, the compute of the SC variant increases, since the number of tokens scales linearly with the number of channels. Due to the fundamentally improved flexibility of the SC version will focus on it in
We train the UNet for 40 epochs, reaching a maximum com- pute budget 2 days.
6

---

PDE-Transforme:: Efficient and Versatile Transformers for Physics Simulations
10a
the diffusion version has a slightly lower accuracy in this evaluation, it is close the supervised baseline. However, the ability to sample from the posterior, even at the cost of more network evaluations, is extremely useful in many downstream tasks and practical engineering applications. We also experiment with the step size At for sampling, see Figure in Appendix and find that performance keeps improving with more steps.
DRMSE
10
Supervised
Flow Matching (100 steps)
2.5
7.5
10.0
12.5
15.0
17.5
109
10a
ARMSE
DAMSE
Window size
Patch size
10-1-
10
Mixed channels Separate channels
10
15
20
10 Time
Time
2.5
10.0
17.5
Time
1.2
16-3
1e-3
Window size
Patch size
Training loss
1.0
16
Figure 4. nRMSE evaluation of supervised training vs. training with flow matching and sampling from the posterior (top). Com- parison of mixed channels (MC) vs. separate channel (SC) with axial attention over channel dimension (bottom).
16
0.8
20
30
40
50
So GFlops
1:00
GFlops
the following, showing its improved versatility and general- ization capabilities in the downstream tasks of Section 4.2.
Figure 5. Impact varying window size (left) and patch size (right) on performance and required compute (GFlops). On the left, the patch size is kept at p and the self-attention window is increased. On the right, the patch size and window size ws chosen such that the product of p and w is constant (p w = 32). Windows cover the same spatial domain.
Supervision versus probabilistic learning We evaluate supervised training against training PDE-Transformer diffusion model via flow matching. While the diffusion TDE-Transformer samples from the posterior distribution, supervised training is incentivized to predict the mean of this posterior. Depending on the PDE, there are better suited evaluation metrics when considering a distribution of possi- ble solutions, however our datasets comprises many differ- ent types of PDEs. Therefore. we only compare using the ARMSE, as shown the top of Figure 4. As expected, the supervised training consistently outperforms samples from the diffusion version, since we consider only single gener- ated sample for each trajectory in the test set. Even though
Influence of patch and window size We evaluate differ- ent patch and window sizes for PDE-Transformer For this experiment, we focus on the S configuration with super- vised training. We first test the effeets of different window sizes w while keeping the patch size fixed at p = This means that the receptive field increases in the windowed self-attention, making interactions less local. While there is decrease in training loss for larger window sizes, these lead to larger compute costs, shown in Figure 5 on the left. At the same time, the nRMSE evaluation on the test set does not improve noticeably, indicating overfitting. We conclude that small window sizes are sufficient for the representative PDE behavior in our dataset. In the second experiment, we adjust window size and patch size at the same time. We keep the receptive field constant, i.e. when changing the patch size, we also modify the window size such that p wr = 32 is constant, as shown in Figure 5 on the right. Smaller patch sizes lead to improved results, however they also signifi- cantly increase the required compute. There is a sweet spot for the accuracy-comput tradeoff at p = and 8.
Table Scaling the token embedding dimension d. nRMSE; shows the results for the supervised training. The training time for 100 epochs is reported on 4x H100 GPUs.
Config
d
nRMSE,
Time (h)
GFlops
S
96
0.045
7h 42m
19.62
B
192
0.038
10h 40m
76.55
L
384
0.035
20h 8m
302.34

---

PDE-Transformer: Efficient and Versatile Transformers for Physics Simulations
4.2. Downstream tasks
Qurs-S Durr-S--retrained
pretraineedd
FNO scOT-s
Galerkin OFormer
Ours-S
In this section, we evaluate the performance of our models on three challenging downstream tasks taken from the Well repository (Ohana et al., 2024): active matter, Rayleigh- Bénard convection (RBC), and shear flow. These tasks con- tain substantially more complex dynamics than any datasets seen during the pre-training phase, involving varied bound- ary conditions (periodic and non-periodic) and geometries (square and non-square domains). Detailed descriptions of these learning tasks provided in Appendix F.1. In Section B.4 of the appendix, we also present preliminary evaluation of the performance of our method in conjunction with low-rank adaption (Hu et al., 2022, LoRA).
Average nRMSE
0.6
0.4
0.0
Active Matter
RBC
Shear flow
Figure 6. Average rollout nRMSE of different models on down- stream tasks.
Table 3. Comparison of the TDDnsforsfssssss with separate chan- nel (SC) and mixed channel (MC) representations. The values are obtained by averaging the nRMSE of the first 20 rollout steps.
For the downstream tasks, we compare Transforrferm to a range SOTA models for PDE prediction: scOT-S (Herde al., 2024), a Galerkin Transformer (Cao. 2021). OFormer (Li al., 2023a), and Fourier Neural Operator (Kossaifi et al., 2024, FNO). All baseline models are trained from scratch. For PDE- Transformer, we train two configuration S models: one initialized from scratch and the other from pre-trained weights. All models have a similar number of trainable parameters, except for OFormer, whose memory consumption increases significantly with model size. It is important to note that the scOT model and Galerkin Trans- former are limited to square simulation domains, making them unsuitable for the RBC and shear flow tasks. Further details on each model can be found in Appendix A.3.
Scratch
Pre-trained
Improvement
Active Matter
SC
0.494
0.455
7.89 %
MC
0.493
0.479
2.84%
SC
0.155
0.104
32.90%
RBC
MC
0.147
0.130
11.56%
Shear Flow
SC
0.199
0.125
37.19%
MC
0.178
0.163
8.43%
Figure shows the results of this comparison. PDE-S with pre-training consistently yields more accurate predictions than the other baselines across the full range of difficult Well tasks. It provides average improvement of 42% in terms of nRMSE, compared the second best model (FNO), demon- strating that PDE-Transformer consistently outperforms the range of baseline models. Notably, although pre-training was conducted with idealized data from a spectral solver, it improves the performance for three downstream tasks featuring highly complex dynamics. In addition, while be- ing pre-trained periodic, square domains, the pre-training contributes to improved performance tasks with non- periodic boundary conditions (RBC) and non-square do- mains (RBC and shear flow). Detailed numerical values of the error during rollouts can be found in Appendix B.2. Additionally, note that the Poseidon model (Herde et al., 2024) provides pre-trained weights for the scOT model at evalu-
predictions. Overall, these results highlight the benefits of the proposed architecture and pre-training approach.
In Table we revisit the performance of the PDE-S using separate channel (SC) and mixed channel (MC) represen- tations with the downstream tasks. While the performance of models without pre-training is similar, the SC version provides greater performance boost on downstream tasks when pre-trained weights are used. Across the three chal- lenging tasks, performance improvements with SC are 2.7x to 4.4x higher than the MC counterparts. This observation further highlights that the proposed -haneel-independentt representation enables the network retain more knowl- edge from pre- pre-training without requiring extensive training or fine-tuning of the PDE-specific layers.
5. Limitations
configuration B. To enable fair comparison, ate the performance of pre-trained PDE-B alongside the pre-trained scOT-B, as both models are of similar size. The PDE-B model demonstrates significant improvements in terms accuracy, with an accumulated nRMSE reduc- tion of 75% compared to Poseidon's 's scOT-B model, which becomes unstable for long-term rollouts. Figure in the appendix presents the average rollout nRMSE of the model
PDE- Transformers currently limited to 2D regular grids, Additionally, we focused on autoregressive prediction tasks. Testing and extending PDE- TDE-Traasformer for tasks with noisy and only partially observed data, data assimilation, or in- verse problems remains open task for future work.
There are several directions for extending PDE- Transformer to unstructured grids. First, TDD-Transformer can be cou- pled with Graph Neural Operator (Li al., 2020, GNO)
8

---

PDE-Transformer: Efficient and Versatile Transformers for Physics Simulations
layers the encoder and decoder, replacing the patchifica- tion, which map between the given geometry and latent regular grid. Alternativlly, it is also possible to general- ize the notion of attention window to be defined on local neighbourhoods of graphs. Correspondingly, token up- and downsampling need to be replaced by respective graph pool- ing and upsampling operations.
R., and Shankar, V. Compulational fluid mechanics and heat transfer. CRC press, 2020.
Awais, M., Naseer, M., Khan, S., Anwer, R. M., Cholakkal, H., Shah, M., Yang. M.-H., and Khan, F. S. Foun- dation models defining a new era vision: a sur- vey and outlook. IEEE Transactions on Pattern Anal- ysis and Machine Intelligence, pp. 1-20, 2025. doi: 10. T.LL1//TAAMI2202435566..
6. Conclusion
Bao, F., Nie, S., Xue, K., Cao, Y., Li, C.,, Su, H., and Zhu, J. All are worth words: A vit backbone for diffusion models. In IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 22669-22679, IEEE, 2023, URL https: /ww.ioiiooooooooooooooo org/ 100 1109 CVVR52729. 1023.0211
We have introduced Traasfaanffrrmr, multi-purpose transformer model that addresses key challenges in physics- based simulations. Its multi-scale architecture modifies the diffusion transformer backbone and limits token interactions to local window without sacrificing performance on learn- ing PDE dynamics. PDE-Traasformrr outperforms state-of- the-art transformer architectures in terms of the accuracy- compute trade-off, demonstrating superior scaling for high- resolution data. By decoupling token embeddings for dis- tinct physical channels, we further enhance performance, especially during fine-tuning on complex downstream tasks. Its focus on scalability makes TDE-Transfrrmrr an excellent backbone for foundation models for high-resolution simula- tion data. In the future, aim to extend TDFF-Taanformm from 2D to 3D simulations.
Bar, L. and Sochen, N. Unsupervised deep learning algo- rithm for pde-based forward and inverse problems. arXiv preprint arXiv: arXiiv/994.080 19009994441111 2019.
Bar-Sinai, Y., Hoyer, S., Hickey, J., and Brenner, M. P. Learning data driven discretizations for partial differen- tial equations. Proceedings of the National Academy of Sciences, 116(1))113444-5539) 2019. ISSN 0027-8424, 1091-6490. doi: 10. 10.1073/pnas.181 18141881144
Bodnar, C.,, Bruinsma, W. P., Lucic, A., Stanley, M., Brand-
Acknowledgements
stetter, J., Garvan, P., Riechert, M., Weyn, J., Dong, H.,
Vaughan, A., et al. Aurora: foundation model of the atmosphere. arXiv preprint arXiv:2405. 13063, 2024.
This work was supported by the ERC Consolidator Grant SpaTe (CoG-2019-863850 The authors gratefully acknowl- edge the scientific support and resources of the Al service infrastructure LRZ AI Systems provided by the Leibniz Su- percomputing Centre (LRZ) of the Bavarian Academy of Sciences and Humanities (BAdW), funded by Bayerisches Staatsministerium für Wissenschaft und Kunst (StMWK).
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. Language models are few-shot learners. Advances in neural information processing systems, 33: 1877-1901, 2020.
Bruna, J., Peherstorfer. B., and Vanden-Eijnden, E. Neural galerkin schemes with active learning for high- nigh-dimensional evolution equations. J. Comput. Phys., 496: 12588, 2024. doi: 10.1016/J.JCP.2223. 112588. URL https: /tt:://ooo http:/doi..o// http://.i.ooo//../00//// jcp.2023 2023. 112588.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
Buchler, E. L. and Buehler, M. J. X-lora: Mixture of low- rank adapter experts, a flexible framework for large lan- guage models with applications in protein mechanics and molecular design, 2024. https: http:///xxxvvrrr/e/// abs/24022 07148.
References
Aleem, S., Dietlmeier, J., Arazo, E., and Little, S. Convlora and adabn based domain adaptation via self-training. In
IEEE International Symposium on Biomedical Imaging, ISBI 2024, Athens, Greece, May 27-30, 2024, pp. 1-5. IEEE, 2024. doi: D0.1009//SB11557722200000666333. URL https: / /ttpi/iiii//// http:p///ee...ogg//...////////// 2024. 10635661.
Burns, K. J., Vasil, G. M., Oishi, J. S., Lecoanet, D., and Brown, B. P. Dedalus: A flexible framework for nu- merical simulations with spectral methods. Physical Review Research, 2(2)023366,, 2020. doi: 10.1103/ PhysRewweesearch2220222
URL https: /do..
Anderson, D., Tannehill, J. C., Pletcher, R. H., Munipalli,
023068.
9

---

PDE-Transformer: Efficient and Versatile Transformers for Physics Simulations
Cao, S. Choose transformer: Fourier Galerkin. In Advances in Neural Information Processing Systems (NeurIPS 2021), volume 34, 2021. URL https:// openreview www.aeview.net/oruueee beterummdd=ssooe id=ssohLcmn4-E
N. An image is worth 16x16 words: Transformers for image recognition scale. In International Conference Learning Representations, 2021. URL https:// openreview net. wtwpr//iwwe.eeeeoooreriiiee.e.. httt://ormmmi--iicaaaaaaa
Chen, T., Korbblith, S., Norouzi, M., and Hinton, G. E. A simple framework for contrastive learning of visual representations. In Proceedings of the 37th International Conference on Machine Learning. ICML 2020, 13-18 July 2020, Viriual Event, volume 119 of Proceedings of Machine Learning Research, pp. 1597-1607. PMLR, 2020. URL http::/// htttt::/proceeeee pttppeeeeeengggsseeeeeeeppss..se malr.press v119/chen20 html.
Dresdner, G., Kochkov, D., Norgaard. P. C., Zepeda-Nunez. L., Smith, J., Brenner, M., and Hoyer, S. Learn- ing to correct spectral methods for simulating turbu- lent flows. Transactions Machine Learning Re- search, 2023. ISSN 2835-8856. URL https:// openreview net wtt.:/ooommmdddBAGeoooll
Duraisamy, K. laccarino, G., and Xiao, H. Tur- bulence modeling in the age of data. Anmual Review of Fluid Mechanics, 51(1):357-377, jan 2019. doi: 10.1 10011////muuc-ttii-0001111----00 URL
Dehghani, M., Djolonga, J., Mustafa, B., Padlewski, P., Heek, J., Gilmer, J., Steiner, A. P., Caron, M., Geirhos, R., Alabdulmohsin, I., Jenatton, R., Beyer, L., Tschan- nen, M., Arnab, A., Wang, X., Ruiz, C. R., Minderer, M.. Puigeerver, J., Evei, U.. Kumar, M., van Steenkiste, S., Elsayed, G. F., Mahendran, A., Yu, F., Oliver, A., Huot, F. Bastings, J., Collier, M., Gritsenko, A. A., Birodkar, V., Vasconcelos, C. N., Tay, Y., Mensink, T., Kolesnikov, A., Pavetic, F., Tran, D., Kipf, T., Lucic, M., Zhai, X., Keysers, D., Harmsen, J. J., and Houlsby, N. Scaling vision transformers to 22 billion parameters. In Krause, A., Brunskill, E., Cho, K., Engelhardt, B., Sabato, S., and Scarlett, J. (eds.), International Confer- ence Machine Learning, ICML 2023, 23-29 July 2023, Honolulu, Hawaii, USA, volume 202 of Proceedings of Machine Learning Research, pp. 7480-7512. PMLR, 2023. htthtts:: www.hhttts./proceeeeieee /ttpp//rceeengssse mal.prress v202 2212//dehhhai... 23a html.
11446
Esser, P., Kulal, S., Blattmann, A., Entezari, R., Müller, J., Saini, H., Levi, Y., Lorenz, D., Sauer, A., Boesel, F,, Podell, D., Dockhorn, T., English, Z., and Rombach, Scaling rectified flow transformers for high-resolution image synthesis. In Forty-first International Conference Machine Learning. ICML 2024, Vienna, Austria, July 21-27, 2024. OpenReview.net, 2024. URL https: openreview net forum? wwwfooommmddoPuuuseeee. id-FFnUhsSJUS..
Falcon, W. and The PyTorch Lightning team. PyTorch Light- ning. March 2019. URL https://giit: com/ Lightning-At/1ightnn
Geneva, N. and Zabaras, N. Quantifying model form uncertainty in reynolds-averaged turbulence models with bayesian deep neural networks. Journal of Computational Physics, 383:125-147,, 2019. ISSN 0021-9991. doi: https://oo..og//000000.......p.. 021. O2L.L//LLttps////ooceeeeeddddddd cciineee/rticcee//////2229999999
Devlin, J., Chang. M., Lee, K., and Toutanova, K. BERT: pre-training of deep bidirectional transformers for lan- guage understanding. In Burstein, J., Doran, C., and Solorio, T. (eds.), Proceedings of the 2019 Conference of the North American Chapter of the Association for Com- putational Linguistics: Human Language Technologies, NAACL-HLT 2019, Minneapolis, MN, USA, June 2-7, 2019, Volume (Long and Short Papers), pp. 4171-4186. Association for Computational Linguistics, 2019. doi: 10. 11885// URL https://doi..oo. 10.18653 1wt.///////11111/1 1423.
Goswami, S., Kontolati, K., Shields, M. D., and Karniadakis, G. E. Deep transfer operator learning for partial differ- ential equations under conditional shift. Nature Machine Intelligence, 4(12):1155-1164 2022.
Dhariwal, P. and Nichol, A. Q. Diffusion models beat gans on image synthesis. In Ranzato, M., Beygelzimer, A., Dauphin, Y. N., Liang, P., and Vaughan, J. W. (eds.), Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurlPS 2021, December 6-14, 2021, vir- tual, pp. 8780-8794 2021.
Goyal, P., Dollár, P., Girshick, R. B., Noordhuis, P. Wesolowski, L., Kyrola, A., Tulloch, A., Jia, Y., and He, K. Accurate, large minibatch SGD: training ima- genet 1 hour. CoRR, abs/1706.02677 2017. URL 02677.
Gupta, D., Bhatti, A., and Parmar, S. Beyond lora: Ex- ploring efficient fine-tuning techniques for time series foundational models, 2024a. URL https: //arxiv. heg/abs/2400 11302.
Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., Uszkoreit, J., and Houlsby,
10

---

Hu, E. J., Shen, Y., Wallis, P., Allen-Zhu, Z., Li, Y., Wang, S., Wang, L., and Chen, W. Lora: Low-rank adapta- tion of large language models. In The Tenth Interna- tional Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net PDE-Transformer: Efficient and Versatile Transformers for Physics Simulations
Gupta, D., Bhatti, A., Parmar, S., Dan, C., Liu, Y., Shen, B., and Lee, S. Low-rank adaptation of time series foundational models for out-of-domain modality fore- casting. In Hung, H., Oertel, C., Soleymani, M., Chas- pari, T., Dibeklioglu, H., Shukla, J., and Truong, K. P. (eds.), Proceedings of the 26th International Confer- ence on Multimodal Interaction, ICMI 2024, San Jose, Costa Rica, November 4-8. 2024, pp. 382-386. ACM, 2024b. doi: 10.1 (0.11445/667899573668777777 URL https: //doi. www./r/.......... 10.
Honolula, Hawaii, USA, volume 202 of Proceedings of Machine Learning Research, pp. 15213-13232 PMLR, mlr mmr.ppu prrsss/
2023. httphttss://pprcceedliee v202
2022. URL https: http::://ppppevveeeeeeee net forum? id-nZeVKeeFYE5,
Gupta, J. K. and Brandstetter, spatiotemporal generalized PDE modeling. Trans. Mach. Learn. Res., 2023, 2023. URL https:// openreview. net/ forum? ddddPSTDDGEE...
Towards multi-
Jacobsen, C.. Zhuang, Y., and Duraisamy, K. Coco- gen: Physically consistent and conditioned score-based generative models for forward and inverse problems. SIAM J. Sci. Comput., 47(2):399, 2025. doi: 10.1137/ 24M1636071. URL /ttp::://///o//e http:////oooooog///////o 24m1 636071.
He, K., Chen, X., Xie, s., Li, Y., Dollár, P., and Gir- shick, R. B. Masked autoencoders are scalable vision learners. In IEEE/CVF Conference on Computer Vision and Puttern Recognition, CVPR 2022, New Orleans, LA,, USA, June 18-24, 2022, pp. 11998. IEEE, 2022. doi: 10. (0.11099CVPS2228822222222 202200055, URL httpts:: doi. org. 10 1109/ CVPR52688 2022. 01553.
Kim, B., Azevedo, V. C., Thuerey, N., Kim, T., Gross, M., and Solenthaler, B. Deep Fluids: A Generative Network for Parameterized Fluid Simulations. Comput. Graph. Fo- rum, 38,2):12, URL http::/wwa/tt httt::::/ww.bbungsoo byungsoo. fluidss.
Herde, M., Raonic, B., Rohner, T., Käppeli, R., Molinaro, R., de Bézenac, E.,, and Mishra, S. Poseidon: Efficient foundation models for pdes. In Globersons, A., Mackey, L., Belgrave, D., Fan, A., Paquet, U., Tomezak, J. M., and Zhang, C. (eds.), Advances in Neural Information Processing Systems 38: Annual Conference on Neural Information Processing Systems 2024, NeurlPS 2024, Vancouver, BC, Canada, December 10 15, 2024, 2024.
Kochkov, D., Smith, J. A., Alieva, A. Wang, Q. Brenner, M. P. and Hoyer, S. Machine learn- ing anc-acceler computational fluid dynamics. Pro- ceedings of the National Academy of Sciences, 11801)):c10188411, 2021. doi: 10.1073/pnas. 2101784118. URL https: ///////ww /www. www..ppas.c... ong. 2101784118.
Ho, J., Kalchbremner, N., Weissenborn, D., and Salimans, T. Axial attention in multidimensional transformers. CoRR,
Koehler, F., Niedermayr, S., Westermann, R., and Thuerey, N. Apebench: A benchmark for autoregressive neural emulators of pdes. In Advances in Neural Information Processing Systems 37 (Datasets and Benchmarks Track), 2024. URL https: / openreview net// forum?
Abs/ 2019. URL abs/1912. 12180.
Ho, J., Jain, A., and Abbeel, P. Denoising diffusion proba- bilistic models. In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M., and Lin, H. (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 6840- 6851. Curran Associates, Inc., 2020.
id=1Mc0qE116.
Kohl, G., Chen, L.-W., and Thuerey, N. Benchmarking autoregressive conditional diffusion models for turbulent flow simulation, 2024. URL https: /ttp:///axivvoooo abs / 0174..
0175
Holzschuh, B. J., Vegetti, S., and Thuerey, N. Solving inverse physics problems with score matching. In Oh, A., Naumann, T., Globerson, A., Saenko, K., Hardt, M., and Levine, S. (eds.), Advances in Neural Information Neural
Kossaifi, J., Kovachki, N., Li, Z., Pitt, D., Liu-Schiaffini, M., George, R. J., Bonev, B., Azizzadenesheli, K., Berner, J., and Anandkumar, A. library for learning neural operators, 2024. URL https: https:t/arxxii..etter httpps///axivv.o////////
Processing Systems 36: Annual Conference Information Processing Systems 2023, NeurIPS 2023, New Orleans, IA, USA, December 10 16, 2023, 2023.
Kovachki, N. B., Li, Z., Liu, B., Azizzadenesheli, K., Bhat- tacharya, K., Stuart, A. M., and Anandkumar, A. Neu- ral operator: Learning maps between function spaces with applications to pdes. J. Mach. Learn. Res., 24:89:1- 89:97, 2023. URL https: : /jirr. wwt./ooppppp.....o 122421-1224 html.
Hoogeboom, E., Heek, J., and Salimans, T. simple diffu- sion: End-to-end difffusion for high resolution images. In Krause, A., Brunskill, E., Cho, K., Engelhardt, B., Sabato, S., and Scarlett, J. (eds.), International Confer- ence on Machine Learning, ICML 2023, 23-29 July 2023,

---

PDE-Transformer: Efficient and Versatile Transformers for Physics Simulations
Li, Z., Kovachki, N. B., Azizzadenesheli, K., Liu, B., Bhat- tacharya, K., Stuart, A. M., and Anandkumar, A. Neu- ral operator: Graph kernel network for partial differ- ential equations. CoRR, abs/2003.003855 2020. URL https: / www.xi.vvoor. www.v.orr//bs.... abs/ 2003. 03485.
2933, 2024. doi: 10.2514/1.1063440 URL https: / //doi.org/10. 2tte:/1..2514/// J063440.
Liu, X., Gong, C., and Liu, Q. Flow straight and fast Learning generate and transfer data with rectified flow. In The Eleventh International Conference Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023. URL https: / //openrevier net/ forum? (asrrummiiix/// d= d=X
Li, Z., Zheng, H., Kovachki, N., Jin, D., Chen, H., Liu, B., Azizzadenesheli, K., and Anandkumar, Physics- informed neural operator for learning partial differential equations. arXiv preprint arXiv: 000/2/11.03777 1.03794, 2021.
Liu, Z., Lin, Y., Cao. Y., Hu, H., Wei, Y., Zhang, Z., Lin, S., and Guo, B. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), 2021.
Li, Z., Meidani, K., and Farimani, A. B. Transformer for partial differential equations' operator learning. Transac- fions on Machine Learning Research, 2023a. ISSN 2835- 8856. URL https: BttKLhttt///pe.ateeee http://pppeeeeeeeeeeeeee n forum? id=EPPqt 3UERT.
Liu, Z.,, Hu, H., Lin, Y., Yao, Z.,, Xie, Z.., Wei, Y., Ning. J., Cao, Y., Zhang, Z., Dong, L., Wei, F,, and Guo, B. Swin transformer V2: scaling up capacity and resolution. In IEEF/CVF Conference on Computer Vision and Pattern Recognition, pp. 11999-12009, IEEE, 2022. doi: 10. 1109/CVPR52688. 2022.01 170. URL https: /doi. org/10. 1tt./111////2//
Li, Z., Shu, D., and Farimani, A. B. Scalable transformer for PDE surrogate modeling. In Oh, A., Naumann, T., Globerson, A., Saenko, K., Hardt, M., and Levine, S. (eds.), Advances in Neural Information Processing Sys- tems 36: Annual Conference on Neural Information Pro- cessing Systems 2023, NeurlPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023b.
2022.0022200
Lu, L., Jin, P., Pang, G., Zhang, Z., and Karniadakis, G. E. Learning nonlinear operators via DeepONet based on the universal approximation theorem of operators. Nature Machine Intelligence, 3(3):218 22222,92222111 2021.
Liang, Y. Wen, H., Nie, Y.. Jiang, Y., Jin, M., Song, D., Pan, S., and Wen, Q. Foundation models for time ries analysis: A tutorial and survey. In Proceedings of the 30th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, KDD '24, pp. 6555-6565, New York. NY, USA, 2024. Association for Comput- ing Machinery. ISBN 9798400704901. doi: 10.1145/
Luo, S., Tan, Y., Patil, S., Gu, D., von Platen, P., Passos, A., Huang, L.,, Li, J., and Zhao, H. Lem-lora: uni- versal stable-diffusion acceleration module, 2023. URL org. https: http:s/://axxi...o//aaa///////// httpg//bbb//1111 05556.
3637528.3671451. URL https: /ttp:://ddo..oo org/10. 1145/3637522 211.13677722.66645..11
Maddu, S., Weady, S., and Shelley, M. J. Learning fast, accurate, and stable closures of a kinetic theory of an active fluid. Journal of Computational Physics, 504: 1 112869, 2024. URL https: /ttp://///.oo.....11111 1016/ 3. j.j.p.220 202pp.202..... 2024 112869.
Lin, L., Fan, H., Zhang, 7.,, Wang, Y., Xu, Y., and Ling, H. Tracking meets lora: Faster training, larger model, stronger performance. In Leonardis, A., Ricci, E., Roth, S., Russakovsky, O., Sattler, T., and Varol, G. (eds.), Computer Vision ECCV 2024 18th Euro- pean Conference, Milan, Italy, September 29-0ctober 4, 2024, Proceedings, Part I, volume 15059 of Lecture Notes in Computer Science, pp. 300-318. Springer, 2024. doi: 10.100779988-3333113222222---2 doi ott..rrg/11. 1007 --27778----31-7332---- 4_13
Mao, Y., Ge, Y., Fan, Y., Xu, W., Mi, Y., Hu, Z., and Gao, Y. A survey on lora of large language models. Frontiers of Computer Science, 1977)119605,, 2025.
Maurício, J., Domingues, I., and Bernardino, J. Com- paring vision transformers and convolutional neural net- works for image classification: A literature review. Ap- plied Sciences, 13(9). 2023. ISSN 2076-3417. doi: 113390/app130955222 URL https: http//wwwwmpppppo
URL https:
Lipman, Y., Chen, R. T. Q., Ben-Hamu, H., Nickel, M., and Le, M. Flow matching for generative modeling. In The Eleventh International Conference on Learning Repre- sentations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023. URL https::/p///. cppnneview. net/ htt.foorrun....PpyyyccCC......
McCabe, M., Blancard, B. R.-S., Parker, L., Ohana, R. Cranmer, M., Bietti, A., Eickenberg, M., Golkar, S., Krawezik, G., Lanusse, F., Pettee, M., Tesileanu, T., Cho, K., and Ho, S. Multiple physics pretraining for physical surrogate models. In NeurlPS 2023 Al for Sci- ence Workshop, 2023. URL https: http///ppp/rre.eee/////e//. net/forum id=M1 200=M111m@@@@mm 1mmkuxx.
Liu, Q. and Thuerey, N. Uncertainty-awart surrogate models for airfoil flow simulations with denoising dif- fusion probabilistic models. AIAA Journal, 62(8):2912-
12

---

Peebles, W. and Xie, S. Scalable diffusion models with transformers. In 2023 IEEE/CVF International Confer- ence on Computer Vision (ICCV), pp. 4172-41822 2023. doi: 10.1 10.100//CCV51000.2023.0000 PDE-Transforme:: Efficient and Versatile Transformers for Physics Simulations
Ohana, R., McCabe, M., Meyer, L.. T., Morel, R., Agocs, F. J., Beneitez, M., Berger, M., Burkhart, B., Dalziel, S. B., Fielding, D. B., Fortunato, D., Goldberg, J. A., Hi- rashima, K., Jiang, Y.-F.. Kerswell, R., Maddu, S., Miller, J. M., Mukhopadhyay, P., Nixon, S. S., Shen, J., Wat- teaux, R., Blancard, B. R.-S., Rozet, F,, Parker, L. H., Cranmer, M., and Ho, The well: a large-scale collec- tion of diverse physics simulations for machine learning. In Advances in Neural Information Processing Systems 37 (Datasets and Benchmarks Track), 2024. URL ht https: // //openreview net/ forum? d=00sx577BBT.
Rodrigo, M., Cuevas, C., and García, N. Comprehensive comparison between vision transformers and convolu- tional neural networks for face recognition tasks. Scien- rific reports, 14(1):21392, 2024.
Rombach, R., Blattmann, A., Lorenz, D., Esser, P., and Ommer, B. High-resolution image synthesis with latent diffusion models. In IEEE/CVF Conference Computer Vision and Pattern Recognition, CVPR 2022, New Or- leans, LA, USA, June 18-24, 2022, pp. 10674-10685. IEEE, 2022. doi: 10.1 10.11099CVPR556688.. 202.2220060..2200444.40 URL https:://d//oo http:///oo..or//..1111/ 2022.01043
Pearson, E. Complex patterns in simple system. Sci- ence, (5118))119111118 1993. doi: 10.1 10/ss2e/s 261.5118. 189. URL https://wwwww/ http::::://ww.scieeee.e..// doi/abs/ 1111/1//ieeeni 261. 11118.1999
wwwenceeeec....ccc org/
Ronneberger, O., Fischer, P., and Brox, T. U-net: Con- volutional networks for biomedical image segmentation. In Medical Image Computing and Computer-Assiste Intervention-MiCCAI 2015: 18th International Confer- ence, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18, pp. 234 241. Springer, 2015.
Sirignano, and MacArt, J. F. Deep learning closure models for large-eddy simulation of flows around bluff bodies. Journal of Fluid Mechanics, 966, jul 2023. doi: 10.1017/ jfim 2em..222.4446. URL https: //oi: org 10.00000000 1017% 2F j Em. 2023. 446.
Perez, E., Strub, F., de Vries, H., Dumoulin, V., and Courville, A. C. Film: Visual reasoning with general conditioning layer. In Mcllraith, S. A. and Weinberger, K. Q. (eds.), Proceedings of the Shcooy- AAAI Con- ference on Artificial Intelligence, (AAAI-18), the 30th in- novative Applications of Artificial Intelligence (IAAI-18), and the 8th AAAI Symposium on Educational Advances in Artificial Intelligence (EAAI-18), New Orleans, Louisiana, USA, February 2-7, 2018, pp. 3942 3951. AAAI Press, 2018. URL https: /woi//oiio///// orgo /00 1609//aai..... v1671.11671.
Smith, G. D. Numerical solution of partial differential equations: finite difference methods. Oxford university press, 1985.
Smith, J. S., Hsu, Y., Zhang, L., Hua, T. Kira, Z. Shen, Y., and Jin, H. Continual diffusion: Contin- ual customization of text-to-image diffusion with c-lora. Trans. Mach. Learn. Res., 2024, 2024. URL https: //oppenrevee.rree.. www.:neee/frmmmmeme
Radford, A., Narasimhan, K., Salimans, T., Sutskever, I., et al. Improving language understanding by generative pre-training. OpenAl blog, 2018
Subramanian, S., Harrington, P., Keutzer, K., Bhimji, W., Morozov, D., Mahoney, M. W., and Gholami, A. To- wards foundation models for scientific machine learning: Characterizing scaling and transfer behavior. In Oh, A., Naumann, T., Globerson, A., Saenko, K., Hardt, M., and Levine, S. (eds.), Advances in Neural Information Pro- cessing Systems 36: Annual Conference on Neural In- formation Processing Systems 2023, NeurlPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023.
Radford, A., Wu, J., Child, R. Luan, D., Amodei, D., Sutskever, I., et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Raissi, M., Perdikaris, P., and Karniadakis, G. E. Physics- informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational physics, 378:686-707, 2019.
Sun. L., Gao, H., Pan, S., and Wang. J.-X. Surro- gate modeling for fluid flows based on physics- constrained deep learning without simulation data. Computer Methods in Applied Mechanics and En- gineering. 361: 12732, 2020. ISSN 0045-7825, doi: URL
Ren, S., Padilla, W. and Malof, J. M. Benchmarking deep in- verse models time, and the neural-adjoint method. In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M., and Lin, H. (eds.), Advances in Neural Information Process- ing Systems 33: Annual Conference on Neural Informa- tion Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020.
13

---

PDE-Transformer: Efficient and Versatile Transformers for Physics Simulations
Takahashi, S., Sakaguchi, Y., Kouno, N., Takasawa, K., Ishizu, K., Akagi, Y., Aoyama, R., Teraya, N., Bolatkan, A., Shinkai, N., et al. Comparison of vision transformers and convolutional neural networks in medical image anal- ysis: a systematic review. Journal of Medical Systems, 48(1):84, 2024.
IEEE/CVF Winter Conference on Applications of Com-
puder Vision, WACV 2025, Tucson, AZ, USA, February 26 March 6, 2025, pp. 6634-6645. IEEE, 2025, doi: 10. 10.11999AACV11041.225..0000 doi.org/10.11000 wwt.:11199////VV11100 041 2025. 00646.
URL https://
Wistuba, M., Sivaprasad. P. T., Balles, L., and Zappella, G. Continual learning with low rank adaptation, 2023. URL https: / aawxiv.iiii..//..i//...aa/aaaava org/ abs 231..11610000
Takamoto, M., Alesiani, F., and Niepert, M. Learning neu- ral PDE solvers with parameter- guided channel attention. In Krause, A., Brunskill, E., Cho, K., Engelhardt, B., Sabato, S., and Scarlett, J. (eds.), International Confer- ence on Machine Learning, ICML 2023, 23-29 July 2023, Honolulu. Hawaii, USA, volume 202 of Proceedings of PMLR,
Wu, H.,, Luo, H., Wang, H., Wang. J., and Long, M. Tran- solver: A fast transformer solver for pdes on general geometries. In Forty-first International Conference on Machine Learning, ICML 2024, Vienna, Austria, July 21-27, 2024. OpenReview.net, 2024a. URL https: / // /openreview net/ forum? d=Ywl 6pDDX
Machine Learning Research, pp. 2023. URL httpttss: http://pooceeeenneneeeee wal...ppes///// v202 2.02/takamcto22333 html.
Thuerey, N., Holzschuh, B., Holl, P., Kohl, G., Lino, M., Liu, Q. Schnell, P., and Trost, F. Physics-based Deep Learning. www, 2021. URL https:// chysicshaseddeepleanningg org.
Wu, P., Li, K., Wang, T., Dong, Y., Leung, V. C. M., and Wang, F. Fedfmsl: Federated learning of foundation models with sparsely activated lora. IEEE Transactions Mobile Computing, 23(12):11677-511)) 2024b. doi: 10.1 N0.11997TMC.222434446644
Tian, Y., Tu, Z., Chen, H., Hu, J., Xu, C., and Wang, Y. U-dits: Downsample tokens in u-shaped diffusion trans- formers. In The Thirty-eighth Annual Conference on Neu- ral Information Processing Systems, 2024. URL https: ///ppnreeiiw... net/ /orte/foruue forum? ? d=
Xu, W., Lu, Y., and Wang, L. Transfer learning enhanced deeponet for long-time prediction of evolution equations. In Proceedings of the AAAI Conference on Artificial In- telligence, volume 37, pp. 10629- 10629, 2023.
Um, K., Brand, R., Fei, Y. R., Holl, P., and Thuerey, N. Solver-in-the-loop: Learning from differentiable physics to interact with iterative pde-solvers. In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M., and Lin, H. (eds.), Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6- 2020, vir- tual, 2020.
Yang, L., Liu, S., Meng, T., and Osher, S. J. In-context operator learning with data prompts for differential equa- tion problems. Proceedings of the National Academy of Sciences, 120(39):c230114222) 2023a. doi: 10.1073/ pnas. 231012210.,., URL https: / www.//wwwwwoo doi/abst 10030.....1111
www.pnws......po.wwmmw.m...m.... pnas org/ 2310142120.
Yang, M., Chen, J., Zhang, Y., Liu, J., Zhang, J., Ma, Q. Verma, H., Zhang, Q., Zhou, M., King, I., and Ying, R. Low-rank adaptation for foundation models: A compre- hensive review, 2024. URL https: http://axxvv.orro abs/2501.0038
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I. Attention is all you need, In Proceedings of the 31st International Conference on Neural Information Processing Systems, NIPS' 17, pp. 6000-6010, Red Hook, NY, USA, 2017. Curran Associates Inc. ISBN 9781510860964.
Yang, S., Nachum, O., Du, Y., Wei, J., Abbeel, P., and Schuurmans, D. Foundation models for decision mak- ing: Problems, methods, and opportunities, 2023b. URL https:
httppp://aaxivvooo http///xivvvv./////////////.aav/ wttpe//w..o/////o///ooooo///.oo. http://www...//////3/.....1..... 04129.
Wang, H.. Planas, R., Chandramowlishwaraan A., and Bostanabad, R. Mosaic flows: transferable deep learn ing framework for solving pdes on unseen domains. Com- puter Methods in Applied Mechanics and Engineering, 389: 114424, 2022.
Yuan, L., Chen, D., Chen, Y.LL., Codella, N., Dai, X., Gao, J.. Hu, H., Huang, X., Li, B., Li, C., Liu, C., Liu, M., Liu, Z., Lu, Y., Shi, Y., Wang, L.,, Wang. J., Xiao, B., Xiao, Z., Yang, J., Zeng, M., Zhou, L., and Zhang, P. Florence: new foundation model for computer vision, 2021. URL hbsp///bs//111 11432.
Wang, S. Li, 7., Khabsa, M., Fang, H., and Ma, H. Linformer: Self-attention with linear complexity. CORR, abs/2006.04768, 2020. URL https:://arxiv..tt abs/2006. 04768.
Zeng, S., Wang, D., Jiang, L., and Xu, D. Parameter- efficient fine-tuning on large protein language models improves signal peptide prediction. Genome research, 34 19):1445-1454, 2024.
Wei, X., Li, G., and Marculescu, R. Online-lora: Trsk-free online continual learning via low rank adaptation.
In
14

---

PDE-Transformer: Efficient and Versatile Transformers for Physics Simulations
Zeng, Y. and Lee, K. The expressive power of low-rank adaptation. In The Twelfth International Conference Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net, 2024. URL https: //openrevies ew. net forum? i d=1 1 km3
Zhang, B. and Sennrich, R. Root mean square layer normal- ization. In Wallach, H. M., Larochelle, H., Beygelzimer, A., d' Alché-Buc F., Fox, E. B., and Garnett, R. (eds.), Advances in Neural Information Processing Systems 32: Anmual Conference on Neural Information Processing Systems 2019, NeurlPS 2019, December 8-14, 2019, Van- couver, BC, Canada, pp. 12360-12371 2019.
Zhang, P., Laiu, M. P., Norman, M., Stefanski, D., and Gounley, J. MATEY: multiscale adaptive foundation mod- cls for spatiotemporal physical systems. arXiv preprint arX1v:2412.20601, 2024a.
Zhang, Z., Xu, J., Liu, P. and Hu, R. Fed-piLot: Opti- mizing lora assignment for efficient federated founda- tion model fine-tuning, 2024b. URL https: / //axxiv. org/abs// abs/ 2410. 10200.
15

---

PDE-Transformer: Efficient and Versatile Transformers for Physics Simulations
A. Training configurations
A.I. Training Hyperrparameters
Table 4 summarizes the key hyperparameters used during training. We employ the Distributed Data Parallel (DDP) strategy, as supported by PyTorch Lightning (Falcon & The PyTorch Lightning team, 2019), to train models across multiple GPUs. To manage memory consumption effectively and maintain a consistent effective batch size across different hardware setups, incorporate the gradient accumulation technique. Additionally, mixed- mixxd-precision training is enabled for most models, except for the FNO and OFormer models in the downstream tasks. The exclusion is due to the inherent incompatibility of fast Fourier transformations within these architectures with the mixed-precision setup.
Table 4. Major hyperparameters for training.
Pre-training
Downstream tasks
Effective batch size
256
256
1.00 10 10-4 (Active matter & RBC)
Learning rate
4.00 10-5
4.00 10 (Shear flow)
Optimizer Epochs
Adam W
AdamW
100
2000
A.2. EMA Gradient Clipping
In this study, we employ an exponential moving average (EMA) of the gradient norm to stabilize training. We compute EMA values for the gradient norm using two different coefficients, as outlined in Algorithm 1. The EMA with larger coefficient places greater emphasis on historical values and serves as the clipping threshold. Meanwhile, the EMA with smaller coefficient acts as the target value for gradient clipping. Notably, the clipping coefficient must be set greater than I to ensure that 91 effectively tracks changes in the gradient norm. Compared to traditional gradient clipping, which applies a fixed threshold and clipping value, EMA gradient clipping offers a more flexible approach by dynamically adjusting the threshold and clipping value. This adaptability is particularly ad udvataggcous, as gradient norms vary significantly across different model sizes and training stages.
Algorithm 1 EMA Gradinet Clip.
Input: first EMA coefficient B1, second EMA coefficient 82, Clip threshold coefficient a, Clip value coefficient k. Initialize 31 0.99, 32 = 0.999, 2, k 1.1, = 0, 0, 92 = 0. repeat
Get gradient g from training step if i! = and g
> a
92
then
g kg 1-
end if
91 = 3191 + (1-8))g)
92 = $192 + (1-8)
= i+1 1
until training is finished
A.3. Model Details
Table. 5 summarizes the sizes of the network in different models. The size definition of scOT models follows the definition in the Poseidon foundation model (Herde et al., 2024). Meanwhile, we rescale the size of FNO by set- m.mooess..ommmm width-16, 6, and hi aeedennchannee--ie
in the official implementation (Kossaifi et al., 2024). For the Galerkin Transformer, we also change the n. n_hidden=96, num num_encode: saroddrrllayrs=5,, dim_feedforwarde 128, and freq_dim=433 432 in the official implementation (Cao, 2021) to extend the network size. We noticed that the OFormer requires much more memory during the training compared with other models. We have to modify
ting the n. modess wwwehhhtttte
16

---

PDE-Transformer: Efficient and Versatile Transformers for Physics Simulations
Table 5. Size of the network of different models.
Models
Number of trainable parameters
PDE-S
46.57M
scOT-S
39.90M
FNO
42.72M
Galerkin Transformer
38.82M
OFormer
0.12M
PDE-B pre-trained
178.97M
scOT-B pre-trained
152.70M
Table 6. Architecture hyperparameters PDE- Transformer.
Name of hyperparameter
Value
window.size depth num_heads mlp-ratio
12, 4, 4, 6, 4, 4, 2]
16
4.0
c class.dropott.prob
0.1 True
qkv bias activation
GELU
the in..em_.de dim=244 and out-seq.ent lile-eemb.ce www.eeemmimiee...e ddim=48 of the original implementation (Li et al., 2023a) to make it have similar memory consumption as other models, although the network size is in different magnitude. We list hyperparameter for the architecture of PDE-Transformer in Table 6.
B. Details and Additional Results
B.1. Training Methodology
Details of the convergence of the training loss for different configurations of PDE Transformer be found in Figure 7. The plots show that B and L. configurations yield consistently lower training losses for both supervised and diffusion based training methodologies. Additionally, the left plot shows the stable behavior of the diffusion training with reduced loss fluctuations across epochs.
17

---

PDE-Transformer: Efficient and Versatile Transformers for Physics Simulations
S
Supervised MSE
B
Diffusion Loss
B
L
10-2
10
10-3
10
20
40
60
80
100
20
40
60
80
100
epoch
epoch
Figure 7. Training loss of supervised (left) and diffusion training (right) for configurations S, B and L.
Figure shows the nRMSE over rollout steps for the PDE Transformer L configuration when using flow matching. After using than 20 flow integration steps the nRMSE values begin to stabilize at lower level in terms of mean errors. Increasing the number of steps gives persistent improvements.
18

---

PDE-Transformer: Efficient and Versatile Transformers for Physics Simulations
100
ARMSE
5
10
20
10-1
30
40
100
0.0
2.5
5.0
7.5
10.0 Time
12.5
15.0
17.5
20.0
Figure evaluate the nRMSE against number inference steps for PDE- TDE-Transformer configuration use explicit Euler method to solve the ODE for sampling.
B.2. Downstream Tasks
Details of the comparison with the pre-trained Poseidon model scOT-B over time are shown in Figure 9. Both scOT-B and PDE-Transformer B have a similar size, show comparable performance during the very first frames of the trajectory. However, the prediction error of scOT-B accumulates rapidly over time, while our PDE-B produces stable trajectories with persistently low errors despite the challenging task. This is reflected in the snapshot for a case at shown on the left of Figure 9. While scOT-B has diverged to state dominated by oscillations, the PDE-B solution stays close to the reference solution.
19

---

PDE-Transformer: Efficient and Versatile Transformers for Physics Simulations
Reference (t=8)
ScOT-B-pretrained
Ours-B-pretrained
Active matter
8
6
nRMSE
SCoT-B-Pretrained
Durs---Pretrained
4
2
0
5
10 Time
15
20
Figure 9. Average trajectory nRMSE of PDE-Trame Tranffoosf and scOT with pro-trained weights active matter tasks. A reprexentative
frame at f = 8 shown left.
B.3. Numerical Values
Table 7 shows the exact numerical values of the nRMSE value in Figure 6. addition, Figure 10 illustrates the nRMSE changes during the rollout for each downstream task. highlights that the pre-trained PDE-Transformer consistently outperforms baselines over the course of the full rollout.
B.4. Fine-tuning with LoRA
Low-Rank Adaptation (LoRA) (Hu et al., 2022) has recently emerged as a popular technique for addressing the computational challenges associated with fine-tuning foundation models (Zhang et al., 2024b; Wu al., 2024b; Yang al., 2024). Like many other parameter-efficient fine-tuning (PEFT) methodologies. LoRA significantly reduces the number of trainable parameters during fine-tuning, achieving this with minimal performance degradation. Building its early success in fine-tuning large language models (Hu et al., 2022; Mao et al., 2025) and diffusion- image models (Luo et al., 2023; Smith et al., 2024), LoRA has gained considerable traction in other domains. These include computer vision (Alcem et al., 2024; Lin et al., 2024), continual learning (Wistuba et al., 2023; Wei et al., 2025), and time-series prediction (Gupta et al., 2024b;a). Within the scientific deep-learning community, LoRA has also been applied to protein and materials engineering
Table 7. Average rollout nRMSE of different models on downstream tasks.
Active Matter
RBC
Shear flow
FNO
0.509
0.269
0.274
OFormer
0.752
0.520
0.780
ScOT-s
0.546
Galerkin
0.638
PDE-S
0.494
0.155
0.199
PDE-S pre-trained
0.455
0.104
0.125
20

---

PDE- Pra-Transforme:: Efficient and Versatile Transformers for Physics Simulations
SCOT
FNO
OFormer Ours-5
Galerkin
Opr----prerrrined
3.0
matter
Rayleigh convection
1.2
2.0
RAMSE
1.5
0.6
1.0
0.4
0.5
0.2
0.0
0.0
2.5
5.0
7.5
10.0
12.5
15.0
17.5
20.0
0.0
2.5
5.0
7.5
10.0
12.5
15.0
17.5
20.0
1.50
Shear flow
1.25
1.00
0.50
0.25
12.5
20.0
Figure 10. Average rollout aRMSE of different models on downstream tasks.
(Buchler Buchler, 2024; Zeng et al., 2024). Despite these advancements, the application of LoRA in foundation models for partial differential equations (PDEs) remains relatively unexplored. Given the input of network and the original pre-trined weight Wn e Rddk LoRA introduces low-rank matrix A W to calculate the network output y:
Wox AWx Wox + " BAx.
(5)
Here, B E Rdxr, A E Rexk is the rank, and (a = r- in the current study) is the coefficient used to rescale the low-rank matrix. By setting r < dk/(d+ k) and freezing the pre-trained weight during the fine-tuning, the trainable parameter during fine-tuning can be significantly reduced. We enable LoRA for the weights of linear layers and convolutional layers in the network and keep other parameters, e.g., biases, unchanged with their full parameters. Fig. illustrates how the size of the network changes with different numbers of ranks. The size of PDE Transformer-S, -B, and -L. all grow linearly with more ranks. We evaluate the performance of LoRA the downstream tasks by training B-size PDB- Tranfrrasformmr with
400
100
350
Numbbr of trainablle perameeter (4)
Ratio of trainable parameterrs (%6
80
250
200
150
100
100
200
300
500
200
500
Rank
Rank
Figure 11. The number of trainable parameters of different sizes of PDE-Transformer w.r. different LoRA ranks.
21

---

PDE-Transformer: Efficient and Versatile Transformers for Physics Simulations
96, which satisfies the theoretical minimum rank required for transformer-based models (Zeng & Lee, 2024). The final fine-tuned network consists of 42.06M trainable parameters, size comparable to the S-size model. Figure 12 presents the average rollout performance of LoRA fine-tuning the B-size model. In the active matter case, the LoRA-B model outperforms the LoRA-S model without pre- pre- training, demonstrating advantages similar to those of the pre-trained S model. However, for the RBC task, the LoRA-B model performs comparably to the non-pre-trained model but falls short of the pre-trained model. These results indicate that while the LoRA fine-tuned model achieves performance similar to that of a model same size trained from scratch, can not always reach the level of fully pre-trained model with a similar size.
pretttaained
B-LoRA
1.4
1.2
0.4
0.2
0.0
2.5
5.0
7.5
10.0
12.5
15.0
17.5
20.0
Time
Rayleigh-Bénarr convection
0.35
0.30
0.20
0.15
0.10
0.05
2.5
5.0
7.5
10.0
12.5
15.0
17.5
20.0
Figure 12, Performance of the LoRA finetuning on downstream tasks.
C. Pre-training Datasets
The sources for our datasets are chosen to ensure wide range of different PDEs, at high spatial resolutions, and with varying physical quantities across the simulations in each dataset. To simulate linear, reaction-diffusion, and nonlinear PDEs, employ the Exponax solver (Koehler et al., 2024). It implements a range of Exponential Time Differencing Kutte-Kutta (ETDRK) methods for numerically solving different PDEs in an efficient and unified manner. Our choice against using the authors benchmark APEBench directly was intentional, to ensure that the datasets are higher in resolution and more diverse in terms of physical behavior across simulations rather than only changing initial conditions. The ETDRK methods operate in Fourier space and as such do not allow for non -on-perioi domains or complex boundaries. We always use the physical solver interface across data sets, rather than the non-dimensionalized difficult-based interface, to provide a simpler, unified way of handling physical quantities for the task embedding.
C.1. Linear PDEs
The Exponax solver framework features a wide range of Exponential Time Differencing Runge- Kutta (ETDRK) methods efficiently simulate different PDEs (Koehler al., 2024). The chosen linear PDEs are comparatively simple and have analytical solutions, but are nevertheless an important building block for more complex PDEs. Here, each linear PDE can be interpreted scalar quantity like density that is affected by different physical processes. Unless specified otherwise, sampling from intervals is always performed uniformly random below. Figure 13 shows example visualizations of every dataset described below.
In addition to varying physical parameters across the dataset, each simulation features randomized initial conditions. By default, the initial conditions are constructed in the following way: One of three initialization types with different spectral
22

---

PDE-Transformer: Efficient and Versatile Transformers for Physics Simulations
Table Overview of datasets simulated with Exponax (top, Koehler al., 2024) featuring linear PDEs, reaction-diffusion PDEs, and nonlinear PDEs. The dimensionality of each dataset is described via the number of simulations s. time steps t. fields (or channels) f. and spatial dimensions and y. addition to varying the specified quantities, the initial conditions of each simulation are different.
Dataset diff
t
Varied Quantities actoss
Test Set
600
30
2048
2048
viscosity (x, y)
8 [500, 600
fisher
600
30
2048
2048
diffusivity, reactivity
[500, 600
sh
600
30
2048
2048
reactivity, critical number
8 E [500, 600
gs-alpha
100
30
2
2048
2048
initial conditions only
separate: = 30, t = 100
gs-beta
100
30
2
2048
2048
initial conditions only
separate: = 30, L 100
gs-ganna
100
30
2
2048
2048
initial conditions only
separate: 30, L 100
gs-delta
100
30
2048
2048
initial conditions only
8 E [80, 100[
ge-epsilon
100
30
2
2048
2048
initial conditions only
separate: 30, t 100
gs-theta
100
30
2
2048
2048
initial conditions only
8 E 180, 100
g9-icta
100
30
2
2048
2048
initial conditions only
E 180, 100
ga-kappa
100
30
2
2048
2048
initial conditions only
E [80, 100
burgers
600
30
2
2048
2048
viscosity
E [500, 600
kdv
600
30
2
2048
2048
domain extent, viscosity
8 [500, 600
ks
600
30
2048
2048
domain extent
separate: 50, t 200
decay-turb
600
30
2048
2048
viscosity
separate: s 50, t 200
kolm-flow
600
30
1
2048
2048
viscosity
separate: = 50, / 200
energy distributions implemented by Exponax is chosen uniformly random: First, the random truncated Fourier series initializer layers multiple Fourier series additively up to a cutoff at certain frequency level. The cutoff threshold is chosen as uniformly random integer from [2, [1 Second, the Gaussian random field initializer creates power-law spectrum in Fourier space, i.e., the energy decays polynomially with the wavenumber. The power-law exponent is chosen uniformly random from the interval [2.3, 3.6 Third, the diffused noise initializer creates a tensor of values with white normally distributed noise, that is diffused afterwards. The resulting spectrum decays exponentially quadratic with an intensity rate chosen uniformly random from the interval [0.00005, 0.01 For all initializers, the resulting values of the initial conditions are normalized to have a maximum absolute value of one after generation. For vector quantities, the randomly selected initializer is sampled independently for each vector component.
Diffusion (diff) features the spatial dissipation of density field due to diffusion. The diffusivity also known as viscosity depends on the coordinate axis direction in our setup. This process is visually similar to blurring the density field over time, as high frequency information is damped. While simple at first glance, diffusion processes occur across domains, for instance in physics, biology, economics, and statistics.
Dimensionality: 600, 30, f 1. 2048, y 2048
Initial Conditions: random truncated Fourier Gaussian random field/ / diffused noise
Boundary Conditions: periodic
Time Step of Stored Data: 0.01
Spatial Domain Size of Simulation: [0, x [0, 1]
Fields: density
Varied Parameters: viscosity [0.005, 0.05( independently for 2; y/
Validation Set: random 15% split of all sequences from E [0, 500
Test Set: all sequences from E [500, 600
C.2.
PDEs
The Exponax solver (Koehler et al., 2024) was also used to simulate different reaction-diffusion PDEs. Such PDEs are most commonly encountered for local chemical reactions, but can also occur in domains such as biology, physics, or geology. They can be used to model traveling waves and pattern formation, and typically describe concentrations of one or more substances.
23

---

PDE-Transformer: Efficient and Versatile Transformers for Physics Simulations
Dittusion
5im pmn.iattoon.oo
0.0
-0.3
Simuistion 24
Densky
-05
Figure 13. Random example simulations from diff.
Fisher-KPP (fisher) contains simulations of cocttoondiffusion system according the Fisher-KPP equation. describes how the concentration of substance changes over time and space due reaction process controlled by reactivity parameter, while also considering the spatial spread of the substance due diffusion according diffusivity parameter. The equation be applied in the context of wave propagation and population dynamics, as well as ecology plasma physics. Figure 14 shows example visualizations from fisher.
Dimensionality: 600, 30, f 1, 2048, = 2048
Initial Conditions: random truncated Fourier Gaussian random field / diffused noise (with clamping to [0, 1)
Boundary Conditions: periodic
Time Step of Stored Data: 0.005
Spatial Domain Size Simulation: [0, 1] x [0, 1]
Fields: concentration
Varied Parameters: diffusivvty [0.00005, 0.01 and reactivity E [5, 15[
Validation Set: random 15% split of all sequences from [0, 500
Test Set: all sequences from [500, 600
Fisher-KPP
Simulation
Concentran
0.50
0.25
0.8
Sinubtinn
[-18
11
Swift Hahenberg
Concentration annanntrationn
1m0
1-3
1-6
1-g
tw12
15
Fx21
Fx24
1n27
Figure 14, Random example simulations from fisher and sh (from top to bottom).
Swift-Hohenberg (sh) features simulations of the Swift-Hooenberg equation, which describes certain pattern formation processes. It can be applied to describe the morphology of wrinkles curved elastic bilayer materials, for example, the formation of human fingerprints, where stresses between layers of skin lead to characteristic wrinkles. Figure 14 shows example visualizations from sh.
24

---

PDE-Transformer: Efficient and Versatile Transformers for Physics Simulations
Dimensionality: 600, 30, f 2048, y 2048
Initial Conditions: random truncated Fourier Gaussian random field / diffused noise
Boundary Conditions: periodic
Time Step of Stored Data: 0.5 (with 5 substeps for the simulation)
Spatial Domain Size of Simulation: [0, 20m] 20
Fields: concentration
Varied Parameters: reactivity E [0.4, 1[ and critical number E [0.8, 1.2[
Validation Set: random 15% split of all sequences from (0,500)
Test Set: all sequences from [500, 600
Gray-Scott (gs) describes a system in which two chemical substances react and diffuse over time. substance sa with concentration Ca consumed by the reaction and is replenished according to a feed rate, while the product of the reaction 8y with concentration с9 is removed from the domain according kill rate. Depending on the configuration of both rates, simulations result in highly different steady unsteady behavior with different patterns. Thus, create various subsets: four with temporally steady configurations, which result in state that does not substantially change anymore (gs-delta, gs-theta, gs-iota, and gs-kappa), and four temporally unsteady configurations, which continuously evolve over time (gs-alpha, gs-beta, gs-çamma, and aS--psillon For the unsteady case, separate test sets with longer temporal rollouts are created. Figure 15 shows example visualizations from the steady configurations, and Figures 16 and 17 from the unsteady configurations and corresponding test sets. For further details, we refer the reader to Pearson (1993). For all simulations, the diffusivity of the substances fixed to d, 0.00002 and d, 0. 0.00001. Furthermore, we random Gaussian blob initializer for these datasets: It creates four Gaussian blobs of random position and variance in the center 60% (20% for gs-kappa) of the domain, where the initialization of ca is the complement of cb i.e. Ca = 1 cb.
Steady Configurations (gs-delta, gs-theta, gs-iota, and gs-kappa):
Dimensionality: 100, = 30, f = 2. = 2048, = 2048 (per configuration)
Initial Conditions: random Gaussian blobs
Boundary Conditions: periodic
Time Step of Simulation: 1.0 (all configurations)
Time Step of Stored Data:
gs-del ta: 130.0
gs-theta: 200.0
as-iota: 240.0
gs ka-kappp: 300.0
Number of Warmup Steps (discarded, in time step of data storage):
gs-delta:
gs-theta:
gs-iota:
- ga-kappa: 15
Spatial Domain Size of Simulation: [0,2.5] [0,
Fields:
concentration c.. concentration C
Varied Parameters: feed rate and kill rate determined by configuration (i.c., initial conditions only within configuration)
gs dd-lelta: feed rate: 0.028, kill rate: 0.056
gs-theta: feed rate: 0.040, kill rate: 0.060
-gsiiot:: feed rate: 0.050. kill rate: 0.0605
- gs -aapppp feed rate: 0.052, kill rate: 0.063
25

---

PDE-Transformer: Efficient and Versatile Transformers for Physics Simulations
Validation Set: random 15% split of all sequences from [0, 80
Test Set: sequences from [80, 100
Unsteady Configurations (gs agsaalpha, gs-beta, gs ga-mam,,, and gs-epsilon):
Dimensionality: 100, 30, f 2048, y 2048 (per configuration)
Initial Conditions: random Gaussian blobs
Boundary Conditions: periodic
Time Step of Simulation: 1.0 (all configurations)
Time Step of Stored Data:
g8 a alah: 30.0
gs-beta: 30.0
gs gs-gamma: 75.0
epseponoon: 15.0
Number of Warmup Steps (discarded, in time step of data storage):
gs-alpha: 75
g8 bett: 50
gs ga-gamma: 70
- g8-epsilon: 300
Spatial Domain Size of Simulation: [0, 2.,
Fields: concentration ca. concentration сj
Varied Parameters: feed rate and kill rate determined by configuration (i.e., initial conditions only within configuration)
ga gs-alpha: feed rate: 0.008, kill rate: 0.046
gs-beta: feed rate: 0.020, kill rate: 0.046
gs ga-gamma: feed rate: 0.024, kill rate: 0.056
gs-eps ilon: feed rate: 0.020, kill rate: 0.056
Validation Set: random 15% split of all sequences from E [0, 100
Test Set: separate simulations with 30, 100, f = 2, 2048, 2048 (per configuration)
C.3. Nonlinear PDEs
Nonlinear PDEs generally difficult to study, the question of the existence of analytical solutions is already hard problem. Furthermore, most general techniques do not work across cases, and single nonlinear PDEs are commonly tackled as individual problems. The Exponax solver (Koehler et al., 2024) provides tools to approach some nonlinear PDEs, however, we also consider other data sources below.
Burgers (burgers) features simulations of Burgers' equation, which similar to an advection-diffusion problem. Rather than the transport of a scalar density, it describes how flow field itself changes due advection and diffusion. This can lead to the development sharp discontinuities or shock waves, making difficult to simulate accurately. Burgers' equation also has applications in nonlinear acoustics and traffic flow. Figure 18 shows example visualizations from burgers.
Dimensionality: 600, = 30, f = 2, x = 2048, = 2048
Initial Conditions: random truncated Fourier Gaussian random field / diffused noise
Boundary Conditions: periodic
Time Step of Stored Data: 0.01 (with 50 substeps for the simulation)
Spatial Domain Size of Simulation: [0, 1] x [0,
26

---

PDE-Transformer: Efficient and Versatile Transformers for Physics Simulations
Gray Scott: Delta
1.0
Sinuiation
-04
Sinulation 7
00
Gray Scott: Theta
Sinilation
04
Sinuiation
t-0
t-3
1-6
t-9
b-12
t-15
t-18
L-21
t-24
t-27
0.0
Gray Scott: lota
1.0
Simuiations
Simutation
Gray Kappa
Simmlettinn 5
Concontration
Simulation
Concentration
b+0
t-3
t-5 t-9
t=12
t-15
1-18
1-21
1-24
Lm27
0.D
Figure 15. Random example simulations from steady configurations of the Gray-Scott model of a reaction-diflusion system: gs-des gs-delta, gs gs-theta, gs-icta, and gs -appa..
Fields: velocity (z, y)
Varied Parameters: viscosity [0.00005, 0.0003[
Validation Set: random 15% split of all sequences from E [0, 500[
Test Set: all sequences from [500, 600
Korteweg-de-Vriee Vries (kdv) contains simulations of the Korteweg-de Vries equation on a periodic domain, which serves as a model waves on shallow water. It challenging as energy transported to high spatial frequencies, leading to individua moving solition waves with unchanged shape and propagation speed. Across simulations, the convection coefficient with value of -6, and the dispersivity coefficient with value of 1 remain constant. Figure 18 shows example visualizations from kdv.
Dimensionality: 600, = 30, f = 2, = 2048, 2048
Initial Conditions: random truncated Fourier Gaussian random field/ / diffused noise
Boundary Conditions: periodic
27

---

PDE-Transforme:: Efficient and Versatile Transformers for Physics Simulations
Gray Scott: Alpha
Smulations
0.75
0.50
Simuiation
0.
Gray Scott: Beta
Sinilation
0.75
Oononnratiin tedion
0.50
0.25
Simuistion
t-0
t-3
L-9
b-12
t-15
t-18
t-21
D-24
t-27
Gray Scott: Gamma
Simuiation
0.75
0.50
0.25
Simutation
Conco
Gray Epsilon
Simalasiion
0.75
Concentration Concontration
0.50
0.25
Simulation
b+0
L-3
L-5
L-9 t=12
t-15
1=
1-21
1-24
L-27
Figure 16. Random example simulations from unsteady configurations of the Gray-Scott model of reaction-diffusion system:gs-ar allpha, as-beta, 3s-ganma, and gs gessiinon..
Time Step of Stored Data: 0.05 (with 10 substeps for the simulation)
Spatial Domain Size of Simulation: varied per simulation
Fields: velocity (x, y)
Varied Parameters: domain extent E [30, 120[ identically for vr y, i.e. a square domain, and viscosity [0.00005, 0.001
Validation Set: random 15% split of all sequences from [0, 500
Test Set: sequences from [500, 600
Kuramoto-Sivashinsky (ks) features simulations of the Kuramoto-Sivashinte equations on periodic domain, which models thermo-diffusive flame instabilities in combustion. It also has applications in reaction-diffusio systems. The equation is well-known for its chaotic behavior, where temporal trajectories with slightly different initial conditions can substantially diverge over time. The initial transient phase of the simulations is discarded. For the ks dataset, test set with longer rollout is used to investigate how well models can deal with this chaotic behavior. Figure 18 shows example visualizations from ks.
Dimensionality: 600, = 30, f 1, 2048, y 2048
28

---

PDE-Transformer: Efficient and Versatile Transformers for Physics Simulations
Gray Scott: Alpha (test)
Siulations
0.50
0.25
00
Smulation
02
Concent
Gray Scott: Beta (test)
0.75
0.50
0.25
0.4
Simuistion
Concentration
t-0
t-10
t-20
t-30
D-40
t-50
t-60
E-70
C-80
t-90
Gray Scott: Gamma (test)
Conco Irreentationn
Simutation
0.23
Simutatio
0.4
5-
Gray Scott: (test)
Concentration A
Simmlastoon
0.75
0.50
0.25
Simulation 2
Concentration
b+0
t-10
t=20
t=30
t=40
t-50
2-60
b-70
1=80
L-90
0.0
Figure 17. Random example simulations from test sets with longer rollout from ga- alla,, gs-beta gw-gget-ee -a-mamm,, and gs -es-epsil...
Initial Conditions:
random truncated Fourier Gaussian random field / diffused noise
Boundary Conditions: periodic
Time Step of Stored Data: 0.5 (with 5 substeps for the simulation)
Number of Warmup Steps (discarded, in time step of data storage): 200
Spatial Domain Size of Simulation: varied per simulation
Fields: density
Varied Parameters: domain extent [10, 130 identically for x. y, i.e. square domain
Validation Set: random 15% split of all sequences from [0, 600
Test Set: separate simulations with 50, 200, f x 2048, y 2048
Decaying Turbulence (decay-turb) contains simulations of the Navier-Stokes equations in a streamfunetion- formulation on a periodic domain. The simulations exhibit swirling turbulent vortices that decay over time. For this dataset, a test set with longer rollout used, where the decay over time is even more pronounced. Figure 19 shows example visualizations from decay-turb.
29

---

PDE-Transformer: Efficient and Versatile Transformers for Physics Simulations
Burgers
30003 X
Sim Simullation
Simaletion
Veledty
Korteweg Vries
0.1
Semsuletion
Simutetion
Velucity
t+0
t+3
L-5
t-12
t-15
C-18
t-21
b-24
t-27
kuramato-Siwashinsky
censtty
--10
Simuletten
Dondky
h-15
Kuramoto-Sivashinsky (test)
Somulatton
-10
Stmuistoon
-10
L-0
t=20
[-40
1-60
1-60
I=100
t-120
t-140
t-160
tw180
Figure 18. Random example simulations from burgers, kdv, ks, and the test set of ks with longer rollout.
Dimensionality: 600, = 30, f 1, 2048, 2048
Initial Conditions: random truncated Fourier Gaussian random field / diffused noise
Boundary Conditions: periodic
Time Step of Stored Data: 3.0 (with 500 substeps for the simulation)
Spatial Domain Size of Simulation: [0, 1] x [0, 1]
Fields: vorticity
Varied Parameters: viscosity& E [0.00005, 0.0001[
Validation Set: random 15% split of all sequences from [0, 600
Test Set: separate simulations with 50, 200, f x 2048, 2048
Kolmogorov Flow (kolm-flow) features simulations of the Navier-Stokes equations in a streamfunction cttprfuunntioooootiiitty formulation on a periodic domain. In contrast to the decaying turbulence above, an additional forcing term ensures that new energy is introduced into the system that sustains the vortices indefinitely, leading to spatiotemporal chaotic behavior. The transient phase of the simulations where the initialization transforms stripe pattern and into vortices afterwards
30

---

PDE-Transformer: Efficient and Versatile Transformers for Physics Simulations
discarded. For this dataset, with longer rollout is used test how well models deal with the chaotic behavior of the flow. Figure 19 shows example visualizations from kolm-ff flow.
Dimensionality: = 600, t = 30, f = x = 2048, 2048
Initial Conditions: random truncated Fourier Gaussian random field / diffused noise
Boundary Conditions: periodic
Time Step of Stored Data: 0.3 (with 1500 substeps for the simulation)
Number of Warmup Steps (discarded, time step of data storage): 50
Spatial Domain Size Simulation: [0, 1] x [0, 1]
Fields: vorticity
Varied Parameters: viscosity E [0.0001, 0.001
Validation Set: random 15% split of all sequences from [0, 600
Test Set: separate simulations with 50, 200, f 2048, y 2048
31

---

PDE-Transformer: Efficient and Versatile Transformers for Physics Simulations
NS: Decayling Turbulence
Simubetio Simuuaatio
-0.
0.0
1-0
1-3
t-5
I-9
t-12
t-15
t-18
t-21
b+24
t-27
NS: Decaying Turbulence (test)
Simutation Simulation 1
0.2
Vorticy
0.0
-0.2
0.05
Vertilty
-0.05
I-20
[-60
[-80
140
NS: Kolmogorov Flow
Simulation
vortict
Vorldlt
10
1+3
1+5
1+9
t-12
t-15
t-18
1-21
1-24
1-27
NS: Kolmogorov Flow (test)
Simusston
Vorticty
Simuuztion 4
-10
Figure 19. Random example simulations from decay-turb, and kolm f1 ow, with examples from each corresponding test set with
longer rollout.
32

---

PDE-Transformer: Efficient and Versatile Transformers for Physics Simulations
D. Autoregressive Predictions
show results for all datasets in Table Table 10 for the nRMSE with rollouts step and 20 steps respectively, Autoregressive predictions of trajectories on the datasets from 0 until f 27 are visualized below, see Figures 20 23, for PDE-L trained with the MSE loss and mixed channels (MC).
Table 9. nRMSE, after step for the pre-training datasets.
ES.hhe
ga-delta
ga-tbeta
koin-flow
DIES
0.0225
06262
036333
DO509
03310
0,088
0.0815
0,09422
010475
0.1254
0,3422
0.365
0.256
0,085
0.2570
0.120
UDITS
(1137))
08191
6,0335
00:73
03242
00224
9.0310
0.0022
0.0149
0.0151
0.1125
0.0500
0.093
0,0519
0.3457
010715
500ES
01874
08355
0.635
00240
0A239
00254
0.0810
010419
IU1771
011232
011112
0.215
0.311
0.559
013114
019977
PactFormer
01140
09455
0,821
01031
0.256
00347
(U15TT
00112
03225
00264
0.410
0/816
01221
01144
UNet
0.0599
06991
0.005
00469
04335
0.441
9.6755
0.8445
0.4133
0.0416
0.308
00355
0.4299
0,0829
0.2855
0.2355
FDBSS
0.0370
06215
0,430
00216
04247
0.048
0.0295
0.0315
0.0322
0.0175
0.0129
0.0141
0.2311
0.063
0.2361
0.0573
FDB-B
0.138
0.362
0.446
00:45
0.039
0.2700
0.0293
0/0256
0.0141
0.0140
0,0332
0.005
0.0055
0,0640
0.2206
0.577
PDBL
0.038
0.1155
0,045
00111
01023
9.91%
0,838
0.0134
001333
0J8376
0J863
0/864
0,1134
0,0647
012113
01447
Table 10. nRMSF20 after 20 steps for the pre-training datasets.
Eisher
ga-delta
ge-tbeta
ge-iota
koin-tlov
DITS
0.2677
0.8003
0,4915
1644
0.704]
1.4305
1.1059
1,0754
1.5140
1.2253
1.5255
1.2355
1.2387
0,9253
0.5148
1.2348
UDITS
0.2835
0.1982
0.3157
019655
06144
0.993
3.7592
(171555
1,0005
0.3591
1J1323
0.9341
0.7568
0.5842
1.0018
015575
SCOTE
037733
0.3965
0.4351
13377
0.7017
0,5744
0.9355
(.9100)
1.1182
0.4704
1J9965
1,0003
1,554
0.6672
10455
119666
FactForerr
L9913
06704
0.224
13988
03370
0.357
3.5107
0.9875
1.1258
0,5425
10534
1.1000
12319
0.526
1,2922
L.570
UNet
0.5253
0.7142
0.5829
13061
0.7127
1,008
1.1585
0.9745
10059
0.7716
0.9579
09197
0.9803
0.532
0.9746
1.2997
PDE-S
0.2055
0.157
033383
0.741
0.568
0.00%
3:643
0.7887
1.0866
0.3156
10247
0.7501
0.7272
0.5737
0.941
0.570
PDB-B
0.2000
0.1142
0.3304
03559
06131
0.116
0.110
017555
1,0253
0.3636
03910
0/632
0.435
0.5354
017333
03005
PDB-L
0.2104
0/821
0.3084
0.7357
0.5731
0.5544
0,5924
0.7083
0.9916
0.1997
0,3382
0,5333
0.5573
0.811
0
0.7102
Prediction
Diffusion
1.5
1.0
Densty
Reference
005
-10
-15
t=0
t=3
t=
t=9
t=12
t=15
t=18
t-21
t=24
1=27
Fihee-KPP
Prediction
Concentration
Reference
t-3
t-6
C-9
1-12
t-15
t-18
t-21
t-21
t-27
Swiit-Hohenberg
Predictin
Concentration
Reference
0-0
tm3
t-
9-
12
1-
t-18
t-21
L-21
1-27
Pigure 20, Visualizations of model's prediction on diff, fisher and sh.
33

---

PDE-Transformer: Efficient and Versatile Transformers for Physics Simulations
Gray Scott: Alpha
Prediction
Concentracion
1-18
1-21
1-27
Gray Scott: Beta
Prediction
Concentration
Reference
t=6
t-9
1-12
1=15
t-18
t-21
t-24
2=
Scott: Gamma
Prediction
Concentration
Reforence
t=3
1=5
1=
1=
L=18
1=21
L-24
1=27
Gray Epsiion
Prediction
Concentration
Reference
C=0
L=3
L=6
t=
1=
1=15
[=16
L=21
L=24
(=27
Figure 21. Visualizations of model's prediction gs-alpha, gs-beta, gs-gamma and gs-epsilon
34

---

PDE-Transformer: Efficient and Versatile Transformers for Physics Simulations
Predictionn
Gray Scott: Gamma
Concentrations A
Reference
1-27
Prediction
Gray Scott: lota
Concentration
0.5
0.0
Reference
-005
-1.0
-15
Prediction
Concentration
Reference
tw3
1-6
t-
1-12
15
1-18
t-21
2-24
Pred ctidiction
Scott: Theta
Concentration
Reference
tw0
Figure 22, Visualizations of model's prediction gs-gamma, gs-iota, ka-kappa and gs-theta.
35

---

PDE-Transformer: Efficient and Versatile Transformers for Physics Simulations
Burgers
Prediction
Velocity
Reference
t=0
t=3
1=
t-9
t=12
t=15
t-18
t-21
t-24
t-27
Prediction
eg
Velocity X
Reference
t-0
tm3
t-6
t-9
t-12
t-15
t-18
t-21
C-21
t-27
Kuramoto Sivashinsky
Prediction
Density
Reference
t=27
NS: Decaying Turbulence
Predictian
Vorticity
Reference
2
tw18
t-21
tw27
NS: Kolmogorov Flow
Prediction
Verticity
Reference
t-0
t=3
t-6
9-
t- 1-12
t=15
tw18
t-21
t-24
t=27
Figure 23, Visualizations of model's prediction on burgers, kay, ks, decay-turb and ko colmmffee Fow.
36

---

PDE-Transformer: Efficient and Versatile Transformers for Physics Simulations
E. Downstream tasks
E.1. Dataset
We introduce three challenging PDE prediction missions, which are active matter, Boyleigh-Bénanne convection, and shear flow, as the downstream tasks. Fig. shows the visual examples of these three tasks. All the simulations are from the Well dataset (Ohana et al., 2024). For each provided dataset, a predefined data split in training, validation, and test already exists in the Well dataset. We randomly select 42, 8, and 10 trajectories from the corresponding split data for training, validation, and testing, respectively. Each trajectory is also truncated randomly to 30 frames. Details for each dataset are discussed as follows:
Active Matter features simulations of rod-like biological active particles immersed in Stokes flow. The active particles transfer chemical energy into mechanical work, leading to stresses that are communicated across the system. Furthermore, particle coordination causes complex behavior inside the flow. The following overview summarizes key characteristics of the dataset (for further details see Maddu al., 2024):
Boundary Conditions: periodic
Time Step of Stored Data: 0.25
Spatial Domain Size of Simulation: [0, 10] [0,
Spatial Resolution: 256, y 256.
Fields: concentration, velocity (z, 3). The orientation (zx,, xy, 222 y2, and strain (zz, ay,, yz. yy) fields in the original Well dataset are dropped in the current test.
CayleighhBénarde Convection contains simulations of Rayleigh-Bénae convection on a horizontally periodic domain. combines fluid dynamics and thermodynamics, by simulating convection cells forming due temperature differences between upper and lower plate. The combination of buoyancy, conduction, and viscosity leads to complex fluid behavior with boundary layers and vortices. The following overview summarizes key characteristies of the dataset (for further details Burns et al., 2020):
Boundary Conditions: periodic (z), wall (y)
Time Step of Stored Data: 0.25
Spatial Domain Size Simulation: [0,4] x [0, 1]
Spatial Resolution: 512, y 128.
Fields: buoyancy, pressure, velocity (x, y)
Shear Flow features simulations of the periodic incompressible Navier-Sookes equations in a shear flow configuration, where fluid layers are sliding past each other different velocities. Predicting the resulting vortices across different Reynolds and Schmidt numbers has important automotive, biomedical, and aerodynamics applications. The following overview summarizes key characteristics of the dataset (for further details see Burns et al., 2020):
Boundary Conditions: periodic
Time Step of Stored Data: 0.1
Spatial Domain Size Simulation: [0, 1] x [0, 2]
Spatial Resolution: x 512, y = 256.
Fields: density, pressure, velocity (z, y)
E.2. Autoregressive Predictions
We show autoregressive predictions with the pre-trained PDE-S and separate channels on the training set for each dataset in Figures 25 27.
37

---

PDE-Transformer: Efficient and Versatile Transformers for Physics Simulations
Rayleigh Benard Convection
Rc
Buoysnnc
-05
t-0
t-20
t:40
t=
t-100
t-120
t-140
t-160
t-180
NS: Shear Flow
0.5
Veledty
0.25
Y
Figure 24. Random example simulations from active matter, Rayleigh-Benard Bénard convection, and shear flow.
38

---

PDE-Transformer: Efficient and Versatile Transformers for Physics Simulations
Prediction coion
Refereno
-0.1
-0.2
t-0
t-3
t-6
t-9
t-12
t-15
18
t-21
t-24
t-27
Prediction
Reference
t=
t=3
t=
t=9
t=12
t=15
t=18
t=21
t=24
t=27
Prediction
Reference
t-0
t-3
t-6
t-9
12
t-15
C- 18
t-
1-24
t-27
Figure 25, Active Matter. Autoregressive prediction with pre-trained PDE-S (SC).
39

---

PDE-Transformer: Efficient and Versatile Transformers for Physics Simulations
Prediction
2.5
2.0
1.5
Buayancy
Refermcce
t=0
L=3
1:6
L=
L=12
t=15
t=10
1=21
L=24
L=27
Prediction
06
Velndlty
000
-02
Reference
t=0
t=3
t=6
t=
t=12
1=15
ts:10
t=21
t=24
t=27
Figure 26. Ehyyleigh-éénard Convection, Autoregressive prediction with pre-trained PDE-S (SC).
40

---

PDE-Transformer: Efficient and Versatile Transformers for Physics Simulations
Prediction
00.5
Reference
-0.5
-1.0
t-12
t-15
tw21
1-24
1-27
2.5
0.0
Prediction
25
-5.0
775
Rel referrecce
-175
t-3
[-12
1-24
1-27
Prediction
Velocity X
Reference
C-0
t-3
C-6
t-9
t-12
1-
t-18
(-21
t-24
t=27
Figure 27, Shear Flow. Autoregressive prediction with pre-trained PDE-S (SC).


Reviews:
[{"review_id": "YIJnjmMUfD", "reviewer": "Reviewer_hdWy", "summary": "This paper introduce PDE-Transformer, a transformer made for PDE data being able to be incorporated in a supervised learning task or in a diffusion model. The model is extensively tests on various benchmarks.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the positive review and feedback. \n\n**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.  \n\n**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. \nEven 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.\n\n[1] Uncertainty-aware Surrogate Models for Airfoil Flow Simulations with Denoising Diffusion Probabilistic Models, https://arxiv.org/pdf/2312.05320v1\n\n[2] From Zero To Turbulence: Generative Modeling for 3D Flow Simulation, https://arxiv.org/pdf/2306.01776\n\n**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. \n\n**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.\nHowever we didn't see any improvements here and using a pure transformer backbone with a single patchification layer was optimal for us.\nOur 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.\n\n**Which is better: MC or SC?**\nFor 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. \nThis 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.\n\n**Comparison with Poseidon/MPP**\nWe 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.\n \nWe 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. \n\nPlease let us know if this could clear your open questions. We are happy to discuss more.", "claims_and_evidence": "Yes, the claims seem supported by convincing evidence", "methods_and_evaluation_criteria": "The authors benchmark their models on many datasets and study the transfer properties to datasets unseen during training.", "theoretical_claims": "No theoretical claims.", "experimental_designs_or_analyses": "The experiments are extensive and look good to me.", "relation_to_broader_scientific_literature": "The contributions of the paper are a good progress towards better architectures in the field.", "essential_references_not_discussed": "None that I can think of.", "other_strengths_and_weaknesses": "Strengths:\n- The paper is very dense with many different ideas. All these ideas are supported with numerical experiments.\n- The results are convincing.\n- I like the downstream tasks study.\n\nWeaknesses:\n- The model is currently limited to 2D regular grids\n- The paper is very dense and makes the architecture not 100% clear in 1 read.\n- The fact that the model can be use in supervised and diffusion training gives me mixed feelings: it is a good idea, but overall, what does a diffusion models additionally bring here? It doesn't seem to be discussed.\n- From the conclusions of Table 3, it is not clear to me which one of MC or SC is better.", "other_comments_or_suggestions": "- Line 158 (right paragraph) should be $S$ is and not S.\n- Line 210 (r.p.) should be separate (typo)", "questions_for_authors": "- Can the authors comment on the use of bf16 precision? This seems like a bold choice to me for this type of data which is usually in fp32. With bf16, don't you lose a lot of precision, especially for autoregressive tasks?\n- One clear advantage of using diffusion models here would be that we would obtain probabilistic samples. It could lead to uncertainties in the prediction, very valuable for scientific machine learning. Can the authors comment on that, and possibly suggest 1 experiment to study this (not mandatory, but I think it would strengthen the paper)?\n- There are some unclear terms in the text, what is token upsampling/downsampling? PixelShuffle/unshuffle?\n- Are images just patchified? No MLP/convolutions? Do the authors know whether their better performance is because they are working in pixel space? \n- Why did the author not compared with other foundation models like MPP and Poseidon?\n\nIf the weaknesses and questions are satisfyingly answered, I am willing to modify my score accordingly.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HVPtz3BzRV", "reviewer": "Reviewer_aHgi", "summary": "The paper presents PDE-transformer, a new transformer approach trained simultaneously on different physical systems. They propose an alternative to the DiT architecture that is suited for PDEs. Specifically, they use a multi-scale architecture and a shifted-window attention to prevent from the quadratic complexity. They propose also a separate channel strategy when dealing with physical systems with different number of physical channels. They evaluate their method on a variety of PDEs with specific behaviors and complex PDEs for finetuning, where they show competitive or sota performance.\n\n## update after rebuttal\nThe authors proposed a nice approach for multi-task PDE solving. There are some results that are surprising to me, that have been mentioned in the rebuttal. Overall, the paper is of good quality and thus I keep my score as it is.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the positive review and feedback. \n\n**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. \n\nWhen 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).\nIn this situation, supervised training is a suitable choice, especially if we use a metric based on the MSE/L2 distance for evaluation and training.\n\nThis 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. \nIn 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. \n\nNow, 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.\n\nTo 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. \n\nOf 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. \nIn 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.\n\n**Generalization to Unstructured Meshes** \n\nIt's possible to generalize the the architecture to irregular grids, which however this is not the focus of the current paper. \n\n- 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] \n\n- 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. \n\nBoth approaches are planned as directions of future work.\n\n[1] Geometry-Informed Neural Operator for Large-Scale 3D PDEs, https://arxiv.org/pdf/2309.00583\n\nPlease let us know if this could clear your open questions. We are happy to discuss more.", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "They used a variety of datasets for their experiments and also tested on very hard datasets  for some additional studies, showing that using a pretrained transformer on more simple datasets can help for downstream tasks such as finetuning on very complex PDEs.", "theoretical_claims": "No theoretical claims", "experimental_designs_or_analyses": "I did check the different experimental designs and analyses and did not find a particular issue. However, I am a bit concerned by the analysis made on the diffusion training. They are known to correctly model complex distributions, which could be the case when pretraining a model on a variety of PDEs.", "relation_to_broader_scientific_literature": "The authors proposed a model in the line of the existing foundation models for PDEs literature.", "essential_references_not_discussed": "The authors discussed the main related works", "other_strengths_and_weaknesses": "The paper tackles an important problem, it learns a single model for solving multiple physical systems. It exhibits strong performance on a diverse number of datasets and have been tested on very complex datasets in a downstream task, showing strong performance.\n\nThe paper is clear and I particularly liked the appendix with all the dataset descriptions.\n\nThe originality of the paper is a bit weak, it extends existing works from computer vision with DiTs for scientific problems. Using U-shape transformers is not novel, but it's application for PDEs with a multi-scale approach makes sense. \n\nOne potential weakness of the paper is the restriction to flexible grids, which should be a key property of neural solvers. I think it could be easily extended to handle irregular grids.", "other_comments_or_suggestions": "No more comments. See above sections.", "questions_for_authors": "I am bit concerned by the conclusion made with MSE training compared to diffusion training. I don't really see why MSE training could yield superior performance, notably when modeling complex distributions such as multi-physics data. In my opinion, diffusion training should exhibit superior performance. Could you elaborate more on why MSE training is superior here please.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "IFLk8kqJ05", "reviewer": "Reviewer_2XVZ", "summary": "The paper presents a transformer model called PDE Transformer designed to solve PDE (partial differential equations), therefore potentially allowing for physical simulations. The model is based on the diffusion transformer architecture (DiT), and as such, can be be trained not only for forecasting, but also for generation. To the best of my understanding, one of the main additions to the architecture is to use down- and upscaling tokens at the end of each transformer stage (instead of on the query- key-value tuple of the self-attention operation like it was introduced before in another version of the architecture) which allows for a much faster training time, and the application of the architecture to a PDE learning setting. The paper claims to beat SOTA architectures on a benchmark PDE forecasting dataset including multiple PDEs, and for so called “downstream tasks”: more complex PDEs with various boundary conditions on specific geometries.\n\n# Update after the rebuttal\n\nI appreciate the efforts made by the authors during the rebuttal and I think the paper is of overall good quality. I am maintaining my score as it is.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the positive review and feedback. \n\n**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.\n\n**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. \n\nBecause 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.\n\nBased 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.\n\n**Graph-based models** \nWe 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. \n\nNonetheless, 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.\n\n| Model    | Epoch | nRMSE(1) | nRMSE(10) | Time (h) | Params | GFlops |\n| -------- | ------- | ------- | ------- | ------- | ------- | ------- |\n| PDE-S  | 100 | 0.044  | 0.36 | 7h 42m | 33.2M | 19.62 |\n| MPNN | 2 | 0.283 | 1.07 | 37h 2m (1 GPU) | 2.44M | 396.16 |\n\nAs 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.\n\n[2] Differentiability in Unrolled Training of Neural Physics Simulators on Transient Dynamics, https://arxiv.org/pdf/2402.12971\n\n**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.\n\n**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.\n\nWe 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. \n\nWe are happy to continue in more detail during the discussion phase.", "claims_and_evidence": "The claims are clear: the suggested model outperforms SOTA models for deep learning based PDE simulation. \n\nThe evidence given is proper but not entirely convincing, considering the following points: \n- The number of SOTA models it is compared to for the main PDEs is quite limited: 3 models, only transformer based. Why not comparing to other types of PDE-learning models that are also quite efficient and always competing such as graph based models and neural operator based models (e.g. Message passing neural pde solvers from Brandstetter et al 202, MagNet boussif et al 2022, or FNOs, which I don't need to introduce) ?\n- I am honestly not sure why different SOTA models are used for the main PDE prediction tasks (trained from scratch) and the “downstream” tasks? They can all do both predictions, so we might as well see the performances of all models on all tasks!\n- Also, one of them (UDiT) is basically performing equally to PDE-Transformer, although it is slower at training time. The suggested model is therefore indeed faster, but not necessarily “outperforming” SOTA. Perhaps this should be highlighted in the main claims!\n- The configurations in which the model is tested is also quite limited: only one spatial resolution is used (and quite a corse one for typical PDE datasets: 256x256) and one time resolution (30 steps for the 600 different PDE trajectories) for a rather low horizon of 10 steps. I appreciated the studies on the patch and window sizes and supervision vs probabilistic learning, but I think time and space resolutions are of paramount importance for PDE learning (if I had to choose, I would put extra studies in the appendix and these on resolution in the main paper if it causes a space problem). Some models might be more efficient for longer horizon and/or higher spatial resolution! I think it would be more fair to do more experiments regarding various resolutions and conclude on the good and bad points of multiple models comparatively (For example FNO and variants are known to be quite good at zero-shot super resolution in both space and time tasks)\n- The same two previous comments can be said for the “downstream tasks”: 4 models are compared, including some models that are not adapted to strange geometries (which is mentioned by the authors, which I appreciate!). Similarly experiences with various resolutions would be interesting.\n- Note that I really appreciated the pretrained experiments on the downstream tasks, I think it is interesting and impressive; however, I did not understand if the other models were also retrained or only trained from scratch? If the latter, the comparison with the pretrained PDE-transfo is a little unfair. (Cf questions for more details)", "methods_and_evaluation_criteria": "The proposed methods do make sense, as it basically builds on top of successful architectures, which may not be the most creative but makes total sense! The benchmark dataset is good and gathers many PDEs (APEBench). Perhaps it would make sense to show the typical multiple metrics for forecasting (MAE, MSE etc rather than only nRMSE ) in order to have a full view on the models performances. Also, more experiments would be needed to have a definite conclusion, as mentioned previously.", "theoretical_claims": "No theoretical claims are given in this paper.", "experimental_designs_or_analyses": "The experimental design seems perfectly sound, except that, as mentioned in the Claims And Evidence section, more experiments regarding spatial and temporal resolutions, higher time horizons and more SOTA models would bring more convincing evidence to the suggested model’s performance.", "relation_to_broader_scientific_literature": "The contribution of the paper seems to be in terms of pure performance in the PDE learning domain, as it competes with state of the art while being apparently faster to train than some SOTA PDE models. In terms of pure architecture, it is however not very original, as it is basically the UDiT architecture but with up and downsample tokens added at the end of transformer stages instead of the attention dot products themselves, applied to a PDE forecasting task. \nIt can be seen from the performance table that the only advantage of the PDE transformer over UDiT is the training time, but the performance is extremely close.", "essential_references_not_discussed": "Relating to the SOTA models previously mentioned in the review, it appears that the authors do not mention PDE learning models based on graph neural networks, while it was of very high significance in the domain at some point, particularly for their super resolution and irregular mesh learning capabilities (and I believe they still are). It is of course necessary to compare to recent transformer models, but at the end of the day all significant SOTA models should be considered, regardless of the architecture type. \n\nSuch graph based models include for example: \nMessage passing neural PDE solvers: Brandstetter, J., Worrall, D., & Welling, M. (2022). Message passing neural PDE solvers. arXiv preprint arXiv:2202.03376.\n\nMAgNet: Boussif, O., Bengio, Y., Benabbou, L., & Assouline, D. (2022). MAgnet: Mesh agnostic neural PDE solver. Advances in Neural Information Processing Systems, 35, 31972-31985.\n\nI am not saying you need to compare to all kinds of models that exist, that would be impossible! But do believe the SOTA models used in comparison needs to include some of these graph based models, in terms of pure forecasting performance, but also super resolution abilities and training time (since it is one of the biggest advantage of the model).", "other_strengths_and_weaknesses": "Strengths:\n- The paper is well written and well presented, and rather clearly explained.\n- I appreciate \n    - the practical idea of using up and downsample tokens in between transformer stages instead of in attention computation to make it more time efficient and\n    - The idea of using a diffusion model to the PDE learning domain (as it seems to work pretty well!)\n\nWeaknesses:\n\nThe main weakness is therefore the originality of the paper, given than its contribution is mainly residing in a very slight practical change in an already existing architecture and the application of the architecture to PDE learning. This could be interpreted as a “real-world” use case since PDE learning can be directly applied for physical simulations, however the paper does not really take it this far, by doing actual physical simulations (combining potentially multiple PDEs) and rather tries it directly on PDE data. \nI do believe it is still a useful contribution, but then I think it would greatly benefit from stronger experiments (on multiple spatial and time resolution experiments, as mentioned previously) to counter balance the lack of contributions in other aspects.", "other_comments_or_suggestions": "- Is the idea of “expansion rate” really useful? It seems to me that it only illustrates the tradeoff between better performance/scalability obtained with lower/higher patch(token) size, which is already known! Also, you are mentioning it in the abstract but you do not define it there but only later in the article. If you want to keep presenting this notion, please define it in the abstract before mentioning it or dont mention it!\n- When the nRMSE is defined, a dependency to horizon time could be added (meaning defining it as nRMSE(t) ) since it is presented like that in table 1, to be more coherent with the notations.", "questions_for_authors": "- Note that I really appreciated the pretrained experiments on the downstream tasks, I think it is interesting and impressive; however, I did not understand if the other models were also retrained or only trained from scratch? If the latter, the comparison with pretrained PDE-transo is a little unfair, and it would be interesting to see the performance of these models pre-trained as well in comparison.\n- In addition to additional experiments, I think it would be interesting to consider super resolution experiments, in space and/or time, as I think it is very important in PDE learning (the ability to have a higher resolution for the prediction, based on coarser resolution models like global weather forecasting models, is more important to me than a faster training time - once a model is trained in real life, what matters is rather the inference time), and since one of the bigger contribution is basically the application of a known architecture in PDEs, it would be significant  to see if it compares to neural operators models and graph neural networks models which are both adapted to these tasks (using respectively implicit functions and mesh agnostic nodes)", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "uRVWUlXe5E", "reviewer": "Reviewer_yvZY", "summary": "This paper presents an enhanced diffusion transformer architecture through the integration of several established techniques, including multi-scale modeling and shifted window attention mechanisms. The improved model is subsequently applied to partial differential equation (PDE) solving tasks. Extensive experimental evaluations demonstrate the effectiveness of the proposed architectural modifications.", "questions": "", "limitations": "", "rating": 1, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the review and feedback. We want to address your remaining concerns in the following:\n\n**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.\n\nThank 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):\n\n| Model    | nRMSE(1) | nRMSE(10) | Time (h) | Params | GFlops |\n| -------- | ------- | ------- | ------- | ------- | ------- |\n| PDE-S  | 0.044  | 0.36 | 7h 42m | 33.2M | 19.62 |\n| FactFormer | 0.069 | 0.65 | 38h 8m | 3.8M | 66.76 |\n\nPDE-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.\n\nFor 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.\n\n\n[1a] Poseidon: Efficient Foundation Models for PDEs, https://arxiv.org/pdf/2405.19101\n\n[2b] U-DiTs: Downsample Tokens in U-Shaped Diffusion Transformers, https://arxiv.org/pdf/2405.02730v1\n\n[3c] Scalable Transformer for PDE Surrogate Modeling, https://arxiv.org/pdf/2305.17560\n\n[4d] PDEformer: Towards a Foundation Model for One-Dimensional Partial Differential Equations, https://arxiv.org/pdf/2402.12652\n\n[5e] Unisolver: Unisolver: PDE-Conditional Transformers Are Universal PDE Solvers, https://arxiv.org/pdf/2405.17527\n\n**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.\n\nWe therefore kindly ask you to reconsider your overall recommendation, and we’d be happy to discuss any remaining open aspects.", "claims_and_evidence": "All claims are clear for me.", "methods_and_evaluation_criteria": "The proposed method appears to primarily incorporate well-established techniques in deep learning, making the specific contributions less distinct. Additionally, the experimental evaluation employs relatively weak baselines. It would be beneficial to include comparisons with recent advancements in PDE foundation models and neural PDE solvers [1,2] for more comprehensive assessment.\n\n[1] Scalable Transformer for PDE Surrogate Modeling\n\n[2] PDEformer: Towards a Foundation Model for One-Dimensional Partial Differential Equations", "theoretical_claims": "This paper does not have any theoretical claim.", "experimental_designs_or_analyses": "The experiments and analyses are extensive and content-rich. The writing is clear and well-structured.", "relation_to_broader_scientific_literature": "N/A", "essential_references_not_discussed": "It is worth noting that UniSolver [1], a recent work also built upon the diffusion transformer architecture, addresses similar PDE generalization challenges. A detailed comparison and discussion with UniSolver should be included to highlight the distinctions and relative advantages of the proposed approach.\n\n[1] Unisolver: PDE-Conditional Transformers Are Universal PDE Solvers", "other_strengths_and_weaknesses": "The main concern is lack of novelty.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "N/A", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The authors introduce PDE-Transformer, an improved transformer-based architecture for modeling various physics simulations on regular grids. The method is limited to 2D regular grids. The method is demonstrated on a variety of PDEs, showcasing the generality and performance of the method. The reviewers are mixed. Three reviewers are in favor of accepting the work, while one objects to it. The three positive reviewers appreciate the paper’s quality, methodology, and added results. The fourth reviewer based some of the criticism on comparisons to very recent arxiv papers. Overall, based on the reviews, I recommend the acceptance of the paper to ICML.
Decision: Accept (poster)