SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Kevin Course
kevin/course@mail.utoronto.ca
Prasanth B. Nair
prasanth.nair@utoronto.ca
We consider the problem of inferring latent stochastic differential equations (SDEs) with a time and memory cost that scales independently with the amount of data, the total length of the time series, and the stiffness of the approximate differential equations. This is in stark contrast to typical methods for inferring latent differential equations which, despite their constant memory cost, have a time complexity that is heavily dependent on the stiffness of the approximate differential equation. We achieve this computational advancement by removing the need to solve differential equations when approximating gradients using a novel amortization strategy coupled with a recently derived reparametrization of expectations under linear SDEs. We show that, in practice, this allows us to achieve similar performance to methods based on adjoint sensitivities with more than an order of magnitude fewer evaluations of the model in training.
Recent years have seen the rise of continuous time models for dynamical system modeling [1]. As compared to traditional autoregressive style models [2], continuous time models are useful because they can deal with non-evenly spaced observations, they enable multilevel/hierarchical and adaptive prediction schemes, and because physics is (mostly) done in continuous time. For example, recent developments in inferring continuous time differential equations from data has been met by a flurry of work in endowing models with physics informed priors [3, 4, 5].
Despite their advantages, continuous time models remain significantly more computationally challenging to train than their autoregressive counterparts due to their reliance on adjoint methods for estimating gradients. Adjoint methods introduce a significant computational burden in training because they require solving a pair of initial value problems to estimate gradients. Solving such initial value problems as a part of an iterative optimization procedure is computationally demanding for the following reasons:
(i) Gradient based updates to models of differential equations can cause them to become extremely stiff. This will have the effect of causing the cost per iteration to explode mid-optimization.
(ii) With the exception of parareal methods [6], differential equation solvers are fundamentally iterative sequential methods. This makes them poorly suited to being parallelized on modern parallel computing hardware.
In accordance with such challenges a number of methods have been introduced to speed up training of continuous time models including regularizing dynamics [7, 8] and replacing ordinary differential equation (ODE) solvers with integration methods where possible [9]. Despite the computational advancements brought about by such innovations, continuous time models remain expensive to train in comparison to discrete time models.
In addition to these computational challenges, it is well-known that adjoint methods suffer from stability issues when approximating gradients of time averaged quantities over long time intervals for chaotic systems [10].
In the current work, we present a memory and time efficient method for inferring nonlinear, latent stochastic differential equations (SDEs) from high-dimensional time-series datasets. In contrast to standard approaches for inferring latent differential equations that rely on adjoint sensitivities [1, 11, 12], our approach removes the requirement of solving differential equations entirely. We accomplish this advancement by coupling a novel amortization strategy with a recently derived reparametrization for expectations under Markov Gaussian processes [13]. We show that our approach can be used to approximate gradients of the evidence lower bound used to train latent SDEs with a time and memory cost that is independent of the amount of data, the length of the time series, and the stiffness of the approximate differential equations. The asymptotic complexity for our approach is compared to well-known methods from the literature in Table 1. We note that our method has a constant cost that is chosen by the user. Moreover, we will show that our method is embarrassingly parallel (i.e. all evaluations of the model can be performed in parallel over each iteration) whereas, we reiterate, differential equation solvers are iterative sequential methods.
Method	Time	Memory
Deterministic adjoints (Neural ODE) [1]	O(J)	O(1)
Stochastic adjoints [12]	O(J log J)	O(1)
Backprop through solver [14]	O(J)	O(J)
Amortized reparametrization (ours)	O(R)	O(R)
Table 1: Asymptotic complexity comparison for approximating gradients. Units are given by the number of evaluations of the differential equation (gradient field for ODEs or drift and diffusion function for SDEs). Here  $J$  is the number of sequential evaluations and  $R$  is the number of parallel evaluations.  $J$  is adaptively chosen by the differential equation solver and is a function of the stiffness of the differential equation meaning it can change (and possibly explode) while optimizing. In contrast,  $R$  is a fixed constant used to control the variance of gradient approximations. While we could choose  $R = 1$  and would still arrive with unbiased approximations for gradients, in practice we found choosing  $R \approx 10^2$  worked well for the problems we considered.
The applications of our method span various standard generative modeling tasks, such as auto-encoding, denoising, inpainting, and super-resolution [15], particularly tailored for high-dimensional time-series. Crucially, the computational efficiency of our approach not only enables the allocation of more computational resources towards hyperparameter tuning but also democratizes access to state-of-the-art methods by making them feasible to train on lower performance hardware.
In the next section we provide a description of the theory underpinning our work with the main result of a stochastic, unbiased estimate for gradients appearing in Lemma 1. In Section 4 we provide a number of numerical studies including learning latent neural SDEs from video and performance benchmarking on a motion capture dataset. Notably we show that we are able to achieve comparable performance to methods based on adjoints with more than one order of magnitude fewer evaluations of the model in training (Section 4.1). We also demonstrate that our approach does not suffer from the numerical instabilities that plague adjoint methods for long time-series with chaotic systems (Section 4.2). Finally, we close with a discussion of the limitations of our approach as well as some suggestions for future work. All code is available at github.com/coursekevin/arlatentsde.
2.1 Problem description
Consider a time-series dataset  $\mathcal{D} = \{(x_i, t_i)\}_{i=1}^N$ , where  $t_i \in \mathbb{R}$  is the time stamp associated with the observation,  $x_i \in \mathbb{R}^D$ . For example  $x_i$  may be an indirect observation of some underlying dynamical system, a video-frame, or a snapshot of a spatio-temporal field. We assume that each observation was generated via the following process: first a latent trajectory,  $z(t)$ , is sampled from a general,
nonlinear SDE with time-dependent diffusion,
$$
d z = f _ {\theta} (t, z) d t + L _ {\theta} (t) d \beta , \tag {1}
$$
with initial condition  $p_{\theta}(z_0)$  and then each  $x_i$  is generated from the conditional distribution  $p_{\theta}(x \mid z(t_i))$  where  $z(t_i)$  is the realization of the latent trajectory at time stamp  $t_i$ . Here  $f_{\theta} : \mathbb{R} \times \mathbb{R}^d \to \mathbb{R}^d$  is called the drift function,  $L_{\theta} : \mathbb{R} \to \mathbb{R}^{d \times d}$  is called the dispersion matrix, and  $\beta$  denotes Brownian motion with diffusion matrix  $\Sigma$ . We will model both the drift function of the SDE and the conditional distribution using neural networks in this work. This combination of SDE and conditional distribution defines the generative model that we wish to infer.
Like in the standard generative modeling setting, the "true" parameters defining the latent SDE and the conditional likelihood, as well as the specific realization of the latent state, remain unknown. In addition, the posterior over the latent state at a point in time,  $p_{\theta}(z(t) \mid \mathcal{D})$ , is intractable. Our objectives are two-fold, we wish to: (i) infer a likely setting for the parameters,  $\theta$ , that is well-aligned with the observed data and (ii) infer a parametric model for the posterior over the latent state at a particular point in time as a function of a small window of the observations, for example  $q_{\phi}(z(t) \mid x_i, x_{i+1}, \ldots, x_{i+M}) \approx p_{\theta}(z(t) \mid \mathcal{D})$  for  $t \in [t_i, t_{i+M}]$  and  $1 \leq M << N$ . This latent variable model is depicted in Figure 1. In this work we will tackle this problem statement using the machinery of stochastic variational inference [16].
Figure 1: The solid lines indicate the data generating process that depends on the parameters,  $\theta$ . The dashed lines indicate the approximate variational posterior that depends on the parameters,  $\phi$ .
Before proceeding with an explanation of our approach, it is worthwhile taking a moment to consider why assuming the latent state is a realization of a SDE, as opposed to an ordinary differential equation (ODE) with random initial conditions [1, 11, 17], merits the additional mathematical machinery. First, there is a long history in using SDEs in the physical sciences to model systems (even deterministic systems) that lack predictability due to super-sensitivity to parameters and or initial conditions [18]. One reason for the preference towards this modeling paradigm is that deterministic ODEs with random initial conditions only permit uncertainty to enter in their initial condition. This means that we are assuming that all randomness which may affect the latent trajectory at all future points in time are precisely encoded in the uncertainty in the initial condition. In contrast, SDEs make the less restrictive assumption that uncertainty accumulates over time.
2.2 Evidence lower bound
As is the case with standard latent variable models, training proceeds by specifying approximate posteriors over latent variables (in this case the latent state,  $z(t)$ ) and then minimizing the Kullback-Leibler (KL) divergence between the approximate and true posterior. An added complication in our specific case comes from the fact that the latent state is defined by a stochastic process rather than a random-vector as is more typical. While it is possible to choose more general approximate posteriors [12], in this work we make the choice to approximate the posterior over the latent state as a Markov Gaussian Process (MGP) defined by the linear SDE,
$$
d z = \left(- A _ {\phi} (t) z + b _ {\phi} (t)\right) d t + L _ {\theta} (t) d \beta , \tag {2}
$$
with Gaussian initial condition,  $z(t_0) \sim \mathcal{N}(m_0, S_0)$  where  $m_0 \in \mathbb{R}^d$  and  $S_0 \in \mathbb{R}^{d \times d}$  is a symmetric positive definite. Here  $A_\phi: \mathbb{R} \to \mathbb{R}^{d \times d}$  and  $b_\phi: \mathbb{R} \to \mathbb{R}^d$  are symmetric matrix and vector-valued functions of time, respectively. Before proceeding, it is worth emphasizing that this choice of approximate posterior is not as restrictive as it may appear since it is only used to approximate the posterior (i.e. the latent state given the observations  $p(z(t) \mid \mathcal{D})$  for  $t \in [t_1, t_N]$ ). When forecasting
beyond the data window, we will use the nonlinear SDE in (1). In this way, the approximate posterior primarily serves as a useful intermediary in training the generative model rather than the end-product.
For our purposes, MGPs are useful because the marginal distribution  $q_{\phi}(z(t)) = \mathcal{N}(m_{\phi}(t),S_{\phi}(t))$  (i.e. the distribution of the latent state at time  $t$ ) is a Gaussian whose mean and covariance are given by the solution to the following set of ODEs [19],
$$
\dot {m} _ {\phi} (t) = - A _ {\phi} (t) m _ {\phi} (t) + b _ {\phi} (t),
$$
$$
\dot {S} _ {\phi} (t) = - A _ {\phi} (t) S _ {\phi} (t) - S _ {\phi} (t) A _ {\phi} (t) ^ {T} + L _ {\theta} (t) \Sigma L _ {\theta} (t) ^ {T}, \tag {3}
$$
where  $m(0) = m_0$  and  $S(0) = S_0$ .
Using the reparametrization trick from [13], the evidence lower bound (ELBO) can be written as,
$$
\begin{array}{l} \begin{array}{c c c} & N \\ \hline & \end{array} \\ \mathcal {L} (\theta , \phi) = \sum_ {i = 1} \mathbb {E} _ {z \left(t _ {i}\right) \sim q _ {\phi} (z (t))} [ \log p _ {\theta} \left(x _ {i} \mid z \left(t _ {i}\right)\right) ] \tag {4} \\ - \frac {1}{2} \int_ {0} ^ {T} \mathbb {E} _ {z (t) \sim q _ {\phi} (z (t))} \left[ | | r _ {\theta , \phi} (z (t), t) | | _ {C _ {\theta} (t)} ^ {2} \right] d t, \\ \end{array}
$$
$$
r _ {\theta , \phi} (z (t), t) = B (t) \left(m _ {\phi} (t) - z (t)\right) + \dot {m} _ {\phi} (t) - f _ {\theta} (z (t), t), \tag {5}
$$
$C_{\theta}(t) = (L_{\theta}(t)\Sigma L_{\theta}(t)^{T})^{-1}$ ,  $B(t) = \mathrm{vec}^{-1}\left((S_{\phi}(t)\oplus S_{\phi}(t))^{-1}\mathrm{vec}(C_{\theta}(t)^{-1} - \dot{S}_{\phi}(t))\right)$ ,  $\oplus$  indicates the Kronecker sum,  $\mathrm{vec}:\mathbb{R}^{d\times d}\to \mathbb{R}^{d^2}$  maps a matrix into a vector by stacking columns, and  $\mathrm{vec}^{-1}:\mathbb{R}^{d^2}\to \mathbb{R}^{d\times d}$  converts a vector into a matrix such that  $\mathrm{vec}^{-1}(\mathrm{vec}(B)) = B\forall B\in \mathbb{R}^{d\times d}$ . See Appendices A and B for a detailed derivation. We note that computing the residual,  $r_{\theta ,\phi}$ , scales linearly in the dimension of the state so long as  $S_{\phi}(t)$  and  $C_{\theta}(t)$  are diagonal. To ensure our approach remains scalable, we will make this assumption throughout the remainder of this work. In this case, we have
$$
B (t) = \frac {1}{2} S _ {\phi} (t) ^ {- 1} \left(C _ {\theta} (t) ^ {- 1} - \dot {S} _ {\phi} (t)\right). \tag {6}
$$
Often we may wish to place additional priors onto a subset of the generative model parameters,  $\theta$ , and infer their posterior using stochastic variational inference as well. In this case we add the KL-divergence between the approximate posterior and prior onto the ELBO,  $KL(q_{\phi}(\theta) || p(\theta))$  and rewrite all expectations with respect to  $q_{\phi}(z(t))$  and  $q_{\phi}(\theta)$ ; details are provided in Appendix C.
Remark 1: Despite having defined the approximate posterior in terms of an SDE with parameters  $A_{\phi}$  and  $b_{\phi}$ , the ELBO only depends on the mean and covariance of the process at a particular point in time,  $m_{\phi}$  and  $S_{\phi}$ . For this reason we can parametrize  $m_{\phi}$  and  $S_{\phi}$  directly while implicitly optimizing with respect to  $A_{\phi}$  and  $b_{\phi}$ . In addition, we can efficiently compute  $\dot{m}_{\phi}$  and  $\dot{S}_{\phi}$  using automatic differentiation.
Remark 2: Despite the fact that the prior and approximate posterior are SDEs, all expectations in the ELBO are taken with respect to normal distributions. Moreover, in contrast to the approach in [20, 21] there are no differential equality constraints – instead we have been left with an integral over the window of observations.
Taken together, these observations allow us to infer the parameters of the generative model (a nonlinear, latent SDE with additive diffusion (1)), without the use of a forward solver.
2.3 Amortization strategy
The implicit assumption in the ELBO in (4) is that the state mean and covariance will be approximated over the entire window of observations. This can pose a serious computational bottleneck with long or complicated time-series. In this section we propose a novel amortization strategy that will allow us to effectively eliminate this cost by requiring that we only approximate the posterior over short partitions of total the data time-window at once.
Rather than attempting to compute and store the posterior over the entire latent trajectory, we will instead construct an approximation to the posterior over a small window of observations as a function of those observations. Consider a reindexing of the dataset by splitting it into  $N / M$  non-overlapping partitions where  $1 \leq M << N$ ,
original indexing:  $[t_1, t_2, \dots, t_M, t_{M+1}, \dots, t_N]$
reindexed dataset:  $[t_1^{(1)}, t_2^{(1)}, \ldots, t_M^{(1)}, t_1^{(2)}, \ldots, t_M^{(N / M)}]$
In the case that  $N$  is not evenly divisible by  $M$  we allow the final split to contain less elements. We approximate the latent state over each partition using only the  $M$  observations in each partition,  $q_{\phi}(z(t) \mid x_1^{(j)}, x_2^{(j)}, \ldots, x_M^{(j)}) \approx p(z(t) \mid \mathcal{D})$  for  $t \in [t_1^{(j)}, t_1^{(j+1)}]$ . This can be interpreted as a probabilistic encoder over the time interval of the partition of observations. Letting  $t_1^{N/M+1} \equiv t_M^{N/M}$ , the ELBO can be compactly rewritten as,  $\mathcal{L}(\theta, \phi) = \sum_{j=1}^{N/M} \mathcal{L}^{(j)}(\theta, \phi)$ , where
$$
\begin{array}{l} \mathcal {L} ^ {(j)} (\theta , \phi) = \sum_ {i = 1} ^ {M} \mathbb {E} _ {q _ {\phi} (z \left(t _ {i} ^ {(j)}\right) \mid x _ {1} ^ {(j)}, \dots , x _ {M} ^ {(j)})} \left[ \log p _ {\theta} \left(x _ {i} ^ {(j)} \mid z \left(t _ {i} ^ {(j)}\right)\right) \right] \tag {7} \\ - \frac {1}{2} \int_ {t _ {1} ^ {(j)}} ^ {t _ {1} ^ {(j + 1)}} \mathbb {E} _ {q _ {\phi} (z (t) | x _ {1} ^ {(j)}, \dots , x _ {M} ^ {(j)})} \left[ | | r _ {\theta , \phi} (z, t) | | _ {C _ {\theta} (t)} ^ {2} \right] d t. \\ \end{array}
$$
An additional advantage of this amortization strategy is that it allows our approach to scale to multiple trajectories without an increase to the overall computational cost. If there are multiple trajectories, we can reindex each trajectory independently and subsequently sum all sub loss functions.
To reiterate, the probabilistic encoder is a function which takes in  $M$  observations from a particular partition along with a time stamp,  $t$ , and outputs a mean vector and covariance matrix as an estimate for the latent state at that particular time. In principle, any function which can transform a batch of snapshots and a time stamp into a mean and covariance could be used as an encoder in our work. In our implementation, we use deep neural networks to encode each  $x_{i}^{(j)}$  using  $i \in \mathcal{I}$  where  $\mathcal{I} \subset [x_1^{(j)}, x_2^{(j)}, \ldots, x_M^{(j)}]$  contains some temporal neighbours of  $x_{i}$  into a single latent vector. This approach yields a set of latent vectors associated with each observation in the partition  $h_{i}$  for  $i = 1, 2, \ldots, M$ . We then interpolate between each latent vector using a deep kernel based architecture to construct the posterior approximation for any time stamp in the partition; see Appendix D for details. We emphasize this is one choice of encoding architecture that we found convenient, it is straightforward to incorporate an autoregressive or transformer based encoder in our methodology [22]
An important consideration is the selection of the partition parameter,  $M$ . In practice,  $M$  should be large enough so that the assumption of a linear SDE for the approximate posterior is appropriate (i.e. we have enough observations in a partition so that the assumption of a Gaussian process over the latent state is reasonable). For example, as we will see in an upcoming numerical study, in the context of inferring latent SDEs from videos, we will need to choose  $M$  to be large enough so that we can reasonably infer both the position and velocity of the object in the video.
2.4 Reparametrization trick
While the previous sections have demonstrated how to eliminate the need for a differential equation solver by replacing the initial value problem with an integral, in this section we show how the reparametrization trick can be combined with the previously described amortization strategy to construct unbiased gradient approximations for the ELBO with a time and memory cost that scales independently with the amount of data, the length of the time series, and the stiffness of the approximation to the differential equations. Consider a reparametrization to the latent state of the form  $z(t) = T(t,\epsilon ,\phi)$  where  $\epsilon \sim p(\epsilon)$  so that  $z(t)\sim q_{\phi}(z(t)\mid x_1^{(j)},x_2^{(j)},\ldots ,x_M^{(j)})$ . We can rewrite the second term in the evidence lower bound as,
$$
\begin{array}{l} \int_ {t _ {1} ^ {(j)}} ^ {t _ {1} ^ {(j + 1)}} \mathbb {E} _ {q _ {\phi} (z (t))} \left[ | | r _ {\theta , \phi} (z (t), t) | | _ {C _ {\theta} (t)} ^ {2} \right] d t = \int_ {t _ {1} ^ {(j)}} ^ {t _ {1} ^ {(j + 1)}} \mathbb {E} _ {p (\epsilon)} \left[ | | r _ {\theta , \phi} (T (t, \epsilon , \phi), t) | | _ {C _ {\theta} (t)} ^ {2} \right] d t \\ = \left(t _ {1} ^ {(j + 1)} - t _ {1} ^ {(j)}\right) \mathbb {E} _ {p (\epsilon) p (t)} \left[ \left| \left| r _ {\theta , \phi} (T (t, \epsilon , \phi), t) \right| \right| _ {C _ {\theta} (t)} ^ {2} \right] \tag {8} \\ \end{array}
$$
where  $p(t)$  is a uniform distribution,  $\mathcal{U}(t_1^{(j)},t_1^{(j + 1)})$  and  $p(\epsilon)\sim \mathcal{N}(0,I)$  is a Gaussian. With this rearrangement, we can derive the main result of this work.
Lemma 1. An unbiased approximation of the gradient of the evidence lower bound, denoted as  $\nabla_{\theta ,\phi}\mathcal{L}(\theta ,\phi)$ , with an  $\mathcal{O}(R)$  time and memory cost can be formulated as follows:
$$
\begin{array}{l} \nabla_ {\theta , \phi} \mathcal {L} (\theta , \phi) \approx \frac {N}{R} \sum_ {i = 1} ^ {M} \sum_ {k = 1} ^ {R} \nabla_ {\theta , \phi} \log p _ {\theta} \left(x _ {i} ^ {(j)} \mid T \left(t _ {i} ^ {(j)}, \epsilon^ {(k)}, \phi\right)\right) \tag {9} \\ - \left(t _ {1} ^ {(j + 1)} - t _ {1} ^ {(j)}\right) \frac {N}{2 R} \sum_ {k = 1} ^ {R} \nabla_ {\theta , \phi} | | r _ {\theta , \phi} (T (t ^ {(k)}, \epsilon^ {(k)}, \phi), t ^ {(k)}) | | _ {C _ {\theta} (t ^ {(k)})} ^ {2}. \\ \end{array}
$$
where each  $t^{(k)}\sim \mathcal{U}(t_1^{(j)},t_1^{(j + 1)})$  and each  $\epsilon^{(k)}\sim \mathcal{N}(0,I)$
The proof follows by applying the standard reparametrization trick [15] to estimating gradients of the amortized objective in (7).
Remark 1: In practice we found choosing  $R \sim 100$  worked well for the problems we considered. Note that in terms of elapsed time, 100 evaluations of this objective, which can be computed in parallel, is far cheaper than 100 evaluations of the SDE forward model evaluated as a part of an iterative sequential SDE solver. Moreover we found that adaptive stepping schemes required far more evaluations of the SDE forward model than our stochastic approach (see Section 4.1).
Remark 2: In the case that evaluations of the SDE drift term were relatively cheap compared to decoder evaluations (for example in the case the dimension of the latent state is much smaller than the dimension of the data), we found it useful to increase the number of samples used to approximate the integral over time without increasing the number of samples from the variational posterior. To do so, we made use of a nested Monte Carlo scheme to approximate the second term in the ELBO,
$$
\begin{array}{l} \left(t _ {1} ^ {(j + 1)} - t _ {1} ^ {(j)}\right) \mathbb {E} _ {p (\epsilon) p (t)} \left[ \left| \left| r _ {\theta , \phi} (T (t, \epsilon , \phi), t) \right| \right| _ {C _ {\theta} (t)} ^ {2} \right] \approx \\ \frac {t _ {1} ^ {(j + 1)} - t _ {1} ^ {(j)}}{R S} \sum_ {k = 1} ^ {R} \sum_ {l = 1} ^ {S} \left\| r _ {\theta , \phi} \left(T \left(t ^ {(k, l)}, \epsilon^ {(k)}, \phi\right), t ^ {(k, l)}\right) \right\| _ {C _ {\theta} \left(t ^ {(k, l)}\right)} ^ {2}, \tag {10} \\ \end{array}
$$