Mean Estimation on the Diagonal of Product Manifolds

Abstract

Computing sample means on Riemannian manifolds is typically computationally costly, as exemplified by computation of the Fréchet mean, which often requires finding minimizing geodesics to each data point for each step of an iterative optimization scheme. When closed-form expressions for geodesics are not available, this leads to a nested optimization problem that is costly to solve. The implied computational cost impacts applications in both geometric statistics and in geometric deep learning. The weighted diffusion mean offers an alternative to the weighted Fréchet mean. We show how the diffusion mean and the weighted diffusion mean can be estimated with a stochastic simulation scheme that does not require nested optimization. We achieve this by conditioning a Brownian motion in a product manifold to hit the diagonal at a predetermined time. We develop the theoretical foundation for the sampling-based mean estimation, we develop two simulation schemes, and we demonstrate the applicability of the method with examples of sampled means on two manifolds.

1. Introduction

The Euclidean expected value can be generalized to geometric spaces in several ways. Fréchet [1] generalized the notion of mean values to arbitrary metric spaces as minimizers of the sum of squared distances. Fréchet’s notion of mean values therefore naturally includes means on Riemannian manifolds. On manifolds without metric, for example, affine connection spaces, a notion of the mean can be defined by exponential barycenters, see, e.g., [2,3]. Recently, Hansen et al. [4,5] introduced a probabilistic notion of a mean, the diffusion mean. The diffusion mean is defined as the most likely starting point of a Brownian motion given the observed data. The variance of the data is here modelled in the evaluation time $T>0$ of the Brownian motion, and Varadhan’s asymptotic formula relating the heat kernel with the Riemannian distance relates the diffusion mean and the Fréchet mean in the $T\to 0$ limit.

Computing sample estimators of geometric means is often difficult in practice. For example, estimating the Fréchet mean often requires evaluating the distance to each sample point at each step of an iterative optimization to find the optimal value. When closed-form solutions of geodesics are not available, the distances are themselves evaluated by minimizing over curves ending at the data points, thus leading to a nested optimization problem. This is generally a challenge in geometric statistics, the statistical analysis of geometric data. However, it can pose an even greater challenge in geometric deep learning, where a weighted version of the Fréchet mean is used to define a generalization of the Euclidean convolution taking values in a manifold [6]. As the mean appears in each layer of the network, closed-form geodesics is in practice required for its evaluation to be sufficiently efficient.

As an alternative to the weighted Fréchet mean, Ref. [7] introduced a corresponding weighted version of the diffusion mean. Estimating the diffusion mean usually requires ability to evaluate the heat kernel making it often similarly computational difficult to estimate. However, Ref. [7] also sketched a simulation-based approach for estimating the (weighted) diffusion mean that avoids numerical optimization and estimation of the heat kernel. Here, a mean candidate is generated by simulating a single forward pass of a Brownian motion on a product manifold conditioned to hit the diagonal of the product space. The idea is sketched for samples in ${\mathbb{R}}^2$ in Figure 1.

Contribution

In this paper, we present a comprehensive investigation of the simulation-based mean sampling approach. We provide the necessary theoretical background and results for the construction, we present two separate simulation schemes, and we demonstrate how the schemes can be used to compute means on high-dimensional manifolds.

2. Background

We here outline the necessary concepts from Riemannian geometry, geometric statistics, stochastic analysis, and bridge sampling necessary for the sampling schemes presented later in the paper.

2.1. Riemannian Geometry

A Riemannian metric g on a d-dimensional differentiable manifold M is a family of inner products ${\left(g_p\right)}_{p\in M}$ on each tangent space $T_{p}M$ varying smoothly in p. The Riemannian metric allows for geometric definitions of, e.g., length of curves, angles of intersections, and volumes on manifolds. A differentiable curve on M is a map $\gamma :[0,1]\to M$ for which the time derivative ${\gamma }^{\prime }\left(t\right)$ belongs to $T_{{\gamma }_t}M$, for each $t\in (0,1)$. The length of the differentiable curve can then be determined from the Riemannian metric by $L\left(\gamma \right):={\int }_0^1\sqrt{g_{{\gamma }_t}\left({\gamma }^{\prime }\left(t\right),{\gamma }^{\prime }\left(t\right)\right)}dt={\int }_0^1{\parallel {\gamma }^{\prime }\left(t\right)\parallel }_{{\gamma }_t}dt$. Let $p,q\in M$ and let $\Gamma$ be the set of differentiable curves joining p and q, i.e., $\Gamma =\{\gamma :[0,1]\to M|\gamma \left(0\right)=p$ and $\gamma \left(1\right)=q\}$. The (Riemannian) distance between p and q is defined as $d(p,q)={min}_{\gamma \in \Gamma }L\left(\gamma \right)$. Minimizing curves are called geodesics.

A manifold can be parameterized using coordinate charts. The charts consist of open subsets of M providing a global cover of M such that each subset is diffeomorphic to an open subset of ${\mathbb{R}}^d$, or, equivalently, ${\mathbb{R}}^d$ itself. The exponential normal chart is often a convenient choice to parameterize a manifold for computational purposes. The exponential chart is related to the exponential map ${exp}_p:T_{p}M\to M$ that for each $p\in M$ is given by ${exp}_p\left(v\right)={\gamma }_v\left(1\right)$, where ${\gamma }_v$ is the unique geodesic satisfying ${\gamma }_v\left(0\right)=p$ and ${\gamma }_v^{\prime }\left(0\right)=v$. For each $p\in M$, the exponential map is a diffeomorphism from a star-shaped subset V centered at the origin of $T_{p}M$ to its image ${exp}_p\left(V\right)\subseteq M$, covering all M except for a subset of (Riemannian) volume measure zero, $Cut\left(p\right)$, the cut-locus of p. The inverse map ${log}_p:M\backslash Cut\left(p\right)\to T_{p}M$ provides a local parameterization of M due to the isomorphism between $T_{p}M$ and ${\mathbb{R}}^d$. For submanifolds $N\subseteq M$, the cut-locus $Cut\left(N\right)$ is defined in a fashion similar to $Cut\left(p\right)$, see e.g., [8].

Stochastic differential equations on manifolds are often conveniently expressed using the frame bundle $FM$, the fiber bundle which for each point $p\in M$ assigns a frame or basis for the tangent space $T_{p}M$, i.e., $FM$ consists of a collection of pairs $(p,u)$, where $u:{\mathbb{R}}^d\to T_{p}M$ is a linear isomorphism. We let $\pi$ denote the projection $\pi :FM\to M$. There exist a subbundle of $FM$ consisting of orthonormal frames called the orthonormal frame bundle $OM$. In this case, the map $u:{\mathbb{R}}^d\to T_{p}M$ is a linear isometry.

2.2. Weighted Fréchet Mean

The Euclidean mean has three defining properties: The algebraic property states the uniqueness of the arithmetic mean as the mean with residuals summing to zero, the geometric property defines the mean as the point that minimizes the variance, and the probabilistic property adheres to a maximum likelihood principle given an i.i.d. assumption on the observations (see also [9] Chapter 2). Direct generalization of the arithmetic mean to non-linear spaces is not possible due to the lack of vector space structure. However, the properties above allow giving candidate definitions of mean values in non-linear spaces.

The Fréchet mean [1] uses the geometric property by generalizing the mean-squared distance minimization property to general metric spaces. Given a random variable X on a metric space $(E,d)$, the Fréchet mean is defined by

$$\mu =\underset{p\in E}{arg\; min}\mathbb{E}\left[d{(p,X)}^2\right].$$

(1)

In the present context, the metric space is a Riemannian manifold M with Riemannian distance function d. Given realizations $x_1,\cdots ,x_n\in M$ from a distribution on M, the estimator of the weighted Fréchet mean is defined as

$$\widehat{\mu }=\underset{p\in M}{arg\; min}\sum _{i=1}^n{w}_{i}d{(p,x_i)}^2,$$

(2)

where $w_1,\cdots ,w_n$ are the corresponding weights satisfying $w_i>0$ and ${\sum }_i{w}_i=1$. When the weights are identical, (2) is an estimator of the Fréchet mean. Throughout, we shall make no distinction between the estimator and the Fréchet mean and will refer to both as the Fréchet mean.

In [6,10], the weighted Fréchet mean was used to define a generalization of the Euclidean convolution to manifold-valued inputs. When closed-form solutions of geodesics are available, the weighted Fréchet mean can be estimated efficiently with a recursive algorithm, also denoted an inductive estimator [6].

2.3. Weighted Diffusion Mean

The diffusion mean [4,5] was introduced as a geometric mean satisfying the probabilistic property of the Euclidean expected value, specifically as the starting point of a Brownian motion that is most likely given observed data. This results in the diffusion t-mean definition

$${\mu }_t=\underset{p\in M}{arg\; min}\mathbb{E}\left[-log{p}_t(p,X)\right],$$

(3)

where $p_t(·,·)$ denotes the transition density of a Brownian motion on M. Equivalently, $p_t$ denotes the solution to the heat equation $\partial u/\partial t=\frac{1}{2}\Delta u$, where $\Delta$ denotes the Laplace-Beltrami operator associated with the Riemannian metric. The definition allows for an interpretation of the mean as an extension of the Fréchet mean due to Varadhan’s result stating that ${lim}_{t\to 0}-2tlog{p}_t(x,y)=d{(x,y)}^2$ uniformly on compact sets disjoint from the cut-locus of either x or y [11].

Just as the Fréchet mean, the diffusion mean has a weighted version, and the corresponding estimator of the weighted diffusion t-mean is given as

$${\widehat{\mu }}_t=\underset{p\in M}{arg\; min}\sum _{i=1}^n-log{p}_{t/w_i}(p,x_i).$$

(4)

Please note that the evaluation time is here scaled by the weights. This is equivalent to scaling the variance of the steps of the Brownian motion [12].

As closed-form expressions for the heat kernel are only available on specific manifolds, evaluating the diffusion t-mean often rely on numerical methods. One example of this is using bridge sampling to numerically estimate the transition density [9,13]. If a global coordinate chart is available, the transition density can be written in the form (see [14,15])

$$p_T(z,v)=\sqrt{\frac{detg\left(v\right)}{{\left(2\pi T\right)}^2}}{e}^{-\frac{{\parallel a\left(z\right)(z-v)\parallel }^2}{2T}}\mathbb{E}\left[\phi \right],$$

(5)

where g is the metric matrix, a a square root of g, and $\phi$ denotes the correction factor between the law of the true diffusion bridge and the law of the simulation scheme. The expectation over the correction factor can be numerically approximated using Monte Carlo sampling. The correction factor will appear again when we discuss guided bridge proposals below.

2.4. Diffusion Bridges

The proposed sampling scheme for the (weighted) diffusion mean builds on simulation methods for conditioned diffusion processes, diffusion bridges. Here, we outline ways to simulate conditioned diffusion processes numerically in both the Euclidean and manifold context.

Euclidean Diffusion Bridges

Let $(\Omega ,\mathcal{F},{\mathcal{F}}_t,\mathbb{P})$ be a filtered probability space, and X a d-dimensional Euclidean diffusion $[0,T]$ satisfying the stochastic differential equation (SDE)

$$d{X}_t=b_t\left(X_t\right)dt+{\sigma }_t\left(X_t\right)d{W}_t,X_0=x,$$

(6)

where W is a d-dimensional Brownian motion. Let $v\in {\mathbb{R}}^d$ be a fixed point. Conditioning X on reaching v at a fixed time $T>0$ gives the bridge process $X|X_T=v$. Denoting this process Y, Doob’s h-transform shows that Y is a solution of the SDE (see e.g., [16])

$$\begin{array}{cc}\hfill d{Y}_t& ={\tilde{b}}_t\left(Y_t\right)dt+{\sigma }_t\left(Y_t\right)d{\tilde{W}}_t,Y_0=x\hfill \\ \hfill {\tilde{b}}_t\left(y\right)& =b_t\left(y\right)+a_t\left(y\right){\nabla }_{y}log{p}_{T-t}(y,v),\hfill \end{array}$$

(7)

where $p_t(·,·)$ denotes the transition density of the diffusion X, $a=\sigma {\sigma }^T$, and where $\tilde{W}$ is a d-dimensional Brownian motion under a changed probability law. From a numerical viewpoint, in most cases, the transition density is intractable and therefore direct simulation of (7) is not possible.

If we instead consider a Girsanov transformation of measures to obtain the system (see, e.g., [17] Theorem 1)

$$\begin{array}{cc}\hfill d{Y}_t& ={\tilde{b}}_t\left(Y_t\right)dt+{\sigma }_t\left(Y_t\right)d{\tilde{W}}_t,Y_0=x\hfill \\ \hfill {\tilde{b}}_t\left(y\right)& =b_t\left(y\right)+{\sigma }_t\left(y\right)h(t,y),\hfill \end{array}$$

(8)

the corresponding change of measure is given by

$$\frac{d{\mathbb{P}}^h}{d\mathbb{P}}{|}_{{\mathcal{F}}_t}=e^{{\int }_0^{t}h{(s,X_s)}^{T}d{W}_s-\frac{1}{2}{\int }_0^t{\parallel h(s,X_s)\parallel }^{2}ds}.$$

(9)

From (7), it is evident that $h(t,x)={\sigma }^T{\nabla }_{x}log{p}_{T-t}(x,v)$ gives the diffusion bridge. However, different choices of the function h can yield processes which are absolutely continuous regarding the actual bridges, but which can be simulated directly.

Delyon and Hu [17] suggested to use $h(t,x)={\sigma }_t^{-1}\left(x\right){\nabla }_{x}log{q}_{T-t}(x,v)$, where q denotes the transition density of a standard Brownian motion with mean v, i.e., $q_t(x,v)={\left(2\pi t\right)}^{-d/2}exp(-{\parallel x-v\parallel }^2/2t)$. They furthermore proposed a method that would disregard the drift term b, i.e., $h(t,x))={\sigma }_t^{-1}\left(x\right){\nabla }_{x}log{q}_{T-t}(x,v)-{\sigma }_t^{-1}\left(x\right)b_t\left(x\right)$. Under certain regularity assumptions on b and $\sigma$, the resulting processes converge to the target in the sense that ${lim}_{t\to T}{Y}_t=v$ a.s. In addition, for bounded continuous functions f, the conditional expectation is given by

$$\mathbb{E}\left[f\left(X\right)|X_T=v\right]=C\mathbb{E}\left[f\left(Y\right)\phi \left(y\right)\right],$$

(10)

where $\phi$ is a functional of the whole path Y on $[0,T]$ that can be computed directly. From the construction of the h-function, it can be seen that the missing drift term is accounted for in the correction factor $\phi$.

The simulation approach of [17] can be improved by the simulation scheme introduced by Schauer et al. [18]. Here, an h-function defined by $h(t,x)={\nabla }_{x}log{\widehat{p}}_{T-t}(x,v)$ is suggested, where $\widehat{p}$ denotes the transition density of an auxiliary diffusion process with known transition densities. The auxiliary process can for example be linear because closed-form solutions of transition densities for linear processes are available. Under the appropriate measure ${\mathbb{P}}^h$, the guided proposal process is a solution to

$$d{Y}_t=b_t\left(Y_t\right)dt+a_t\left(Y_t\right){\nabla }_{x}log{\widehat{p}}_{T-t}(x,v){|}_{x=Y_t}dt+{\sigma }_t\left(Y_t\right)d{W}_t.$$

(11)

Note the factor $a(t,y)$ in the drift in (7) which is also present in (11) but not with the scheme proposed by [17]. Moreover, the choice of a linear process grants freedom to model. For other choices of an h-functions see e.g., [19,20].

Marchand [19] extended the ideas of Delyon and Hu by conditioning a diffusion process on partial observations at a finite collection of deterministic times. Where Delyon and Hu considered the guided diffusion processes satisfying the SDE

$$d{Y}_t=b_t\left(Y_t\right)dt-\frac{Y_t-v}{T-t}dt+{\sigma }_t\left(Y_t\right)d{w}_t,$$

(12)

for $v\in {\mathbb{R}}^d$ over the time interval $[0,T]$, Marchand proposed the guided diffusion process conditioned on partial observations $v_1,\cdots ,v_N$ solving the SDE

$$d{Y}_t=b_t\left(Y_t\right)dt-\sum _{k=1}^n{P}_t^k\left(Y_t\right)\frac{Y_t-u_k}{T_k-t}{1}_{(T_k-{\epsilon }_k,T_k)}dt+{\sigma }_t\left(Y_t\right)d{w}_t,$$

(13)

where $u_k$ is be any vector satisfying $L_k\left(x\right)u_k=v_k$, $L_k$ a deterministic matrix in $M_{m_k,n}\left(\mathbb{R}\right)$ whose $m_k$ rows form a orthonormal family, $P_t^k$ are projections to the range of $L_k$, and $T_k-{\epsilon }_k<T_k$. The ${\epsilon }_k$ only allow the application of the guiding term on a part of the time intervals $[T_{k-1},T_k]$. We will only consider the case $k=1$. The scheme allows the sampling of bridges conditioned on $L{Y}_T=v$.

2.5. Manifold Diffusion Processes

To work with diffusion bridges and guided proposals on manifolds, we will first need to consider the Eells–Elworthy–Malliavin construction of Brownian motion and the connected characterization of semimartingales [21]. Endowing the frame bundle $FM$ with a connection allows splitting the tangent bundle $TFM$ into a horizontal and vertical part. If the connection on $FM$ is a lift of a connection on M, e.g., the Levi–Civita connection of a metric on M, the horizontal part of the frame bundle is in one-to-one correspondence with M. In addition, there exist fundamental horizontal vector fields $H_i:FM\to HFM$ such that for any continuous ${\mathbb{R}}^d$-valued semimartingale Z the process U defined by

$$d{U}_t=H_i\left(U_t\right)\circ d{Z}_t^i,$$

(14)

is a horizontal frame bundle semimartingale, where ∘ denotes integration in the Stratonovich sense. The process $X_t:=\pi \left(U_t\right)$ is then a semimartingale on M. Any semimartingale $X_t$ on M has this relation to a Euclidean semimartingale $Z_t$. $X_t$ is denoted the development of $Z_t$, and $Z_t$ the antidevelopment of $X_t$. We will use this relation when working with bridges on manifolds below.

When $Z_t$ is a Euclidean Brownian motion, the development $X_t$ is a Brownian motion. We can in this case also consider coordinate representations of the process. With an atlas ${\left\{(D_{\alpha },{\varphi }_{\alpha })\right\}}_{\alpha }$ of M, there exists an increasing sequence of predictable stopping times $0\le T_k\le T_{k+1}$ such that on each stochastic interval $〚T_k,T_{k+1}〛=\{(\omega ,t)\in \Omega \times {\mathbb{R}}_{+}|T_k\left(\omega \right)\le t\le T_{k+1}\left(\omega \right)\}$ the process $x_t\in D_{\alpha }$, for some $\alpha$ (see [22] Lemma 3.5). Thus, the Brownian motion x on M can be described locally in a chart $D_{\alpha }\subset M$ as the solution to the system of SDEs, for $(\omega ,t)\in 〚T_k,T_{k+1}〛\cap \{T_k<T_{k+1}\}$

$$d{x}_t^i\left(\omega \right)=b^i\left(x_t\left(\omega \right)\right)dt+{\sigma }_j^i\left(x_t\left(\omega \right)\right)d{W}_t^j\left(\omega \right),$$

(15)

where $\sigma$ denotes the matrix square root of the inverse of the Riemannian metric tensor $\left(g^{ij}\right)$ and $b^k\left(x\right)=-\frac{1}{2}{g}^{ij}\left(x\right){\Gamma }_{ij}^k\left(x\right)$ is the contraction over the Christoffel symbols (see, e.g., [11] Chapter 3). Strictly speaking, the solution of Equation (15) is defined by $x_t^i={\varphi }_{\alpha }{\left(x_t\right)}^i$.

We thus have two concrete SDEs for the Brownian motion in play: The $FM$ SDE (14) and the coordinate SDE (15).

Throughout the paper, we assume that M is stochastically complete, i.e., the Brownian motions does not explode in finite time and, consequently, ${\int }_M{p}_t(x,y)d{Vol}_M\left(y\right)=1$, for all $t>0$ and all $x\in M$.

2.6. Manifold Bridges

The Brownian bridge process Y on M conditioned at $Y_T=v$ is a Markov process with generator $\frac{1}{2}\Delta +\nabla log{p}_{T-t}(·,v)$. Closed-form expressions of the transition density of a Brownian motion are available on selected manifolds including Euclidean spaces, hyperbolic spaces, and hyperspheres. Direct simulation of Brownian bridges is therefore possible in these cases. However, generally, transition densities are intractable and auxiliary processes are needed to sample from the desired conditional distributions.

To this extent, various types of bridge processes on Riemannian manifolds have been described in the literature. In the case of manifolds with a pole, i.e., the existence of a point $p\in M$ such that the exponential map ${exp}_p:T_{p}M\to M$ is a diffeomorphism, the semi-classical (Riemannian Brownian) bridge was introduced by Elworthy and Truman [23] as the process with generator $\frac{1}{2}\Delta +\nabla log{k}_{T-t}(·,v)$, where

$$k_t(x,v)={\left(2\pi t\right)}^{-n/2}{e}^{-\frac{d{(x,v)}^2}{2t}}{J}^{-1/2}\left(x\right),$$

and $J\left(x\right)=|det{D}_{exp{\left(v\right)}^{-1}}{exp}_v|$ denotes the Jacobian determinant of the exponential map at v. Elworthy and Truman used the semi-classical bridge to obtain heat-kernel estimates, and the semi-classical bridge has been studied by various others [24,25].

By Varadhan’s result (see [11] Theorem 5.2.1), as $t\to T$, we have the asymptotic relation $((T-t)log{p}_{T-t}(x,y)\sim -\frac{1}{2}d{(x,y)}^2$. In particular, the following asymptotic relation was shown to hold by Malliavin, Stroock, and Turetsky [26,27]: $(T-t)\nabla log{p}_{T-t}(x,y)\sim -\frac{1}{2}\nabla d{(x,y)}^2$. From these results, the generators of the Brownian bridge and the semi-classical bridge differ in the limit by a factor of $-\frac{1}{2}\nabla logJ\left(x\right)$. However, under a certain boundedness condition, the two processes can be shown to be identical under a changed probability measure [8] Theorem 4.3.1.

To generalize the heat-kernel estimates of Elworthy and Truman, Thompson [8,28] considered the Fermi bridge process conditioned to arrive in a submanifold $N\subseteq M$ at time $T>0$. The Fermi bridge is defined as the diffusion process with generator $\frac{1}{2}\Delta +\nabla log{q}_{T-t}(·,N)$, where

$$q_t(x,N)={\left(2\pi t\right)}^{-n/2}{e}^{-\frac{d{(x,N)}^2}{2t}}.$$

For both bridge processes, when $M={\mathbb{R}}^d$ and N is a point, both the semi-classical bridge and the Fermi bridge agree with the Euclidean Brownian bridge.

Ref. [15] introduce a numerical simulation scheme for conditioned diffusions on Riemannian manifolds, which generalize the method by Delyon and Hu [17]. The guiding term used is identical to the guiding term of the Fermi bridge when the submanifold is a single point v.

3. Diffusion Mean Estimation

The standard setup for diffusion mean estimation described in the literature (e.g., [13]) is as follows: Given a set of observations $x_1,\cdots ,x_n\in M$, for each observation $x_i$, sample a guided bridge process approximating the bridge $X_{i,t}|X_{i,T}=x_i$ with starting point $x_0$. The expectation over the correction factors can be computed from the samples, and the transition density can be evaluated using (5). An iterative maximum likelihood approach using gradient descent to update $x_0$ yielding an approximation of the diffusion mean in the final value of $x_0$. The computation of the diffusion mean, in the sense just described, is, similarly to the Fréchet mean, computationally expensive.

We here explore the idea first put forth in [7]: We turn the situation around to simulate n independent Brownian motions starting at each of $x_1,\dots ,x_n$, and we condition the n processes to coincide at time T. We will show that the value $x_{1,T}=\cdots =x_{n,T}$ is an estimator of the diffusion mean. By introducing weights in the conditioning, we can similarly estimate the weighted diffusion mean. The construction can concisely be described as a single Brownian motion on the n-times product manifold $M^n$ conditioned to hit the diagonal $diag\left(M^n\right)=\left\{(x,\dots ,x)\right|x\in M\}\subset M^n$. To shorten notation, we denote the diagonal submanifold N below. We start with examples with M Euclidean to motivate the construction.

Example 1.

Consider the two-dimensional Euclidean multivariate normal distribution

$$\left(\begin{array}{c}X\\ Y\end{array}\right)\sim N\left(\left(\begin{array}{c}{\mu }_1\\ {\mu }_2\end{array}\right),\left(\begin{array}{c}{\sigma }_{11}{\sigma }_{12}\\ {\sigma }_{21}{\sigma }_{22}\end{array}\right)\right).$$

The conditional distribution of X given $Y=y$ follows a univariate normal distribution

$$X|Y=y\sim N\left({\mu }_1+{\sigma }_{12}{\sigma }_{22}^{-1}(y-{\mu }_2),{\sigma }_{11}-{\sigma }_{12}{\sigma }_{22}^{-1}{\sigma }_{21}\right).$$

This can be seen from the fact that if $X\sim N\left(\mu ,\Sigma \right)$ then for any linear transformation $AX+b\sim N\left(b+A\mu ,A\Sigma A^T\right)$. Defining the random variable $Z=X-{\sigma }_{12}{\sigma }_{22}^{-1}Y$, the result applied to $(Z,X)$ gives $Z\sim N\left({\mu }_1-{\sigma }_{12}{\sigma }_{22}^{-1}{\mu }_2,{\sigma }_{11}-{\sigma }_{12}{\sigma }_{22}^{-1}{\sigma }_{21}\right)$. The conclusion then follows from $X=Z+{\sigma }_{12}{\sigma }_{22}^{-1}Y$. Please note that X and Y are independent if and only if ${\sigma }_{12}={\sigma }_{21}=0$ and the conditioned random variable is in this case identical in law to X.

Let now $x_1,\cdots ,x_n\in M$ be observations and let $x=(x_1,\cdots ,x_n)\in M^n$ be an element of the n-product manifold $M\times \cdots \times M$ with the product Riemannian metric. We again first consider the case $M={\mathbb{R}}^d$:

Example 2.

Let $Y_i\sim N\left(x_i,\frac{T}{w_i}{I}_d\right)$ be independent random variables. The conditional distribution $Y_1|Y_1=\cdots =Y_n$ is normal $N\left(\frac{{\sum }_i{w}_i{x}_i}{{\sum }_i{w}_i},\frac{T}{{\sum }_i{w}_i}\right)$. This can be seen inductively: The conditioned random variable $Y_1|Y_1=Y_2$ is identical to $Y_1|Y_1-Y_2=0$. Now let $X=Y_1$ and $Y=Y_1-Y_2$ and refer to Example 1. To conclude, assume $Z_n:=Y_1|Y_1=\cdots ,=Y_{n-1}$ follows the desired normal distribution. Then $Z_n|Z_n=Y_n$ is normally distributed with the desired parameters and $Z_n|Z_n=Y_n$ is identical to $Y_1|Y_1=\cdots =Y_n$.

The following example establishes the weighted average as a projection onto the diagonal.

Example 3.

Let x be a point in ${\left({\mathbb{R}}^d\right)}^n$ and let P be the orthogonal projection to the diagonal of ${\left({\mathbb{R}}^d\right)}^n$ such that $Px={\left(\frac{1}{nd}{\sum }_{i=1}^{nd}{x}_i\cdots \frac{1}{nd}{\sum }_{i=1}^{nd}{x}_i\right)}^T$. We see that the projection yields n copies of the arithmetic mean of the coordinates. This is illustrated in Figure 2.

The idea of conditioning diffusion bridge processes on the diagonal of a product manifold originates from the facts established in Examples 1–3. We sample the mean by sampling from the conditional distribution $Y_1|Y_1=\cdots =Y_n$ from Example 2 using a guided proposal scheme on the product manifolds $M^n$ and on each step of the sampling projecting to the diagonal as in Example 3.

Turning now to the manifold situation, we replace the normal distributions with mean $x_i\in {\mathbb{R}}^d$ and variance $T/w_i$ with Brownian motions started at $x_i\in M$ and evaluated at time $T/w_i$. Please note that the Brownian motion density, the heat kernel, is symmetric in its coordinates: $p_t(x,y)=p_t(y,x)$. We will work with multiple process and indicate with superscript the density with respect to a particular process, e.g., $p_T^X$. Note also that change of the evaluation time T is equal to scaling the variance, i.e., $p_{\alpha T}^X(x,y)=p_T^{X^{\alpha }}(x,y)$ where $X^{\alpha }$ is a Brownian motion with variance of the increments scaled by $\alpha >0$. This gives the following theorem, first stated in [7] with sketch proof:

Theorem 1.

Let $X_t=(X_{1,t}^{w_1^{-1}},\cdots ,X_{n,t}^{w_n^{-1}})$ consist of n independent Brownian motions on M with variance $w_i^{-1}$ and $X_{i,0}=x_i$, and let ${\mathbb{P}}^{*}$ the law of the conditioned process $Y_t=X_t|X_T\in N$, $N=diag\left(M^n\right)$. Let v be the random variable $Y_{1,T}$. Then v has density $p_v^Y\left(y\right)\propto {\prod }_{i=1}^n{p}_{T/w_i}(x_i;y)$ and $v=Y_{i,T}$ for all i a.s. (almost surely).

Proof. 

$p_T^X((x_1,\cdots ,x_n),(y,\cdots ,y))={\prod }_{i=1}^n{p}_T^{X^{w_i^{-1}}}(x_i,y)$ because the processes $X_{i,t}$ are independent. By symmetry of the Brownian motion and the time rescaling property, $p_T^{X_i^{w_i^{-1}}}(x_i,y)$$=p_{T/w^i}(y,x_i)$. For elements $(y,\cdots ,y)\in diag\left(M^n\right)$ and $x\in M^n$, $p_v\left(y\right)=p_T^Y(x,y)\propto p_T^X(x,y)$. As a result of the conditioning, $v=Y_{1,T}=\cdots =Y_{n,T}$. In combination, this establishes the result.    □

Consequently, the set of modes of $p_v$ equal the set of the maximizers for the likelihood $L(y;x_1,\cdots ,x_n)={\prod }_{i=1}^n{p}_{T/w_i}(x_i;y)$ and thus the weighted diffusion mean. This result is the basis for the sampling scheme. Intuitively, if the distribution of v is relatively well behaved (e.g., close to normal), a sample from v will be close to a weighted diffusion mean with high probability.

In practice, however, we cannot sample $Y_t$ directly. Instead, we will below use guided proposal schemes resulting in processes ${\tilde{Y}}_t$ with law $\tilde{\mathbb{P}}$ that we can actually sample and that under certain assumptions, will be absolutely continuous with respect to $Y_t$ with explicitly computable likelihood ratio so that ${\mathbb{P}}^{*}=\frac{\phi \left({\tilde{Y}}_T\right)}{{\mathbb{E}}^{\tilde{\mathbb{P}}}\left[\phi \left({\tilde{Y}}_T\right)\right]}\tilde{\mathbb{P}}$.

Corollary 1.

Let $\tilde{\mathbb{P}}$ be the law of ${\tilde{Y}}_t$ and φ be the corresponding correction factor of the guiding scheme. Let $\tilde{v}$ be the random variable ${\tilde{Y}}_{1,T}$ with law $\frac{\phi \left({\tilde{Y}}_T\right)}{{\mathbb{E}}^{\tilde{\mathbb{P}}}\left[\phi \left({\tilde{Y}}_T\right)\right]}\tilde{\mathbb{P}}$. Then $\tilde{v}$ has density $p_{\tilde{v}}\left(y\right)\propto {\prod }_{i=1}^n{p}_{T/w_i}(x_i;y)$.

We now proceed to actually construct the guided sampling schemes.

3.1. Fermi Bridges to the Diagonal

Consider a Brownian motion $X_t=(X_{1,t},\cdots X_{n,t})$ in the product manifold $M^n$ conditioned on $X_{1,T}=\cdots =X_{n,T}$ or, equivalently, $X_T\in N$, $N=diag\left(M^n\right)$. Since N is a submanifold of $M^n$, the conditioned diffusion defined above is absolutely continuous with respect to the Fermi bridge on $[0,T)$[8,28]. Define the $FM$-valued horizontal guided process

$$d{U}_t=H_i\left(U_t\right)\circ \left(d{W}_t^i-\frac{H_i{\tilde{r}}_N^2\left(U_t\right)}{2(T-t)}dt\right),$$

(16)

where $\tilde{r}$ denotes the lift of the radial distance to N defined by ${\tilde{r}}_N\left(u\right):=r_N\left(\pi \left(u\right)\right)=d(\pi \left(u\right),N)$. The Fermi bridge $Y^F$ is the projection of U to M, i.e., $Y_t^F:=\pi \left(U_t\right)$. Let ${\mathbb{P}}^F$ denotes its law.

Theorem 2.

For all continuous bounded functions f on $M^n$, we have

$$\mathbb{E}\left[f\left(X\right)|X_{1,T}=\cdots =X_{n,T}\right]=\underset{t\uparrow T}{lim}C{\mathbb{E}}^{{\mathbb{P}}^F}\left[f\left(Y\right)\phi \left(Y\right)\right],$$

(17)

with a constant $C>0$, where

$$dlog\phi \left(Y_s^F\right)=\frac{r_N\left(Y_s^F\right)}{T-s}\left(d{\eta }_s+d{L}_s\right)\mathrm{with}d{\eta }_s=\frac{\partial }{\partial r_N}log{\Theta }_N^{-\frac{1}{2}}ds,$$

$d{L}_s:=d{\mathbb{L}}_s\left(Y^F\right)$ with $\mathbb{L}$ being the geometric local time at $Cut\left(N\right)$, and ${\Theta }_N$ the determinant of the derivative of the exponential map normal to N with support on $M^n\backslash Cut\left(N\right)$ [8].

Proof. 

From [15] Theorem 8 and [28],

$$\mathbb{E}\left[f\left(X\right)|X_T\in N\right]=\underset{t\uparrow T}{lim}C{\mathbb{E}}^{{\mathbb{P}}^F}\left[f\left(Y^F\right)\phi \left(Y_t^F\right)\right].$$

   □

Since N is a totally geodesic submanifold of dimension d, the results of [8] can be used to give sufficient conditions to extend the equivalence in (17) to the entire interval $[0,T]$. A set A is said to be polar for a process $X_t$ if the first hitting time of A by X is infinity a.s.

Corollary 2.

If either of the following conditions are satisfied

(i) 

the sectional curvature of planes containing the radial direction is non-negative or the Ricci curvature in the radial direction is non-negative;

(ii) 

$Cut\left(N\right)$ is polar for the Fermi bridge $Y^F$ and either the sectional curvature of planes containing the radial direction is non-positive or the Ricci curvature in the radial direction is non-positive;

then

$$\mathbb{E}\left[f\left(X\right)|X_{1,T}=\cdots =X_{n,T}\right]=C{\mathbb{E}}^{{\mathbb{P}}^F}\left[f\left(Y^F\right)\phi \left(Y_T^F\right)\right].$$

In particular, $\frac{\phi \left(Y_T^F\right)}{{\mathbb{E}}^{{\mathbb{P}}^F}\left[\phi \left(Y_T^F\right)\right]}d{\mathbb{P}}^F\propto d{\mathbb{P}}^{*}$.

Proof. 

See [8] (Appendix C.2).    □

For numerical purposes, the equivalence (17) in Theorem 2 is sufficient as the interval $[0,T]$ is finitely discretized. To obtain the result on the full interval, the conditions in Corollary 2 may at first seem quite restrictive. A sufficient condition for a subset of a manifold to be polar for a Brownian motion is its Hausdorff dimension being two less than the dimension of the manifold. Thus, $Cut\left(N\right)$ is polar if $dim(Cut(N\left)\right)\le nd-2$. Verifying whether this is true requires specific investigation of the geometry of $M^n$.

The SDE (16) together with (17) and the correction $\phi$ gives a concrete simulation scheme that can be implemented numerically. Implementation of the geometric constructs is discussed in Section 4. The main complication of using Fermi bridges for simulation is that it involves evaluation of the radial distance $r_N$ at each time-step of the integration. Since the radial distance finds the closest point on N to $x_1,\cdots ,x_n$, it is essentially a computation of the Fréchet mean and thus hardly more computationally efficient than computing the Fréchet mean itself. For this reason, we present a coordinate-based simulation scheme below.

3.2. Simulation in Coordinates

We here develop a more efficient simulation scheme focusing on manifolds that can be covered by a single chart. The scheme follows the partial observation scheme developed [19]. Representing the product process in coordinates and using a transformation L, whose kernel is the diagonal $diag\left(M^n\right)$, gives a guided bridge process converging to the diagonal. An explicit expression for the likelihood is given.

In the following, we assume that M can be covered by a chart in which the square root of the cometric tensor, denoted by $\sigma$, is $C^2$. Furthermore, $\sigma$ and its derivatives are bounded; $\sigma$ is invertible with bounded inverse. The drift b is locally Lipschitz and locally bounded.

Let $x_1,\cdots ,x_n\in M$ be observations and let $X_{1,t},\cdots ,X_{n,t}$ be independent Brownian motions with $X_{1,0}=x_1,\cdots ,X_{n,0}=x_n$. Using the coordinate SDE (15) for each $X_{i,t}$, we can write the entire system on $M^n$ as

$$d\left(\begin{array}{c}{X}_{1,t}^1\\ ⋮\\ X_{1,t}^d\\ ⋮\\ X_{n,t}^1\\ ⋮\\ X_{n,t}^d\end{array}\right)=\left(\begin{array}{c}{b}^1\left(X_{1,t}\right)\\ ⋮\\ b^d\left(X_{1,t}\right)\\ ⋮\\ b^1\left(X_{n,t}\right)\\ ⋮\\ b^d\left(X_{n,t}\right)\end{array}\right)dt+\left[\begin{array}{c}{\Sigma }^1(X_{1,t},\cdots ,X_{n,t})\\ ⋮\\ {\Sigma }^d(X_{1,t},\cdots ,X_{n,t})\end{array}\right]d{W}_t.$$

(18)

In the product chart, $\Sigma$ and b satisfy the same assumptions as the metric and cometric tensor and drift listed above.

The conditioning $X_T\in N$ is equivalent to the requiring $X_T\in diag\left({\left({\mathbb{R}}^d\right)}^n\right)$ in coordinates. $diag\left({\left({\mathbb{R}}^d\right)}^n\right)$ is a linear subspace of ${\left({\mathbb{R}}^d\right)}^n$, we let $L\in M^{d\times nd}$ be a matrix with orthonormal rows and $kerL=diag\left({\left({\mathbb{R}}^d\right)}^n\right)$ so that the desired conditioning reads $L{X}_T=0$. Define the following oblique projection, similar to [19],

$$P_t\left(x\right)=a\left(x\right)L^{T}A\left(x\right)L$$

(19)

where

$$a\left(x\right)=\Sigma \left(x\right)\Sigma {\left(x\right)}^T\mathrm{and}{A}_t\left(x\right)={\left(La\left(x\right)L^T\right)}^{-1}.$$

Set $\beta \left(x\right)=\Sigma {\left(x\right)}^T{L}^{T}A\left(x\right)$. The guiding scheme (13) then becomes

$$d{Y}_t=b\left(Y_t\right)dt+\Sigma \left(Y_t\right)d{W}_t-\Sigma \left(Y_t\right)\beta \left(Y_t\right)\frac{L{Y}_t}{T-t}{1}_{(T-\epsilon ,T)}\left(t\right)dt,Y_0=u.$$

(20)

We have the following result.

Lemma 1.

Equation (20) admits a unique solution on $[0,T)$. Moreover, ${\parallel L{Y}_t\parallel }^2\le C\left(\omega \right)(T-t)loglog[{(T-t)}^{-1}+e]$ a.s., where C is a positive random variable.

Proof. 

Since $LP=L$, the proof is similar to the proof of [19] Lemma 6.    □

With the same assumptions, we also obtain the following result similar to [19] Theorem 3.

Theorem 3.

Let $Y_t$ be a solution of (20), and assume the drift b is bounded. For any bounded function f,

$$\mathbb{E}\left[f\left(X\right)|X_T\in N\right]=C\mathbb{E}\left[f\left(y\right)\phi \left(Y\right)\right],$$

(21)

where C is a positive constant and

$$\begin{array}{cc}\hfill \phi \left(Y_t\right)& =\sqrt{det\left(A\left(Y_T\right)\right)}exp\{-\frac{{\parallel {\beta }_{T-\epsilon }\left(Y_{T-\epsilon }\right)L{Y}_{T-\epsilon }\parallel }^2}{2\epsilon }\hfill \\ \hfill -& {\int }_{T-\epsilon }^T\frac{2{\left(L{Y}_s\right)}^{T}Lb\left(Y_s\right)ds-{\left(L{Y}_s\right)}^{T}d\left(A\left(Y_s\right)\right)L{Y}_s+d[A{\left(Y_s\right)}^{ij},{\left(L{Y}_s\right)}_i{\left(L{Y}_s\right)}_j]}{2(T-s)}\}\hfill \end{array}$$

Proof. 

A direct consequence of [19] Theorem 3, for $k=1$, and Lemma 1.    □

The theorem can also be applied for unbounded drift by replacing b with a bounded approximation and performing a Girsanov change of measure.

3.3. Accounting for $\phi$

The sampling schemes (16) and (20) above provides samples on the diagonal and thus candidates for the diffusion mean estimates. However, the schemes sample from a distribution which is only equivalent to the bridge process distribution: We still need to handle the correction factor in the sampling to sample from the correct distribution, i.e., the rescaling $\frac{\phi }{\mathbb{E}\left[\phi \right]}$ of the guided law in Theorem 1. A simple way to achieve this is to do sampling importance resampling (SIR) as described in Algorithm 1. This yields an approximation of the weighted diffusion mean. For each sample $y^i$ of the guided bridge process, we compute the corresponding correction factor $\phi \left(y^i\right)$. The resampling step then consists of picking $y_T^j$ with a probability determined by the correction terms, i.e., with J the number of samples we pick sample j with probability $P_j=\frac{\phi \left(y_T^j\right)}{{\sum }_{i=1}^J\phi \left(y_T^i\right)}$.

Algorithm 1:weighted Diffusion Mean

It depends on the practical application if the resampling is necessary, or if direct samples from the guided process (corresponding to $J=1$) are sufficient.

4. Experiments

We here exemplify the mean sampling scheme on the two-sphere ${\mathbb{S}}^2$ and on finite sets of landmark configurations endowed with the LDDMM metric [29,30]. With the experiment on ${\mathbb{S}}^2$, we aim to give a visual intuition of the sampling scheme and the variation in the diffusion mean estimates caused by the sampling approach. In the higher-dimensional landmark example where closed-form solutions of geodesics are not available, we compare to the Fréchet mean and include rough running times of the algorithms to give a sense of the reduced time complexity. Note, however, that the actual running times are very dependent on the details of the numerical implementation, stopping criteria for the optimization algorithm for the Fréchet mean, etc.

The code used for the experiments is available in the software package Jax Geometry http://bitbucket.org/stefansommer/jaxgeometry (accessed on 4 February 2022). The implementation uses automatic differentiation libraries extensively for the geometry computations as is further described [31].

4.1. Mean Estimation on ${\mathbb{S}}^2$

To illustrate the diagonal sampling scheme, Figure 3 displays a sample from a diagonally conditioned Brownian motion on ${\left({\mathbb{S}}^2\right)}^n$, $n=3$. The figure shows both the diagonal sample (red point) and the product process starting at the three data points and ending at the diagonal. In Figure 4, we increase the number of samples to $n=256$ and sample 32 mean samples $(T=0.2)$. The population mean is the north pole, and the samples can be seen to cluster closely around the population mean with little variation in the mean samples.

4.2. LDDMM Landmarks

We here use the same setup as in [13], where the diffusion mean is estimated by iterative optimization, to exemplify the mean estimation on a high-dimensional manifold. The data consists of annotations of left ventricles cardiac MR images [32] with 17 tlandmarks selected from the annotation set from a total of 14 images. Each configuration of 17 landmarks in ${\mathbb{R}}^2$ gives a point in a 34-dimensional shape manifold. We equip this manifold with the LDDMM Riemannian metric [29,30]. Please note that the configurations can be represented as points in ${\mathbb{R}}^{34}$, and the entire shape manifold is the subset of ${\mathbb{R}}^{34}$ where no two landmarks coincide. This provides a convenient Euclidean representation of the landmarks. The cometric tensor is not bounded in this representation, and we therefore cannot directly apply the results of the previous sections. We can nevertheless explore the mean simulation scheme experimentally.

Figure 5 shows one landmark configuration overlayed the MR image from which the configuration was annotated, and all 14 landmark configurations plotted together. Figure 6 displays samples from the diagonal process for two values of the Brownian motion end time T. Please note that each landmark configuration is one point on the 34-dimensional shape manifold, and each of the paths displayed is therefore a visualization of a Brownian path on this manifold. This figure and Figure 3 both show diagonal processes, but on two different manifolds.

In Figure 7, an estimated diffusion mean and Fréchet mean for the landmark configurations are plotted together. On a standard laptop, generation of one sample diffusion mean takes approximately 1 s. For comparison, estimation of the Fréchet mean with the standard nested optimization approach using the Riemannian logarithm map as implemented in Jax Geometry takes approximately 4 min. The diffusion mean estimation performed in [13] using direct optimization of the likelihood approximation with bridge sampling from the mean candidate to each data point is comparable in complexity to the Fréchet mean computation.

5. Conclusions

In [7], the idea of sampling means by conditioning on the diagonal of product manifolds was first described and the bridge sampling construction sketched. In the present paper, we have provided a comprehensive account of the background for the idea, including the relation between the (weighted) Fréchet and diffusion means, and the foundations in both geometry and stochastic analysis. We have constructed two simulation schemes and demonstrated the method on both low and a high-dimensional manifolds, the sphere ${\mathbb{S}}^2$ and the LDDMM landmark manifold, respectively. The experiments show the feasibility of the method and indicate the potential high reduction in computation time compared to computing means with iterative optimization.