SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Valerii Iakovlev Markus Heinonen Harri Lahdesmaki
Department of Computer Science, Aalto University, Finland {valerii.iakovlev, markus.o.heinonen, harri.lahdesmaki}@aalto.fi
We introduce a novel grid-independent model for learning partial differential equations (PDEs) from noisy and partial observations on irregular spatiotemporal grids. We propose a space-time continuous latent neural PDE model with an efficient probabilistic framework and a novel encoder design for improved data efficiency and grid independence. The latent state dynamics are governed by a PDE model that combines the collocation method and the method of lines. We employ amortized variational inference for approximate posterior estimation and utilize a multiple shooting technique for enhanced training speed and stability. Our model demonstrates state-of-the-art performance on complex synthetic and real-world datasets, overcoming limitations of previous approaches and effectively handling partially-observed data. The proposed model outperforms recent methods, showing its potential to advance data-driven PDE modeling and enabling robust, grid-independent modeling of complex partially-observed dynamic processes.
Modeling spatiotemporal processes allows to understand and predict the behavior of complex systems that evolve over time and space (Cressie and Wikle, 2011). Partial differential equations (PDEs) are a popular tool for this task as they have a solid mathematical foundation (Evans, 2010) and can describe the dynamics of a wide range of physical, biological, and social phenomena (Murray, 2002; Hirsch, 2007). However, deriving PDEs can be challenging, especially when the system's underlying mechanisms are complex and not well understood. Data-driven methods can bypass these challenges (Brunton and Kutz, 2019). By learning the underlying system dynamics directly from data, we can develop accurate PDE models that capture the essential features of the system. This approach has changed our ability to model complex systems and make predictions about their behavior in a data-driven manner.
While current data-driven PDE models have been successful at modeling complex spatiotemporal phenomena, they often operate under various simplifying assumptions such as regularity of the spatial or temporal grids (Long et al., 2018; Kochkov et al., 2021; Pfaff et al., 2021; Li et al., 2021; Han et al., 2022; Poli et al., 2022), discreteness in space or time (Seo et al., 2020; Pfaff et al., 2021; Lienen and Gunnemann, 2022; Brandstetter et al., 2022), and availability of complete and noiseless observations (Long et al., 2018; Pfaff et al., 2021; Wu et al., 2022). Such assumptions become increasingly limiting in more realistic scenarios with scarce data and irregularly spaced, noisy and partial observations.
We address the limitations of existing methods and propose a space-time continuous and grid-independent model that can learn PDE dynamics from noisy and partial observations made on irregular spatiotemporal grids. Our main contributions include:
Source code and datasets can be found in our github repository.
- Development of an efficient generative modeling framework for learning latent neural PDE models from noisy and partially-observed data;
- Novel PDE model that merges two PDE solution techniques – the collocation method and the method of lines – to achieve space-time continuity, grid-independence, and data efficiency;
- Novel encoder design that operates on local spatiotemporal neighborhoods for improved data-efficiency and grid-independence.
Our model demonstrates state-of-the-art performance on complex synthetic and real-world datasets, opening up the possibility for accurate and efficient modeling of complex dynamic processes and promoting further advancements in data-driven PDE modeling.
In this work we are concerned with modeling of spatiotemporal processes. For brevity, we present our method for a single observed trajectory, but extension to multiple trajectories is straightforward. We observe a spatiotemporal dynamical system evolving over time on a spatial domain  $\Omega$ . The observations are made at  $M$  arbitrary consecutive time points  $t_{1:M} \coloneqq (t_1, \ldots, t_M)$  and  $N$  arbitrary observation locations  $\pmb{x}_{1:N} \coloneqq (\pmb{x}_1, \ldots, \pmb{x}_N)$ , where  $\pmb{x}_i \in \Omega$ . This generates a sequence of observations  $\pmb{u}_{1:M} \coloneqq (\pmb{u}_1, \ldots, \pmb{u}_M)$ , where  $\pmb{u}_i \in \mathbb{R}^{N \times D}$  contains  $D$ -dimensional observations at the  $N$  observation locations. We define  $\pmb{u}_i^j$  as the observation at time  $t_i$  and location  $\pmb{x}_j$ . The number of time points and observation locations may vary between different observed trajectories.
We assume the data is generated by a dynamical system with a latent state  $\boldsymbol{z}(t, \boldsymbol{x}) \in \mathbb{R}^d$ , where  $t$  is time and  $\boldsymbol{x} \in \Omega$  is spatial location. The latent state is governed by an unknown PDE and is mapped to the observed state  $\boldsymbol{u}(t, \boldsymbol{x}) \in \mathbb{R}^D$  by an unknown observation function  $g$  and likelihood model  $p$ :
$$
\frac {\partial \boldsymbol {z} (t , x)}{\partial t} = F (\boldsymbol {z} (t, \boldsymbol {x}), \partial_ {\boldsymbol {x}} \boldsymbol {z} (t, \boldsymbol {x}), \partial_ {\boldsymbol {x}} ^ {2} \boldsymbol {z} (t, \boldsymbol {x}), \dots), \tag {1}
$$
$$
\boldsymbol {u} (t, \boldsymbol {x}) \sim p (g (\boldsymbol {z} (t, \boldsymbol {x}))), \tag {2}
$$
where  $\partial_{\pmb{x}}^{\bullet}\pmb {z}(t,\pmb {x})$  denotes partial derivatives wrt  $\pmb{x}$
In this work we make two assumptions that are highly relevant in real-world scenarios. First, we assume partial observations, that is, the observed state  $\pmb{u}(t,\pmb{x})$  does not contain all information about the latent state  $z(t,\pmb{x})$  (e.g.,  $z(t,\pmb{x})$  contains pressure and velocity, but  $\pmb{u}(t,\pmb{x})$  contains information only about the pressure). Second, we assume out-of-distribution time points and observation locations, that is, their number, positions, and density can change arbitrarily at test time.
Here we describe the model components (Sec. 3.1) which are then used to construct the generative model (Sec. 3.2).
3.1 Model components
Our model consists of four parts: space-time continuous latent state  $\mathbf{z}(t, \mathbf{x})$  and observed state  $\mathbf{u}(t, \mathbf{x})$ , a dynamics function  $F_{\theta_{\mathrm{dyn}}}$  governing the temporal evolution of the latent state, and an observation function  $g_{\theta_{\mathrm{dec}}}$  mapping the latent state to the observed state (see Figure 1). Next, we describe these components in detail.
Figure 1: Model sketch. Initial latent state  $z(t_{1}, \boldsymbol{x})$  is evolved via  $F_{\theta_{\mathrm{dyn}}}$  to the following latent states which are then mapped to the observed states by  $g_{\theta_{\mathrm{dec}}}$ .
Latent state. To define a space-time continuous latent state  $\pmb{z}(t, \pmb{x}) \in \mathbb{R}^d$ , we introduce  $\pmb{z}(t) \coloneqq (\pmb{z}^1(t), \dots, \pmb{z}^N(t)) \in \mathbb{R}^{N \times d}$ , where each  $\pmb{z}^i(t) \in \mathbb{R}^d$  corresponds to the observation location  $\pmb{x}_i$ . Then, we define the latent state  $\pmb{z}(t, \pmb{x})$  as a spatial interpolant of  $\pmb{z}(t)$ :
$$
\boldsymbol {z} (t, \boldsymbol {x}) := \operatorname {I n t e r p o l a t e} (\boldsymbol {z} (t)) (\boldsymbol {x}), \tag {3}
$$
where  $\operatorname{Interpolate}(\cdot)$  maps  $z(t)$  to an interpolant which can be evaluated at any spatial location  $x \in \Omega$  (see Figure 2). We do not rely on a particular interpolation method, but in this work we use linear interpolation as it shows good performance and facilitates efficient implementation.
Latent state dynamics. Given a space-time continuous latent state, one can naturally define its dynamics in terms of a PDE:
$$
\frac {\partial \boldsymbol {z} (t , x)}{\partial t} = F _ {\theta_ {\mathrm {d y n}}} (\boldsymbol {z} (t, \boldsymbol {x}), \partial_ {\boldsymbol {x}} \boldsymbol {z} (t, \boldsymbol {x}), \partial_ {\boldsymbol {x}} ^ {2} \boldsymbol {z} (t, \boldsymbol {x}), \dots), \tag {4}
$$
where  $F_{\theta_{\mathrm{dyn}}}$  is a dynamics function with parameters  $\theta_{\mathrm{dyn}}$ . This is a viable approach known as the collocation method (Kansa, 1990; Cheng, 2009), but it has several limitations. It requires us to decide which partial derivatives to include in the dynamics function, and also requires an interpolant which has all the selected partial derivatives (e.g., linear interpolant has only first order derivatives). To avoid these limitations, we combine the collocation method with
another PDE solution technique known as the method of lines (Schiesser, 1991; Hamdi et al., 2007), which approximates spatial derivatives  $\partial_{\pmb{x}}^{\bullet}\pmb {z}(t,\pmb {x})$  using only evaluations of  $\pmb {z}(t,\pmb {x})$ , and then let the dynamics function approximate all required derivatives in a data-driven manner. To do that, we define the spatial neighborhood of  $\pmb{x}$  as  $\mathcal{N}_{\mathrm{S}}(\pmb {x})$ , which is a set containing  $\pmb{x}$  and its spatial neighbors, and also define  $\pmb {z}(t,\mathcal{N}_{\mathrm{S}}(\pmb {x}))$ , which is a set of evaluations of the interpolant  $\pmb {z}(t,\pmb {x})$  at points in  $\mathcal{N}_{\mathrm{S}}(\pmb {x})$ :
Figure 2: Latent state  $z(t, \boldsymbol{x})$  defined as an interpolant of  $z(t) \coloneqq (z^1(t), \dots, z^4(t))$ .
$$
\mathcal {N} _ {\mathrm {S}} (\boldsymbol {x}) := \left\{\boldsymbol {x} ^ {\prime} \in \Omega : \boldsymbol {x} ^ {\prime} = \boldsymbol {x} \text {o r} \boldsymbol {x} ^ {\prime} \text {i s a s p a t i a l n e i g h b o r o f} \boldsymbol {x} \right\}, \tag {5}
$$
$$
\boldsymbol {z} (t, \mathcal {N} _ {\mathrm {S}} (\boldsymbol {x})) := \left\{\boldsymbol {z} \left(t, \boldsymbol {x} ^ {\prime}\right): \boldsymbol {x} ^ {\prime} \in \mathcal {N} _ {\mathrm {S}} (\boldsymbol {x}) \right\}, \tag {6}
$$
and assume that this information is sufficient to approximate all required spatial derivatives at  $\pmb{x}$ . This is a reasonable assumption since, e.g., finite differences can approximate derivatives using only function values and locations of the evaluation points. Hence, we define the dynamics of  $z(t, \pmb{x})$  as
$$
\frac {\partial \boldsymbol {z} (t , \boldsymbol {x})}{\partial t} = F _ {\theta_ {\mathrm {d y n}}} \left(\mathcal {N} _ {\mathrm {S}} (\boldsymbol {x}), \boldsymbol {z} (t, \mathcal {N} _ {\mathrm {S}} (\boldsymbol {x}))\right), \tag {7}
$$
which is defined only in terms of the values of the latent state, but not its spatial derivatives.
One way to define the spatial neighbors for  $\pmb{x}$  is in terms of the observation locations  $x_{1:N}$  (e.g., use the nearest ones) as was done, for example, in (Long et al., 2018; Pfaff et al., 2021; Lienen and Gunnemann, 2022). Instead, we utilize continuity of the latent state  $z(t, \pmb{x})$ , and define the spatial neighbors in a grid-independent manner as a fixed number of points arranged in a predefined pattern around  $\pmb{x}$  (see Figure 3). This allows to fix the shape and size of the spatial neighborhoods in advance, making them independent of the observation locations. In this work we use the spatial neighborhood consisting of two concentric circles of radius  $r$  and  $r/2$ , each circle contains 8 evaluation points as in Figure 3. In Appendix D we compare neighborhoods of various shapes and sizes.
Equation 7 allows to simulate the temporal evolution of  $z(t, \boldsymbol{x})$  at any spatial location. However, since  $z(t, \boldsymbol{x})$  is defined only in terms of a spatial interpolant of  $z(t)$  (see Eq. 3), with  $z^i(t) = z(t, \boldsymbol{x}_i)$ , it is sufficient to simulate the latent state dynamics only at the observation locations  $\boldsymbol{x}_{1:N}$ . Hence, we can completely characterize the latent state dynamics in terms of a system of  $N$  ODEs:
Figure 3: Example of  $\mathcal{N}_{\mathrm{S}}(\pmb{x}_i)$ . Instead of using the observation locations (dots) to define spatial neighbors, we use spatial locations arranged in a fixed predefined pattern (crosses).
$$
\frac {d \boldsymbol {z} (t)}{d t} := \left( \begin{array}{c} \frac {d \boldsymbol {z} ^ {- 1} (t)}{d t} \\ \vdots \\ \frac {d \boldsymbol {z} ^ {N} (t)}{d t} \end{array} \right) = \left( \begin{array}{c} \frac {\partial \boldsymbol {z} (t , \boldsymbol {x} _ {1})}{\partial t} \\ \vdots \\ \frac {\partial \boldsymbol {z} (t , \boldsymbol {x} _ {N})}{\partial t} \end{array} \right) = \left( \begin{array}{c} F _ {\theta_ {\mathrm {d y n}}} \left(\mathcal {N} _ {\mathrm {S}} (\boldsymbol {x} _ {1}), \boldsymbol {z} (t, \mathcal {N} _ {\mathrm {S}} (\boldsymbol {x} _ {1}))\right) \\ \vdots \\ F _ {\theta_ {\mathrm {d y n}}} \left(\mathcal {N} _ {\mathrm {S}} (\boldsymbol {x} _ {N}), \boldsymbol {z} (t, \mathcal {N} _ {\mathrm {S}} (\boldsymbol {x} _ {N}))\right) \end{array} \right). \tag {8}
$$
For convenience, we define  $\boldsymbol{z}(t; t_1, \boldsymbol{z}_1, \theta_{\mathrm{dyn}}) \coloneqq \mathrm{ODESolve}(t; t_1, \boldsymbol{z}_1, \theta_{\mathrm{dyn}})$  as the solution of the ODE system in Equation 8 at time  $t$  with initial state  $\boldsymbol{z}(t_1) = \boldsymbol{z}_1$  and parameters  $\theta_{\mathrm{dyn}}$ . We also define  $\boldsymbol{z}(t, \boldsymbol{x}; t_1, \boldsymbol{z}_1, \theta_{\mathrm{dyn}})$  as the spatial interpolant of  $\boldsymbol{z}(t; t_1, \boldsymbol{z}_1, \theta_{\mathrm{dyn}})$  as in Equation 3. We solve the ODEs using off the shelf differentiable ODE solvers from torchdiffeq package (Chen, 2018). Note that we solve for the state  $\boldsymbol{z}(t)$  only at the observation locations  $\boldsymbol{x}_{1:N}$ , so to get the neighborhood values  $\boldsymbol{z}(t, \mathcal{N}_{\mathrm{S}}(\boldsymbol{x}_i))$  we perform interpolation at every step of the ODE solver.
Observation function. We define the mapping from the latent space to the observation space as a parametric function  $g_{\theta_{\mathrm{dec}}}$  with parameters  $\theta_{\mathrm{dec}}$ :
$$
\boldsymbol {u} (t, \boldsymbol {x}) \sim \mathcal {N} \left(g _ {\theta_ {\mathrm {d e c}}} (\boldsymbol {z} (t, \boldsymbol {x})) , \sigma_ {u} ^ {2} I _ {D}\right), \tag {9}
$$
where  $\mathcal{N}$  is the Gaussian distribution,  $\sigma_u^2$  is noise variance, and  $I_D$  is  $D$ -by- $D$  identity matrix.
3.2 Generative model
Training models of dynamic systems is often challenging due to long training times and training instabilities (Ribeiro et al., 2020; Metz et al., 2021). To alleviate these problems, various heuristics have been proposed, such as progressive lengthening and splitting of the training trajectories (Yildiz et al., 2019; Um et al., 2020). We use multiple shooting (Bock and Plitt, 1984; Voss et al., 2004), a simple and efficient technique which has demonstrated its effectiveness in ODE learning applications (Jordana et al., 2021; Hegde et al., 2022). We extent the multiple shooting framework for latent ODE models presented in (Iakovlev et al., 2023) to our PDE modeling setup by introducing spatial dimensions in the latent state and designing an encoder adapted specifically to the PDE setting (Section 4.2).
Multiple shooting splits a single trajectory  $\{z(t_i)\}_{i = 1,\dots ,M}$  with one initial state  $z_{1}$  into  $B$  consecutive non-overlapping subtrajectories  $\{z(t_i)\}_{i\in \mathcal{I}_b}$ ,  $b = 1,\ldots ,B$  with  $B$  initial states  $s_{1:B}:= (s_1,\ldots ,s_B)$  while imposing a continuity penalty between the subtrajectories (see Figure 4). The index set  $\mathcal{I}_b$  contains time point indices for the  $b$ th sub-trajectory. We also denote the temporal posi
tion of  $s_b$  as  $t_{[b]}$  and place  $s_b$  at the first time point preceding the  $b$ 'th sub-trajectory (except  $s_1$  which is placed at  $t_1$ ). Note that the shooting states  $s_b$  have the same dimension as the original latent state  $z(t)$  i.e.,  $s_b \in \mathbb{R}^{N \times d}$ . Multiple shooting allows to parallelize the simulation over the sub-trajectories and shortens the simulation intervals thus improving the training speed and stability. In Appendix D we demonstrate the effect of multiple shooting on the model training and prediction accuracy.
We begin by defining the prior over the unknown model parameters and initial states:
$$
p \left(\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 {10}
$$
where  $p(\theta_{\mathrm{dyn}})$  and  $p(\theta_{\mathrm{dec}})$  are zero-mean diagonal Gaussians, and the continuity inducing prior  $p(\boldsymbol{s}_{1:B}|\theta_{\mathrm{dyn}})$  is defined as in (Iakovlev et al., 2023)
$$
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). \tag {11}
$$
Intuitively, the continuity prior  $p(\boldsymbol{s}_b|\boldsymbol{s}_{b-1}, \theta_{\mathrm{dyn}})$  takes the initial latent state  $\boldsymbol{s}_{b-1}$ , simulates it forward from time  $t_{[b-1]}$  to  $t_{[b]}$  to get  $\boldsymbol{\mu}_{[b]} = \mathrm{ODESolve}(t_{[b]}; t_{[b-1]}, \boldsymbol{s}_{b-1}, \theta_{\mathrm{dyn}})$ , and then forces  $\boldsymbol{\mu}_{[b]}$  to approximately match the initial state  $\boldsymbol{s}_b$  of the next sub-trajectory, thus promoting continuity of the full trajectory. We assume the continuity inducing prior factorizes across the grid points, i.e.,
$$
\begin{array}{l} p \left(\boldsymbol {s} _ {1: B} \mid \theta_ {\mathrm {d y n}}\right) = \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], (12) \\ = \left[ \prod_ {j = 1} ^ {N} p \left(\boldsymbol {s} _ {1} ^ {j}\right) \right] \left[ \prod_ {b = 2} ^ {B} \prod_ {j = 1} ^ {N} \mathcal {N} \left(\boldsymbol {s} _ {b} ^ {j} \mid \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 _ {d}\right) \right], (13) \\ \end{array}
$$
where  $p(s_1^j)$  is a diagonal Gaussian, and parameter  $\sigma_c^2$  controls the strength of the prior. Note that the term  $\boldsymbol{z}(t_{[b]},\boldsymbol{x}_j;t_{[b-1]},\boldsymbol{s}_{b-1},\theta_{\mathrm{dyn}})$  in Equation 13 equals the ODE forward solution ODESolve  $(t_{[b]};t_{[b-1]},\boldsymbol{s}_{b-1},\theta_{\mathrm{dyn}})$  at grid location  $\boldsymbol{x}_j$ .
Figure 4: Multiple shooting splits a trajectory with one initial state (top) into two subtrajectories with two initial states (bottom) and tries to minimize the gap between subtrajectories (orange arrow).
Finally, we define our generative in terms of the following sampling procedure:
$$
\theta_ {\mathrm {d y n}}, \theta_ {\mathrm {d e c}}, s _ {1: B} \sim p \left(\theta_ {\mathrm {d y n}}\right) p \left(\theta_ {\mathrm {d e c}}\right) p \left(s _ {1: B} \mid \theta_ {\mathrm {d y n}}\right), \tag {14}
$$
$$
\boldsymbol {z} \left(t _ {i}\right) = \boldsymbol {z} \left(t _ {i}; t _ {[ b ]}, \boldsymbol {s} _ {b}, \theta_ {\mathrm {d y n}}\right), \quad b \in \{1, \dots , B \}, i \in \mathcal {I} _ {b}, \tag {15}
$$
$$
\boldsymbol {u} _ {i} ^ {j} \sim p \left(\boldsymbol {u} _ {i} ^ {j} \mid g _ {\theta_ {\mathrm {d e c}}} (\boldsymbol {z} \left(t _ {i}, \boldsymbol {x} _ {j}\right)), \quad i = 1, \dots , M, j = 1, \dots , N, \right. \tag {16}
$$
with the following joint distribution (see Appendix A for details about the model specification.):
$$
p \left(\boldsymbol {u} _ {1: M}, \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} \left[ p \left(\boldsymbol {u} _ {i} ^ {j} \mid \boldsymbol {s} _ {b}, \theta_ {\mathrm {d y n}}, \theta_ {\mathrm {d e c}}\right) \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 {17}
$$
4 Parameter Inference, Encoder, and Forecasting
4.1 Amortized variational inference
We approximate the true posterior over the model parameters and initial states  $p(\boldsymbol{s}_{1:B},\theta_{\mathrm{dyn}},\theta_{\mathrm{dec}}|\boldsymbol{u}_{1:M})$  using variational inference (Blei et al., 2017) with the following approximate posterior:
$$
q \left(\theta_ {\mathrm {d y n}}, \theta_ {\mathrm {d e c}}, s _ {1: B}\right) = q \left(\theta_ {\mathrm {d y n}}\right) q \left(\theta_ {\mathrm {d e c}}\right) q \left(s _ {1: B}\right) = q _ {\psi_ {\mathrm {d y n}}} \left(\theta_ {\mathrm {d y n}}\right) q _ {\psi_ {\mathrm {d e c}}} \left(\theta_ {\mathrm {d e c}}\right) \prod_ {b = 1} ^ {B} \prod_ {j = 1} ^ {N} q _ {\psi_ {b} ^ {j}} \left(s _ {b} ^ {j}\right), \tag {18}
$$
where  $q_{\psi_{\mathrm{dyn}}}, q_{\psi_{\mathrm{dec}}}$  and  $q_{\psi_b^j}$  are diagonal Gaussians, and  $\psi_{\mathrm{dyn}}, \psi_{\mathrm{dec}}$  and  $\psi_b^j$  are variational parameters. To avoid direct optimization over the local variational parameters  $\psi_b^j$ , we use amortized variational inference (Kingma and Welling, 2013) and train an encoder  $h_{\theta_{\mathrm{enc}}}$  with parameters  $\theta_{\mathrm{enc}}$  which maps observations  $\pmb{u}_{1:M}$  to  $\pmb{\psi}_b^j$  (see Section 4.2). For brevity, we sometimes omit the dependence of approximate posteriors on variational parameters and simply write e.g.,  $q(\pmb{s}_b^j)$ .
In variational inference the best approximation of the posterior is obtained by minimizing the Kullback-Leibler divergence: KL  $\left[q(\theta_{\mathrm{dyn}},\theta_{\mathrm{dec}},s_{1:B})\| p(\theta_{\mathrm{dyn}},\theta_{\mathrm{dec}},s_{1:B}|u_{1:N})\right]$ , which is equivalent to maximizing the evidence lower bound (ELBO), defined for our model as:
$$
\mathcal {L} = \underbrace {\sum_ {b = 1} ^ {B} \sum_ {i \in \mathcal {I} _ {b}} \sum_ {j = 1} ^ {N} \mathbb {E} _ {q \left(\boldsymbol {s} _ {b} , \theta_ {\mathrm {d y n}} , \theta_ {\mathrm {d e c}}\right)} \left[ \log p \left(\boldsymbol {u} _ {i} ^ {j} \mid \boldsymbol {s} _ {b} , \theta_ {\mathrm {d y n}} , \theta_ {\mathrm {d e c}}\right) \right]} _ {(i) \text {o b s e r v a t i o n m o d e l}} - \underbrace {\sum_ {j = 1} ^ {N} \mathrm {K L} \left[ q \left(\boldsymbol {s} _ {1} ^ {j}\right) \| p \left(\boldsymbol {s} _ {1} ^ {j}\right) \right]} _ {(i i) \text {i n i t i a l s t a t e p r i o r}}
$$
$$
-\underbrace{\sum_{b = 2}^{B}\sum_{j = 1}^{N}\mathbb{E}_{q(\theta_{\mathrm{dyn}},\boldsymbol{s}_{b - 1})}\left[\mathrm{KL}\left[q(\boldsymbol{s}_{b}^{j})\| p(\boldsymbol{s}_{b}^{j}|\boldsymbol{s}_{b - 1},\theta_{\mathrm{dyn}})\right]\right]}_{(iii)\text{continuity prior}} - \underbrace{\mathrm{KL}\left[q(\theta_{\mathrm{dyn}})\|p(\theta_{\mathrm{dyn}})\right]}_{(iv)\text{dynamics prior}} - \underbrace{\mathrm{KL}\left[q(\theta_{\mathrm{dec}})\|p(\theta_{\mathrm{dec}})\right]}_{(v)\text{decoder prior}}.
$$
The terms  $(ii)$ ,  $(iv)$ , and  $(v)$  are computed analytically, while terms  $(i)$  and  $(iii)$  are approximated using Monte Carlo integration for expectations, and numerical ODE solvers for initial value problems. See Appendix A and B approximate posterior details and derivation and computation of the ELBO.
4.2 Encoder
Here we describe our encoder which maps observations  $\pmb{u}_{1:M}$  to local variational parameters  $\psi_{b}^{j}$  required to sample the initial latent state of the sub-trajectory  $b$  at time point  $t_{[b]}$  and observation location  $\pmb{x}_j$ . Similarly to our model, the encoder should be data-efficient and grid-independent.
Similarly to our model (Section 3.1), we enable grid-independence by making the encoder operate on spatial interpolants of the observations  $\pmb{u}_{1:M}$  (even if they are noisy):