Vector-valued Gaussian Processes on Riemannian Manifolds via Gauge Independent Projected Kernels

1 Introduction . Gaussian processes are an effective model class for learning unknown functions . They are particularly attractive for use within data-efficient decision systems , including Bayesian optimization [ 3 , 32 , 39 ] , model-based reinforcement learning [ 34 , 7 ] , and active learning [ 21 ] . In these settings , Gaussian processes can represent and propagate uncertainty , as well as encode inductive biases as prior information in order to drive data efficiency . A key aspect of prior information is the geometry of the domain on which the Gaussian process is defined , which often encodes key properties , such as symmetry . Following the growing deployment of Gaussian processes , a number of recent works have focused on how to define Gaussian processes on non-Euclidean domains in ways that reflect their geometric structure [ 2 , 1 ] . In many applications , such as climate science , quantities of interest are vector-valued . For example , global wind velocity modeling must take into account both speed and direction , and is represented by a vector field . On geometric domains , the mathematical properties of vector fields can differ noticeably from their Euclidean counterparts : for instance , one can prove that every smooth vector field on a sphere must vanish in at least one point [ 26 ] . Behavior such as this simultaneously highlights the need to represent geometry correctly when modeling vector-valued data , and presents a number of non-trivial technical challenges in constructing models that are mathematically sound . ∗Equal contribution . Code : https : //github.com/MJHutchinson/ExtrinsicGaugeIndependantVectorGPs . For a general implementation , see https : //github.com/GPflow/GeometricKernels/ . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . In particular , even the classical definition of a vector-valued Gaussian process—that is , a random function with multivariate Gaussian marginals at any finite set of points—already fails to be a fully satisfactory notion when considering smooth vector fields on a sphere . This is because tangent vectors at distinct points live within different tangent spaces , and it is not clear how to construct a cross-covariance between them that does not depend on a completely arbitrary choice of basis vectors within each space . Constructions that are independent of this choice of basis are called gauge independent , and recent work [ 45 , 5 , 15 ] in geometric machine learning has focused on satisfying this key property for convolutional neural networks that deal with non-Euclidean data . Our contributions include the following . We ( a ) present a differential-geometric formalism for defining Gaussian vector fields on manifolds in a coordinate-free way , suitable for Gaussian process practitioners with minimal familiarity with differential geometry , ( b ) present a universal and fully constructive technique for defining prior Gaussian vector fields on Riemannian manifolds , which we term the projected kernel construction , and ( c ) discuss how to adapt key components in the computational Gaussian process toolkit , such as inducing point methods , to the vector field setting . The structure of the paper is as follows . In Section 2 , we define vector-valued Gaussian processes on smooth manifolds . We start by reviewing the multi-output Gaussian process set-up , which is typically used in machine learning . We then detail a differential-geometric formalism for defining vectorvalued Gaussian processes on smooth manifolds . In Section 3 , we provide a concrete construction for these Gaussian processes on Riemannian manifolds and discuss how they can be trained using variational sparse approximations . Section 4 showcases Gaussian vector fields on two tasks , namely weather imputation from satellite observations and learning the dynamics of a mechanical system . 2 Vector-valued Gaussian Processes on Smooth Manifolds . A vector-valued Gaussian process ( GP ) is a random function f : X → Rd such that , for any finite set of points x ∈ Xn , the random variable f ( x ) ∈ Rn×d is jointly Gaussian . Every such GP is characterized by its mean function µ : X → Rd and matrix-valued covariance kernel k : X ×X → Rd×d , which is a positive-definite function in the matrix sense . These functions satisfy E ( f ( x ) ) = µ ( x ) and Cov ( f ( x ) , f ( x′ ) ) = k ( x , x′ ) for any x , x′ ∈ X . Here , dependence between function values is encoded in the kernel ’ s variability along its input domain , and correlations between different dimensions of the vector-valued output are encoded in the matrix that the kernel outputs . Consider a function f with y = f ( x ) + ε , where ε ∼ N ( 0 , σ2I ) and training data ( x , y ) . Placing a GP prior f ∼ GP ( 0 , k ) on the unknown function results in a GP posterior , whose mean and covariance are given by E ( f | y ) = K ( · ) x ( Kxx + σ2I ) −1y Cov ( f | y ) = K ( · , · ) −K ( · ) x ( Kxx + σ2I ) −1Kx ( · ) . ( 1 ) Here , ( · ) denotes an arbitrary set of test locations , Kxx = k ( x , x ) is the kernel matrix , and K ( · ) x = k ( x , · ) is the cross-covariance matrix between function values evaluated at the training and test inputs . The GP posterior can also be written as ( f | y ) ( · ) = f ( · ) + K ( · ) x ( Kxx + σ2I ) −1 ( y − f ( x ) − ε ) , ε ∼ N ( 0 , σ2I ) ( 2 ) where f ( · ) is the prior GP , and equality holds in distribution [ 47 , 48 ] . These expressions form the foundation upon which Gaussian-process-based methods in machine learning are built . Recent works have studied techniques for working with the expressions ( 1 ) and ( 2 ) when the input domain X is a Riemannian manifold , focusing both on defining general classes of kernels [ 2 ] , and on efficient computational techniques [ 47 , 48 ] . In this setting , namely for f : X → R , defining kernels already presents technical challenges : the seemingly-obvious first choice one might consider , namely the geodesic squared exponential kernel , is ill-defined in general [ 13 ] . We build on these recent developments to model vector fields on manifolds using GPs . We do not consider manifold-valued generalizations of Gaussian processes , for instance f : R→ X : various constructions in this setting are instead studied by Stroock [ 40 ] , Émery [ 11 ] , Mallasto and Feragen [ 29 ] , and Mallasto et al . [ 30 ] . To begin , we review what a vector field on a manifold actually is . 2.1 Vector Fields on Manifolds . Let X be a d-dimensional smooth manifold with TxX denoting its tangent space at x . Let TX = { ( x , v ) |x ∈ X , v ∈ TxX } be its tangent bundle , and let T ∗X = { ( x , φ ) |x ∈ X , φ ∈ T ∗xX } Basis coefficients Frame Vector field be its cotangent bundle—endow both spaces with the structure of smooth manifolds . Define the projection map projX : TX → X by projX ( x , v ) = x . A vector field on X is a map that assigns each point in X to a tangent vector attached to that point . More formally , a vector field is a cross-section , or simply a section , of the tangent bundle , which is a map f : X → TX , such that projX f ( x ) = x for all x.1 A vector field is called smooth if this map f is smooth . To represent a vector field on a manifold numerically , one must choose a basis in each tangent space , which serves as a coordinate system for vectors in the tangent space . On many manifolds it is impossible to choose these basis vectors in a way that they vary smoothly in space.2 This can be handled by working with local coordinates , or with bases that are non-smooth . Any chosen set of basis vectors is arbitrary , so objects constructed using them should not depend on this choice . Constructions that satisfy this property are called gauge independent . This notion is illustrated in Figure 1 , and will play a key role in the sequel . 2.2 Gaussian Vector Fields . Upon reflecting on the above considerations in the context of GPs , the first issue one encounters is that , for a random vector field f : X → TX , it is not clear what it means for finite-dimensional marginal distributions to be multivariate Gaussian given that f takes its values in a bundle rather than a vector space . The first step towards constructing Gaussian vector fields , therefore , involves adapting the notion of finite-dimensional marginal distributions appropriately . Definition 1 . Let X be a smooth manifold . We say that a random vector field f is Gaussian if for any finite set of locations x1 , . . . , xn ∈ X , the random vector ( f ( x1 ) , . . . , f ( xn ) ) ∈ Tx1X⊕ . . .⊕TxnX is Gaussian ( either in the sense of duality or in any basis : see Appendix A for details ) . This definition is near-identical to the Euclidean case : the only difference is that finite-dimensional marginals are now supported in a direct sum of tangent spaces , instead of Rn×d . With this definition , the standard multi-output GP properties , such as conditioning , carry over , virtually unmodified . Definition 1 is a natural choice : if we embed our manifold into Euclidean space , the induced GP is a vector-valued GP as defined in the beginning of Section 2 . 1A vector field is not the same as a map f̃ : X → Rd : an output value f ( x ) ∈ TX formally consists of both a copy of the input point x , and vector within the tangent space TxX at this point . This encodes the geometric structure of the underlying manifold . The algebraic requirement projX f ( x ) = x for all x ensures that the tangent vector chosen correctly corresponds to the point at which it is attached . 2If a smooth choice of basis vectors existed , it would define a smooth non-vanishing vector field . On the sphere , by the hairy ball theorem , all smooth vector fields vanish in at least one point , so no such bases exist . Proposition 2 . Let X be a manifold and emb : X → Rp be a smooth embedding . Let f be a Gaussian vector field ( as defined in Definition 1 ) , and let femb : emb ( X ) → Rp be its pushforward along the embedding . Then femb is a vector-valued Gaussian process in the Euclidean sense . All proofs in this work can be found in Appendix A . Having established the notion of a vector-valued Gaussian process on a smooth manifold , we proceed to deduce what mathematical objects play the role of a mean function and kernel , so that it is clear what ingredients are needed to construct and determine such a process . The former is clear : the mean of a Gaussian vector field should be an ordinary vector field , and will determine the mean vector at all finite-dimensional marginals . The kernel , on the other hand , is less obvious : because distinct tangent vectors live in different tanget spaces , it is unclear whether or not a Gaussian vector field is characterized by an appropriate notion of a matrix-valued kernel , or by something else . Now , we define the right notion for kernel in this setting . Definition 3 . We call a symmetric function k : T ∗X × T ∗X → R fiberwise bilinear if for all pairs of points x , x′ ∈ X k ( λαx + µβx , γx′ ) = λk ( αx , γx′ ) + µk ( βx , γx′ ) ( 3 ) holds for any αx , βx ∈ T ∗xX , γx′ ∈ T ∗x′X and λ , µ ∈ R , and positive semi-definite if for any set of covectors αx1 , . . . , αxn ∈ T ∗X , we have ∑n i=1 ∑n j=1 k ( αxi , αxj ) ≥ 0 . We call a symmetric fiberwise bilinear positive semi-definite function a cross-covariance kernel . This coordinate-free function should be viewed as analogous to ( ( x , v ) , ( x′ , v′ ) ) 7→ vTKx , x′v′ in the Euclidean setting , where v , v′ multiply the matrix-valued kernel from both sides . Its coordinate representation , which more closely matches the Euclidean case , will be explored in the sequel . To show that this is indeed the right notion , we prove the following result . Theorem 4 . The system of marginal distributions of a Gaussian vector field on a smooth manifold X is uniquely determined by a mean vector field µ : X → TX and a cross-covariance kernel k : T ∗X × T ∗X → R. Moreover , this correspondence is one-to-one . By virtue of defining and characterizing all Gaussian vector fields , Theorem 4 assures us the definition of a kernel introduced is the correct mathematical notion . The constructions presented here are all intrinsic or , in other words , coordinate-free , and do not involve the use of bases . To understand how to perform numerical calculations with these kernels we proceed to study their coordinate representations with respect to a specific choice of basis .

Summary: The paper suggests a construction of Covariance functions of multi-output Gaussian processes for vector fields on manifolds. The construction is independent of coordinate choices, but does depend on a Nash embedding of the manifold. In addition, the paper describes a clear mathematical framework of coordinate independent constructions of Gaussian processes. The approach is extended to variational approximations of such Gaussian processes, where the independence of the approach w.r.t. coordinate choices is again taken into consideration.