SlowGuess/NeurIPS_2023_OCR huggingface dataset:
In terms of computational resources, all experiments have been performed on NVIDIA V100 GPUs with 16GB memory. For the experiments on the KS equation, each model was trained on a single NVIDIA V100 for 1 to 2 days. We note that since the model errors are becoming close to the float32 precision limit, the results may differ on other GPU architectures (e.g. A100), especially when different precision like tensorfloat (TF32) or float16 are used in the matrix multiplications. This can artificially limit the possible accurate rollout time a model can achieve. For reproducibility, we recommend using V100 GPUs. For the 2D Kolmogorov Flow dataset, we parallelized the models across 4 GPUs, with a training time of 2 days. The speed comparison for the 2D Kolmogorov Flow were performed on an NVIDIA A100 GPU with 80GB memory. Overall, the experiments in this paper required roughly 250 GPU days, with additional 400 GPU days for development, hyperparameter search, and the supplementary results in Appendix E.
C PDE-Refiner - Pseudocode
In this section, we provide pseudocode to implement PDE-Refiner in Python with common deep learning frameworks like PyTorch [62] and JAX [6]. The hyperparameters to PDE-Refiner are the number of refinement steps  $K$ , called num_steps in the pseudocode, and the minimum noise standard deviation  $\sigma_{\mathrm{min}}$ , called min_noise_std. Further, the neural operator N0 can be an arbitrary network architecture, such as a U-Net as in our experiments, and is represented by MyNetwork / self.neural_operator in the code.
The dynamics of PDE-Refiner can be implemented via three short functions. The train_step function takes as input a training example of solution  $u(t)$  (named u_t) and the previous solution  $u(t - \Delta t)$  (named u Prev). We uniformly sample the refinement step we want to train, and use the classical MSE objective if  $k = 0$ . Otherwise, we train the model to denoise  $u(t)$ . The loss can be used to calculate gradients and update the parameters with common optimizers. The operation randn_like samples Gaussian noise of the same shape as u_t. Further, for batch-wise inputs, we sample  $k$  for each batch element independently. For inference, we implement the function predict_nextSolution, which iterates through the refinement process of PDE-Refiner. Lastly, to generate a trajectory from an initial condition u_init, the function rollout autoregressively predicts the next solutions. This gives us the following pseudocode:
class PDERefiner: def__init__(self，num_steps,min_noise_std): self.num_steps  $=$  num_steps self.min_noise_std  $=$  min_noise_std self.neural_operator  $=$  MyNetwork(..）   
deftrain_step(self,u_t,u Prev): k  $=$  randint(0,self(num_steps  $+1$  1 if  $\mathrm{k} = = 0$  pred  $=$  self.neural_operator(zeros_like(u_t)，u Prev,k) target  $=$  u_t else: noise_std  $=$  self.min_noise_std \*\* (k/ self(num_steps) noise  $=$  randn_like(u_t) u_t.noised  $=$  u_t  $^+$  noise \*noise_std pred  $=$  self.neural_operator(u_t.noised,u Prev,k) target  $=$  noise loss  $=$  mse(pred,target) return loss   
defpredict_next_solution(self,u Prev): u_hat_t  $=$  self.neural_operator(zeros_like(u Prev)，u Prev,0) for k in range(1,self(num_steps  $+1$  ： noise_std  $=$  self.min_noise_std \*\* (k/ self(num_steps) noise  $=$  randn_like(u_t) u_hat_t.noised  $=$  u_hat_t  $^+$  noise \*noise_std pred  $=$  self.neural_operator(u_hat_t.noised,u Prev,k) u_hat_t  $=$  u_hat_t.noised - pred \*noise_std return u_hat_t   
def rollout(self,u Initial,timesteps): trajectory  $=$  [u INITIAL] for t in range(timesteps): u_hat_t  $=$  self.predict_next_solution(trajecory[-1]) trajectory.append(u_hat_t) return trajectory
As discussed in Section 3.1, PDE-Refiner can be alternatively implemented as a diffusion model. To demonstrate this implementation, we use the Python library diffusers [65] (version 0.15) in the pseudocode below. We create a DDPM scheduler where we set the number of diffusion steps to the number of refinement steps and the prediction type to v_prediction [70]. Further, for simplicity,
we set the betas to the noise variances of PDE-Refiner. We note that in diffusion models and in diffusers, the noise variance  $\sigma_k^2$  at diffusion step  $k$  is calculated as:
$$
\sigma_ {k} ^ {2} = 1 - \bar {\alpha} _ {k} = 1 - \prod_ {\kappa = k} ^ {K} (1 - \beta_ {\kappa}) = 1 - \prod_ {\kappa = k} ^ {K} (1 - \sigma_ {\min } ^ {2 \kappa / K})
$$
Since we generally use few diffusion steps such that the noise variance falls quickly, i.e.  $\sigma_{\mathrm{min}}^{2k / K} \gg \sigma_{\mathrm{min}}^{2(k + 1) / K}$ , the product in above's equation is dominated by the last term  $1 - \sigma_{\mathrm{min}}^{2k / K}$ . Thus, the noise variances in diffusion are  $\sigma_k^2 \approx \sigma_{\mathrm{min}}^{2k / K}$ . Further, for  $k = 0$  and  $k = K$ , the two variances are always the same since the product is 0 or a single element, respectively. If needed, one could correct for the product terms in the intermediate variances. However, as we show in Appendix E.7, PDE-Refiner is robust to small changes in the noise variance and no performance difference was notable. With this in mind, PDE-Refiner can be implemented as follows:
from diffusers.schedulers import DDPMScheduler
class PDERefinerDiffusion: def__init__(self，num_steps,min_noise_std): betas  $=$  [min_noise_std \*\* (k/ num_steps) for k in reversed(range(num_steps + 1))] self}scheduler  $=$  DDPMScheduler(num_train_timesteps=num_steps + 1, trained_betas  $\equiv$  betas, prediction_type  $\equiv$  v_prediction', clip_sample  $\equiv$  False) self.num_steps = num_steps self.neural_operator  $=$  MyNetwork(..）   
deftrain_step(self,u_t,u Prev): k  $=$  randint(0,self.num_steps +1) #The scheduler uses  $t = K$  for first step prediction,and  $t = 0$  for minimum noise. #To be consistent with the presentation in the paper,we keep k and the #scheduler time separate.However,one can also use the scheduler time step #as k directly and acts as conditional input to the neural operator. scheduler_t  $=$  self.num_steps -k noise_factor  $=$  self.scheduler.alphas_cumprod[scheduler_t] signal_factor  $= 1$  -noise_factor noise  $=$  randn_like(u_t) u_t_noised  $=$  self.scheduler.add_noise(u_t, noise,scheduler_t) pred  $=$  self.neural_operator(u_t_noised,u Prev,k) target  $=$  (noise_factor  $\ast \ast$  0.5)\*noise-(signal_factor  $\ast \ast$  0.5)\*u_t loss  $=$  mse(pred,target) return loss   
defpredict_next_solution(self,u Prev): u_hat_t_noised  $=$  randn_like(u Prev) for scheduler_t in self.scheduler.timesteps: k  $=$  self.num_steps-scheduler_t pred  $=$  self.neural_operator(u_hat_t_noised,u Prev,k) out  $=$  self.scheduler_STEP(pred,scheduler_t,u_hat_t_noised) u_hat_t_noised  $=$  outPrev_sample u_hat_t  $=$  u_hat_t_noised return u_hat_t   
def rollout(self,u Initial,timesteps): trajectory  $=$  [u INITIAL] for t in range(timesteps): u_hat_t  $=$  self.predict_next_solution(trajctory[-1]) trajectory.append(u_hat_t) return trajectory
Figure 8: Dataset examples of the Kuramoto-Sivashinsky dataset. The training trajectories are generated with 140 time steps, while the test trajectories consist of 640 time steps. The spatial dimension is uniformly sampled from  $[0.9\cdot 64, 1.1\cdot 64]$ , and the time step in seconds from  $[0.18, 0.22]$ .
Frequency spectrum of Kuramoto-Sivashinsky trajectories
Figure 9: Frequency spectrum of the Kuramoto-Sivashinsky dataset under different precisions. Casting the input data to ffloat32 precision removes the high frequency information due to adding noise with higher amplitude. Neural surrogates trained on ffloat64 did not improve over ffloat32, showing that it does not affect models in practice.
In this section, we provide a detailed description of the data generation, model architecture, and hyperparameters used in our three datasets: Kuramoto-Sivashinsky (KS) equation, parameter-dependent KS equation, and the 2D Kolmogorov flow. Additionally, we provide an overview of all results with corresponding error bars in numerical table form. Lastly, we show example trajectories for each dataset.
D.1 Kuramoto-Sivashinsky 1D dataset
Data generation. We follow the data generation setup of Brandstetter et al. [8], which uses the method of lines with the spatial derivatives computed using the pseudo-spectral method. For each trajectory in our dataset, the first 360 solution steps are truncated and considered as a warmup for the solver. For further details on the data generation setup, we refer to Brandstetter et al. [8].
Our dataset can be reproduced with the public code of Brandstetter et al. [8]. To obtain the training data, the data generation command in the repository needs to be adjusted by setting the number of training samples to 2048, and 0 for both validation and testing. For validation and testing, we increase the rollout time by adding the arguments --nt=1000 --nt_EFFECTive=640 --end_time=200, and setting the number of samples to 128 each. We provide training and test examples in Figure 8.
<sup>4</sup>https://github.com/brandstetter-johannes/LPSDA#produce-datasets-for-kuramoto-shivashinsky-ks-equation
Table 2: Detailed list of layers in the deployed modern U-Net. The parameter channels next to a layer represents the number of feature channels of the layer's output. The U-Net uses the four different channel sizes  $c_{1}, c_{2}, c_{3}, c_{4}$ , which are hyperparameters. The skip connection from earlier layers in a residual block is implemented by concatenating the features before the first GroupNorm. For the specifics of the residual blocks, see Figure 10.
Index	Layer
Encoder
1	Conv(kernel size=3, channels=c1, stride=1)
2	ResidualBlock(channels=c1)
3	ResidualBlock(channels=c1)
4	Conv(kernel size=3, channels=c1, stride=2)
5	ResidualBlock(channels=c2)
6	ResidualBlock(channels=c2)
7	Conv(kernel size=3, channels=c2, stride=2)
8	ResidualBlock(channels=c3)
9	ResidualBlock(channels=c3)
10	Conv(kernel size=3, channels=c3, stride=2)
11	ResidualBlock(channels=c4)
12	ResidualBlock(channels=c4)
Middle block
13	ResidualBlock(channels=c4)
14	ResidualBlock(channels=c4)
Decoder
15	ResidualBlock(channels=c4, skip connection from Layer 12)
16	ResidualBlock(channels=c4, skip connection from Layer 11)
17	ResidualBlock(channels=c3, skip connection from Layer 10)
18	TransposeConvolution(kernel size=4, channels=c3, stride=2)
19	ResidualBlock(channels=c3, skip connection from Layer 9)
20	ResidualBlock(channels=c3, skip connection from Layer 8)
21	ResidualBlock(channels=c2, skip connection from Layer 7)
22	TransposeConvolution(kernel size=4, channels=c3, stride=2)
19	ResidualBlock(channels=c2, skip connection from Layer 6)
20	ResidualBlock(channels=c2, skip connection from Layer 5)
21	ResidualBlock(channels=c1, skip connection from Layer 4)
22	TransposeConvolution(kernel size=4, channels=c3, stride=2)
23	ResidualBlock(channels=c1, skip connection from Layer 3)
24	ResidualBlock(channels=c1, skip connection from Layer 2)
25	ResidualBlock(channels=c1, skip connection from Layer 1)
26	GroupNorm(channels=c1, groups=8)
27	GELU activation
28	Convolution(kernel size=3, channels=1, stride=1)
The data is generated with float64 precision, and afterward converted to float32 precision for storing and training of the neural surrogates. Since we convert the precision in spatial domain, it causes minor artifacts in the frequency spectrum as seen in Figure 9. Specifically, frequencies with wavenumber higher than 60 cannot be adequately represented. Quantizing the solution values in spatial domain introduce high-frequency noise which is greater than the original amplitudes. Training the neural surrogates with float64 precision did not show any performance improvement, besides being significantly more computationally expensive.
Model architecture. For all models in Section 4.1, we use the modern U-Net architecture from Gupta et al. [22], which we detail in Table 2. The U-Net consists of an encoder and decoder, which are implemented via several pre-activation ResNet blocks [25, 26] with skip connections between encoder and decoder blocks. The ResNet block is visualized in Figure 10 and consists of Group Normalization [90], GELU activations [27], and convolutions with kernel size 3. The conditioning parameters  $\Delta t$  and  $\Delta x$  are embedded into feature vector space via sinusoidal embeddings, as for
Figure 10: ResNet block of the modern U-Net [22]. Each block consists of two convolutions with GroupNorm and GELU activations. The conditioning features, which are  $\Delta t$ ,  $\Delta x$  for the KS dataset and additionally  $\nu$  for the parameter-dependent KS dataset, influence the features via a scale-and-shift layer. Residual blocks with different input and output channels use a convolution with kernel size 1 on the residual connection.
Table 3: Hyperparameter overview for the experiments on the KS equation. Hyerparameters have been optimized for the baseline MSE-trained model on the validation dataset, which generally worked well across all models.
Hyperparameter	Value
Input Resolution	256
Number of Epochs	400
Batch size	128
Optimizer	AdamW [52]
Learning rate	CosineScheduler(1e-4 → 1e-6)
Weight Decay	1e-5
Time step	0.8s / 4Δt
Output factor	0.3
Network	Modern U-Net [22]
Hidden size	c1 = 64, c2 = 128, c3 = 256, c4 = 1024
Padding	circular
EMA Decay	0.995
example used in Transformers [84]. We combine the feature vectors via linear layers and integrate them in the U-Net via AdaGN [61, 64] layers, which predicts a scale and shift parameter for each channel applied after the second Group Normalization in each residual block. We represent it as a 'scale-and-shift' layer in Figure 10. We also experimented with adding attention layers in the residual blocks, which, however, did not improve performance noticeably. The implementation of the U-Net architecture can be found in the public code of Gupta et al. [22].<sup>5</sup>
Hyperparameters. We detail the used hyperparameters for all models in Table 3. We train the models for 400 epochs on a batch size of 128 with an AdamW optimizer [52]. One epoch corresponds to iterating through all training sequences and picking 100 random initial conditions each. The learning rate is initialized with 1e-4, and follows a cosine annealing strategy to end with a final learning rate of 1e-6. We did not find learning rate warmup to be needed for our models. For regularization, we use a weight decay of 1e-5. As mentioned in Section 4.1, we train the neural operators to predict 4 time steps ahead via predicting the residual  $\Delta u = u(t) - u(t - 4\Delta t)$ . For better output coverage of the neural network, we normalize the residual to a standard deviation of about 1 by dividing it with 0.3.
$^{5}$ https://github.com/microsoft/pdearena/blob/main/pdearena/modules/conditioned/twod_unet.py
Figure 11: Left: Visualizing the average MSE error over rollouts on the test set for four methods: the baseline MSE-trained model (blue), the pushforward trick (green), the diffusion model with standard cosine scheduling (orange), and PDE-Refiner with 8 refinement steps. The markers indicate the time when the method's average rollout correlation falls below 0.8. The y-axis shows the logarithmic scale of the MSE error. While all models have a similar loss for the first 20 seconds, PDE-Refiner has a much smaller increase of loss afterwards. Right: showing the average correlation over rollout time. Different thresholds would lead to the same findings in this paper.
Thus, the neural operators predict the next time step via  $\hat{u}(t) = u(t - 4\Delta t) + 0.3 \cdot \mathsf{N}\mathsf{O}(u(t - 4\Delta t))$ . We provide an ablation study on the step size in Appendix E.3. For the modern U-Net, we set the hidden sizes to 64, 128, 256, and 1024 on the different levels, following Gupta et al. [22]. This gives the model a parameter count of about 55 million. Crucially, all convolutions use circular padding in the U-Net to account for the periodic domain. Finally, we found that using an exponential moving average (EMA) [41] of the model parameters during validation and testing, as commonly used in diffusion models [30, 38] and generative adversarial networks [16, 91], improves performance and stabilizes the validation performance progress over training iterations across all models. We set the decay rate of the moving average to 0.995, although it did not appear to be a sensitive hyperparameter.
Next, we discuss extra hyperparameters for each method in Figure 3 individually. The history 2 model includes earlier time steps by concatenating  $u(t - 8\Delta t)$  with  $u(t - 4\Delta t)$  over the channel dimension. We implement the model with  $4\times$  parameters by multiplying the hidden size by 2, i.e. use 128, 256, 512, and 2048. This increases the weight matrices by a factor of 4. For the pushforward trick, we follow the public implementation of Brandstetter et al. [9] and increase the probability of replacing the ground truth with a prediction over the first 10 epochs. Additionally, we found it beneficial to use the EMA model weights for creating the predictions, and rolled out the model up to 3 steps. We implemented the Markov Neural Operator following the public code of Li et al. [50]. We performed a hyperparameter search over  $\lambda \in \{0.2, 0.5, 0.8\}$ ,  $\alpha \in \{0.001, 0.01, 0.1\}$ ,  $k \in \{0, 1\}$ , for which we found  $\lambda = 0.5$ ,  $\alpha = 0.01$ ,  $k = 0$  to work best. The error correction during rollout is implemented by performing an FFT on each prediction, setting the amplitude and phase for wavenumber 0 and above 60 to zero, and mapping back to spatial domain via an inverse FFT. For the error prediction, in which one neural operator tries to predict the error of the second operator, we scale the error back to an average standard deviation of 1 to allow for a better output scale of the second U-Net. The DDPM Diffusion model is implemented using the diffusers library [65]. We use a DDPM scheduler with squaredcos_cap_v2 scheduling, a beta range of 1e-4 to 1e-1, and 1000 train time steps. During inference, we set the number of sampling steps to 16 (equally spaced between 0 and 1000) which we found to obtain best results while being more efficient than 1000 steps. For our schedule, we set the betas the same way as shown in the pseudocode of Appendix C. Lastly, we implement PDE-Refiner using the diffusers library [65] as shown in Appendix C. We choose the minimum noise variance  $\sigma_{\mathrm{min}}^2 = 2\mathrm{e - }7$  based on a hyperparameter search on the validation, and provide an ablation study on it in Appendix E.7.
Results. We provide an overview of the results in Figure 3 as table in Table 4. Besides the high-correction time with thresholds 0.8 and 0.9, we also report the one-step MSE error between the prediction  $\hat{u}(t)$  and the ground truth solution  $u(t)$ . A general observation is that the one-step MSE is not a strong indication of the rollout performance. For example, the MSE loss of the history 2 model is twice as low as the baseline's loss, but performs significantly worse in rollout. Similarly,
$^{6}$ https://github.com/brandstetter-johannes/MP-Neural-PDE-Solvers/
<sup>7</sup>https://github.com/neuraloperator/markov_neural_operator/
Table 4: Results of Figure 3 in table form. All standard deviations are reported over 5 seeds excluding Ensemble, which used all 5 baseline model seeds and has thus no standard deviation. Further, we include the average one-step MSE error of each method on the test set. Notably, lower one-step MSE does not necessarily imply longer stable rollouts (e.g. History 2 versus baseline).
Method	Corr. > 0.8 time	Corr. > 0.9 time	One-step MSE
MSE Training
Baseline	75.4 ± 1.1	66.5 ± 0.8	2.70e-08 ± 8.52e-09
History 2	61.7 ± 1.1	54.3 ± 1.8	1.50e-08 ± 1.67e-09
4× parameters	79.7 ± 0.7	71.7 ± 0.7	1.02e-08 ± 4.91e-10
Ensemble	79.7 ± 0.0	72.5 ± 0.0	5.56e-09 ± 0.00e+00
Alternative Losses
Pushforward [9]	75.4 ± 1.1	67.3 ± 1.7	2.76e-08 ± 5.68e-09
Sobolev norm k = 0 [50]	71.4 ± 2.9	62.2 ± 3.9	1.33e-07 ± 8.70e-08
Sobolev norm k = 1 [50]	66.9 ± 1.8	59.3 ± 1.5	1.04e-07 ± 3.28e-08
Sobolev norm k = 2 [50]	8.7 ± 0.9	7.3 ± 0.5	7.84e-04 ± 9.30e-05
Markov Neural Operator [50]	66.6 ± 1.0	58.5 ± 2.1	2.66e-07 ± 1.08e-07
Error correction [58]	74.8 ± 1.1	66.2 ± 0.9	1.46e-08 ± 1.99e-09
Error Prediction	75.7 ± 0.5	67.3 ± 0.6	2.96e-08 ± 2.36e-10
Diffusion Ablations
Diffusion - Standard Scheduler [30]	75.2 ± 1.0	66.9 ± 0.7	3.06e-08 ± 5.24e-10
Diffusion - Our Scheduler	88.9 ± 1.0	79.7 ± 1.1	2.85e-09 ± 1.65e-10
PDE-Refiner
PDE-Refiner - 1 step (ours)	89.8 ± 0.4	80.6 ± 0.2	3.14e-09 ± 2.85e-10
PDE-Refiner - 2 steps (ours)	94.2 ± 0.8	84.2 ± 0.4	5.24e-09 ± 1.54e-10
PDE-Refiner - 3 steps (ours)	97.5 ± 0.5	87.0 ± 0.9	5.80e-09 ± 1.65e-09
PDE-Refiner - 4 steps (ours)	98.3 ± 0.8	87.8 ± 1.6	5.95e-09 ± 1.95e-09
PDE-Refiner - 8 steps (ours)	98.3 ± 0.1	89.0 ± 0.4	6.16e-09 ± 1.48e-09
PDE-Refiner - 3 steps mean (ours)	98.5 ± 0.8	88.6 ± 1.1	1.28e-09 ± 6.27e-11
the Ensemble has a lower one-step error than PDE-Refiner with more than 3 refinement steps, but is almost 20 seconds behind in rollout.
As an additional metric, we visualize in Figure 11 the mean-squared error loss between predictions and ground truth during rollout. In other words, we replace the correlation we usually measure during rollout with the MSE. While PDE-Refiner starts out with similar losses as the baselines for the first 20 seconds, it has a significantly smaller increase in loss afterward. This matches our frequency analysis, where only in later time steps, the non-dominant, high frequencies start to impact the main dynamics. Since PDE-Refiner can model these frequencies in contrast to the baselines, it maintains a smaller error accumulation.
Speed comparison. We provide a speed comparison of an MSE-trained baseline with PDE-Refiner on the KS equation. We time the models on generating the test trajectories (batch size 128, rollout length  $640\Delta t$ ) on an NVIDIA A100 GPU with a 24 core AMD EPYC CPU. We compile the models in PyTorch 2.0 [62], and exclude compilation and data loading time from the runtime. The MSE model requires 2.04 seconds ( $\pm 0.01$ ), while PDE-Refiner with 3 refinement steps takes 8.67 seconds ( $\pm 0.01$ ). In contrast, the classical solver used for data generation requires on average 47.21 seconds per trajectory, showing the significant speed-up of the neural surrogates. However, it should be noted that the solver is implemented on CPU and there may exist faster solvers for the 1D Kuramoto-Sivashinsky equation.
D.2 Parameter-dependent KS dataset
Data generation. We follow the same data generation as in Appendix D.1. To integrate the viscosity  $\nu$ , we multiply the fourth derivative estimate  $u_{xxxx}$  by  $\nu$ . For each training and test trajectory, we uniformly sample  $\nu$  between 0.5 and 1.5. We show the effect of different viscosity terms in Figure 12.
Training trajectories
Test trajectories
Figure 12: Dataset examples of the parameter-dependent Kuramoto-Sivashinsky dataset. The viscosity is noted above each trajectory. The training trajectories are 140 time steps, while the test trajectories are rolled out for 1140 time steps. Lower viscosities generally create more complex, difficult trajectories.
Table 5: Results of Figure 7 in table form. All standard deviations are reported over 5 seeds.
Method	Viscosity	Corr. > 0.8 time	Corr. > 0.9 time
MSE Training	[0.5, 0.7)	41.8 ± 0.4	35.6 ± 0.6
[0.7, 0.9)	57.7 ± 0.6	50.7 ± 1.3
[0.9, 1.1)	73.3 ± 2.3	66.0 ± 2.5
[1.1, 1.3)	88.0 ± 1.5	76.7 ± 2.2
[1.3, 1.5]	97.0 ± 2.7	85.5 ± 2.2
PDE-Refiner	[0.5, 0.7)	53.1 ± 0.4	46.7 ± 0.4
[0.7, 0.9)	71.4 ± 0.3	64.3 ± 0.6
[0.9, 1.1)	94.5 ± 0.6	84.9 ± 0.6
[1.1, 1.3)	112.2 ± 0.9	98.5 ± 1.5
[1.3, 1.5]	130.2 ± 1.5	116.6 ± 0.7
Model architecture. We use the same modern U-Net as in Appendix D.1. The conditioning features consist of  $\Delta t$ ,  $\Delta x$ , and  $\nu$ . For better representation in the sinusoidal embedding, we scale  $\nu$  to the range [0, 100] before embedding it.
Hyperparameters. We reuse the same hyperparameters of Appendix D.1 except reducing the number of epochs to 250. This is since the training dataset is twice as large as the original KS dataset, and the models converge after fewer epochs.