Title: Latent SDEs on Homogeneous Spaces
Venue: NeurIPS 2023
Accepted: yes
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

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 a 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, we take a step back and study a 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 n -sphere are arguably the most relevant incarnation of this setup. Notably, for variational inference, the sphere not only facilitates using a truly uninformative prior, but we also obtain a particularly simple and intuitive expression for the Kullback-Leibler 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 interpolation/classification problems.

1 Introduction

Learning models for the dynamic behavior of a system from a collection of recorded observation sequences over time is of great interest in science and engineering. Applications can be found, e.g., in finance [5], physics [46] or population dynamics [2]. While the traditional approach of adjusting a prescribed differential equation (with typically a 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 one parametrizes the drift and diffusion vector fields. Such neural SDEs can, by construction, account for potential stochasticity in a system and offer increased flexibility, but also come at the price of numerical instabilities, high computational/storage costs, or technical difficulties when backpropagating gradients through the typically more involved numerical integration schemes.

While several recent developments to address these challenges exist, such as [30] or [26], it is important to note that neural ODEs or SDEs are most often used in settings where they serve as a model for some sort of (unobserved) latent dynamics. Such latent ODE or latent SDE models, cf. [10; 30; 48], hinge on the premise that the available observations are determined by an underlying continuous-time latent state (that might be subject to uncertainty). If one postulates that the observed dynamics are indeed governed by simpler latent dynamics that are easier to model, the question naturally arises whether

a 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) a rich theory [1; 12; 22], and (3) easy-to-implement 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 on a 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 'discretize-then-optimize' strategy. In the variational Bayesian learning regime, we instantiate our construction by learning latent SDEs on the unit n -sphere where the formula for the Kullback-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-art performance on various tasks.

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 is the Bayesian learning paradigm, where learning is understood as 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 earlier 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; 30], 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 on 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 a 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 a strong sense of order 1. Stochastic Runge-Kutta-Munthe-Kaas schemes [40] of higher order were introduced just recently, however, at the cost of being numerically more involved and perhaps less appealing from a 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 to larger, high-dimensional datasets.

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 a 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 1 , . . . , x N . In search for a reasonable model that describes the data-generating system, we consider a continuous stochastic process X : Ω × [0 , T ] → R d where each data sample refers to a continuous path x i : [0 , T ] → R d . 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 i , but only to observations x i ( t ) ∈ R d at a finite number of time points t = t k that can vary between the x i and from which we 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, a latent path z i is drawn from a parametric prior distribution p θ ∗ ( z ) with true parameters θ ∗ ; second, an observation path x i is drawn from the conditional distribution p θ ∗ ( x | z i ) .

N

z

x

θ

ϕ

Essentially, this is the problem setting of [29]; however, the latent variables z , as well as the observations x i are path-valued and the true posterior p θ ( z | x ) is intractable in general. To efficiently perform variational inference (e.g., with respect to z ) and to fit the model to observations, we choose a tractable approximate posterior q ϕ ( z | x ) from some family of parametric distributions, a suitable prior p θ ( z ) , and seek to maximize the evidence lower bound (ELBO)

This optimization problem is well understood in the context of vector-valued random variables, and can be efficiently implemented (provided the choice of prior and approximate posterior allows a 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 inference, we need to select a prior and approximate posterior in a reasonable and tractable manner from parametric families of distributions. To this end, we consider path distributions [11] of stochastic processes that are implicitly defined as strong solutions to linear Itô SDEs of the form

W is a vector-valued Wiener process, f and σ denote vector-valued drift and matrix-valued diffusion coefficients, respectively, and Z 0 is a multivariate random variable that follows some distribution P . To ensure the existence of a unique, strong solution to Eq. (2), we require that f, σ and Z 0 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 to θ 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 1, 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 SO( n ) acting on 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 i (e.g., a class-label) per time series or to one target variable per time point and time series, i.e., y i t 1 , . . . , y i t m (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 ELBO from Eq. (1) with an additive task-specific loss term. In practice, the overall optimization objective is then a weighted sum of (i) the KL divergence, (ii) the log-likelihood, and (iii) a task-specific supervised loss.

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 ⊂ R n . If M admits a transitive action by a Lie group G , i.e., if it is a homogeneous G -space, then variational Bayesian inference can be performed over a family of processes that are induced by the group action. Specifically, we consider one-point motions , cf. [31], i.e., processes in M of the form Z t = G t · Z 0 for some stochastic process G in a matrix Lie group G ⊂ GL( n ) that acts by matrix-vector multiplication.

We define this stochastic process G as solution to the linear SDE

Herein, w i are independent scalar Wiener processes, V 1 , . . . , V m ∈ g are fixed (matrix-valued) tangent vectors in the Lie algebra g , and I n is the n × n identity matrix. The drift function V 0 = K + K ⊥ : [0 , T ] → R n × n consists of a tangential ( g -valued) component K plus an orthogonal component K ⊥ 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 i and only depends on the noise components V 1 , . . . , V m (and the geometry of G ). In case the Lie group G is quadratic, i.e., 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

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 Z 0 ∼ P drawn from some distribution P on the manifold M , the SDE for G induces an SDE for the one-point motion Z = G · Z 0 via d Z t = d G t · Z 0 . Therefore, Z solves

Numerical solutions. To sample paths from a process Z defined by an SDE as in Eq. (5), we numerically solve the SDE by a one-step geometric Euler-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 1 [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 [ t j , t j +1 ] of length ∆ t , using the relation

Notably, G t G -1 t j follows the same SDE as G t , see Eq. (3), but with initial value I n at t = t j instead of at t = 0 . Moreover, if ∆ t is sufficiently small, then G t G -1 t j stays close to the identity and reparametrizing G t G -1 t j = exp(Ω t ) 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 an Euler-Maruyama scheme. Specifically, Ω t j +1 = V 0 ( t j )∆ t + ∑ m i =1 ∆ w i j V i , where ∆ w i j ∼ N (0 , ∆ t ) and we can differentiate through the numeric solver by leveraging the well-known reparametrization trick for the Gaussian distribution.

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 p θ ( z ) and approximate posterior distribution q ϕ ( z | 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 , i > 0 , then the KL divergence D KL ( q ϕ ( z | x ) ∥ p θ ( z )) in the ELBO is finite and can be computed via Girsanov's theorem.

Specifically, if V prior 0 and V post 0 denote the respective drift functions and V ∆ 0 = V prior 0 -V post 0 , then

where q Z t is the marginal distribution of the approximate posterior process Z at time t and Σ + ( z ) is the pseudo inverse of Σ ( z ) = ∑ m i =1 V i zz ⊤ V ⊤ i . For a derivation of this result, we refer to [43] and Appendix E.5.

3.3 SDEs on the unit n -sphere

The unit n -sphere S n -1 frequently occurs in 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, S n -1 is appealing due to its intrinsic connection to trigonometric functions and, in a Bayesian learning setting, it is appealing since compactness allows selecting the uniform distribution U S n -1 as an uninformative prior on the initial SDE state Z 0 .

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 = I n . Its Lie algebra so ( n ) = { A : A + A ⊤ = 0 } consists of skew-symmetric n × n matrices, and has a basis { E kl : 1 ≤ k < l ≤ n } with E kl = e k e ⊤ l -e l e ⊤ k , where e k is the k -th standard basis vector of R n . Importantly, if

̸

then Eq. (4) is fulfilled and the SDE in Eq. (5) becomes, after some transformations,

whose solution is a rotation symmetric and driftless 1 stochastic process on the unit n -sphere, i.e., the perfect uninformative prior process in our setting.

In fact, in case of α θ = 1 , Eq. (8) is known as the Stroock equation [23; 55] for a spherical Brownian motion.

Similarly, we define an approximate posterior process as solution to an SDE with learnable drift and learnable (parametric) distribution P ϕ 0 on the initial state Z 0 , i.e.,

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 ( n ) (as part of the drift coefficient) and the parametric distribution P ϕ 0 , the diffusion term does not depend on the evidence x but only on ϕ . This allows selecting the same diffusion components α ϕ = α θ in the posterior and prior process to ensure a finite KL divergence.

The latter has a particularly simple form (cf. Appendix E.5)

As customary in practice, we may replace the second summand with a Monte Carlo estimate and sum over all available time steps. Alternatively, using the upper bound

is equally possible. Notably, the second summand in the KL divergence of Eq. (10) has a quite intuitive interpretation: the squared norm is akin to a weight decay penalty, only that it penalizes a large (in norm) drift or, in other words, a group action that represents a large rotation.

3.4 Parameterizing the approximate posterior

As indicated in Eq.(9), ϕ enters the SDE associated with the approximate posterior q ϕ ( z | x ) in two ways: first , through the parametric distribution on initial states Z 0 , and second , via the drift component K ϕ ( x ) . For P ϕ 0 , we choose the power spherical [8] distribution S with location parameter µ ϕ ( x ) and concentration parameter κ ϕ ( x ) , i.e., P ϕ 0 ( Z 0 | x ) = S ( µ ϕ ( x ) , κ ϕ ( x )) . We realize µ ϕ , κ ϕ and K ϕ by a task specific recognition network capable of learning a representation for time series x i (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 of time-indexed observations to the parameters of the approximate posterior SDE. Vectorvalued sequences are fed directly to mTAND, whereas image-valued sequences are first passed through a CNN feature extractor. Both variants are illustrated as types (A, ) and (B, ), resp., in Figure 2 and are used for interpolation tasks (i.e., imputation of missing observations at desired times) and in (per-time-point) classification problems. In situations where we need to extrapolate from just one single initial observation, we use type (C, ) from Figure 2, which directly maps this image to a representation. In all cases, the so obtained representations are linearly mapped to the parameters ( µ ϕ , κ ϕ ) of the power spherical distribution and to a 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 )( t ) is continuous and can be evaluated for every t ∈ [0 , T ] , we use Chebyshev polynomials 2 as suggested by [37]. We found this choice suitable for our experiments in Section 4, but in principle any differentiable function [0 , T ] → so ( n ) can be used. In our implementation, a neural network maps a time series x to K skew-symmetric matrices K ϕ i ( x ) ∈ so ( n ) and subsequently K ϕ ( x )( t ) is computed by

where p i denotes the i -th Chebyshev polynomial.

4 Empirical evaluation

We evaluate our latent SDE approach on different tasks from four benchmark datasets, ranging from time series

h

t

1

, . . . ,

h

t

n

{

µ

ϕ

, κ

ϕ

,

K

ϕ

}

z

image-valued

vector-valued

h

CNN

x

t

1

x

t

n

CNN

mTAND

GRU

SDE Solver

desired

t

i

Type (A)

Type (B)

Type (C)

Linear map

only

x

0

available

Time

(

t

)

interpolation to per-time-point classification and regression tasks. Full dataset descriptions are listed in Appendix A. While we stay as close as possible to the experimental setup of prior work for a fair comparison, we re-evaluated (using available reference implementations) selected results in one consistent regime, specifically in cases where we 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 an order of magnitude smaller and can be controlled by the weighting of the KL divergence in the ELBO.

Training & hyperparameters. 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, 1e-3]). The weighting of the KL divergences in the ELBO is selected on validation splits (if available) or set to 1e-5 without any annealing schedule. We fix the dimensionality of the latent space to n = 16 and use K = 6 polynomials in Eq. (11), unless stated otherwise.

Comparison to prior work. We compare against several prior works that advocate for continuoustime models of latent variables. This includes the latent ODE approach of [10; 48] (ODE-ODE),

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] is conceptually closest to our work but differs in the implementation of the recognition network (i.e., using an ODE that maps a time series to the initial state of the latent ODE), we also assess the performance of a mTAND-ODE variant; i.e., 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

HumanActivity. This motion capture dataset 3 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 to one of seven (motion) categories. The dataset is split into 4,194 sequences for training, 1,311 for testing, and 1,049 for validation. We use a recognition model of type (B, ) from Figure 2 and linearly map latent states to class predictions. For training, we minimize the 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 on the latent paths, the posterior SDEs need to account for possible label switches over time, e.g., a change of the motion class from standing to walking . Latent paths of sequences with only a 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.

Wevisually assess this intuition after repeating the experiment with a decreased latent space dimension of n = 3 (i.e., S 2 ), which reduces recognition accuracy only marginally to 90.2%.

A collection of three latent paths is shown in Figure 3. Each path is drawn from a posterior distribution computed from a different input time series at test time; the coloring visualizes 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 primarily 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 to a 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

0

50

100

150

200

KL divergence

0.00

0.02

0.04

0.06

0.08

w switch

(mean=44.3)

w/o switch (mean=12.6)

switches exhibit higher KL divergences, i.e., they deviate more from the driftless prior.

Pendulum regression. In this regression task [49], we 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, ) 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 an additional 1,000 images for validation and hyperparameter selection.

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 S5, can, by construction, condition on the data at observed time points and thus better 'pin down' the temporal dynamic of the pendulum motion. Latent ODE/SDE models, on the other hand, need to integrate 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-)CRU). Nevertheless, Table 2 also reveals that performance on this task largely hinges on the network variant that encodes the initial latent state: in particular, although

†

,

‡

indicate results from [49] and [52], resp.

ODE-ODE [48] fails, switching the initial state encoding (to mTAND) yields substantially better results. Notably, our experiments 4 with mTAND-Enc [50] (which fails in the runs reported in [49]), yields results on a par with the state-of-the-art.

4.2 Interpolation

PhysioNet (2012). We consider the interpolation task of [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 is understood as observing a 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] or [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 = 1 not directly comparable to 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 to the time interval [0 , 1] at a quantization level of either 1 [20; 48; 50] or 6 minutes [49]. As in the experiments of Section 4.1, the recognition model is parameterized by type (B, ) from Figure 2 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 no subsampling , i.e., p = 1 , the f-CRU and mTAND-Full 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.e., p < 1 , truly allows assessing the interpolation capability on previously unseen data. In this situation, we achieve a 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 prediction task of Section 4.1 in that we observe only around half of 50 randomly chosen, but uncorrupted, time points per sequence. Interpolation is understood as predicting the pendulum images at the unseen time points. As remarked earlier, this is notably more challenging for latent ODE/SDE type approaches than for models that rely on recurrent units, as we cannot pin down the latent path at the observed measurements other than through the loss. In terms of recognition network ar-

MSE

(

×

10

-

3

)

†

mTAND-Full [50]

15.40

±

0.07

†

ODE-ODE [48]

15.06

±

0.14

†

f-CRU [49]

1.39

±

0.07

†

CRU [49]

1.00

±

0.07

8.23

±

0.10

f-CRU

2.61

±

0.19

CRU

2.06

±

0.05

Ours

8.15

±

0.06

†

indicates results from [49].

† indicates results from [50] (which only lists the 1-min setting).

chitecture, we use type (A, ) from Figure 2 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-grained dynamic by integrating forward from t = 0 and guiding the learner only through the loss at the available time points.

Rotating MNIST. This dataset [9] consists of sequences of 28 × 28 grayscale images of gradually clockwise rotated versions 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 [0 , 1] . During training, 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 at time t = 0 , the task is to interpolate the image at the left-out time point t = 3 / 16 .

We follow the evaluation protocol of [60], where 360 images

† indicates results from [60].

are used for training, 360 for testing, and 36 for validation and hyperparameter selection. The recognition model is of type (C, ) 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 5 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 a 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 mTAND 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, we refer the reader to Appendix C.1.

t

∈

[0

,

1)

t

∈

[1

,

2)

t

∈

[2

,

3)

t

∈

[3

,

4)

5 Discussion & Limitations

We have presented an approach to learning (neural) latent SDEs on homogeneous spaces with a Bayesian learning regime. Specifically, we focused on one incarnation of this construction, i.e., the unit n -sphere, due to 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 in terms of overall model complexity, as backpropagating through the SDE solver can be done efficiently (in a 'discretize-thenoptimize' manner) without having to resort to adjoint sensitivity methods or 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 to 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 do use neural networks as high-capacity function approximators, but in a setting where one prescribes the space on which the SDE evolves. Finally, whether it 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 be preferable.

Acknowledgments

This work was supported by the Austrian Science Fund (FWF) under project P31799-N38 and the Land Salzburg within the EXDIGIT project 20204-WISS/263/6-6022 and projects 0102-F1901166KZP, 20204-WISS/225/197-2019.

Source code is available at https://github.com/plus-rkwitt/LatentSDEonHS .

References

D. Applebaum. Probability on Compact Lie Groups . 1st edition. V ol. 70. Probability Theory and Stochastic Modelling. Springer, 2014.

M. Arató. 'A famous nonlinear stochastic equation (Lotka-Volterra model with diffusion)'. In: Mathematical and Computer Modelling 38 (2003), pp. 709-726.

C. Archambeau, M. Opper, Y. Shen, D. Cornford, and J. Shawe-Taylor. 'Variational Inference for Diffusion Processes'. In: NIPS . 2007.

P. Becker, H. Pandya, G. Gebhardt, C. Zhao, J. Taylor, and G. Neumann. 'Recurrent Kalman Networks: Factorized Inference in High-Dimensional Deep Feature Spaces'. In: ICML . 2019.

F. Black and M. Scholes. 'The Pricing of Options and Corporate Liabilities'. In: Journal of Political Economy 81.3 (1973), pp. 637-654.

R. W. Brockett. 'Lie Algebras and Lie Groups in Control Theory'. In: Geometric Methods in System Theory . Springer, 1973, pp. 43-82.

B. C. Brown, A. L. Caterini, B. L. Ross, J. C. Cresswell, and G. Loaiza-Ganem. 'Verifying the Union of Manifold Hypothesis for Image Data'. In: ICLR . 2023.

N. D. Cao and W. Aziz. 'The Power Spherical distribution'. In: arXiv abs/2006.04437 (2020).

F. P. Casale, A. V. Dalca, L. Saglietti, J. Listgarten, and N. Fusi. 'Gaussian Process Prior Variational Autoencoders'. In: NeurIPS . 2018.

R. T. Chen, Y. Rubanova, J. Bettencourt, and D. Duvenaud. 'Neural Ordinary Differential Equations'. In: NeurIPS . 2018.

R. Chetrite and H. Touchette. 'Nonequilibrium Markov Processes Conditioned on Large Deviations'. In: Annales Henri Poincaré 16 (2015), pp. 2005-2057.

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.

T. S. Cohen, M. Geiger, J. Köhler, and M. Welling. 'Spherical CNNs'. In: ICLR . 2018.

M. Connor, G. Canal, and C. Rozell. 'Variational Autoencoder with Learned Latent Structure'. In: AISTATS . 2021.

T. R. Davidson, L. Falorsi, N. DeCao, T. Kipf, and J. M. Tomczak. 'Hyperspherical Variational Autoencoders'. In: UAI . 2018.

E. Dupont, A. Doucet, and Y. W. Teh. 'Augmented Neural ODE'. In: NeurIPS . 2019.

C. Fefferman, S. Mitter, and H. Narayanan. 'Testing the manifold hypothesis'. In: Journal of the American Mathematical Society 29.4 (2016), pp. 983-1049.

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.

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.

M. Horn, M. Moor, C. Bock, B. Rieck, and K. Borgwardt. 'Set Functions for Time Series'. In: ICML . 2020.

A. Iserles, H. Z. Munthe-Kaas, S. P. Nørsett, and A. Zanna. 'Lie-group methods'. In: Acta Numerica 9 (2000), 215--365.

K. Itô. 'Brownian motions in a Lie group'. In: Proceedings of the Japan Academy, Series A, Mathematical Sciences 26.8 (1950).

K. Itô. 'Stochastic Calculus'. In: Lect. Notes Phys. 39 (1975). Ed. by H. Araki, pp. 218-223.

J. Jia and A. R. Benson. 'Neural Jump Stochastic Differential Equations'. In: NeurIPS . 2019.

P. Kidger, J. Foster, X. Li, and T. Lyons. 'Efficient an Accurate Gradients for Neural SDEs'. In: NeurIPS . 2021.

P. Kidger, J. Foster, X. Li, H. Oberhauser, and T. Lyons. 'Neural SDEs as Infinite-Dimensional GANs'. In: ICML . 2021.

P. Kidger, J. Morrill, J. Foster, and T. Lyons. 'Neural Controlled Differential equations for irregular time series'. In: NeurIPS . 2020.

D. P. Kingma and J. Ba. 'Adam: A Method for Stochastic Optimization'. In: ICLR . 2015.

D. P. Kingma and M. Welling. 'Auto-Encoding Variational Bayes'. In: ICLR . 2014.

X. Li, T.-K. L. Wong, R. T. Chen, and D. Duvenaud. 'Scalable Gradients for Stochastic Differential Equations'. In: AISTATS . 2020.

M. Liao. Levy processes in Lie groups . Providence, R.I.: Cambridge University Press, 2004.

W. Liu, R. Lin, Z. Liu, L. Xiong, B. Schölkopf, and A. Weller. 'Learning with Hyperspherical Uniformity'. In: AISTATS . 2021.

W. Liu, Y. Zhang, X. Li, Z. Liu, B. Dai, T. Zhao, and L. Song. 'Deep Hyperspherical Learning'. In: NeurIPS . 2017.

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/1906.02355 (2019).

G. Marjanovic, M. Piggott, and V. Solo. 'A Simple Approach to Numerical Methods for Stochastic Differential Equations in Lie Groups'. In: CDC . 2015.

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 (2018), pp. 4035-4050.

S. Massaroli, M. Poli, J. Park, A. Yamashita, and H. Asama. 'Dissecting Neural ODEs'. In: NeurIPS . 2020.

P. Mettes, E. van der Pol, and C. Snoek. 'Hyperspherical Prototype Networks'. In: NeurIPS . 2019.

J. Morrill, C. Salvi, P. Kidger, and J. Foster. 'Neural rough differential equations for long time series'. In: ICML . 2021.

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.

H. Munthe-Kaas. 'High order Runge-Kutta methods on manifolds'. In: Applied Numerical Mathematics 29.1 (1999). Proceedings of the NSF/CBMS Regional Conference on Numerical Analysis of Hamiltonian Differential Equations, pp. 115-127.

H. Munthe-Kaas. 'Runge-Kutta methods on Lie groups'. In: BIT Numerical Mathematics 38.1 (1998), pp. 92-111.

M. Opper. 'Variational Inference for Stochastic Differential Equations'. In: Annalen der Physik 531.3 (2019), p. 1800233.

S. W. Park, H. Kim, K. Lee, and J. Kwon. 'Riemannian Neural SDE: Learning Stochastic Representations on Manifolds'. In: NeurIPS . 2022.

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.

G. A. Pavliotis. Stochastic Processes and Applications . Springer, 2014.

M. Piggott and V. Solo. 'Geometric Euler-Maruyama Schemes for Stochastic Differential Equations in SO(n) and SE(n)'. In: SIAM Journal on Numerical Analysis 54.4 (2016), pp. 24902516.

Y. Rubanova, R. T. Chen, and D. Duvenaud. 'Latent ODE for Irregularly-Sampled Time Series'. In: NeurIPS . 2019.

M. Schirmer, M. Eltayeb, S. Lessmann, and M. Rudolph. 'Modeling Irregular Time Series with Continuous Recurrent Units'. In: ICML . 2022.

S. N. Shukla and B. M. Marlin. 'Multi-Time Attention Networks for Irregularly Sampled Time Series'. In: ICLR . 2021.

I. Silva, G. Moody, D. J. Scott, L. A. Celi, and R. G. Mark. 'Predicting in-hospital mortality of ICU patients: The PhysioNet/Computing in cardiology challenge 2012'. In: Computing in Cardiology 39 (2012), pp. 245-248.

J. T. Smith, A. Warrington, and S. W. Linderman. 'Simplified State Space Layers for Sequence Modeling'. In: ICLR . 2023.

V. Solo. 'An approach to stochastic system identification in Riemannian manifolds'. In: CDC . 2014.

Y. Song, J. Sohl-Dickstein, D. P. Kingma, A. Kumar, S. Ermon, and B. Poole. 'Score-based generative modeling through stochastic differential equations'. In: ICLR . 2021.

D. W. Stroock. 'On the growth of stochastic integrals'. In: Zeitschrift für Wahrscheinlichkeitstheorie und Verwandte Gebiete 18.4 (1971), pp. 340-344.

C. Tan, Z. Gao, L. Wu, S. Li, and S. Z. Li. 'Hyperspherical Consistency Regularization'. In: CVPR . 2022.

B. Øksendal. Stochastic Differential Equations . Universitext. Springer, 1995.

B. Tzen and M. Raginsky. 'Neural Stochastic Differential Equations: Deep Latent Gaussian Models in the Diffusion Limit'. In: arXiv abs/1905.09883 (2019).

B. Tzen and M. Raginsky. 'Theoretical guarantees for sampling and inference in generative models with latent diffusions'. In: COLT . 2019.

Ç. Yildiz, M. Heinonen, and H. Lahdesmäki. 'ODE 2 VAE: Deep generative second order ODEs with Bayesian neural networks'. In: NeurIPS . 2019.

Supplementary material

In this supplementary material, we report full dataset descriptions, list architecture details, present additional experiments, and report any left-out derivations.

Contents

Societal impact. This work mainly revolves around a novel type of neural SDE model. We envision no direct negative societal impact but remark that any application of such a 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 are interpolated 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.

A Datasets

This section contains a description of 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 PhysioNet (2012)). Hence, we find it essential to highlight the exact data-loading and preprocessing implementations we use: for Human Activity and PhysioNet (2012) experiments, we rely on publicly available code from [50] (which is based upon the reference

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 5 , 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 is 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 is split into 4,194 sequences for training, 1,311 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 a 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 are shown in Figure 5. For both tasks, we use 4,000 images for training, 1,000 images for testing and 1,000 images for validation.

1

1

/

A.3 PhysioNet (2012)

The PhysioNet (2012) [51] dataset has been used many times in the literature. However, a 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.

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 set 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 , 1] (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 on the experimental regime in prior works, observations are either aggregated within 1 min [48; 50] or within 6 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.

While prior work typically focuses on mortality prediction , i.e., predicting one binary outcome per time series, we are interested in interpolating observations at unseen time points. To assess this interpolation capability of a model, we adopt the interpolation task of [50]. In this task, interpolation is 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 is 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 x 1 , . . . , x N denote all values of one variable in the training dataset, and x min , x max denote the corresponding minimum and maximum values of that variable, [50] normalizes by ( x i -x min ) /x max , whereas [49] normalizes by ( x i -x min ) / ( x max -x min ) . 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 × 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 at time t = 0 , 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.

3MM

WL

3

3

9

3

33M

3

3

W

3

M

3

3

7

>

L

330

C

E

m333

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 n -sphere remains the same across all experiments. In particular, given a representation

of an input sequence, say h , we linearly map to (1) the parameters of the power spherical distribution, i.e., the posterior distribution on the initial state, and (2) to 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-symmetric and has, in case of SO( n ) , n ( n -1) / 2 free parameters, we linearly map to Kn ( n -1) / 2 coefficients overall. Unless noted otherwise, we consistently use K = 6 polynomials and set n = 16 . In all experiments with ODE/SDE type models, we use ADAM [28] for optimization, with initial learning rate set to 1e-3 and subject to a cyclic cosine learning rate schedule (cycle length of 60 epochs) [18]. Finally, when replacing our SDE on the n -sphere by an ODE in R n (denoted as mTAND-ODE), we use an 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, ) from Figure 2 with an mTAND encoder (mTAND-Enc [50]) configured as in the publicly-available reference implementation 6 . When used as a standalone module (i.e., without any subsequent ODE/SDE), mTAND-Enc uses the outputs of its final GRU unit (which yields representations in R 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 ∈ R 128 to the linear maps mentioned above. When replacing our SDE on the n -sphere by an ODE in R n (denoted as mTANDODE), the last GRU state is linearly mapped to the location parameter of a Gaussian distribution and to the entries of a diagonal covariance matrix. For mTAND-ODE, the ODE is parameterized as in [48] by a 3-layer MLP (with ELU activations) with 32 hidden units, but the dimensionality of the latent space is increased to 16 (for a 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., an ODE in [48], vs. an mTAND encoder here).

Loss components. In addition to the KL 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, ) from Figure 2. Images are first fed through a convolutional neural network (CNN) with the same architecture as in [49] and the so obtained representations are then input to an mTAND encoder (mTAND-Enc). Different to the mTAND configuration in Appendix B.1, the final GRU unit is set to output representations h ∈ R 32 . 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 task , we compute the mean-squared-error (MSE) on the (sine, cosine) encodings of the pendulum angle and use a Gaussian log-likelihood on 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-Enc approach from [50] fails in the experimental runs reported in [49] (see † mTAND-Enc 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.

B.3 PhysioNet (2012)

In the PhysioNet (2012) interpolation experiment, we parameterize our recognition network by type (B, ) from Figure 2 using representations h ∈ R 3 2 as input to the linear layers implementing the parameters of the drift and the initial distribution. The mTAND encoder is configured as suggested in the mTAND reference implementation for this experiment. We further adjusted the reference implementation 7 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 to the KL divergence terms, the ELBO includes the Gaussian loglikelihood of reconstructed observations at the available time points (using a fixed standard deviation of 0.01). The weightings of the KL divergence terms is fixed to 1e-5, without any further tuning on a 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, ) from Figure 2. In particular, we rely on the convolutional encoder (and decoder) from [60], except that we only use the first image as 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 [0 , 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 use only one polynomial ( K = 1 ), which is reasonable due to the underlying 'simple' dynamic of a constant-speed rotation; setting K > 1 only marginally changes the results listed in Table 5.

Loss components. Gaussian log-likelihood is used on image reconstructions and weights for the KL divergences are selected based on the validation set (and no KL annealing is performed).

B.5 Computing infrastructure

All experiments were executed on systems running Ubuntu 22.04.2 LTS (kernel version 5.15.0-71-generic x86_64 ) with 128 GB of main memory and equipped with either NVIDIA GeForce RTX 2080 Ti or GeForce RTX 3090 Ti GPUs. All experiments are implemented in PyTorch (v1.13 and also tested on v2.0). In this hardware/system configuration, all runs reported in the tables of the main manuscript terminated within a 24-hour time window.

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 n . 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 as a function of the integration steps, moving from 16 steps on Rotating MNIST to 480 on PhysioNet (2012) at a 6 min quantization level and to 2,880 steps on PhysioNet (2012) at a 1 min quantization level. In all experiments, we exclusively use Euler's method as ODE/SDE solver with a fixed step size. Results are summarized in Table 6. In particular, we list the elapsed real time for a forward+backward pass (batch size 50), averaged over 1,000 iterations.

Overall, although our approach implements a 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 the sphere (i.e., as a consequence of the group action) lends itself to a quite efficient implementation where large parts of the computation can be carried out as a single tensor operation (see Appendix D).

↚⊬⊬

⊮∦⊬⊬⊬

Number of epochs

⊬

⊮⊬⊬

⊭⊬⊬

⊯⊬⊬

⋭⊬⊬

⋬⊬⊬

⋫⊬⊬

↚⊬⊬

Training MSE

Accumulated

training time

(forward+backward pass) [s]

⊬glyph[notshortbar]⊬⋪⊬

.

⊬glyph[notshortbar]⊬⋭⊬

⊬glyph[notshortbar]⊬⊭⊬

.

⊬glyph[notshortbar]⊬⊬⋬

Training MSE

⊬

⊭⊬⊬

⋭⊬⊬

⊬glyph[notshortbar]⊬⋪⊬

⊬glyph[notshortbar]⊬⋭⊬

⊬glyph[notshortbar]⊬⊭⊬

⊬glyph[notshortbar]⊬⊬⋬

CNN∤ODE

CNN∤SDE

Ours

CNN∤ODE

CNN∤SDE

Ours

Aside from the direct runtime comparison above, we also point out that while a shorter (forward) runtime per batch ensures a faster inference time, it does not necessarily result in faster overall training time as the convergence speed might differ among methods. To account for the latter, we compare the evolution of training losses in Figure 7 when training on the Rotating MNIST dataset. The upper plot shows training MSE vs. number of epochs and reveals that our method needs fewer updates to converge, whereas the same model but with latent dynamics replaced by a latent ODE [10] or latent SDE [30] in R n converge at approximately the same rate. The bottom plot accounts for

different runtimes per batch and reveals that our method needs only around half of the training time than the latent SDE [30] variant in R n .

We want to point out, though, that runtime is, to a large extent, implementation-dependent, and therefore, the presented results should be taken with caution. For the baselines, we used publicly available implementations of ODE/SDE solvers, i.e., https://github.com/rtqichen/torchdiffeq for the latent ODE [10] models and https://github.com/google-research/torchsde for the latent SDE [30] models.

Comparison to non-ODE/SDE models. Finally, we present a runtime comparison with methods that do not model an underlying (stochastic) differential equation, i.e., mTAND-Full and CRU. This relative runtime comparison was performed on the PhysioNet (2012) dataset and the runtimes/batch are averaged over batches of size 50; results are summarized in Table 7. The PhysioNet (2012) dataset is particularly interesting concerning runtime, as it contains the most time points (480 potential time points for a quantization level of 6 minutes and 2,880 time points for a quantization level of 1 minute).

Table 7: Relative timing results ( forward+backward pass) on PhysioNet (2012) for both quantization levels (reported in the main manuscript).

All runtime experiments were performed on the same hardware (see Appendix B.5 with a 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 8 shows 75 posterior sample paths.

z

x

y

2

0

-

2

2

2

0

0

-

2

-

2

C.3 Uncertainty quantification

In the main part of this manuscript, we primarily focussed on the interpolation and regression/classification capabilities of our method on existing benchmark data. It 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/classification performance, we present preliminary qualitative results on the Pendulum angle regression experiment in Figure 9. This figure shows density estimates for the angle predictions 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 a weight of 1e-3.

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

-

1

0

1

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

-

1

0

1

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

-

1

0

1

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

-

1

0

1

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

-

1

0

1

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

-

1

0

1

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

-

1

0

1

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

-

1

0

1

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

-

1

0

1

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

-

1

0

1

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

-

1

0

1

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

-

1

0

1

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

-

1

.

00

-

0

.

75

-

0

.

50

-

0

.

25

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

D Implementation

To highlight that our approach can be implemented in a few lines of code, Listings 1, 2 and 3 provide a minimal working example (in PyTorch) of how to obtain solutions to an SDE on S n -1 . 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 8 from [8] as well as the einops package 9 are needed.

Listing 1 SDE solver in Lie algebra so ( n ) .

10 11 12 13 14 Listing 2 Generate basis in so ( n ) . Listing 3 Generate SDE solution on S n -1 . 10 11 12 13 14

Listing 4 Example usage with dummy K , µ , κ .

E Theory

In this part, we provide further details on 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 manuscript.

E.1 Background on stochastic differential equations (SDEs)

We first make some initial remarks on stochastic processes and on stochastic differential equations (SDEs) in R n and introduce the concepts and notations that are used throughout this final chapter of the supplementary material.

Given some probability space (Ω , F , P ) , a stochastic process on [0 , T ] with state space R n is a function X : Ω × [0 , T ] → R n such that for each point of time t ∈ [0 , T ] the map X t = X ( · , t ) : Ω → R n , ω ↦→ X ( ω, t ) is measurable. In other words, { X t : t ∈ [0 , T ] } is a collection of random variables X t indexed by t , cf. [66, Chapter I.1; 57, Chapter 2.1; 63, Chapter 1.1]. Conversely, each sample ω ∈ Ω defines a function X ( ω, · ) : [0 , T ] → R n , t ↦→ X ( ω, t ) called a realization or sample path. Thus, a stochastic process models randomness in the space of functions [0 , T ] → R n , or in a "reasonable" subspace thereof. In more formal terms, the law or path distribution P X associated to a stochastic process X 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 a small particle inside a fluid caused by thermal fluctuations, see [68, Chapter 3] for a heuristic approach to the topic. For stochastic processes X that are sufficiently well-behaved, one can define a stochastic integral ∫ t 0 X s d W s , called Itô integral of X with respect to W , cf. [66, Chapter IV.2; 57, Chapter 3.1; 63, Chapter 3.2]. Notably, the result, is again a stochastic process so that one can implicitly define stochastic processes via an Itô integral equation of the form

often denoted as an Itô stochastic differential equation (SDE)

Here, f : R n × [0 , T ] → R n is called the drift coefficient, σ : R n × [0 , T ] → R n × m 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.e., 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 a treatment of SDEs from a numerical perspective, we recommend [64].

In this work, we consider time series data as realizations of a stochastic process, i.e., as samples from a function-valued random variable as described in Section 3.1. To perform variational inference in such a setting, we utilize SDEs to construct function-valued distributions. Specifically, we consider non-autonomous linear Itô SDEs with multiplicative noise that evolve on a 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]

A remark on notation. As outlined in Section 3.4, we aim to learn the parameters ϕ of an approximate posterior. Thus, the SDEs we consider are not specified by an initial random variable X 0 = ξ , but by an initial distribution P X 0 such that X 0 ∼ P X 0 . Therefore, we utilize the notation

A (strong) solution to such an SDE is then understood as a stochastic process X that (i) is adapted (see [66, Definition 4.2]), (ii) has continuous sample paths t ↦→ X ( ω, · ) , (iii) almost surely satisfies the integral equation in Eq. (12), and (iv) there is a random variable ξ ∼ P X 0 that is independent of the Wiener process W such that X 0 = ξ 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 X 0 = ξ might define distinct strong solutions X .

E.2 SDEs in quadratic matrix Lie groups

Definition 1 (cf. [65]) . Let R n × n be the vector space of real n × n matrices. We denote the general linear group, i.e., the group of invertible real n × n matrices equipped with the matrix product, as

GL( n ) is a Lie group with Lie algebra

i.e., R n × n together with the commutator bracket ( A , B ) ↦→ [ A , B ] = AB -BA .

In the context of this work, we specifically consider quadratic matrix Lie groups.

Definition 2 (cf. [62]) . Given some fixed matrix P ∈ R n × n , the quadratic matrix Lie group G is defined as

The Lie algebra g of G then consists of the n × n 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 an 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 3 (cf. [6, Chapter 4.1], [12]) . Let G be a stochastic process that solves the matrix Itô SDE

where V 0 : [0 , T ] → R n × n , V 1 , . . . , V m ∈ R n × n and w 1 , . . . , w m are independent scalar-valued Wiener processes. The process G stays in the quadratic matrix Lie group G ( P ) defined by P ∈ R n × n if

Remark. Although the matrices V 1 , . . . , V m ∈ g can be chosen arbitrarily, it is natural to choose an orthogonal basis of the Lie algebra g . Then m = dim g and each direction in the vector space g 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 ⊤ t P G t is constant, i.e., if d( G ⊤ t P G t ) = 0 . The latter is to be understood in the Itô sense, i.e.,

Here, d w i t d w j t = δ ij d t , as w i and w j are independent scalar Wiener processes. Eq. (19) implies the following two conditions:

̸

As the following calculation shows, these conditions are satisfied if V 1 , . . . , V m ∈ g and V 0 ( t ) = K ( t ) + 1 2 ∑ m i =1 V 2 i where K : [0 , T ] → g .

Numerical solution. Algorithm 1 lists our choice of a numerical solver, which is the one-step geometric Euler-Maruyama scheme from [36] (see also [40, Section 3.1]).

This scheme consists of three parts. First, we divide [0 , T ] into subintervals [ t j , t j +1 ] , starting with t 0 = 0 until t j +1 = T . Second, we solve Eq. (18) pathwise in G over successive intervals [ t j , t j +1 ] by applying a McKean-Gangolli Injection [12] in combination with a numerical solver. Third, the (approximate) path solution is eventually computed as G j +1 = G j exp(Ω j +1 ) .

Algorithm 1 Geometric Euler-Maruyama Algorithm (g-EM) [36].

1:

2:

3:

4:

5:

6:

7:

8:

g

∈ G

0

t

0

G

0

I

repeat over successive subintervals

δ

j

δw

j

t

[

, t

+1

j

t

i

-

j

Ω

t

j

∼ N

j

+1

G

until

j

+1

t

j

(0

+1

=

, δ

(

V

for

j

)

,

0

G

j

T

Remark. We remark that in the original derivation of the algorithm 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, we see that line 6 reduces to

E.3 SDEs in homogeneous spaces of quadratic matrix Lie groups

In the following, we use our results of Appendix E.2 to define stochastic processes in a homogeneous space M of a quadratic matrix Lie group G . The main idea is, given a starting point Z 0 ∈ M , to translate a path G : [0 , T ] → G into a path Z = G · Z 0 in M via the Lie group action. Notably, this construction can be done at the level of stochastic processes, i.e., a stochastic process G : [0 , T ] × Ω →G induces a stochastic process Z = G · Z 0 : [0 , T ] × Ω → M called a one-point motion , see [31, Chapter 2]. In particular, if Z 0 follows some distribution P in M and G solves an SDE of the form as in Eq. (18), then Z solves d Z t = d G t · Z 0 , i.e.,

Notably, Eq. (3) guarantees that a solution to the SDE in Eq. (18) remains in the group G , and consequently a solution to Eq. (24) remains in M . As we only consider Lie group actions that are given by matrix-vector multiplication, we substitute the notation G · Z by the more concise GZ from now on.

j

t

(

)

exp(Ω

j

i

= 1

-

1

∑

2

+1

, . . . , m

m

i

=1

)

j

V

2

i

)

δ

+

n

j

+1

]

,

j

m

i

=1

∑

0

δw

i

j

V

i

▷

▷

Matrix exp. ensures that

Ω

j

G

j

+1

+1

∈

.

.

Remark. For subsequent results, it is convenient to represent Eq. (24) in the form of

i.e., with diffusion coefficient σ : ( z , t ) ↦→ [ V 1 z , . . . , V m z ] ∈ R n × m and an m -dimensional Wiener process W = [ w 1 , w 2 , . . . , w m ] ⊤ . Notably, the matrix σ ( z , t ) σ ( z , t ) ⊤ is given by

Numerical solution. As Z = GZ 0 , a numerical solution to Eq. (24) is easily obtained from the numerical approximation of G , i.e., the numerical solution to Eq. (18) via Algorithm 1. In fact, it suffices to replace line 2 by Z 0 ∼ P Z 0 and line 7 by Z j +1 ← exp(Ω j +1 ) Z j . Notably, we do not need to compute the matrices G j but only the iterates Ω j .

E.4 SDEs on the unit n -sphere

Next, we consider the homogeneous space ( S n -1 , SO( n )) , i.e., the unit n -sphere S n -1 together with the action of the group SO( n ) of rotations.

Notably, the orthogonal group O( n ) is the quadratic matrix Lie group G ( P ) defined by P = I n , i.e., O( n ) = { A ∈ R n × n | A ⊤ A = I n } and the special orthogonal group SO( n ) is its connected component that contains the identity I n . Thus, Lemma 3 ensures that a solution to Eq. (24) remains in S n -1 if K : [0 , T ] → so ( n ) and V 1 , . . . , V m ∈ so ( n ) .

̸

In the variational setting of the main manuscript, we consider stochastic processes that solve such an SDE for the choice of { V i } m i =1 = { α ( e k e ⊤ l -e l e ⊤ k ) } 1 ≤ k<l ≤ n , where α = 0 and e k denotes the k -th standard basis vector of R n . With that choice, Eq. (24) becomes

The diffusion term in Eq. (27) consists of a matrix-valued noise d U t where u ii = 0 and u ij = -u ji are independent Wiener processes (for i < j ). The following lemma (Lemma 4) converts the diffusion term to the standard form of only a vector-valued Wiener process. This facilitates (in Appendix E.5) deriving the KL-divergence between two solutions of such a SDE from existing results.

Lemma 4. Let W be an n -dimensional Wiener process and let P Z 0 be a distribution on R n with finite second moment. Further, let K : [0 , T ] → so ( n ) be continuous. The solutions to the SDE in Eq. (27) and

have the same probability distribution. Furthermore, for a given initial value Z 0 a.s. = ξ ∼ P Z 0 , a sample path solution of each equation also is a sample path solution of the other one.

Proof. W.l.o.g. assume α = 1 . The SDEs from Eq. (27) and Eq. (28) have the same drift coefficient but a 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 σ 1 : ( z , t ) ↦→ [ V 1 z , . . . , V m z ] and σ 2 : ( z , t ) ↦→ I n -zz ⊤ , respectively. The lemma then follows from [61, Theorem 2.1] which states that in the present setting, the solution to Eq. (27) and Eq. (28) have the same distribution and equal sample paths if σ 1 σ ⊤ 1 = σ 2 σ ⊤ 2 , i.e., if for every z ∈ R n ,

To prove this equality, first, note that I n -zz ⊤ is the projection onto the orthogonal complement of z and so

To calculate the left-hand side, we insert our choice of { V i } m i =1 = { α ( e k e ⊤ l -e l e ⊤ k ) } 1 ≤ k<l ≤ n , which yields

We want to point out that for vanishing drift K = 0 , the SDE in Eq. (28) becomes the well-known Stroock equation [23; 55] for a spherical Wiener process, i.e., a diffusion process on the sphere S n -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 ∈ SO( n ) , we consider the process R Z , cf. [23, p. 220f.]. For brevity, we introduce f : z ↦→-α 2 ( n -1) / 2 z and σ : z ↦→ α ( I n -zz ⊤ ) so that Eq. (28) becomes d Z t = f ( Z t )+ σ ( Z t ) dW t . Notably, the identities R f ( Z t ) = f ( R Z t ) and R σ ( Z t ) R ⊤ = σ ( R Z t ) hold. Thus, the process R Z t satisfies

As ˜ W t = R W t is again a Wiener process, the processes Z t and R Z t follow the same SDE. The invariance of Z w.r.t. rotations is also apparent from the SDE in Eq. (27) as U t is already invariant under rotations, i.e., at each t ∈ [0 , T ] the random variables ˜ U t = R ⊤ U t R and U t are equal in distribution.

E.5 Kullback-Leibler (KL) divergences

First, in Lemma 5, we state the general form of the KL divergence between path distributions of stochastic processes in R n .

Lemma 5. Let X and Y be stochastic processes with (path) distributions P X and P Y , respectively, and let W be a vector-valued Wiener process; all with state space R n . Further, let P X 0 and P Y 0 be distributions on R n with finite second moments. If X and Y solve 10 the following Itô SDEs on [0 , T ] ,

then the KL divergence between P Y and P X is given by

Here, f and g denote vector-valued drift coefficients, and σ denotes the matrix-valued diffusion coefficient. Furthermore, q Y t denotes the marginal density of Y t (i.e., the process Y at time t ) and Σ + denotes the pseudo inverse of Σ = σσ ⊤ .

For a proof via Girsanov's theorem , we refer the reader to [69, Section A] and [67, Lemma 2.1]. For a heuristic derivation via the limit of an Euler discretization, we refer to [43, Section 4].

The following corollary specifies Lemma 5 to a particular class of SDEs that is compatible with the SDEs introduced in Appendix E.3.

Corollary 6. Let V X 0 , V Y 0 : [0 , T ] → R n × n be continuous maps and let V 1 , . . . , V m ∈ R n × n . Further, let w 1 , . . . , w m be independent scalar-valued Wiener processes and let P X 0 and P Y 0 be distributions on R n with finite second moments. If X and Y are stochastic processes on [0 , T ] with (path) distributions P X and P Y , resp., that solve

then, the KL divergence between P Y and P X is given by

where V ∆ 0 = V X 0 -V Y 0 , and Σ + ( z ) denotes the pseudo inverse of Σ ( z ) = ∑ m i =1 V i zz ⊤ V ⊤ i .

Proof. The noise term of the SDEs can be written as σ ( Z t , t ) dW t where σ : R n × [0 , T ] → R n × n is specified in Eq. (25). Thus, the corollary is a direct consequence of Lemma 5, as σ ( z , t ) σ ( z , t ) ⊤ = ∑ m i =1 V i zz ⊤ V ⊤ i , see Eq. (26).

Finally, the following proposition (Proposition 7) considers exactly the SDEs discussed in the main manuscript. The process X is 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 (KL divergence for one-point motions on S n -1 ) . Let K ( t ) : [0 , T ] → so ( n ) be a continuous map and let α = 0 . If X and Y are stochastic processes on [0 , T ] with (path) distributions P X and P Y , resp., that solve the following Itô SDEs

where P X 0 S n -1 and P Y 0 S n -1 are distributions supported on the unit n -sphere S n -1 with finite second moment. Then, the KL divergence between P Y and P X is given by

Proof. By Lemma 4, we can replace Eq. (36) and Eq. (37) by

as the solutions follow the same distribution P X and P Y , respectively.

The proposition is thus a consequence of Lemma 5 for a specific choice of f, g and σ , i.e.,

To finish the proof, we highlight two points. First, Lemma 4 ensures that the process Y evolves on the unit n -sphere S n -1 , so it suffices to integrate q Y t over S n -1 ⊂ R n instead of over the entire R n as in Eq. (41). Second, the integrand simplifies to

where the first equality holds because ( I n -zz ⊤ ) is a projection, and the second equality holds because z ⊤ K ( t ) z = 0 as K ( t ) ∈ g is skew-symmetric.

References (Supplementary material)

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.

P. Becker, H. Pandya, G. Gebhardt, C. Zhao, J. Taylor, and G. Neumann. 'Recurrent Kalman Networks: Factorized Inference in High-Dimensional Deep Feature Spaces'. In: ICML . 2019.

A. M. Bloch, P. E. Crouch, J. E. Marsden, and A. K. Sanyal. 'Optimal Control and Geodesics on Quadratic Matrix Lie Groups'. In: Foundations of Computational Mathematics 8 (4 2008), pp. 469-500.

R. W. Brockett. 'Lie Algebras and Lie Groups in Control Theory'. In: Geometric Methods in System Theory . Springer, 1973, pp. 43-82.

N. D. Cao and W. Aziz. 'The Power Spherical distribution'. In: arXiv abs/2006.04437 (2020).

F. P. Casale, A. V. Dalca, L. Saglietti, J. Listgarten, and N. Fusi. 'Gaussian Process Prior Variational Autoencoders'. In: NeurIPS . 2018.

R. T. Chen, Y. Rubanova, J. Bettencourt, and D. Duvenaud. 'Neural Ordinary Differential Equations'. In: NeurIPS . 2018.

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.

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.

K. Itô. 'Stochastic Calculus'. In: Lect. Notes Phys. 39 (1975). Ed. by H. Araki, pp. 218-223.

I. Karatzas and S. E. Shreve. Brownian motion and stochastic calculus . 2nd edition. Springer, 1998.

D. P. Kingma and J. Ba. 'Adam: A Method for Stochastic Optimization'. In: ICLR . 2015.

P. Kloeden and E. Platen. Numerical methods for stochastic differential equations . 1991.

J. M. Lee. Introduction to smooth manifolds . Graduate texts in mathematics. Springer, 2003.

X. Li, T.-K. L. Wong, R. T. Chen, and D. Duvenaud. 'Scalable Gradients for Stochastic Differential Equations'. In: AISTATS . 2020.

M. Liao. Levy processes in Lie groups . Providence, R.I.: Cambridge University Press, 2004.

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 (2018), pp. 4035-4050.

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.

M. Opper. 'Variational Inference for Stochastic Differential Equations'. In: Annalen der Physik 531.3 (2019), p. 1800233.

D. Revuz and M. Yor. Continuous Martingales and Brownian Motion . 3rd edition. Springer, 1999.

Y. Rubanova, R. T. Chen, and D. Duvenaud. 'Latent ODE for Irregularly-Sampled Time Series'. In: NeurIPS . 2019.

D. Sanz-Alonso and A. M. Stuart. 'Gaussian approximations of small noise diffusions in Kullback-Leibler divergence'. In: Communications in Mathematical Sciences 15.7 (2017), pp. 2087-2097.

S. Särkkä and A. Solin. Applied stochastic differential equations . Cambridge University Press, 2019.

M. Schirmer, M. Eltayeb, S. Lessmann, and M. Rudolph. 'Modeling Irregular Time Series with Continuous Recurrent Units'. In: ICML . 2022.

S. N. Shukla and B. M. Marlin. 'Multi-Time Attention Networks for Irregularly Sampled Time Series'. In: ICLR . 2021.

I. Silva, G. Moody, D. J. Scott, L. A. Celi, and R. G. Mark. 'Predicting in-hospital mortality of ICU patients: The PhysioNet/Computing in cardiology challenge 2012'. In: Computing in Cardiology 39 (2012), pp. 245-248.

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.

D. W. Stroock. 'On the growth of stochastic integrals'. In: Zeitschrift für Wahrscheinlichkeitstheorie und Verwandte Gebiete 18.4 (1971), pp. 340-344.

B. Øksendal. Stochastic Differential Equations . Universitext. Springer, 1995.

Ç. Yildiz, M. Heinonen, and H. Lahdesmäki. 'ODE 2 VAE: Deep generative second order ODEs with Bayesian neural networks'. In: NeurIPS . 2019.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\n- 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)\n\nQUESTIONS:\nThe 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?\n\nLIMITATIONS:\nThe authors have adequately addressed potential social implications.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-14"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nStrengths :\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\nWEAKNESSES:\nWeaknesses :\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\nQUESTIONS:\n- 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.\n\nLIMITATIONS:\n- 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1. (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.\n\nWEAKNESSES:\n1. (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\nQUESTIONS:\ntypos:\n\n\n* lines 25, 58: Paramaterizing -> parameterizing\n* line 178: extra $≤$\n* line 180: arithmetics -> arithmetic\n\nLIMITATIONS:\nLimmitations are well discussed explicitly in the paper.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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\nSTRENGTHS:\nThis 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.\n\nWEAKNESSES:\n1. 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\nQUESTIONS:\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?\n\nLIMITATIONS:\nThe paper includes a discussion on the limitations and possible societal impact.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-23"}