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 authors propose a class of vector Gaussian processes (GP) on a class of Riemannian manifolds.  An important feature of the proposed GP is its ability to model vector fields on the Riemannian manifold. This is a broad application in physical sciences. One of the main novelties of the work is the development of gauge-equivariant kernels operators. Variation inference algorithms are extended for posterior inference.  Numerical examples demonstrate the practical performance of the proposed GP.