P3D: Highly Scalable 3D Neural Surrogates for Physics Simulations with Global Context
Authors: Benjamin Holzschuh
Georg Kohl
Florian Redinger
Nils Thuerey
Venue: ICLR 2026 Poster (2026)
Keywords: neural surrogates
physics simulations
transformers
3D
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: We present a scalable framework for learning neural surrogates for high-resolution 3D physics simulations

Abstract: We present a scalable framework for learning deterministic and probabilistic neural surrogates for high-resolution 3D physics simulations. We introduce P3D, a hybrid CNN-Transformer backbone architecture targeted for 3D physics simulations, which significantly outperforms existing architectures in terms of speed and accuracy. Our proposed network can be pretrained on small patches of the simulation domain, which can be fused to obtain a global solution, optionally guided via a scalable sequence-to-sequence model to include long-range dependencies. This setup allows for training large-scale models with reduced memory and compute requirements for high-resolution datasets. We evaluate our backbone architecture against a large set of baseline methods with the objective to simultaneously learn 14 different types of PDE dynamics in 3D. We demonstrate how to scale our model to high-resolution isotropic turbulence with spatial resolutions of up to $512^3$. Finally, we show the versatility of our architecture by training it as a diffusion model to produce probabilistic samples of highly turbulent 3D channel flows across varying Reynolds numbers, accurately capturing the underlying flow statistics.

Published as a conference paper at ICLR 2026
P3D: HIGHLY SCALABLE 3D NEURAL SURROGATES FOR PHYSICS SIMULATIONS WITH GLOBAL CONTEXT
Benjamin Holzschuh"
Georg Kohl
Florian Redinger
Nils Thuerey
School of Computation, Information and Technology Technical University of Munich
ABSTRACT
We present scalable framework for learning deterministic and probabilistic neural surrogates for high-resolution 3D physics simulations. We introduce P3D, hy- brid Transforsfrr backbone architecture targeted for 3D physics simulations, which significantly outperforms existing architectures in terms of speed and accu- racy. Our proposed network can be pretrained on small patches of the simulation domain, which can be fused to obtain global solution, optionally guided via a scal- able
model to include long-range dependencies. This setup allows for training large-scale models with reduced memory and compute require- ments for high nigh-resolution datasets. We evaluate our backbone architecture against a large set of baseline methods with the objective simultancously learn 14 different types of PDE dynamics in 3D. We demonstrate how to scale our model to high- resolution isotropic turbulence with spatial resolutions of up to 5123. Finally, we show the versatility of our architecture by training it as diffusion model to produce probabilistic samples of highly turbulent 3D channel flows across varying Reynolds numbers, accurately capturing the underlying flow statistics. Our source code is available at /tttps://githuss com/ tttptbs//pppppppppppppppppeesset wttprr-oppppeeeetfffooeeeeee
INTRODUCTION
Training neural networks on high-resolution data substantially increases the required GPU memory and compute costs. Scaling models and their input dimensions typically requires substantial engineer- ing effort, posing a major barrier to the widespread and cost-effective adoption of machine learning across application domains. Scientific machine learning and engineering are especially affected due to the multi-scale nature of relevant phenomena whose modeling often requires specialized and highly computationally demanding numerical solutions. this paper, we focus on learning surrogate models for simulations focusing on fluid dynamics that have downstream applications in fields such as acrospace (Arranz et al., 2024). climate science (Bodnar et al. 2024), energy systems (Degrave al. 2022), and biomedical engineering (Morris et al. 2016). Machine learning models inherently compete with existing solvers, which are often employed to create the reference targets for learning. As such, they need either significantly outperform the corresponding solvers while maintaining an acceptable level of accuracy (Kochkov et al. 2021 Pestourie et al., 2023), or yield solutions where traditional solvers fall short, for example working with noisy (Franz al.) (023) only partial input data (Shu et al. (203)) or by providing uncertainty estimates (Jacobsen et al. 2023). A large fraction of papers in this area address learning problems for either low-dimensional or comparatively smooth data in 2D). In this paper, we deliberately focus on high-resolution 3D) phenomena, covering wide variety of different types of PDEs.
We propose P3D: hybrid CNN- INN-Transformer backbone built on PDE- DDEntsaasfrrrer (Holzschuh et al. 2025) that combines fast processing local features via convolutions and windowed attention mechanisms for learning generalizable token representations. P3D makes important extensions for scaling to very high resolutions in 3D: crop-based pre-training combined with a designated module for processing global dependencies. This design achieves superior scalability and accuracy compared to existing 3D baselines. We evaluate the P3D backbone architecture through an extensive comparison
"Correspondence ben jenjaminn hol i.schuhttum

---

Published as a conference paper at ICLR 2026
with existing architectures for modeling physics simulations in 3D, training 14 different types of dynamics simultaneously. We then employ P3D as a surrogate model for isotropic turbulence at resolution 512 and demonstrate how the P3D architecture pretrained on crops of 128 can be scaled to the entire domain, while achieving high accuracy with temporally stable autoregressive rollouts.
A significant difficulty when modeling large systems is how aggregate and distribute information globally across the network. We propose to link the bottleneck layers of P3D with a sequence-to- sequence model, called context model, for an efficient global processing and information aggregation utilizing highly optimized self-attention mechanisms in LLM layers, which is combined with different scalable and memory-awar training and inference strategies. Additionally, we propose a direct mechanism to let aggregated global information flow back to the decoder modules via adaptive instance normalization layers. In our final experiment, we train P3D as a diffusion model to learn the distribution of velocity and pressure fields of turbulent channel flow on non-equidistant grid. This setup requires access to global information like the relative position to the walls, and cannot be addressed by learning local representations alone. We verify that velocity profiles from the generated samples of P3D closely match the ground truth, demonstrating that high distributional accuracy can be obtained even when the solution fields are constructed from smaller regions, which are only coordinated through the information flow from the context model. To summarize our contributions:
We introduce P3D, a hybrid CNN-Transformer architecture for autoregressive prediction of high-resolution physics simulations in 3D, combining convolutions for fast learning of local features and windowed self-attention for deep representation learning.
We demonstrate the efficiency and versatility of P3D in three experiments: (1) jointly learning multiple types of simulations (2) scaling P3D pretrained on crops of 128 to a high-resolution simulation of isotropic turbulence 512 and, (3) generating probabilistic samples from P3D trained via flow matching for the velocity and pressure fields of a turbulent channel flow, closely matching the ground truth flow statisties.
We propose mechanisms for efficient global information processing, including linking global context model and injecting global
bottleneck layers with information into adaptive instance normalization layers via region tokens.
We evaluate different setups for finetuming, which enable a more fine-grained control of precomputation and gradient backpropagation through encoder and decoder blocks to reduce VRAM and compute requirements.
RELATED WORK
PDEs and machine learning Machine learning has sparked much resonance in accelerating and
improving numerical PDE solvers as well as fully replacing them (Kim et al. 2019 Thuerey et al.) 2020). Approaches that are combined with existing PDE solvers can replace components of the solver
(Bar-Sinai al. 2019). learn closure models (Duraisamy et al., 20199; Siingnano & MacArt 2023) or learn corrections to a fast solver on a coarse grid (Um et al. 2220 Kochkov et al. [021] Dresdner al. 2023). Other directions target problems that are more difncult to address with numerical
solvers, such as inverse problems (Raissi et al. 2019 Bruna et al. 2022 Holzschuh et al.) (023).
or uncertainty quantification (Xiao & Cinnella 2019 Liu & Thuerey [2024). [002)) A variety of tailored
learned representations was likewise proposed for PDEs Hagnberger al. (2024); Koupaï al.
(2025). Leveraging diffusion models for autoregressive prediction and inverse problems for PDEs
has been explored by several works in the past (Lippe al. 2023 Kohl al. 2024 Shu et al.
2023
Shehata et al. 2025). albeit limited to data in 2D.
Pretraining and 3D neural surrogates Subramanian et al. (2023) Yang et al. (2023) ); McCabe et al. (2023) pretrain foundation models for PDEs on multiple PDE dynamics simultaneously. Such models can be used either for zero-shot-predd sero-shot-reedictions or for finetuning when adapting to new dynamics, which allows for improved results with fewer training samples. Previous works have demonstrated
learning PDEs in 3D. e.g., learning large-eddy simulations (Peng et al. 2023) Li et al., (2023a] 2024a Jiang et al. 2025). , and elastic wave propagation (Lehmann et al.) 2024). Smoke buoyancy in 3D was targeted as a test case by Li et al. (2023b) Notably, most previous work targets resolutions of up to 643, an exception being probabilistic experiments at a resolution of 1283 (Molinaro et al. 2024). and
surrogate training with up to 128% x 128 x 256 (Ohana al.) 2024). This motivates contributions for scalable architectures, as surrogate models for truly high-resolution 3D physics simulations are of paramount interest in different scientific areas.
2

---

Published as a conference paper at ICLR 2026
Embedding
Input
Conv Block
X2
Layer
BoDonncmmpl
Convolutional Encoder
Coro Dnraneampl
Conv Block
X2
3D Window Attention
Transformer
Encoder
Skp connection
Scale
Code
Transformer Siage
x
Dowisomppin ing
Transformer Decoder
Foebiuille Eivennmmple
Laver Norm
Scall, Shift
Transformer Stage
Convolutional Decoder
Conditioning
MLP
Conditioning
Transformer Stage
Scale
Projection
Figure 1: Overview of P3D. Convolutional blocks for local feature processing are combined with transformers for deep representation learning, yielding a U-shaped multi-scale architecture. The transformer backbone combines windowed attention and conditioning via adaptive instance normal- ization, which are modified and optimized for 3D.
Scalable transformer architectures Transformers have become one of the dominant backbone architectures in deep learning due to their high computational efficiency and their ability to model long-range causal relationships (Vaswani et al.] 2017 Devlin et al. 2018). Transformers have also become popular competitor CNNs in vision and understanding tasks (Dosovitskiy al. 2020 Rodrigo et al. 2024). scale to large images (Gupta et al.] 2024), and have recently been adopted for learning surrogate models for physics simulations (McCabe al. 2023 Wu et al. 2024 Alkin et al. 2024). A major computational difficulty is the quadratic complexity of the global self-attention mechanism. Liu et al. (2021) address this limitation by restricting the computation of the attention operation to a local window and Ho et al. (2019) by computing the attention only across the data axes.
3 METHOD
Problem formulation Let Q denote a spatial domain with n physical quantities u(x, t) Ω x [0,7] + R" that are discretized in time and space and described by the temporal sequence [uo, Ust, ...; Ur]. We consider additional information about the sequence such as the type of PDE or hyperparameters of the simulator to be encoded in m-dimensional conditioning vector c E RTm. We assume the availability of many such sequences as training data, representing the temporal evolution of different types of PDEs with varying initial conditions coefficients. We denote our proposed network architecture by Me with weights e.
We address two main tasks in this paper. The first is autoregressive prediction: For given sequence of P preceding states [u, P&t. ..., ut], denoted by Uin our target is to predict the next state u. == Uout- The second task train a probabilistic sampler draw samples from a distribution of states representing solutions for a PDE as specified by the parameter vector c. In this case uin = 0.
3.1 BACKBONE ARCHITECTURE: P3D TRANSFORMER
The key components of our proposed hybrid CNN- Tranfrranformer architecture P3D are the hierarchical U-shape structure with the hybrid encoder and decoder-pair based on convolutional and transformer blocks. A visual overview is given in Figure [ In the following. we highlight the main components of the architecture and explain how they support its central goal to enable the efficient training and inference for high-resolution 3D simulations.
Hybrid encoder/decodee We utilize convolutional en- and decoders with skip connections. Fully transformer-basee architectures that work in the pixel space for 2D data and images like ViTs rely on patchification operation to transform patches of size p X into tokens. A corresponding approach in 3D would transform patches of size pd into a single token, significantly increasing the amount of information encoded in each token. To balance both the number of tokens for the transformer as well as the information density of each token, we learn local features via the convolutional encoder to obtain optimized compressed representation.
3

---

Published as a conference paper at ICLR 2026
Attention and positional encoding The self-attention operation used by transformers has quadratic complexity in the number of tokens. For 3D data, this becomes a major computational issue as the number of tokens grows cubically with the spatial discretization, leading computational blow-up as the domain size increases. The central building block of our transformer encoder is the windowed multi-head self-attention (Liu et al., 202T), which only computes self-attention within local windows. For computing the attention scores between tokens, we use the log-spaced relative positions of tokens inside the same window.
The architecture of the P3D combines Swin transformers (Liu et al. 2021), diffusion transformers (Peebles Xie 2023] DiT) and UNets (Ronneberger et al., 2015) into a 3D variant, which can be seen as an extension of PDE- TDD--Tansformer (Holzschuh et al., 2025) in 3D. The main alteration is the replacement of the patchification with large convolutional en- and decoders. We also removed the shifting of windows during the computation of self-attention. The convolutional en- and de- coder follow the design of modern UNet blocks, using adaptive instance normalization and group normalization. We give a detailed description of the architecture in Section A.4
3.2 CONTEXT NETWORK
The P3D backhone architecture does not include any absolute positional embeddings as well as no opera- tions aggregating and distributing information glob- ally. Thus, has to rely on learning local features and dynamics within its perceptual field. This promotes translation-qqiivrrancc which is an important induc- tive bias for surrogate modeling of PDEs. Our design choice to not include other physical inductive biases intentional to keep the P3D backbone architecture flexible. For learning large-scale simulations, we pre- train P3D on smaller crops of the simulation domain and then scale the pretrained network to larger in- puts. However, global information and long-range dependencies often play crucial role to obtain cor- rect solutions, which is also a major limitation of TDE-Tranfform.r. To address this shortcoming, introduce the context network.
Conditioning
Transformer Encoder
MLP
Skip cornection
Sequence
positoona encoding
-------------------- Latent tokens
Inject global
Transformer Decoder
Global context via a sequence model. The bottleneck layers are connected to the sequence model, which embeds the bottle- neck representation latent tokens. Region tokens are used to inject global information directly into the decoder.
Figure 2:
Token embeddings The bottleneck layer of the transformer encoder consist of tokens, which are em- bedded into latent tokens via linear layer. The P3D encoder compresses crop of size 16 into a single latent token. Then, a frequency- bedding vector is added to each latent token, similar to Dosovitskiy et al. (2020), which constitutes an absolute positional encoding of latent tokens. In ad-
Region tokens
Condittoring
Regions
positional em-
Scale, shift
MLP
dition, we partition the domain into regions and we match the size of regions with the size of the domain crops P3D was pretrained on. For each region, we include corresponding so-called region token in the sequence of latent tokens, similar to the classification token in ViTs. Each region token initialized via a learnable embedding layer and we add frequency- based positional embedding vector. The purpose of region tokens is to serve as a more direct feedback mechanism to the decoder, which we describe the next paragraph. Our implementation processes the token sequence via n = 6 transformer layers. provies an overview of this setup. In principle, any efficient sequence model can be used. After the sequence of region and latent tokens is processed, the latent tokens are added to the input of the decoder via skip connection. See Section A.5 for a full description.
Figure 3: Conditioning via the region to- kens. The input domain is partitioned into re- gions/crops, each of which has a correspond- ing messenger token. The regions are mod- ulated individually via learnable scale and shifts based on the region tokens.
4

---

Published as a conference paper at ICLR 2026
active gradient backpropagation
inactive gradient backpropagation
noon
(a) full domain
(b) cropped domain
(c) global context
(d) randomized backpropagation
(e) decoder only
Figure 4: Different training and inference setups. (a) shows training the full domain and (b) domain crops. (c) includes the context network for global information processing, which can also be trained by randomly disabling gradient backpropagation for percentage of the encoders and decoders, see (d). In (e) the latent codes from pretrained encoder can be precomputed and only the context network and decoder trained.
Region tokens The region tokens retained and used more direct mechanism let informa- tion flow through the decoder network. Region tokens are initialized as a learnable embedding vector with frequency-basse
positional encoding and are processed in the sequence model together with the latent tokens. Each region token corresponds to crop that is processed independently of other crops by the encoder and decoder blocks. We use scale and shift operations condition the decoder block of each region the region token. Within each decoder block, for each adaptive instance normalization layer with scale or shift operations, we transform the region token via a linear layer to region embedding vector which is added to the embedding vector of the conditioning Each region gets modulated differently based on the region token. This visualized Figure
3.3 SCALING OUTPUT DOMAINS
We consider different setups for training and inference, see Figure 4 which include training on the full domain (a), on crops (b) and different training variants in combination with the context network (c) to (e). While is preferable to train on small crops due compute requirements, for inference, generally want process the full domain. We consider two strategies: (1) we scale to the full domain via the translation equivariance the P3D architecture, i.e., we combine the domain crops and process them single input, and, (2) encode and decode each crop of the full domain independently and combine the network outputs. We tag a model that trained on crops of resolution r and which internally runs inference resolution yd by <.rly>. For example, a network trained on crops of size 643 that is scaled via strategy (1) to resolution 128 is tagged <64 128>, while the same network scaled via strategy (2) tagged <6414>> For strategy (2), if we the context network for communication between the latent codes, we use the tag <XrlXy/>
EXPERIMENTS
We evaluate P3D as well as different scaling and finetuning setups in three experiments. P3D has 3 different configurations: S, and that determine the embedding dimension (32, and 128 respectively) of the first layer. The corresponding models have 11M to 180M parameters. We denote configurations with P3D-S for our model with the config, changed accordingly and L.
4.1 JOINTLY LEARNING MULTIPLE PDES
Our dataset for this task comprises 14 different types of PDE dynamics including Burger's equation, Shivashinsyy...... Gray-Scott, Sohhteeo and many others. The dataset is based on APEBench (Koehler et al. 2024), and full description of each PDE with visualizations can be found in Section B8.2 in the appendix. For all PDEs except Gray-Scott, we consider 60 different simulations with varying initial conditions and PDE-specific parameters such as viscosity, domain extent or diffusivity. For Gray-Scott, periodic simulation contains 20 snapshots discretized at resolution 320 for Gray-Scot and 3843 for all other PDEs. We evaluate and benchmark models random crops of the simulation domain of size 1283 64 and 323 using supervised training, see Section A.1 Even with full information about simulation hyperparameters and the type of PDE, the behavior not fully deterministic quantities
include 10 simulations for each of its hyperparameters. Each spatially

---

Published as a conference paper at ICLR 2026
(ours)
10a1
10
P3D.B (aurs)
P3D-L (ours)
ARMSE
6x 10
ARMSE
6x
Swin3D
FactFormer30
40
P30.5
4x
P1D-5
P3D-B
UNer3D
3x
PSD-L
PID-L
AVIT
10a
10
10
100
Memory Consumption (MB)
Computational Cost (GFLOPs)
Figure 5: Comparison of model accuracy vs. (left) memory usage during backpropagation and (right) computational costs for inference for jointly learning different types of PDEs with crops of size 64 for P3D and baselines.
beyond the cropped regions influence the solution inside it. Simulations have different numbers of channels and zero-pad data with fewer channels than the number of maximum channels NC 3.
Training on cropped data Cropped data has an implicit time-dependent boundary condition, which is not known by the model. This can be seen as an extension and more difficult variant of the multi-phe multi-physics training (Subramanian et al. 2023 Yang et al., 2023] McCabe et al. 2023), where in addition to not knowing the PDE or simulation hyperparameters, the model has to estimate the boundary conditions in data-driven manner. Mathematically, the model Me is trained to regress
argmin E(se. [Bgger crop o Mal crop c) - crop S2trr122]
(I)
where we sample (s,, St+ c) E Drrsin from the training dataset and apply a random cropping
obtain (sprop St)4 The input Uin corresponds to 8t crop and Wout St4Ar. The mapping
crop St + S At not deterministic since the boundary conditions are not prescribed. The model Me
has to learn a prediction that minimizes the prediction error w.r.t. all possible simulation states that are outside the cropped domain, i.e., the optimal prediction for 8 At minimizes
argmin ae [[(rro((++ $| crop( =8
8,
where crop(-) is the crop operation used for st crop, The performance depends on how well the model is able to extrapolate the dynamics outside the cropped input for short prediction horizon.
mRMSE evaluation We evaluate the nRMSE, see Section A.3 and consider wide range of sota base- line architectures. Specifically, we include Swin3D, our own implementation of the Swin architec- ture (Liu et al. 2021) extended to 3D, AVIT, ax- ial vision transformer (McCabe et al. 2002). Adap- tive FNOs (Guibas et al.] 2021 AFNO), Tucker- Factorized FNOs (Kossaifi et al., 2023 TFNO), and con-
Table 1:
Comparison of the nRMSE (x x10-2) the test dataset (averaged over all PDEs).
Method
Crop size
epochs 1000
323
643
1283
TENO
8.46
8.37
FactFormer
6.24
4.62
FactFormer (Li al. (024b). Additionally, sider two different convolutional UNet architectures, UNetConNNeXt as used in Ohana al. (2024) and UNetGerCFD (Molinaro et al., 2024). We train all models for 1000 epochs on four H100 GPUs. We use a fixed learning rate of 2.0 10 for all models with the AdamW optimizer with weight decay 10-15 and batch size 256 in bf m-xieed precision. See Table for the results and Tables B3 to [5r detailed evalu- ations of each PDE type. As train bigger crop sizes, the observed simulation domains becomes larger, increasing the amount of information that
UNetConv
8.59
7.09
UNetGenCFD
7.61
8.04
8.27
AVIT
20.9
25.0
15.1
Swin3D
7.92
7.04
5.03
AFNO
9.95
4.98
4.79
P3D-S
6.27
3.76
3.33
P3D-B
4.69
3.03
2.52
P3D-L
4.13
2.49
2.08
becomes available to the network. At the same time, the relative volume of the boundary becomes smaller, thus decreasing its relative weight and causing an according decrease in the nRMSE. P3D performs best across all crop sizes. Performance significantly improves when scaling the model size from S to L.

---

Published as a conference paper at ICLR 2026
P3D-S < 512>
P3D-S < <128| 1128>
Reference
velocity X
Figure 6: Foreed isotropic turbulence. Prediction the test set at resolution 5123 with autore- gressive rollout of 15 steps. The model is pretrained on patches of size 128 without finetuning on 512 P3D-S <1281512> successfully scales the full domain, whereas for P3D-S <1281128> discontinuities at the borders of crop regions become apparent.
Memory and compute For scaling architecture high-resolution 3D simulations, the memory requirements as well as the inference speed are essential. Transformer architectures have been shown achieve improved accuracy as the number of parameters and floating point operations increase, therefore comparing different architectures needs to take both factors into account. In Figure 5 we compare the nRMSE patch size 643 against (1) the computational cost for inference measured in GFLOPs, and (2) the VRAM consumption in MB for backward pass with batch size 1. The P3D networks provide the best tradeoff between accuracy and computational cost/memory requirements. See Table A2 for additional information.
Ablation on network design We empirically verify our network design and consider: (1) P3D without the transformer backbone, which purely relies on the con- volutional encoder and decoder, denoted P3D-conv, (2) patch-based P3D with linear tokemizer for patch size p = 4 and no convolutional encoder and decoder. denoted P3DD-aachh This baseline (*) extends the mixed channel version of PDE- TDF=-ransformer to 3D. (3) disabling PDE type conditioning by setting O, denoted P3D-no-c, (4) P3D with window shifting en- abled, P3D-shift. Additionally,
Table 2:
Validation MSE (x10-3 3 time- weighted EMA with X = 0.99).
Method
Crop size
epochs 400
323
643
1283
P3D-5-conv
8.33
5.40
3.25
P3D-S-pater
6.48
3.78
2.16
P3D-S-no-c
5.69
2.94
1.41
P3D-S-shift
5.41
2.84
1.37
evaluate different values of the window size w = 2, w = 4 (default) and 8. P3D achieves relative improvement of 16.0% (323), 26.7% % (64") and 37.5% (1283) over 3D PDE-TTaanffrmerr" ), Table Shifting and larger window sizes give no significant improvements here.
P3D-S
5.44
2.77
1.35
P3D-S =2
5.68
2.96
1.49
P3D-S w = 4
5.44
2.77
1.35
P3D-S = 5.44
2.90
1.32
4.2 ISOTROPIC: TURRULENCE
The goal of the this experiment is scale P3D to high-resolution simulation involving complex dynamics. For this, we consider forced isotopic turbulence simulated via direct numerical simulation (DNS) resolution 1024 provided by the John Hopkins Turbulence Database (Perlman al. 2007). The dataset is cropped from the original resolution to 5123 with total of 500 snapshots, which are saved after reaching a statistical stationary state and comprises channels for the velocity X/Y/Z and pressure. The data is split into test and training sets, where the first 420 snapshots are used for training and the last snapshots for testing.
RMSE and enstrophy error for crops 1283 We first evaluate the performance of P3D compared to baselines from the previous experiment trained domain crops of size 1283. We consider an evaluation of autoregressive rollouts from to 15 steps. We evaluate the RMSE and a spectral error based the enstrophy spectrum, which we compute based the vorticity that is derived from the velocity fields of the data, see Section C33 We train all models for 4000 epochs using the same setup in Section 4.1 but reduce the batch size to 32. P3D performs best and achieves the lowest RMSE
7

---

Published as a conference paper at ICLR 2026
for few autoregressive rollout steps. It consistenly has by far the best spectral error, Table and Table [C for an extended evaluation. For many autoregressive rollout steps, the RMSE becomes less informative, the prediction starts deviate from the reference due to the uncertainty from the boundary of the crop regions. See Figures C6land and C7 for visualizations.
and
Scaling P3D to 512s We scale P3D pre- trained crops of size 128 the full do- main 512". Training P3D on cropped data re- quires less than a day on four A100 GPUs. See Section C.4 for details. We compare the per- formance P3D when scaling the network to larger domain sizes using the scaling strate- gies introduced in Section 3.3 P3D (128) 128) processes blocks independently using the orig- inal training resolution. Thus the domain 5123
Table 3: RMSE (x10-²) and the enstrophy spec- trum error (x (x10²) at different autoregressive roll-
out steps the test set with crop size 128
Method
RMSE
Enstrophy
15
15
UNetGenCFD
5.48
67.7
4.2
Swin3D
3.22
24.3
4.0
156
AVIT
9.45
37.7
26.7
112
is split into 64 blocks of size 128 that processed independent of each other. P3D (128|512) processes the full domain, leverag- the architec-
AFNO
3.69
29.8
7.6
190
P3D-S
2.81
28.2
2.15
31.9
ing the mranslatio-equiiariinne ture. See Table an evaluation of the RMSE. Increasing the domain size during inference con- sistently gives improvements. With increasing domain size, the relative volume of the bound- ary shrinks, thus the uncertainty turbulent mo- tions is reduced and networks are able to provide more accurate predictions. Similar to previous results, larger networks improve performance well. Note that the RMSE different between Tables and 4 since do consider longer rollouts in Table which affects the distribution of samples in the test set. In this experiment, due to the isotropic and homogeneous nature of the simulation, we achieve highly accurate results without finetuming or global information.
P3D-B
2.04
31.5
0.72
19.2
Table 4: RMSE (x x10-2) on the test set for P3D
trained on crop size 1283 and evaluated at 512 for different scaling strategies.
Method
Scaling strategy
<1281128>
<1281256>
<128|512>
P3D-S P3D-B
1.90
1.68
1.60
1.38
1.24
Long auto-regressive rollouts
Enstrophy error (x102)
300
To evaluate accuracy over pro- longed periods, we evaluated ad- ditional tests comparing the full- scale an individual crop models for 50 rollout steps. We per-
0.6
200
RMSE
150
100
P3D-5 (128[512)
formed inference with the scaled-
P30-5 (128)1281
up model P3D-S (128) 512) the full domain, while for evalua- tion, computed the spectral er- ror and RMSE on the 128 crops ensure the metries are directly comparable to those in Table 3
Mean-field prediction
0
0.0
0
20
40
40 Auto-regressive rollouts
20
Auto-regressive rollouts
Figure 7: Enstrophy spectrum error (x10²) and RMSE over auto-
regressive rollout steps at resolution 5123.
P3D-S (128 512) maintains an excellent enstrophy spectrum error below 10 102 for the full period, whereas P3D-S (128)128), which processes individual regions, achieves this for only 10 steps. We note that the 5123 domain itself cropped from an even larger 1024 resolution, introducing boundary uncertainties that complicate auto-regressive unrolling, even for the scaled-up model. Not being able to reliably compare small scale structures, the RMSE values are still slightly lower for the scaled-up (128) 512) model. This P3D-S version outperforms the mean-field prediction for up to 40 to steps. In particular, the enstrophy comparison highlights the advantages of applying P3D models pre- pre-trained on 1283 crops the full 512 domain for ruto-segressive rollouts.
8

---

Published as a conference paper at ICLR 2026
P3D-L pretrained on 48
P3D-L <48 1922
P3D-L <X48 <X48|X483 X488>
xesiie
Velacit X
Figure 8: Turbulent channel flow. We pretrain PDE-L on crops of size 48 For P3D-L <481192> scaled to the full domain 96² x 192, the relative positions of crops to the wall is critical. Without absolute positional encodings and global information passing, the generated samples are incorrect. P3D-T. <X48IX48 finetuned with the context network resolves this issue.
4.3 TURBULENT CHANNEL FLOW
For the last experiment we train P3D generative model for learning sample from turbulent channel flow simulation with periodic channel with no-slip boundaries at +y and is driven by dynamic forcing to prevent loss of energy. The seemingly simply geometry represents well-studied and relevant scenario that is highly challenging as it requires fine spatial and temporal discretizations with correspondingly long simulation times to provide converged turbulence statisties (Hoyas & Jiménez 2008).
Due to the very costly initial transient phase of these simulations. it is especially attractive to phrase the turbulent channel flow (TCF) prob- lem probabilistic learning problem, where states from the relevant equilibrium phase can be sampled directly, i.e. without resolving the initial warm-up phase. We generate a dataset comprising 20 simulations with Reynolds num- bers Re within the interval [400, 800] spaced
Reynolds Number
400
500
600
700
800
Generated Samples
Reference
0.6
0.6
Velorhy
0.4
0.4
25
50
75
25
50
75
equidistantly, see Section D.I. The computa- tional grid comprises 96 x 96 x 192 spatially adaptive cells with finer discretization near the wall. The data contains channels for the velocity in X/Y/Z. direction as well as pressure.
Channel Index
Channel Index
Figure 9: Flow statistics of samples produced by P3D-L L conditioned with varying Reynolds number (left) and the time-resolved DNS (right).
P3D-L pretrained on 483 We train P3D-L as diffusion model following Section A.2 on crops of size 483. Pretraining crops requires significantly less compute and VRAM and the network converges much faster. For P3D-L with scaling strat- egy <481192>, the network is aware of the position of the wall relative to the position of the individual regions, causing the global structure of the generated flow to incorrect, see Figure X For quantitative statistical eval- uation including training on the full domain and baselines, see Table
Table 5: L2 error for the mean (x10-5) and variance o² (x 10-3) of the velocity profiles along the channel flow direc- tion, see Section D.2 for details. For finetuning, we report the mean and std. dev. over different seeds.
Model
Mean 1.2
Variance 1.2
VRAM
epochs
UNetQ full domain
132,38
17.66
17.4 GB
400
AFNO domain
1849.3
3.4
P3D-L full domain
3.02
13,20
14,9 GB
400
P3D-I <481192>
5862.81
233.77
2.8 GB
2000
finetune w/o region tokens
4541 :: 495
2026 = 267
15.8 GB
20
finetune
23.6 21.4
21.4
40.4 49.44
15.8 1588B
20
finetune, decoder only
1170
6.0 GB
finetune, decoder only
97.7 + 102.6
131 + 149
6.0 GB
100
finetune, decoder only
16.8 2.
24.1 17.2
6.
500
Ablation finetuning and context We evaluate finetuning P3D pretrained on the cropped domain
38 using the context model. see Table Finetuning corresponds to P3D-L <X48IX48> with different training setups, see Figure finetune and and finetune w/o region tokens follows the setup of Figure

---

Published as a conference paper at ICLR 2026
which backpropagates gradients through all crops and network modules. While this requires more VRAM than training on the full domain, only few epochs to achieve good results are necessary. For finetume. decoder only, only backpropagate gradients through the decoder and context network, and also only backpropagate through 10% randomly selected decoder blocks, which corresponds to the setup in Figure 4p. We do precompute the bottleneck representations from the frozen encoders, but this could be done to further reduce the VRAM and compute requirements. This setup achieves significant reduction in VRAM, but requires more training epochs. The generated sample for P3D-L <X48IX48>, see Figure 8 qualitatively shows the correct global structure and content of individual regions. See also Figure 9 and Figures Do D comparisons of the velocity profiles. While finetuning matches the statistics well, there are still visible discontinuities between generated crop regions, which leaves further room for improvement.
Benchmarking speedups We compare the runtime per sample (s) and the mean L2 error of the different methods, see Figure 10 We also compare against running the sim- ulation at half resolution 48 96. The warmup phase of the DNS takes on average 2 hours and 24 minutes, which simulate total time of = 20. To calculate statistics, we take samples with At 1 avoid high autocorrelation. The reference simulation at full resolution has runtime of 454s per sam- ple. For the generative models, we consider 20 samples for each Reynolds number and 100 inference steps. Timings are obtained on a L40 GPU. P and AFNO represent the Pareto frontier for the generative models regarding the mean L2 error.
10
Simulation
Runtime per sample
10
UNetguroe
P3D-L
cannot be skipped. After the warumup,
100
AFNO
10a
10a
10
Mean L2
Figure 10: Runtime per sample (s) vs.
Mcan L2 error, Table 5 w.r.t. to reference simulation at 962 x 192.
5 LIMITATIONS
P3D relies on grids to achieve high computational efficiency. Extending P3D to arbitrary geometries not trivial and requires careful modification of each architectural layer. Also, training P3D on datasets featuring diverse data modalities, resolutions, and dimensionalities remains unaddressed task. The finetuning process using the context network decodes each region independently. which cause discontinuities at the boundary of each region. Further strategies to prevent and mitigate discontinuities are an interesting direction for future work.
6 CONCLUSION
We have presented P3D: efficient hybrid CNN- CNN-Transformer architecture for learning surrogates for high-resolution 3D physics simulations. We demonstrated the strong advantages of P3D over a comprehensive list of baselines for simultaneously learning different types of PDE dynamics, showing improved accuracy and stable training while at the same time being faster and more memory efficient than the strongest competitors. We scaled P3D to a high-resolution isotropic turbulence simulation by pretraining on smaller crops from the domain, and demonstrated its capabilities as a probabilistic generative model. The P3D model accurately predicts distributions of high- high-resolution velocity and pressure fields for turbulent channel flow with varying Reynolds numbers, demonstrating how to include global information and coordinated pretrained networks via global context model. Our architecture enables the scaling of neural surrogate models to very high resolutions, unlocking their potential to deliver weal-world impact across scientific domains.
ACKNOWLEDGEMENT
This work was supported by the ERC Consolidator Grant SpaTe (CoG-2019-863850) The authors also gratefully acknowledge the support of the Erlangen National High Performance Computing Center (NHR @ FAU) of the Friedrich-Alexander Alexander-Univerrstati BаyernK project B278BB. BayerKII funding is provided by Bavarian state authorities. NHR funding is provided by federal and Bavarian state authorities. NHR @ FAU hardware is partially funded by the German Research Foundation (DFG) 44044440880 446199633.
Universität Grlangen-Nürmeed (FAU) under the
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Benedikt Alkin, Andreas Fürst, Simon Schmid, Lukas Gruber, Markus Holzleitner, and Johannes Brandstetter. Universal physies transformers: A framework for efficiently scaling neural operators. Advances in Neural Information Processing Systems, 37:25152-25194 2024.
Gonzalo Arranz, Yuenong Ling, Sam Costa, Konrad Goc, and Adrian Lozano-Duran. Building-bledd flow model for computational fluids. arXiv:2433.09000 2024.
Yohai Bar-Sinai, Stephan Hoyer, Jason Hickey, and Michael Brenner. Learning data driven discretizations for partial differential equations. Proceedings of the National Academy of Sciences, 116631:115344-554) 2019. ISSN 0027-8424, 1091-6490. doi: 10.1073/pnas. 1810588116
Cristian Bodnar, Wessel Bruinsma, Ana Lucic, Megan Stanley, Johannes Brandstetter, Patrick Garvan, Maik Riechert, Jonathan Weyn, Haiyu Dong, Anna Vaughan, et al. Aurora: A foundation model of the atmosphere. arXiv preprint 3rXiv:240S.13063, 2024.
Joan Bruna, Benjamin Peherstorfer, and Eric Handen-Fijnden. Neural galerkin scheme with active learning for high-dimensional evolution equations, 2022.
Yifan Chen, Mark Goldstein, Mengjian Hua, Michael S Albergo, Nicholas M Boffi, and Eric Vanden-Fijndens Probabilistic forecasting with stochastic interpolants and föllmer processes. In Proceedings of the s International Conference Machine Learning, pp. 6728-67887777 6756, 2024.
Jonas Degrave. Federico Felici, Buchli, Michael Neunert, Brendan Tracey, Francesco Carpanese, Timo Ewalds, Roland Hafner, Abbas Abdolmaleki, Diego de Las Casas, et al. Magnetic control of tokamak plasmas through deep reinforcement learning. Nature, 602(7897):414 2022.
Jacob Devlin, Ming-Wei Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint wwxivv:810.04888 2018.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprins arXiv:2010.11929, 2020.
Gideon Dresdner, Dmitrii Kochkov, Peter Christian Norgaard, Leonardo Zepeda-Nuned Jamie Smith, Michael Brenner, and Stephan Hoyer. Learning to correct spectral methods for simulating turbulent flows. Transactions Machine Learning Research, 2023. ISSN 2835-8856. URL
net w wdd=NBAAGO
Karthik Duraisamy, Gianluca laccarino, and Heng Xiao. Turbulence modeling in the age of data. Annual Review of Fluid Mechanics, 51(1):357-377, jan 2019. doi: flaidammmrre- FFanurev-ffuudd-01001----000
URL
Patrick Esser. Sumith Kulal, Andreas Blattmann. Rahim Entezari. Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, Dustin Podell, Tim Dockhorn, Zion En- glish, and Robin Rombach. Scaling rectified flow transformers for high-resolution image synthe- sis. In Forty-first International Conference Machine Learning. ICML 2024, Vienna, Austria, July 21-27, 2024. OpenReview.net, 2024. URL http:::/ /ttpn////eepeeeeiewweeo net forum?id FPrUhsQJ5B
Erik Franz, Barbara Solenthaler, and Nils Thuerey. Learning estimate single-view volumetric flow motions without 3d supervision. International Conference on Learning Representations, 2023.
John Guibas, Morteza Mardani, Zongyi Li, Andrew Tao, Anima Anandkumar, and Bryan Catanzaro. Adaptive fourier neural operators: Efficient token mixers for transformers. CoRR, 2021.
Ritwik Gupta, Shufan Li, Tyler Zhu, Jitendra Malik, Trevor Darrell, and Karttikeya Mangalam. xt: Nested tokenization for larger context in large images. In Forty-first International Conference Machine Learning. ICML 2024, Vienna, Austria, July 21-27. 2024. OpenReview.net, 2024. URL
net wtt.:/oruummimoo wttp://ooommmidd/orotttaa

---

Published as a conference paper at ICLR 2026
Jan Hagnberger, Marimuthu Kalimuthu, Daniel Musekamp, and Mathias Niepert. Vectorized con- ditional neural fields: A framework for solving iimedeepndentt parametric partial differential equations. In Ruslan Salakhutdinov, Zico Kolter, Katherine A. Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp (eds.), Forty-first International Conference on Machine Learning, ICMI. 2024, Vienna, Austria, July 21-27, 2024, volume 235 of Proceed- ings of Machine Learning Research, pp. 17189-17223. PMLR Open 2024. URL https: htttp:://proceeiinnge mlr wttps//w......ee////////aaaeeaea
html
Insu Han, Rajesh Jayaram, Amin Karbasi, Vahab Mirrokni, David P Woodruff, and Amir Zandieh. Hyperattention: Long-conters attention in near-linear time, arXiv:2310.05869, 2023.
Jonathan Ho, Nal Kalchbrenner, Dirk Weissenborn, and Tim Salimans. Axial attention in multidi-
mensional transformers, 2019.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 6840-6851. Curran Asso- ciates, Inc., 2020. URL httpsst///poooceeiiiiie www.diiinnnnn.erpppo.. neurips...... cc pccppaper/202//ii1
Benjamin Holzschuh, Simona Vegetti, and Nils Thuerey. Solving inverse physics problems with score matching. Alice Oh. Tristan Naumann, Amir Globerson. Kate Saenko, Moritz Hardt, and Sergey Levine (eds.), Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurIPS 2023, New Orleans, IA, USA, December 10 16, 2023, 2023.
Benjamin J. Holzschuh, Qiang Liu, Georg Kohl, and Nils Thuerey. Pde-transformer: Efficient and versatile transformers for physics simulations. In Aarti Singh, Maryam Fazel, Daniel Hsu, Simon Lacoste-Julien, Felix Berkenkamp, Tegan Maharaj, Kiri Wagstaff, and Jerry Zhu (eds.), Forty- second International Conference Machine Learning. ICML 2025, Vancouver, BC, Canada, July 13-19, 2025, volume 267 of Proceedings of Machine Learning Research. PMLR / OpenReview.net, 2025. URL htttps::/proceeeddiittto ngs ml wwwppppwww..pppppppppppppppp wwwhhhllsschhhhhcccccc html
Sergio Hoyas and Javier Jiménez. Reynolds number effects on the Reynolds-stress budgets in turbulent channels. Physics of Fluids, 20(10): 101511, 10 2008. ISSN 1070-6631. doi: 10.1063/1.000586..
Christian Jacobsen, Yilin Zhuang, and Karthik Duraisamy. Cocogen: Physically-consistent and conditioned score-based generative models for forward and inverse problems. arXiv preprint arXiv:2312. 10527, 2023.
Yuchi Jiang, Zhijie Li, Yunpeng Wang, Huiyu Yang, and Jianchun Wang. An implicit adaptive fourier neural operator for long-term predictions of diree-dimensional turbulence. arXiv preprint 1rXi::2501.12740, 2025.
Byungsoo Kim, Vinicius C. Azevedo, Nils Thuerey, Theodore Kim, Markus H. Gross, and Barbara Solenthaler. Deep fluids: A generative network for parameterized fluid simulations. Comput. Graph. Forum, 38(2):59-70, 2019. doi: 10.1111/CGF.1360 URL
1111/ 1a11/cgf.1o 13619
Dmitrii Kochkov, Jamie A. Smith, Ayya Alieva, Qing Wang, Michael P. Brenner, and Stephan Hoyer. Machine learning accelerated computational fluid dynamics. Proceedings of the National Academy of Sciences, 118021))210118811188 2021. doi: 10. 1a..0777//naa22178811111 URL https
7/ pnas. /ttp://////////11111//////////// /pass//p/.s..///aa////aa/////a.. 21017 2101184111 784118
Felix Koehler, Simon Niedermayr, Rüdiger Westermann, and Nils Thuerey, Apebench: benchmark for autoregressive neural emulators of pdes. In Advances in Neural Information Processing Systems 37 (Datasets and Benchmarks Track), 2024. URL
iwc0gE1 16u
Georg Kohl, Lei Chen, and Nils Thuerey. Cenchmaaking conditioon diffusion models for turbulent flow simulation, 2024. wwl.nntpss/aarxii.rrr/////////00 745
12

---

Published as a conference paper at ICLR 2026
Jean Kossaifi, Nikola Kovachki, Kamyar Azizzadenesheli, and Anima Anandkumar. Multi-grid tensorized fourier neural operator for high-resolution pdes. arXiv preprint arXiv:2310.00/20, 2023.
Armand Kassai Koupaï, Lise Le Boudec, Louis Serrano, and Patrick Gallinari. ENMA: tokenwise autoregression for generative neural PDE operators. CoRR, abx/2506.06158, 2025. doi: 10.48550/
ARXIV.2506.061555 URI. htttstt:: /ttppp::::ddd./ http://bb..ogg1.....55/////r1...
48550/arXiv..
2506 061
Fanny Lehmann, Filippo Gatti, Michaël Bertin, and Didier Clouteau. 3d elastic wave propagation with a factorized fourier neural operator (f-fno). Computer Methods in Applied Mechanics and Engineering, 420: 16718 18, 2024.
Zhijie Li, Wenhui Peng, Zelong Yuan, and Jianchun Wang. Long-term predictions of turbulence by implicit u-net enhanced fourier neural operator. Physies of Fluids, 35(7), 2023a.
Zhijie Li, Tianyuan Liu, Wenhui Peng, Zelong Yuan, and Jianchun Wang. A transformer- based neural operator for large-eddy simulation of turbulence. Physics of Fluids, 36(6), 2024a.
Zijie Li, Dule Shu, and Amir Barati Farimani. Scalable transformer for pde surrogate modeling. Advances in Neural Information Processing Systems, 36:28010-2800 2023b.
Zijie Li, Kazem Meidani, and Amir Barati Farimani. Scalable transformer for pde surrogate modeling. Advances in Neural Information Processing Systems, 36, 2024b.
Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. 2023. URL htttthttt///atitit wwwpxxiii.orra//o. abs 221..22242111
Phillip Lippe. Bas Veeling, Paris Perdikaris, Richard Turner, and Johannes Brandstetter. Pde- refiner: Achieving accurate long rollouts with neural pde solvers. Advances in Neural Information Processing Systems, 0127398--6433, 2023.
Qiang Liu and Nils Thuerey. Uncertainty-aware surrogate models for airfoil flow simulations with denoising diffusion probabiliiticc models. AIAA Journal, 62(8):2912-299 2933, 2024. doi: 10.2514/1.J063440 /RRLattss: org 10.1000000 2514/1111111111 J043440
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023. URL net
htttps://ppenreriii www.:nne///ooooooooooooo htt////oummmmm//TT/// nww5z
Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEF/CVF International Conference on Computer Vision (ICCV), 2021.
Michael McCabe, Bruno Régaldo-Saiee Saint Blancard, Liam Parker, Ruben Ohana, Miles Cranmer, Alberto Bietti, Michael Eickenberg, Siavash Golkar, Geraud Krawezik, Francois Lanusse, Mariel Pettee, Tiberiu Tesileanu, Kyunghyun Cho, and Shirley Ho. Multiple physics pretraining for physical surrogate models. In NeurIPS 2023 AI for Science Workshop, 2023. URL https //openreview.oo net/ Eorum 1id22ddd-1-- 1
Roberto Molinaro, Samuel Lanthaler, Bogdan Raonic, Tobias Rohner, Victor Armegioiu, Zhong Yi Wan, Fei Sha, Siddhartha Mishra, and Leonardo Zépeda-Náñez. Generative AI for fast and accurate statistical computation of fluids. CoRR, abs/2409, 18359, 2024. doi: 10.4850/ARXXV.2409.8855 URL https: /oi. www.aoo.......ooooooo..ooooooe.. 10. (11550/axXi10 wtw.rriii..4..90.. 2409. 18359
Paul D Morris, Andrew Narracott, Hendrik von Tengg-Kobligk, Daniel Alejandro Silva Soto, Sarah Hsiao, Angela Lungu, Paul Evans, Neil W Bressloff, Patricia V Lawford. D Rodney Hose, et al. Computational fluid dynamics modelling in cardiovascular medicine. Heart, 102(1):18-22 2016.
Ruben Ohana, Michael McCabe, Lucas Thibaut Meyer, Rudy Morel, Fruzsina Julia Agocs, Miguel Beneitez, Marsha Berger, Blakesley Burkhart, Stuart B. Dalziel. Drummond Buschman Fielding, Daniel Fortunato, Jared A. Goldberg, Keiya Hirashima, Yan-Fei Jiang, Rich Kerswell, Surya- narayana Maddu, Jonah M. Miller, Payel Mukhopadhyay, Stefan S. Nixon, Jeff Shen, Romain
13

---

Published as a conference paper at ICLR 2026
Watteaux, Bruno Régaldo-Sain Saint Blancard, François Rozet, Liam Holden Parker, Miles Cranmer, and Shirley Ho. The well: a large-scale collection of diverse physics simulations for machine learning. In Advances in Neural Information Processing Systems 37 (Datasets and Benchmarks Track), 2024. URL
net/ 0=00x5577055
John E. Pearson. Complex patterns a simple system. Science, 26111111006611111111111111119 1993. doi: 00.112//cciecce2211.51110
Selence.261.5111
William Peebles and Saining Xie. Scalable diffusion models with transformers. In 2023 IEEE/CVF International Conference on Computer Vision (ICCV), pp. 4172-4182, 2023. doi: 10.1109/ ICCV51000.2023.00033
Wenhui Peng, Zelong Yuan, Zhijie Li, and Jianchun Wang. Linear attention coupled fourier neural operator for simulation of dinee-dimensional turbulence. Physics of Fluids, 35(1), 2023.
Eric A. Perlman, Randal C. Burns, Yi Li, and Charles Meneveau. Data exploration of turbulence
simulations using a database cluster. In Proceedings of the ACM/IEEE Conference on High Performance Networking and Computing, SC 2007, pp. 23, 2007. doi: 10.1145/136222211366655
URL htt.tttpp://doo www.ddoi.orrrgooooooooooooooee 10. 1145/114513555 622. 3626544
Raphaël Pestourie, Youssef Mroueh, Chris Rackauckas, Payel Das, and Steven Johnson. Physics- enhanced deep surrogates for partial differential equations. Nature Machine Intelligence, 5(12): 1458-1465, 2023.
Maziar Raissi, Paris Perdikaris, and George Karniadakis. 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-7077 2019.
Marcos Rodrigo, Carlos Cuevas, and Narciso García. Comprehensive comparison between vision transformers and convolutional neural networks for face recognition tasks. Scientific reports, 14 (1):21392, 2024.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In Medical Image Computing and Computer-Assisee Intervention-MICC 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18, pp. 234-241. Springer, 2015.
Youssef Shehata, Benjamin Holzschuh, and Nils Thuerey. Improved sampling of diffusion models in fluid dynamics with tweedie's formula. In The Thirteenth International Conference on Learning Representations, 2025.
Dule Shu, Zijie Li, and Amir Barati Farimani. A physics-informed diffusion model for high-fidelity flow field reconstruction. Journal of Computational Physics, 478: 18:1117 19972, 2023.
Justin Sirignano and Jonathan F. MacArt. Deep learning closure models for large-eddy simulation of flows around bluff bodies. Journal of Fluid Mechanics, 966, jul 2023. doi: 10.1017/jfm.2223.446 URL ponttss://do war.i.org1100 10. 1 01 7 ( fm. 2023. 446
Shashank Subramanian, Peter Harrington, Kurt Keutzer, Wahid Bhimji. Dmitriy Morozov, Michael Mahoney, and Amir Gholami. Towards foundation models for scientific machine learning: Charac- terizing scaling and transfer behavior, 2023.
Nils Thuerey, Konstantin Weißenow, Lukas Prantl, and Xiangyu Hu. Deep learning methods for reynolds-averaged navier-stokes simulations airfoil flows. AIAA Journal, 58(1):25-36, 2020.
doi:
URL URLaatps:://oa wrt/ 10. 2514 J058291
Kiwon Um, Robert Brand, Yun (Raymond) Fei, Philipp Holl, and Nils Thuerey. Solver-in- the-loop: Learning from differentiable physics to interact with iterative pde-solvers. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan- Tien Lin (eds.), Advances in Neural Information Processing Systems 33: Anmual Con- ference Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020,
2020. / neurips............er......... http://iww..appppp/////////e/eea 133446a6f344 e00671 e 237 137 40e01 9a8 Abst - html
14

---

Published as a conference paper at ICLR 2026
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
Haixu Wu, Huakun Luo, Haowen Wang, Jianmin Wang, and Mingsheng Long. Transolver: fast transformer solver for pdes on general geometries. arXiv preprini arXiv:2402.00336. 2024.
Heng Xiao and Paola Cinnella. Quantification of model uncertainty in rans simulations: A review. Progress in Aerospace Sciences, 108:1-31, 2019.
Liu Yang, Siting Liu, Tingwei Meng, and Stanley Osher. In-context operator learning for differential equation problems. arXiv preprint mxXiv:2304.07993 2023.
15

---

Published as a conference paper at ICLR 2026
A TRAINING METHODOLOGY AND NETWORK ARCHITECTURES
A.1 SUPERVISED TRAINING
For tasks that have deterministic solution, such as training surrogate model for numerical solver, the P3D can be trained in a supervised manner using mean squared error (MSE) loss, enabling fast, single-step inference. In this case the network is directly optimized with the MSE
Ls = E [IMe c) -uou
(3)
A.2 PROBABILISTIC TRAINING
To allow for sampling from the full posterior distribution rather than approximating an averaged outcome, diffusion training is preferable for probabílistic solutions. For such cases, we employ the flow matching formulation of diffusion models (Lipman et al. 2023 Liu et al. 2023 Ho et al.
2020 Given the imput un and a conditioning vector c, samples Xy drawn from noise distribution Po (0, Z) are transformed into samples from the posterior distribution
by solving an
ordinary differential equation (ODE) of the form dx,, = v(x,, t); dt. Then the model MM learns the velocity field v by regressing a vector field that defines probabilistic trajectory from po to pr. Samples along this trajectory are produced through the forward process
X Wout [1 - min )t}
(4)
for t E [0, 1] with e ~ (0, and a time threshold of m = 10 10- 4. The velocity u can be regressed by training via
CFM = E [[| (um, [[]
c,, Hout + (1 - min
(5)
After training, samples can be drawn from the posterior conditioned on uin and c, by sampling Xg (O, 1) and integrating the corresponding ODE dx, M(um, M1, C, I) di over the time interval = 0 to t = 1. We typically employ explicit Euler steps with a suitable, chosen step size At.
A.3 NORMALIZED RMSE
The normalized RMSE (nRMSE) is defined as
M RRMSE = M Σ
MSE( Hout: Uout)
MSE(O, uout.)
where Nout is the network prediction and M corresponds to the number of trajectories for a specific PDE in the test dataset. In our evaluation, we calculate the nRMSE for each PDE dataset and report the nRMSE averaged over all PDE datasets.
We can also define the nRMSE for autoregressive rollouts over the entire simulation trajectory. We define the nRMSE, at time / by comparing the predicted state ú, for given system at time with the reference u²a We average over all systems S for each PDE dataset.
A.4 P3D ARCHITECTURE
We provide additional details the backbone architecture of P3D below.
Embedding of time, class labels and physical parameters We combine all three types of con- ditionings within a combined embedding layer. Time for flow matching/diffusson and physical parameters are implemented via timestep embeddings. Class labels are implemented via label em- beddings. The embedding vectors of all three types are added and used as the joint embedding. The embedding dimension for each in the convolutional encoder/decoder is 64. In our experiments, class labels used in Section 4.1 Physical parameters and timestep embeddings are used in Section 4.3
16

---

Published as a conference paper at ICLR 2026
Convolutional encoder The convolution encoder first embeds the input using Conv3D layer (kernel size 3. padding 1) with filters that correspond to the embedding dimensions of the configuration. This is followed by downsampling layers implemented via Conv3D layers (kerell size 3, padding 1, stride 2). Intermediate states before each downsampling operations are saved for residual connections. Encoder blocks and consecutive downsampling are applied twice. For each layer, the corresponding number of filters is shown in Table [A]] Encoder blocks are repeated twice. Each encoder block consists of GroupNormalization layers, followed by GELU activations, Conv3D layer (kernel size 3, padding 1), GroupNormalizationn modulation via shift and scale operations depending on the conditioning, GELU and an additional Conv3D layer (kernel size 3, padding 1). The input and output of each encoder block are connected via skip connections. The shift and scale vectors are learned via linear layers from the embedding vectors of the convolutional encoder/decoder.
Convolutional decoder The design of the convolutional decoder mirrors the convolutional encoder in a U-shape architecture with residual connections. Upsampling layers are implemented via a combination of Conv3D layers to increase the number of filters and PixelShuffle3D layers. For an input number of channels Cin and a target number of channels in the upsampled output Cout, the Conv3D operation first expands the number of channels Cin H x x D -> Cout H x x D and PixelShuffle3D spatially rearranges the pixels 8 Cout X H x W x D > Cout x 2H x 2W x 2D.
A.5
CONTEXT MODEL
The context module is cunsformer-basee architecture for processing 3D volumetric data in sequence-like fashion. The model's design and operation can be broken down into the following key components:
Latent tokens The model accepts a 5D tensor of shape (B, C, H, W, D) with batch size B, token embedding dimension C, height H, width W, and depth D, corresponding the token representations of the transformer encoder block arranged on a 3D spatial grid according to their positions. This input is first reshaped into 2D sequence of tokens with shape (B, (H x W D), by flattening the spatial dimensions (H, W, D) into a single sequence length. This token sequence is then passed through an input projection layer map it to the model's internal hidden size de 512. The new tokens are called latent tokens.
Region tokens The model utilizes learnable region tokens, which are a set of nn. Pa rameter tensors. The number of tokens is defined by the number of regions. The region tokens are initialized from a normal distribution with a standard deviation of 10-6, Regions tokens are appended to the input sequence and are used obtain embeddings for each region, which injected into the adaptive instance normalization layers of the decoder modules.
Positional encoding Positional information is incorporated using 3D sine-cosine positional em- beddings. The embeddings added to the input token sequence after the initial projection and concatenation with region tokens. The positional embedding for a given position (w, h, d) and dimension i is calculated by concatenating three separate ID sine-cosine embeddings for each spatial axis. The 1D positional embedding at position and dimension are calculated via:
PE(p, 21) = sin
(7)
T2//do
PE(p, 2i + 1) cos
(8)
(2//dp)
where dp = dc/3 is the embedding dimension per axis, and = 10 is a temperature hyperparameter. Both latent tokens and regions tokens have their own positional encoding.
Core architectural blocks The sequence of tokens is processed by stack of nt = 6 transformer blocks. Each block consists of:
RMSNorm: A root mean square normalization layer that normalizes the input to unit RMS, followed by a learned scaling factor. is a more computationally efficient alternative to traditional layer normalization.
17

---

Published as a conference paper at ICLR 2026
Self-Attention: A multi-head self-attention module that computes attention weights to model relationships between tokens in the input sequence. This block receives the input query tokens and calculates query, key, and value vectors. There is no causal masking. We use hyper attention (Han et al. (223), this is a specialized attention variant designed for efficient sequence processing.
MLP: A gated multi-layer perceptron with GELU activation function.
Output The processed latent tokens are projected back to the original token size C' via linear layer and reshaped to match the input of the context network. The resulting tensor is added to the input of the transformer decoder, forming skip connection. For each adaptive instance normalization module in the decoder, the regions tokens are embedded via unique linear layer. The region embeddings are added the input of the adaptive instance normalization layer of their corresponding region, resulting in different modulation for each region.
Table Al: Different configurations S, and L of P3D. Table shows the total number of weights, the number of filters within the convolutional encoder/decoder and the number of groups for GroupNor- malization layers.
Configuration Number parameters
Embedding dimensions
Number groups
11.2M
[32, 32, 64]
46.2M
[64, 128, 128]
181M
[128, 256, 256]
32
A.6 BASELINE ARCHITECTURES
Performance comparison of baseline architectures Below, we summarize the numbers of the different architectures used in the experiments for an input of size 128 with a single channel, see Table A2
Table A2: Performance comparison of architectures.
Model
Params
GFLOPS
Memory
Throughput
Swin3D
50.3M
289.6
2.9GB
9.86it/s
FactFormer
5.0M
20.4GB
0.82it/s
UNCIGenCFD
100.0M
11039.6
4.8GB
2.03it/s
UNelCunvNext
9.2M
335.6
2.4GB
11.9it/s
TENO
75.6M
139.4
5.1GB
4.7it/s
AVIT
60.0M
143.0
0.3GB
188.3it/s
AFNO
64.1M
2116.4
0.4GB
31.7it/s
P3D-S
11.2M
238.7
0.9GB
35.2it/s
P3D-B
46.1M
2330.6
2.1GB
10.0it/s
P3D-L
181.2M
9276.2
4.6GB
3.7it/s
Regarding the number of parameters, GFLOPs, memory and throughput, the three configurations S. B and of P3D are well positioned compared to the baselines we chose. Timings were obtained on RTX A5000 GPU with batch size 1. Importantly, across all tasks, P3D outperforms the baselines in terms of accuracy. We fixed the training setup for all architectures, using learning rate/optimizer/bathh size that are common for training large transformer and UNet models (Esser et al., (024). Due the number and size of the different models, we do not perform any hyperparameter tuning for individual models.
Swin3D Our implementation of extends the SwinV2 (Liu et al., 2021 transformer in 3D. The code based Sw..v2PreeTrainedMMode
the Python trans transfornee formers package, with pixel shuffling in 3D for upsampling. For the B (default) configuration of Swin3D, we token embedding dimension of 96 with patch size 2, four stages with two layers each, window size 4 and 12 attention heads for each stage.
18

---

Published as a conference paper at ICLR 2026
FactFormer We the official code repository of FactFormer https://qithuttt com/ BaraatrLabFFactoorner We adapt the FactFormer3h 3D module used for the 3D Navier- Stokes experiment (Li al. 2023b). Hyperparameters are taken from the configuration file in the repository: di dim=1188 o dii__mmm=66,,,... heads=6.
UNDDGenCFD We use the official code repository for UNDD bttps://qithub: com/ camlab-ethz etteetthh///nDD GenCFD, We remove any modifications of the core UNet structure spe- cific to the diftusion training setup in Molinaro et al. (2024). Hyperparameters are taken from the repository: rum_channel [128, 256, 256], num num_blocks=un
2],
eeleeeiiiruuii circular', wee_posittioneenoodiiee
and num heads-8.
UNetComNNoot We use the implementation of Neet with ConvNeXt blocks provided by the GitHub repository https: /top:://gttt htt.tttbbbb.....o/ thub. com/ Pol wwwtticct/ttteeeeee.e chitecture was used as a benchmark Ohana et al. (204).. Hyperparameters are kept the same: apatial_resolutttto pepeer_neckkieee b, 3
This ar-
[128, 128, 128]. stages-4, 4,
TENO
We use the implementation of TENO (Kossaifi et al. 2023) from the Python neura lop package https: com neura operator neura operator Hy- perparameters arc: cootiial_eesluttion [128, 128, 128], modes 1-1 6, modes 2-16, modes_3-1 3-16, hidden_channeeleeeeeeeeeeeeeeee
AVIT We use the implementation of axial vision transformers (Ho et al., 2019) provided by the GitHub repository /atps:/// com/ molmmtthtiiA/theee well This architec- ture was used as a benchmark in Ohana et al. (2024). Hyperparameters for the B configu- ration kept the same: spatial resolution- [126, 128, 128], hidden. hidden_dim=765 768
num num_heads=12, processor sittesssorbloccsst blocks=8.
AFNO We use the implementation of AFNO (Guibas et al.] 2021 provided by the GitHub repository https: /ttp://ggthbbb/t com. Poly htt://oooaaatiiiittrreeraaaaaaae
well This architecture was used as a benchmark in Ohana et al. 2024). Hyperparameters for the configuration are kept the same: spatiial resolutions in= [128, 128, 128], wddddenddmmete h_ blockks=12, cmlp_diagona blocks=8 and padcchtc__te pat.c_ssiz=8..
19

---

Published as a conference paper at ICLR 2026
B EXPERIMENT 1: JOINTLY LEARNING PDES
The datasets used in this experiment were carefully selected for a diverse list of partial differential equations (PDEs) with a focus on high spatial resolution. This combined with variations of the physical parameters of the PDEs and different initial conditions, creating diverse dynamics across the different simulations for each type of PDE.
The dataset encompasses linear, reaction- reaction-diffusion diffusion, and nonlinear PDEs. We utilized the Exponax solver, as detailed by Kochler et al. (2024) in the APEBench benchmark. The solver utilizes Exponential Time Differencing Runge-Kutta (ETDRK) methods. We intentionally opted not to the APEBench dataset directly from the original authors. This decision was driven by our goal to create datasets with enhanced resolution and greater diversity in the underlying physical behaviors, rather than relying solely on variations in initial conditions done in APEBench. It is characteristic of the ETDRK methods that they operate within the Fourier domain. Consequently, their application is limited to scenarios with periodic domains and cannot accommodate complex boundary conditions.
Table B1: Summary of datasets produced for the joint PDE learning task, covering linear, reaction- diffusion, and nonlinear PDEs. The table details the dimensions of each dataset: number of simula- tions (s), time steps (t), fields/channels (f), and spatial dimensions (x, y. z). Beyond the explicitly varied quantities listed for each dataset, the initial conditions for every simulation (8) are also distinct.
Dataset
S
y
z
Varied Quantities across
Test Set
hyp
60
30
1
384
384
384
hyper-diffusivity
E
[50, 60]
fisher
60
30
1
384
384
384
diffusivity, reactivity
E [50, 60]
sh
60
30
384
384
384
reactivity, critical number
E [50, 60]
gs-alpha
10
30
2
320
320
320
initial conditions only
separate: t=100
gs-beta
10
30
2
320
320
320
initial conditions only
separate: s=3, t=100
g8-gamma
10
30
2
320
320
320
initial conditions only
separate: s=3, t=100
g8-delta
10
30
2
320
320
320
initial conditions only
E [8, 10
g8-epsilon
10
30
2
320
320
320
initial conditions only
separate: s=3, 1=100
g8-theta
10
30
2
320
320
320
initial conditions only
e [8, 10
gs-iota
10
30
2
320
320
320
initial conditions only
e [8, 10
gs-kappa
10
30
2
320
320
320
initial conditions only
[8, 10]
burgers
60
30
3
384
384
384
viscosity domain extent, viscosity domain extent
E [50, 60
kdv
60
30
3
384
384
384
E [50, 60
ks
60
30
1
384
384
384
separate: s=5, 1=200
B.1 DATA GENERATION SETUP
key aspect of the simulations, in addition to parameter variation, is the use randomized initial conditions. The standard approach for constructing these conditions involves randomly selecting one of three initialization methods, each providing a unique spectral energy distribution. The first method, a random truncated Fourier series initializer, involves summing multiple Fourier series up cutoff frequency, chosen as a uniformly random integer between 2 and 10 (exclusive of 1 1). The second, the Gaussian random field initializer, produces a pawer-law spectrum Fourier space, where energy diminishes polynomially with the wavenumber; its exponent is uniformly randomly selected from [2.3.3.6[. The third method, the diffused noise initializer, generates tensor of values from normally distributed white noise, subsequently diffusing it. This results in a spectrum that decays exponentially quadratically, with an intensity rate uniformly random from [0.00005,0.00 0.00 After generation, all initializers ensure the initial conditions' values normalized to maximum absolute value of one. For vector quantities, the randomly chosen initializer applied independently to each component.
B.2 PDE TYPES
We make use of Exponential Time Differencing Kutte-Kutta (ETDRK) methods to efficiently simulate different PDEs via Exponax. While the chosen linear PDEs are simple and analytically
20

---

Published as a conference paper at ICLR 2026
solvable, the underlying dynamics essential for complicated PDEs. These linear PDEs be understood as representing a scalar attribute, such as density. Unless stated otherwise, sampling from intervals is consistently performed using uniform random distribution. In the following, describe a range of different three-dimensional PDE problems that are employed in our experiments ranging from linear, particularly challenging, as these cases more closely resemble real-world problems. Table B2 gives an overview of the different PDE equations and dynamics.
reaction-diffusion, non-linear PDEs. The class of non-linear PDEs is
Hyper-Diffusion (hyp) behaves similarly to diffusion, where density dissipates inside periodic domain due to the effects hyper-diffusion. Unlike diffusion, hyper-diftusion does not treat all wavelengths equally, analogous to the relation between dispersion and advection. Higher frequency components are damped even more aggressively compared to normal diffusion, leading to visually stronger blur effect of the density field over time.
Dimensionality: 60, / 30. 1, 384, 384, = 384
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, 1] x [0, 1]
Fields: density
Varied Parameters: hyper-diffusiiity E [0.00005, 0.0005
Validation Set: random 15% split of all sequences from [0, 50
Test Set: all sequences from [50, 60
hyp Simulation
02
-0
Figure B1: Random example simulation from hyp.
21

---

Published as a conference paper at ICLR 2026
Coefficients
Mo 0.002
1/1 0.001
Ç vvaried
1/ varried
r varied
y varied
h varied
f varied
k varied
varied
b -60
a 1.0
M variied
= 0.01
b 1=0.0
==10
1=10 1.0
an - - - o V ( Q
Table B2: Overview of the PDE dynamics. Table adapted from Koehler et al. (022))
81 = 444 (u yuu m - Ie A1 4 ( wM O V
V A(1 & V VAA V) M
1 = 1 Vu Fnne f(f +mm )
f(1 - uo)
- (u u + VAAAAA v/u 8h
rm - (I: + V) + 2u 12- - nn
o ( V)
2( + m)
Vuo - fnn
= - - V + +1
(V)
= - <(V
vi = vo
PDE
Fu 8t
88
0t
0t
gs gs-al aappha gs beta gs os-delta gs-iota gs-t.heta aset gs-epsilon sa-kappa
Dynamies
burgers
fisher
hyp
Vries káv
ks
Hyper-Diffusion
Korteweg de-
Fisher-KPP KPP
Hohenberg
Gray-Scott
Kuramoto-
Sivashinsky
Burgers
Name
Swift-
22

---

Published as a conference paper at ICLR 2026
Fisher-KPP (fisher) is foundational reaction-diffusee PDE. These systems used model biological or geological processes, often resulting in pattern formation. This equation details how the concentration of substance varies over time and space, influenced by reaction process governed by a reactivity parameter, and its dispersal through diffusion, which is defined by diffusivity parameter. Its applications extend various domains, including wave propagation, population dynamics, ecology, and plasma physics. Figure B2 shows example visualizations from f f sher.
Dimensionality: 60, t = 30, f = 1, 384, 384, 384
Initial Conditions: random truncated Fourier / Gaussian random field diffused noise (with
clamping to [0, 1D)
Boundary Conditions: periodic
Time Step of Stored Data: 0.005
Spatial Domain Size of Simulation: [0, [0, 1]
Fields: concentration
Varied Parameters: diffusivity и E [0.0001, 0.02| and reactivity r [5, 15
Validation Set: random 15% split of all sequences from E [0,
Test Set: all sequences from [50, 60
fisher Simulation
0.6
0.4
Figure B2: Random example simulation from fisher.
23

---

Published as a conference paper at ICLR 2026
Swift-Hohenberg (sh) is known for depicting various pattern formation processes. This equation can be applied to illustrate the structure wrinkles curved elastic bilayer materials. prime example is the formation of human fingerprints, where tensions between skin layers generate their unique wrinkling. Figure B3 shows example visualizations from sh.
Dimensionality: 60, t = 30, f = 1, = 384, 384, 2 = 384
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, 20T] x 20m]
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, 50
Test Set: all sequences from [50, 60[
sh mulation 40
2000
1.0
0.8
02
1=0
1-3
[-6
1-9
t- 12
1- 15
1 10
t= 21
1 24
1
0.0
Figure B3: Random example simulation from sh.
24

---

Published as a conference paper at ICLR 2026
Gray-Scott (gs) illustrates the dynamic interplay of two reacting and diffusing chemical substances. Substance Sa with concentration Ca is depleted through reaction but resupplied based on defined feed rate. Mcanwhile, substance Sb, the reaction's 's product with concentration ca. is eliminated from the domain at a given kill rate. The balance between these two rates profoundly influences the simulation outcomes, leading diverse stable evolving patterns. We simulate several cases: four with temporally steady configurations, which result in a state that does not substantially change anymore sgs-delta, gs-theta, gs ggs-icta, and ga kappa), and four temporally unsteady configurations, which continuously evolve over time (gs-alpha, gs-beta, gs-qamma, and gs-epsi inn). For the unsteady case, separate test sets with longer temporal rollouts are created. Figure B4 shows example visualizations from the steady configurations, and Figure B5 from the unsteady configurations and corresponding test sets. For further details, see PPersso 1993).
For all simulations, the diffusivity of the substances is fixed to d, 0.00002 and d, 0.00001. In addition to that, datasets are initialized with a Gaussian blob initializer. The initializer creates four Gaussian blobs random positions and variances in the center 60% (20% for gs-kappa) of the Ch.
domain, where the initialization of ca is the complement of ch, i.e. Ca
Steady Configurations (gs-delta, gs-theta, gs-iota, and gs-kappa):
Dimensionality: 10, / = 30, f = 2, x 320, 320, 320 (per configuration)
Initial Conditions: random Gaussian blobs
Boundary Conditions: periodic
Time Step of Simulation: 1.0 (all configurations)
Time Step of Stored Data:
gs-delta: 130.0
gs-theta: 200.0
gs-iota: 240.0
gg- 3a-kappa: 300.0
Number of Warmup Steps (discarded, in time step of data storage):
gs rs-della:: 0
gs-theta: 0
gs-iota: 0
gs- kss-kppa:
Spatial Domain Size of Simulation: 2.5) [0,2.5]
Fields: concentration ca. concentration c
Varied Parameters: feed rate and kill rate determined by configuration (i.e., initial conditions only within configuration)
gs-delta: feed rate: 0.028, kill rate: 0.056
gs-thetar feed rate: 0.040, kill rate: 0.060
gs-iota: feed rate: 0.050, kill rate: 0.0605
gss-kppa: feed rate: 0.052, kill rate: 0.063
Validation Set: random 15% split of all sequences from 8 E [0,8]
Test Set: all sequences from [8, 10]
Unsteady Configurations (gs-alpha, gs-beta, gs-gamma, and gs-epsilon);
Dimensionality: 10, / 30, 2, x 320,
320, = 320 (per configuration)
Initial Conditions: random Gaussian blobs
Boundary Conditions: periodic
Time Step of Simulation: 1.0 (all configurations)
Time Step of Stored Data:
asaalpha: 30.0
gs-beta: 30.0
ga-gamma: 75.0
25

---

Published as a conference paper at ICLR 2026
ps-epsilon: 15.0
Number of Warmup Steps (discarded, in time step of data storage):
as-alphh: 75
gs-beta: 50
gs-ganma: 70
gs-epsilon: 300
Spatial Domain Size of Simulation: [0, [0,2.5]
Fields: concentration Ca. concentration
Varied Parameters: feed rate and kill rate determined by configuration (i.e., initial conditions only within configuration)
gs-alpha: feed rate: 0.008, kill rate: 0.046
gs-beta: feed rate: 0.020, kill rate: 0.046
gammam feed rate: 0.024, kill rate: 0.056
gs-epsilon: feed rate: 0.020, kill rate: 0.056
Validation Set: random 15% split of all sequences from 8 [0, 10
Test Set: separate simulations with 30, 100, f 2, 320, y 320, 320
(per configuration)
26

---

Published as a conference paper at ICLR 2026
02,deta
yske
8
ghetheta Simudaiion
2100
00
MOR
Figure B4: Random example simulations from steady configurations of the Gray-Scott model of system: gs ge, gh-theta.
27

---

Published as a conference paper at ICLR 2026
006
Smulation
Figure B5: Random example simulations from unsteady configurations of the Gray-Scott model of reaction-diffe Paacton-diifusion system: ga-a asaa, gs -beta.
28

---

Published as a conference paper at ICLR 2026
Burgers (burgers) bears resemblance the transport scalar density, this equation describes how a flow field itself evolves due to the combined effects of advection and diffusion. This process result in the formation of abrupt discontinuities, often referred to as shock waves, which present a significant challenge for accurate simulation. Burgers' equation also finds utility in fields such nonlinear acoustics and the modeling of traffic flow. Figure B6 shows example visualizations from burgers.
divection-diffusion problem. Instead of modeling
Dimensionality: 60, t = 30, f = 3, z 384, 384, = 384
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, 1]
Fields: velocity (I, y)
Varied Parameters: viscosity E [0.001, 0.005[
Validation Set: random 15% split of all sequences from 50
Test Set: all sequences from [50, 60
burgers Simuletion 24
0.0
Figure B6: Random example simulation from burgers.
29

---

Published as a conference paper at ICLR 2026
Korteweg-de-Van Vries (kdv) presents simulations of the Korteweg- Vries equation within pe- riodic domain. This equation models dispersive, non-dissipative wave propagation and a classic example of an integrable PDE. It poses challenge because energy is transferred to high spatial frequencies, resulting in distinct, moving soliton waves that maintain their shape and propagation speed. Throughout these simulations, the convection coefficient remains constant -6, and the dispersivity coefficient consistently 1. Figure B7 shows example visualizations from kdv.
Dimensionality: 60, t = 30, f = 2, z 384, 384, 384
Initial Conditions: random truncated Fourier Gaussian random field / diffused noise
Boundary Conditions: periodic
Time Step of Stored Data: 0.05 (with substeps for the simulation)
Spatial Domain Size of Simulation: varied per simulation
Fields: velocity (z, y)
Varied Parameters: domain extent [30, 120 identically for 4, Mr i.e. square domain, and viscosity [0.1, 0.25[
Validation Set: random 15% split of all sequences from E [0, 50
Test Set: all sequences from [50, 60
04
-02
8--
5--3
5-
1-
17
1- 15
t- 10
1-21
t-24
t-
Figure B7: Random example simulation from kdv.
30

---

Published as a conference paper at ICLR 2026
Srrramoto-Sivasinnse (ks) models thermo-diffusiig flame instabilities in combustion and also finds use in -iffuuoonif systems on a periodic domain. It's notable for its chaotic behavior, where even slightly different initial conditions can lead to wildly divergent temporal trajectories time. The initial transient phase of the simulations is not included the dataset. Figure B8 shows example visualizations from ks.
Dimensionality: 60, t 30, f = 1, 1, 384, y 384, 2 384
Initial Conditions: random truncated Fourier Gaussian random field diffused noise
Boundary Conditions: periodic
Time Step of Stored Data: 0.2 (with 2 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 E [10, 130 identically for æ, y, i.c. a square domain
Validation Set: random 15% split of all sequences from & E [0, 600
Test Set: separate simulations with 50, I. 200, f 1, 384, y = 384, 384
10
-20
t-0
t-3
t-6
t-9
t-
t-
t-
t-
1- 24
(-
Figure B8: Random example simulation from ks.
31

---

Published as a conference paper at ICLR 2026
9.2
02
-04
-04
Figure B9: hyp. Reference (left) and autoregressive prediction for t = 8 steps with P3D-L <128138>> (right) the test set at resolution 3843.
B.3 VISUALIZATIONS OF PREDICTIONS
Below, we visualize several example predictions from the P3D-L network trained crop size 1283 on the test datasets. See Figures |9 to B17 During inference, we apply the network to larger crops of domain size 320 for the Gray-Scott PDEs and 3843 for all other PDEs. The resolution is significantly higher than what the network was originally trained on. In addition, since the full domain is periodic, change the padding mode of convolutional layers within the convolutional encoder and decoder from "zeros" to "circular". consider autoregressive rollouts of up to = 8 steps.
32

---

Published as a conference paper at ICLR 2026
X-monn
y-nen
z-menn
0.6
ysto
Fw0
t=7
1=4
Tw6
T=0
0.0
tw2
t=
t=6
0.0
Figure B10: fisher. Reference (left) and autoregressive prediction for steps with P3D-L <1281384> (right) on the test set at resolution 3843.
man
1.
vdm
0.5
03
2-00i
-05
1=2
1=6
Figure B11: sh. Reference (left) and autoregressive prediction for t steps with P3D-L <1281384> (right) on the test set at resolution 3843.
33

---

Published as a conference paper at ICLR 2026
Xmmoo
yroin
-06
-0.
y
2-00
0.0
Figure B12: ga ga-agpllppe alha. Reference (left) and autoregressive prediction for 8 steps with P3D-L < 1281320> (right) the test set resolution 320
X-mon
0.6
0.0
k-tcc
6-4
1-6
1-08
0.0
1=0
1=2
1w4
Figure B13: gs-epsilon Reference (left) and autoregressive prediction for steps with P3D-L <128132>> (right) the test set resolution 3203
34

---

Published as a conference paper at ICLR 2026
100
>-ww
0.4
004
Figure B14 gs-delta. Reference (left) and autoregressive prediction for = 8 steps with P3D-L <1281320> (right) the test set resolution 320
35

---

Published as a conference paper at ICLR 2026
000
2900
1000 20
t-0
1-2
1-4
t-0
1-6
1-2
t-4
C-6
C=
Figure B15; burgers. Reference (left) and autoregressive prediction for t = 8 steps with P3D-L < 1281384> (right) on the test set resolution 3843.
36

---

Published as a conference paper at ICLR 2026
X-7944
upperc
t=
Figure B16: kdv. Reference (left) and autoregressive prediction for = steps with P3D-L < 1281384> (right) the test set resolution 3843.
smmo
15
10
5
2100
Figure B17: ks. Reference (left) and autoregressive prediction for / = 8 steps with P3D-L 281384x (right) the test at resolution 38
37

---

Published as a conference paper at ICLR 2026
Table B3: Normalized RMSE (x10-3) for crop size 128.
Model
GS
GS
Hyp
Burgers
Average
P3D-S
8.3
15.7
29.3
14.1
13.4
47.4
12.9
16.0
6.3
73.4
92.0
18.9
91.4
27:2
33.3
19.0
75.6
P3D-L
6.9
5.7
8.9
5.7
5.4
16.0
6.0
7.1
2.3
69.1
66.8
18.1
62.7
11.1
20.8
AFNO
30.4
81.5
32.4
183.5
47.9
AVIT
124.6
09.7
60.8
51.7
62.7
101.0
47.6
48.1
56.7
619.2
191.2
65.1
470.7
177.0
151.9
Swin3D
11.7
21.3
30.7
22.0
23.0
50.0
22.8
17.5
5.2
153.6
90.1
27.0
190.8
36.7
50.1
UNelewweFF
24.9
43.2
Table B4: Normalized RMSE (x101 (x10-3) for crop size 64.
Model Name
Fisher
GS
GS B
GS
GS
GS
GS
GS x
G8 6
Hyp
KDV
KS
SH
Burgers
Average
P3D.5
21.0
25.4
5.2
18.1
43.3
4.6
4.8
92.2
100.4
50.5
111.8
34.4
37.6
P3D-B
26.2
P3D-I
6.8
7.3
10.1
7.1
6.6
17.5
9.8
8.3
1.4
09.6
73.7
28.5
88.1
15.0
24.9
AFNO
19.7
20.0
22.7
3.1
21.0
38.0
2.3
2.3
189.
90.8
39.4
210.7
36.3
19.9
33.6
Swin3D
20.7
67.4
58.9
5.0
50.7
98.4
4.5
4.5
155.6
115..9
61.2
285.3
46.2
70.4
FactFormer
19.06
18.01
13.91
3.75
16.10
31.76
3.75
3.75
3.75
182.78
95.15
31.53
184.20
08
16.2
UNelowCF
36.6
35.1
72.7
14.0
42.5
65.5
14.7
13.6
13.7
272.4
128.4
37.0
324.4
56.2
80.4
40.3
49.8
36.2
06.8
6.5
190.8
49.0
70.9
TENO
23.2
111.6
56.2
6.3
48.5
96.7
8.0
6.0
6.0
175.5
113.8
222.1
260.1
36.2
83.7
Table B5: Normalized RMSE (x 10 10-3) for crop size 32.
Model Name
Fisher
GS
GS 8
GS
CSS
G8
GS
GS
G8
Hyp
KDV
KSS
SH
Burgers
Average
49.7
33.7
62.7
P3D-I
8.3
24.5
14.8
2.5
15.9
40.1
2.5
2.5
2.5
131.1
126.4
117.5
131.7
37.9
46.9
P3D-T.
7.9
18.5
11.4
1.9
14.22
32.4
1.9
1.9
1.9
109.3
110.8
110.9
129.4
27.3
41.4
AVIT
237.1
126.7
50.7
13.6
107.6
122.7
13.6
13.6
13.8
1089.4
258.7
177.6
476.0
290.1
209.2
Swin3D
16.6
154.5
38.6
3.8
56.2
02.6
3.8
3-8
3.8
189.1
143.6
83-4
263.5
36.7
79.2
UNNtcaccio
18.5
38.6
22.1
3.5
20.9
53.6
3.5
3.5
3.6
196.3
105.8
37.0
517.2
41.9
76.1
20.3
111.6
38.3
7.4
49.2
89.9
7.4
7.4
7.4
166.3
128.5
77.4
442.5
48.9
85.9
TENO
176.8
5.6
252.4
38.3
84.7
38

---

Published as a conference paper at ICLR 2026
C EXPERIMENT 2: ISOTROPIC TURBULENCE
For the second experiment, we make use the Johns Hopkins Turbulence Database (JHTDB). contains data from various direct numerical simulations of homogeneous and wall-bounded turbulent flows (Perlman el al. 2007). The simulations are stored with space-time history and allow for arbitrary spatiotemporal query points.
Isotropic Turbulence (iso-turb) direct numerical simulation of the Navier-Stokes equations at Reynolds number around 433 simulated on a 1024* periodic grid via pseudo-spectral parallel code. It contains homogeneous isotropic turbulence, i.e., idealized version of realistic turbulence with statistical properties that are invariant to translations and rotations of the coordinate axes. The following overview summarizes key characteristics of the dataset (for further details see PPellman et al. 2007):
Dimensionality: = 1, / = 500, / = 4, x = 512, 512, 2 = 512
Boundary conditions: periodic
Time step of stored data: 0.002
Spatial domain size: [0, 2T] [0, 2mJ [0,
Fields: velocity X/Y/Z., pressure
Validation set: random 15% split all timesteps from E [2.4
Test set: all sequences from t [420, 500]
C.I EXTENDED EVALUATION
See Table CI for an extended evaluation expanding upon Table 3 The B config corresponds to the previous size of the baselines models for Swin3D, AVIT and AFNO. The S config decreases the hidden dimensionality of the model architecture, which we change from 768 (B) to 384 (S) for AVIT and AFNO. For Swin3D, we decrease it from 96 (B) 48 (5). For Swin3D and AVIT, the hidden dimensionality corresponds to the token embedding dimension. We also trained an L config for the AVIT model with hidden dimensionality 1536. All baseline models were trained for 4000 epochs. Additionally, we kept training P3D-S and P3D-B for a total of 20000 epochs. The evaluation shows that the RMSE and spectral error keep improving, showing no indication of overfitting.
We also include a benchmark comparison regarding the number of parameters, GFLOPs, VRAM and throughput for inference of P3D and the different baseline architectures Table C2
39

---

Published as a conference paper at ICLR 2026
meb
1-00
5-3
1-2
1-3
L-
1=3
1
1-r
C-8
L-9
Figure Cl: Isotropic Turbulence. Training dataset visualization resolution 512 showing the velocity X/Y/Z. and pressure from / = 0 until / = 420. The first six rows are 2D mean projections and slices of the velocity x-direction. The 3D renderings show velocity X/Y/Z and pressure.
40

---

Published as a conference paper at ICLR 2026
Table CI: Performance comparison on the test set with crop size 128 for RMSE (x10 3) and the 1.2 enstrophy graph error (x 109) at different autoregressive rollout steps.
Model
RMSE
L2 Enstrophy
epochs
1
5
10
15
10
15
UNDD
5.48
25.42
48.60
67.72
4.25
14.5
22.7
140
4000
Swin3D-S
4.04
11.53
18.66
24.30
6.28
59.7
120
158
4000
Swin3D-B
3.22
10.45
18.15
24.36
4.06
52.1
112
156
4000
AViT-S
10.09
21.47
32.60
40.36
31.60
44.1
99.2
144
4000
AViT-B
9.45
19.57
30.00
37.77
26.70
49.1
84.6
112
4000
AViT-L
9.18
18.80
29.76
39.20
24.20
54.6
96.8
129
4000
AFNO-S
4.25
14.80
23.58
30.40
8.30
99.9
171
196
4000
AFNO-B
3.69
13.33
23.52
29.80
7.69
88.7
158
190
4000
P3D-S
2.81
9.87
20.50
28.25
2.15
8.23
21.6
31.9
4000
2.17
8.99
19.40
27.40
1.29
6.68
16.0
25.3
20000
P3D-B
2.04
8.79
20.23
31.52
0.72
1.39
3.38
19.2
4000
1.54
8.11
21.09
44.92
0.21
0.71
3.49
14.7
20000
Table C2: Architecture benchmark comparison crop size 128 measured on H100 GPU using CUDA 12.8. VRAM is measured at inference with batch size Throughput with batch size 16.
Model
Params
GFLOPs
VRAM
Throughput
Swin3D-S
18.9M
79.6
2.2G
51.2it/s
Swin3D-B
50.3M
289.6
2.9G
34.
AViT-S
15.1M
48.8
0.16G
859it/s
AVIT-B
60.0M
143.0
0.3G
346it/s
AViT-L
239.2M
467.0
1.1G
128it/s
AFNO-S
17.0M
942.2
0.2G
167it/s
AFNO-B
64.1M
2116.4
0.4G
58.6it/s
P3D-S
11.2M
238.7
0.9G
88 98/9st/s
P3D-B
46.1M
2330.6
2.1G
30.9it/s
41

---

Published as a conference paper at ICLR 2026
C.2 VORTICITY
The vorticity w, is measure of the local rotation in fluid flow. It is defined as the curl of the velocity vector field (Uz, Myy Wz), where Uny by and Uz are the velocity fields in æ, y and z direction. The vorticity is defined as
Ouz Juy By dz
Ju, Oz
Ouz
Juy dx
Our dy
c1
e2 +
E3
(9)
dx
where 61, e2 and eg represent the unit vectors in æ, y and z direction.
The vorticity is not part of the data. Only the velocity channels in X/Y/Z are available.
Approximation by finite differences To compute the components of the vorticity numerically from the velocity, we use finite difference approximations for the partial derivatives in Equation (
For grid point ji k) located in the interior of the discretized simulation domain, approximate the partial derivates as
Our
1,1. W1
++1jk
~
Oz
2Ar
Our
u 1,j+1
21 1,,
θy
2Ay
Our
u.J.kt
2
Oz
2Ac
and analogously for By and
C.3 ENSTROPHY GRAPH
We consider the magnitude |سي of the vorticity w, which is scalar field defined as
|
w.l2
(10)
The enstrophy of the magnitude of the vorticity field |w) wavenumber R.4, is defined as
Enstrophy(4)
Σ
(a(m))
(11)
ks|m|sk+1
where á with E 2 are the Fourier coefficients of (w/, see Chen et al. 2024 C.2) for reference.
The enstrophy spectrum the graph the function + Enstrophy average the enstrophy spectrum over different vorticity fields.
compute the enstrophy error between the reference enstrophy spectrum and the enstrophy spectrum of generated vorticities
2
Enstrophy
=
Σ OSECK
Enstrophy
avg reference (k) Enstrophy
avg generated( (13)
(12)
K +
with K = 16.
The data on the cropped domain of size 128 is periodic, thus there artifacts at the boundary of the crop. This leads to problems when calculating the Fourier coefficients w We therefore smoothen |سس towards the boundary by multiplying the data with the Hann window as defined in torch. hann_ wowwwnd indow.
Hanning filter
In Figure C2 we show the magnitude of the vorticity |w| generated by P3D-B and the reference after applying the Hanning filter (t autoregressive steps).
See Figure C3 for the enstrophy spectrum Enstrophy generated of the predicted vorticities by P3D-B and the enstrophy spectrum Enstrophyave avg reference the reference for autoregressive unrolling steps I and 15 on the test set.
42

---

Published as a conference paper at ICLR 2026
Prediction Vort Nagnitude (alice (1=b4)
Reference Vorticky Magnituce (slice (i=ba)
320
0.40
120
930
20
0.20
80
0.20
20
20
803
0.00
9.0
200
300
Figure C2: Generated vorticity magnitude |w]] by P3D-B (left) and reference (right) after applying the Hanning filter.
Enstrophy error: 2.10e+01
Enstrophy L2 error: 1.47e+03
10a
Reference
10a
Reference
10a
Prediction
10a
Prediction
10
10c
10a
10
10
10
10
10
10 3
10
10-4
10
10-5
10a
10a
10a
enstrophy spectrum (v.s. k) at time 1
enstrophy spectrum (v.s. k) at time 15
Figure C3: Enstrophy spectrum of P3D-B B and the reference simulation at resolution 1283 for and 15 autoregressive prediction steps.
C.4 TRAINING AND EVALUATION
The model weights for the evaluation in Table are the EMA weights at epoch 2000. We train both the S and B configurations of P3D on crop size 1283, Training P3D-S for 4000 epochs took 1 48m and training P3D-B took 20h 25m on four A100 GPUs. See Figure C4 for the validation loss curve.
In Figure CS we show a comparison between the reference and P3D-S <1281512> with an autoregres- sive rollout until / 16 the test resolution 5123.
43

---

Published as a conference paper at ICLR 2026
P3D-S
P3D-8
Validation MSE
10-2
10-1
500
1000
1500
2000 epoch
2500
3000
3500
4000
Figure C4: Isotropic Turbulence. Validation MSE for P3D-S and P3DDB B during training.
-02
(a) Reference
(b) P3D-S <1281512>
Figure C5: Isotropic Turbulence. Autoregressive prediction fort 16 steps with P3D-S <1281512> on the test set at resolution 5123. The first six rows are 2D mean projections and slices of the velocity in x-direction. The 3D renderings show velocity X/Y/Z and pressure.
44

---

Published as a conference paper at ICLR 2026
2/700
20000
1.
0.4
zdcce
-32
0.
vige
Reference
(b) P3D-S
xm00
14
1.2
0.6
2-iic
wpe
t-0
1-4
1-5
12
1-
9.
1-0
1-4
L-5
1-
1-
0.
(c) P3D-B
(d) UNetGenCPD
Figure C6: Vorticity magnitude |سm| caleulated from the predicted velocity X/Y for different models (1/2).
45

---

Published as a conference paper at ICLR 2026
2/700
1.
0.4
zdic
4.2
0.
vige
Reference
(b) Swin3D-B
xmo0
14
vimoon
1.2
0.6
2-iic
wpp
1-0
1-4
L-5
1-
1-
9.
1-0
1-4
1-5
1-
1-
03
(c) AFNO-B B
(d) AVIT-B B
Figure C7: Vorticity magnitude |سm| caleulated from the predicted velocity X/Y for different models (2/2).
46

---

Published as a conference paper at ICLR 2026
D EXPERIMENT 3: TURBULENT CHANNEL FLOW
The dataset for the turbulent channel flow (TCF) represents a periodic channel with no-slip boundaries dy that is driven by dynamic forcing re-inject energy lost due wall friction, and prevent the flow from slowing down. This results in a continuous production of vortex structures at the walls, which have very characteristic and well-studied, ppatiall-vvarying distribution (Hoyas Jiménez, 2008). Due to the complexity of the flow, these flows require very long transient phases to develop the characteristic structures, We target this scenario by learning with generative model, in the context of which the TCF problem represents probabilistic learning problem to infer turbulent stats from the equilibrium phase, bypassing the costly transient warm-up phase.
D.1 DATASET
We generate a dataset comprising 20 simulations with Reynolds numbers within the interval [400, 800] spaced equidistantly. After the initial- phase, we simulate ETT = 20 eddy turnover cycles, which we save 200 snapshots with At 0.1. The computational grid comprises 96 96 192 spatially adaptive cells with finer discretization near the wall. The data contains channels for the velocity in X/Y/Z direction as well as pressure. The computational grid is shown in Figure Di D3 we show visualizations of the turbulent channel flow for Reynold numbers
In Figures D2 and
Re 400 and Re 640 respectively.
Dimensionality: 20, t 200, f 4, x = 96, = 96, 2 192
Initial conditions: noise
Boundary conditions: periodic (x), wall (y.z)
Time step of stored data: 0.1
Number of warmup steps (discarded, in time step of data storage): 200
Spatial domain: [-1, x [-1, 1] x [-
Fields: velocity X/Y/Z. pressure
Varied parameters: Reynolds number E [400, 800]
Validation set: random 15% split of Reynolds number
1.0
4090
407
3.00
Figure D1: Computational grid of the turbulent channel flow simulation. The spatial discretization is refined in the near-wall region to resolve the boundary layer.
47

---

Published as a conference paper at ICLR 2026
1.0
X-meaa
y-mean
0.8
0.6
2-slice
vape
vape-c
vepecc
Figure D2: Turbulent channel flow with Reynolds number 400.
48

---

Published as a conference paper at ICLR 2026
X-mean
y-mean
0.8
0.6
y-siics
0.4
2-slice
vape
0.2
vape-a
0.0
vape-c
19pe.0
-0.
t=
Figure D3: Turbulent channel flow with Reynolds number 640.
49

---

Published as a conference paper at ICLR 2026
Training S, B and L configs We train P3D with the configurations S. B and L for 400 epochs the full domain of size 962 x 192. The training loss is shown Figure D4 is important to choose large architectures in generative modeling setup based flow matching. The network size, specifically the embedding dimension critical for this task with the config reaching significantly lower loss values compared the S config trained with the same number of epochs. All models were trained on 4 A100 GPUs with 80GB VRAM. Training took Ihh 4m, 14h 43m and 27h 55m for the S, B and configs respectively. We train P3D-L both with conditioning on the Reynolds number enabled and with conditioning disabled. We use the former for the visualization in Figure 9 and the latter in the comparison against baselines.
100
P3D-L
P3D-B
6x
Training loss
40
30-0--
myhh WWWm
1x
50
100
150
200
250
300
350 400
epach
Figure D4: Turbulent channel flow. Training curve for different configurations of P3D.
D.2 STATISTICAL EVALUATION
Let (I1, 32, T3) denote the spatial coordinates and u(x, I,, Re) denote the velocity of the flow direction X. The reference simulations reach an equilibrium phase after the initial transient phase of the warmup. Therefore, for the reference simulations, the moments Um (x, Re) should be the same for all i. Additionally, the setup combining periodic boundary conditions and no-slip boundaries for the wall implies that the statistics only depend on the distance to the wall on the flow axis æ1, ie.. um (m=1 12, ), f, Re) is the same for all I2 and 23 inside the domain. Thus is reasonable to consider um (11, Re) and calculate the moments by sampling over 22 23 t. Since the baseline methods cannot be conditioned on the Reynolds number Re, also compute the velocity profile over Rc. We compute the L2 distance between mean (m = 0) and variance o² (m = 1) of the velocity profile graph 21 + Um (I1)& of the time-resolved DNS reference and generated samples. The resulting I.2 distance for the mean l and variance o² is shown in Table 5 which provides meaningful statistical metric evaluate accuracy. We additionally train two baselines, AFNO and UNetGeaCFD with identical training setups as P3D-L on the full domain.
Moments of the flow field We included comparisons the first three moments (mean, variance, skewness) Um of the flow direction (velocity in x x-direction) averaging over the standard deviation when estimating the moments of the reference based on randomly drawing 20 simulations states of the equilibrium phase per Reynolds number as done for the velocity profiles to properly assess how close the generated samples should match the reference. For the finetuned P3D-L.
well. We report
<X48IX48>, picked the best model out of the five finetuned models. See Table DI Overall, P3D-L trained on the full domain and P3D-L <X48IX48> achieve the best results.
Velocity profile comparison See Figure D5 for comparison between the mean channel flow of the reference simulation for different Reynolds number and the mean flow from the generative model P3D-/. trained on the full domain. We show additional comparisons between the different
50

---

Published as a conference paper at ICLR 2026
Table DI: First three moments of the velocity field in the flow direction for the reference, P3D-L and the two baseline methods.
Moment
Mean
Variance
Skewness
Reference
0.5034+0.0007
0.051110.0001
-0.776:10.007
UNC full domain
0,5002
0.0532
-0.723
AFNO full domain
0.5040
0.0930
-0.361
P3D-L full domain P3D-L <X48IX48>
0,5009
0.0513
-0.789
0.5044
0.0510
-0.802
training and inference strategies in Figure D6 Scaling the P3D-L network trained on the small crops of size 48 to the full domain does not work well and results in incorrect velocity profiles. By finetuning with the context network, region crops
coordinate and obtain information about their relative position to each other as well as to the wall. As a results, the flow statistics improve significantly, more closely matching the reference and samples from P3D-L trained the full domain.
We show samples from P3D-L trained the full domain in Figure D7 when applying P3D-L pretrained on the full domain without any finetuning in Figure D8 and with finetuning via the context network and learned region-dependent conditioning in Figure D9 respectively.
Reference
Generated samples
0.8
0.7
0.6
0.5
Reynolds Number
Reynolds Number
Re=400
Re=400
0.2
Re=500
Re=500
Re=600
Re=600
0.1
Re=700
Re=700
Re-800
800
0.0
20
40
60
80
20
40
60
80
Velocity profile (x)
Velocity profile (x)
Figure D5: Comparison of the mean channel flow of the reference and of generated samples from P3D-7 trained the full domain.
51

---

Published as a conference paper at ICLR 2026
Reference
P3D-L full domain
P3D-L =X48|X48>
P3D-L <48/192>
0.8
0.6
Rc-300
0.2
Reference
P3D-L full domain
P3D-L <X48XX88>
P3D-L <48/192>
0.6
Rc-800
0.4
0.2
0.0
Reference
P3D-L domain
P3D-L <X48[X48>
P3D-L <48/192>
0.6
Re=100
0.4
0.0
Reference
full domain
P3D-L <X48|X48>
<48/192>
0.8
0.6
Rc-800
0.4
0.0
Velocity profile (x)
Velocity profile
Velocity profile (x)
Velacity profile (x)
Figure D6: Comparison of the first two moments of the velocity profile between the reference DNS and generated samples from P3D-7. with different training and inference strategies.
52

---

Published as a conference paper at ICLR 2026
x-mean
2meen
0.4
vo2c c0
0.2
vape cl
vape
0.0
voge-c3
02
5-2 2
t-3
t-5
t-7
Figure D7: Samples from P3D-L trained on the full domain at Re 800.
53

---

Published as a conference paper at ICLR 2026
Xmmcon
0.8
0.5
0.4
vape
t= 0
t= 1
t=2
t=3
t= 4
t=5
t=
t=7
Figure D8: Samples from P3D-L <481192> pretrained on crops of size 483. Inference on the full domain at Re 800 produces incorrect samples, as information on relative positions between region crops is not available.
54

---

Published as a conference paper at ICLR 2026
1.0
0.8
0.5
0.4
0 Podte
vape
t= 0
t= 1
t=2
t=3
t= 4
t=
t=
t=7
Figure D9: Samples from P3D-L <X48IX48> pretrained on crops of size 48 and finetuned with the context network. Inference on the full domain at Re 800 produces samples that exhibit the correct flow statistics.
55


Reviews:
[{"review_id": "G9ishL75IB", "reviewer": "Reviewer_1qCk", "summary": "This is a solid engineering effort with many moving parts: a UNet encoder–decoder, local attention, a small global context stack, a region‑token MLP conditioning, and a probabilistic training scheme. I am slightly concerned that the paper is an incremental 3D variant of the earlier \"PDE-Transformer\" approach which the authors cite as a predecessor. I think this is an interesting paper and the insights can prove useful to the community once disentangled. However, I think many of the architectural components are not well-justified by the experiments (particularly those choices which have fallen out of practice in the field), and the current evaluation has a few asymmetries that make the improvements over baselines harder to trust. I appreciate the amount of work that went in, but I would like to see some improvements to these points before acceptance, so that the paper and experiments can present insights in a more effective way to the community. I have compiled some notes below to help you strengthen the justification for your claims and make the evaluation fairer.", "questions": "[Included in Weaknesses]", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "* Comprehensive set of experimental datasets and baselines across different physical regimes.\n* Reasonable scaling argument via crop training and stitched inference.\n* Dataset visualisations look quite polished.\n* The context finetuning idea for channel flow is interesting when wall information is missing.", "weaknesses": "Major:\n\n1. Novelty\n   1. The paper reads as a 3D variant of prior \"PDE-Transformer\" work. It is a bit challenging to identify which architectural modification is the major contribution to the performance improvements, as this is something I would be quite interested in reading about. I recommend you isolate which component it is, and demonstrate effect sizes. Right now there are a lot of moving parts with insufficient ablation support.\n   2. LLMLayer and \"hyper attention\" are mentioned but seemingly never defined despite being part of the \"core architectural blocks\". The only text I found on \"hyper attention\" in the paper is \"We use hyper attention , this is a specialized attention variant designed for efficient long-sequence processing\". Please spell out exactly what these are (I'm assuming these are from another paper you could simply cite), including complexity and memory scaling, or simply use a standard transformer block.\n\n2. Baselines\n   1. The crop axis of the evaluation feels hard to trust. P3D‑L at 32^3 beats Swin3D at 128^3 on the multi‑PDE suite. To me, this suggests under‑tuned baselines or other asymmetries in how the evaluation is being performed. It is perhaps also an issue with this type of evaluation, done by masking the data in a way that is suitable for the presented model. In addition, there is a fixed learning rate and weight regularization across all models, rather than using the reported training settings presented in each baseline's paper - which likely vary based on architecture settings. I would recommend at the very least using the same training settings as presented in baseline papers, or, even better, doing modest hyperparameter sweeps for the strongest baselines. Otherwise, baseline comparisons may lose their scientific value.\n   2. Comparing speed against an accurate DNS is misleading. Comparisons need to be done either (1) on accuracy, with fixed compute budget for all models (likely resulting in using an LES, not a DNS), or, (2) on speed, at a fixed accuracy target, which is likely to result in using a coarser LES at comparable accuracy to your model. See panel A of Figure 1 of https://www.pnas.org/doi/10.1073/pnas.2101784118 for an example of what this should look like. At the moment, you are comparing diagonally in this accuracy-runtime plot, so it is impossible to draw meaningful conclusions. (Nevertheless, I commend the authors for running the DNS on the GPU, as doing a CPU comparison is the other common mistake.)\n   3. For channel flow, from what I can tell, your model is conditioned on Re, but the baselines are not. Either condition baselines or train per‑Re baselines. In general I do not consider providing a model with an additional scalar context to be a novel contribution worthy of highlighting in a comparison. Also, the grid is non‑uniform in the wall‑normal direction; please explain how convolutions and positional information handle this and ensure baselines get the same treatment (assuming this is something you are controlling for).\n   4. In general, I would like to see some spectra for the tasks. It is hard to know if your model is doing well by staring at thumbnail images of the density fields.\n   5. The EMA reporting was a bit confusing to me. Please state which tables use EMA weights.\n\n3. Ablations\n   1. In general, there are many moving parts in this model. Any additional complexity needs to be justified with quantitative evidence, which I do not see as strong enough at the moment.\n   2. Region‑token conditioning: please cite the original AdaIN/FiLM papers (I think AdaIN is mentioned by name, but not cited); please also clarify novelty. It would also be important to quantify effect size on at least two datasets.\n   3. Context token density: each 16^3 crop becomes one latent token in the context model, so it would be useful to show sensitivity to token density (8^3/16^3/32^3) and where accuracy breaks.\n   4. This is a point I am not confident about so please correct me if I am wrong. But I am confused about the positional encoding; the paper states \"P3D does not include any absolute positional embeddings,\" yet the context model injects 3D sine–cosine PEs. But, since sine-cosine PEs are absolute, the global context actually seems to have full positional information. For relative positional encodings you would need to use RoPE.\n   5. Regarding this \"Inactive gradients\" idea, this seems like another complex idea, which is also not completely necessary for the model itself, introduced in the paper without sufficient quantitative justification. I think this would not only result in a biased gradient estimator to save some VRAM (which does not seem to be theoretically explored), but would also introduce extra stochasticity into the training process. Why not use more standard approaches like gradient checkpointing or, e.g., LayerDrop?\n   6. PixelShuffle3D: would be helpful to compare to more standard approaches to see how much this helps. It might be interesting for other contexts but at the moment it is hard to know whether this actually helps.\n   7. The splits vary by time and by simulation index s which is a bit hard to follow. It might be nice to document the splits in a small table.\n\n\nMinor:\n\n1. In general, I felt like there was too much emphasis on the datasets and not enough emphasis on the model. These datasets are indeed visually nice to look at, but this paper is about the model. Only if the figures are part of the story should they be in the main text of the paper. For example, Figure 7, despite the nice 3D renders of this fluid, is completely unnecessary. I would include some spectra there instead so I can actually see how well the model is doing. Or other analysis on the model itself. Not to mention there is no reference flow for this figure, because it is crammed in to include the 3D render.\n2. I felt like some of the math notation was a bit performative. For example, there was this $u(x, t)$ function defined with set theory notation, but then it was immediately dropped, even in the context of defining the $\\mathbf{u}$ fields (it is not concrete how these two are related). If you need this to define the fields, then do it, but otherwise I would drop it.\n3. The appendix feels too long and low-density in terms of information, it is a lot of the same figure generated for different datasets, many of which there is no discernible pattern at all that the reader should be following. The PDF is 40MB as a result of all of the figures. The dataset does not need to be a literal part of the PDF itself. The reader does not need to see all 9 timesteps, sliced through every axis, for a dataset that was (I think) not even introduced first by this paper. \n4. For filtering choices, a Hanning window is applied before some vorticity visuals. Please state clearly whether spectra are computed on filtered or raw fields and show both.\n5. \"weight regularization\" -> \"weight decay\".", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2rYWsu7mrw", "reviewer": "Reviewer_EzRu", "summary": "The paper proposes a hybrid 3D CNN–Transformer surrogate that learns on small crops and fuses them with a lightweight global context sequence model (via adaln) to capture long-range dependencies while keeping memory/computation manageable. It scales to high-resolution 3D flows, and can also be trained probabilistically (using flow-matching) to generate channel-flow fields with faithful statistics across Reynolds numbers.", "questions": "On patch artifacts in Fig. D9. I still notice some patch artifacts. Could this stem from padding choices in the conv layer (e.g., zero vs. reflect vs. circular)?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. An innovative domain-decomposition neural surrogates with an effective and efficient module to fuse in global context, which is a meaningful step towards large-scale PDE modelling\n\n2. The empirical performance of the model is quite strong. After training on local crop it can be deployed to larger domain.\n\n3. Clear ablations isolate the global-context module as the source of gains over crop-only baselines.", "weaknesses": "The paper’s design and \"why it works” intuition are not fully clear to me\n\n1. Decoder/attention choice. The text says a Transformer decoder with “LLMAttention.” Does this mean causal attention over latent/region tokens? If so, why prefer causal over bi-directional for spatial/latent context aggregation (is it solely because of saving up computeduring fine-tuning)? A more in-depth discussion (e.g., compute/memory in backprop, training stability) would clarify the design choice.\n\n2. Region token vs. simpler conditioning. For the region token, did the authors try a simpler variant: feeding a region index/ID embedding (analogous to a timestep/position embedding) directly into the diffusion model via adaln instead of introducing a learnable token?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "MikcSINZYm", "reviewer": "Reviewer_rjE6", "summary": "The authors introduce a new architecture that uses convolution and attention modules in an encoder-decoder UNet structure and combine it with a latent global attention module that is used to exchange global information and inject region specific information into the decoder.\nThey evaluate various training and scaling strategies focusing on three setups: jointly learning multiple PDEs, pretraining on smaller patches and scaling to larger patches at inference, and a probabilistic training scenario.\nAcross all scenarios P3D generally outperforms the compared baselines, showing superior accuracy and efficiency.", "questions": "1. How is data normalization handled? I think this can be a bit nuanced in the mulit-PDE setting: could be normalised per channel per PDE, or per channel across all PDEs with the same channels, or different strategy?\n\n2. How does the P3D-shift ablation perform without the region tokens? It would be interesting to see if the information flow across windows is enough or this global information injection in the decoder is essential.\n\n3. I think an important downstream task of such scalable architectures is finetuning on an unseen PDE: It would be interesting to have experiments that show what advantage pretraining can bring if we e.g. have a new PDE that is extremely expensive to simulate and we can only afford a few samples. Is there positive transfer?\n\n4. Could you provide a breakdown of how the presented P3D model ablations (e.g. Table 2) compare with respect to their efficiency (Params, GFLOPs, Memory, Throughput)?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "## Originality\n- The paper introduces a novel combination of architectural components essential for scaling to larger grids and systematically analyzes their contributions through ablations.\n- The learned region tokens introduced in the bottleneck and the specific conditioning of the decoder regions is, to the best of my knowledge, a novel architectural design that, according to the presented results, helps to propagate global information.\n\n## Quality\n- Experiments are well designed and a rigorous description of the PDE setups as well as the architectural design and training hyperparameters are provided in the Appendix.\n- The authors benchmark against multiple strong baselines and report consistent improvements in accuracy as well as efficiency.\n\n## Clarity\n- The paper’s organization is clear and logical, moving from background and motivation to contributions, ablations and results, which makes the narrative easy to follow.\n- Main concepts are well described and most additional information that is crucial for understanding is provided in the Appendix.\n\n## Significance\n- Very interesting paper for the neural surrogate community, where scaling is essential to move towards real world problems.\n- The authors outline design choices and findings (e.g.: domain-crop pretraining or region token based conditioning) that can be integrated into other methofs. Elements of this paper can therefore can serve as building blocks for future work.", "weaknesses": "1. Lack of uncertainty/variance reporting for the first two experiment setups. While I understand this is computationally expensive, I think it would be an important addition if the authors could at least include std across a few seeds in the main tables where they compare against other architectures.\n\n2. The paper only reports 1-step errors for the mulit-PDE setting in Section 4.1: Including rollout errors or spectral metrics (as done in Section 4.2) would provide deeper insights.\n\n3. Minor inconsistencies:\n    - In Appendix A.5, the latent token dimensions notation is slightly confusing: (B, C, H, W, D) -> (N, (H × W × D), C). Why do you use B for batch and then N?\n    - Section 4.1 mentions a maximum number of 3 channels. However Table B1 lists a maximum of two fields per PDE. Clarification would be helpful.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "UncIbg9Hut", "reviewer": "Reviewer_HKuH", "summary": "The paper proposes a network architecture and strategy for training on high resolution 3D data. The overall approach uses entirely translation equivariant blocks in the encoder and decoder to allow the model to take full advantage of cropped training. The authors then propose multiple approaches for scaling these models to high resolution 3D data. They experimentally evaluate their architecture on a multi-task learning problem, on a high resolution autoregressive task, and on a generation problem.", "questions": "1. Is the \"conv\" model in the ablations in 4.1 replacing the attention layers with more convolutional blocks or is it just removing components? \n2. Could the context network have eliminated the boundary artifacts in 4.2 for <128|128>? \n3. How does this approach compare to baselines in the presence of boundary conditions or geometric influence?\n4. How does the model perform over longer autoregressive rollouts?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 4, "contribution": 3, "strengths": "1. Overall this is a very well written paper. It makes the problem clear and proposes a solution along with varying approaches to implement that solution.\n2. Apart from the one note below, the experiments are well designed and well suited to showing the scaling performance of the proposed method. \n3. This is an enormous challenge in the field and this work is showing promising results on some extremely difficult tasks.", "weaknesses": "Overall, this is a strong paper that makes important progress on a challenging problem in the space. Right now there are a few issues that I'd consider to be large enough that I can't mark it as a clear accept, but with some minor changes, I'd be happy to increase my score. I've included comments on how these issues could be addressed experimentally which would lead to a larger increase, but to be clear, my primary asks to recommend acceptance are the inclusion of rollout metrics and several adjustments to claims and presentation.\n\nMajor:\n\n1. 4.1 is definitely a weak spot in the paper.  The rest of the experiments are strong enough that I'm not marking down too far for this, but either experiments or claims could be adjusted to address several issues:\n    1. The biggest issue is that that the claim of 14+ PDEs does not appear to be true. It's possible I'm missing something, but Appendix B lists 7 PDEs. It's not reasonable to consider Gray-Scott multiple PDEs due to the varying macroscale behavior.\n   2. Second, there is limited evidence of performance as a \"foundation model\". While previous works like Poseidon have shown the conv+swin backbone is a strong basis for foundation models, the experimental evidence provided is insufficient to make that argument. These are largely fairly simple systems defined on periodic cubes, so this is a case where the model has a strong inductive bias and then is only evaluated for multi-task learning on problems that benefit from that inductive bias. To make this claim, I'd recommend including varying boundary conditions or geometric forcings or just generally more real-world inspired data. Otherwise, I'd recommend just being more explicit on what the limitations of the experiment are. \n2. The difference in metrics between table 3 and 4 also makes it difficult to evaluate whether the results in table 4 are \"good\" beyond the accomplishment of running a model at all on such high resolution data. While it doesn't make sense to ask for ML baselines, I think simple baselines like persistence and predicting the field average would help put the results in context. Additionally, I think evaluating step 1 on the same distribution in both spaces would allow for some comparison between the models.\n3. Generally speaking, the paper's claim of stable autoregressive rollouts does not seem well supported.  Metrics are provided for 1 and 15 steps in 4.2. Some evaluation of rollout on the scaled up model is necessary to make claims that the approach can make stable predictions on high resolution. \n4. The paper could engage more with the weaknesses of the method. For instance, many of the architectural choices are not compatible with non-uniform grids so there isn't a natural extension to many realistic problems. \n\nMinor:\n1. For data generated for a specific project, there should really be more precise information about the generating conditions in the appendix, specifically precise mathematical descriptions for each task. \n2. The context network is introduced, but it's not clear if its being used outside of 4.3. Is it being used? Would it have helped with the independent patch model <128|128> in 4.2? Similarly, there's not much comparison of the training approaches in Figure 4. \n3. This might be a personal issue, so I'm not marking down for this either, but it's a bit weird to use llama library names for standard network components, particularly given this is not a NLP submission. It would be more informative to say the standard names for the model components, then mention you use the llama implementations. \n4. Appendix C.1. links to table 4 - it seems to be an extended version of table 3 based on the caption and numbers.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)