SlowGuess/NeurIPS_2023_OCR huggingface dataset:
This motion capture dataset<sup>5</sup>, 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 \times 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.
Figure 5: Illustration of one exemplary training time series from the (left) regression and (right) pendulum position interpolation task, respectively, from [49]. Note that the images shown here are inverted for better visualization.
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.
<sup>5</sup>https://doi.org/10.24432/C57G8X
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  $\approx 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}, \ldots, 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 \times 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.
Figure 6: Five (counted in rows) exemplary training sequences from the rotating MNIST dataset.
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  $\mathbf{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  $\mathbf{K}^{\phi}$ . Here, the superscript  $\phi$  denotes the dependency of  $\mathbf{K}$  on the (parameters of the) recognition network. As  $\mathbf{K}^{\phi}$  is skew-symmetric and has, in case of  $\mathrm{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  $\mathbb{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 torchdiffreq package).
B.1 Human Activity
As recognition network we use type (B,  $\odot$ ) from Figure 2 with an mTAND encoder (mTAND-Enc [50]) configured as in the publicly-available reference implementation. 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  $\mathbb{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  $\mathbf{h} \in \mathbb{R}^{128}$  to the linear maps mentioned above. When replacing our SDE on the  $n$ -sphere by an ODE in  $\mathbb{R}^n$  (denoted as mTAND-ODE), 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, O) 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  $\mathbf{h} \in \mathbb{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  $\dagger$  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.
$^{6}$ https://github.com/reml-lab/mTAN
B.3 PhysioNet (2012)
In the PhysioNet (2012) interpolation experiment, we parameterize our recognition network by type  $(\mathbf{B},\circ)$  from Figure 2 using representations  $\mathbf{h}\in \mathbb{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<sup>7</sup> 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 log-likelihood 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  $(\mathbf{C},\circ)$  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  $\mathbb{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.
<sup>7</sup>https://github.com/boschresearch/Continuous-Recurrent-Units
Table 6: Comparison of the elapsed real time of a forward+backward pass over a batch (of size 50) across different tasks/datasets (with varying number of integration steps). For each task, the overall model sizes (in terms of the number of parameters) are comparable. Listed values are averages over 1,000 iterations  $\pm$  the corresponding standard deviations. We remark that the runtime of our approach when using  $(\mathbb{R}^n,\mathrm{GL}(n))$  instead of  $(\mathbb{S}^{n - 1},\mathrm{SO}(n))$  differs only marginally.
Rotating MNIST	PhysioNet (2012) (6 min)	PhysioNet (2012) (1 min)
Number of integration steps	16	480	2,880
CNN/mTAND-ODE ([10; 48])	0.053 ± 0.004	0.926 ± 0.017	3.701 ± 0.033
CNN/mTAND-SDE ([30])	0.112 ± 0.009	2.075 ± 0.042	10.273 ± 0.173
Ours - (Sn-1, SO(n))	0.055 ± 0.005	0.179 ± 0.008	2.693 ± 0.027
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).
Figure 7: Training time comparison on Rotating MNIST for five different training runs per model. Bottom: Training MSE vs. (accumulated) training time of each model. Top: Training MSE (dashed) and standard deviations (shaded) vs. number of training epochs.
Aside from the direct runtime comparison above, we also point out that while 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  $\mathbb{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  $\mathbb{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/torchdiffreq 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).
6 min	1 min
CRU	22×	25×
f-CRU	17×	20×
mTAND-Full	1×	1×
Ours	5×	55×
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.
Figure 8: Shown are 75 posterior sample paths from our model on the stochastic Lorenz attractor data of [30].
C.3 Uncertainty quantification
In the main part of this manuscript, we primarily focussed on the interpolation and 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.
Figure 9: Top: Visualization of quality of fit and uncertainty for the first 12 testing time series in the pendulum regression task. Shown is the ground truth (in blue) and the predicted angle (sine). Shaded regions correspond to densities, estimated from quantiles of 1,000 draws from the posterior distribution. Bottom: Zoomed-in plot of the first time series.
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  $\mathbb{S}^{n - 1}$ . The only part left out in this example is the generation of the time evolution of  $\mathbf{K}$ . To execute the code snippet in Listing 4, the reference implementation of the power spherical distribution<sup>8</sup> from [8] as well as the einops package<sup>9</sup> are needed.
Listing 1 SDE solver in Lie algebra  $\mathfrak{so}(n)$ .
from einops import repeat   
class PathSampler(nnModule): def__init__(self，dim): super(PathSampler，self).__init__(） self.register_buffer("skew_sym BASIS"，gen BASIS(dim)) self.alpha  $\equiv$  nn.Parameters(torch.tensor(1./dim，dtype  $\equiv$  torch.float32))   
def forward(self,Kt，delta): bs，nt  $=$  Kt.shape[0]，Kt.shape[1] Et  $=$  repeat(self.skew_sym BASIS,'kab->bs nt kab'，bs=bs，nt=nt) noise  $=$  torch.randn(bs，nt，Et.size(2)，1，1，device  $\equiv$  Kt_device) Vt  $=$  (noise\*Et\*delta.sqrt()\*self.alpha).sum(2) return Kt\*delta+Vt
Listing 2 Generate basis in  $\mathfrak{so}(n)$ .
def gen BASIS(n):  
    r = int(n*(n-1)/2)  
    B = torch.zeros(r, n, n)  
    for nth, (i,j) in enumerate(torch.combinations(torch.arange(n)):  
        B[nth, i, j] = 1.  
    return B - B.transpose(1, 2)
Listing 3 Generate SDE solution on  $\mathbb{S}^{n - 1}$
from torch.distributions import Distribution   
from torch import Tensor   
def sample(t: Tensor, pzx: Distribution, Kt: Tensor, sampler: PathSampler):   
z0 = pzx.rsample().unsqueeze(-1)   
delta = torch(diff(t)[0]   
Ot = sampler(Kt, delta)   
Qt = torch.matrix_exp(Ot.contiguous())   
zi = [z0]   
for tidx in range(0, len(t)-1):   
Qtz = torch/matmul(Qt[:,tidx,::,], zi[-1])   
zi.append(Qtz)   
return torch.stack(zi, dim=1).squeeze()
<sup>8</sup>https://github.com/nicola-decao/power_spherical  
<sup>9</sup>https://einops.rocks/
Listing 4 Example usage with dummy  $\mathbf{K},\pmb {\mu},\kappa$
1 import torch
2 import torch.nn as nn
3 from power_spherical import PowerSpherical
4
5 n = 16
6 batch_size = 4
7 num_time_points = 100
8
9 device = 'CUDA' if torch.cuda.is-available() else 'cpu'
10 t = torch.linspace(0, 1, num_timepoints, device=device)
11 sampler = PathSampler(n).to(device)
12
13 mu = torch.randn(batch_size, n, device=device) # batch of  $\mu^{\prime}s$ 
14 kappa = torch RAND(batch_size, 1, device=device) # batch of  $\kappa^{\prime}s$ 
15 pzx = PowerSpherical(
16 torch(nn.Functionalnormalize(mu),
17 kappa.squeeze())
18
19 # Kt just contains random skew-symmetric matrices in this example
20 Kt = torch.randn(batch_size, len(t), n, n, device=device)
21 Kt = Kt - Kt.transpose(2,3)
22 zis = sample(t, pzx, Kt, sampler)
23 # zis is now a tensor of size (bs, num_timepoints, n)
E Theory
In this part, we provide further details on the construction of SDEs on homogeneous spaces of quadratic matrix Lie groups  $\mathcal{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  $\mathbb{R}^n$  and introduce the concepts and notations that are used throughout this final chapter of the supplementary material.
Given some probability space  $(\Omega, \mathcal{F}, P)$ , a stochastic process on  $[0, T]$  with state space  $\mathbb{R}^n$  is a function  $X: \Omega \times [0, T] \to \mathbb{R}^n$  such that for each point of time  $t \in [0, T]$  the map  $X_t = X(\cdot, t): \Omega \to \mathbb{R}^n$ ,  $\omega \mapsto X(\omega, t)$  is measurable. In other words,  $\{X_t: t \in [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  $\omega \in \Omega$  defines a function  $X(\omega, \cdot): [0, T] \to \mathbb{R}^n$ ,  $t \mapsto X(\omega, t)$  called a realization or sample path. Thus, a stochastic process models randomness in the space of functions  $[0, T] \to \mathbb{R}^n$ , or in a "reasonable" subspace thereof. In more formal terms, the law or path distribution  $\mathbb{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].