Latent SDEs on Homogeneous Spaces
Authors: Sebastian Zeng
Florian Graf
Roland Kwitt
Venue: NeurIPS 2023 poster (2023)
Keywords: Variational Bayesian inference
stochastic differential equation
homogeneous spaces
geometric Euler-Maruyama
time series

Abstract: We consider the problem of variational Bayesian inference in a latent variable model where a (possibly complex) observed stochastic process is governed by the unobserved solution of a latent stochastic differential equation (SDE). Motivated by the challenges that arise when trying to learn a latent SDE in $\mathbb{R}^n$ from large-scale data, such as efficient gradient computation, we take a step back and study a specific subclass instead. In our case, the SDE evolves inside a homogeneous latent space and is induced by stochastic dynamics of the corresponding (matrix) Lie group. In the context of learning problems, SDEs on the $n$-dimensional unit sphere are arguably the most relevant incarnation of this setup. For variational inference, the sphere not only facilitates using a uniform prior on the initial state of the SDE, but we also obtain a particularly simple and intuitive expression for the KL divergence between the approximate posterior and prior process in the evidence lower bound. We provide empirical evidence that a latent SDE of the proposed type can be learned efficiently by means of an existing one-step geometric Euler-Maruyama scheme. Despite restricting ourselves to a less diverse class of SDEs, we achieve competitive or even state-of-the-art performance on a collection of time series interpolation and classification benchmarks.

Latent SDEs on Homogeneous Spaces
Sebastian Zeng, Florian Graf, Roland Kwitt University of Salzburg, Austria
{sebastian.zeng aooffloooo florian. graf, roland kwitt]/plus..... ac.at
Abstract
We consider the problem of variational Bayesian inference in a latent variable model where a (possibly complex) observed stochastic process is governed by the solution of latent stochastic differential equation (SDE). Motivated by the challenges that arise when trying to learn an (almost arbitrary) latent neural SDE from data, such as efficient gradient computation, specific subclass instead. In our case, the SDE evolves on a homogeneous latent space and is induced by stochastic dynamics of the corresponding (matrix) Lie group. In learning problems, SDEs on the unit r--phere are arguably the most relevant incarnation of this setup. Notably, for variational inference, the sphere not only facilitates using truly uninformative prior, but we also obtain particularly simple and intuitive expression for the Kullback-Leiller divergence between the approximate posterior and prior process in the evidence lower bound. Experiments demonstrate that a latent SDE of the proposed type can be learned efficiently by means of an existing one-step geometric Euler-Maruyama scheme. Despite restricting ourselves to a less rich class of SDEs, we achieve competitive or even state-of-the-art results on various time series interpolationvllasffcciiio problems.
take step back and study a
Introduction
I Learning models for the dynamic behavior of a system from collection of recorded observation sequences over time of great interest in science and engineering. Applications can be found, c.g., in finance [5], physics [46] population dynamics [2]. While the traditional approach of adjusting prescribed differential equation (with typically few parameters) to an observed system has a long history, the advent of neural ordinary differential equations (ODEs) [10] has opened the possibility for learning (almost) arbitrary dynamics by parameterizing the vector field of an ODE by a neural network. This has not only led to multiple works, e.g., [16; 27; 37; 60], elucidating various aspects of this class of models, but also to extensions towards stochastic differential equations (SDEs) [26; 30; 34; 58; 59], where construction, account for potential stochasticity in a system and offer increased flexibility, but also come the price of numerical instabilities, high computationallstoragg costs, technical difficulties when backpropagating gradients through the typically more involved numerical integration schemes.
parametrizes the drift and diffusion vector fields. Such neural SDE can, by
While several recent developments to address these chal- lenges exist, such as [30] or [26], is important to note that neural ODEs SDEs
most often used in set- tings where they serve as a model for some sort of (unob- served) latent dynamics. Such latent ODE latent SDE models, cf. [10; 30; 48], hinge on the premise that the available observations continuous-tin ttme latent state (that might be subject to un- certainty). If one postulates that the observed dynamics are indeed governed by simpler latent dynamics that are casier model, the question naturally arises whether
determined by an underlying
Figure 1: Sample paths from prior (left) and posterior (right) process the sphere. Initial values are marked by
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

smaller subclass of SDEs, which can be solved efficiently and with less technical difficulty, may suffice for accurately modeling real-world phenomena. Specifically promising are SDEs that arise from the action of a Lie group on a homogeneous space, as they come with (1) an intuitive geometric interpretation, (2) rich theory []; 12; 22], and (3) easy-to-impleemen numerical solvers with strong convergence guarantees [35; 36; 40; 47].
Contribution. In this work, we take a step back from trying to learn (in theory) arbitrary SDEs and instead consider the very subclass mentioned above, i.e., SDEs that evolve homogeneous latent space as the result of the Lie group action. Numerical solutions to such SDEs are computed with a one-step geometric Euler-Maruyama scheme. The latter, combined with the assumption of a low latent space dimension enables backpropagating gradients through the SDE solver during learning via the besscretize-then-optiiiin instantiate our construction by learning latent SDEs on the unit n-sphere where the formula for the Kullback-Leite Leibler (KL) divergence (within the evidence lower bound) between approximate posterior processes and an uninformative prior becomes surprisingly simple. Experiments on a variety of datasets provide empirical evidence that the subclass of considered SDEs is sufficiently expressive to achieve competitive or state-of-the-an performance on various tasks.
strategy. In the variational Bayesian learning regime, we
2 Related work
Our work is mainly related to advances in the neural ODE/SDE literature and recent developments in the area of numerical methods for solving SDEs on Lie groups.
Neural ODEs and SDEs. Since the introduction of neural ODEs in [10], many works have proposed extensions to the paradigm of parameterizing the vector fields of an ODE by neural networks, ranging from more expressive models [16] to higher-order variants [60]. As ODEs are determined by their initial condition, [27; 39] have also introduced a variant that can adjust trajectories from subsequent observations. Neural SDEs [19; 24; 34; 45; 54; 58; 59]. meaning that an SDE's drift and diffusion coefficient are parametrized by neural networks, represent another natural extension of the neural ODE paradigm and can account for uncertainty in the distribution over paths. From a learning perspective, arguably the most common way of fitting neural ODEs/SDES to data the Bayesian learning paradigm, where learning is understood posterior inference [29], although neural SDEs might as well be trained as GANs [26]. In the Bayesian setting, one may choose to model uncertainty (1) only over the initial latent state, leading to latent ODEs [10; 48], or (2) additionally over all chronologically subsequent latent states in time, leading to latent SDES [19; 30]; in the first case, one only needs to select a suitable approximate posterior over the initial latent state whereas, in the second case, the approximate posterior is defined over paths (e.g.. in carlier work [3] chosen as a Gaussian process). Overall, with recent progress along the line of (memory) efficient computation of gradients to reversible SDE solvers [25; 30J, neural SDEs have become a powerful approach to model real-world phenomena under uncertainty.
SDEs on matrix Lie groups. In the setting of this work, SDEs are constructed via the action of a matrix Lie group the corresponding homogeneous space. In particular, an SDE on the Lie group will translate into an SDE in the homogeneous space. Hence, numerical integration schemes that retain the Lie group structure are particularly relevant. Somewhat surprisingly, despite vast amount of literature on numerical integration schemes for ODEs that evolve in Lie groups and which retain the Lie group structure under discretization, e-g., [21: 41; 42], similar schemes for SDEs have only appeared recently. One incarnation, which we will use later on, is the geometric Euler-Maruyama scheme from [35; 36] for Itô SDEs. For the latter, [47] established convergence in strong sense of order 1. Stochastic
schemes [40] of higher order were introduced just recently, however, at the cost of being numerically more involved and perhaps less appealing from learning perspective. Finally, we remark that [44] introduces neural SDEs on Riemannian manifolds. However, due to the general nature of their construction, it is uncertain whether this approach can be transferred to a Bayesian learning setting and whether the intricate optimization problem of minimizing a discrepancy between measures scales larger, high-dimensional datasets.
2

---

3 Methodology
3.1 Preliminaries
In the problem setting of this work, we assume access to a repeated (and partially observed) record of the continuous state history of random dynamical system in observable space over time. The collected information is available in the form of a dataset of N multivariate time series x ........ x In search for a reasonable model that describes the data- aata-generating system, we consider a continuous stochastic process X Ω x [0. T] - IR where each data sample refers to a continuous path x² [0, T + Rd. This path is assumed to be part of a collection of i.i.d. realizations sampled from the path distribution of the process X. As common in practice, we do not have access to the entire paths x', but only to observations x (t) E IR at a finite number of time points = ty that can vary between the x' and from which seek to learn X.
Fitting such a process to data can be framed as posterior inference in a Bayesian learning setting, where the data generating process is understood as a latent variable model of the form shown in the figure below: first, latent path z is drawn from a parametric prior distribution po- (z) with true parameters 0" second, observation path x" is drawn from the conditional distribution pe (x|2")
Essentially, this is the problem setting of [29]; however, the latent variables as well as the observations x² are path-valued and the true posterior po x) intractable in general. To efficiently perform variational inference (e.g..
0
with respect to z) and to fit the model to observations, we choose tractable approximate posterior qs (z|x) from some family of parametric distributions, suitable prior po (z), and seek to maximize the evidence lower bound (ELBO)
log (x 2 C(O, Φ: x') = - DKL (ga(z)x")ll pl + (x)e)(x)xe) [log pe (x' [z)
(1)
This optimization problem well understood in the context of vector-valued random variables, and can be efficiently implemented (provided the choice of prior and approximate posterior allows reparameterization trick). However, the setting of a path-valued random latent variable comes with additional challenges.
Parametric families of path distributions. To efficiently perform approximate variational infer- ence, we need to select prior and approximate posterior in reasonable and tractable manner from parametric families of distributions. To this end, we consider path distributions [1 ]] of stochastic processes that are implicitly defined as strong solutions to linear Itô SDEs of the form
dZ,- fZZ,, t) dt + oZ,, t) dW, E [0, Zo 7
(2)
W is vector-valued Wiener process, f and a denote vector-valued drift and matrix-valuee diffusion coefficients, respectively, and Zy multivariate random variable that follows some distribution P. To ensure the existence of a unique. strong solution to Eq. (2), we require that f. a and Zo satisfy the assumptions of [57, Theorem 5.5). As in the seminal work on neural ODEs [10], one may parameterize the drift and diffusion coefficients via neural networks, yielding neural SDEs. Given these prerequisites, maximizing the ELBO in Eq. (1) requires computing gradients with respect and Φ, which hinges on sample paths from the approximate posterior process, generated by some (possibly adaptive step-size) SDE solver. Unfortunately, as mentioned in Section I, modeling (2) with flexible drift and diffusion coefficients comes at high computational and memory cost. In the following (Section 3.2), we introduce our construction for SDEs on homogeneous spaces, discuss a simple SDE solver, and elaborate on the specific form of the KL divergence term in the ELBO of Eq. (1). In Section 3.3, we then instantiate the SDE construction on the example of the rotation group (O(n) acting the unit n-sphere.
Solving auxiliary tasks. While learning SDEs of the form outlined above is an unsupervised problem, the overall approach can easily be extended to solve an auxiliary supervised task. For instance, we may have access to one target variable y² (e.g.. a class-label) per time series or one target variable per time point and time series, i.e., Vi ....,... ...,.... (e.g., a class-label or some real-valued quantity per time point). In these scenarios, we will additionally learn to map latent states z to target variables, and extend the BLBO from Eq. (1) with an additive task-specific loss term. In practice, the overall optimization objective is then a weighted sum of (i) the KI. divergence, (ii) the likeliood, and (iii task-specific supervised loss.
3

---

3.2 SDEs on homogeneous spaces
Drawing inspiration from manifold learning, e.g., [7: 14; 17], we assume the unobservable latent stochastic process Z evolves on a low dimensional manifold M. C RR. If M. admits a transitive action by a Lie group 9. i.e., if it is a homogeneous G-space, then variational Bayesian inference can be performed over family of processes that are induced by the group action. Specifically, we consider one-point motions, cf. [31], i.c., processes in M of the form Z, = G. Zo for some stochastic process G' in matrix Lie group G C (L( that acts by matrix-vector multiplication.
We define this stochastic process G as solution to the linear SDE
dG, = Vo((t) di + Σ d G2,
Go In
(3)
Herein, w² are independent scalar Wiener processes, V....., Vn tangent vectors the Lie algebra g. and I, is the n X n identity matrix. The drift function Vo K + K; [J7 -> RRxx consists of tangential (g-valued) component K plus an orthogonal component KP called pinning drift [36; 53]. It offsets the stochastic drift away from the Lie group g due to the quadratic variation of the Wiener processes w" and only depends the noise components V Vi...., Vm (and the geometry of g). In case the Lie group G is quadratic, i.c., defined by some fixed matrix P via the equation G PG P. then the pinning drift K can be deduced directly from the condition d(G PG) = 0 which implies that
are fixed (matrix-valued)
m
(o(t) = (t) +
M V?
and V PV. =0
(4)
2
We include the main arguments of the derivation of Eq. (4) in Appendix E.2 and refer to [6] and [12] for broader details.
Finally, given an initial value Zo ~ P drawn from some distribution P on the manifold M. the SDE for G induces an SDE for the one-point motion Z G Zo via dZt dG,, Zu. Therefore, Z solves
m Vo(t) + Σ d
d7,-
V.
Zc Zo P.
Numerical solutions. To sample paths from a process Z defined by an SDE as in Eq. (5), numer- ically solve the SDE by a one-step geometric Maler-Maruyama scheme [36] (listed in Appendix E). Importantly, this numerical integration scheme ensures that the approximate numerical solutions reside in the Lie group. has convergence of strong order I [47], and can be implemented in a few lines of code (see Appendix D).
The idea of the scheme is to iteratively solve the SDE in intervals [4, (yl of length At, using the relation
Z, G.Zo (G,GT, )Z1,
(6)
Notably, G follows the same SDE G., see Eq. (3), but with initial value I, instead of at 6 O. Moreover, if AL sufficiently small, then G,G,,' stays close the identity and reparametrizing G,C,, exp(12, by the matrix (or Lie group) exponential is feasible. The resulting SDE in the Lie algebra is linear and can thus be solved efficiently with Euler-Maruyama scheme. Specifically, 24 Vo(ty) At + Em Auj Vi, where Awú ~ N (O, At) and we can differentiate through the numeric solver by leveraging the well-known reparametrization trick for the Gaussian distribution.
b
Kullback-Leibler (KL) divergence. For approximate inference in the variational Bayesian learning setting of Section 3.1, we seek to maximize the ELBO in Eq. (1). In our specific setting, the prior (z) and approximate posterior distribution yo (z) x/|x of latent paths z are determined by an SDE of the form as in Eq. (5). In case these SDEs have the same diffusion term, i.e., the same V., i> 0, then the KI. divergence DKL. (1o (z|x)] po (z)) in the EI ELBO is finite and can be computed via Girsanov's theorem.

---

Specifically, if Virior and vona denote the respective drift functions and Vo Varior
Vo then
1 T /' / 92, (z) (t) Σ
DKL (g0(z)x)l
(t) dz dt
where 92, is the marginal distribution the approximate posterior process at time and Σ is the pseudo inverse (z) mm Vizz VV. For a derivation of this result, we refer [43] and Appendix E.5.
3.3 SDEs on the unit r2-sphere
The unit n-sphere Sa-1 frequently occurs many learning settings, e.g., see [13; 15; 32; 33; 38; 56]. It is a homogeneous space in context of the group SO(n) of rotations, For modeling time series data, Sn-1 is appealing due to intrinsic connection trigonometric functions and, in a Bayesian learning setting, it is appealing since compactness allows selecting the uniform distribution Ugn-1 as an uninformative prior on the initial SDE state Zo-
The unit n-sphere also perfectly fits into the above framework, as SO(n) is the identity component of O(n), i.e. the quadratic Lie group defined by P In. Its Lie algebra sn {A A A 0 consists of skew-symmetrie n x matrices, and has basis {Ex: I k Egi e e.eg where eg is the k-th standard basis vector of R". Importantly, if
< 1 S nj with
a # (V.)tn 1f/2-1//2 a ERI hickeisn and Vo(t) = (aa))=aaa))-000
then Eq. (4) fulfilled and the SDE in Eq. (5) becomes, after some transformations,
dz = (a))²"
7z dt + a (I, - -,,7 dd Zo ~ po( Ug=-1
(8)
2
whose solution is rotation symmetric and driftless' stochastic process the unit n-sphere, i.e., the perfect uninformative prior process in our setting.
In fact, in case of a = 1, Eq. (8) is known the Stroock equation [23; 55] for spherical Brownian motion.
Similarly, we define an approximate posterior process as solution to an SDE with learnable drift and learnable (parametric) distribution Pi on the initial state Zo. i.e.,
dz (x) (a n2 1- I, Z. dt (I, - -ZZ7 ) dW,,
Zo Zo~pi (Zol))
(9)
Exemplary sample paths from solutions of Eq. (8) and Eq. (9) with constant drift are shown in
Figure 1. We want to highlight that, in contrast to the parametric function K*(x): [0, T] + so
(as part of the drift coefficient) and the parametric distribution Pdd the diffusion term does not depend on the evidence but only on Φb This allows selecting the same diffusion components の a in the posterior and prior process to ensure a finite KL divergence.
The latter has particularly simple form (cf. Appendix E.5)
['] qz, (2)] dt. g
DKL (z)(x) po(z)) DKL. (Zo
2(00)2
(10)
As customary in practice, may replace the second summand with Monte Carlo estimate and sum over all available time steps. Alternatively, using the upper bound
K (x)()2l22 dz <|| (x)(1l
qz,
is equally possible. Notably, the second summand in the KL divergence of Eq. (10) has a quite intuitive interpretation: the squared norm is akin weight decay penalty, only that it penalizes large (in norm) drift or, in other words, a group action that represents large rotation.
'The drift only consists of the pinning component that is orthogonal to the Lie algebra so
5

---

3.4 Parameterizing the approximate posterior
As indicated in Eq.( enters the SDE associated with the approximate posterior qø (z x) in two ways: first, through the parametric distribution on initial states Zo. and second, via the drift component (x)). For Po. we choose the power spherical [8] distribution S with location parameter 1 and concentration parameter n² (x), i.e., Po (Zo (x) S(u (x), (x)). We realize yr k and K6 by task specific recognition network capable of learning representation for time series x' (with possibly missing time points and missing observations per time point).
In our implementation, we primarily use multi-time attention networks (mTAND) [50] to map a sequence time-indexed observations to the parameters of the approximate posterior SDE. Vector- valued sequences are fed directly to mTTAND, whereas image-valued sequences are first passed through a CNN feature extractor. Both variants are illustrated as types (A, O) and (B, O), resp., in Figure 2 and are used for interpolation tasks (i.c., imputation of missing observations at desired times) and in (per-time-poim classification problems. In situations where we need to extrapolate from just one single initial observation, we use type (C, O) from Figure 2, which directly maps this image to representation. In all cases, the so obtained representations are linearly mapped to the parameters (1t x)) of the power spherical distribution and skew-symmetric matrix K representing the drift at time 0 (and evolving in time as described below).
Time evolution of the drift. To ensure that K$ (x)(() is continuous and can be evaluated for every & E [0. 7). we use Chebyshev polynomials' this choice suitable for our experiments Section 4, but in principle any differentiable function [0, T] + so()) can be used. In our implementation, a neural network maps a time series to K skew-symmetric matrices K? (x) E so and subsequently K (t) is computed by
Nia
Time (6)
suggested by [37]. We found
imnge--ulued
only xa avnilable
vector-nllaed
ITAND
CNN
CNN
Type (A)
K (x)
Σ (x)xx(
(11)
Type (BI
Type (C)
Linear map
where pj denotes the i-th Chebyshev polynomial.
desired (u² ,n" {*)
4 Empirical evaluation
SDE Solver
We evaluate our latent SDE approach on different tasks from four benchmark datasets, ranging from time series Figure 2: Recognition network types. interpolation to per-time-pur
classification and regression tasks. Full dataset descriptions are listed in Appendix A. While we stay close as possible to the experimental setup of prior work for fair comparison, we re-evaluated (using available reference implementations) selected results in one consistent regime, specifically in cases where could not match the original results. In addition to the published performance measures, we mark the results of these re-runs in blue. For each method, we train ten different models and report the mean and standard deviation of the chosen performance measure. We note that if one were to report the same measures but over samples from the posterior, the standard deviations would usually be order of magnitude smaller and can be controlled by the weighting of the KL divergence in the ELBO.
Training hyperparrametes. Depending on the particular (task-specific) objective, we optimize all model parameters using Adam [28] with a cyclic cosine learning rate schedule [18] (990 epochs with cycle length of 60 epochs and learning rate within [1e-6, Ie-3]). The weighting of the KL divergences in the ELBO is selected on validation splits (if available) or set to le-5 without any annealing schedule. We fix the dimensionality of the latent space to 16 and use K polynomials in Eq. (1 1), unless stated otherwise.
Comparison to prior work. We compare against several prior works that advocate for continuous- time models of latent variables. This includes the latent ODE approach of [10; 48] (ODE-ODE).
Chebysheev polynomials are recursively defined via (!) pt) pt) with (o 1. p.)

---

the continuous recurrent units (CRU, f-CRU) work from [49], as well as the multi-time attention (mTAND-Enc, mTAND-Full) approach of [50]. On rotated MNIST, we additionally compare against the second-order latent ODE model from [60] and on the pendulum regression task, we also list the recent S5 approach from [52]. As [48] conceptually closest to our work but differs in the implementation of the recognition network (i.e., using an ODE that maps time series to the initial state of the latent ODE), we also assess the performance of a mTANN-ODE variant; i.., we substitute the ODE encoder from [48] by the same mTAND encoder we use throughout. Finally, we remark that we only compare to approaches trainable on our hardware (see Appendix B) within 24 hours.
4.1 Per-time-point classification & regression
Human Activity. This motion capture dataset consists of time series collected from five individuals performing different activities. The time series are 12-dimensional, which corresponds to 3D coordinates captured by four motion sensors. Following the pre-processing regime of [48], we have 6,554 sequences available with 211 unique time points overall. The task is to assign each time point one of seven (motion) categories. The dataset is split into 4,194 sequences for training, 1,311 I for testing, and 1,049 for validation. We use a recognition model of type (B, O) from Figure 2 and linearly map latent states to class predictions. For training, we minimize the cross- cross-entropy loss evaluated on class predictions per available time point. Results are listed in Table 1.
Our primary focus in this experiment is to assess how a per-time-point classification loss affects our latent SDE model. As the linear classifier operates directly the latent paths, the posterior SDEs need to account for possible label switches over time, e.g., change of the motion class from standing to walking. Latent paths of sequences with only constant class label across all time points must remain in a single decision region to minimize their error, while paths of sequences with label switches must cross decision boundaries and consequently increase their distance from the starting point; they literally must drift away.
Table 1: Human Activity.
Accuracy [%]
ODEE-DE [48]
87.0 1 2.8
''TTADDEE [50]
90.7 + 0.2
"mTANDDFI Full [50]
91.1 + 0.2
nTAND-ODE
90.4 1 0.2
Ours
90.6 + 0.4
indicates results from [50].
We visually assess this intuition after repeating the experiment with decreased latent space dimension of n = 3 (i.e., 8 which reduces recognition accuracy only marginally to 90.2%.
A collection of three latent paths shown in Figure 3. Each path is drawn from posterior distribution computed from a different input time series at test time: the coloring visu- alizes the predicted class. Clearly visible, the latent paths corresponding to input samples with constant labels remain within the vicinity of their initial state (marked and depending on the label). This suggests that the SDE is pri- marily driven by its diffusion term. Conversely, the latent path corresponding to an input sample with a label switch clearly shows the influence of the drift component. Our conjecture can be supported quantitatively by analyzing the path-wise KL divergences (which we measure relative driftless prior). Figure 3 shows the distributions of the latter over input samples with and without label switches. We see that latent paths corresponding to input samples with label switches exhibit higher KL divergences, i.c., they deviate more from the driftless prior.
14
0.08
0.06
0.04
0.02
switch (mcan==H4_3) w/o switch (mean=12.6)
0.00
50
100
150
200
divergence
Figure 3: Latent space paths with and without label switches (on 3?) and dis- tribution of KL divergences.
Pendulum regression. In this regression task [49], are given (algorithmically generated [4] time series of pendulum images at 50 irregularly spaced time points with the objective to predict the pendulum position, parameterized by the sine/cosine of the pendulum angle. The time series exhibit noise in the pixels and in the pendulum position in form of a deviation from the underlying physical model. We parameterize the recognition network by type (A, O) from Figure 2 using the same CNN as in [49] and map latent states to (sine, cosine) tuples by the same decoder as in the latter work. 2,000 images are used for training. 1,000 images for testing and additional 1, 1,000 images for validation and hyperparameter selection.
hhttps://doi.c org/ 24432//5768x
7

---

Table 2 lists the MSE computed on the (sine, cosine) encoding of the pendulum angle, averaged over all time points in the testing sequences. Notably, the recurrent architectures in Table 2, i.e., (f-)CRU and $5, can, by construction, condition on the data at ob- served time points and thus better "pin down" the temporal dynamic of the pendulum motion. Latent ODE/SDE models, on the other hand, need to inte- grate forward from one initial state and decode the angle from potentially erroneous latent states. In light of this observation, it is encouraging that the proposed latent SDE model performs only slightly worse than models based upon recurrent units (S5 and (f-)CRRU). Nevertheless, Table 2 also reveals that performance on this task largely hinges the network variant that encodes the initial latent state: in particular, although ODE-ODE [48] fails, switching the initial state encoding (to mTAND) yields substantially better results. Notably, our experiments" with mTAND-Enc [50] (which fails in the runs reported in [49]),
Table 2: Pendulum regression.
MSE (x10-9)
TODE-ODE [48]
15.40 +: 2.85
'mTAND- Enc [50]
65.64 + 4.05
*CRU [49]
4.63 = 1.07
'f-CRU [49]
6,16 1: 0.88
1S5 [52]
3,41 +: 0.27
mTAAN--DD
4.65 + 1.21
mTTAN--Ene
3.20 0.60
CRU
3.74 - 0.27
[-CRU
3.25 1: 0.26
Ours
3.84 - 0.35
1. indicate results from [49] and [52]. resp.
yields results on a par with the state-of-the-ard.
4.2 Interpolation
PhysioNet (2012). We consider the interpolation task [50] with subsampled time points. The dataset [51] contains 8,000 time series of 41 patient variables, collected over a time span of 48 hours. Each variable is observed at irregularly spaced time points. In the considered experimental setup of [50], interpolation understood observing fraction (p) of the available time points per sequence and interpolating observations at the remaining fraction (1 - p) of time points. Importantly, this is more challenging than the setting in [49] [48], as the experimental setup in the latter works corresponds to p 1. Furthermore, the pre-processing (see Appendix A for details) of [50] renders results for p
directly comparable the ones presented in [49]. As customary in the literature [20; 49; 50], we use 80% of all time series for training, 20% for testing, scale all variables to a range of [0, 1). and map the full time span of 48 hours the time interval [0, 1] at a quantization level of either [20; 48; 50] or 6 minutes [49]. As in the experiments of Section 4,1, the recognition model is parameterized by type (B, O) from Figure and latent states are decoded to observations as in [50].
Table 3 lists the MSE obtained on varying fractions (p) of time points available to the recognition model. In the situation of subsampling, i.e., p = 1, the f-CRU and mTANN--ull models perform substantially better than ours. However, in that case, the task effectively boils down to reconstructing the observed input sequences. In contrast, subsampling, i.c., p 1, truly allows assessing the interpolation capability on previously unseen data. In this situation, we achieve new state-of-the art as our method improves upon its underlying mTAND module.
Pendulum interpolation. We evaluate on the same pendulum interpolation task as in [49]. This task differs from the position pre- diction task of Section 4.1 in that we observe only around half of randomly chosen, but uncorrupted, time points per sequence. In- terpolation is understood as predicting the pendulum images at the unseen time points. As remarked earlier, this is notably more chal- lenging for latent ODE/SDE type approaches than for models that rely on recurrent units, cannot pin down the latent path the observed measurements other than through the loss. In terms of recognition network ar-
Table 4: Pendulum interpolation.
MSE (x100*
TTTAND [50]
15.40 0.07
II
TODE-ODE [48]
15.06 + 0.14
*f-CRU [49]
1.39 + 0.07
II
CRU [49]
1.00 +
mTAND-ODE
0.10
11
F-CRU
2.61 0.19
CRU
2.06 - 0.05
11
Ours
8.15 0.06
W
f indicates results from [49].
I1
11
11
11
'We set the mTAND encoder such that its GRU returns an output for each point in time. These latent states then mapped angle prediction by the same decoder as in [49].
8

---

Table 3: Results the PhysioNet (2012) interpolation task from [50] for two quantization levels (Top: 6 min; Bottom: I min). Percentages indicate the fraction of observed time points. We report the MSE (x10 ") on the observations at unseen time points, averaged across all testing sequences. Results for m mTAND-ODE at min quantization (bottom part) are not listed, as training time exceeded our 24 hour threshold.
Observed % +
50%
60%
70%
80%
90%
100%
CRU [49]
5.11 + 0.40
4.81 + 0.07
5.60
+ 0.70
5,54
+ 0.61
5,85 + 0.82
1.77 + 0,92
f-CRU [49]
5.24 1 0.49
4.96 1 0.36
4.85 0.10
0.71 5.66
6.11 1 0,82
1.10 1 0.95
mTTAND-Full [50]
3.61 + 0.08
3.53 0.06
3.48 1 0.08
3
3.58 1: 0.09
0.49 :: 0.03
mTANN-ODE
3.38 + 0.03
3.33 0.02
3.32 0.03
3.33
3.35 :: 0.02
1.82 1: 0.04
Ours
3.14 0.03
3.06 = 0.05
3.01 + 0.05
2.94 + 0.07
3.00 + 0.08
1.55 = 0.01
CRU [49]
5.20 + 0.44
5.25 + 0.55
5.22 + 0,60
5.23 + 0.59
5.59 + 0.78
1.82 + 0.97
[-CRU [49]
5.53 + 0.54
5.22 1 0.58
5,41 1 0.65
5.25 1 0.60
5,80 1 0.76
0.84 1 0.16
mmTAND-Full [50]
4.19 - 0.03
4.02 + 0.05
4.16 + 0.05
4.41 + 0.15
4.80 + 0.04
0.47 + 0.01
Ours
3.14 1 0.04 3.03 1 0.03
3.01 + 0.02
3.04
1 0.09
3.03 1 0.14
1.54 1 0.02
indicates results from [50] (which only lists the I-min setting).
chitecture, we use type (A, 0) from Figure and decode latent states via the convolutional decoder of [49]. Hyperparameters are selected based on the validation set.
As can be seen from Table 4, which lists the MSE across all images in the testing sequences. approaches based on recurrent units perform exceptionally well on this task, while the latent ODE approach, as well as mTAND-Full perform mediocre at best. Our approach yields decent results, but visual inspection of the reconstructed sequences (see visualization next to Table 4) shows slight phase differences, which underscores our remark on the difficulty of learning this fine- ine-grained dynamic by integrating forward from = 0 and guiding the learner only through the loss at the available time points.
Rotating MNIST. This dataset [9] consists of sequences of 28 x 28 grayscale images of gradually clockwise rotated ver- sions of the handwritten digit "3". One complete rotation takes place over 16 evenly spaced time points. This gives a total sequence length of 16, where we interpret each partial rotation as a time step of 1/16 in the time interval [o, 1]. During train- ing, four images in each sequence are dropped randomly, and one rotation angle (i.e., the testing time point) is consistently left out. Upon receiving the first image time = O, the task is to interpolate the image at the left-out time point = 3/16.
Table S: Rotating MNIST.
MSE (x10-9)
GPPVAE-dis GPPVAE-joinl VOE
30.9 0.02
28.8 1 0.05
19.4 + 0.06
VODE"VAE-KL CNN-ODE
18.8 + 0.31
14.5 + 0.73
Ours
11.5 + 0.38
indicates results from [60].
We follow the evaluation protocol of [60], where 360 images are used for training, 360 for testing, and 36 for validation and hyperparameter selection. The recognition model is of type (C, O) from Figure 2 and uses the same convolutional encoder (CNN) as in [60], as well as the same convolutional decoder from latent states to prediction images. Table reports the MSE at the testing time point, averaged across all testing sequences. Our model achieves the lowest MSE on this benchmark. Figure 4 additionally highlights the extrapolation quality of the model on the example of integrating the SDE learned for the interpolation task forward from t = 0 to t= 4 in steps of 1/16.
Notably, this is the only experiment where we do not use K = 6 Chebyshev polynomials for the time evolution of the drift, see Eq. (11), but only K 1, as we want to assess whether constant velocity on the sphere is a suitable model for the constant rotation velocity in the data. Figure 4 confirms this intuition. In fact, using K >> 1 polynomials yields better interpolation results but at a loss of extrapolation quality as the higher-order polynomials are less well-behaved for large t. Regarding the baselines, we do not compare against (F-)CRU and mAND as they are not directly applicable here. As only the image at the initial time point is used as model input, learning any reasonable attention mechanism (mTAND) or conditioning on observations spread out over the observation time interval (CRU) is impossible.
Runtime measurements. For a detailed runtime analysis across different approaches that model an underlying (stochastic) differential equation, refer the reader to Appendix C.

---

E 10
3
3
3
3
3 3
3
3
3
E 21
3
3
3
c ()
c [3,
3
3
3
Figure 4: Exemplary reconstructions on the Rotating MNIST data. Shown are the results (on one testing
sequence) by integrating forward from & (marked red) to 4, i.e., three times longer than what observed [0, 1)) during training.
Discussion & Limitations
We have presented approach learning (neural) latent SDEs homogeneous spaces with Bayesian learning regime. Specifically, we focused on one incarnation of this construction, i.e., the unit n-sphere, due its abundance in many learning problems. Our SDEs are, by design, less expressive than in earlier works. Yet, when modeling a latent variable, this limitation appears to be of minor importance (as shown in our experiments). We do, however, gain terms of overall model complexity, as backpropagating through the SDE solver can be done efficiently (in a "discretize -hen- optimize" manner) without having resort to adjoint sensitivity methods variable step-size solvers. While our approach could benefit from the latter, we leave this for future research. Overall, the class of neural SDEs considered in this work presents a viable alternative existing neural SDE approaches that seek to learn (almost) arbitrary SDEs by fully parameterizing the vector fields of an SDEs drift and diffusion coefficient by neural networks. In our case, we use neural networks as high-capacity function approximators, but in setting where one prescribes the space which the SDE evolves. Finally, whether is reasonable to constrain the class of SDEs depends on the application type. In situations where one seeks to learn a model in observation space directly, approaches that allow more model flexibility may preferable.
Acknowledgments
This work was supported by the Austrian Science Fund (FWF) under project P31799-N38 and the Land Salzburg within the EXDIGIT project W120/-WWS//223/666022 and projects 0102-F1901 166- KZP, 20204-
Source code is available at
References
[1] D. Applebaum. Probability Compact Lie Groups. 1st edition. Vol. 70. Probability Theory and Stochastic Modelling. Springer, 2014.
[2] M. Arató. "A famous nonlinear stochastic equation (Lotka-Volterra Volterra model with diffusion)". In: Mathematical and Computer Modelling 38 (2003), pp. 709-726.
[3] Archambeau, M. Opper, Y. Shen, D. Cornford, and Shawe-Taylor, "Variational Inference
for Diffusson Processes". In: NIPS. 2007.
[4] P. Becker, H. Pandya, G. Gebhardt, C. Zhao, J. Taylor, and G. Neumann. "Recurrent Kalman Networks: Factorized Inference High-Dimensional Deep Feature Spaces". In: ICML. 2019.
[5] F. Black and M. Scholes. "The Pricing of Options and Corporate Liabilities". In: Journal of Political Economy 81.3 (1973), pp. 637-654
[6] R. W. Brockett. "Lic Algebras and Lie Groups in Control Theory". In: Geometric Methods in System Theory. Springer, 1973, pp. 43-82.
[7] B. C. Brown, A. L. Caterini, B. Ross, C. Cresswell, and G. Loaiza-Ganem. "Verifying the Union of Manifold Hypothesis for Image Data". In: ICLR. 2023.
10

---

[8] N. D. Cao and W. Aziz. "The Power Spherical distribution". In: arXiv abs/2006.0443) (2020).
[9] F. P. Casale, A. V. Dalca, L. Saglietti, J. Listgarten, and N. Fusi. "Gaussian Process Prior Variational Autoencoders". In: NeurIPS. 2018.
[10] R. T. Chen, Y. Rubanova, J. Bettencourt, and D. Duvenaud. "Neural Ordinary Differential
Equations". In: NeurIPS. 2018.
[11] Chetrite and H. Touchette. "Nonequilibrium Markov Processes Conditioned on Large Deviations". In: Annales Henri Pvincaré 16 (2015), pp. 2005-2057.
[12] G. S. Chirikjian. "Stochastic models, information theory, and Lie groups". : In: Applied and
numerical harmonic analysis. Boston, Basel, and Berlin: Birkhäuser, 2012, pp- 361-388.
[13] T. S. Cohen, M. Geiger, J. Köhler, and M. Welling. "Spherical CNNs". In: /CLR. 2018.
[14] M. Connor, G. Canal, and C. Rozell. "Variational Autoencoder with Learned Latent Structure".
In: AISTATS, 2021.
[15] T. R. Davidson, L. Falorsi, N. DeCao, T. Kipf, and J. M. Tomczak. "Hyperspherical Variational Autoencoders". In: UAI. 2018.
[16] E. Dupont, A. Doucet, and W. Teh. "Augmented Neural ODE". In: NeurIPS. 2019.
[17] C. Fefferrman, S. Mitter, and H. Narayanan. "Testing the manifold hypothesis". In: Journal of the American Mathematical Society 29.4 (2016), pp. 983-1049.
[18] H. Fu, C. Li, X. Liu, J. Gao, A. Celikyilmaz, and L. Carin. "Cyclical Annealing Schedule: A Simple Approach to Mitigating KL Vanishing". In: NAACL. 2019.
[19] A. Hasan, J. M. Pereira, S. Farsiu, and V. Tarokh. "Identifying Latent Stochastic Differential
Equations". In: IEEE Transactions on Signal Processing 70 (2022), pp. 89-104.
[20] M. Horn, M. Moor, C. Bock, B. Rieck, and K. Borgwardt. "Set Functions for Time Series". In:
ICML. 2020.
[21] A. Iserles, H. Z. Munthe-Kaas, S. P. Nørsett, and A. Zanna. "Lie-group methods". In: Acta Numerica 9 (2000), 215-365.
[22] K. Itó. "Brownian motions in Lie group". In: Proceedings of the Japan Academy, Series A, Mathematical Sciences 26.8 (1950).
[23] K. Itô. "Stochastic Calculus". In: Lect. Notes Phys. 39 (1975). Ed. by H. Araki, pp. 218-223.
[24] J. Jia and A. R. Benson. "Neural Jump Stochastic Differential Equations". In: NeurIPS. 2019.
[25] P. Kidger, J. Foster, X. Li, and T. Lyons. "Efficient an Accurate Gradients for Neural SDEs".
In: NeurIPS. 2021.
[26] P. Kidger, J. Foster, X. Li, H. Oberhauser, and T. Lyons. "Neural SDEs as Infnite-Diimesiional
GANs". In: ICMI.. 2021.
[27] P. Kidger, J. Morrill, J. Foster, and T. Lyons. "Neural Controlled Differential equations for irregular time series". In: NeurlPS. 2020.
[28] D. P. Kingma and J. Ba. "Adam: A Method for Stochastic Optimization". In: ICLR. 2015.
[29]
D. P. Kingma and M. Welling. "Auto-Encoding Variational Bayes". In: ICLR. 2014.
[30] Li, T.-K. L. Wong, T. Chen, and D. Duvenaud. "Scalable Gradients for Stochastic
Differential Equations". In: AISTATS. 2020.
[31] M. Liao. Levy processes in Lie groups. Providence, R.I.: Cambridge University Press, 2004.
[32] W. Liu, R. Lin, Z. Liu, L. Xiong, B. Schölkopf, and A. Weller. "Learning with Hyperspherical
Uniformity". In: AISTATS. 2021.
[33] W. Liu, Y. Zhang. X. Li, 7. Liu, B. Dai, T. Zhao, and I. Song. "Deep Hyperspherical Learning".
In: NeurIPS. 2017.
[34] X. Liu, T. Xiao, S. Si, Q. Cao, S. Kumar, and C.-J. Hsieh. "Neural SDE: Stabilizing Neural
ODE Networks with Stochastic Noise". In: arXiv abs/ Abs/1906.02355 (2019).
[35] Marjanovic, M. Piggott, and V. Solo. "A Simple Approach to Numerical Methods for Stochastic Differential Equations in Lie Groups". In: CDC. 2015.
[36] G. Marjanovic and V. Solo. "Numerical Methods for Stochastie Differential Equations in Matrix Lie Groups Made Simple". In: IEEE Transactions on Automatic Control 63.12 (2018),
pp. 4035-4050.
[37] S. Massaroli, M. Poli, Park, A. Yamashita, and H. Asama. "Dissecting Neural ODEs". In: NeurIPS. 2020.

---

[38] P. Mettes, E. van der Pol, and C. Snoek. "Hyperspherical Prototype Networks". In: NeurIPS. 2019.
[39] J. Morrill, C. Salvi, P. Kidger, and J. Foster. "Neural rough differential equations for long time series". In: ICML. 2021.
[40] M. Muniz, M. Ehrhardt, M. Günther, and R. Winkler. "Higher strong order methods for linear
Itō SDEs on matrix Lie groups". In: BIT Numerical Mathematics 62.4 (2022), pp. 1095-1119.
[41] H. Kaat "High order Runge- Kutta methods manifolds". In: Applied Numerical Muthematics 29.1 (1999). Proceedings of the NSF/CBMS Regional Conference on Numerical Analysis of Hamiltonian Differential Equations, pp. 115-127.
[42] H. Munthe-Kast "Runge-Kutte methods Lie groups". In: BIT Numerical Mathematics 38.1
(1998), pp. 92-111.
[43] M. Opper. "Variational Inference for Stochastic Differential Equations". In: Annalen der Physik 531.3 (2019). p. 1800233.
[44] S. W. Park, Kim, K. Lee, and J. Kwon. "Riemannian Neural SDE: Learning Stochastic Representations on Manifolds". In: NeurlPS, 2022.
[45] S. W. Park, K. Lee, and J. Kwon. "Neural Markov Controlled SDE: Stochastic Optimization for Continuous-Time Data". In: The Tenth International Conference on Learning Representations, ICLR. OpenReview.net, 2022.
[46] A. Pavliotis. Stochastic Processes and Applications. Springer, 2014.
[47] M. Piggott and V. Solo. "Geometric Euler- Malle-Maruyama Schemes for Stochastic Differential Equations in SO(n) and SE(n)". In: SIAM Journal on Numerical Analysis 54.4 (2016), pp. 2490- 2516.
[48] Y. Rubanova, R. Т. Chen, and D. Duvenaud. "Latent ODE for Pregularly-Samplee Time Series". In: NeurIPS. 2019.
[49] M. Schirmer, M. Eltayeb, S. Lessmann, and M. Rudolph. "Modeling Irregular Time Series
with Continuous Recurrent Units". In: ICML. 2022.
[50] S. N. Shukla and B. M. Marlin. "Multi-Time Attention Networks for Irregularly Sampled Time
Series". In: ICLR. 2021.
[51] I. Silva, G. Moody, D. J. Scott, L. A. Celi, and R. G. Mark. "Predicting in-hospital mortality of ICU patients: The PhysioNer/Computin in cardiology challenge 2012". In: Computing in Cardiology 39 (2012), pp. 245-248.
[52] J. T. Smith, A. Warrington, and S. W. Linderman. "Simplified State Space Layers for Sequence Modeling". In: ICLR. 2023.
[53] V. Solo. "An approach to stochastic system identification in Riemannian manifolds". In: CDC. 2014.
[54] Song. J. Sohl-Dickstein, P. Kingma, A. Kumar, S. Ermon, and B. Poole. "Score-based generative modeling through stochastic differential equations". In: ICLR. 2021. [55] D. W. Stroock. "On the growth of stochastic integrals". In: Zeitschrift für Wahnscheinlichkeit- theorie und Verwandte Gebiete 18.4 (1971), pp. 340-344.
[56] C. Tan, 7. Gao, I.. Wu, S. Li, and S. 7. Li. "Hyperspherical Consistency Regularization". In: CVPR. 2022.
[57]
B. Øksendal. Stochastic Differential Equations. Universitext. Springer, 1995.
[58] B. Tzen and M. Raginsky. "Neural Stochastic Differential Equations: Deep Latent Gaussian Models in the Diffusion Limit". In: arXiv abs/1905. 0008009933.. (2019).
[59] B. Tzen and M. Raginsky. "Theoretical guarantees for sampling and inference in generative models with latent diffusions". In: COLT. 2019.
[60] Ç. Yildiz, M. Heinonen, and H. Lahdesmäki. "ODE? Deep generative second order ODEs
with Bayesian neural networks". In: NeurlPS. 2019.
12

---

Supplementary material
In this supplementary material, report full dataset descriptions, list architecture details, present additional experiments, and report any left-out derivations.
Contents
A Datasets
13
Human Activity
Pendulum
A.3 PhysioNet (2012)
14
A.4
Rotating MNIST
15
B
Architecture details
15
B.1
Human Activity
16
B.2
Pendulum regression interpolation
PhysioNet (2012)
B.4 Rotating MNIST
B.5
Computing infrastructure
17
C Additional empirical results
17
C.I Runtime measurements
17
C.2 Stochastic Lorenz attractor
19
Uncertainty quantification
D Implementation
21
E.
Theory
23
E.1
Background stochastic differential equations (SDEs)
23
E.2
SDEs in quadratic matrix Lie groups
24
E.3 SDEs homogeneous spaces of quadratic matrix groups
25
SDEs on the unit n-sphere
26
E.5
Kullback-Leibler (KL) divergences
27
Societal impact. This work mainly revolves around novel type of neural SDE model. We envision direct negative societal impact but remark that any application of such model, e.g., in the context of highly sensitive personal data (such as ICU stays in PhysioNet (2012)), should be inspected carefully for potential biases. This is particularly relevant if missing measurements and eventually used in support of treatment decisions. In this scenario, biases (most likely due to the training data) could disproportionately affect any minority group, for instance, as imputation results might be insufficiently accurate.
interpolated
A Datasets
This section contains a description each dataset we use in our experiments. While all datasets have been used extensively in the literature, reported results are often not directly comparable due to inconsistent preprocessing steps and/or subtle differences in terms of data partitioning as well as variable selection (e.g.. for PhysioNer (2012)). Hence, we find it essential to highlight the exact data-loading and preprocessing implementations we use: for Human Activity and PhysioNet (2012) experiments, rely publicly available code from [50] (which is based upon the reference
13

---

implementation of [48]). For the Pendulum experiments, we rely on the reference implementation from [49], and for Rotating MNIST, we rely on the reference implementation of [60].
A.1 Human Activity
This motion capture dataset', used in Section 4.1, consists of time series collected from five individuals performing different activities. The time series are 12-dimensional, corresponding to 3D coordinates captured by four motion sensors. Following the preprocessing regime of [48] (who group multiple motions into one of seven activities), we have 6,554 sequences available with 211 unique time points overall and a fixed sequence length of 50 time points. The time range of each sequence is mapped to the time interval [0, 1,, and no normalization of the coordinates performed. The task is to assign each time point to one of the seven (motion/activity) categories (i.e., "lying". "lying down", "sitting", "sitting down". "standing up from lying". "standing up from sitting". "standing up from sitting on the ground"). The dataset split into 4,194 sequences for training, 1,31 11 for testing, and 1,049 for validation. Notably, the majority of time series have one constant label, with only 239 (18.2%) of the time series in the testing portion of the dataset containing actual label switches (e.g., from "lying" to "standing up from lying").
A.2 Pendulum
The pendulum dataset, used in Sections 4.1 and 4.2, contains algorithmically generated [4] (grayscale) pendulum images of size 24 24. For the regression task from [49]. we are given time series of these pendulum images at 50 irregularly spaced time points with the objective to predict the pendulum position, parameterized by the (sine, cosine) of the pendulum angle. The time series exhibit noise in the pixels and in the pendulum position (in form of deviation from the underlying physical model). For the interpolation task from [49], uncorrupted images are available, but we only observe a smaller fraction of the full time series (50% of time points randomly sampled per time series). The task is to interpolate images at the unobserved time points. Two exemplary time series (per task) from the training portion of the dataset shown in Figure 5. For both tasks, 4,000 images for training, 1,000 images for testing and 1,000 images for validation.
Figure 5: Illustration of one exemplary training time series from the (leff) regression and (right) pendulum position interpolation task, respectively, from [49]. Note that the images shown here inverted for better visualization.
A.3 PhysioNet (2012)
The PhysioNet (2012) [51] dataset has been used many times in the literature. However, clean comparison is difficult as the preprocessing steps, as well as the split and variable selection, often differ in subtle details. Hence, we focus on one evaluation setup [50] and re-evaluate prior work precisely in this setting instead of adopting the reported performance measures.
2rg/10024432//57GB8
14

---

The original challenge dataset contains time series for 12,000 ICU stays over a time span of 48 hours. 41 variables are available, out of which a small is only measured at admission time (e.g., age. gender), and up to 37 variables are sparsely available within the 48-hour window. To the best of our knowledge, almost all prior works normalize each variable to [0. (e.g., [48; 49; 50] and many more). In [50], and hence in our work as well, all 41 variables are used, and only the training portion of the full challenge dataset is considered. The latter consists of 8,000 time series, split into 80% training sequences and 20% testing sequences. Depending observations are either aggregated within min [48; 50J or within min [49] time slots. This yields 480 and 2,880, resp., possible (unique) time points. Overall, the available measurements in this dataset are very sparse, i.e., only = 2% of all potential measurements are available.
the experimental regime in prior works,
While prior work typically focuses on mortality prediction, i.e., predicting one binary outcome per time series, we are interested in interpolating observations unseen time points. To assess this interpolation capability of model, we adopt the interpolation task of [50]. In this task, interpolation understood as observing a fraction (p) of the (already few) available time points per sequence and interpolating observations at the remaining fraction (1 p) of time points. From our perspective, this considerably more challenging than the setting in [49] or [48], as the experimental setup in the two latter works corresponds to p 1.
Furthermore, the preprocessing regime in [50] slightly differs from [49] in two aspects: (1) in the way each variable is normalized and (2) in the number of variables used (41 in [50] and 37 in [49]). Regarding the normalization step: letting Z1...., .I... denote all values of one variable in the training dataset, and Tmin Tmax denote the corresponding minimum and maximum values of that variable, [50] normalizes by (x: fmmn) /max. whereas [49] normalizes by (r; - Fminn/ /Tmmx Tmin). For fair comparison, we re-ran all experiments with the preprocessing of [50].
A.4 Rotating MNIST
This dataset [9] consists of sequences of 28 x 28 grayscale images of gradually clockwise rotated versions of the handwritten digit "3" One complete rotation takes place over 16 evenly spaced time points. During training, one time point is consistently left-out (according to the reference implementation of [60], this is the 4th time point, i.e. at time t = 3/16) and four additional time points are dropped randomly per sequence. Upon receiving the first image time O, the task is to interpolate the image at the left-out (4th) time point. Exemplary training sequences are shown in Figure 6. We exactly follow the dataset splitting protocol of [60], where 360 images are used for training. 360 for testing and 36 for validation and hyperparameter selection.
3 334
3
3
3
33
3
3
9
38
G
CG
9
3
WWW
3
3
7
3
M
1
3
3
1
}
<
1
7
7
3
3
w
E
3
3
3
Figure 6: Five (counted in rows) exemplary training sequenees from the rotating MNIST datasct.
B Architecture details
The three types of recognition networks (i.e., A, B, or C) we use throughout all experiments are shown in Figure 2. In the following, we specify the configuration of the networks (from Figure 2) per experiment, discuss the choice of loss function (in addition to the KL divergence terms from the ELBO), and provide details about the architecture(s) in works we compare to.
Irrespective of the particular type of network that processes an input sequence, our parameterization of the SDE on the st-shere remains the same across all experiments. In particular, given a representation
15

---

of an input sequence, say h, we linearly map to (1) the parameters of the power spherical distribu- tion, i.e., the posterior distribution on the initial state, and (2) the coefficients of K Chebyshev polynomials, cf. Eq. (11), that control the time evolution of K Here, the superscript & denotes the dependency of K on the (parameters of the) recognition network, As K is skew-symmetrie and has, in case of (0),,, n2nn 1) free parameters, we linearly map to Kn(n 1)/2 coefficients overall. Unless noted otherwise, we consistently use K = 6 polynomials and set 71 = 16. In all experiments with ODE/SDE type models, we use ADAM [28] for optimization, with initial learning rate set to le-3 and subject to a cyclic cosine learning rate schedule (cycle length of 60 epochs) [18]. Finally, an
when replacing our SDE on the n-sphere by an ODE in R (denoted mTAND-ODE), Euler scheme as ODE solver and compute gradients via the adjoint sensitivity method (implemented in the publicly available torchdiffeq package).
B.1 Human Activity
As recognition network we use type (B, O) from Figure 2 with mTTAND encoder (mTAADD-Enc [50]) configured as in the avalibly-vvallabee reference implementation' When used as a standalone module (i.e., without any subsequent ODE/SDE), mTAND-Ene uses the outputs of its final GRU unit (which yields representations in IR 128 at all desired time points) to linearly map to class predictions. In our setting, we take the last hidden state of this GRU unit as input h E IR 1 to the linear maps mentioned above. When replacing our SDE on the n-sphere by an ODE in R (denoted as mTAND- ODE), the last GRU state linearly mapped to the location parameter of Gaussian distribution and to the entries of diagonal covariance matrix. For mTAND-ODE, the ODE is parameterized as in [48] by a 3-layer MI.P (with ELU activations) with 32 hidden units, but the dimensionality of the latent space is increased to 16 (for fair comparison). As in our approach, latent states are linearly mapped to class predictions. Notably, when comparing ODE-ODE [48] to mTAND-ODE the only conceptual difference is in the choice of recognition network (i.e., ODE in [48], mTAND encoder here).
Loss components. In addition to the KI. divergence terms in the ELBO, the cross-entropy loss is computed on class predictions at the available time points. For our approach, as well as mTAND-ODE, the selected weighting of the KL divergence terms is based on the validation set accuracy (and no KL annealing is performed).
B.2 Pendulum regression & interpolation
In the pendulum experiments of Section 4.1 and Section 4.2, we parameterize the recognition network by type (A, 0) from Figure 2. Images are first fed through convolutional neural network (CNN) with the same architecture as in [49] and the so obtained representations are then input to an mTAND encoder (mTANDD-Enc). Different to the mTANND configuration in Appendix B. the final GRU unit is set to output representations h e 3232, For the regression task, latent states are decoded via the same two-layer MLP from [49]. For interpolation, we also use the convolutional decoder from [49].
Loss components. For the regression lask, we compute the mean-squarrd-errro (MSE) on the (sine, cosine) encodings of the pendulum angle and use Gaussian log-likelihood the image reconstructions. Although the images for the regression task are corrupted (by noise), we found that adding a reconstruction loss helped against overfitting. The weighting of the KL divergence terms in the ELBO is selected on the validation set (and no KL annealing is performed), For the interpolation task, the MSE term is obviously not present in the overall optimization objective.
Remark on mTAND-Enc. While the mTAND-Ene approach from [50] fails in the experimental runs reported in [49] (see 'mTAND-Ene in Table 2). we highlight that simply using the outputs of the final GRU unit as input to the two-layer MLP that predicts the (sine, cosine) tuples worked remarkably well in our experiments.
'nttps://githuu htatp:://github.coom/mmm.. www.cooomooo//ee.. httpr/remm-aabm/1a. lab/mTAN
16

---

B.3 PhysioNet (2012)
In the PhysioNet (2012) interpolation experiment, we parameterize our recognition network by type (B, 0) from Figure using representations h E I as input to the linear layers implementing the parameters of the drift and the initial distribution. The mTAND encoder configured as suggested in the mTAND reference implementation for this experiment. We further adjusted the reference implementation of (f-)CRU from [49] to use exactly the same data preprocessing pipeline as in [50] and configure (f-)CRU as suggested by the authors.
Loss components. In addition the KI divergence terms, the ELBO includes the Gaussian log- likelihood of reconstructed observations at the available time points (using fixed standard deviation of 0.01). The weightings of the KL divergence terms fixed le-5, without any further tuning on separate validation portion of the dataset (and no KL annealing is performed).
B.4 Rotating MNIST
In the rotating MNIST experiment, the recognition network is parameterized by type (C. O) from Figure 2. In particular, we rely on the convolutional encoder (and decoder) from [60], except that only use the first image input (whereas [60] also need the two subsequent frames for their momentum encoder). Unlike the other experiments, where the Chebyshev polynomials are evaluated on the time interval [0, 1), we deliberately set this range to [o, 20] in this experiment. This is motivated by the fact that we also seek to extrapolate beyond the time range observed during training, see Figure 4. Furthermore, we underlying "simple" dynamic of a constant-speed rotation; setting > 1 only marginally changes the results listed in Table 5.
only one polynomial (K 1), which is reasonable due to the
Loss components. Gaussian log- likelihoeh is used on image reconstructions and weights for the KI divergences are selected based on the validation set (and no KI. annealing is performed).
B,5 Computing infrastructure
All experiments
executed systems running Ubuntu 22.04.2 LTS (kernel version (86_64) with 128 GB of main memory and equipped with either NVIDIA
5. 1..50.01 1.5.0-71-ge0 0-71-generic GeForce RTX 2080 Ti or GeForce RTX 3090 Ti GPUs. All experiments (v1.13 and also tested on v2.0). In this hardware/systemn configuration, all runs reported in the tables of the main manuscript terminated within 24-hour time window.
implemented in Py PyTorch
C Additional empirical results
C.1 Runtime measurements
We compare the runtime of our method to two close alternatives, where we substitute our latent (neural) SDE on the unit n-sphere, by either a latent (neural) ODE [10; 48] or a latent (neural) SDE [30] in R". In other words, we exchange the model components for the latent dynamics in the approaches we use for Rotating MNIST and PhysioNet (2012), see Section 4. This provides insights into runtime behavior
function of the integration steps, moving from 16 steps on Rotating MNIST to 480 on PhysioNet (2012) at a min quantization level and to 2,880 steps on PhysioNet (2012) at exclusively use Euler's method as ODE/SDE solver
I min quantization level. In all experiments, with fixed step size. Results are summarized in Table 6. particular, we list the elapsed real time for forwwrr++backwaar pass (batch size 50), averaged over 1,000 iterations.
htte::://gthubb.ccc///cccrrrrrrr hotpraachcccccccoooooooooeetertr Recurrent-Units
17

---

Table 6: Comparison of the elapsed real time of a +ranwwad pass batch (of size 50) across different tasks/datasets (with varying number of integration steps). For each task, the overall model sizes (in terms of the number of parameters) are comparable. Listed values are averages over 1,000 iterations = the corresponding standard deviations. We remark that the runtime of our approach when using (R", GL(n)) instead of (O(())) differs only marginally.
Rotating MNIST
PhysioNet (2012)
PhysioNet (2012) (1 min)
(6 min)
Number of integration steps
16
480
2,880
CNNmTANN -ODE (f10; 48))
0.053 + 0.004
0.926 + 0.017
3.701 + 0.033
CNN/mTAND SDE ([30])
0.112 + 0.009
2.075 1 0.042
10.273 0.173
Ours (8"=",S0(mm)
0.055 + 0.005
0.179 + 0.008
2.693 + 0.027
Overall, although our approach implements latent SDE, runtime per batch is on a par (or better) with the latent ODE variant on rotating MNIST and on PhysioNet (2012) at both quantization levels. Furthermore, we observe a substantially lower runtime of our approach compared to the (more flexible) latent SDE work of [30], which is closest in terms of modeling choice. It is, however, important to point out that the way our SDE evolves on sphere (i.e., as consequence the group action) lends itself to quite efficient implementation where large parts of the computation can be carried out as single tensor operation (see Appendix D).
0.060
ODE CNN-SDE
Training MSE
Ours
0.040
0.020
0.005
200
400
600
800
1,000
Number of epochs
0.060
ONN-ODE
CNN-SDE
Training MSE
Ours
0.040
0.020
0.005
100
200
300
400
500
600
700
800
Accumulated training time (forward+backwar pass) [s]
Figure Training time comparison on Rotating MNIST for five different training runs per model. Bottom: Training MSE vs. (accumulated) training time of each model. Top: Training MSE (dashed) and standard deviations (shaded) vs. number of training epochs.
Aside from the direct runtime comparison above, we also point out that while shorter (forward) runtime per batch ensures a faster inference time, it does not necessarily result faster overall training time as the convergence speed might differ among methods. To account for the latter, we compare the evolution of training losses Figure when training on the Rotating MNIST dataset. The upper plot shows truining MSE number of epochs and reveals that our method needs fewer updates to converge, whereas the same model but with latent dynamics replaced by latent ODE [10] or latent SDE [30] in R" converge at approximately the same rate. The bottom plot accounts for
18

---

different runtimes per batch and reveals that our method needs only around half of the training time than the latent SDE [30] variant in IR..
We want to point out, though, that runtime is, to a large extent, implementatioo--deeeddeen and there- fore, the presented results should be taken with caution. For the baselines, used publicly available implementations of ODE/SDE solvers, i.., https: /ttp:////////ttt///ee/epe con//riiichen/ooochhooo.ooo ffeeffeq for the latent ODE [10] models and https: /tt:///gttttb/tt/ com reeegggereeeearcccttrrhhhhheee e for the latent SDE [30] models.
Comparison to non-ODE/SDE models. Finally, Table 7: Relative timing results (for- wardwbackward pass) on PhysioNet (2012) for both quantization levels (re- ported in the main manuscript).
present a runtime comparison with methods that do not model an underlying (stochastic) differential equation, i.c., mTAND-Full and CRU. This relative runtime comparison was performed on the PhysioNet (2012) dataset and the run- times/batch are averaged over batches size 50; results are summarized in Table 7. The PhysioNet (2012) dataset particularly interesting concerning runtime, as it contains the most time points (480 potential time points for quantization level of 6 minutes and 2,880 time points for quantization level of minute).
min
min
CRU
22x
25
F-CRU mTAND-Ful
17x
20x
1x
1x
Ours
All runtime experiments were performed on the same hardware (see Appendix B.5 with GeForce RTX 2080 Ti GPU).
C.2 Stochastic Lorenz attractor
To illustrate that our model is expressive enough to model difficult dynamics, we replicate the experiment by [30] of fitting a stochastic Lorenz attractor. Figure shows 75 posterior sample paths.
2
2
Figure 8: Shown are 75 posterior sample paths from our model on the stochastic Lorenz attractor data of [30].
C.3 Uncertainty quantification
In the main part of this manuscript, we primarily focussed on the interpolation and regres- sion/elassification capabilities of our method on existing benchmark data. is, however, worth pointing out that our generative model can be utilized for uncertainty quantification. We can sample from the probabilistic model outlined in Section 3 and use these samples to estimate the underlying distribution. While we did not evaluate the uncertainty quantification aspect as thoroughly as the regression/llasiffcatiin performance, we present preliminary qualitative results on the Pendulum angle regression experiment in Figure 9, This figure shows density estimates for the angle predic- tions on testing data. The spread of the angle predictions is (as one would expect) sensitive to the KL-divergence weight in the loss function. The present model was trained with weight of le-3.
19

---

-1
-1
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.00
0.25
0.50
-1
-1
-1
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
0-
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
1.00
0.75
0.50
0.25
0.00
0.0
0.2
0.6
0.8
Figure 9: Top: Visualization of quality of fit and uncertainty for the first 12 testing time series in the pendulum regression task. Shown is the ground truth (in blue) and the predicted angle (sine). Shaded regions correspond to densities, estimated from quantiles of 1,000 draws from the posterior distribution. Bottom: Zoomed-in plot of the first time series.
20

---

D Implementation
To highlight that our approach can be implemented in few lines of code, Listings 1, and provide minimal working example (in PyTorch) of how to obtain solutions to an SDE on s The only part left out in this example is the generation of the time evolution of K. To execute the code snippet in Listing 4, the reference implementation of the power spherical distribution" R from [8] as well as the einops package" are needed.
Listing SDE solver in Lie algebra so (n).
from einops import repeat
class PathSampler (oddulll):
def (senf, dim)
super (PathSample,
self register_buffer buffer (aaffer("akev_syy_bba gen_ sen_basis(dim)) self.alphe mn nn. Parameter (torch. tensor (1./dim, ttyph
def forward(self, Kt, delta):
shapeele[0], Kt shape
Et repeat wwwtsteeee..ssssymmmmeese wwww_sym_ basis, 'k b -> bs nt k a b', bs-bs, nt-nt) noise torch.randn (bs, nt, Et (2) 1, 1, device device) Vt (noise=@t-deeta.ssgrr =self alpha) (2)
14
return Kt*delta + Vt
Listing 2 Generate basis in so (n).
def gen_ basis ;
in=
(n* (n-1)/2)
torch. zeros (r, n)
for nth, (i,j) in enumerate (torch. combinations (torch. arange (aange(n))) )n))):
B [[tth,i,]] 1.
return B (1,))posee1,2)
Listing 3 Generate SDE solution on S
from torch. distributions import Distribution from torch import Tensor
def sample (t: Tensor, pex: Distribution, Kt: Tenaor, sampler: PathSampler)
20 pzx. rsample() unsqueeze (-1)
delta torch diff (t) [o]
7
0t sampler (Kt, delta)
Qt torch. www.mmieexet.tt...omooooo con.iloooontiuuus(()))
zi [z0]
for t_idx in range (0, len(t)-1):
Qtz torch, matmul (Qt(: t_idxx;,;i] zi (-1)
13
zi.append (Qtz)
14
return torch. stack dim=1) squeeze()
httppp///gthhbbbcc///oooladedddd corraeerrarrreeeeecccalll huttp:://einops.roook//
21

---

Listing 4 Example usage with dummy K, n.R.
import torch
2.
import torch. nn as nn from power_spherical import PowerSpherical
16
batch_size = 4
num_time hum_time_pointe 100
device 'cuda' if torch. cuda. is_available() else 'cpu' torch. linspace 1, him_timm__ponta,, device=device) sampler PathSampler (n) (deviice)
torch. randn (batch_size, n, device-device) batch p's kappa torch rand (batch_size, 1, device=device) batch of k's pzx = PoverSpherical
15
16
torch. mn functional. (tt normalize
17
kappa.squeeze s)ueeze())
Xt just contains random skew-symmetric matrices this emample Kt torch. rachbranddnnaan (batchhbah_ssizee len(t), n, device-device) Kt Kt Kt. (2,
zis sample pzx, Kt, sampler)
23
zús is now tensor of size (bs, num, tointto n)
22

---

E
Theory
In this part, we provide further details the construction of SDEs on homogeneous spaces of quadratic matrix Lie groups g, and derive the formula for the KL divergence between solutions of such SDEs that appeared in the main manuseript.
E. Background on stochastic differential equations (SDEs)
We first make some initial remarks on stochastic processes and on stochastic differential equations (SDEs) in R" and introduce the concepts and notations that are used throughout this final chapter of the supplementary material.
Given some probability space (R, P), a stochastic process on [0. 7] with state space R is a function X 22 [) R such that for each point of time t E [0, T] the map X; = x(;
Ω - R w + X (w, t) is measurable. In other words, {X₁ : t E [0, 7] is a collection of random variables X, indexed by cf. [66, Chapter I.1; 57, Chapter 2.1; 63, Chapter 1.1]. Conversely, each sample w E Ω defines a function X(w,) [0, T] + R", L + X (w, () called a realization or sample path. Thus, a stochastic process models randomness in the space of functions [o, T7 - R", or in a "reasonable" subspace thereof. In more formal terms, the law or path distribution F. associated to stochastic process is defined as the push-forward of the measure P to this function space [66, Chapter I.3].
Arguably the most important stochastic process is the Wiener process W which formalizes the physical concept of the Brownian motion, i.e., the random motion of small particle inside a fluid caused by thermal fluctuations, see [68, Chapter 3] for heuristic approach the topic. For stochastic processes X that are sufficiently well-behaved, one can define a stochastic integral 0 X. dW., a. called Itō integral of X with respect W, cf. [66, Chapter IV.2; 57, Chapter 3.1; 63, Chapter 3.2]. Notably, the result, is again stochastic process so that one can implicitly define stochastic processes via an Itō integral equation of the form
x, Xo+ L. ()x,,)ds ds + L. (Xx ()
(12)
often denoted as an Itô stochastic differential equation (SDE)
dX, = f(X,, t) dt + oX,, dW,, Xo E
(13)
Here. / R [0, + R is called the drift coefficient, R 10, T) Raxm is called the diffusion coefficient, and the initial value & is a random variable that is independent of the Wiener process W. Further information about different notions of solutions to an SDE, i... strong and weak solutions, as well as a comprehensive consideration of their existence and uniqueness can be found in [66, Chapter 9; 57, Chapter 5; 63, Chapter 5]. Additionally, for treatment of SDEs from numerical perspective, we recommend [64].
In this work, we consider time series data as realizations of stochastic process, i.e., as samples from a function-valued random variable as described in Section 1. To perform variational inference in such a setting, utilize SDEs construct function-valued distributions. Specifically, we consider non-autonomous linear Itô SDES with multiplicative noise that evolve on homogeneous space of a (quadratic) matrix Lie group. To ensure the existence and uniqueness of strong solutions, the drift and diffusion coefficients that appear in the considered SDEs satisfy the Lipschitz and growth conditions of [57, Theorem 5.2.1]
remark on notation. As outlined in Section 3.4, we aim to learn the parameters & of an approximate posterior. Thus, the SDEs consider are not specified by an initial random variable Xu = E, but by an initial distribution Xe such that X, ~ PNo, Therefore, we utilize the notation
dX2, = [(X,, () di + o(x,, dW, Xo ~ xx
(14)
(strong) solution to such an SDE is then understood as a stochastic process that (i) is adapted (see [66, Definition 4.2l). (ii) has continuous sample paths t + X(w,)), -), (iii) almost surely satisfies the integral equation Eq. (12), and (iv) there is a random variable E ~ P- that is independent of the Wiener process W such that X, = E holds almost surely. In accordance with [57, Theorem 5.2.1], we demand that and the initial distribution & has finite second moment, but, of course, distinct choices of Xo might define distinct strong solutions X.
23

---

E.2 SDEs in quadratic matrix Lie groups
Definition (cf. [65]), Let Rux" be the vector space of real n. x matrices. We denote the general linear group, i.e. the group of invertible real nu matrices equipped with the matrix product, as
GL(n) (A
R det A + #0.
(15)
GL(n) is Lie group with Lie algebra
gl(n)= R 11.-1)
(16)
i.e., Raxx together with the commutator bracket (A, B) + [A, AB BA.
In the context of this work, we specifically consider quadratic matrix Lie groups
Definition 2 (cf. [62]). Given some fixed matrix P E <RaxK, the quadratic matrix Lie group G is defined as
G G {A e GL(n): A PA-PP
(17)
The Lie algebra of G then consists of the matrices V that satisfy V P + PV = 0.
Moreover, we consider non-autonomous linear Itô SDEs in quadratic matrix Lie groups G. The following lemma (Lemma 3) provides sufficient conditions to ensure that SDE solution evolves in G. This will pave the way to introduce a class of SDEs on homogeneous spaces in Appendix E.3 via the Lie group action.
Lemma (ef. [6, Chapter 4.1], [12]). Let G be stochastic process that solves the matrix Itô SDE
dG. = Vo(t) dt Σ V G.
Go In
(18)
where Vo [0, + Raxn Vi...., Vm E Ruxn and w ...., t are independent scalar-valued
Wiener processes. The process G stays in the quadratic matrix Lie group G(P) defined by BRxx if
(i) Vo(t) = k(t) + EE V?
with
K: K; g. and
(ii) Vi...., Em
Remark. Although the matrices V....., Vn E g can be chosen arbitrarily, it is natural to choose an orthogonal basis of the Lie algebra g. Then dim g and each direction in the vector space provides an independent contribution to the overall noise in Eq. (18).
Proof. The process G stays in the quadratic matrix Lie group G(P) if G. PG, is constant, i.c., d(G, PG.) O. The latter is to understood in the Itô sense, i.e..
0 = .d.... = dG. PG, + G,PdG,+ dG, PdG.
PV)(t)) = -G. (VT (V,(t)))vv())) dt Σ dwf (V. PV,) + Σ dt V PV, G. (19)
V.V
Here, dw, du? = dey dt, as w" and wi are independent scalar Wiener processes. Eq. (19) implies the following two conditions:
m
v. v.P+PVV PV. for i I,..., and ( = Σ v PV, (20)
24

---

As the following calculation shows, these conditions are satisfied if V1. (t) +
Vm E g and Vo(t)-
Di=1 V? where K [0. g.
m
m
K (1) (K(1)
Σ
V?
P+
P
Σ
V?
(21)
2
=0
m
V
- Σ
(V, + V PV,)
(22)
(23)
Numerical solution. Algorithm 1 lists choice of numerical solver, which is the one-step geometric Euler- Marlaa-aaruyama scheme from [36] (see also [40, Section 3.1)).
This scheme consists of three parts. First, we divide TJ into subintervals [ty, fy+1], starting with to = 0 until lj+1 T. Second, we solve Eq. (18) pathwise in G over successive intervals [t,, G+11 by applying McKean- Ganenllla Injection [12] in combination with numerical solver. Third, the (approximate) path solution is eventually computed as Gj+1 = G, exp(I2x11).
Algorithm 1 Geometric Euler-Maruyama Algorithm (g-EM) [36].
1:
ln <
2: Go I,,
3: repeat over successive subintervals [ty, ty+1l.
4:
òj 6 ts+1 t,
5:
ōwj ~ NN(,0,)). for = I,...,,
6:
2/+1 (Vo(t)) 2/2 - V?)
D ow
b Q2f1 1 g. Matrix exp. ensures that Gj+1 E G.
7:
Gj+1 + G, exp (9z+1)
8:
until =
Remark. We remark that in the original derivation of the algorithu in [36], the group acts from the right (as in Algorithm 1). However, the same derivation can be done for a left group action which is what we use in our implementation and the presentation within the main manuscript. Furthermore, under condition (i) from Lemma 3, that line reduces
m
2/11 + K(t))6g +
V
E.3 SDEs in homogencous spaces quadratic matrix Lie groups
In the following, we use our results of Appendix E.2 to define stochastic processes in homogeneous space MI of a quadratic matrix Lie group g. The main idea is, given a starting point Zo E M,, to translate path G [0, T] G into path Z G Zo in M via the Lie group action. Notably, this construction can be done the level of stochastic processes, i.e., a stochastic process G [0, 2" x Ω + G induces stochastic process Z G. Zo J0, 7'I x Ω - MI called a one-point motion, see [31, Chapter 2]. particular, if Zo follows some distribution P in M and G solves an SDE the form as in Eq. (18), then Z solves dZ; = dGG- Zo. i.c..
dz. (Vo()) d
Edwiv, Zz, Zo ~ P 1=1
(24)
Notably, Eq. (3) guarantees that a solution to the SDE in Eq. (18) remains in the group G, and consequently solution to Eq. (24) remains in M. As we only consider Lie group actions that are given by matrix-veett multiplication, we substitute the notation G by the more concise GZ from now
25

---

Remark. For subsequent results, it is convenient to represent Eq. (24) in the form of
dZ Vu(t) dt VV V2h......
Z d. Zo plo
(25)
i.e., with diffusion coefficient o: (z,t) + IVi,.... m E Raxm and m-dimensional Wiener process W = [wo w² um Notably, the matrix a (z, (2) is given by
a a(2,1)0(2,1) = Σ V Viz
(26)
Numerical solution. As Z GZo. numerical solution Eq. (24) is easily obtained from the numerical approximation of G. i.e. the numerical solution to Eq. (18) via Algorithm 1. In fact, suffices replace 1ine by Zo p and line 7 by Zj+1 exp(S)+1)5 Notably, do not need to compute the matrices Gj but only the iterates 22j.
E.4 SDEs on the unit n-sphere
Next, we consider the homogeneous space (S"-1, (0(n)) i.e. the unit n-sphere S together with the action of the group (n) of rotations.
Notably, the orthogonal group O(n) the quadratic matrix Lie group ((P) defined by P I. i.e. O(n) {A TR A A } and the special orthogonal group SO(n) is its connected component that contains the identity I,. Thus, Lemma 3 ensures that a solution to Eq. (24) remains in 1- K [o, T| + so(n) and V v....., Vm so(n).
In the variational setting of the main manuscript, we consider stochastic processes that solve such an SDE for the choice of {V.) {a(exe, eie, ) Jiskiinn where a A 0 and eş denotes the k-th standard basis vector of R". With that choice, Eq. (24) becomes
dz: K a² -- In) Ze df + ZU, Zo PZu
(27)
The diffusion term in Eq. (27) consists of a matrix-valuce noise dU, where U4s 0 and u4j Uja are independent Wiener processes (for i < j). The following lemma (Lemma 4) converts the diffusion term to the standard form of only cector-valued Wiener process. This facilitates (in Appendix E.5) deriving the -i-diiergence between two solutions of such a SDE from existing results.
Lemma 4. Let W be an n-dimensional Wiener process and let P be a distribution on IK" with finite second moment. Further, let K [0, T - so(n) be continuous. The solutions to the SDE in Eq. (27) and
dZ. - (K(()) In) 2, dt + a(In -Z,4, dW,, Zo pZo
(28)
have the same probubility distribution. Furthermore, for a given initial value Zo an ~ pZo, sample path solution of each equation also is sumple puth solution of the other one.
Proof. W.l.o.g. assume a = 1. The SDEs from Eq. (27) and Eq. (28) have the same drift coefficient but different diffusion term. Recall that Eq. (27) can be written as in Eq. (25), and denote the diffusion coefficients of Eq. (25) and Eq. (28) by a. (z, t) + IViz.... ..., Vnz and 02 (z, t) + I zzz" respectively. The lemma then follows from [61, Theorem 2. which states that in the present setting, the solution to Eq. (27) and Eq. (28) have the same distribution and equal sample paths if T 0202 i.e.. if for every z E R",
VV. V
(In -zz -) (In -zz TT
(29)
To prove this equality, first, note that I, - zz T is the projection onto the orthogonal complement of and so
(I, -zzz (In -zz )
I, -zz
26

---

To calculate the left-hand side, we insert our choice of [V.)== = {alere) - -ceo which yields
iskelgn 1
Viz V
-e,ok) zz
(erel c
km Im1
e
et (e, z)
ek e
e e)
1=1
||0||==
In ZZ
We want to point out that for vanishing drift K 0, the SDE in Eq. (28) becomes the well-known Strvock equation [23; 55] for a spherical Wiener process, i.c., a diffusion process on the sphere ga-1 that is invariant under rotations. In the variational setting of the main manuscript, we use this particular process to define an uninformative prior.
To illustrate that the condition K = 0 defines a process Z that is invariant under rotations R E SO(n). we consider the process R.Z,, cf. [23, p. 220f.]. For brevity, we introduce f +> an2(n - 1)/2 and z + alI, -zz so that Eq. (28) becomes dZ,
(dW, Notably, the identities (Z, = fRZZ) and Ro R o (RZ) hold. Thus, the process R.Z, satisfies
d(RZ.) = RdZ.
RR((Z)) di (Z. dW, R.f(Z,) di + (R( RoTZ)))RRddd RdW,
=f(RZ) + (RR, d(RW))
((RZ) di (RZ,) dW.
As W. = R.,, is again a Wiener process, the processes Z₁ and RZ, follow the same SDE. The invariance of Z w.r.t rotations is also apparent from the SDE in Eq. (27) as U2 is already invariant under rotations, i.e., each [0, T] the random variables U. R.U. and U, equal in distribution.
E.5 Kullback-Leiber (KL) divergences
First, in Lemma 5, we state the general form of the KL divergence between path distributions of stochastic processes in R".
Lemma 5. Let X and Y be stochastic processes with (path) distributions p. and P.Y respectively, and let W be a vector-valued Wiener process; all with state space R". Further, let P.x and p be
distributions on I with finite second moments. If and solve" 10 the following Itô SDEs on [0, T1.
dX,- = f(x,,t) t) aw,
Xo ~ PXx.
(30)
dY, = g(Y,, dt (Y dW,
Yo ~ p%%
(31)
then the Kl. divergence between PY and xpX is given by
DKL (F')) (|xx)
-((]t,,))))((())) DKL (pYo) pxo) + /. qr. (x.(2) o(2,t)) (2, (Fa g(z, f) dz dt
(32)
gg
Here, I and denote vector-valued drift coefficients, and a denotes the matrix-valued diffusion coefficient. Furthermore, qr, denotes the marginal density of Ye (ie., the process Y at time and Σ denotes the pseudo inverse of oo0
We further assume that g and satisfy the assumptions of [57, Theorem 5.2.1] such that unique solution exists (for given Xo, Yo ab E
27

---

For proof via Girsanov's theorem, we refer the reader to [69, Section A] and [67, Lemma 2.1]. For heuristic derivation via the limit of Euler discretization, we refer [43, Section 4].
The following corollary specifies Lemma 5 to particular class of SDEs that is compatible with the SDEs introduced in Appendix E.3.
Corollary 6. Let V Vy [0.T] + R be continuous maps and let Vi..... V m E Rax". Further, let w', ....,. w be independent scalar-valued Wiener processes and let p and P be distributions IRn with finite second moments. X and Y are stochastic processes [0, T with (path) distributions P. and TPY resp., that solve
dx,
V (t)
Σ dwiV,
Xt, Xo ~ pxx
(33)
m
dY Vi (1) Σ dw Vi
Ye Yo ~ NDko
(34)
then, the KL. divergence between PY and Tpx is given by
DKL (P 11=x) DKL (p%i + J.///So 9r
1
(o) E (z) (V) dz dt
2 RR_
(35)
where Vaa = Vo - Vy and Σ Ez(z) denotes the pseudo inverse of Ezz) > =i=1 Vizz V.
Proof. The noise term of the SDEs can be written as a (Zz, L)dW, where o R x [0, T] + R Rexn is
specified in Eq. (25). Thus, the corollary is direct consequence of Lemma 5, as o (z. Σ V,zz T V.,, see Eq. (26).
Finally, the following proposition (Proposition 7) considers exactly the SDEs discussed in the main manuscript. The process X the uninformative prior process that we select in advance and the process Y is the posterior process that we learn from the data. Thus, Eq, (38) is the KL-term that appears in the ELBO in our variational setting.
Proposition 7 (KI. divergence for one-point motions S=-1). Let K(I) [0, so(n) be continuous map and let a + O. If X and Y are stochastic processes on [0, 7 with (path) distributions P.X and TPY resp., that solve the following Itô SDEs
dX, - X, di + a dU,x,
Xo PNo
P
(36)
dY; (K(t)- I.) X, dt dU Yo pYd
(37)
No where g--1 and PYa g- are distributions supported on the unit n-sphere Sn-1 with finte second moment. Then, the KL divergence between IPY and y. is given by
DKL. (PX (P*) = DRt. (p% (1 gu-1
1 [
Tr (z) [(())| dzdi
(38)
20
Proof. By Lemma 4, replace Eq. (36) and Eq. (37)bs
dX,- -a 2 n-1 I, di + a (I, X,X,) dW,
Xo ~ Doxo
(39)
dY; = (K(t) 2 I..) + a (I, -Y,, aw, as the solutions follow the same distribution p and PY, respectively. The proposition is thus a consequence of Lemma 5 for a specific choice of ʃ.g and a, i.e..
Yo P
(40)
DKL (T'Y) I
= =DL (D% 11 ) 2
T / [ [
f gv. [K(t)2]] a²I -zzzzz (I, zz 1)7)
[K(t)z] dt
(41)
28

---

To finish the proof, we highlight two points. First, Lemma 4 ensures that the process Y evolves on the unit n-sphere $a-1, so suffices to integrate ςr, over $a-1 R instead of over the entire IR as in Eq. (41). Second, the integrand simplifies to
[Kt)z]]
(I zz (In zz )
[K(t)z] =a [(t)z]] (In -zzz [t(zzz]
(42)
a 2 |(t)/l|
(43)
where the first equality holds because (In -22 T) a projection, and the second equality holds because (t)zz ((t) e g skew-symmetrie.
References (Supplementary material)
[61] E. J. Allen, L. J. S. Allen, A. Arciniega, and P. E. Greenwood. "Construction of equivalent stochastic differential equation models". In: Stochastic Analysis and Applications 26.2 (2008),
Pp. 274-297.
[4] P. Becker, H. Pandya, G. Gebhardt, C. Zhao, J. Taylor, and G. Neumann. "Recurrent Kalman Networks: Factorized Inference High-Dimensional Deep Feature Spaces". In: ICML. 2019.
[62] A. M. Bloch, P. E. Crouch, J. E. Marsden, and A. K. Sanyal. "Optimal Control and Geodesics Quadratic Matrix Lie Groups". In: Foundations of Computational Mathematics (4 2008),
pp. 469-500.
[6] R. W. Brockett. "Lic Algebras and Lie Groups in Control Theory". In: Geometric Methods in System Theory. Springer, 1973, pp. 43-82.
[8] D. Cao and W. Aziz. "The Power Spherical distribution". In: arXiv abs/2006.044 044447 (2020).
[9] F. P. Casale, A. V. Dalca, L. Saglictti, Listgarten, and N. Fusi. "Gaussian Process Prior Variational Autoencoders". In: NeurlPS. 2018.
[10] R. T. Chen, Y. Rubanova, J. Bettencourt, and D. Duvenaud, "Neural Ordinary Differential
Equations". In: NeurlPS. 2018.
[12] G. S. Chirikjian. "Stochastic models, information theory, and Lie groups". _ In: Applied and numerical harmonic analysis. Boston, Basel, and Berlin: Birkhäuser, 2012, pp. 361-388.
[18] Fu, C. Li, X. Liu, J. Gao, A. Celikyilmaz, and L. Carin. "Cyclical Annealing Schedule: A Simple Approach to Mitigating KL Vanishing". In: NAACL. 2019.
[23] K. Itô. "Stochastic Calculus". In: Lect. Notes Phys. 39 (1975). Ed. by H. Araki, pp. 218-223.
[63] I. Karatzas and S. E. Shreve. Brownian motion and stochastic calculus. 2nd edition. Springer,
1998.
[28]
D. P. Kingma and J. Ba. "Adam: A Method for Stochastic Optimization". In: ICLR. 2015.
[64]
P. Kloeden and E. Platen. Numerical methods for stochastic differential equations. 1991.
[65] J. M. Lee. Introduction to smooth manifolds. Graduate texts in mathematics. Springer, 2003.
[30] X. Li, T.-K. L. Wong, R. T. Chen, and D. Duvenaud. "Scalable Gradients for Stochastic Differential Equations". In: AISTATS. 2020.
[31] M. Liao. Levy processes in Lie groups. Providence, R.I.: Cambridge University Press, 2004.
[36] G. Marjanovic and V. Solo. "Numerical Methods for Stochastic Differential Equations in Matrix Lie Groups Made Simple". In: IEEE Transactions on Automatic Control 63.12 12 (2018),
pp. 4035
[40] M. Muniz, M. Ehrhardt, M. Günther, and R. Winkler. "Higher strong order methods for linear Itō SDEs on matrix Lie groups". In: B/T Numerical Mathematics 62.4 (2022), pp. 1095-1119,
[43] M. Opper. "Variational Inference for Stochastic Differential Equations". In: Annalen der Physik 531.3 (2019), p. 1800233.
[66] D. Revuz and M. Yor. Continuous Martingales and Brownian Motion. 3rd edition. Springer,
1999.
[48] Y. Rubanova, T. Chen, and D. Duvenaud. "Latent ODE for Irregularly-Samphe Time Series". In: NeurIPS, 2019.
[67] D. Sanz-Alonso and A. M. Stuart, "Gaussian approximations of small noise diffusions in Kullhack-Leeble divergence". In: Communications in Mathematical Sciences 15.7 (2017), pp. 2087-2097.
29

---

[68] S. Särkkä and A. Solin. Applied stochastic differential equations. Cambridge University Press, 2019.
[49] M. Schirmer, M. Eltayeb, S. Lessmann, and M. Rudolph. "Modeling Irregular' Time Series with Continuous Recurrent Units". In: ICML. 2022.
[50] S. N. Shukla and B. M. Marlin. "Multi-Time Attention Networks for Irregularly Sampled Time Series". In: ICLR. 2021.
[51] I. Silva, G. Moody, D. J. Scott, L. A. Celi, and R. G. Mark. "Predicting in-hospital mortality of ICU patients: The PhysioNe:/Computing in cardiology challenge 2012". In: Computing in Cardiology 39 (2012). pp. 245-248.
[69] T. Sottinen and S. Särkkä. " "Application of Girsanov theorem to particle filtering of discretely observed continuous-time non-linear systems". In: Bayesian Analysis 3.3 (2008), pp. 555 -584.
[55] D. W. Stroock. "On the growth of stochastic integrals". In: Zeitschrift für Wahrscheinlichkeits theorie und Verwandte Gehiete 18.4 (1971), pp. 340-344.
[57]
B. Øksendal. Stochastic Differential Equations. Universitext. Springer, 1995.
[60]
Yildiz, M. Heinonen, and H. Lahdesmäki. "ODE? Deep generative second order ODEs with Bayesian neural networks". In: NeurIPS. 2019.
30


Reviews:
[{"review_id": "SQ7GxlhooN", "reviewer": "Reviewer_GYzQ", "summary": "The paper provides an affirmative answer to a very natural and intriguing question: Can we simplify the underlying latent model describing the dynamics of a temporal process so that it can overcome the computational and technical challenges with neural ODEs/SDEs while accurately modeling the real-world phenomenon? The authors propose SDE models that arise from the action of a Lie group on a homogeneous space instead of an arbitrary SDE. They also demonstrate at-par performance of their approach on benchmark regression, classification and interpolation problems when compared to the existing methods. \n", "questions": "\n1. Why do we need such a specific SDE form for G_t in the display (4)?\n2. Why do we not need a reparametrization trick for the proposed approach? I found no discussion on this in the paper (except a standard comment in line 111).\n\nMinor.\nYou followed [47] (in the paper)or [48](in the appendix) for preprocessing the human activity dataset. Which one did you follow?", "limitations": "The paper includes a discussion on the limitations and possible societal impact.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This work is really interesting as it opens a new direction of research that could reduce the model's degrees of freedom and yet achieve (or nearly) the sota. It may encourage machine learners to leverage recent developments in the SDE literature and further simplify learning a time series phenomenon. The paper is well written.", "weaknesses": "1. The paper claims in the introduction that their approach significantly reduced computing efforts. \n                    “ However, .......... computing gradients.”  \nHowever, there is no comparison with other approaches I could find in the numerical section. Also, a mathematical discussion on the reason behind computational gain is absent. \n\n2. I think the paper would benefit from adding more discussion on the statement below in the main paper. \n\n“Numerical solutions to such an SDE are computed with a simple one-step geometric Euler-Maruyama scheme for which the “discretize-then-optimize” strategy of backpropagating gradients during learning is not a limiting factor.\"\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Reviewer GYzQ identified two weaknesses in our submission that we will address below.\n\n> The paper claims in the introduction that their approach significantly reduced computing efforts ...  \n\nTo address the remark on significantly reduced computing effort, we reran our experiments on the rotated MNIST dataset and carefully compared the training time of our method to the latent ODE and SDE models from [Rubanova et al., 2019] and [Xi et al., 2020], respectively. We report the results in **R-Tbl. 1** & **R-Fig. 3** of the attached PDF, and refer the reviewer to our **General Response** section for a detailed discussion.\n\n>I think the paper would benefit from adding more discussion on the statement below in the main paper.  \n>“Numerical solutions to such an SDE are computed with a simple one-step geometric Euler-Maruyama scheme for which the “discretize-then-optimize” strategy of backpropagating gradients during learning is not a limiting factor.\"\n\nThe word choice in this paragraph is somewhat unfortunate. The phrasing is too general and should rather refer to the actual use case scenario. We will adjust this accordingly in the final version. As we need to backpropagate gradients through each solver step, memory complexity scales linearly with the number of (fixed) time steps multiplied by the amount of memory consumed per step (which depends on the latent space dimensionality). Our statement in the manuscript should be understood in the context of the **latent** dynamics setting, where typically, a low dimensionality of the latent space is expected to suffice for good performance on downstream tasks (as shown in our experiments). While each step of the solver is quite simple (due to the construction of our SDE), there may well be limitations in using the “discrete-then-optimize” strategy once the latent space dimensionality reaches a certain point. We will rephrase (i.e., tone down) this statement and add more discussion on that issue.\n\n>Why do we need such a specific SDE form for $G_t$ in the display (4)?\n\nWe want to clarify that we do not state that this form of an SDE for $G_t$ is necessary. We restrict our SDEs to the form as in Eq. (4) because it constitutes an easy and natural way to define a stochastic process $G_t$ that evolves in the Lie group in terms of drift and diffusion residing in the Lie algebra. Moreover, for SDEs of this form, we can rely on numerical solvers from [Marjanovic et al., 2018]. On the tasks and datasets considered in our work, SDEs of such a form appear to be sufficient to achieve competitive performance (as noted by the reviewer). Nevertheless, the formulation in Eq. (4) is quite general. The restrictions on the coefficients $V_0,V_1,\\dots,V_m$ specified in Eq. (5) only state that $V_1,\\dots,V_m$ are in the Lie algebra and that $V_0$ is determined by an additional element in the Lie algebra that needs to be adjusted by a correction term to cancel out the stochastic drift away from the Lie group.\n\n> Why do we not need a reparametrization trick for the proposed approach? I found no discussion on this in the paper (except a standard comment in line 111).\n\nWe apologize for not being clear enough on this point. The short answer is 'yes'; we do need a reparametrization trick. To be more specific, sampling from our posterior (path-) distribution on the sphere consists of two steps. First, we sample an initial value from a power-spherical distribution with location parameter $\\boldsymbol{\\mu}$ and concentration parameter $\\kappa$. Importantly, this distribution allows for a reparametrization trick which we use for training. Second, we numerically solve an SDE in the homogeneous space that starts at the sampled initial value with a geometric Euler-Maruyama scheme. This also uses a reparametrization trick, as with each update step, a random matrix with elements $\\sim \\mathcal N(0,1)$ is multiplied with $\\sqrt{\\Delta t}$ and a learnable parameter $\\sigma^{\\boldsymbol \\phi}$ to realize matrix elements $\\sim \\mathcal N(0,(\\sigma^\\boldsymbol \\phi)^2 \\Delta t)$. We will add this discussion to the appendix.\n\n> Minor. You followed [47] (in the paper) or [48] (in the appendix) for preprocessing the human activity dataset. Which one did you follow?\n\nThank you for pointing this out. We followed [47] and will update the manuscript accordingly."}, {"review_id": "TiMPWqcIz1", "reviewer": "Reviewer_tqmQ", "summary": "The authors are interested in learning neural SDE models. Instead of parameterising arbitrary latent SDEs, the authors restrict their attention to homogeneous spaces, and in particular the unit sphere, in order that they can leverage the transitive group (the Lie group) to construct an SDE in the space in terms of an SDE in the Lie group, whose logarithm is a linear SDE, which leads to convenient solutions. The major advantage of this is using a straightforward discretise-then-optimise approach without an explosion of computational cost. They perform a thorough evaluation against other neural ODE/SDE methods, and show competitive performance despite the more restrictive form of the latent SDE.\n\nPost-discussion: The major concern has been addressed. The authors have pledged to extend their evaluation of the efficiency of their method to other experiments in the paper, which will be sufficient to address that part.", "questions": "typos:\n\n\n* lines 25, 58: Paramaterizing -> parameterizing\n* line 178: extra $≤$\n* line 180: arithmetics -> arithmetic", "limitations": "Limmitations are well discussed explicitly in the paper.", "rating": 7, "confidence": 2, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1. (major) To my knowledge the approach is novel and markedly different to other latent SDE methods, but is applicable to many of the same tasks, so is highly relevant.\n2. (major) The evaluation is generally very thorough (though see weaknesses), with a wide range of related methods and different tasks evaluated.\n3. (major) The code is available, which improves reproducibility, and is based on a widely used framework (pytorch) which improves possible impact.\n4. (minor) The paper is quite clearly written, including good discussion of the empirical findings.", "weaknesses": "1. (major) The method is claimed to be efficient for learning, but no evidence is provided to support this claim -- the experiments as they are demonstrate that the method can produce a test performance which is competitive with apparently more flexible approaches, and that the relative performances may depend somewhat on the task. For the reader to judge how far this method is more efficient, a quantitative evaluation of the time for learning for different models would help.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ">(major) The method is claimed to be efficient for learning, but no evidence is provided to support this claim -- the experiments as they are demonstrate that the method can produce a test performance which is competitive with apparently more flexible approaches, and that the relative performances may depend somewhat on the task. For the reader to judge how far this method is more efficient, a quantitative evaluation of the time for learning for different models would help.\n\nWe agree that additional empirical evidence is required to demonstrate and quantify that our method is more efficient. We will add a subsection to discuss computational/runtime aspects. \n\nImportantly, during the author's response period, we already ran a careful runtime comparison on the Rotated MNIST interpolation task, with results providing evidence in support of our claim. We kindly refer the reviewer to our **General Response** section for a detailed discussion and, in particular, to **R-Tbl. 1** in the attached PDF. Overall, our method is on par in terms of runtime with a latent ODE model of comparable size (#parameters) and substantially faster than the conceptually closer (and more flexible) latent SDE approach of [Xi et al., 2020], both evaluated using an Euler(-Maruyama) scheme.\n\nWe also like to thank the reviewer for pointing out typos in the manuscript."}, {"review_id": "YckvwaOoCX", "reviewer": "Reviewer_XRoq", "summary": "This paper deals with the problem of variational inference for sequential data, i.e. time series, using latent SDEs. The idea of this class of methods is to assume the stochastic process that is observed is related to a generative SDE in a latent space, whose parameters need to be inferred from data in a Bayesian fashion. The novelty of this paper is to consider SDE priors that live on homogeneous spaces, and in particular on the sphere (acted upon transitively by matrix multiplication by $SO(n)$). Many works have taken the path of having either a spherical or sequential latent space recently, but have never considered latent SDEs on the sphere as priors for variational inference of the posterior process. This is made possible by designing an uniformative sequential prior process on the sphere (with a tractable KL divergence) and a parametric class of posterior SDEs that provably lead to solutions on the sphere, and using an associated geometric Euler Maruyama scheme for the numerical integration of the SDE. The model is then optimized using an ELBO loss that is adapted to sequential variables and the SDE at hand, using the Power Spherical law as a reparametrization-friendly prior distribution on initial condition, and furthermore fitting the drift of the latent SDE. The method is then tested on a number of datasets and machine learning tasks, for which having both a sequential and spherical prior may or may not be a natural choice. Results indicate that the proposed method is at least comparable to other SOTA approaches on those problems on the tasks adressed.", "questions": "- Are the error bars reported in the tables computed though multiple training sessions or simply several samples of the posterior, for a given trained model ?\n\n- On a related note, I am not sure how the model is adapted for tasks such as per-frame classification. I get that the purpose of the h function is to potentially map to another space, e.g. using a softmax for classification problems, and I imagine the second term of the ELBO is the one that needs to be adapted depending on the model. It could be useful to detail one example (e.g. classification) and explaining what eq. (2) becomes in that case. Could h be learned jointly with the posterior instead of being assumed to be known ?\n\n- In figure 1, I would have been interested in seeing a few trajectories of posterior samples that do not have a constant label, to see if the proposed intuition for the trajectories still holds in that case.\n\n- How does the neural net used to fit the parameters of the initial condition given by the spherical power law constrain them to be positive for the concentration parameter and on the sphere for the location parameter ?\n\n- Note that Figure 2 does not display properly on all the different pdf readers I tried. For one of them, I cannot see the boxes and arrows, but just the text, which makes the figure hard to read.", "limitations": "- One of the main limitations of the work is the lack of motivation of using SDEs/generative models instead of simply ODEs/deterministic models, if the sampling capacity of the model is never really put to use, while in many domains involving time series uncertainty quantification is a key scientific issue (see above for detailed comments on this point).\n\n- Another minor limitation is the limitation of a geometric Euler Maruyama solver for the SDE, while more accurate solvers (of Runge Kutta Munthe Kaas type) exist. Those are not used so as not to end up with a too much computational burden, but I wonder how much of a problem this would be. More generally, I also would like to have an idea of the required computational time/complexity of the proposed model, and how much it would depend on the used solver.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "Strengths :\n- The theory of ODEs on Lie Groups and the associated geometric integration schemes is nicely put to use, and is potentially applicable to any homogeneous space. It would be quite nice to try and extend all this work to hyperbolic spaces, which have been used a lot in ML recently, and happen to also be homogeneous spaces.\n- I found the paper well written and enjoyable to read.\n- The example of the sphere is compelling and leads to the specification of stochastic processes defined through SDEs on the sphere both for the prior and posterior.  The tools introduced lead to a new prior distribution for stochastic processes on the sphere, and tractable posterior distributions shown to be usable in a variational inference context.\n- An extensive experimental study is carried out on different standard datasets. The proposed framework is applicable to a number of different tasks involving time series : classification (of each time step or of the global sequence), interpolation, regression. The method proves competitive with respect to SOTA approaches for each task.\n\n", "weaknesses": "Weaknesses :\n\n- As often in works involving spherical, or more generally hyperbolic latent space, the motivation of using them is not always crystal clear for all applications. For data involving some kind of periodicity, as most of those tested here (pendulum, rotating MNIST), this makes sense, but for others it is not so obvious why one would want to use spherical latent spaces. Would this method work for e.g. time series prediction of chaotic data, e.g. the Lorenz system ? However this is a minor remark. Once the need for a spherical latent space has been established, the proposed work extends successfully variational inference for sequential data to this setting.\n\n- I find it a bit disappointing not to find an application of the proposed method to uncertainty quantification, or at least one where the capacity to sample from the posterior process is not put to use. To me, this is the main interest of working with latent SDEs instead of ODEs : one gets a sequential generative model capable of sampling, computing expectations, covariance or other statistics. This could be showcased e.g. through a simple experiment of forecasting time series, showing predictions and confidence intervals.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "First, we thank the reviewer for the positive feedback! The suggestion to extend our method to hyperbolic spaces $\\mathbb H^n$ is quite interesting and opens up a new direction of future applications that we did not think of so far, e.g., relativistic dynamics in Minkowski space-time. The hyperboloid model of $\\mathbb H^n = O(1,n)/(O(1) \\times O(n))$ enables defining a stochastic process on $\\mathbb H^n$ via an SDE in $O(1,n)$ or its identity component $SO^+(1,n)$, similar to the presented case of $\\mathbb S^n$ and $O(n)$. In fact, $O(1,n)$ is also a quadratic Lie group and nicely fits into our setting.\n\n> ... Would this method work for e.g. time series prediction of chaotic data, e.g. the Lorenz system? ... \n\nWe fully agree that the choice of latent space geometry should be informed by the data. For data that involves some kind of periodicity, a spherical latent space is a good choice, but for chaotic data, this is less clear. To investigate the reviewer's suggestion regarding the Lorenz system, we used our latent dynamic model to replicate the experimental setup of [Xi et al., 2020] for fitting a *stochastic* Lorenz attractor. As can be seen from **R-Fig. 1** in the attached PDF (showing 75 posterior samples), our method is expressive enough to model the dynamics of this system. As a side remark, a second motivation for using a spherical or, more generally, a homogeneous latent space is that it enables using Lie group solvers with stronger convergence guarantees (see [Muniz et al., 2022], [Marjanovic et al., 2018]).\n\n> I find it a bit disappointing not to find an application of the proposed method to uncertainty quantification, or at least one where the capacity to sample from the posterior process is not put to use ...\n\nWe agree with the reviewer that leveraging the capability of the SDE model for uncertainty quantification would undoubtedly be very interesting. In the current work, we primarily focused on developing the methodological foundation and limited our experiments to a thorough evaluation of interpolation and regression/classification capabilities on existing benchmark data. A similarly extensive evaluation of uncertainty quantification seemed out of scope. Nevertheless, following the reviewer's suggestion, **R-Fig. 2** in the attached PDF now includes a visualization of uncertainty in the angle predictions of the **Pendulum angle regression** experiment (for one testing instance). While this is a regression and not a forecasting problem, it highlights XRoq's point of the model's capability to assess uncertainty. Qualitatively, uncertainty is higher in regions where the angle prediction is less accurate. We will include this figure (and additional visualizations of this kind) in a final version and will more prominently point out the possibility for uncertainty assessment.     \n\n> Are the error bars reported in the tables computed through multiple training sessions or simply several samples of the posterior, for a given trained model?\n\nThe error bars in the tables (e.g., 0.5e-3 for pendulum regression) are standard deviations computed over **5** training runs (with different random seeds). The standard deviations with respect to sampling from the posterior are lower (0.05e-3). We will make this clear in a final version.\n\n> On a related note, I am not sure how the model is adapted for tasks such as per-frame classification ...    \n> Could h be learned jointly with the posterior instead of being assumed to be known?\n\nWe apologize for the confusion and refer the reviewer to our **General Response** section for full clarification. In short, $h$ can be removed from the preliminaries section, as our variational setting can be described without it. We intended to introduce the observation model in liaison with the more common approach of simply stating the generative model, but this caused more confusion than it helped. \n\nRegarding a detailed description of our setting on one example: let's take the pendulum angle regression task: in that case, the decoder part of our model (1) yields $p(\\mathbf{x}(t)|\\mathbf{z}(t))$ modeled as a Gaussian distribution with its mean representing the reconstructed observation (i.e., pendulum images) and (2) an additional neural network (as in [Schirmer et al., 2022]) maps latent states $\\mathbf{z}(t)$ to the pendulum's angle. Deviations from the desired ground truth angle are measured via an MSE loss. Hence, the overall optimization objective becomes *MSE + negative ELBO*.\n\n>In figure 1, I would have been interested in seeing a few trajectories of posterior samples that do not have a constant label, to see if the proposed intuition for the trajectories still holds in that case.\n\nFollowing your suggestion, we selected two trajectories with constant class labels and **one** trajectory with a label switch. **R-Fig. 4** (right) in the attached PDF shows that the latter has a larger drift component (as it needs to cross decision boundaries). The coloring indicates the *predicted* class label (with a 30\\% error in the label-switch case and 0\\% error in the constant-label case). Quantitatively, **R-Fig. 4** (left) shows the distribution of path KL divergences *with* and *without* label switches, highlighting that trajectories *with* label switches indeed deviate more from the driftless prior.\n\n> How does the neural net used to fit the parameters of the initial condition given by the spherical power law constrain them to be positive for the concentration parameter and on the sphere for the location parameter?\n\nTo ensure positivity of the concentration parameter ($\\kappa$) of the power-spherical distribution, we take the square. Initially, we experimented with an exponential mapping but found that to be too aggressive. The location parameter is divided by its norm. \n\n> Note that Figure 2 does not display properly on all the different pdf readers I tried ...\n\nThank you for pointing this out. We will obviously fix this issue in a final version."}, {"review_id": "BJ9SHRhFwO", "reviewer": "Reviewer_uKNv", "summary": "In this paper the authors develop machinery for performing variational inference on latent functions in models where observations are generated by a latent stochastic process.  That is, the generative model is that a path in some latent space is generated according to a prior, and then we observe (noisy) values of this path at a subset of the points.  In principle this paper deals with a general case where the latent space is a \"homogeneous space\" which is a manifold such that there is a Lie group such that for any pair of points in the manifold one can translate one of the points to the other using a single element of the Lie group.  In practice, the paper focuses on the specific case of the manifold being an n-sphere (so that the Lie group is the rotations, SO(n)), and the generative model is that the initial point is sampled uniformly from the sphere and then evolves according to a (scaled) Brownian motion on the sphere.  Once the authors develop their inference machinery, they apply their method to a number of regression, classification, interpolation, and extrapolation problems, showing that this latent SDE formulation can be useful in some settings.", "questions": "The actual generative model is quite simple, and very little of it is learned.  In particular, the prior distribution has only a single learnable parameter.  Presumably most of the modeling flexibility comes from learning the mapping from the latent space to the observable data (i.e., $p(\\mathbf{x}(t) | \\mathbf{z}(t))$).  Is there a reason for not learning a drift term for the prior?  Would one run into identifiability issues with more complex priors?\n\nHow difficult would it be to extend the implementation and framework to other homogeneous spaces?  Are there any applications where such a generalization would be obviously useful?", "limitations": "The authors have adequately addressed potential social implications.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 1, "contribution": 3, "strengths": "The proposed method is very elegant and the mathematics is nice.  Moving to SDEs on a compact latent space allows for many nice features like being able to define uninformative priors and having nice correspondences between the initial distribution and the stationary distribution of the SDE.\n\nThe simplicity and performance on the considered tasks is strong motivation for the usefulness of the construction, and having access to posteriors over entire latent trajectories allows for a number of interesting tasks like interpolation.", "weaknesses": "- The main weakness for me is the presentation of the material.  I will try to provide more specific feedback in the following points, but overall there was a large disconnect between the setup in equations (1) and (2) and then the actual applications.  Specifically, (2) is listed as the objective function, but then it's not clear what the objective function would be for the regression and classification tasks.  The classification task presumably has a different generative model than the one presented between equations (1) and (2), and it's not clear where the class label would go, and what data is used where.  What parts of the models are learnable?\n-  The introduction is a bit confusing.  For example, it is unclear while reading the introduction what is path-valued and what is not.  Line 105 says everything is path-valued but then lines 93-94 say that we only have a finite number of discrete observations of the path.\n- I found the condition on $h$ in line 103 confusing.  Is this at all a restriction?  It seems like equation (1) automatically makes $X_t | Z_t$ normally distributed with mean $\\mu_t = h(Z_t)$ and variance $\\mathbf{R}$.\n- I defer to the authors on how they want to present it, but I was confused by the presentation of the ELBO around (2) because in general the KL divergence between distributions over functions determined by SDEs will be infinite unless the diffusion terms are the same.  The authors address this thoroughly and clearly later in the manuscript, but sweep this (to me) important point under the rug in the introduction.\n- Equation (12) and Figure 2 make it clear that the inference network takes in the $\\mathbf{x}$'s and then learns the hyperparameters of the initial distribution and the SDE.  Is the inference network necessary because of the sample sizes considered here?  Wouldn't it be possible to just directly optimize the set of $\\mathbf{K}_i^\\phi$'s that specify the SDE (and the parameters of the initialization) for each set of observations?  Some comment on the distinction of what is being used for amortization (and why) vs. modeling flexibility would be useful.\n- The notation around equation (7) is very confusing to me.  It appears that the time points $t_k$ represent the time that has passed since time $0$, but then it feels like $Z_{t_j}$ should be $G(t_j)G^{-1}(t_{j-1})Z_{t_{j-1}}$.  Perhaps relatedly, it's not obvious to me what the connection is between $G(t_j)$ and the SDE in (4). In particular it feels like to solve the SDE for the interval $[t_{j-1}, t_j]$ you would need a different initial condition (the distribution over $G(t_{j-1})$ obtained by solving the SDE up to that point.  Apologies if I'm being slowing here.\n- The example of extrapolating the rotated MNIST was a bit confusing.  Doesn't the extrapolation rely exclusively on the Chebyshev polynomials behaving well after the end of the interval over which they are trained?  Relatedly, how is the end time point chosen during the VI optimization and does it matter?  The KL term seems sensitive to how long of an interval is considered.  E.g., if we consider data sampled on [0, 1], but then compute the KL on paths from [0, T] with T>>1, we would want the drift terms to eventually relax back to matching the prior.\n- Lines 371-372 are confusing: what is meant by saying that only the initial time point is observed?  The prior is driftless, so it feels like if one just observes the initial state it should be hard to get directional rotations using the latent SDE model.\n\nTypos:\n- \"on a various time series interpolation\" --> \"on various time series interpolation\"\n- \"the paradigm of Parameterizing the vector fields\" --> \"the paradigm of parameterizing the vector fields\"\n- Line 171: \"in context of\" --> \"in the context of\"\n- Line 186: \"are show in\" --> \"are shown in\"\n- Line 189: \"This allows to select\" --> \"This allows selecting\"\n- I believe the equation prior to line 192 is missing a square on the Frobenius norm term since the vector norm on the lefthand side is squared\n- Line 265: \"with label switches need\" --> \"with label switches needing\"\n- Line 358: \"Upon receiving first\" --> \"Upon receiving the first\"\n- I believe equation (35) has an erroneous $\\mathbf{z}$ on the lefthand side (the righthand side is matrix valued)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Below, we address all points and outline how we will revise our manuscript. If there are further questions, we are more than happy to answer.\n> ... (2) is listed as the objective function, but then it's not clear what the objective function would be for the regression and classification tasks ... What parts of the models are learnable?\n\nWe kindly refer to our **General Response** (♦) to clarify the optimization objective & learnable model components.\n> ... it is unclear while reading the introduction what is path-valued and what is not. Line 105 says everything is path-valued but then lines 93-94 say that we only have a finite number of discrete observations of the path.\n\nIn the mentioned paragraph, we wanted to express that our setting differs from a typical VAE setting in that the random variables are stochastic processes, i.e., path-valued. To clarify, in `l105`, \"observations\" refer to realizations $\\mathbf x = X(\\omega):[0,T]\\to \\mathbb R^d$ of the process $X$. In practice, not the entire $\\mathbf x$ is available, but only a finite tuple $(\\mathbf x(t_1), \\dots, \\mathbf x(t_m))$ of evaluations of $\\mathbf x$. We will clarify this in a final version.\n> I found the condition on $h$ in line 103 confusing. Is this at all a restriction? It seems like equation (1) automatically makes $X_t|Z_t$\n normally distributed with mean $\\mu_t = h(Z_t)$ and variance $\\boldsymbol{R}$.\n\nOverall, Eq. (1) in our \"Preliminaries\" section is unnecessary and (apparently) caused confusion; please see our **General Response** (♥). \n > I was confused by the presentation of the ELBO around (2) because ...\n\nWe agree that even at this (early) point in the manuscript, a remark would help to clarify that when distributions are determined by SDEs, the diffusion terms need to agree to avoid infinite KL divergence. We will add a remark.\n> Equation (12) and Figure 2 make it clear that the inference network takes in the $\\mathbf x$'s and then learns the hyperparameters of the initial distribution and the SDE. Is the inference network necessary because of the sample sizes considered here? ...\n\nThe inference network learns a mapping from observations to the initial state parameterization of the SDE and to the coefficients in $\\mathbf K^\\phi_i$. In principle, these parameters could be optimized directly and separately for each time series, but this would not scale well: in fact, at test time, each new sequence of observations would require optimizing over the SDE parameterization again. With an *inference network*, one simply needs one forward pass through the model.\n> The notation around equation (7) is very confusing to me ...\n\nThe reviewer is right that the notation of $G(t_j)$ around Eq. (7) is inconsistent with the SDE on $G_t$ in Eq. (4) and that $G(t_j)$ needs to be replaced with $G_{t_j} G_{t_{j-1}}^{-1}$, i.e., with $\\exp(\\Omega_{j})$ from Alg. 1 in the suppl. material. In this respect, for $t \\in [t_{j-1}, t_j]$, the SDE in Eq. (4) is to be understood as an SDE on $G_t$ with initial value $G_{t_{j-1}}$, or, equivalently, as an SDE on $G_t G_{t_{j-1}}^{-1}$ starting at the identity. We will clarify and update the notation.\n> The example of extrapolating the rotated MNIST was a bit confusing. Doesn't the extrapolation rely exclusively on the Chebyshev polynomials ...\n\nFor extrapolation, we use the *same* model that was trained on the interpolation task, which has only seen data in the time $[0,1]$ (i.e., the 1st full rotation, w/o knowledge of extrapolation times $t>1$). Accordingly, the path KL div. is only computed on $[0,1]$. At test time, predictions at future time points are generated by integrating over a longer time range. It is true that extrapolation quality primarily depends on how well the drift term behaves for $t>1$. To verify the intuition that a constant velocity on the sphere is a suitable model for the constant rotation velocity in the data, we use only the first $K=1$ Chebyshev polynomial, i.e., *a constant*. $K>1$ yields better results but at a loss of extrapolation quality as the higher-order polynomials are less well-behaved for large $t$.\n> Lines 371-372 are confusing: what is meant by saying that only the initial time point is observed? ...\n\nWe agree that the phrase \"only the initial time point is observed\" is misleading: for *loss computation*, we **do** have information about 11 out of 16 images rotated by multiples of 22.5°. However, for training & testing, only the initial upright '3' is available as model **input**. This strategy works as the rotation speed in the data is constant. For rotation speeds that vary across time series, we would likely need more input images. \n> ... Is there a reason for not learning a drift term for the prior? Would one run into identifiability issues with more complex priors? ...\n\nWe use an uninformative prior with zero drift to express our limited knowledge about the underlying latent dynamics. In case of additional information, learning a drift component for the prior might be beneficial, e.g., learning a prior with constant velocity might be a good idea on Rotated MNIST.   \nRegarding identifiability, note that we are not interested in interpreting the parameters of the latent SDE (in isolation) but only in its utility for downstream tasks. \nIdentifiability would be more relevant when a specific parametric SDE is prescribed (e.g., from a physical model) and one is interested in the values of the fitted parameters.\n> How difficult would it be to extend the implementation and framework to other homogeneous spaces? Are there any applications where such a generalization would be obviously useful?\n\nFor more modeling flexibility, SDEs on $\\mathbb R^n$, induced by matrix multiplication with $G\\in GL(n)$, can be used. We already implemented this variant, and **R-Tbl. 1** in the attached PDF lists results on Rotated MNIST.  Another application (cf. reviewer XRoq) is latent dynamics in a hyperbolic space induced by $O(1,n)$, e.g., for modeling relativistic dynamics."}]

Author rebuttal:
We like to thank **all** reviewers for their overall positive feedback, valuable comments and suggestions!

While we address all issues point by point per reviewer, we first like to clarify issues that are common to (almost) all reviews: (&#9824;) to substantiate the computational claims from the manuscript with a **runtime comparison**, (&#9829;) to resolve shortcomings in the **presentation of our preliminaries** and (&#9830;) to clarify the **optimization objective** for per time-point classification/regression experiments.

---
*We refer to the following works in our rebuttal and will denote figures and tables from the attached PDF as R-Fig. & R-Tbl. XXX, respectively.*

**[Rubanova et al., 2019]** Y. Rubanova, R.T.Q. Chen, and D. Duvenaud.
*Latent ODE for irregularly-sampled time series*. In: NeurIPS 2019.

**[Xi et al., 2020]** X. Li, T.-K. L. Wong, R.T.Q. Chen, and D. Duvenaud. *Scalable gradients for stochastic differential equations*. In: AISTATS 2020.

**[Schirmer et al., 2022]** M. Schirmer, M. Eltayeb, S. Lessmann, and M.  Rudolph. *Modeling irregular time series with continuous recurrent units*. In: ICML 2022.

**[Marjanovic et al., 2018]** G. Marjanovic and V. Solo. *Numerical methods for stochastic differential equations in matrix Lie groups made simple*. IEEE Trans. Autom. Control., 63(12):4035–4050, 2018.

**[Muniz et al., 2022]** M. Muniz, M. Ehrhardt, M. Günther, and R. Winkler. *Higher strong order methods for linear Itô SDEs on matrix lie groups. BIT Numer. Math., 62(4):1095–1119, 2022.*

**[Itô, 1975]** Kiyosi Itô. Stochastic Calculus. Lect. Notes Phys., 39:218–223, 1975.

--- 

(&#9824;) **Runtime comparison.** To back up our computational claims, we ran additional experiments to assess the runtime and efficiency of our method, summarized in the table below (and **R-Tbl. 1**). 

| | Runtime/Batch [s] | Test MSE $\left(\times 10^{-3}\right)$ |
| ---------------------------- | ----------------- | -------- |
| LODE (Rubanova et al., 2019) | 0.053 $\pm$ 0.004 | 14.9 $\pm$ 0.275 |
| LSDE (Xi et al., 2020)       | 0.112 $\pm$ 0.009 | 14.0 $\pm$ 0.543 |
| Ours - $(\mathbb{S}^{n-1},\mathrm{SO}(n))$ | 0.055 $\pm$ 0.005 | 11.2 $\pm$ 0.573 |
| | | |
| Ours - $(\mathbb{R}^{n},\mathrm{GL}(n))$ | 0.056 $\pm$ 0.008 | 12.9 $\pm$ 0.854 |

This comparison is done on _Rotating MNIST_, using the same architecture (encoder/decoder) as in Sec. 4.2 of the manuscript, but the latent dynamic models vary. We compare our approach to 1) a latent ODE (_LODE_, as in [Rubanova et al., 2019] and 2) a latent SDE (_LSDE_, as in [Xi et al., 2020]. All models have approx. 450k parameters. Runtime/Batch refers to the wall clock time per **forward+backward pass** (with batch size 50), computed from a list of all SGD update steps from 5 randomly initialized training runs. We also list the final test MSEs, averaged over these 5 runs. For a fair comparison, we always use *Euler's method* as ODE/SDE solver with a fixed step size of 1/16 (since 16 time points are available in [0,1]). 

Overall, although our approach implements a latent **SDE**, runtime per batch is on par with the latent ODE variant and substantially lower than the more flexible latent SDE approach from [Xi et al., 2020] (which is closest in terms of modeling choice). Notably, decreasing the step size (e.g., to 1e-2) for latent ODE/SDE did not noticeably improve performance (MSE) but increased runtime linearly.

To account for potentially different convergence speeds, we also present loss curves in **R-Fig. 3**. The left plot in this figure shows *training MSE *vs.* epoch* and reveals that our method needs fewer SGD updates to converge, whereas LSDE and LODE converge at approx. the same rate. The right plot accounts for different runtimes per batch and reveals that our method needs only around half of the training time as LSDE.

*Similar runtime measurements for all experiments will be added to a final version.*

(&#9829;) **Presentation of preliminaries.** We agree that our current presentation of the preliminaries in Sec. 3.1 could be misunderstood, especially with the introduction of $h$ in Eq. (1). The arguably more common approach (which we will follow) is to state the data generation process as (1) sampling a realization of a path-valued latent variable $\mathbf{z}$ from a suitable parametric prior $p_{\boldsymbol{\theta}^*}(\mathbf{z})$ and subsequently (2) sampling a realization of a path-valued observation $\mathbf{x}$ from some conditional distribution $p_{\boldsymbol{\theta}^*}(\mathbf{x}|\mathbf{z})$. This makes the use of $h$ obsolete. We apologize for the confusion and will adjust Sec. 3.1 accordingly.

(&#9830;) **Clarification of optimization objective(s).**
We first note that our data (irrespective of the task) always contains multiple series of time-indexed observations and that during training, we minimize the negative ELBO, consisting of a KL divergence term between posterior and prior latent path distributions and a log-likelihood term $\log p_{\boldsymbol{\theta}}(\mathbf{x}|\mathbf{z})$. If additional supervision is available, a corresponding loss term is added, e.g., (per time-point) cross-entropy for classification or (per time-point) MSE for regression. These losses are computed on the output of an additional neural network (typically a two-layer MLP) that maps the **latent states** $\mathbf{z}$ (*not* the reconstructions) to the response variable(s). In practice, when evaluating $\log p_{\boldsymbol{\theta}}(\mathbf{x}|\mathbf{z})$, we model the conditional distribution at each time point as a Gaussian distribution. However, this is just one of many possible choices.

The only exception to the statement above is the **Human Activity** experiment, where we follow [Rubanova et al., 2019] for fair comparison and train solely with cross-entropy and the KL term, but without $\log p_{\boldsymbol{\theta}}(\mathbf{x}|\mathbf{z})$.


Meta-review: The reviewers agree that this would make an interesting contribution, with the author responses having generally strengthened the reviewers' opinions of the paper. Please read the reviews carefully, as well as your responses, and make the necessary changes in the camera-ready version.
Decision: Accept (poster)