SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Sebastian Zeng, Florian Graf, Roland Kwitt
University of Salzburg, Austria
{sebastian.zeng, florian.graf, roland.kwitt}@plus.ac.at
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.
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
Figure 1: Sample paths from a prior (left) and posterior (right) process on the sphere. Initial values are marked by  $\bullet$ .
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.
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 Ito 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.
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  $\mathbf{x}^1,\ldots ,\mathbf{x}^N$ . In search for a reasonable model that describes the data-generating system, we consider a continuous stochastic process  $X:\Omega \times [0,T]\to \mathbb{R}^{d}$  where each data sample refers to a continuous path  $\mathbf{x}^i:[0,T]\to \mathbb{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  $\mathbf{x}^i$ , but only to observations  $\mathbf{x}^i(t)\in \mathbb{R}^d$  at a finite number of time points  $t = t_k$  that can vary between the  $\mathbf{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  $\mathbf{z}^i$  is drawn from a parametric prior distribution  $p_{\theta^*}(\mathbf{z})$  with true parameters  $\theta^*$ ; second, an observation path  $\mathbf{x}^i$  is drawn from the conditional distribution  $p_{\theta^*}(\mathbf{x}|\mathbf{z}^i)$ .
Essentially, this is the problem setting of [29]; however, the latent variables  $\mathbf{z}$ , as well as the observations  $\mathbf{x}^i$  are path-valued and the true posterior  $p_{\theta}(\mathbf{z}|\mathbf{x})$  is intractable in general. To efficiently perform variational inference (e.g., with respect to  $\mathbf{z}$ ) and to fit the model to observations, we choose a tractable approximate posterior  $q_{\phi}(\mathbf{z}|\mathbf{x})$  from some family of parametric distributions, a suitable prior  $p_{\theta}(\mathbf{z})$ , and seek to maximize the evidence lower bound (ELBO)
$$
\log p _ {\boldsymbol {\theta}} \left(\mathbf {x} ^ {i}\right) \geq \mathcal {L} (\boldsymbol {\theta}, \phi ; \mathbf {x} ^ {i}) = - D _ {\mathrm {K L}} \left(q _ {\phi} \left(\mathbf {z} \mid \mathbf {x} ^ {i}\right) \| p _ {\boldsymbol {\theta}} (\mathbf {z})\right) + \mathbb {E} _ {\mathbf {z} \sim q _ {\phi} \left(\mathbf {z} \mid \mathbf {x} ^ {i}\right)} \left[ \log p _ {\boldsymbol {\theta}} \left(\mathbf {x} ^ {i} \mid \mathbf {z}\right) \right]. \tag {1}
$$
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 Ito SDEs of the form
$$
\mathrm {d} Z _ {t} = f \left(Z _ {t}, t\right) \mathrm {d} t + \sigma \left(Z _ {t}, t\right) \mathrm {d} W _ {t}, \quad t \in [ 0, T ], \quad Z _ {0} \sim \mathcal {P}. \tag {2}
$$
$W$  is a vector-valued Wiener process,  $f$  and  $\sigma$  denote vector-valued drift and matrix-valued diffusion coefficients, respectively, and  $Z_{0}$  is a multivariate random variable that follows some distribution  $\mathcal{P}$ . To ensure the existence of a unique, strong solution to Eq. (2), we require that  $f, \sigma$  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  $\theta$  and  $\phi$ , 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  $\mathrm{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_{t_1}^i,\ldots ,y_{t_m}^i$  (e.g., a class-label or some real-valued quantity per time point). In these scenarios, we will additionally learn to map latent states  $\mathbf{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  $\mathbb{M} \subset \mathbb{R}^n$ . If  $\mathbb{M}$  admits a transitive action by a Lie group  $\mathcal{G}$ , i.e., if it is a homogeneous  $\mathcal{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  $\mathbb{M}$  of the form  $Z_{t} = G_{t} \cdot Z_{0}$  for some stochastic process  $G$  in a matrix Lie group  $\mathcal{G} \subset \mathrm{GL}(n)$  that acts by matrix-vector multiplication.
We define this stochastic process  $G$  as solution to the linear SDE
$$
\mathrm {d} G _ {t} = \left(\mathbf {V} _ {0} (t) \mathrm {d} t + \sum_ {i = 1} ^ {m} \mathrm {d} w _ {t} ^ {i} \mathbf {V} _ {i}\right) G _ {t}, \quad G _ {0} = \mathbf {I} _ {\mathrm {n}}. \tag {3}
$$
Herein,  $w^i$  are independent scalar Wiener processes,  $\mathbf{V}_1, \ldots, \mathbf{V}_m \in \mathfrak{g}$  are fixed (matrix-valued) tangent vectors in the Lie algebra  $\mathfrak{g}$ , and  $\mathbf{I}_n$  is the  $n \times n$  identity matrix. The drift function  $\mathbf{V}_0 = \mathbf{K} + \mathbf{K}_{\perp} : [0, T] \to \mathbb{R}^{n \times n}$  consists of a tangential (g-valued) component  $\mathbf{K}$  plus an orthogonal component  $\mathbf{K}_{\perp}$  called pinning drift [36; 53]. It offsets the stochastic drift away from the Lie group  $\mathcal{G}$  due to the quadratic variation of the Wiener processes  $w^i$  and only depends on the noise components  $\mathbf{V}_1, \ldots, \mathbf{V}_m$  (and the geometry of  $\mathcal{G}$ ). In case the Lie group  $\mathcal{G}$  is quadratic, i.e., defined by some fixed matrix  $\mathbf{P}$  via the equation  $\mathbf{G}^\top \mathbf{P}\mathbf{G} = \mathbf{P}$ , then the pinning drift  $\mathbf{K}_{\perp}$  can be deduced directly from the condition  $\mathrm{d}(\mathbf{G}^\top\mathbf{P}\mathbf{G}) = 0$  which implies that
$$
\mathbf {V} _ {0} (t) = \mathbf {K} (t) + \frac {1}{2} \sum_ {i = 1} ^ {m} \mathbf {V} _ {i} ^ {2}, \quad \text {a n d} \quad \mathbf {V} _ {i} ^ {\top} \mathbf {P} + \mathbf {P} \mathbf {V} _ {i} = \mathbf {0}. \tag {4}
$$
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 \sim \mathcal{P}$  drawn from some distribution  $\mathcal{P}$  on the manifold  $\mathbb{M}$ , the SDE for  $G$  induces an SDE for the one-point motion  $Z = G \cdot Z_0$  via  $\mathrm{d}Z_t = \mathrm{d}G_t \cdot Z_0$ . Therefore,  $Z$  solves
$$
\mathrm {d} Z _ {t} = \left(\mathbf {V} _ {0} (t) \mathrm {d} t + \sum_ {i = 1} ^ {m} \mathrm {d} w _ {t} ^ {i} \mathbf {V} _ {i}\right) Z _ {t}, \quad Z _ {0} \sim \mathcal {P}. \tag {5}
$$
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  $\Delta t$ , using the relation
$$
Z _ {t} = G _ {t} Z _ {0} = \left(G _ {t} G _ {t _ {j}} ^ {- 1}\right) Z _ {t _ {j}}. \tag {6}
$$
Notably,  $G_{t}G_{t_{j}}^{-1}$  follows the same SDE as  $G_{t}$ , see Eq. (3), but with initial value  $\mathbf{I}_{\mathrm{n}}$  at  $t = t_{j}$  instead of at  $t = 0$ . Moreover, if  $\Delta t$  is sufficiently small, then  $G_{t}G_{t_{j}}^{-1}$  stays close to the identity and reparametrizing  $G_{t}G_{t_{j}}^{-1} = \exp (\Omega_{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,  $\Omega_{t_{j + 1}} = \mathbf{V}_0(t_j)\Delta t + \sum_{i = 1}^m\Delta w_j^i\mathbf{V}_i$ , where  $\Delta w_j^i\sim \mathcal{N}(0,\Delta 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_{\theta}(\mathbf{z})$  and approximate posterior distribution  $q_{\phi}(\mathbf{z}|\mathbf{x})$  of latent paths  $\mathbf{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  $\mathbf{V}_i$ ,  $i > 0$ , then the KL divergence  $D_{\mathrm{KL}}(q_{\phi}(\mathbf{z}|\mathbf{x})||p_{\theta}(\mathbf{z}))$  in the ELBO is finite and can be computed via Girsanov's theorem.
Specifically, if  $\mathbf{V}_0^{\mathrm{prior}}$  and  $\mathbf{V}_0^{\mathrm{post}}$  denote the respective drift functions and  $\mathbf{V}_0^\Delta = \mathbf{V}_0^{\mathrm{prior}} - \mathbf{V}_0^{\mathrm{post}}$ , then
$$
D _ {\mathrm {K L}} \left(q _ {\phi} (\mathbf {z} | \mathbf {x}) \| p _ {\boldsymbol {\theta}} (\mathbf {z})\right) = \frac {1}{2} \int_ {0} ^ {T} \int_ {\mathbb {M}} q _ {Z _ {t}} (\mathbf {z}) \left[ \mathbf {V} _ {0} ^ {\Delta} (t) \mathbf {z} \right] ^ {\top} \boldsymbol {\Sigma} ^ {+} (\mathbf {z}) \left[ \mathbf {V} _ {0} ^ {\Delta} (t) \mathbf {z} \right] \mathrm {d} \mathbf {z} \mathrm {d} t, \tag {7}
$$
where  $q_{Z_t}$  is the marginal distribution of the approximate posterior process  $Z$  at time  $t$  and  $\pmb{\Sigma}^{+}(\mathbf{z})$  is the pseudo inverse of  $\pmb{\Sigma}(\mathbf{z}) = \sum_{i=1}^{m} \mathbf{V}_i \mathbf{z} \mathbf{z}^{\top} \mathbf{V}_i^{\top}$ . 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  $\mathbb{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  $\mathrm{SO}(n)$  of rotations. For modeling time series data,  $\mathbb{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  $\mathcal{U}_{\mathbb{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  $\mathrm{SO}(n)$  is the identity component of  $\mathbf{O}(n)$ , i.e., the quadratic Lie group defined by  $\mathbf{P} = \mathbf{I}_n$ . Its Lie algebra  $\mathfrak{so}(n) = \{\mathbf{A}:\mathbf{A} + \mathbf{A}^\top = \mathbf{0}\}$  consists of skew-symmetric  $n\times n$  matrices, and has a basis  $\{\mathbf{E}_{kl}:1\leq k < l\leq n\}$  with  $\mathbf{E}_{kl} = \mathbf{e}_k\mathbf{e}_l^\top -\mathbf{e}_l\mathbf{e}_k^\top$ , where  $\mathbf{e}_k$  is the  $k$ -th standard basis vector of  $\mathbb{R}^n$ . Importantly, if
$$
\alpha^ {\boldsymbol {\theta}} \neq 0, \quad \left\{\mathbf {V} _ {i} \right\} _ {i = 1} ^ {n (n - 1) / 2} = \left\{\alpha^ {\boldsymbol {\theta}} \mathbf {E} _ {k l} \right\} _ {1 \leq k <   l \leq n} \quad \text {a n d} \quad \mathbf {V} _ {0} (t) = - (\alpha^ {\boldsymbol {\theta}}) ^ {2} \frac {n - 1}{2} \mathbf {I} _ {n},
$$
then Eq. (4) is fulfilled and the SDE in Eq. (5) becomes, after some transformations,
$$
\mathrm {d} Z _ {t} = - \left(\alpha^ {\boldsymbol {\theta}}\right) ^ {2} \frac {n - 1}{2} Z _ {t} \mathrm {d} t + \alpha^ {\boldsymbol {\theta}} \left(\mathbf {I} _ {\mathrm {n}} - Z _ {t} Z _ {t} ^ {\top}\right) \mathrm {d} W _ {t}, \quad Z _ {0} \sim p _ {\boldsymbol {\theta}} \left(Z _ {0}\right) = \mathcal {U} _ {\mathbb {S} ^ {n - 1}}, \tag {8}
$$
whose solution is a rotation symmetric and driftless stochastic process on the unit  $n$ -sphere, i.e., the perfect uninformative prior process in our setting.
In fact, in case of  $\alpha^{\theta} = 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  $\mathcal{P}_0^\phi$  on the initial state  $Z_0$ , i.e.,
$$
\mathrm {d} Z _ {t} ^ {\phi} = \left(\mathbf {K} ^ {\phi} (\mathbf {x}) (t) - \left(\alpha^ {\phi}\right) ^ {2} \frac {n - 1}{2} \mathbf {I} _ {\mathrm {n}}\right) Z _ {t} \mathrm {d} t + \alpha^ {\phi} \left(\mathbf {I} _ {\mathrm {n}} - Z _ {t} Z _ {t} ^ {\top}\right) \mathrm {d} W _ {t}, \quad Z _ {0} \sim \mathcal {P} _ {0} ^ {\phi} \left(Z _ {0} | \mathbf {x}\right). \tag {9}
$$
Exemplary sample paths from solutions of Eq. (8) and Eq. (9) with constant drift are shown in Figure 1. We want to highlight that, in contrast to the parametric function  $\mathbf{K}^{\phi}(\mathbf{x}):[0,T]\to \mathfrak{so}(n)$  (as part of the drift coefficient) and the parametric distribution  $\mathcal{P}_0^\phi$ , the diffusion term does not depend on the evidence  $\mathbf{x}$  but only on  $\phi$ . This allows selecting the same diffusion components  $\alpha^{\phi} = \alpha^{\theta}$  in the posterior and prior process to ensure a finite KL divergence.
The latter has a particularly simple form (cf. Appendix E.5)
$$
D _ {\mathrm {K L}} \left(q _ {\phi} (\mathbf {z} | \mathbf {x}) \| p _ {\theta} (\mathbf {z})\right) = D _ {\mathrm {K L}} \left(\mathcal {P} _ {0} ^ {\phi} \left(Z _ {0} | \mathbf {x}\right) \| \mathcal {U} _ {\mathbb {S} ^ {n - 1}}\right) + \frac {1}{2 \left(\alpha^ {\theta}\right) ^ {2}} \int_ {0} ^ {T} \int_ {\mathbb {S} ^ {n - 1}} q _ {Z _ {t}} (\mathbf {z}) \left\| \mathbf {K} ^ {\phi} (\mathbf {x}) (t) \mathbf {z} \right\| ^ {2} \mathrm {d} \mathbf {z} \mathrm {d} t. \tag {10}
$$
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
$$
\int_ {\mathbb {S} ^ {n - 1}} q _ {Z _ {t}} (\mathbf {z}) \left\| \mathbf {K} ^ {\phi} (\mathbf {x}) (t) \mathbf {z} \right\| ^ {2} \mathrm {d} \mathbf {z} \leq \left\| \mathbf {K} ^ {\phi} (\mathbf {x}) (t) \right\| _ {F}
$$
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.
<sup>1</sup>The drift only consists of the pinning component that is orthogonal to the Lie algebra  $\mathfrak{so}(n)$ .
3.4 Parameterizing the approximate posterior