SlowGuess/NeurIPS_2023_OCR huggingface dataset:
SST UNet For the SST dataset, we use a UNet implementation commonly used as backbone architecture of diffusion models<sup>5</sup>. The UNet for the SST dataset consists of three downsampling and three upsampling blocks. Each block consists of two convolutional residual blocks (ResNet blocks), followed by an attention layer, and a downsampling (or upsampling) module. Each ResNet block can be further divided into two sub-blocks. The first one consists of convolution  $\rightarrow$  normalization  $\rightarrow$  time-embedding scale-shift  $\rightarrow$  activation function, and the second sub-block is the same but without the time-embedding scale-shift. The downsampling module is a 2D convolution that halves the spatial size of the input (with a  $4 \times 4$  kernel and stride  $= 2$ ). The upsampling module doubles the spatial size via nearest neighbor upsampling followed by a 2D convolution (with  $3 \times 3$  kernel). At the end of each downsampling (upsampling) block the spatial size is halved (doubled) and the channel dimension is doubled (halved). We use 64 channels for the initial downsampling block, which means that the channel dimensionalities are  $64 \rightarrow 128 \rightarrow 256$  and the spatial dimensions  $(60, 60) \rightarrow (30, 30) \rightarrow (15, 15)$  in the corresponding downsampling blocks (reversed for the upsampling blocks). We use three different dropout rates for the SST UNet: 1) before the query-key-value projection of each attention layer,  $dr_{at}$ ; 2) After the first sub-block of each ResNet block,  $dr_{bl_1}$ ; 3) After the
$^{4}$ https://xskillscore.readthedocs.io/  
 $^{5}$ https://github.com/lucidrains/denoising-diffusion-pytorch/blob/main/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
Table 6: Navier-Stokes results, including on the out-of-distribution (OOD) test set. Left columns are same as in Table 1. As observed by Fig. 12 in [44], the differences across main vs OOD test set are minor.
Method	Test	Out of Distribution
CRPS	MSE	SSR	CRPS	MSE	SSR
Dropout	0.078 ± 0.001	0.027 ± 0.001	0.715 ± 0.005	0.078 ± 0.001	0.027 ± 0.001	0.715 ± 0.006
DDPM	0.180 ± 0.004	0.105 ± 0.010	0.573 ± 0.001	0.189 ± 0.012	0.113 ± 0.013	0.555 ± 0.023
MCVD	0.154 ± 0.043	0.070 ± 0.033	0.524 ± 0.064	0.170 ± 0.046	0.082 ± 0.036	0.489 ± 0.061
DYffusion	0.067 ± 0.003	0.022 ± 0.002	0.877 ± 0.006	0.069 ± 0.002	0.023 ± 0.002	0.882 ± 0.002
(a) Navier-Stokes
(b) Spring Mesh
Figure 6: Comparison against single-step deterministic baselines from [44]. We plot the MSE as a function of the rollout step (predicted horizon). It is clear that both our stochastic multi-step baseline, Dropout-multi-step, and DYffusion outperform the UNet (CNN) models from [44] on the Navier-Stokes (spring mesh) datasets of [44]. As reported in [44], the single-step CNN diverges after a few autoregressive steps on the spring mesh dataset. In addition, DYffusion performs especially well on long-range forecasts relative to the baselines. For spring mesh we plot the three forecasted rollout samples of each model separately due to higher variance between samples.
second sub-block of each ResNet block,  $dr_{bl_2}$ . For all models except the Dropout baseline and  $\mathcal{I}_{\phi}$  in DYfusion, which use higher dropout rates, we use  $dr_{at} = 0.1$ ,  $dr_{bl_1} = 0$ ,  $dr_{bl_2} = 0.3$ .
Navier-Stokes UNet and spring mesh CNN For the Navier-Stokes and spring mesh benchmark datasets from [44], we simply re-use their proposed UNet and CNN architecture for the respective dataset. The only modification is the integration of the SST UNet time embedding module into the architectures of [44], as described below.
Time embedding module The time-embedding scale-shift, taken from the SST UNet, is a key component of all architectures since it enables them to condition on the diffusion step (for DDPM and MCVD) or the dynamical timestep (for the time-conditioned barebone models as well as for both  $F_{\theta}$  and  $\mathcal{I}_{\phi}$  in DYffusion). It is implemented by a sine and cosine-based featurization of the scalar diffusion step/dynamical timestep. These features are projected by a linear layer, followed by a GeLU activation, and another linear layer, which results in a "time embedding". Then, separately for each convolutional (or ResNet) block of the neural architecture the "time embedding" is further processed by a SiLU activation and another linear layer whose output is interpreted as two vectors that are used to scale and shift the block's inputs. In all architectures, the scale-shift operation is performed after the convolution and normalization layers, but before the activation function and dropout layer.
D.4 Additional results
D.4.1 Navier-Stokes out-of-distribution
In Table 6 we show the out-of-distribution test results for the Navier-Stokes datasets (analogously to Table 2 for the spring mesh dataset). As mentioned in the main text and observed by [44], the differences compared to the main test set are marginal.
D.4.2 Benchmark against deterministic models
In this set of experiments, we directly compare our method against the deterministic single-step forecasting baselines presented in the benchmark dataset for the Navier-Stokes and spring mesh experi
Table 7: DYffusion ablation. We change one component at a time in DYffusion, starting with all components enabled (first row). For SST, we perform the ablation only on a subset of the test dataset (box 88 in Fig. 5). For the Navier-Stokes and spring mesh datasets, we use the full test sets. The second to fifth rows are ablations at sampling/inference time: No refinement refers to not using the refinement step in line 6 of Alg. 2. No  $\mathcal{I}_{\phi}$  dropout refers to disabling the inference dropout of the interpolator network,  $\mathcal{I}_{\phi}$ . No  $\mathcal{I}_{\phi}$  dropout &  $\sigma_{\epsilon}$  refers to disabling the inference dropout of  $\mathcal{I}_{\phi}$  and perturbing the inputs by  $\sigma_{\epsilon} = 0.05$  (as for the Perturbation baseline). Naive sampling refers to swapping cold sampling (Alg. 2) with the naive sampling algorithm from [2]. The bottom three rows apply to both training and inference and ablate the choice of the forecaster conditioning  $\mathbf{c}(\mathbf{x}_t,n)$ , where for the last row  $\alpha_{n} = \frac{n}{N - 1}$ . Here, the best choice varies across datasets.
SST	Navier-Stokes	Spring Mesh
CRPS	MSE	SSR	CRPS	MSE	SSR	CRPS	MSE	SSR
Base	0.181 ± 0.002	0.111 ± 0.002	1.04 ± 0.01	0.067 ± 0.003	0.022 ± 0.002	0.88 ± 0.01	0.0103 ± 0.0025	4.2e-4 ± 2.1e-4	1.13 ± 0.08
No refinement	0.181 ± 0.002	0.111 ± 0.002	1.08 ± 0.00	0.069 ± 0.003	0.024 ± 0.002	1.12 ± 0.02	0.0246 ± 0.0012	6.8e-4 ± 2.8e-4	2.05 ± 0.13
No Ωφ dropout	0.320 ± 0.009	0.206 ± 0.012	0.00 ± 0.00	0.098 ± 0.005	0.028 ± 0.003	0.00 ± 0.00	0.0337 ± 0.0038	2.5e-3 ± 6.4e-4	0.01 ± 0.00
No Ωφ dropout & σε	0.308 ± 0.009	0.197 ± 0.012	0.40 ± 0.01	0.096 ± 0.005	0.028 ± 0.003	0.23 ± 0.00	0.0282 ± 0.0031	2.7e-3 ± 6.5e-4	0.98 ± 0.05
Naive sampling	0.681 ± 0.062	0.945 ± 0.117	0.52 ± 0.03	0.088 ± 0.004	0.029 ± 0.002	0.54 ± 0.01	0.0115 ± 0.0035	4.5e-4 ± 2.6e-4	0.76 ± 0.06
c = null	0.182 ± 0.002	0.111 ± 0.002	0.94 ± 0.02	0.067 ± 0.003	0.022 ± 0.002	0.88 ± 0.01	104 ± 111	9.9e+7 ± 1.4e+8	1.75 ± 0.16
c = xt	0.196 ± 0.002	0.128 ± 0.002	1.13 ± 0.01	0.077 ± 0.006	0.028 ± 0.003	0.77 ± 0.02	0.0103 ± 0.0022	4.2e-4 ± 2.1e-4	1.13 ± 0.08
c = αnxt + (1 - αn)e	0.181 ± 0.002	0.111 ± 0.002	1.04 ± 0.01	0.074 ± 0.003	0.026 ± 0.001	0.83 ± 0.02	2.2044 ± 1.6889	3.6e+2 ± 3.5e+2	1.63 ± 0.04
iments [44] in terms of MSE. It is important to emphasize that our method is specifically designed to generate ensemble-based probabilistic forecasts, which is why our baselines and evaluation procedure focus on the probabilistic forecasting setting. Besides, MSE can only be part of the evaluation of a probabilistic forecasting model, since it does not capture the skill of the forecasted distribution like the CRPS or SSR metrics do. In Fig 6 it is visible that the baselines from [44] either degrade or diverge during the inference rollouts for the Navier-Stokes and spring mesh test evaluations, respectively. Both our method as well as our multi-step Dropout baseline outperform [44], especially in the long-range forecasting regime. This is likely because our models are trained to forecast multiple timesteps, while the models from [44] only learn to forecast the next timestep. As a result, the training objective significantly deviates from the evaluation procedure. The struggle with long-term stability and long-range forecasts by the benchmark datasets baselines was already noted by the dataset paper itself [44]. Note that architecturally all these models are almost identical, using the neural backbones proposed in [44].
D.5 Ablations
D.5.1 Cold sampling versus naive sampling
As mentioned in section C.2, we find clear evidence that the DDIM-like cold sampling (Alg. 2) proposed by [2] for "generalized diffusion models" is very important to attain good samples from DYffusion. In Table 7, cold sampling corresponds to the first row, Base, and the Naive sampling row refers to swapping cold sampling (Alg. 2) with the naive sampling algorithm from [2] (see Alg. 4). The finding that naive sampling yields poor results for DYffusion and, more generally, "generalized diffusion models" is consistent with the findings from [2].
D.5.2 Inference dropout in the interpolator network
In Table 7 we show that disabling the inference dropout in the interpolator network,  $\mathcal{I}_{\phi}$ , results in considerably worse scores. This is to be expected, since without stochasticity in  $\mathcal{I}_{\phi}$ , our current framework collapses to generating deterministic forecasts (since the sampling algorithm and forecaster network are deterministic, and we assume that the given initial conditions are fixed). In such a case, computing the CRPS collapses to the mean absolute error, and the SSR becomes 0 since there is no spread in the predictions. To attain an ensemble of forecasts, but keeping the interpolator dropout disabled, we also include an ablation row where we perturb the initial conditions with small random noise  $\epsilon \sim \mathcal{N}(0,\sigma_{\epsilon}\mathbf{I})$ , where we use  $\sigma_{\epsilon} = 0.05$ .
D.5.3 Refining the interpolation forecasts after sampling
We find in Table 7 (No ref. row), that the addition of line 6 to the cold sampling algorithm (see Alg. 2) can sometimes improve performance. However, this is not consistent across datasets: While for the
SST dataset, we hardly observe any difference, the scores degrade considerably when not refining the forecasts for the spring mesh dataset. A reason could be the relatively long training horizon used for spring mesh (134 for spring mesh versus 16 or 7 for Navier-Stokes or SST). In practice, we recommend that practitioners train DYffusion with the refinement step disabled to accelerate inference time (since the refinement step requires an additional forward pass per output timestep). Then, during evaluation it is encouraged to perform inference with DYffusion with both disabled and enabled refinement step to analyze whether enabling the refinement step can meaningfully improve the forecasts.
D.5.4 Forecaster network conditioning
In the bottom three rows of Table 7 we ablate the three forecaster conditioning options (see B for detailed definitions),  $\mathbf{c}(\mathbf{x}_t,n)$ . The optimal choice of  $\mathbf{c}$  varies across datasets. In particular, for the spring mesh experiments, it turns out to be crucial to use the clean initial conditions (variant 2, where  $\mathbf{c}(\mathbf{x}_t,\cdot) = \mathbf{x}_t$ ) as conditioning (the two other variants diverge when forecasting the full test trajectories). A likely explanation for this is that the training horizon of  $h = 134$  is relatively long for spring mesh. As a result, the forecaster network can greatly benefit from being explicitly conditioned on the initial conditions during all diffusion states (as opposed to only being indirectly influenced by  $\mathbf{x}_t$  through the interpolator network). In practice and if compute allows, it may be worthwhile to train DYfusion with at least the first two variants for a few iterations to analyze if any conditioning variant strongly outperforms the others.
D.5.5 Choosing the training horizon
In any multi-step forecasting model, the training horizon,  $h$ , is a key hyperparameter choice. Usually, its choice is constrained by the number of timesteps that fit into GPU memory, and it is expected that larger training horizons will improve performance when evaluated on long (autoregressive) rollouts. However, for continuous-time models including ours, where the number of timesteps needed in GPU memory does not change as a function of  $h$  (see Table 10), the choice of  $h$  is flexible. In Table 8, we explore using three different training hori
Table 8: Navier-Stokes ablation of the training horizon,  $h$ . All methods are evaluated on the 64-step test trajectories. For example, for  $h = 1$  ( $h = 16$ ) the corresponding methods are unrolled autoregressively 64 (4) times.
Method	CRPS	MSE	SSR
Dropout (h=1)	0.132 ± 0.006	0.046 ± 0.006	0.002 ± 0.00
Dropout (h=8)	0.086 ± 0.012	0.026 ± 0.002	0.416 ± 0.29
Dropout (h=16)	0.078 ± 0.001	0.027 ± 0.001	0.715 ± 0.01
Dropout (h=32)	0.078 ± 0.001	0.025 ± 0.001	0.651 ± 0.01
DYffusion (h=8)	0.076 ± 0.002	0.027 ± 0.001	0.701 ± 0.02
DYffusion (h=16)	0.067 ± 0.003	0.022 ± 0.002	0.877 ± 0.01
DYffusion (h=32)	0.075 ± 0.003	0.028 ± 0.001	0.862 ± 0.04
zons  $(h\in \{8,16,32\})$  for the Navier-Stokes dataset for both the barebone time-conditioned Dropout model as well as DYffusion. For DYffusion, this means that we train both an interpolator network as well as a corresponding forecaster net with the same training horizon. Note that  $h = 16$  corresponds to the main results, and is a sweet spot for DYffusion for this task. However, any of the used horizons results in better scores than the best baseline (for any baseline training horizon). We also include a next-step forecasting baseline  $(h = 1)$ , which results in significantly worse results than any multi-step forecasting approach. This is likely due to the compounding effect of autoregressive errors.
D.5.6 Using auxiliary diffusion steps can improve performance
In Table 9 we ablate the number of artificial diffusion steps,  $k$ , used to train/evaluate DYffusion on the SST dataset. As in Table 7 we evaluate on the test box 88 only. It is clear, that  $k > 0$  is important for optimal CRPS and SSR performance. Here, the  $k$  auxiliary diffusion steps correspond to the floating-point interpolation timesteps  $\left\{\frac{j}{k + 1}\right\}_{j = 1}^{k}$ , as visualized for a toy example in Fig. 9b. Interestingly, we did not observe this phenomenon for the Navier-Stokes or spring mesh datasets, where using  $k > 0$  did not improve performance over using the simplest schedule  $S = [i_n]_{n = 0}^{N - 1} = [j]_{j = 0}^{h - 1}$ , where the number of diffusion steps and the horizon are equal ( $k = 0$ ).
D.5.7 Forecaster network iteratively refines its forecasts
In this set of experiments, we verify that the forecaster's prediction of  $\mathbf{x}_{t + h}$  for a given diffusion step  $n$  and associated input  $\hat{\mathbf{x}}_{t + i_n}$  usually improves as  $n$  increases, i.e. as we approach the end of
Table 9: For the SST dataset, using  $k$  extra artificial diffusion steps corresponding to floating-point  $i_k$ , especially such that  $i_k < 1$  for all  $k$  (similarly to Fig. 9b), improved performance significantly. The simplest possible schedule with  $k = 0$  corresponds to  $S = [i_n]_{n=0}^{N-1} = [j]_{j=0}^{h-1}$ , where the number of diffusion steps and the horizon are equal. The CRPS does not vary significantly when  $k > 0$ , while the SSR increases steadily with larger  $k$ . This can be explained by larger  $k$  creating a longer sampling sequence with  $k$  extra diffusion steps that may lead to more diverse samples. We evaluate only on a subset of the test dataset (box 88 in Fig. 5).
k	CRPS	MSE	SSR
k=0	0.208 ± 0.002	0.107 ± 0.001	0.49 ± 0.00
k=10	0.184 ± 0.002	0.113 ± 0.003	0.84 ± 0.02
k=25	0.181 ± 0.001	0.111 ± 0.001	0.99 ± 0.02
k=40	0.184 ± 0.002	0.113 ± 0.002	1.06 ± 0.03
k=45	0.189 ± 0.004	0.119 ± 0.005	1.09 ± 0.03
(a) OISSTv2
Figure 7: We verify that the forecasts  $\hat{\mathbf{x}}_{t + h} = F_{\theta}(\hat{\mathbf{x}}_{t + i_n},i_n)$  of the forecaster network (i.e. the diffusion model backbone) usually improve as the reverse process comes closer in time to  $t + h$  (i.e. as  $n$  and  $i_n$  increase) in terms of CRPS and spread-skill ratio (SSR). Interestingly, the MSE scores tend to degrade for the SST and spring mesh datasets (7a and 7c). For the SST dataset (7a), we also visualize a dotted vertical, orange line at diffusion step 36, which corresponds to the end of the 35 auxiliary diffusion steps (i.e. with floating-point  $i_n < 1$ ), which are followed by the diffusion steps corresponding to the data resolution (i.e.  $i_n\in \{1,2,\ldots ,6\}$ ). The SSR score improves significantly after that diffusion step. To achieve a clearer visualization of the spring mesh dataset, where the variance across random seeds is relatively large, we plot each run separately. Interestingly, for spring mesh, the CRPS starts degrading after around 30 to 70 diffusion steps, which may indicate that a shorter training horizon could be more amenable for performance.
(b) Navier-Stokes
(c) Spring Mesh
Figure 8: There is a clear trade-off between inference speed (x-axis) and performance (y-axis) as a function of the number of diffusion steps used for inference by DYffusion. Here, we show the SST test scores for one run of DYffusion, which was trained with 35 auxiliary diffusion steps (on top of the 7 given by the data). Each dot from left to right represents performing inference with an increasing number of diffusion steps. Each dot uses the base schedule  $S_{base} = \{0,1,\dots,h - 1\}$ , where  $h = 7$ , plus  $N_{aux}$  additional diffusion steps drawn from the ones used for training.  $N_{aux} = (1,2,3,4,5,6,7,8,9,10,12,14,16,18,20,23,26,29,35)$ , and each such additional diffusion step corresponds to an implicit dynamical timestep in (0,1). Interestingly, almost equivalent (for CRPS) or slightly better (for MSE) scores can be sometimes obtained by using fewer diffusion steps than used for training (right-most dots of each subplot), which immediately benefits inference speed.
the reverse process and  $\mathbf{x}_{t + i_n}$  gets closer in time to  $\mathbf{x}_{t + h}$  (see Fig. 7). This analysis is conducted on a subset of the SST test set (box 112 in Fig. 5) and validation sets of Navier-Stokes and spring mesh, for the respective datasets, using a 20-member ensemble. We use the validation sets of the physical systems benchmark because the scope of the analysis only extends up to the training horizon (as opposed to the full test set trajectories).
D.5.8 Accelerated DYffusion sampling
In Fig. 8 we study how sampling from DYffusion can be accelerated in a similar way to how DDIM [61] can accelerate sampling from Gaussian diffusion models. The continuous-time nature of the backbone networks in DYffusion invites the use of arbitrary dynamical timesteps as diffusion states during inference. Thus, to accelerate sampling, we can skip some of the  $N$  diffusion steps used for training,  $S_{train} = \{i_n\}_{n=0}^{N-1}$ , which automatically results in fewer neural network forward passes. For example, we can only use the base schedule  $S_{base} = \{0,1,\dots,h-1\}$ , where the diffusion states correspond to the temporal resolution of the dynamical data in a one-to-one mapping (see black arrows in Fig. 3). In Fig. 8, we start with  $S_{base}$  as inference schedule (left-most dots in each subplot) for the SST dataset, and then incrementally add more diffusion steps from  $S_{train} \setminus S_{base}$  to it, until reaching the full training schedule (right-most dots). We find that sampling can be significantly accelerated with marginal drops in CRPS and MSE performance. Note that the dynamical timesteps needed as outputs of DYffusion or for downstream applications pose a lower bound (here,  $S_{base}$ ) on how much we can accelerate our method since any such output timestep needs to be included in the sampling schedule or in the set of output timesteps,  $J$  (line 6 in Alg. 2). For that reason, it is not straightforward to accelerate the Navier-Stokes and spring mesh DYffusion runs, since their training schedules already correspond to  $S_{base}$ .
D.6 Modeling Complexity of DYffusion and Baselines
In Table 10 we enumerate the different modeling and computing requirements needed for each of the baselines and our method. Dropout (multi-step) refers to the barebone backbone network that forecasts all  $h$  timesteps  $\mathbf{x}_{t + 1:t + h}$  in a single forward pass. Dropout (continuous) refers to the barebone backbone network that forecasts one timestep  $\mathbf{x}_{t + k}$  for  $k \in (1,h)$  in a single forward pass, conditioned on the time,  $k$ . Both methods perform similarly in our exploratory experiments (not shown), and in our experiments, we always report the scores of the time-conditioned (i.e. continuous) variant. The multi-step approach corresponds to the way the backbone model of a video diffusion model operates, while the continuous variant is similar to how the forecaster network in DYffusion
Table 10: We report the requirements needed to train a method to forecast up to  $h$  steps into the future, where  $c$  refers to the number of input/output channels (e.g. 1 for SST data), and  $w$  to the window size (here,  $w = 1$  for all experiments). In the second and third columns, we report the input and output channel dimensions, respectively, that the (backbone) neural network needs to have (assuming that the window dimension is concatenated to the channel dimension).  $|\mathrm{Mem}(\mathbf{x}_t)|$  refers to the number of timesteps that need to be present in (GPU) memory to compute the training objective. The last column denotes how many network forward passes are needed to get the forecasts for all  $h$  timesteps. Here,  $N_1$ ,  $N_2$  refers to the number of (sampling) diffusion steps used by DDPM/MCVD and DYffusion, respectively. Usually,  $N_1 > N_2 \geq h$ , since Gaussian noise diffusion models will require more diffusion steps to attain comparable predictive skill. For Navier-Stokes and spring mesh  $N_2 = h$ , for SST  $N_2 = 35$ . For MCVD,  $N_1 = 1000$ . The factor of 3 for DYffusion is a result of the two extra interpolator network forward passes needed in line 4 of Alg. 2. For large horizons, the model size and memory requirements of multi-step models and conventional diffusion models can be prohibitive. It is clear that (video) diffusion models do not scale well for long horizons.
Modeling complexity
Method	\(c_{in}\)	\(c_{out}\)	\(|Mem(x_t)|\)	#Forward
Dropout (continuous)	\(w*c\)	\(c\)	\(w+1\)	\(h\)
Dropout (multi-step)	\(w*c\)	\(h*c\)	\(w+h\)	1
DDPM/MCVD	\((h+w)*c\)	\(h*c\)	\(w+h\)	\(N_1\)
DYffusion	\(w*c\)	\(c\)	\(w+1\)	\(3*N_2\)
operates. Video diffusion models have higher modeling complexity because they need to model the full "videos",  $\mathbf{x}_{t + 1:t + h}$ , at each diffusion state (or corrupted versions of it). Especially for long horizons,  $h$ , and high-dimensional data with several channels, this complicates the learning task for the neural network. Meanwhile, our method is only slightly impacted by the choice of  $h$ .
D.7 Example Sampling Trajectories Visualized
Example sampling trajectories of our approach, as a function of the schedule  $i_n$ , are visualized in Fig. 9, where the top one corresponds to the simplest case where we use a one-to-one mapping between diffusion steps,  $n$ , and interpolation/dynamical timesteps,  $i_n$ . Each of the intermediate  $\hat{\mathbf{x}}_i$  can be used as a forecast for timestep  $i$ . The forecaster network,  $F_{\theta}$ , repeatedly forecasts  $\mathbf{x}_h$ , but does so with increasing levels of skill (analogously to how conventional diffusion models iteratively denoise/refine the predictions of the "clean" data,  $\mathbf{s}^{(0)}$ ; see Fig. 7).