SlowGuess/NeurIPS_2023_OCR huggingface dataset:
As indicated in Eq.(9),  $\phi$  enters the SDE associated with the approximate posterior  $q_{\phi}(\mathbf{z}|\mathbf{x})$  in two ways: first, through the parametric distribution on initial states  $Z_{0}$ , and second, via the drift component  $\mathbf{K}^{\phi}(\mathbf{x})$ . For  $\mathcal{P}_0^\phi$ , we choose the power spherical [8] distribution  $\mathcal{S}$  with location parameter  $\mu^{\phi}(\mathbf{x})$  and concentration parameter  $\kappa^{\phi}(\mathbf{x})$ , i.e.,  $\mathcal{P}_0^\phi(Z_0|\mathbf{x}) = \mathcal{S}(\pmb{\mu}^\phi(\mathbf{x}), \kappa^\phi(\mathbf{x}))$ . We realize  $\pmb{\mu}^{\phi}$ ,  $\kappa^{\phi}$  and  $\mathbf{K}^{\phi}$  by a task specific recognition network capable of learning a representation for time series  $\mathbf{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. Vector-valued 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,  $\circ$ ) and (B,  $\circ$ ), 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,  $\circ$ ) from Figure 2, which directly maps this image to a representation. In all cases, the so obtained representations are linearly mapped to the parameters  $(\mu^{\phi}, \kappa^{\phi})$  of the power spherical distribution and to a skew-symmetric matrix  $\mathbf{K}^{\phi}$  representing the drift at time 0 (and evolving in time as described below).
Time evolution of the drift. To ensure that  $\mathbf{K}^{\phi}(\mathbf{x})(t)$  is continuous and can be evaluated for every  $t\in [0,T]$ , we use Chebyshev polynomials as suggested by [37]. We found this choice suitable for our experiments in Section 4, but in principle any differentiable function  $[0,T]\to \mathfrak{so}(n)$  can be used. In our implementation, a neural network maps a time series  $\mathbf{x}$  to  $K$  skew-symmetric matrices  $\mathbf{K}_i^\phi (\mathbf{x})\in \mathfrak{so}(n)$  and subsequently  $\mathbf{K}^{\phi}(\mathbf{x})(t)$  is computed by
$$
\mathbf {K} ^ {\phi} (\mathbf {x}) (t) = \sum_ {i = 1} ^ {K} \mathbf {K} _ {i} ^ {\phi} (\mathbf {x}) p _ {i} (t), \tag {11}
$$
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
Figure 2: Recognition network types.
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 continuous-time models of latent variables. This includes the latent ODE approach of [10; 48] (ODE-ODE),
2 Chebyshev polynomials are recursively defined via  $p_{i + 1}(t) = 2tp_i(t) - p_{i - 1}(t)$  with  $p_0(t) = 1, p_1(t) = t$ .
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
Human Activity. This motion capture dataset consists of time series collected from five individuals performing different activities. The time series are 12-dimensional, which corresponds to 3D coordinates captured by four motion sensors. Following the pre-processing regime of [48], we have 6,554 sequences available with 211 unique time points overall. The task is to assign each time point 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, O) 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.
Table 1: Human Activity.
Accuracy [%]
†ODE-ODE [48]	87.0 ± 2.8
†mTAND-Enc [50]	90.7 ± 0.2
†mTAND-Full [50]	91.1 ± 0.2
mTAND-ODE	90.4 ± 0.2
Ours	90.6 ± 0.4
$\dagger$  indicates results from [50].
We visually assess this intuition after repeating the experiment with a decreased latent space dimension of  $n = 3$  (i.e.,  $\mathbb{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  $\bullet$  and  $\bullet$ , 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
Figure 3: Latent space paths with and without label switches (on  $\mathbb{S}^2$ ) and distribution of KL divergences.
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, O) from Figure 2 using the same CNN as in [49] and map latent states to (sine, cosine) tuples by the same decoder as in the latter work. 2,000 images are used for training, 1,000 images for testing and an additional 1,000 images for validation and hyperparameter selection.
<sup>3</sup>https://doi.org/10.24432/C57G8X
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
Table 2: Pendulum regression.
MSE (×10-3)
†ODE-ODE [48]	15.40 ± 2.85
†mTAND-Enc [50]	65.64 ± 4.05
†CRU [49]	4.63 ± 1.07
†f-CRU [49]	6.16 ± 0.88
‡S5 [52]	3.41 ± 0.27
mTAND-ODE	4.65 ± 1.21
mTAND-Enc	3.20 ± 0.60
CRU	3.74 ± 0.27
f-CRU	3.25 ± 0.26
Ours	3.84 ± 0.35
$\dagger ,\ddagger$  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<sup>4</sup> 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, O) 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
Table 4: Pendulum interpolation.
MSE (×10-3)	I I I I I I
†mTAND-Full [50]	15.40 ± 0.07	I I I I I I
†ODE-ODE [48]	15.06 ± 0.14	I I I I I I
†f-CRU [49]	1.39 ± 0.07	I I I I I I
†CRU [49]	1.00 ± 0.07	I I I I I I
mTAND-ODE	8.23± 0.10	I I I I I I
f-CRU	2.61 ± 0.19	I I I I I I
CRU	2.06 ± 0.05	I I I I I I
Ours	8.15 ± 0.06	I I I I I I
$\dagger$  indicates results from [49].
4We set the mTAND encoder such that its GRU returns an output for each point in time. These latent states are then mapped to an angle prediction by the same decoder as in [49].
Table 3: Results on the PhysioNet (2012) interpolation task from [50] for two quantization levels (Top: 6 min; Bottom: 1 min). Percentages indicate the fraction of observed time points. We report the MSE  $(\times 10^{-3})$  on the observations at unseen time points, averaged across all testing sequences. Results for mTAND-ODE at 1 min quantization (bottom part) are not listed, as training time exceeded our 24-hour threshold.
Observed % →	50%	60%	70%	80%	90%	100%
CRU [49]	5.11 ± 0.40	4.81 ± 0.07	5.60 ± 0.70	5.54 ± 0.61	5.85 ± 0.82	1.77 ± 0.92
f-CRU [49]	5.24 ± 0.49	4.96 ± 0.36	4.85 ± 0.10	5.66 ± 0.71	6.11 ± 0.82	1.10 ± 0.95
mTAND-Full [50]	3.61 ± 0.08	3.53 ± 0.06	3.48 ± 0.08	3.55 ± 0.12	3.58 ± 0.09	0.49 ± 0.03
mTAND-ODE	3.38 ± 0.03	3.33 ± 0.02	3.32 ± 0.03	3.33 ± 0.04	3.35 ± 0.02	1.82 ± 0.04
Ours	3.14 ± 0.03	3.06 ± 0.05	3.01 ± 0.05	2.94 ± 0.07	3.00 ± 0.08	1.55 ± 0.01
CRU [49]	5.20 ± 0.44	5.25 ± 0.55	5.22 ± 0.60	5.23 ± 0.59	5.59 ± 0.78	1.82 ± 0.97
f-CRU [49]	5.53 ± 0.54	5.22 ± 0.58	5.41 ± 0.65	5.25 ± 0.60	5.80 ± 0.76	0.84 ± 0.16
†mTAND-Full [50]	4.19 ± 0.03	4.02 ± 0.05	4.16 ± 0.05	4.41 ± 0.15	4.80 ± 0.04	0.47 ± 0.01
Ours	3.14 ± 0.04	3.03 ± 0.03	3.01 ± 0.02	3.04 ± 0.09	3.03 ± 0.14	1.54 ± 0.02
$\dagger$  indicates results from [50] (which only lists the 1-min setting).
chitecture, we use type (A,  $\circ$ ) 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 \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. 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$ .
Table 5: Rotating MNIST.
MSE (×10-3)
†GPPVAE-dis	30.9 ± 0.02
†GPPVAE-joint	28.8 ± 0.05
†ODE²VAE	19.4 ± 0.06
†ODE²VAE-KL	18.8 ± 0.31
CNN-ODE	14.5 ± 0.73
Ours	11.5 ± 0.38
$\dagger$  indicates results from [60].
We follow the evaluation protocol of [60], where 360 images
are used for training, 360 for testing, and 36 for validation and hyperparameter selection. The recognition model is of type  $(\mathbf{C},\mathbf{\Theta})$  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.
Figure 4: Exemplary reconstructions on the Rotating MNIST data. Shown are the results (on one testing sequence) by integrating forward from  $t = 0$  (marked red) to  $t = 4$ , i.e., three times longer than what is observed ( $t \in [0,1)$ ) during training.
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-then-optimize" 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.
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-F1901166-KZP, 20204-WISS/225/197-2019.
Source code is available at https://github.com/plus-rkwitt/LatentSDEonHS.
[1] D. Applebaum. Probability on Compact Lie Groups. 1st edition. Vol. 70. Probability Theory and Stochastic Modelling. Springer, 2014.
[2] M. Arato. “A famous nonlinear stochastic equation (Lotka-Volterra model with diffusion)". In: Mathematical and Computer Modelling 38 (2003), pp. 709-726.
[3] C. Archambeau, M. Opper, Y. Shen, D. Cornford, and J. Shawe-Taylor. "Variational Inference for Diffusion Processes". In: NIPS. 2007.
[4] 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.
[5] F. Black and M. Scholes. “The Pricing of Options and Corporate Liabilities”. In: Journal of Political Economy 81.3 (1973), pp. 637–654.
[6] R. W. Brockett. "Lie Algebras and Lie Groups in Control Theory". In: Geometric Methods in System Theory. Springer, 1973, pp. 43-82.
[7] B. C. Brown, A. L. Caterini, B. L. Ross, J. C. Cresswell, and G. Loaiza-Ganem. “Verifying the Union of Manifold Hypothesis for Image Data”. In: ICLR. 2023.
In this supplementary material, we report full dataset descriptions, list architecture details, present additional experiments, and report any left-out derivations.
A Datasets 13
A.1 Human Activity 14
A.2 Pendulum 14
A.3 PhysioNet (2012) 14
A.4 Rotating MNIST 15
B Architecture details 15
B.1 Human Activity 16
B.2 Pendulum regression & interpolation 16
B.3 PhysioNet (2012) 17
B.4 Rotating MNIST 17
B.5 Computing infrastructure 17
C Additional empirical results 17
C.1 Runtime measurements 17
C.2 Stochastic Lorenz attractor 19
C.3 Uncertainty quantification 19
D Implementation 21
E Theory 23
E.1 Background on stochastic differential equations (SDEs) 23
E.2 SDEs in quadratic matrix Lie groups 24
E.3 SDEs in homogeneous spaces of quadratic matrix Lie groups 25
E.4 SDEs on the unit  $\pmb{n}$  -sphere 26
E.5 Kullback-Leibler (KL) divergences 27
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.
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