SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \mathcal {L} = \int q \left(\theta_ {\mathrm {d y n}}, \theta_ {\mathrm {d e c}}, s _ {1: B}\right) \ln \frac {p \left(\boldsymbol {u} _ {1 : M} , \boldsymbol {s} _ {1 : B} , \theta_ {\mathrm {d y n}} , \theta_ {\mathrm {d e c}}\right)}{q \left(\theta_ {\mathrm {d y n}} , \theta_ {\mathrm {d e c}} , \boldsymbol {s} _ {1 : B}\right)} d \theta_ {\mathrm {d y n}} d \theta_ {\mathrm {d e c}} d s _ {1: B} (37) \\ = \int q \left(\theta_ {\mathrm {d y n}}, \theta_ {\mathrm {d e c}}, s _ {1: B}\right) \ln \frac {p \left(\boldsymbol {u} _ {1 : M} \mid \boldsymbol {s} _ {1 : B} , \theta_ {\mathrm {d y n}} , \theta_ {\mathrm {d e c}}\right) p \left(\boldsymbol {s} _ {1 : B} \mid \theta_ {\mathrm {d y n}}\right) p \left(\theta_ {\mathrm {d y n}}\right) p \left(\theta_ {\mathrm {d e c}}\right)}{q \left(\boldsymbol {s} _ {1 : B}\right) q \left(\theta_ {\mathrm {d y n}}\right) q \left(\theta_ {\mathrm {d e c}}\right)} d \theta_ {\mathrm {d y n}} d \theta_ {\mathrm {d e c}} d s _ {1: B} (38) \\ = \int q \left(\theta_ {\mathrm {d y n}}, \theta_ {\mathrm {d e c}}, s _ {1: B}\right) \ln p \left(\boldsymbol {u} _ {1: M} \mid \boldsymbol {s} _ {1: B}, \theta_ {\mathrm {d y n}}, \theta_ {\mathrm {d e c}}\right) d \theta_ {\mathrm {d y n}} d \theta_ {\mathrm {d e c}} d s _ {1: B} (39) \\ - \int q \left(\theta_ {\mathrm {d y n}}, \theta_ {\mathrm {d e c}}, s _ {1: B}\right) \ln \frac {q \left(s _ {1 : B}\right)}{p \left(s _ {1 : B} \mid \theta_ {\mathrm {d y n}}\right)} d \theta_ {\mathrm {d y n}} d \theta_ {\mathrm {d e c}} d s _ {1: B} (40) \\ - \int q \left(\theta_ {\mathrm {d y n}}, \theta_ {\mathrm {d e c}}, s _ {1: B}\right) \ln \frac {q \left(\theta_ {\mathrm {d y n}}\right)}{p \left(\theta_ {\mathrm {d y n}}\right)} d \theta_ {\mathrm {d y n}} d \theta_ {\mathrm {d e c}} d s _ {1: B} (41) \\ - \int q \left(\theta_ {\mathrm {d e c}}, \theta_ {\mathrm {d e c}}, s _ {1: B}\right) \ln \frac {q \left(\theta_ {\mathrm {d e c}}\right)}{p \left(\theta_ {\mathrm {d e c}}\right)} d \theta_ {\mathrm {d y n}} d \theta_ {\mathrm {d e c}} d s _ {1: B} (42) \\ = \mathcal {L} _ {1} - \mathcal {L} _ {2} - \mathcal {L} _ {3} - \mathcal {L} _ {4}. (43) \\ \end{array}
$$
Next, we will look at each term  $\mathcal{L}_i$  separately.
$$
\begin{array}{l} \mathcal {L} _ {1} = \int q \left(\theta_ {\mathrm {d y n}}, \theta_ {\mathrm {d e c}}, \boldsymbol {s} _ {1: B}\right) \ln p \left(\boldsymbol {u} _ {1: M} \mid \boldsymbol {s} _ {1: B}, \theta_ {\mathrm {d y n}}, \theta_ {\mathrm {d e c}}\right) d \theta_ {\mathrm {d y n}} d \theta_ {\mathrm {d e c}} d \boldsymbol {s} _ {1: B} (44) \\ = \int q \left(\theta_ {\mathrm {d y n}}, \theta_ {\mathrm {d e c}}, \boldsymbol {s} _ {1: B}\right) \ln \left[ \prod_ {b = 1} ^ {B} \prod_ {i \in \mathcal {I} _ {b}} \prod_ {j = 1} ^ {N} p \left(\boldsymbol {u} _ {i} ^ {j} \mid \boldsymbol {s} _ {b}, \theta_ {\mathrm {d y n}}, \theta_ {\mathrm {d e c}}\right) \right] d \theta_ {\mathrm {d y n}} d \theta_ {\mathrm {d e c}} d \boldsymbol {s} _ {1: B} (45) \\ = \sum_ {b = 1} ^ {B} \sum_ {i \in \mathcal {I} _ {b}} \sum_ {j = 1} ^ {N} \int q \left(\theta_ {\mathrm {d y n}}, \theta_ {\mathrm {d e c}}, s _ {1: B}\right) \ln \left[ p \left(\boldsymbol {u} _ {i} ^ {j} \mid \boldsymbol {s} _ {b}, \theta_ {\mathrm {d y n}}, \theta_ {\mathrm {d e c}}\right) \right] d \theta_ {\mathrm {d y n}} d \theta_ {\mathrm {d e c}} d s _ {1: B} (46) \\ = \sum_ {b = 1} ^ {B} \sum_ {i \in \mathcal {I} _ {b}} \sum_ {j = 1} ^ {N} \int q \left(\theta_ {\mathrm {d y n}}, \theta_ {\mathrm {d e c}}, \boldsymbol {s} _ {b}\right) \ln \left[ p \left(\boldsymbol {u} _ {i} ^ {j} \mid \boldsymbol {s} _ {b}, \theta_ {\mathrm {d y n}}, \theta_ {\mathrm {d e c}}\right) \right] d \theta_ {\mathrm {d y n}} d \theta_ {\mathrm {d e c}} d \boldsymbol {s} _ {b} (47) \\ = \sum_ {b = 1} ^ {B} \sum_ {i \in \mathcal {I} _ {b}} \sum_ {j = 1} ^ {N} \mathbb {E} _ {q \left(\theta_ {\mathrm {d y n}}, \theta_ {\mathrm {d e c}}, \boldsymbol {s} _ {b}\right)} \ln \left[ p \left(\boldsymbol {u} _ {i} ^ {j} \mid \boldsymbol {s} _ {b}, \theta_ {\mathrm {d y n}}, \theta_ {\mathrm {d e c}}\right) \right]. (48) \\ \end{array}
$$
where KL is Kullback-Leibler (KL) divergence. Both of the KL divergences above have a closed form but the expectation w.r.t.  $q(\theta_{\mathrm{dyn}}, s_{b-1})$  does not.
$$
\mathcal {L} _ {3} = \mathrm {K L} \left(q \left(\theta_ {\mathrm {d y n}}\right) \| p \left(\theta_ {\mathrm {d y n}}\right)\right), \quad \mathcal {L} _ {4} = \mathrm {K L} \left(q \left(\theta_ {\mathrm {d e c}}\right) \| p \left(\theta_ {\mathrm {d e c}}\right)\right). \tag {60}
$$
B.2 Computation of ELBO.
We compute the ELBO using the following algorithm:
1. Sample  $\theta_{\mathrm{dyn}}$ ,  $\theta_{\mathrm{dec}}$  from  $q_{\psi_{\mathrm{dyn}}}(\theta_{\mathrm{dyn}})$ ,  $q_{\psi_{\mathrm{dec}}}(\theta_{\mathrm{dec}})$ .
2. Sample  $s_{1:B}$  by sampling each  $s_b^j$  from  $q_{\pmb{\psi}_b^j}(\pmb{s}_b^j)$  with  $\pmb{\psi}_b^j = h_{\theta_{\mathrm{enc}}}(\pmb{u}[t_{[b]},\pmb{x}_j])$ .
3. Compute  $\pmb{u}_{1:M}$  from  $s_{1:B}$  as in Equations 14-16.
4. Compute ELBO  $\mathcal{L}$  (KL terms are computed in closed form, for expectations we use Monte Carlo integration with one sample).
Sampling is done using reparametrization to allow unbiased gradients w.r.t. the model parameters.
C Appendix C
SHALLOW WATER. The shallow water equations are a system of partial differential equations (PDEs) that simulate the behavior of water in a shallow basin. These equations are effectively a depth-integrated version of the Navier-Stokes equations, assuming the horizontal length scale is significantly larger than the vertical length scale. Given these assumptions, they provide a model for water dynamics in a basin or similar environment, and are commonly utilized in predicting the propagation of water waves, tides, tsunamis, and coastal currents. The state of the system modeled by these equations consists of the wave height  $h(t,x,y)$ , velocity in the  $x$ -direction  $u(t,x,y)$  and velocity in the  $y$ -direction  $v(t,x,y)$ . Given an initial state  $(h_0,u_0,v_0)$ , we solve the PDEs on a spatial domain  $\Omega$  over time interval  $[0,T]$ . The shallow water equations are defined as:
$$
\frac {\partial h}{\partial t} + \frac {\partial (h u)}{\partial x} + \frac {\partial (h v)}{\partial y} = 0, \tag {61}
$$
$$
\frac {\partial u}{\partial t} + u \frac {\partial u}{\partial x} + v \frac {\partial u}{\partial y} + g \frac {\partial h}{\partial x} = 0, \tag {62}
$$
$$
\frac {\partial v}{\partial t} + u \frac {\partial v}{\partial x} + v \frac {\partial v}{\partial y} + g \frac {\partial h}{\partial y} = 0, \tag {63}
$$
where  $g$  is the gravitational constant.
The spatial domain  $\Omega$  is a unit square with periodic boundary conditions. We set  $T = 0.1$  sec. The solution is evaluated at randomly selected spatial locations and time points. We use 1089 spatial locations and 25 time points. The spatial end temporal grids are the same for all trajectories. Since we are dealing with partially-observed cases, we assume that we observe only the wave height  $h(t,x,y)$ .
For each trajectory, we start with zero initial velocities and the initial height  $h_0(x,y)$  generated as:
$$
\tilde {h} _ {0} (x, y) = \sum_ {k, l = - N} ^ {N} \lambda_ {k l} \cos (2 \pi (k x + l y)) + \gamma_ {k l} \sin (2 \pi (k x + l y)), \tag {64}
$$
$$
h _ {0} (x, y) = 1 + \frac {\tilde {h} _ {0} (x , y) - \operatorname* {m i n} \left(\tilde {h} _ {0}\right)}{\operatorname* {m a x} \left(\tilde {h} _ {0}\right) - \operatorname* {m i n} \left(\tilde {h} _ {0}\right)}, \tag {65}
$$
where  $N = 3$  and  $\lambda_{kl},\gamma_{kl}\sim \mathcal{N}(0,1)$
The datasets used for training, validation, and testing contain 60, 20, and 20 trajectories, respectively.
We use scikit-fdiff (Cellier, 2019) to solve the PDEs.
NAVIER-STOKES. For this dataset we model the propagation of a scalar field (e.g., smoke concentration) in a fluid (e.g., air). The modeling is done by coupling the Navier-Stokes equations with the Boussinesq buoyancy term and the transport equation to model the propagation of the scalar field. The state of the system modeled by these equations consists of the scalar field  $c(t,x,y)$ , velocity in  $x$ -direction  $u(t,x,y)$ , velocity in  $y$ -direction  $v(t,x,y)$ , and pressure  $p(t,x,y)$ . Given an initial state
$(c_0, u_0, v_0, p_0)$ , we solve the PDEs on a spatial domain  $\Omega$  over time interval  $[0, T]$ . The Navier-Stokes equations with the transport equation are defined as:
$$
\frac {\partial u}{\partial x} + \frac {\partial v}{\partial y} = 0, \tag {66}
$$
$$
\frac {\partial u}{\partial t} + u \frac {\partial u}{\partial x} + v \frac {\partial u}{\partial y} = - \frac {\partial p}{\partial x} + \nu \left(\frac {\partial^ {2} u}{\partial x ^ {2}} + \frac {\partial^ {2} u}{\partial y ^ {2}}\right), \tag {67}
$$
$$
\frac {\partial v}{\partial t} + u \frac {\partial v}{\partial x} + v \frac {\partial v}{\partial y} = - \frac {\partial p}{\partial y} + \nu \left(\frac {\partial^ {2} v}{\partial x ^ {2}} + \frac {\partial^ {2} v}{\partial y ^ {2}}\right) + c, \tag {68}
$$
$$
\frac {\partial c}{\partial t} = - u \frac {\partial c}{\partial x} - v \frac {\partial c}{\partial y} + \nu \left(\frac {\partial^ {2} c}{\partial x ^ {2}} + \frac {\partial^ {2} c}{\partial y ^ {2}}\right), \tag {69}
$$
where  $\nu = 0.002$
The spatial domain  $\Omega$  is a unit square with periodic boundary conditions. We set  $T = 2.0$  sec, but drop the first 0.5 seconds due to slow dynamics during this time period. The solution is evaluated at randomly selected spatial locations and time points. We use 1089 spatial locations and 25 time points. The spatial and temporal grids are the same for all trajectories. Since we are dealing with partially-observed cases, we assume that we observe only the scalar field  $c(t,x,y)$ .
For each trajectory, we start with zero initial velocities and pressure, and the initial scalar field  $c_{0}(x,y)$  is generated as:
$$
\tilde {c} _ {0} (x, y) = \sum_ {k, l = - N} ^ {N} \lambda_ {k l} \cos (2 \pi (k x + l y)) + \gamma_ {k l} \sin (2 \pi (k x + l y)), \tag {70}
$$
$$
c _ {0} (x, y) = \frac {\tilde {c} _ {0} (x , y) - \min  (\tilde {c} _ {0})}{\max  (\tilde {c} _ {0}) - \min  (\tilde {c} _ {0})}, \tag {71}
$$
where  $N = 2$  and  $\lambda_{kl},\gamma_{kl}\sim \mathcal{N}(0,1)$
We use PhiFlow (Holl et al., 2020) to solve the PDEs.
SCALAR FLOW. This dataset, proposed by Eckert et al. (2019), consists of observations of smoke plumes rising in hot air. The observations are postprocessed camera images of the smoke plumes taken from multiple views. For simplicity, we use only the front view. The dataset contains 104 trajectories, where each trajectory has 150 time points and each image has the resolution  $1080 \times 1920$ . Each trajectory was recorded for  $T = 2.5$  seconds.
To reduce dimensionality of the observations we sub-sample the original spatial and temporal grids. For the temporal grid, we remove the first 50 time points, which leaves 100 time points, and then take every 4th time point, thus leaving 20 time points in total. The original  $1080 \times 1920$  spatial grid is first downsampled by a factor of 9 giving a new grid with resolution  $120 \times 213$ , and then the new grid is further sub-sampled based on the smoke density at each node. In particular, we compute the average smoke density at each node (averaged over time), and then sample the nodes without replacement with the probability proportional to the average smoke density (thus, nodes that have zero density most of the time are not selected). See example of a final grid in Figure 9. This gives a new grid with 1089 nodes.
We further smooth the observations by applying Gaussian smoothing with the standard deviation of 1.5 (assuming domain size  $120 \times 213$ ).
Figure 9: Spatial grid used for SCALAR FLOW dataset.
We use the first 60 trajectories for training, next 20 for validation and next 20 for testing.
In this case the spatial domain is non-periodic, which means that for some observation location  $\pmb{x}_i$  some of its spatial neighbors  $\mathcal{N}_{\mathrm{S}}(\pmb{x}_i)$  might be outside of the domain. We found that to account for such cases it is sufficient to mark such out-of-domain neighbors by setting their value to -1.
Time grids used for the three datasets are shown in Figure 10.
Figure 10: Time grids used for SHALLOW WATER (top), NAVIER-STOKES (middle), and SCALAR FLOW (bottom).
C.2 Model architecture and hyper-parameters.
Dynamics function. For all datasets we define  $F_{\theta_{\mathrm{dyn}}}$  as an MLP. For SHALLOW WATER/NAVIER-STOKES/SCALAR FLOW we use 1/3/3 hidden layers with the size of 1024/512/512, respectively. We use ReLU nonlinearities.
Observation function. For all datasets we define  $g_{\theta_{\mathrm{dec}}}$  as a selector function which takes the latent state  $\pmb {z}(t,x)\in \mathbb{R}^d$  and returns its first component.
Encoder. Our encoder  $h_{\theta_{\mathrm{enc}}}$  consists of three functions:  $h_{\theta_{\mathrm{spatial}}}$ ,  $h_{\theta_{\mathrm{temporal}}}$ , and  $h_{\theta_{\mathrm{read}}}$ . The spatial aggregation function  $h_{\theta_{\mathrm{spatial}}}$  is a linear mapping to  $\mathbb{R}^{128}$ . The temporal aggregation function  $h_{\theta_{\mathrm{temporal}}}$  is a stack of transformer layers with temporal attention and continuous relative positional encodings (Iakovlev et al., 2023). For all datasets, we set the number of transformer layers to 6. Finally, the variational parameter readout function  $h_{\theta_{\mathrm{read}}}$  is a mapping defined as
$$
\psi_ {b} ^ {j} = h _ {\theta_ {\text {r e a d}}} \left(\alpha_ {[ b ]} ^ {\mathrm {T}}\right) = \binom {\gamma_ {b} ^ {j}} {\tau_ {b} ^ {j}} = \binom {\operatorname {L i n e a r} \left(\alpha_ {[ b ]} ^ {\mathrm {T}}\right)} {\exp \left(\operatorname {L i n e a r} \left(\alpha_ {[ b ]} ^ {\mathrm {T}}\right)\right)}, \tag {72}
$$
where Linear is a linear layer (different for each line), and  $\gamma_{b}^{j}$  and  $\tau_{b}^{j}$  are the variational parameters discussed in Appendix A.
Spatial and temporal neighborhoods. We use the same spatial neighborhoods  $\mathcal{N}_{\mathrm{S}}(\pmb{x})$  for both the encoder and the dynamics function. We define  $\mathcal{N}_{\mathrm{S}}(\pmb{x})$  as the set of points consisting of the point  $\pmb{x}$  and points on two concentric circles centered at  $\pmb{x}$ , with radii  $r$  and  $r/2$ , respectively. Each circle contains 8 points spaced 45 degrees apart (see Figure 11 (right)). The radius  $r$  is set to 0.1. For SHALLOW WATER/NAVIER-STOKES/SCALAR FLOW the size of temporal neighborhood  $(\delta_T)$  is set to 0.1/0.1/0.2, respectively.
Multiple Shooting. For SHALLOW WATER/NAVIER-STOKES/SCALAR FLOW we split the full training trajectories into 4/4/19 sub-trajectories, or, equivalently, have the sub-trajectory length of 6/6/2.
C.3 Training, validation, and testing setup.
Data preprocessing. We scale the temporal grids, spatial grids, and observations to be within the interval [0, 1].
Training. We train our model for 20000 iterations using Adam (Kingma and Ba, 2017) optimizer with constant learning rate 3e-4 and linear warmup for 200 iterations. The latent spatiotemporal dynamics are simulated using differentiable ODE solvers from the torchdiffeq package (Chen, 2018) (we use dopri5 with rtol=1e-3, atol=1e-4, no adjoint). The batch size is 1.
Validation. We use validation set to track the performance of our model during training and save the parameters that produce the best validation performance. As performance measure we use the mean absolute error at predicting the full validation trajectories given some number of initial observations. For SHALLOW WATER/NAVIER-STOKES/SCALAR FLOW we use the first 5/5/10 observations. The predictions are made by taking one sample from the posterior predictive distribution (see Appendix C.4 for details).
Testing. Testing is done similarly to validation, except that as the prediction we use an estimate of the expected value of the posterior predictive distribution (see Appendix C.4 for details).
C.4 Forecasting.
$$
p \left(\tilde {\boldsymbol {u}} _ {n} \mid \tilde {\boldsymbol {u}} _ {1: m}, \boldsymbol {u} _ {1: M}\right) \approx \int p \left(\tilde {\boldsymbol {u}} _ {n} \mid \tilde {\boldsymbol {s}} _ {m}, \theta_ {\mathrm {d y n}}, \theta_ {\mathrm {d e c}}\right) q \left(\tilde {\boldsymbol {s}} _ {m}\right) q \left(\theta_ {\mathrm {d y n}}\right) q \left(\theta_ {\mathrm {d e c}}\right) d \tilde {\boldsymbol {s}} _ {m} d \theta_ {\mathrm {d y n}} d \theta_ {\mathrm {d e c}}. \tag {73}
$$
The expected value is estimated via Monte Carlo integration, so the algorithm for predicting  $\tilde{u}_n$  is:
1. Sample  $\theta_{\mathrm{dyn}}$ ,  $\theta_{\mathrm{dec}}$  from  $q(\theta_{\mathrm{dyn}})$ ,  $q(\theta_{\mathrm{dec}})$ .
2. Sample  $\tilde{s}_m$  from  $q(\tilde{s}_m) = \prod_{j=1}^N q_{\psi_m^j}(\tilde{s}_m^j)$ , where the variational parameters  $\psi_m^j$  are given by the encoder  $h_{\theta_{\mathrm{enc}}}$  operating on the initial observations  $\tilde{\pmb{u}}_{1:m}$  as  $\psi_m^j = h_{\theta_{\mathrm{enc}}}(\tilde{\pmb{u}}[t_m, \pmb{x}_j])$ .
3. Compute the latent state  $\tilde{z}(t_n) = z(t_n; t_m, \tilde{s}_m, \theta_{\mathrm{dyn}})$ .
4. Sample  $\tilde{u}_n$  by sampling each  $\tilde{u}_n^j$  from  $\mathcal{N}(\tilde{u}_n^j |g_{\theta_{\mathrm{dec}}}(\tilde{z}(t_n,\boldsymbol{x}_j)),\sigma_u^2 I)$ .
5. Repeat steps  $1 - 4n$  times and average the predictions (we use  $n = 10$ ).
C.5 Model comparison setup.
NODE. For the NODE model the dynamics function was implemented as a fully connected feedforward neural network with 3 hidden layers, 512 neurons each, and ReLU nonlinearities.
FEN. We use the official implementation of FEN. We use FEN variant without the transport term as we found it improves results on our datasets. The dynamics were assumed to be stationary and autonomous in all cases. The dynamics function was represented by a fully connected feedforward neural network with 3 hidden layers, 512 neurons each, and ReLU nonlinearities.
NSPDE. We use the official implementation of NSPDE. We set the number of hidden channels to 16, and set modes1 and modes2 to 32.
DINO. We use the official implementation of DINO. The encoder is an MLP with 3 hidden layers, 512 neurons each, and Swish non-linearities. The code dimension is 100. The dynamics function is an MLP with 3 hidden layers, 512 neurons each, and Swish non-linearities. The decoder has 3 layers and 64 channels.
MAgNet. We use the official implementation of MAgNet. We use the graph neural network variant of the model. The number of message-passing steps is 5. All MLPs have 4 layers with 128 neurons each in each layer. The latent state dimension is 128.
D Appendix D
D.1 Spatiotemporal neighborhood shapes and sizes.
Here we investigate the effect of changing the shape and size of spatial and temporal neighborhoods used by the encoder and dynamics functions. We use the default hyperparameters discussed in Appendix C and change only the neighborhood shape or size. A neighborhood size of zero implies no spatial/temporal aggregation.
Initially, we use the original circular neighborhood displayed in Figure 11 for both encoder and dynamics function and change only its size (radius). The results are presented in Figures 12a and 12b. In Figure 12a, it is surprising to see very little effect from changing the encoder's spatial neighborhood size. A potential explanation is that the dynamics function shares the spatial aggregation task with the encoder. However, the results in Figure 12b are more intuitive, displaying a U-shaped curve for the test MAE, indicating the importance of using spatial neighborhoods of appropriate size. Interestingly, the best results tend to be achieved with relatively large neighborhood sizes. Similarly, Figure 12c shows U-shaped curves for the encoder's temporal neighborhood size, suggesting that latent state inference benefits from utilizing local temporal information.
We then examine the effect of changing the shape of the dynamics function's spatial neighborhood. We use  $n$  circle neighborhoods, which consist of  $n$  equidistant concentric circular neighborhoods (see examples in Figure 11). Effectively, we maintain a fixed neighborhood size while altering its density. The results can be seen in Figure 13. We find that performance does not significantly improve when using denser (and presumably more informative) spatial neighborhoods, indicating that accurate predictions only require a relatively sparse neighborhood with appropriate size.
Figure 11: Left: original circular neighborhood (1circle). Center: circular neighborhood with increased size. Right: circular neighborhood of a different shape (2circle).
(c)
Figure 12: (a), (b): Test MAE vs spatial neighborhood sizes of the encoder and dynamics function, respectively. (c): Test MAE vs temporal neighborhood size of the encoder. Note that the spatial and temporal domains are normalized, so their largest size in any dimension is 1.
Figure 13: Test MAE vs spatial neighborhood shape.
D.2 Multiple shooting.
Here we demonstrate the effect of using multiple shooting for model training. In Figure 14 (left), we vary the sub-trajectory length (longer sub-trajectories imply more difficult training) and plot the test errors for each sub-trajectory length. We observe that in all cases, the best results are achieved when the sub-trajectory length is considerably smaller than the full trajectory length. In Figure 14 (right) we further show the training times, and as can be seen multiple shooting allows to noticeably reduce the training times.
Figure 14: Test MAE vs training sub-trajectory length.
E Appendix E
Noisy Data. Here we show the effect of observation noise on our model and compare the results against other models. We train all models with data noise of various strengths, and then compute test MAE on noiseless data (we still use noisy data to infer the initial state at test time). Figure 15 shows that our model can manage noise strength up to 0.1 without significant drops in performance. Note that all observations are in the range [0, 1].
Figure 15: Test MAE vs observation noise  $\sigma_{u}$
F Appendix F
F.1 Model Predictions
We show (Fig. 16) predictions of different models trained on different datasets (synthetic data is partially observed).
F.2 Visualization of Prediction Uncertainty
Figures 17, 18, and 19 demonstrate the prediction uncertainty across different samples from the posterior distribution.
Figure 16: Predictions from different models. Only forecasting is shown.
Figure 17: Left: Test prediction for a single trajectory on SHALLOW WATER dataset. Show are data (top), mean prediction (middle), and standard deviation of the predictions (bottom). Columns show predictions at consecutive time points. Right: Ground truth observations (dashed black) and predictions (colored) with standard deviation plotted over time at two spatial locations.
Figure 18: Left: Test prediction for a single trajectory on NAVIER-STOKES dataset. Show are data (top), mean prediction (middle), and standard deviation of the predictions (bottom). Columns show predictions at consecutive time points. Right: Ground truth observations (dashed black) and predictions (colored) with standard deviation plotted over time at two spatial locations.
Figure 19: Left: Test prediction for a single trajectory on SCALAR FLOW dataset. Show are data (top), mean prediction (middle), and standard deviation of the predictions (bottom). Columns show predictions at consecutive time points. Right: Ground truth observations (dashed black) and predictions (colored) with standard deviation plotted over time at two spatial locations.