SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\boldsymbol {u} _ {i} (\boldsymbol {x}) := \operatorname {I n t e r p o l a t e} \left(\boldsymbol {u} _ {i}\right) (\boldsymbol {x}), \quad i = 1, \dots , M, \tag {19}
$$
where spatial interpolation is done separately for each time point  $i$ . We then use the interpolants  $\pmb{u}_i(\pmb{x})$  to define the spatial neighborhoods  $\mathcal{N}_{\mathrm{S}}(\pmb{x})$  in a grid-independent manner.
Figure 5: Spatiotemporal neighborhood of a multiple shooting time point  $t_{[b]} = t_i$  and location  $\boldsymbol{x}_j$ ,  $\boldsymbol{u}[t_{[b]},\boldsymbol{x}_j]$  (denoted by green, blue and orange crosses and the dots inside), is mapped to the variational parameters  $\psi_b^j$  via the encoder.
To improve data-efficiency, we assume  $\psi_b^j$  does not depend on the whole observed sequence  $\pmb{u}_{1:M}$  but only on some local information in a spatiotemporal neighborhood of  $t_{[b]}$  and  $\pmb{x}_j$ . We define the temporal neighborhood of  $t_{[b]}$  as
$$
\mathcal {N} _ {\mathrm {T}} \left(t _ {[ b ]}\right) := \left\{k: \left| t _ {k} - t _ {[ b ]} \right| \leq \delta_ {T}, k = 1, \dots , M \right\}, \tag {20}
$$
where  $\delta_T$  is a hyperparameter controlling the neighborhood size, and then define the spatiotemporal neighborhood of  $t_{[b]}$  and  $\pmb{x}_j$  as
$$
\boldsymbol {u} \left[ t _ {[ b ]}, \boldsymbol {x} _ {j} \right] := \left\{\boldsymbol {u} _ {k} (\boldsymbol {x}): k \in \mathcal {N} _ {\mathrm {T}} \left(t _ {[ b ]}\right), \boldsymbol {x} \in \mathcal {N} _ {\mathrm {S}} \left(\boldsymbol {x} _ {j}\right) \right\}. \tag {21}
$$
Our encoder operates on such spatiotemporal neighborhoods  $\pmb{u}[t_{[b]},\pmb{x}_j]$  and works in three steps (see Figure 5). First, for each time index  $k\in \mathcal{N}_{\mathrm{T}}(t_{[b]})$  it aggregates the spatial information  $\{\pmb{u}_k(\pmb {x})\}_{\pmb {x}\in \mathcal{N}(\pmb{x}_j)}$  into a vector  $\alpha_{k}^{\mathrm{S}}$ . Then, it aggregates the spatial representations  $\alpha_{k}^{\mathrm{S}}$  across time into another vector  $\alpha_{[b]}^{\mathrm{T}}$  which is finally mapped to the variational parameters  $\psi_b^j$  as follows:
$$
\boldsymbol {\psi} _ {b} ^ {j} = h _ {\theta_ {\mathrm {e n c}}} \left(\boldsymbol {u} \left[ t _ {[ b ]}, \boldsymbol {x} _ {j} \right]\right) = h _ {\text {r e a d}} \left(h _ {\text {t e m p o r a l}} \left(h _ {\text {s p a t i a l}} \left(\boldsymbol {u} \left[ t _ {[ b ]}, \boldsymbol {x} _ {j} \right]\right)\right)\right). \tag {22}
$$
Spatial aggregation. Since the spatial neighborhoods are fixed and remain identical for all spatial locations (see Figure 5), we implement the spatial aggregation function  $h_{\mathrm{spatial}}$  as an MLP which takes elements of the set  $\{\pmb{u}_k(\pmb{x})\}_{\pmb{x} \in \mathcal{N}_{\mathrm{S}}(\pmb{x}_j)}$  stacked in a fixed order as the input.
Temporal aggregation. We implement  $h_{\mathrm{temporal}}$  as a stack of transformer layers (Vaswani et al., 2017) which allows it to operate on input sets of arbitrary size. We use time-aware attention and continuous relative positional encodings (Iakovlev et al., 2023) which were shown to be effective on data from dynamical systems observed at irregular time intervals. Each transformer layer takes a layer-specific input set  $\{\xi_k^{\mathrm{in}}\}_{k\in \mathcal{N}_{\mathrm{T}}(t_{[b]})}$ , where  $\xi_k^{\mathrm{in}}$  is located at  $t_k$ , and maps it to an output set  $\{\xi_k^{\mathrm{out}}\}_{k\in \mathcal{N}_{\mathrm{T}}(t_{[b]})}$ , where each  $\xi_k^{\mathrm{out}}$  is computed using only the input elements within distance  $\delta_T$  from  $t_k$ , thus promoting temporal locality. Furthermore, instead of using absolute positional encodings the model assumes the behavior of the system does not depend on time and uses relative temporal distances to inject positional information. The first layer takes  $\{\alpha_k^{\mathrm{S}}\}_{k\in \mathcal{N}_{\mathrm{T}}(t_{[b]})}$  as the input, while the last layer returns a single element at time point  $t_{[b]}$ , which represents the temporal aggregation  $\alpha_{[b]}^{\mathrm{T}}$ .
Variational parameter readout. Since  $\alpha_{i}^{\mathrm{T}}$  is a fixed-length vector, we implement  $h_{\mathrm{read}}$  as an MLP.
Figure 6: Top: SHALLOW WATER dataset contains observations of the wave height in a pool of water. Middle: NAVIER-STOKES dataset contains observations of the concentration of a species transported in a fluid via buoyancy and velocity field. Bottom: SCALAR FLOW dataset contains observations of smoke plumes raising in warm air.
4.3 Forecasting
Given initial observations  $\tilde{\pmb{u}}_{1:m}$  at time points  $t_{1:m}$ , we predict the future observation  $\tilde{\pmb{u}}_n$  at a time point  $t_n > t_m$  as the expected value of the approximate posterior predictive distribution:
$$
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 {23}
$$
The expected value is estimated via Monte Carlo integration (see Appendix C.4 for details).
We use three challenging datasets: SHALLOW WATER, NAVIER-STOKES, and SCALAR FLOW which contain observations of spatiotemporal system at  $N \approx 1100$  grid points evolving over time (see Figure 6). The first two datasets are synthetic and generated using numeric PDE solvers (we use scikit- diff (Cellier, 2019) for SHALLOW WATER, and PhiFlow (Holl et al., 2020) for NAVIER-STOKES), while the third dataset contains real-world observations (camera images) of smoke plumes raising in warm air (Eckert et al., 2019). In all cases the observations are made at irregular spatiotemporal grids and contain only partial information about the true system state. In particular, for SHALLOW WATER we observe only the wave height, for NAVIER-STOKES we observe only the concentration of the species, and for SCALAR FLOW only pixel densities are known. All datasets contain 60/20/20 training/validation/testing trajectories. See Appendix C for details.
We train our model for 20k iterations with constant learning rate of 3e-4 and linear warmup. 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). Training is done on a single NVIDIA Tesla V100 GPU, with a single run taking 3-4 hours. We use the mean absolute error (MAE) on the test set as the performance measure. Error bars are standard errors over 4 random seeds. For forecasting we use the expected value of the posterior predictive distribution. See Appendix C for all details about the training, validation, and testing setup.
Latent state dimension. Here we show the advantage of using latent-space models on partially observed data. We change the latent state dimension  $d$  from 1 to 5 and measure the test MAE. Note that for  $d = 1$  we effectively have a data-space model which models the observations without trying to reconstruct the missing states. Figure 7 shows that in all cases there is improvement in performance as the latent dimension grows. For SHALLOW WATER and NAVIER-STOKES the true latent dimension is 3. Since SCALAR FLOW is a real-world process, there is no true latent dimension. As a benchmark, we provide the performance of our model trained on fully-observed versions of the synthetic datasets (we use the same architecture and hyperparameters, but fix  $d$  to 3). Figure 7 also shows examples of model predictions (at the final time point) for different values of  $d$ . We see a huge difference between  $d = 1$  and  $d = 3, 5$ . Note how apparently small difference in MAE at  $d = 1$  and  $d = 5$  for SCALAR FLOW corresponds to a dramatic improvement in the prediction quality.
Figure 7: Left: Test MAE vs latent state dimension  $d$ . Black lines are test MAE on fully-observed versions of the datasets ( $\pm$  standard error). Right: Model predictions for different  $d$ .
Grid independence. In this experiment we demonstrate the grid-independence property of our model by training it on grids with  $\approx 1100$  observation locations, and then testing on a coarser, original, and finer grids. We evaluate the effect of using different interpolation methods by repeating the experiment with linear, k-nearest neighbors, and inverse distance weighting (IDW) interpolants. For SHALLOW WATER and NAVIER-STOKES the coarser/finer grids contain 290/4200 nodes, while for SCALAR FLOW we have 560/6420 nodes, respectively. Table 1 shows the model's performance for different spatial grids and interpolation methods. We see that all interpolation methods perform rather similarly on the original grid, but linear interpolation and IDW tend to perform better on finer/coarser grids than k-NN. Performance drop on coarse grids is expected since we get less accurate information about the system's initial state and simulate the dynamics on coarse grids. Figure 8 also shows
examples of model predictions (at the final time point) for different grid sizes and linear interpolant.
Figure 8: Predictions on spatial grids of different density (linear interpolant, test data).
Comparison to other models. We test our model against recent spatiotemporal models from the literature: Finite Element Networks (FEN) (Lienen and Gunnemann, 2022), Neural Stochastic PDEs (NSPDE) (Salvi et al., 2021), MAgNet (Boussif et al., 2022), and DINO (Yin et al., 2023). We also use Neural ODEs (NODE) (Chen et al., 2018) as the baseline. We use the official implementation for all models and tune their hyperparameters for the best performance (see App.
C for details). For SHALLOW WATER and NAVIER-STOKES we use the first 5 time points to infer the latent state and then predict the next 20 time points, while for SCALAR FLOW we use the first 10 points for inference and predict the next 10 points. For synthetic data, we consider two settings: one where the data is fully observed (i.e., the complete state is recorded) – a setting for which most models are designed – and one where the data is partially observed (i.e., only part of the full state
Table 1: Test MAE for different interpolation methods.
Dataset	Grid	k-NN	Linear	IDW
Shallow Water	Coarser	0.046 ± 0.002	0.034 ± 0.001	0.038 ± 0.002
Original	0.017 ± 0.002	0.016 ± 0.002	0.017 ± 0.003
Finer	0.031 ± 0.003	0.017 ± 0.003	0.030 ± 0.002
Navier Stokes	Coarser	0.087 ± 0.006	0.069 ± 0.009	0.066 ± 0.006
Original	0.048 ± 0.009	0.041 ± 0.003	0.045 ± 0.010
Finer	0.054 ± 0.009	0.044 ± 0.004	0.049 ± 0.002
Scalar Flow	Coarser	0.041 ± 0.021	0.032 ± 0.009	0.035 ± 0.012
Original	0.019 ± 0.001	0.018 ± 0.000	0.018 ± 0.001
Finer	0.040 ± 0.016	0.026 ± 0.006	0.028 ± 0.007
is given, as discussed at the beginning of this section). The results are shown in Table 2. We see that some of the baseline models achieve reasonably good results on the fully-observed datasets, but they fail on partially-observed data, while our model maintains strong performance in all cases. Apart from the fully observed SHALLOW WATER dataset where FEN performs slightly better, our method outperforms other methods on all other datasets by a clear margin. See Appendix C for hyperparameter details. In Appendix E we demonstrate our model's capability to learn dynamics from noisy data. In Appendix F we show model predictions on different datasets.
Table 2: Test MAE for different models.
Model	Shallow Water (Full)	Shallow Water (Partial)	Navier Stokes (Full)	Navier Stokes (Partial)	Scalar Flow
NODE	0.036 ± 0.000	0.084 ± 0.001	0.053 ± 0.001	0.109 ± 0.001	0.056 ± 0.001
FEN	0.011 ± 0.002	0.064 ± 0.005	0.031 ± 0.001	0.108 ± 0.002	0.062 ± 0.005
SNPDE	0.019 ± 0.002	0.033 ± 0.001	0.042 ± 0.004	0.075 ± 0.002	0.059 ± 0.002
DINo	0.027 ± 0.001	0.063 ± 0.003	0.047 ± 0.001	0.113 ± 0.002	0.059 ± 0.001
MAgNet	NA	0.061 ± 0.001	NA	0.103 ± 0.003	0.056 ± 0.003
Ours	0.014 ± 0.002	0.016 ± 0.002	0.024 ± 0.003	0.041 ± 0.003	0.042 ± 0.001
Closest to our work is Ayed et al. (2022), where they considered the problem of learning PDEs from partial observations and proposed a discrete and grid-dependent model that is restricted to regular spatiotemporal grids. Another related work is that of Nguyen et al. (2020), where they proposed a variational inference framework for learning ODEs from noisy and partially-observed data. However, they consider only low-dimensional ODEs and are restricted to regular grids.
Other works considered learning the latent space PDE dynamics using the "encode-process-decode" approach. Pfaff et al. (2021) use GNN-based encoder and dynamics function and map the observations to the same spatial grid in the latent space and learn the latent space dynamics. Sanchez et al. (2020) use a similar approach but with CNNs and map the observations to a coarser latent grid and learn the coarse-scale dynamics. Wu et al. (2022) use CNNs to map observations to a low-dimensional latent vector and learn the latent dynamics. However, all these approaches are grid-dependent, limited to regular spatial/temporal grids, and require fully-observed data.
Interpolation has been used in numerous studies for various applications. Works such as (Alet et al., 2019; Jiang et al., 2020; Han et al., 2022) use interpolation to map latent states on coarse grids to observations on finer grids. Hua et al. (2022) used interpolation as a post-processing step to obtain continuous predictions, while Boussif et al. (2022) used it to recover observations at missing nodes.
Another approach for achieving grid-independence was presented in neural operators (Li et al., 2021; Lu et al., 2021), which learn a mapping between infinite-dimensional function spaces and represent the mapping in a grid-independent manner.
We proposed a novel space-time continuous, grid-independent model for learning PDE dynamics from noisy and partial observations on irregular spatiotemporal grids. Our contributions include an efficient generative modeling framework, a novel latent PDE model merging collocation and method of lines, and a data-efficient, grid-independent encoder design. The model demonstrates state-of-the-art performance on complex datasets, highlighting its potential for advancing data-driven PDE modeling and enabling accurate predictions of spatiotemporal phenomena in diverse fields. However, our model and encoder operate on every spatial and temporal location which might not be the most efficient approach and hinders scaling to extremely large grids, hence research into more efficient latent state extraction and dynamics modeling methods is needed.
Cagatay Yildiz, Markus Heinonen, and Harri Lähdesmäki. Ode2vae: Deep generative second order odes with bayesian neural networks. ArXiv, abs/1905.10994, 2019.
Yuan Yin, Matthieu Kirchmeyer, Jean-Yves Franceschi, Alain Rakotomamonjy, and patrick gallinari. Continuous PDE dynamics forecasting with implicit neural representations. In The Eleventh International Conference on Learning Representations, 2023. URL https://openreview.net/forum?id=B73niNjbPs.
A.1 Model specification.
Here we provide all details about our model specification. The joint distribution for our model is
$$
p \left(\boldsymbol {u} _ {1: M}, \boldsymbol {s} _ {1: B}, \theta_ {\mathrm {d y n}}, \theta_ {\mathrm {d e c}}\right) = p \left(\boldsymbol {u} _ {1: N} \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). \tag {24}
$$
Next, we specify each component in detail.
Parameter priors. The parameter priors are isotropic zero-mean multivariate normal distributions:
$$
p \left(\theta_ {\mathrm {d y n}}\right) = \mathcal {N} \left(\theta_ {\mathrm {d y n}} \mid \mathbf {0}, I\right), \tag {25}
$$
$$
p \left(\theta_ {\mathrm {d e c}}\right) = \mathcal {N} \left(\theta_ {\mathrm {d e c}} \mid \mathbf {0}, I\right), \tag {26}
$$
where  $\mathcal{N}$  is the normal distribution,  $\mathbf{0}$  is a zero vector, and  $I$  is the identity matrix, both have an appropriate dimensionality dependent on the number of encoder and dynamics parameters.
Continuity prior. We define the continuity prior as
$$
\begin{array}{l} p \left(\boldsymbol {s} _ {1: B} \mid \theta_ {\mathrm {d y n}}\right) = p \left(\boldsymbol {s} _ {1}\right) \prod_ {b = 2} ^ {B} p \left(\boldsymbol {s} _ {b} \mid \boldsymbol {s} _ {b - 1}, \theta_ {\mathrm {d y n}}\right), (27) \\ = \left[ \prod_ {j = 1} ^ {N} p \left(\boldsymbol {s} _ {1} ^ {j}\right) \right] \left[ \prod_ {b = 2} ^ {B} \prod_ {j = 1} ^ {N} p \left(\boldsymbol {s} _ {b} ^ {j} \mid \boldsymbol {s} _ {b - 1}, \theta_ {\mathrm {d y n}}\right) \right], (28) \\ = \left[ \prod_ {j = 1} ^ {N} \mathcal {N} \left(\boldsymbol {s} _ {1} ^ {j} | \boldsymbol {0}, I\right) \right] \left[ \prod_ {b = 2} ^ {B} \prod_ {j = 1} ^ {N} \mathcal {N} \left(\boldsymbol {s} _ {b} ^ {j} | \boldsymbol {z} \left(t _ {[ b ]}, \boldsymbol {x} _ {j}; t _ {[ b - 1 ]}, \boldsymbol {s} _ {b - 1}, \theta_ {\mathrm {d y n}}\right), \sigma_ {c} ^ {2} I\right). \right], (29) \\ \end{array}
$$
where  $\mathcal{N}$  is the normal distribution,  $\mathbf{0} \in \mathbb{R}^d$  is a zero vector,  $I \in \mathbb{R}^{d \times d}$  is the identity matrix, and  $\sigma_c \in \mathbb{R}$  is the parameter controlling the strength of the prior. Smaller values of  $\sigma_c$  tend to produce smaller gaps between the sub-trajectories.
Observation model
$$
\begin{array}{l} p \left(\boldsymbol {u} _ {1: N} \mid \boldsymbol {s} _ {1: B}, \theta_ {\mathrm {d y n}}, \theta_ {\mathrm {d e c}}\right) = \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) (30) \\ = \prod_ {b = 1} ^ {B} \prod_ {i \in \mathcal {I} _ {b}} \prod_ {j = 1} ^ {N} p \left(\boldsymbol {u} _ {i} ^ {j} \mid g _ {\theta_ {\mathrm {d e c}}} \left(\boldsymbol {z} \left(t _ {i}, \boldsymbol {x} _ {j}; t _ {[ b ]}, \boldsymbol {s} _ {b}, \theta_ {\mathrm {d y n}}\right)\right)\right) (31) \\ = \prod_ {b = 1} ^ {B} \prod_ {i \in \mathcal {I} _ {b}} \prod_ {j = 1} ^ {N} \mathcal {N} \left(\boldsymbol {u} _ {i} ^ {j} \mid g _ {\theta_ {\mathrm {d e c}}} \left(\boldsymbol {z} \left(t _ {i}, \boldsymbol {x} _ {j}; t _ {[ b ]}, \boldsymbol {s} _ {b}, \theta_ {\mathrm {d y n}}\right)\right), \sigma_ {u} ^ {2} I\right), (32) \\ \end{array}
$$
where  $\mathcal{N}$  is the normal distribution,  $\sigma_u^2$  is the observation noise variance, and  $I\in \mathbb{R}^{D\times D}$  is the identity matrix. Note again that  $z(t_{i},\pmb{x}_{j};t_{[b]},\pmb{s}_{b},\theta_{\mathrm{dyn}})$  above equals the ODE forward solution ODESolve  $(t_i;t_{[b]},s_b,\theta_{\mathrm{dyn}})$  at grid location  $\pmb{x}_j$ .
A.2 Approximate posterior specification.
Here we provide all details about the approximate posterior. We define the approximate posterior as
$$
q \left(\theta_ {\mathrm {d y n}}, \theta_ {\mathrm {d e c}}, \boldsymbol {s} _ {1: B}\right) = q \left(\theta_ {\mathrm {d y n}}\right) q \left(\theta_ {\mathrm {d e c}}\right) q \left(\boldsymbol {s} _ {1: B}\right) = q _ {\boldsymbol {\psi} _ {\mathrm {d y n}}} \left(\theta_ {\mathrm {d y n}}\right) q _ {\boldsymbol {\psi} _ {\mathrm {d e c}}} \left(\theta_ {\mathrm {d e c}}\right) \prod_ {b = 1} ^ {B} \prod_ {j = 1} ^ {N} q _ {\boldsymbol {\psi} _ {b} ^ {j}} \left(\boldsymbol {s} _ {b} ^ {j}\right). \tag {33}
$$
Dynamics parameters posterior. We define  $q_{\psi_{\mathrm{dyn}}}(\theta_{\mathrm{dyn}})$  as
$$
q _ {\psi_ {\mathrm {d y n}}} \left(\theta_ {\mathrm {d y n}}\right) = \mathcal {N} \left(\theta_ {\mathrm {d y n}} \mid \gamma_ {\mathrm {d y n}}, \operatorname {d i a g} \left(\tau_ {\mathrm {d y n}} ^ {2}\right)\right), \tag {34}
$$
where  $\gamma_{\mathrm{dyn}}$  and  $\tau_{\mathrm{dyn}}^2$  are vectors with an appropriate dimension (dependent on the number of dynamics parameters), and  $\mathrm{diag}(\tau_{\mathrm{dyn}}^2)$  is a matrix with  $\tau_{\mathrm{dyn}}^2$  on the diagonal. We define the vector of variational parameters as  $\psi_{\mathrm{dyn}} = (\gamma_{\mathrm{dyn}},\tau_{\mathrm{dyn}}^2)$ . We optimize directly over  $\psi_{\mathrm{dyn}}$  and initialize  $\gamma_{\mathrm{dyn}}$  using Xavier (Glorot and Bengio, 2010) initialization, while  $\tau_{\mathrm{dyn}}$  is initialized with each element equal to  $9\cdot 10^{-4}$ .
Decoder parameters posterior. We define  $q_{\psi_{\mathrm{dec}}}(\theta_{\mathrm{dec}})$  as
$$
q _ {\psi_ {\mathrm {d e c}}} \left(\theta_ {\mathrm {d e c}}\right) = \mathcal {N} \left(\theta_ {\mathrm {d e c}} \mid \gamma_ {\mathrm {d e c}}, \operatorname {d i a g} \left(\tau_ {\mathrm {d e c}} ^ {2}\right)\right), \tag {35}
$$
where  $\gamma_{\mathrm{dec}}$  and  $\tau_{\mathrm{dec}}^2$  are vectors with an appropriate dimension (dependent on the number of decoder parameters), and  $\mathrm{diag}(\tau_{\mathrm{dec}}^2)$  is a matrix with  $\tau_{\mathrm{dec}}^2$  on the diagonal. We define the vector of variational parameters as  $\psi_{\mathrm{dec}} = (\gamma_{\mathrm{dec}},\tau_{\mathrm{dec}}^2)$ . We optimize directly over  $\psi_{\mathrm{dec}}$  and initialize  $\gamma_{\mathrm{dec}}$  using Xavier (Glorot and Bengio, 2010) initialization, while  $\tau_{\mathrm{dec}}$  is initialized with each element equal to  $9\cdot 10^{-4}$ .
Shooting variables posterior. We define  $q_{\psi_b^j}(s_b^j)$  as
$$
q _ {\psi_ {b} ^ {j}} \left(\boldsymbol {s} _ {b} ^ {j}\right) = \mathcal {N} \left(\boldsymbol {s} _ {b} ^ {j} \mid \boldsymbol {\gamma} _ {b} ^ {j}, \operatorname {d i a g} \left(\left[ \boldsymbol {\tau} _ {b} ^ {j} \right] ^ {2}\right)\right), \tag {36}
$$
where the vectors  $\gamma_b^j, \tau_b^j \in \mathbb{R}^d$  are returned by the encoder  $h_{\theta_{\mathrm{enc}}}$ , and  $\mathrm{diag}([\tau_b^j]^2)$  is a matrix with  $[\tau_b^j]^2$  on the diagonal. We define the vector of variational parameters as  $\psi_b^j = (\gamma_b^j, [\tau_b^j])$ . Because the variational inference for the shooting variables is amortized, our model is trained w.r.t. the parameters of the encoder network,  $\theta_{\mathrm{enc}}$ .
B Appendix B
B.1 Derivation of ELBO.
For our model and the choice of the approximate posterior the ELBO can be written as