Score matching for sub-Riemannian bridge sampling
Erlend Grong Karen Habermann Stefan Sommer



    Simulation of conditioned diffusion processes is an essential tool in inference for stochastic processes, data imputation, generative modelling, and geometric statistics. Whilst simulating diffusion bridge processes is already difficult on Euclidean spaces, when considering diffusion processes on Riemannian manifolds the geometry brings in further complications. In even higher generality, advancing from Riemannian to sub-Riemannian geometries introduces hypoellipticity, and the possibility of finding appropriate explicit approximations for the score of the diffusion process is removed. We handle these challenges and construct a method for bridge simulation on sub-Riemannian manifolds by demonstrating how recent progress in machine learning can be modified to allow for training of score approximators on sub-Riemannian manifolds. Since gradients dependent on the horizontal distribution, we generalise the usual notion of denoising loss to work with non-holonomic frames using a stochastic Taylor expansion, and we demonstrate the resulting scheme both explicitly on the Heisenberg group and more generally using adapted coordinates. We perform numerical experiments exemplifying samples from the bridge process on the Heisenberg group and the concentration of this process for small time.


Supplementary Material for TalkingGaussian
    Jiahe Li1 Jiawei Zhang1 Xiao Bai1Corresponding author: Xiao Bai (baixiao@buaa.edu.cn). Jin Zheng1 Xin Ning2 Jun Zhou3 Lin Gu4,5

    Received xx; accepted xx
====================================================================================================================================





§ INTRODUCTION


The problem of simulating bridges for diffusion processes has been studied extensively in the literature, such as in <cit.>, and effective simulation schemes for diffusion bridge processes have become paramount for data imputation and in parameter inference for stochastic processes as they provide a stochastic method for approximating the intractable likelihood or when sampling from a posteriori distribution of parameters.
In the parameter estimation from given sets of time-discrete observations, bridge simulations feature in the data imputation performed when missing data is simulated through diffusion bridge processes that connect the observed data, see e.g. Golightly and Wilkinson <cit.>, Papaspiliopoulos, Roberts and Stramer <cit.>, and van der Meulen and Schauer <cit.>.

Bridge sampling further plays a prominent role in geometric statistics and has found applications in areas such as medical image analysis and shape analysis, see <cit.>. As discussed in <cit.>, simulated diffusion bridge processes have been used to approximate the heat kernel of a diffusion process on a Riemannian manifold and to estimate the underlying metric structure.

Most previous work on diffusion bridge simulations focuses on elliptic diffusion processes on Euclidean spaces where the diffusivity matrix is uniformly invertible. The simulations introduced by Bierkens, van der Meulen and Schauer <cit.> apply both to elliptic diffusion processes and to those hypoelliptic diffusion processes whose hypoellipticity arises from an interaction of the drift term with the diffusivity, such as for Langevin dynamics.
Although the recent papers <cit.> have addressed bridge simulations on manifolds, these schemes do not apply directly to simulate bridge processes associated with diffusion processes on sub-Riemannian manifolds, that is, those hypoelliptic diffusion processes where the hypoellipticity is induced by the diffusivity term itself, without any need to interact with the drift term. Phrased differently, as far as the authors are aware, it has not been known how to effectively simulate bridge processes on sub-Riemannian manifolds. Such spaces are described as a triple (M,E,g), where M is a smooth manifold, E⊆ TM is a bracket-generting subbundle that is called horizontal bundle and g is a smoothly varying inner product on E. The subbundle E represents the directions in which we allow the stochastic process on M to diffuse and the bracket-generating assumption on E ensures that the induced diffusion process is hypoelliptic. Sub-Riemannian structures naturally appear in all sciences, e.g. in the study of constrained physical systems such as the motion of robot arms or the orbital dynamics of satellites.

Motivated by the apparent gap in existing bridge simulation schemes, the immersiveness of sub-Riemannian geometry, and the importance of bridge simulation for data imputation and parameter inference, we here develop a method for bridge simulation on sub-Riemannian manifolds.



A central object in the study of diffusion bridge processes is the score of the density of the corresponding unconditioned diffusion process. To introduce the score associated with a diffusion process on a manifold M, we need a gradient ∇^E on M. On a Riemannian manifold, ∇^E is simply taken to be the usual gradient, whilst on a sub-Riemannian manifold (M,E,g), the horizontal gradient ∇^E will be defined with respect to the subbundle E.
For a time-homogeneous diffusion process on M starting from x_0∈ M and having transition probability density p_t M× M→ for t>0, the score S_t(x_0,·) M→ E is given by, for y∈ M,

    S_t(x_0,y) = ∇^y,Elog p_t(x_0, y).

Notably, the score is generally intractable due to its dependency on the transition probability densities.


Bridge simulation schemes often employ guiding terms that approximate the intractable score of the diffusion process. As we do not have such approximations available explicitly in the sub-Riemannian context, we employ the approach of <cit.> building on score matching techniques <cit.> to learn the score parameterised by a neural network and use this score in the simulation of the conditioned diffusion processes. The guiding terms used in previous works on Riemannian manifolds include radial vector fields, see <cit.>, or exploit the heat kernel on comparison manifolds as in <cit.>. Neither of these options are computationally tractable in the sub-Riemannian context. For training the neural network score approximation, we show how the score divergence as well as a denoising term can be used for the loss function. Both approaches are used in the Euclidean setting, but we will see that the sub-Riemannian geometry influences the terms resulting in non-trivial differences between the Euclidean and geometric losses, see Figure <ref>.







 §.§ Related work

In Euclidean space, simulation of conditioned diffusion processes has been treated extensively in the literature in works including <cit.>.
The use of explicit score approximations has extended bridge simulations to Riemannian manifolds, see <cit.>. Instead of using explicit score approximations, the work <cit.> has demonstrated how score matching in <cit.> can be exploited to train a neural network to approximate the score and use this in the simulation. Score matching and generative modelling on manifolds has been the focus of <cit.>.
In this paper, we build on elements of these ideas to show how a neural network can learn an approximated score on sub-Riemannian manifolds, which subsequently can be used for bridge sampling.



 §.§ Paper outline

In Section <ref>, we discuss the needed background material for the paper, and, in Section <ref>, we provide further details on bridge processes and their time reversals in the sub-Riemannian setting linking their generators and the score. We apply this in Section <ref> to construct the neural network score approximator using divergence loss, followed by a derivation of the corresponding denoising loss in the sub-Riemannian context in Section <ref>. 
Combined together, the stochastic differential equation for the conditioned process and the score approximation allow for numerical simulations of the bridge processes, which we use in Section <ref> to exemplify score simulations and bridge sampling on the Heisenberg group.



§ DIFFUSION PROCESSES ON SUB-RIEMANNIAN MANIFOLDS


We here discuss relevant background material on sub-Riemannian geometry and diffusion processes on sub-Riemannian manifolds needed in the remainder of the paper.


 §.§ Sub-Riemannian manifolds

Let M be a connected manifold. We consider a fixed subbundle E ⊆ TM that we denote the horizontal bundle.
We refer to the elements of E as horizontal vectors and we use (E) to denote the space of vector fields that take values in E. Let g = ⟨· , ·⟩_g be a smoothly varying inner product defined only on E. The triple (M, E,g) is called a sub-Riemannian manifold, and the smoothly varying tensor g is called a sub-Riemannian metric. Such an inner product gives rise to a sharp map

    ♯ T^*M → E     defined by    α(v) = ⟨♯α, v⟩_g,     α∈ T^*M, v ∈ E.

Note that ♯ is not invertible when E ≠ TM. A sub-Riemannian metric g can also be described in terms of its cometric g^*.
This cometric g^*=⟨· , ·⟩_g^* is defined for any pair of covectors α, β∈ T_x^*M for x ∈ M by

    ⟨α, β⟩_g^* = ⟨♯α, ♯β⟩_g.

Whilst the sub-Riemannian cometric g^* is defined on all of T^*M, it is degenerate whenever E ≠ TM because it vanishes on ♯ = (E), which consists of all covectors vanishing on E.

We say that an absolutely continuous curve γ[a,b]→ M is horizontal if it is the case that γ̇(t) ∈ E_γ(t) for almost every t ∈ [a,b]. Such a curve has a well-defined length given by

    (γ) = ∫_a^b γ̇(t)  _g  t :=∫_a^b ⟨γ̇(t) , γ̇(t) ⟩_g^1/2 t.

Furthermore, the metric defines a distance d_g M× M→ by setting d_g(x,y) for x,y∈ M to be the infimum over the lengths of all those horizontal curves which connect x to y. In order to guarantee the property that d_g(x,y) is finite for all x,y∈ M, we assume for the remainder of the paper that E is bracket-generating, that is, we suppose that all directional derivatives can be recreated from E. To be more precise, let us define, for any x∈ M,

    E̅_x = { [σ_1,[σ_2,…, [σ_l-1,σ_l]⋯ ]] (x) : σ_j ∈(E), j∈{1,…, l}, l≥ 1 },

where we interpret the bracket for the case l=1 simply as σ_1. The vector bundle E is then said to be bracket-generating if E̅ = TM. The Chow–Rashevskiĭ theorem in <cit.> tells us that the bracket-generating assumption is a sufficient condition not only for d_g to be finite, but furthermore for its metric topology to have the same open sets as the original topology of the manifold M.

Observe that we can consider a Riemannian manifold (M,g) as a special case of a sub-Riemannian manifold (M,E,g) with E = TM.




 §.§ Sub-Laplacians and diffusion processes

We consider a sub-Riemannian manifold (M, E, g).
For a function f ∈ C^∞(M), define ∇^E f = ♯ f ∈(E), the horizontal gradient of f. In the special case where E =TM, we simply write ∇^E f= ∇ f. Assuming that we also have a smooth volume density μ on M, we define a corresponding second order operator Δ = Δ_E,g,μ by

    Δ f = _μ(∇^E f),      f∈ C^∞(M).

The operator Δ is the sub-Laplacian of (M,E,g,μ). It is the unique second order operator on M satisfying, for any pair of compactly supported functions f_1, f_2 ∈ C^∞_0(M),

    ⟨∇^E f_1, ∇^E f_2 ⟩_g     = 1/2( Δ (f_1 f_2) - f_1 Δ f_2 - f_2 Δ f_1),
    ∫_M ⟨∇^E f_1, ∇^E f_2 ⟩_g μ    = - ∫_M f_1 Δ f_2 μ = - ∫_M f_2 Δ f_1 μ.

If E=k and (σ_1, …, σ_k) is a local orthonormal frame for E with respect to the sub-Riemannian metric g, then

    Δ = ∑_j=1^k (σ_j^2 + (_μσ_j) σ_j).

If the subbundle E is a proper subbundle of the tangent bundle TM, then the operator Δ is not elliptic. However, if E is bracket-generating, then both the operator Δ and its heat operator ∂_t - 1/2Δ are hypoelliptic, see Hörmander <cit.>.

For the special case where E  = TM and (M,g) is a Riemannian manifold, the operator Δ is called the Witten–Laplacian. If μ_g is the Riemannian volume form of the Riemanian metric g and Δ_g is the Laplace–Beltrami operator corresponding to μ_g, then there exists a function ϕ on M such that μ = ^ϕμ_g, which is related to Δ by

    Δ = Δ_g + ∇ϕ,

where ∇ϕ is defined with respect to the metric g.

For the remainder of the paper, we consider an operator L of the form

    L = Δ+ 2Z,

where Δ is the sub-Laplacian of (M,E,g,μ) and Z is a smooth vector field on M. For such a second order partial differential operator L and x_0∈ M, there exists a unique semi-martingale (X_t)_ζ>t≥ 0 = (X_t^x_0)_ζ>t≥ 0 on M starting from x_0, with infinitesimal generator 1/2L and defined up to some maximal explosion time ζ, see e.g. Elworthy <cit.>, Émery <cit.> and Hsu <cit.>. Throughout the paper, we will assume that the process does in fact not explode, that is, ζ(x_0)=∞ almost surely. For details on some criteria for non-explosion, see e.g. Bichteler and Jacod <cit.> or Norris <cit.> as well as Grigor’yan <cit.> and <cit.> for geometric conditions.

 
Consider L = Δ_μ+2Z for Δ_μ = Δ_E,g,μ defined with respect to a sub-Riemannian structure (E,g) and a smooth volume density μ on a manifold M. We observe that if ν = ^ϕμ is another smooth measure on M then Δ_ν = Δ_E,g,ν satisfies

    Δ_ν=Δ_μ + ∇^E ϕ,

and hence, we can write

    L = Δ_ν - ∇^E ϕ+ 2Z =: Δ_ν + 2Z̃.

It follows that we are free to change our definition of μ as long as we update the vector field Z accordingly.

Conversely, if we have L = Δ + 2Z where 2Z = ∇^E ϕ is a horizontal gradient then L = Δ_E,g,μ̃ is the sub-Laplacian with respect to some measure μ̃ on M.




 §.§ Local description of the operators

The following discussion concerning the local description of Brownian motion on a Riemannian manifold provides a motivation for studying the more general case of sub-Riemannian manifolds. Let (W_t)_t≥ 0 be a standard Brownian motion on ℝ^d and consider the unique strong solution (X_t)_t≥ 0, assumed to exist for all times, to the Stratonovich stochastic differential equation

    X_t = σ(X_t) ∘ W_t + σ_0(X_t)  t,

where σ_0 ℝ^d →ℝ^d is a smooth vector field and σℝ^d →ℝ^d× d, x ↦σ(x) = (σ^i_j(x)) a smooth function. If σσ^⊤(x) = (∑_l=1^d σ^i_lσ^j_l(x)) is positive definite for all x∈^d, then we can define a Riemannian metric g on ^d by setting g = (g_ij) = (σσ^⊤ )^-1. The infinitesimal generator 1/2 L of the stochastic process (X_t)_t≥ 0 is then in terms of the Laplace–Beltrami operator Δ_g given by

    L = ∑_i,j,l=1^d σ^i_lσ^j_l∂_x^i∂_x^j + ∑_l=1^d (2σ_0^l + ∑_i,j=1^d  σ^i_j (∂_x^iσ^l_j) )∂_x^l =: Δ_g + 2Z,

where Δ_g = ∑_i,j,l=1^d σ_l^i σ_l^i ∂_x^i∂_x^j -∑_i,j,l=1^d g^ijΓ_ij^l ∂_x^l,

    2Z = ∑_l=1^d 2Z^l ∂_x^l = ∑_l=1^d ( 2 σ^l_0 + ∑_i,j=1^d (   σ^i_j (∂_x^i σ^l_j) +g^ij Γ_ij^l   )    ) ∂_x^l

and Γ_ij^l are the Christoffel symbols with respect to the Riemannian metric g on ^d.
 
If σσ^⊤ is merely positive semi-definite of constant rank, then we cannot invert it. However, we can still use g^* = σσ^⊤ as a cometric for a sub-Riemannian structure (E,g). If we consider g^* as a matrix, which is the same as identifying it with ♯, then E will be the image of g^*.



 §.§ Densities and conditional probability

Let Δ be the sub-Laplacian of the sub-Riemannian manifold (M, E,g) equipped with the smooth volume density μ and let Z be a smooth vector field on M. Define (X_t)_t≥ 0 to be the unique stochastic process starting from x_0∈ M and having generator 1/2 L = 1/2Δ + Z. Recall that we assume this stochastic process does not explode.
For T>0 given, let =^x_0,T be the measure on the path space C([0,T], M) induced by (X_t)_t∈ [0,T]. For s,t∈[0,T] with s<t, we further let _s,t denote the push-forward measure on M × M of the map (x_r)_r ∈ [0,T]↦ (x_s, x_t). Similarly, we set _t to be the push-forward measure of the evaluation (x_r)_r ∈ [0,T]↦ x_t.
Since E is assumed to be bracket-generating, it follows by <cit.> that _s,t and _t have smooth densities with respect to μ⊗μ and μ, respectively.
For any x,y∈ M and for any s,t∈[0,T] with s<t, we write

    _s,t(x,y) = p_s(x) p_s,t(x,y) (μ⊗μ)(x,y)
      and  _t(x) = p_t(x) μ(x),

where p_s,t M× M→ is the transition probability density such that

    p_t(y)  =∫_M p_s(x) p_s,t(x,y) μ(x)   and   p_t(y) = p_0,t(x_0,y).

Furthermore, since L is a time-homogeneous operator, we know that

    p_s,t(x,y) = p_0,t-s(x,y) =: p_t-s(x,y).

From our assumption that E is bracket-generating, it further follows from <cit.> that p_t(y) > 0 and p_s,t(x,y) >0 for all x,y∈ M and all s,t∈[0,T] with s<t. By definition, for any open set U⊆ M and for s,t∈(0,T] with s<t, we have

    (X_t ∈ U) = ∫_U p_t(y) μ(y)
      and  (X_t ∈ U | X_s =x) = ∫_U p_s,t(x,y) μ(y).


Let L^* denote the adjoint of the operator L with respect to the smooth volume density μ, which is given by

    L^* = Δ- 2Z - 2_μ Z.

From the definition of the generator, we have the relation
∂_t [f(X_t)] = 1/2[Lf(X_t)] for f∈ C^∞(M), which in turn implies that, for x,y∈ M and t∈[0,T],

    (∂_t- 1/2 L^*) p_t(y) = 0,   (∂_t - 1/2 L^y,*) p_s,t(x,y) =0

as well as

    (∂_s + 1/2 L^x) p_s,t(x,y) = 0,

subject to initial conditions given by p_0(y) = δ_x_0(y) and p_t,t (x,y)= δ_x(y). For the details, see e.g. Haussmann and Pardoux <cit.>. 


If we define p_s,t^*(x,y) = p_s,t(y,x) and P_t^*f(x) = ∫_M f(y) p_t^*(x,y)μ(y) then by the above equations, we have (∂_t - 1/2L^*) P_t f =0 and P_0 f = f. It follows that as long as _μ Z =0 there then exists a stochastic process (X_t^*)_t≥ 0 with generator 1/2 L^* and such that

    p_s,t^*(x,y) = p_t-s^*(x,y) = p_t-s(y,x)

are its transition probability densities.




§ BRIDGE PROCESSES AND TIME-REVERSED PROCESSES 


We determine the generator for Riemannian and sub-Riemannian time-reversed diffusion bridge processes, with our main objective being to demonstrate how this generator is related to the score defined in (<ref>).



 §.§ Bridge processes

As before, let (X_t)_t≥ 0 on (M,E,g,μ) be the stochastic process with generator 1/2L=1/2Δ +Z and initial value x_0∈ M. For T>0 and x_T∈ M, the diffusion bridge process (Y_t)_t∈[0,T] on (M,E,g,μ) obtained by conditioning (X_t)_t∈[0,T] on X_T=x_T can formally be defined using disintegration of measure techniques, see e.g. <cit.> and <cit.>. If there exists a unique path of minimal energy connecting the starting point x_0 to the final point x_T then in small time, as shown in <cit.>, the laws of the diffusion bridge processes concentrate near the unique path of minimal energy. This result has been extended in <cit.> with x_0 and x_T being connected by multiple minimal paths, subject to the additional assumption that none of them contains a segment which is so-called abnormal. We remark that this concentration property for the diffusion bridge processes in small time heavily relies on the underlying geometric structure and, as demonstrated in <cit.>, may fail for more general diffusion bridge processes.

Using the formulation of the Doob h-transform, see e.g. <cit.>, an expression for the generator of the diffusion bridge process (Y_t)_t∈[0,T] in terms of the logarithmic derivative of the transition probability densities p_t,T = p_T-t and the generator 1/2L of the unconditioned process (X_t)_t∈[0,T] can be obtained.
 
The diffusion bridge process (Y_t)_t∈[0,T] has the infinitesimal generator 

    1/2 L + ∇^x,Elog p_T-t(· , x_T).




    For t∈[0,T), let the function h_t M→ be defined by h_t(x) = p_t,T(x, x_T).
    According to (<ref>), we know that
    
    (∂_t + 1/2 L)h_t = 0,

    and by the Chapman–Kolmogorov equation, we further have, for s∈(0,T-t),
    
    h_t(x) = ∫_M p_t,t+s(x,y) h_t+s(y)μ(y).

    For any open set U⊆ M, we obtain
    
        (X_t+s∈U|X_t=x,X_T=x_T)
        =∫_Up_t,t+s(x,y)p_t+s,T(y,x_T)/p_t,T(x,x_T)μ(y).
    
    Hence, the diffusion bridge process (Y_t)_t∈[0,T] on (M,E,g,μ) has the same law as the Markov process (X_t^h)_t∈[0,T] starting from X_0^h = x_0 whose transition kernel p^h is given by, for x,y∈ M,
    
    p_t,t+s^h(x,y) = p_t,t+s(x,y) h_t+s(y)/h_t(x).

    It follows that, with p_t^h(y) = p_0,t^h(x_0, y), we have
    
    (Y_t ∈ U) = ∫_U p_t^h(y) μ(y).

    This is indeed all well-defined due to h_t(x)=p_t,T(x,x_T)>0 for all x∈ M and all t∈[0,T) 
    and, by the space-time regularity property (<ref>), further defines a genuine transition probability density and thus a Markov process.
    The expression (<ref>) implies that, for f∈ C^∞_0(M),
    
        𝔼[f(X_t+s^h)|X_t^h =x]
        = 𝔼[f(X_t+s) h_t+s(X_t+s)| X_t =x]/h_t(x).
    
    Using this formula, we can find the generator 1/2 L^h_t for the Markov process (X_t^h)_t∈[0,T], and hence for (Y_t)_t∈[0,T], by first deducing that
    
    1/2L^h_t f(x)
            = lim_s ↓ 0𝔼[f(X_t+s^h)| X_t^h = x] -f(x)/s
        = lim_s ↓ 0𝔼[f(X_t+s)h_t+s(X_t+s)| X_t = x] -f(x)h_t(x)/s h_t(x)
        = (∂_t+1/2L) (f h_t)(x)/h_t(x).

    By applying (<ref>) and using 1/2L=1/2Δ +Z, we conclude that
    
        (∂_t+1/2L) (f h_t)/h_t
        =1/2Lf+ ⟨∇^E f, ∇^E h_t ⟩/h_t,
    
    which establishes 1/2L^h_t = 1/2L + ∇^E log h_t. Due to h_t(x) = p_t,T(x, x_T)=p_T-t(x,x_T), the claimed result follows.




 §.§ Time-reversals
 
We consider a general process (Y_t)_t≥ 0 on (M,E,g,μ) with a time-dependent generator 1/2L_t = 1/2Δ + Z_t for a time-dependent vector field Z_t and with initial value x_0∈ M. We shall again assume non-explosion. For T>0, we denote the associated measure induced on the path space C([0,T],M) by . Let us further define q_t on M for t∈ [0,T] through, for any open set U⊆ M,

    (Y_t ∈ U) = ∫_U q_t(y) μ(y).


We are now interested in looking at the time-reversed process (Y̅_t)_t∈[0,T] given by Y̅_t = Y_T-t. We observe the following about its infinitesimal generator.
 
The time-reversed processes (Y̅_t)_t∈[0,T] defined by setting Y̅_t = Y_T-t for t∈[0,T] has infinitesimal generator

    1/2L̅_t = 1/2Δ- Z_T-t + ∇^E logq_T-t.

Furthermore, we have that ∇^E q_t is in L^2(E; g,μ) for t∈[0,T].

To produce our result, we modify the proof of <cit.> for a coordinate independent formulation.


In order to prove that the time-reversed process (Y̅_t)_t∈[0,T] has infinitesimal generator 1/2L̅_t, it suffices to show that, for s,t∈ [0,T] with s<t and for any compactly supported function f ∈ C^∞_0(M), the mapping

    t ↦ f(Y̅_t) - f(Y̅_s) - 1/2∫_s^t (L̅_r f)(Y̅_r)  r

defines a martingale.
As argued in <cit.>, the above is equivalent to the statement that, for all f, ϕ∈ C^∞_0(M),

    𝔼[ (f(Y̅_t) - f(Y̅_s) - 1/2∫_s^t (L̅_r f)(Y̅_r)  r ) ϕ(Y̅_s)] = 0,

which after applying Y̅_t = Y_T-t and changing variables in the time integral amounts to

    𝔼[ (f(Y_T-t) - f(Y_T-s) - 1/2∫_T-t^T-s (L̅_T-r f)(Y_r)  r ) ϕ(Y_T-s)] = 0.

Since T-t<T-s for s<t, it therefore suffices to establish that, for all s,t∈ [0,T] with s<t and for all f, ϕ∈ C^∞_0(M),

    𝔼[ (f(Y_t) - f(Y_s)  + 1/2∫_s^t (L̅_T-r f)(Y_r)  r ) ϕ(Y_t)] = 0.


Let _s,t with, for x,y∈ M,

    _s,t(x,y)=q_s(x)q_s,t(x,y) (μ⊗μ)(x,y)

denote the push-forward measure of the map (x_r)_r ∈ [0,T]↦ (x_s,x_t). We observe

    [ f(Y_s) ϕ(Y_t) ] = ∫_M × M f(x) ϕ(y) _s,t(x,y) 
        = ∫_M f(x) q_s(x) ( ∫_M q_s,t(x,y) ϕ(y) μ(y) ) μ(x) 
        =: ∫_M f(x) q_s(x) Φ_s,t(x) μ(x) = [ f(Y_s) Φ_s,t(Y_s) ],

where Φ_s,t(x) = ∫_M q_s,t(x,y) ϕ(y) μ(y) = [ϕ(Y_t)| Y_s = x].
By using the properties that (∂_r + 1/2L_r) Φ_r,t =0 and (∂_r - 1/2L_r^*) q_r =0, we further obtain

    [f(Y_t) Φ_t,t(Y_t)] - [f(Y_s) Φ_s,t(Y_s)]
       =∫_s^t∫_M f(x)(q_r(x)∂_rΦ_r,t(x) + Φ_r,t(x)∂_r q_r(x))μ(x) r
       =1/2∫_s^t ∫_M f(x)(-q_r(x)  L_r Φ_r,t(x) + Φ_r,t(x) L_r^* q_r(x))μ(x) r
       =1/2∫_s^t ∫_M  Φ_r,t(x) (- L_r^*( f q_r)(x)   + f(x) L_r^* q_r(x))μ(x) r.

Since the adjoint L^*_t of the operator L_t=Δ+2Z_t with respect to μ is given by

    L^*_t = Δ - 2Z_t - 2_μ Z_t,

we have

    1/2L^*_t(fq_t) - 1/2fL^*_t q_t
        = (1/2Δ  - Z_t - _μ Z_t) (fq_t) - f (1/2Δ  - Z_t - _μ Z_t) q_t 
        = q_t (1/2Δ - Z_t) f +⟨∇^E f, ∇^E q_t ⟩
    = 1/2q_t L̅_T-t f.

Hence, it follows that

    [f(Y_t) ϕ(Y_t)] - [f(Y_s) ϕ(Y_t)]=
      [f(Y_t) Φ_t,t(Y_t)] - [f(Y_s) Φ_s,t(Y_s)]
       = -1/2∫_s^t ∫_M q_r(x) Φ_r,t(x) (L̅_T-rf)(x)μ(x) r 
       = -1/2∫_s^t ∫_M× M q_r(x) q_r,t(x,y) ϕ(y) (L̅_T-rf)(x) (μ⊗μ)(x,y)  r 
        = - 1/2[ ϕ(Y_t) ∫_s^t (L̅_T-rf)(Y_r)  r ],

as required.


By combining Lemma <ref> and Lemma <ref>, we can determine the infinitesimal generator of the stochastic process obtained by time-reversing the diffusion bridge process corresponding to (X_t)_t≥ 0 with infinitesimal generator 1/2 L = 1/2Δ + Z and starting point x_0∈ M conditioned on X_T=x_T for T>0 and x_T∈ M. Recall that we use p_t(x,y) = p_0,t(x,y) for x,y∈ M to denote the transition probability density of the semi-martingale (X_t)_t≥ 0.

    For (Y_t)_t∈[0,T] the diffusion bridge process obtained by conditioning (X_t)_t∈[0,T] with infinitesimal generator 1/2 L = 1/2Δ + Z and starting point x_0∈ M on X_T=x_T, we define the time-reversed process (Y̅_t)_t∈[0,T] by Y̅_t = Y_T-t for t∈[0,T]. Then the infinitesimal generator of (Y̅_t)_t∈[0,T] is given by
    
    1/2L̅_t
            =1/2Δ-Z+S_T-t(x_0,·).




    By Lemma <ref>, the infinitesimal generator of (Y_t)_t∈[0,T] takes the form
    
        1/2 Δ+Z + ∇^x,E logp_T-t(·, x_T).
    
    Lemma <ref> then implies that the time-reversed process (Y̅_t)_t∈[0,T] has infinitesimal generator
    
    1/2L̅_t
            =1/2Δ-Z-∇^x,Elog p_t(·,x_T)
            +∇^Elog q_T-t,

    where q_t is still defined by (<ref>). It remains to show that the linear term reduces to the claimed expression. Since (Y_t)_t∈[0,T] is the diffusion bridge process associated with (X_t)_t≥ 0 from x_0 to x_T in time T>0, we have, for y∈ M,
    
        q_T-t(y)
        =p_0,T-t(x_0,y)p_T-t,T(y,x_T)/p_0,T(x_0,x_T)
        =p_T-t(x_0,y)p_t(y,x_T)/p_T(x_0,x_T).
    
    It follows that
    
        ∇^Elogq_T-t
        =∇^y,E logp_T-t(x_0,·)+∇^x,E logp_t(·,x_T).
    
    This combined with (<ref>) and the definition (<ref>) establishes
    
        1/2L̅_t
        =1/2Δ-Z+∇^y,E logp_T-t(x_0,·)
        = 1/2 Δ- Z + S_T-t(x_0,·),
    
    as required.



We use the reversed generator from Lemma <ref> for simulating bridge processes. However, we could have also applied it for generative modelling as is often the focus in the machine learning literature on diffusion models, i.e., given a stochastic process (Y_t)_t∈[0,T] with some final distribution Y_T, reversing the process enables sampling from the initial distribution for Y_0. Some examples of generative modelling on manifolds include <cit.>.





§ APPROXIMATING THE SCORE WITH NEURAL NETWORKS


We now build the setup for approximating the score using neural networks in the sub-Riemannian setting. 
The ideas follow the ones presented in <cit.>, where diffusion processes on ^d with uniformly positive definite diffusivity matrices are considered. The complexities arising in our context correctly handle the sub-Riemannian geometry, particularly the hypoellipticity of the sub-Riemannian diffusion processes.

In this section, we discuss loss functions involving the divergence of the score and a basic Euler–Maruyama sampling scheme leading to the score learning algorithm. We also give some details on the sampling scheme, for further use in Section <ref>, where we discuss the alternative denoising loss.

As previously, let (E,g) be a bracket-generating sub-Riemannian structure on a manifold M and let Δ be the sub-Laplacian with respect to the volume form μ. We consider the stochastic process (X_t)_t≥ 0=(X_t^x_0)_t≥ 0 with infinitesimal generator 1/2 L = 1/2Δ + Z started from the initial value x_0∈ M and assumed to not explode. Let p_t(x_0, ·) for t>0 denote the density of (X_t)_t≥ 0 with respect to μ whose score S_t(x_0,·) is given by (<ref>).

Recall that according to Lemma <ref>, the bridge process associated with (X_t)_t≥ 0 from x_0 to x_T∈ M in time T>0 has generator 1/2 L + ∇^x,Elog p_T-t(· , x_T). Since, if Z = 0, the operator L^* = L = Δ is symmetric, it follows from Remark <ref> that then the generator of the bridge process can be expressed as 1/2L + S_T-t(x_T, ·). This would allow us to simulate bridge processes by approximating the score S_T-t(x_T, ·) from the target point x_T.
However, in general, we have p_t^*(y,x) = p_t(x,y) with p_t^*(y,·) a fundamental solution of ∂_t - 1/2L^* where the adjoint operator

    L^* = Δ- 2Z - 2_μ Z

with respect to μ may have a non-zero zeroth-order term. To avoid this problem, we instead consider the reversed bridge process (Y̅_t)_t∈[0,T] as in Corollary <ref> which has the generator

    1/2L̅_t = 1/2Δ-Z + S_T-t(x_0, · ).

Thus, we will first need to determine the score from the starting point x_0 and then simulate bridges backwards from the endpoint x_T.

We henceforth write S_t(·) = S_t(x_0, ·).



 §.§ Score representation and loss function

For estimating the score S, we let S^θ be a time-dependent vector field in E given as the output of a neural network whose network parameters are denoted by θ. If (σ_1, …, σ_k) is frame for E, we can consider this problem as searching for functions, for j∈{1,…,k},

    S^θ,j[0,T]× M   →ℝ
    
        (t,y)    ↦ S_t^θ,j(y)

such that S^θ = ∑_j=1^k S^θ,jσ_j. Further details of the network architecture are given in Section <ref>.

When modifying the neural network weights to make the network match the score, we work with the squared L^2-distance as the loss function

    (θ)      =∫_0^T ∫_M S^θ_t(y)- ∇^y,Elog p_t(x_0,y)_g^2  p_t(x_0,y) μ(y)  t 
        = ∫_0^T 𝔼^x_0[ S^θ_t(X_t)- ∇^y,Elog p_t(x_0,X_t)_g^2 ]  t.

For short-time approximations and for being able to compute the loss at every step, we will subdivide the energy (θ) into shorter time-intervals.
Observe that, for s∈[0,t),

    ∫_M ⟨ S_t^θ(y) , ∇^y,Elog p_t(x_0,y) ⟩_g   p_t(x_0, y)μ(y) =  ∫_M ⟨ S_t^θ(y) , ∇^y,E p_t(x_0,y) ⟩_g μ(y) 
        = ∫_M × M⟨ S_t^θ(y) , ∇^y,Elog p_t-s(z,y) ⟩_g   p_s(x_0,z) p_t-s(z,y) (μ⊗μ)(z,y)  
        =∫_M × M⟨ S_t^θ(y) , ∇^y,Elog p_t-s(z,y) ⟩_g ℙ_s,t(z,y)  = ^x_0[⟨ S^θ_t(X_t), S_t-s(X_s, X_t) ⟩_g ].

From these computations, we deduce that for e_s,t(θ) defined by

    e_s,t(θ)      = ^x_0[⟨ S^θ_t(X_t), S^θ_t(X_t) - 2 S_t-s(X_s, X_t) ⟩_g ],

the expression

    𝔼^x_0[ S^θ_t(X_t)- ∇^y,Elog p_t(x_0,X_t)_g^2 ] -e_s,t(θ)

is constant with respect to θ. This implies in particular that, for any subdivision 0=t_0 < t_1 < ⋯ < t_n= T of the interval [0,T], there exists a constant C such that

    (θ) = _0,T(θ) + C = ∑_i=0^n-1_t_i,t_i+1(θ) + C,
      where  _s,t(θ) := ∫_s^t e_s,r(θ)  r.

It follows that letting our neural network minimise relative to (θ) is equivalent to minimising with respect to _0,T(θ). Note that the definition for _s,t(θ) still contains the expression of the true score, which is usually intractable. In the subsequent sections, we discuss how we can consider minima of the loss functions without knowing the true score explicitly.





 §.§ Loss function with divergence

By exploiting an integration-by-parts trick, the loss function can be computed without explicitly knowing the true score. We prove the result here for the sub-Riemannian case that has been previously derived for the Euclidean setting in <cit.> and for Riemannian manifolds in <cit.>, based on earlier work in <cit.>.


For any s,t∈[0,T] with s<t and for e_s,t(θ) defined as in (<ref>), we have

    e_s,t(θ)     = ∫_M( S^θ_t(y)^2_g + 2(_μ S_t^θ)(y) ) _t(y) 
        = ^x_0[S^θ_t(X_t)^2_g + 2(_μ S_t^θ)(X_t)]  =: e_t^(2)(θ).


We remark that this rewritten expression for e_t^(2)(θ) does not depend explicitly on the score itself. Since then _s,t(θ) = ∫_s^t e_r^(2)(θ)  r, we can use the formulation of Theorem <ref> to minimise _0,T(θ) without knowing the score.
In computations, the expectations are approximated by averaging over samples X^(1),…,X^(K) from the stochastic process X=(X_t)_t∈ [0,T] through, with δ = T/n,

    _0,T(θ)
    ≈δ/K∑_l=1^K ∑_i=1^n (S^θ_iδ(X_iδ^(l))^2_g + 2(_μ S_iδ^θ)(X_iδ^(l))).




Recall that on the manifold M without boundary we have, for f∈ C^∞(M) and a smooth vector field V on M,

    ∫_M  f(V) μ = -∫_M f _μ(V) μ.

We apply (<ref>) to deduce

    ∫_M × M⟨ S_t^θ(y), ∇^y,Elog p_t-s(z,y) ⟩_g _s,t(z,y) 
        = ∫_M × M p_s(x_0, z) (^y p_t-s(z,y)(S_t^θ(y)))   (μ⊗μ)(z,y) 
        = - ∫_M × M p_s(x_0,z) p_t-s(z,y)(_μ S_t^θ)(y)   (μ⊗μ)(z,y)  
        = - ∫_M (_μ S_t^θ)(y) _t(y).

The claimed result that e_s,t(θ) = e^(2)_t(θ) then follows.




  §.§.§ Expression in local coordinates
 
Suppose we choose a local coordinate system (x^1, …, x^d) on M and let x =  x^1 … x^d denote the Euclidean volume measure in these coordinates. Consider a second order operator L = Δ + 2Z on M where Δ is a sub-Laplacian of a sub-Riemannian structure (E,g). By Remark <ref>, we know that with the appropriate choice for Z we can  assume that Δ = Δ_E,g, x is the sub-Laplacian with respect to x. Let (σ_1, …, σ_k) be a local orthonormal frame for (E,g), which is described by a d × k matrix (σ_j^i) such that σ_j = ∑_i=1^d σ_j^i ∂_x^i for j∈{1,…,k}. With Z = ∑_i=1^d Z^i ∂_x^i, we have

    L = Δ +2Z    = ∑_j=1^k ∑_i,l=1^d σ_j^l σ^i_j ∂_x^l∂_x^i + ∑_j=1^k ( ∑_l=1^d ∂_x^lσ_j^l ) ∑_i=1^d σ_j^i ∂_x^i
            + ∑_j=1^k ∑_i,l=1^d σ^l_j   ∂_x^lσ_j^i ∂_x^i  + 2∑_i=1^d Z^i ∂_x^i.


Let (W_t)_t≥ 0 be a standard Brownian motion on ^k and set

    σ_0 =  ∑_j=1^k (_ xσ_j) σ_j + 2Z = ∑_j=1^k ( ∑_l=1^d ∂_x^lσ_j^l ) ∑_i=1^d σ_j^i ∂_x^i  + 2∑_i=1^d Z^i ∂_x^i

as well as τ_0 = σ_0 + 1/2∑_j=1^k (∇̅_σ_jσ_j), where
∇̅ denotes the flat connection given by

    ∇̅_V_1 V_2 = ∑_i=1^d (V_1 V_2^i) ∂_x^i = ∑_i,l=1^d V_1^l (∂_x^l V_2^i) ∂_x^i .

Then the stochastic process (X_t)_t≥ 0 with generator 1/2 L is the unique strong solution to the Stratonovich stochastic differential equation

    X_t = ∑_j=1^k σ_j(X_t) ∘ W_t^j + σ_0(X_t)  t,     X_0=x_0,

or to the Itô stochastic differential equation

    X_t = ∑_j=1^k σ_j(X_t)  W_t^j + τ_0(X_t)  t ,     X_0=x_0.

Writing the output of the neural network as S^θ = ∑_j=1^k S^θ,jσ_j, we obtain

    e_t^(2)(θ)    =  ∫_M ( ∑_j=1^k S_t^θ,j (y)^2 ) p_t(x_0,y)  y 
          + 2 ∫_M (  ∑_j=1^k 
    ∑_i=1^d(S_t^θ,j(y) (∂_x^iσ^i_j)(y) + σ_j^i(y) (∂_x^i S_t^θ,j)(y) )) p_t(x_0,y)  y.

If we now have sample paths X^(1), …, X^(K) and discretise the interval [0,T] with δ = T/n, then (<ref>) gives rise to

    _0,T(θ)    ≈δ/K∑_l=1^K ∑_i=1^n  ∑_j=1^k S^θ, j_i δ(X_iδ^(l))^2 
          + δ/K∑_l=1^K ∑_i=1^n  ∑_j=1^k  
    ∑_m=1^d( S_iδ^θ,j(X_iδ^(l)) (∂_x^mσ^m_j)(X_i δ^(l)) + σ_j^m(X_iδ^(l)) (∂_x^m S_iδ^θ,j)(X_iδ^(l)) ).




 §.§ Time integration


We now explicitly describe an Euler–Maruyama simulation scheme for finding paths of the stochastic process (X_t)_t∈[0,T]=(X_t^x_0)_t∈[0,T] with generator 1/2 L = 1/2Δ+Z. Later on, we will further exploit this approximation for the denoising version of the loss function. 

In order to describe the method, we assume that M = ℝ^d topologically.
If the manifold M has a different topology, we can either work in local coordinates or use geometric Euler–Maruyama methods. For instance, see <cit.> for some examples of methods for matrix Lie groups.

Assume that (X_t)_t∈[0,T] can, at least locally, be considered as the solution of (<ref>) in a coordinate system (x^1,…,x^d). For a fixed n∈, we set δ = T/n and t_i = i δ for i∈{0,1,2,…, n}. The Euler–Maruyama approximation (X̂_t)_t∈[0,T] = (X̂_t^n)_t∈[0,T] is then defined by, for i∈{0,1,2,…, n-1} and t∈(0,δ],

    X̂_t_i+t = X̂_t_i + ∑_j=1^k σ_j(X̂_t_i)  (W_t_i+t^j-W_t_i^j) + (X̂_t_i) t,     X̂_0 = x_0.

The Euler–Maruyama approximation (X̂_t^n)_t∈[0,T] converges strongly to (X_t)_t∈[0,T] as n→∞ under appropriate growth conditions on the vector fields ,σ_1,…, σ_k, see <cit.> for details.



 §.§ Score estimation

Together the loss e_s,t(θ) given by (<ref>) and the sampling scheme (<ref>) provide the tools needed to train the score approximation. The explicit algorithm is listed in Algorithm <ref> below.




§ APPROXIMATING LOSS FUNCTIONS FOR SHORT TIME STEPS


Even though the loss function in _s,t(θ) = ∫_s^t e_r^(2)(θ)  r defined by (<ref>) has no explicit dependency on the score, using the term (_μ S_t^θ)p_t can in practice be problematic because the integral often is approximated with finite samples at the steps of the training, and the training might therefore at each step minimise the divergence at the sample points only, resulting in unstable convergence. Moreover, the computational expense of finding derivatives of the neural network and taking gradients of those for the optimisation can be a problem if the dimension of M is high. These issues have also been observed in the Euclidean case, and a common way to deal with these problems is the denoising loss, see <cit.>. In this section, we show how this approach can be generalised to the sub-Riemannian setting.

With the notation in this paper, the denoising loss arises from (<ref>) using explicit approximations of the true score for s and t close, e.g. for steps from t_i to t_i+1. Such steps can be approximated, in the Euclidean setting, by normal distributions and, in the Riemannian situation, by heat kernel approximations or tangent space normal distributions as in <cit.>. In the sub-Riemannian context, we take horizontal derivatives of the score, but we also need to take account of the change of the distribution between tangent spaces and the hypoellipticity of the diffusion process. We therefore look at short-time approximations for the steps of diffusion processes on sub-Riemannian manifolds.



 §.§ Short-time score approximations with Euler steps
  
We choose a local coordinate system (x^1, …, x^d) such that we can identify M with ^d. As in Section <ref>, let us consider the approximation (X̂_t)_t∈[0,T] given as the solution to (<ref>) from the Euler–Maruyama integration. The equation is stated using the global orthonormal frame (σ_1,…,σ_k) represented by the matrix σ = (σ^i_j). We write Σ = σσ^⊤ and observe, for t∈(0,δ), that X̂_t ∼ N(t (x_0), t Σ(x_0)) which is supported on t (x_0) + E_x_0. It follows that the probability measure _t =_t^n of X̂_t is not even absolutely continuous with respect to _t whenever t∈ (0,δ) and E≠ TM. If we write _t = p̂_t μ then p̂_t cannot be described as a continuous function.



We further encounter the problem that in general the tangent space at the point y ∈ t(x_0) + E_x_0 is not included in the horizontal distribution at y, i.e., we generally have T_y(t  (x_0) + E_x_0) ⊈E_y. In particular, it becomes problematic to define a horizontal gradient. We instead consider the following solution. We take S^θ = ∑_j=1^k S^θ,jσ_j with associated vector representation ^θ = (S^θ,1, …, S^θ,k) and observe that

    e_t_i,t_i+t(θ)     = ^x_0[⟨ S^θ_t_i+t(X_t_i+t), S^θ_t_i+t(X_t_i+t) - 2 S_t(X_t_i, X_t_i+t) ⟩_g ] 
        =  ∑_j=1^k ∫_M× M( S^θ,j_t_i+t(y)^2 - 2 S^θ,j_t_i+t(y) ^y log p_t(z,y)(σ_j(y))  ) ℙ_t_i,t_i+t(z,y).

As a result of the discussed problems with the horizontal gradient in the approximation scheme due to only having access to derivatives of p̂_t(z,·) in the E_z directions, we evaluate the differential above at the initial point instead of the target point and use the approximation

    ê_t_i,t_i+t(θ) := ∑_j=1^k ∫_M× M( S^θ,j_t_i+t(y)^2 - 2 S^θ,j_t_i+t(y) ^y logp̂_t(z,y)(σ_j(z)) )
    ℙ̂_t_i,t_i +t(z,y),

where _s,t is the probability measure on M × M of (X̂_s, X̂_t).
For t∈(0,δ), we have

    p̂_t(z,y) = δ_t(z) + E_z(y) 1/√((2π t)^k ∏_i λ_i)exp(- ⟨ y-z- t(z), y-z- t(z) ⟩_g(z)/2t),

where ∏_i λ_i is the product of all non-zero eigenvalues of Σ(z). These functions have well-defined derivatives in the directions of E_z which yields

    ^y logp̂_t(z,y)(σ_j(z)) = -1/t⟨σ_j(z) , y-z - t(z) ⟩_g(z).

Recalling now that X̂_t_i+t - X̂_t_i = σ(X̂_t_i)  (W_t_i+t-W_t_i) + (X̂_t_i) t, we obtain

    ê_t_i,t_i+t(θ) 
        = ∑_j=1^k ^x_0[ S^θ,j_t_i+t(X̂_t_i+t)^2 + 2/t⟨ (S^θ,j_t_i+tσ_j)(X̂_t_i) , X̂_t_i+t - X̂_t_i - t(X̂_t_i) ⟩_g(X_t_i)] 
        =  ^x_0[ ⟨^θ_t_i+t(X̂_t_i+t), ^θ_t_i+t(X̂_t_i+t) + 2/t (W_t_i+t-W_t_i) ⟩_ℝ^k].


In summary, if we approximate the diffusion process (X_t)_t∈[0,T] with (X̂_t)_t∈[0,T] which relies on the Euler–Maruyama scheme, then we can use as a loss function _0,T(θ) = ∑_i=0^n-1∫_t_i^t_i+1ê_t_i,s(θ)  s that tries to predict each step. However, since each Euler step merely has positive probability on a proper subspace, we can only predict the gradient of logp̂_t( X̂_t_i, X̂_t_i+t) at the initial point X̂_t_i. We again use a time-discretisation with n intervals and write δ = T/n. We generate K × n random vectors Δ_i W^(l) for i∈{1, …, n} as well as l∈{1, …, K} that are drawn from a N(0,δ I)-distribution in ℝ^k. We define first X^(l)_0 = x_0 and then iteratively

    X̂_i+1^(l)= X̂_i^(l)+∑_j=1^k σ_j(X̂_i^(l)) Δ_i+1 W^(l),j+ δτ_0(X_i^(l))

and compute the loss function by

    _0,T(θ)    ≈1/K∑_l=1^K ∑_i=1^n  ⟨^θ_i δ(X̂_i^(l)), δ ^θ_i δ(X̂_i^(l)) + 2 Δ_i W^(l)⟩_ℝ^k
        = δ/K∑_l=1^K ∑_i=1^n  ^θ_i δ(X̂_i^(l))+ 1/δΔ_i W^(l)_ℝ^k^2 +C.

We can use the loss _0,T(θ) to train a neural network by applying Algorithm <ref> with the loss _0,T(θ) replaced by _0,T(θ).



 §.§ Taylor expansion and approximations of stochastic integrals


We improve on the method in Section <ref> by considering a more complicated approximation scheme for each step. Contrary to what is the case for the Euler–Maruyama approximation, each step in our new approximation scheme gives us densities with respect to Lebesgue measure that are smooth and positive. As a tool in this new approximation, we use the stochastic Taylor expansion. We will describe the resulting approach in Section <ref>.

We first look at the Stratonovich Taylor expansion for a stochastic process found in <cit.>.
Let (X_t)_t≥ 0 be the unique strong solution of the Stratonovich stochastic differential equation (<ref>) with vector fields σ_0, σ_1, …, σ_k. We use the convention W^0_t =t and write W_s,t^j = W_t^j - W_s^j. For a multi-index α = (α_1, …, α_l) ∈{0,…, k}^l for l∈_0, we set l(α) = l and we let n(α) denote the number of zeros in the multi-index α. We further define the differential operators

    σ_α = σ_α_1σ_α_2⋯σ_α_l

as well as the stochastic processes given by

    J^α_s,t = ∫_s<t_1< ⋯ < t_l<t∘ W_t_1^α_1∘⋯∘ W_t_l^α_l.

Let (x^1, …, x^d) be a choice of local coordinate system for the manifold M and define a function _α =(σ_α^i)_i with values in ^d such that

    σ_α^i = σ_α x^i,

i.e., the differential operator σ_α applied to the i-th coordinate function.
For further reference, we state below <cit.>.


Let ⌊·⌋ denote the floor function, i.e., ⌊ a ⌋ = max_n ∈ℤ, n ≤ a  n. Let _γ be the collection of all multi-indices α satisfying l(α) +n(α) ≤ 2 γ, and set _γ^+ = { (β_1, α) : α∈_γ, β_1 ∈{0,1,…, k}}. We say that a vector-valued function 𝐟 satisfies (<ref>) if almost surely

    H_α∫_0^T ∫_0<t_1< ⋯ <t_l-1<s𝐟(X_t_1)  W_t_1^α_1⋯ W_t_l-1^α_l-1^2-δ_0,α_l s < ∞.

Let integers γ and n be given such that γ≥ 0 and n>T. Set δ = T/n and t_i = i δ for i∈{0,1, …,n}, and define X̂^i iteratively by X̂^0 = 0 as well as

    X̂^i+1 = ∑_l(α) + n(α)≤ 2γ J_t_i,t_i+1^α·_α(X̂^i),

which we extend to (X̂_t)_t∈[0,T] through

    X̂_t_i+t = ∑_l(α) + n(α)≤ 2γ J_t_i,t_i+t^α·_α(X̂^i) ,      for t∈ [0, δ].

If we assume that we have constants C_1, C_2 such that both


  * _α(x) - _α(y)_^d≤ C_1 x-y_^d for any α∈_γ, and

  * _α is C^1,1, satisfies condition (<ref>) and furthermore the inequality

    _α(x) _^d≤ C_2(1+x_^d)

for any α∈_γ∪_γ^+,

then we have the following bound, for some constant C,

    ^x_0(sup_0 ≤ t ≤ T X_t - X̂_t_^d) ≤ C δ^γ.




Let us now consider the result of Lemma <ref> for the particular case where γ=1. For t ∈ [t_i, t_i+1], we obtain

    X̂_t    = X̂_t_i + (t-t_i) _0(X̂_t_i) + ∑_j=1^k W^j_t_i,t_j(X̂_t_i)  + ∑_j,l=1^k  ∫_t_i^t W^j_t_i,s∘ W_s^l _(j,l)(X̂_t_i)  
       = X̂_t_i + (t-t_i) _0(X̂_t_i) + ∑_j=1^k W^j_t_i,t_j(X̂_t_i)  
            + 1/2∑_j,l=1^k  W_t_i,t^j W_t_i,t^l_(j,l)(X̂_t_i)   + ∑_1≤ j<l≤ k A^j,l_t_i,t (_(j,l)-_(l,j))(X̂_t_i),

where A_s,t^j,l = 1/2∫_s^t (W_s,r^j∘ W^l_r-  W_s,r^l∘ W^j_r) is the so-called Lévy area.

We continue by discussing several methods for approximating this Lévy area. One option, see e.g. <cit.> exploits the Fourier expansion of the involved Brownian motions and eliminates a dependency to give, for j,l∈{1,…,k},

    A^j,l_t,t+ h = 
        ∑_m=1^∞(a_l,m W^j_t,t+h - a_j,m W^l_t,t+h)+  π∑_m=1^∞ m(a_j,m b_l,m -a_l,m b_j,m )

in terms of independent Gaussian random variables with W^j_t,t+h∼ N(0, h) and

    a_j,m   =2/h∫_0^h (W_s^j -s/h W_h^j )
        cos(2π ms/h) s
        ∼ N(0, h/2π^2 m^2),
    
        b_j,m   =2/h∫_0^h (W_s^j -s/h W_h^j )
        sin(2π ms/h) s
        ∼ N(0, h/2π^2 m^2).

We can then approximate the Lévy area by truncating the series in (<ref>).

A second option which only requires the simulation of independent Gaussian random variables makes use of the polynomial decomposition of Brownian motion developed in <cit.> and results in the expansion, see <cit.>,

    A_t,t+h^j,l
        =1/2(c_l,1 W_t,t+h^j - c_j,1 W_t,t+h^l )
        +1/2∑_m=1^∞(c_j,mc_l,m+1-c_j,m+1c_l,m),

with independent Gaussian random variables given by W_t,t+h^j∼ N(0,h) as well as

    c_j,m∼N(0,h/2m+1).

Truncating the series in (<ref>) then again yields an approximation for the Lévy area.

If we have a discretisation 0 =t_0 < ⋯ < t_n=T with δ = t_i+1- t_i then we can obtain the bound in Lemma <ref> with approximated Lévy area under appropriate bounds. If Â_t,t+h^j,l is an approximation for the Lévy area A_t,t+h^j,l such that ^x_0(sup_t_i ≤ t ≤ t_i+1| A^j,l_t_i,t - Â^j,l_t_i,t| ) ≤ C_0 δ^3 on each interval [t_i,t_i+1], and if we replace A^j,l_t_i,t by Â^j,l_t_i,t in the definition of (X̂_t)_t∈[0,T], then by <cit.>, we still have

    ^x_0(sup_0 ≤ t ≤ T X_t - X̂_t_^d) ≤ C δ.



Since we are working with local coordinates, we can identify vectors v = ∑_i=1^d v^i ∂_x^i with the point (v^1,…, v^d). Under this correspondence, we can identify _j with σ_j and _(j,l) with ∇̅_σ_jσ_l for ∇̅ defined by (<ref>). It follows that X̂_t for t ∈ [t_i, t_i+1] can be expressed as

    X̂_t 
       = X̂_t_i + (t-t_i) σ_0(X̂_t_i) + ∑_j=1^k W^j_t_i,tσ_j(X̂_t_i)  
            + 1/2∑_j,l=1^k  W_t_i,t^j W_t_i,t^l∇̅_σ_jσ_l(X̂_t_i)  + ∑_1≤ j<l≤ k A^j,l_t_i,t [σ_j, σ_l](X̂_t_i).





 §.§ Heisenberg group
 
Let us consider the space M=ℝ^2k+1 with coordinates q =(x,y,z) for x,y ∈ℝ^k and z ∈ℝ. We equip M with a multiplication rule such that if q = (x,y,z) and q̃ = (x̃, ỹ, z̃), then

    ℓ_q̃(q) =q̃· q
            = (x̃+x, ỹ+y, z̃+ z+1/2(⟨x̃, y ⟩_ℝ^k - ⟨ x, ỹ⟩_ℝ^k) ).
 
Here, ℓ_q̃ denotes the left translation with respect to q̃, i.e., multiplication by q̃ on the left. Observe that this multiplication is not abelian. With respect to this group structure, we have the identity element 0 =(0,0,0) and inverses given by (x,y,z)^-1 = (-x,-y,-z). We further define the sub-Riemanian structure (E,g) on M where E is the rank 2k subbundle spanned by the vector fields

    σ_j = ∂_x^j - y^j/2∂_z,     τ_j = ∂_y^j + x^j/2∂_z,      j∈{1, …, k},

and the sub-Riemannian metric g is defined uniquely by requiring (σ_1,τ_1,…,σ_k,τ_k) to be an orthonormal frame for E. The subbundle E is bracket-generating because [σ_i, τ_j] = δ_i,j∂_z.
We additionally observe that the above vector fields are left-invariant, meaning that, for any smooth function φ M→ and for q∈ M, we have

    σ_j (φ∘ℓ_q) = (σ_j φ) ∘ℓ_q,    τ_j(φ∘ℓ_q) = (τ_j φ) ∘ℓ_q.

As a consequence, if d_g denotes the sub-Riemannian distance on (M,E,g) and we write f(q) := d_g(0,q) then

    d_g(q̃, q) = d_g(0, q̃^-1·  q) = f( q̃^-1· q)  .

If we choose μ to be the usual Lebesgue measure on ℝ^2k+1 then the corresponding sub-Laplacian Δ=Δ_E,g,μ on M=^2k+1 is given by

    Δ = ∑_j=1^k σ_j^2 + ∑_j=1^k τ_j^2.

This sub-Laplacian Δ is also left-invariant, i.e., for a smooth function φ M→ and q∈ M, we have Δ (φ∘ℓ_q) = (Δφ) ∘ℓ_q. We now consider the stochastic process (X_t)_t≥ 0 = (X_t^0)_t≥ 0 on M with infinitesimal generator 1/2Δ and starting from 0∈ M. Alternatively, (X_t)_t≥ 0 can be characterised as the unique strong solution to the Stratonovich stochastic differential equation

    X_t = ∑_j=1^k σ_j(X_t) ∘ W^j_t + ∑_j=1^k τ_j(X_t) ∘ W^k+j_t,

subject to X_0=0 and where (W_t^1, …, W^2k_t)_t≥ 0 is a standard Brownian motion on ^2k. The stochastic differential equation (<ref>) subject to X_0=0 has the explicit solution existing for all times t≥ 0 and defined by

    X_t     = (W_t^1, …, W^k_t, W^k+1_t,…,W_t^2k, 1/2∑_j=1^k ∫_0^t (W^j_s ∘ W_s^k+j - W^k+j_s ∘ W_s^j)  ) 
        = (W_t^1, …, W_t^2k, ∑_j=1^k A^j,k+j_0,t).

If we take another initial point q ∈ M, then the solution (X_t^q)_t≥ 0 to (<ref>) subject to X_0^q=q is given by X_t^q = ℓ_q(X_t) for t≥ 0. Furthermore, we have

    X_t_i +t = X_t_i· X_t_i,t_i+t  with   X_t_i,t_i+t = (W_t_i,t_i+t^1,…, W_t_i,t_i+t^2k, ∑_j=1^k A_t_i,t_i+t^j,k+j).


From the explicit expression above, we see that the stochastic process (X_t)_t≥ 0 and the associated density p_t(0,q) = p_t(q) can be simulated using approximations of the Lévy area described in Section <ref>. We further note that since the sub-Laplacian Δ is left-invariant, we have p_t(q̃, q) = (p_t ∘ℓ_q̃^-1)(q).

Let us now look at an explicit formula for p_t. By adapting the expression derived in <cit.> to our normalisation for the vector fields used to define the sub-Riemannian structure (E,g), we obtain
 
    p_t(q) = 4/(2π t)^k+1∫_-∞^∞( 2λ/sinh 2λ)^k  exp(4λ z/t-2λ(2λ) x_^k^2 +y_^k^2/2t)λ.

As further shown in <cit.>, we then have the approximation, as t↓ 0,

    log p_t(q) ≈ -d_g(0,q)^2/2t = - f(q)^2/2t,

and consequently

    log p_t(q̃, q) = log p_t(q̃^-1· q) ≈ - f(q̃^-1· q)^2/2t = -d_g(q̃,q)^2/2t.

Unfortunately, computing the distance f explicitly is very expensive. However, it is well known, see e.g. <cit.>, that

    f(x,y,0) = √(x_ℝ^k^2 + y_ℝ^k^2)    and     f(0,0,z) = 2 √(π |z|).

Define f̂ M→[0,∞) by

    f̂(x,y,z)^2 = f(x,y,0)^2 + f(0,0,z)^2 =  x_ℝ^k^2 +  y_ℝ^k^2 +4 π |z|.

Following the proof of <cit.>, we also have f(x,y,0) ≤ f(x,y,z) and f(0,0,z) ≤ 2f(x,y,z). Hence, we obtain that

    1/8f̂(x,y,z)^2    ≤1/4max{ f(x,y,0)^2, f(0,0,z)^2}
       ≤ f(x,y,z)^2  ≤ ( f(x,y,0) + f(0,0,z))^2 ≤ 2f̂(x,y,z)^2.

If we further consider d̂ M× M→ given by d̂(q̃, q) = f̂(q̃^-1· q), then d̂ and d_g are Lipschitz equivalent left-invariant metrics agreeing on the distance from 0 to points on the z-axis and to points in the (x,y)-plane. We therefore introduce the following score approximation

    Ŝ_t(q̃, q) := - 1/2t∇^q,Ed̂(q̃, q)^2 = -1/2t∇^q,Ef̂(q̃^-1· q)^2.

We further note that

    ∇^E(f̂^2∘ℓ_q̃^-1)(q) = ∇^q,Ef̂(q̃^-1· q  )^2 
        = ∑_j=1^k ( 2 (x^j-x̃^j)- 2π(y^j - ỹ^j)(z - z̃ - 1/2 (⟨x̃,  y ⟩_ℝ^k - ⟨ x, ỹ⟩_ℝ^k))) σ_j( q)
          + ∑_j=1^k (2  (y^j - ỹ^j) + 2π (x^j - x̃^j)(z - z̃ - 1/2 (⟨x̃,  y ⟩_ℝ^k - ⟨ x, ỹ⟩_ℝ^k))) τ_j(q).

It follows that if writing Ŝ_t(q̃, q) = ∑_j=1^k Ŝ_t^j(q̃, q) σ_j(q) + ∑_j=1^k Ŝ_t^k+j(q̃, q) τ_j(q)
then, for l∈{1,…,2k},

    Ŝ_t^l(q̃, q) = Ŝ_t^l(0,q̃^-1· q) =: Ŝ_t^l(q̃^-1· q),

which are given by, for j∈{1, …, k},

    -t Ŝ_t^j(q) =  x^j- πy^j(z )    and    -t Ŝ_t^k+j(q) = y^j  + πx^j(z).


We now continue with the following procedure. Suppose we have a discretisation  0 = t_0 < t_1 < ⋯ < t_n = T with t_i+1 - t_i =δ = T/n. For each t∈(0, δ], we define

    X̂_t_i+t = X̂_t_i·X̂_t_i,t_i+t,

with

    X̂_t_i,t_i+t = (W_t_i,t_i+t^1, …, W_t_i,t_i+t^2k,  ∑_j=1^k Â^j,k+j_t_i,t_i+t).

Letting _s,t denote the probability measure on M × M associated with (X̂_s,X̂_t), we then define the approximation ê_s,t(θ) to e_s,t(θ) by

    ê_s,t(θ) = ∫_M × M⟨ S^θ_t(y), S^θ_t(y) - 2Ŝ_t-s(z,y)⟩^2_g(y) _s,t(z,y).

We write S^θ_t = ∑_j=1^k S_t^θ,jσ_j + ∑_j=1^k S_t^θ,k+jτ_j. In terms of _t = (Ŝ^1_t, …, Ŝ^2k_t) as well as _t^θ = (S^θ,1_t, …, S^θ,2k_t), we obtain

    ê_t_i,t_i+t(θ)    = [⟨^θ_t_i+t(X̂_t_i+t), ^θ_t_i+t(X̂_t_i+t) - 2_t(X̂_t_i,t_i+t ) ⟩_ℝ^2k] 
       = [⟨^θ_t_i+t(ℓ_X̂_t_iX̂_t_i,t_i+t), ^θ_t_i+t(ℓ_X̂_t_iX̂_t_i,t_i+t) - 2_t(X̂_t_i,t_i+t ) ⟩_ℝ^2k].

We remark that

    -_t(q) = 1/t[ x; y ] + π/t(z) [ -y;  x ]  =: - 1/t(q).


Using these observation in summary, we describe how _0,T(θ) is presented with a discretisation of [0,T] for δ = T/n. We generate K · n · 2k random samples Δ_i W^(l),j from the distribution N(0,δ), and further K · K_2 · n · 2k numbers Δ_i c_j,m^(l) sampled from N(0, δ/2m+1). Here 1 ≤ i ≤ n, 1≤ j ≤ 2k, 1 ≤ l≤ K and 1 ≤ m ≤ K_2. We then determine

    Δ_i A^(l), j,k+j   = 1/2(Δ_i c_k+j,1^(l)Δ_i W^(l),j - Δ_i c_j,1^(l)Δ_i W^(l),k+j) 
            +1/2∑_m=1^K_2-1( Δ_i c_j,m^(l)Δ_i c_k+j,m+1^(l) - Δ_i c_j,m+1^(l)Δ_i c_k+j,m^(l))

and set

    Δ_i A^(l) = ∑_j=1^k Δ_i A^(l),j,k+j,     Δ_i X̂^(l) = (Δ_i W^(l),1,…, Δ_i W^(l),2k, Δ_i A^(l)).

We further define X̂_0^(l) =x_0 and iteratively

    X̂_i+1^(l) = X̂_i^(l)·Δ_i+1X̂^(l).

We compute the loss function by

    _0,T(θ)
       ≈1/K∑_l=1^K ∑_i=1^n ⟨^θ_i δ(X̂_i^(l)), δ ^θ_i δ(X̂_i^(l)) - 2(Δ_i X̂^(l) ) ⟩_ℝ^2k
        = δ/K∑_l=1^K ∑_i=1^n ^θ_i δ(X̂_i^(l)) - 1/δ(Δ_i X̂^(l) ) _ℝ^2k^2+ C

and can now apply Algorithm <ref> with the loss _0,T(θ) replaced by _0,T(θ).






 §.§ General case and adapted coordinates
 
Despite us using very particular properties of the Heisenberg group in Section <ref>,
it is possible to perform local approximations for a general sub-Riemannian manifold in a similar manner.

Let (M,E,g) be a sub-Riemannian manifold where the manifold M has dimension d. Let p_t be the Dirichlet heat kernel for an operator 1/2 L =1/2Δ + Z.
We will make use of small-time asymptotics, see <cit.>, saying that, for any x_1, x_2 ∈ M, the transition density p_t satisfies the limit

    lim_t→ 0 2t log p_t(x_1,x_2) = - d_g(x_1,x_2)^2.

In order to use this result, we need an approximation of the distance d_g.

For the remaining results in this section, we follow Montgomery <cit.>. We consider adapted coordinates, which are also called privileged coordinates in <cit.>. For a given point x_0 ∈ M, we define a collection of integers k_1,k_2,…,k_s which satisfy k= k_1 ≤  k_2 ≤⋯≤ k_s=d such that k_1 is the rank of E, k_2 is the rank of the space spanned by vector fields in E along with their first order Lie brackets at x_0, k_3 is the rank of the space which also includes their second order Lie brackets at x_0, and so on. Due to the bracket-generating assumption on E, there exists some minimal s∈ such that k_s=d, which is the step of E. The vector (k_1, …, k_s) is then called the growth vector of E at x_0.
For instance, in the (2k+1)-dimensional Heisenberg group, the growth vector is (2k,2k+1) at all points. With the convention that k_0=0, we further define weights ν_1, …, ν_d such that ν_i for i∈{1,…, d} is the maximal number satisfying

    k_ν_i-1≤ i.

In the Heisenberg group, we have (ν_1,…,ν_2k,ν_2k+1) =(1,…, 1,2).

A coordinate system (y^1,…,y^d) centered at x_0 is then called adapted at x_0 if, for any l ≤ν_i-1 and for any selection of vector fields V_1,…, V_l in E,

    (V_1… V_l y^i)(x_0) = 0.

Note that the condition for ν_i =1 is already satisfied for all coordinates because y(x_0) = 0 by assumption. Relative to an adapted coordinate system, we define the box metric d_box by, for x_1 and x_2 in the domain of the coordinate system y,

    d_box(x_1, x_2) =max_1≤ i ≤ d{ |y^i(x_1) - y^i(x_2)|^1/ν_i}.

The ball box theorem says, for all x in the domain of y and for constants c and C,

    c d_box(x_0,x) ≤ d_g(x_0,x) ≤ C d_box(x_0,x).

In order to have a smoother metric, we instead consider d̂ defined by

    d̂(x_1,x_2)^2 = ∑_i=1^d |y^i(x_1) - y^i(x_2)|^2/ν_i.

The previous inequalities imply that there are positive constants c̃ and C̃ such that

    c̃d̂(x_0,x)^2 ≤ d_g(x_0,x)^2 ≤C̃d̂(x_0,x)^2.

Motivated by (<ref>) and (<ref>), we introduce the approximation

    log p_t(x_0,x) ≈  -d̂(x_0,x)^2/2t,

leading to our final approximated score

    Ŝ_t(x_0,x) := -1/2t∇^x,Ed̂(x_0,x)^2.



We conclude that if (σ_1, …, σ_k) is an orthonormal frame for E with respect to g and we write σ_j = ∑_i=1^d σ^i_j ∂_y^i for j∈{1,…,k}, then our approximation of the score is given by

    Ŝ_t(x_0,x) = - 1/2t∇^x,Ed̂(x_0,x)^2 = - 1/t∑_j=1^k ( ∑_i=1^d σ_j^i(x)  y^i |y^i|^2/ν_i-2/ν_i) σ_j(x).



Note that, for any selection of strictly positive constants c^1,…, c^d, the metric d̂^c defined by

    d̂^c(x_1,x_2)^2 = ∑_i=1^d c^i |y^i(x_1) - y^i(x_2)|^2/ν_i

will also satisfy (<ref>). Hence, for each case, we can scale this metric to make it match the true sub-Riemannian metric close to x_0 for each coordinate, as we did for the Heisenberg group in Section <ref>. A recommended method for finding a good choice of c^1, …, c^d is to compute explicit lengths along each coordinate axis according to the sub-Riemannian metric of the nilpotent homogeneous approximation at x_0, as described in <cit.>.





 §.§ Simulation of a step in the general geometry
 
For simulating the stochastic process (X_t)_t∈[0,T] in the general sub-Riemannian setting, we use the stochastic Taylor expansion discussed in Section <ref> and the adapted coordinates reviewed in Section <ref>. For the sake of simplicity, we reduce our considerations to the case where the bracket-generating distribution E has step 2, that is, where the tangent bundle TM is spanned by E and its first order Lie brackets.

As before, let (σ_1, …, σ_k) be an orthonormal frame for E, let (W_t)_t≥ 0 be a standard Brownian motion on ^k and assume that (X_t)_t∈[0,T] is the unique strong solution to the Stratonovich stochastic differential equation

    X_t = ∑_j=1^k σ_j(X_t) ∘ W_t^j + σ_0(X_t)  t,     X_0=x_0.


Additionally consider vector fields σ_k+1, …, σ_d on M such that (σ_1, …, σ_d) is a frame for the full tangent bundle TM.
Let (x^1, …, x^d) be any choice of local coordinate system for M around the point x_0 = (x^1_0, …, x^d_0)∈ M. We introduce a matrix-valued function σ =(σ^i_l) such that σ_l = ∑_i=1^d σ^i_l∂_x^i for l∈{1,…,d}, and we denote its inverse by σ^-1 = (σ_-l^i). We then define a new local coordinate system (y^1, …, y^d) for M centred at x_0 by setting

    y^i = ∑_l=1^d σ^i_-l(x_0) (x^l- x^l_0).

This yields an adapted coordinate system at x_0.
Letting (Y_t)_t∈[0,T] defined by Y_t = y(X_t) be the process in this adapted coordinate system and taking the Taylor expansion of order 1, we obtain that, for t∈[0,δ] and i∈{1,…,k} as well as m∈{k+1,…,d},

    Y_t^i ≈  W_t^i + ∑_l=1^d ∑_j_1,j_2=1^k (∫_0^t W_s^j_1∘ W_s^j_2) σ^i_-l(x_0) (∇̅_σ_j_1σ_j_2)^l(x_0) + ∑_l=1^d t σ^i_-l(x_0) σ^l_0(x_0)

as well as

    Y_t^m≈∑_l=1^d ∑_j_1,j_2=1^k (∫_0^t W_s^j_1∘ W_s^j_2) σ^m_-l(x_0) (∇̅_σ_j_1σ_j_2)^l(x_0) + ∑_l=1^d t σ^m_-l(x_0) σ^l_0(x_0).

Returning to our original coordinate system (x^1,…,x^d) and identifying vectors in the basis (∂_x^1, …, ∂_x^d) with points in the domain of (x^1,…,x^d), we get the approximation

    X_t - x_0≈∑_j=1^k W_t^j σ_j(x_0)  + ∑_j_1,j_2=1^k (∫_0^t W_s^j_1∘ W_s^j_2) (∇̅_σ_j_1σ_j_2)(x_0) + t σ_0(x_0).

We now have the following beneficial proposition.

Suppose that {σ_j, [σ_j, σ_l] : j,l∈{1,…, k}} equals the full tangent space at every point on M. For a given point x_0∈ M,  we define a stochastic process (X̂_t)_t≥ 0 by X̂_0 = x_0 and

    X̂_t - x_0 = ∑_j=1^k σ_j(x_0) W_t^j + ∑_j_1,j_2=1^k  (∇̅_σ_j_1σ_j_2)(x_0) ∫_0^t W_s^j_1∘ W_s^j_2+ σ_0(x_0)t.

Then X̂_t  for t>0 has a positive smooth density with respect to Lebesgue measure.

We remark that by using the identity ∫_0^t W_s^j_1∘ W_s^j_2 = 1/2 W_t^j_1 W_t^j_2 + A^j_1,j_2_0,t we can rewrite the expression (<ref>) as

    X̂_t - x_0    = ∑_j=1^k σ_j(x_0) W_t^j + 1/2∑_j_1,j_2=1^k  (∇̅_σ_j_1σ_j_2)(x_0) W_t^j_1 W_t^j_2
            + ∑_1≤ j_1<j_2≤ k  [σ_j_1,σ_j_2](x_0) A^j_1,j_2_0,t + σ_0(x_0)t.




We consider the space ^k ×^d with coordinates (w,x) and we introduce the vector fields, for j∈{1,…,k},

    σ̂_j(w,x)     = ∂_w_j + ∑_i=1^d σ_j^i(x_0) ∂_x_i + ∑_j_1=1^k ∑_i=1^d w^j_1(∇̅_σ_j_1σ_j)^i(x_0) ∂_x_i
        = ∂_w_j +σ_j(x_0)  + ∑_j_1=1^k w^j_1 (∇̅_σ_j_1σ_j)(x_0).

We notice that, for j,l∈{1,…,k} and at any point (w,x)∈^k ×^d,

    [σ̂_j, σ̂_l](w,x) = ∑_i=1^d [ σ_j, σ_l]^i(x_0) ∂_x^i.


Let us now consider the set O⊆^k ×^d of points that can be reached from (0,x_0) by moving tangentially to the subbundle Ê = {σ̂_j:j∈{1,…,k}}, and further define the constant vector space R = {∑_i=1^d [σ_j_1, σ_j_2]^i(x_0) ∂_x^i}_1≤ j_1 < j_2≤ k. By the Orbit theorem <cit.>, the set O is a manifold of dimension k +  R, whose tangent space is given by

    T_(w,x)O = Ê_(w,x)⊕ R_(w,x),      (w,x) ∈ O.
 
Let π^k ×^d →^d be the projection on the second factor. By our assumption on the brackets, we have π_* (Ê_(w,x)⊕ R_(w,x)) = T_x ^d for any (w,x) ∈^k ×^d. It follows that π|_O O →^d is a surjective submersion with each fiber π^-1(x) being an embedded manifold.

We observe that π_*,(w,x) maps (π_*,(w,x))^⊥ bijectively to T_x^d. Hence, we can define a vector field σ̂_0 on ^k ×^d determined uniquely by being in (π_*)^⊥ and satisfying π_*,(w,x)σ̂_0(w,x) = σ_0(x_0), where we identify the vector σ(x_0) and the corresponding constant vector field.

Define a stochastic process (Φ_t)_t≥ 0 on ^k×^d as the unique strong solution to

    Φ_t = ∑_j=1^k σ̂_j(Φ_t) ∘ W_t^j + σ̂_0(Φ_t)  t,    Φ_0 =(0,x_0).

Since Ê|_O is bracket-generating as a subbundle of TO, we deduce from <cit.> and <cit.> that Φ_t for t>0 has a positive smooth density φ_t with respect to Lebesgue measure μ^O on O.
Furthermore, by definition we have π(Φ_t) = X̂_t for t≥ 0. Hence, the density
p̂_t of X̂_t given by
p̂_t(x) = ∫_π^-1(x)φ_t(w,x) μ^O|_π^-1(x) is the average of φ_t on every fiber, which is then also smooth and positive for t>0.




 §.§ General summary

To summarise the above results, assume that (E,g) is a step 2 sub-Riemannian structure with an orthonormal frame (σ_1, …, σ_k) for E, extended with further vector fields σ_k+1, …, σ_d to a frame for the tangent bundle TM, which in local coordinates (x^1, …, x^d) is represented as a matrix (σ^i_j) with inverse (σ_-j^i). Let us divide [0,T] into n intervals of length δ = T/n and set t_i = i δ for i∈{0,…,n}. We then define an approximation (X̂_t)_t∈[0,T] to (X_t)_t∈[0,T] by X̂_0 = x_0 as well as, for t∈(0,δ] and i∈{0,…, n-1},

    X̂_t_i +t    = X̂_t_i+ ∑_j=1^k σ_j(X̂_t_i) W_t_i,t_i+t^j + 1/2∑_j_1,j_2=1^k  (∇̅_σ_j_1σ_j_2)(X̂_t_i)  W_t_i,t_i+t^j_1 W_t_i,t_i+t^j_2
            + ∑_1 ≤ j_1 <j_2 ≤ k  [σ_j_1, σ_j_2](X̂_t_i)  Â^j_1,j_2_t_i,t_i+t + σ_0(X̂_t_i) t,

where Â^j_1,j_2_t_i,t_i+t is a truncation of (<ref>) or (<ref>), respectively. We apply this model to sample from the density p̂_t(x_0,·) of X̂_t.

Moreover, we need a corresponding loss function for the score. For this, we use the adapted coordinates defined in (<ref>) for the approximation (<ref>) of the score, that is, with

    f̂(x)^2 = d̂(x_0, x)^2     = ∑_i=1^k |y^i|^2 + ∑_i=k+1^d |y^i| 
        = ∑_i=1^k ( ∑_l=1^d σ_-l^i(x_0) (x^l -x_0^l))^2 + ∑_i=k+1^d|∑_l=1^d σ_-l^i(x_0) (x^l -x_0^l) |,

we take Ŝ_t(x_0,x) = - 1/2t∇^E f̂(x)^2.
Writing

    Ŝ_t(x_0,x)=- 1/2t ∑_j=1^k σ_j(f̂^2)(x) σ_j(x)
    =∑_j=1^k Ŝ_t^j(x_0,x) σ_j(x),

we then have

    Ŝ_t^j(x_0,x) 
        = - 1/t∑_i=1^k ∑_l_2=1^d σ^l_2_j(x) σ_-l_2^i(x_0) ( ∑_l=1^d σ_-l^i(x_0) (x^l -x_0^l)) 
            - 1/2t∑_i=k+1^d ∑_l_2=1^d σ^l_2_j(x) σ_-l_2^i(x_0) (∑_l=1^d σ_-l^i(x_0) (x^l -x_0^l) ).

We remark that the expression above differs from (<ref>) in that there the coordinate expressions σ_j^i for the vector fields are defined relative to the adapted coordinate system (y^1,…,y^d), whereas here these are relative to a given coordinate system (x^1,…, x^d). A particular advantage of this formulation is that we do not need to change coordinate system if we change x_0.

We now define the loss function _0,T(θ) for the neural network S^θ =∑_j=1^k S^θ,jσ_j using ^θ = (S^θ,1, …, S^θ,k)  and = (Ŝ^1, …, 
Ŝ^k) by

    _0,T(θ)     = ∫_0^T ^x_0[⟨_t^θ(X̂_t), _t^θ(X̂_t)-2_t(x_0,X̂_t) ⟩_ℝ^k] t 
        = ∫_0^T ^x_0[ ^θ_t(X̂_t) - _t(x_0, X̂_t) _ℝ^k^2 ]  t +C.





§ EXPERIMENTS


We now exemplify the bridge simulation code for the Heisenberg group whose sub-Riemannian structure is described in Section <ref>. We use an Euler–Heun to simulate the unconditioned Brownian motion, and we employ this to train score approximators using the denoising loss (<ref>) and the divergence loss (<ref>). Code for representing the Heisenberg geometry, geometric computations, and training of the score approximations is available in the Jax Geometry library[<https://github.com/computationalevolutionarymorphometry/jaxgeometry/>].

The neural net used is a 3-layer dense network with 15 nodes in each layer and exponential linear units (ELU) activation functions. For the loss functions depending on the divergence, it is important that the activation functions are differentiable. In contrast, with the denoising loss, the ELU can be replaced with e.g. rectified linear units (ReLU).
The nets are trained for 2500 epochs with batch size of 64 sample paths and 8 batches per epoch.

In Figure <ref>, we sample a single sample path (blue) starting at x_0=(0.5,0,0.8) and conditioned on hitting x_T=(0,0,0) at T=0.1. For comparison, the minimising geodesic (red) between the two points x_0 and x_T is included in the figure. We additionally plot the norm of the (x,y)-component and the z-component, respectively, as a function of t∈[0,0.1], particularly to illustrate the progression of the z coordinate from its starting value towards 0.


Since we know from short-time asymptotics established in <cit.> for diffusion processes on sub-Riemannian manifolds  that sample paths should concentrate around a geodesic if we reduce the conditioning time towards 0, we plot in Figure <ref> median curves and quartiles with conditioning for T=0.1, T=0.2, T=0.5 and T=1, respectively. The concentration effect is clearly visible.


Finally, Figure <ref> shows plots of the estimated score fields in the unit cube in ^3 as well as evaluated only on the (x,y)-slice of the unit cube satisfying z=0.2. The figure shows estimated score fields with the denoising loss (<ref>) (top row) and the divergence loss (<ref>) (bottow row). For values close to the origin, the estimated scores appear visually close. Our experience is that training with the divergence loss is in practice less stable than using the denoising loss which is exemplified in the figure when the score is evaluated for higher values of z where the stochastic process has lower probability density and the training therefore is performed with fewer samples.



abbrv




