Gauge Equivariant Mesh CNNs: Anisotropic convolutions on geometric graphs

A common approach to define convolutions on meshes is to interpret them as a graph and apply graph convolutional networks ( GCNs ) . Such GCNs utilize isotropic kernels and are therefore insensitive to the relative orientation of vertices and thus to the geometry of the mesh as a whole . We propose Gauge Equivariant Mesh CNNs which generalize GCNs to apply anisotropic gauge equivariant kernels . Since the resulting features carry orientation information , we introduce a geometric message passing scheme defined by parallel transporting features over mesh edges . Our experiments validate the significantly improved expressivity of the proposed model over conventional GCNs and other methods . 1 INTRODUCTION . Convolutional neural networks ( CNNs ) have been established as the default method for many machine learning tasks like speech recognition or planar and volumetric image classification and segmentation . Most CNNs are restricted to flat or spherical geometries , where convolutions are easily defined and optimized implementations are available . The empirical success of CNNs on such spaces has generated interest to generalize convolutions to more general spaces like graphs or Riemannian manifolds , creating a field now known as geometric deep learning ( Bronstein et al. , 2017 ) . A case of specific interest is convolution on meshes , the discrete analog of 2-dimensional embedded Riemannian manifolds . Mesh CNNs can be applied to tasks such as detecting shapes , registering different poses of the same shape and shape segmentation . If we forget the positions of vertices , and which vertices form faces , a mesh M can be represented by a graph G. This allows for the application of graph convolutional networks ( GCNs ) to processing signals on meshes . ∗Equal Contribution †Qualcomm AI Research is an initiative of Qualcomm Technologies , Inc . However , when representing a mesh by a graph , we lose important geometrical information . In particular , in a graph there is no notion of angle between or ordering of two of a node ’ s incident edges ( see figure 1 ) . Hence , a GCNs output at a node p is designed to be independent of relative angles and invariant to any permutation of its neighbours qi ∈ N ( p ) . A graph convolution on a mesh graph therefore corresponds to applying an isotropic convolution kernel . Isotropic filters are insensitive to the orientation of input patterns , so their features are strictly less expressive than those of orientation aware anisotropic filters . To address this limitation of graph networks we propose Gauge Equivariant Mesh CNNs ( GEM-CNN ) , which minimally modify GCNs such that they are able to use anisotropic filters while sharing weights across different positions and respecting the local geometry . One obstacle in sharing anisotropic kernels , which are functions of the angle θpq of neighbour q with respect to vertex p , over multiple vertices of a mesh is that there is no unique way of selecting a reference neighbour q0 , which has the direction θpq0 = 0 . The reference neighbour , and hence the orientation of the neighbours , needs to be chosen arbitrarily . In order to guarantee the equivalence of the features resulting from different choices of orientations , we adapt Gauge Equivariant CNNs ( Cohen et al. , 2019b ) to general meshes . The kernels of our model are thus designed to be equivariant under gauge transformations , that is , to guarantee that the responses for different kernel orientations are related by a prespecified transformation law . Such features are identified as geometric objects like scalars , vectors , tensors , etc. , depending on the specific choice of transformation law . In order to compare such geometric features at neighbouring vertices , they need to be parallel transported along the connecting edge . In our implementation we first specify the transformation laws of the feature spaces and compute a space of gauge equivariant kernels . Then we pick arbitrary reference orientations at each node , relative to which we compute neighbour orientations and compute the corresponding edge transporters . Given these quantities , we define the forward pass as a message passing step via edge transporters followed by a contraction with the equivariant kernels evaluated at the neighbour orientations . Algorithmically , Gauge Equivariant Mesh CNNs are therefore just GCNs with anisotropic , gauge equivariant kernels and message passing via parallel transporters . Conventional GCNs are covered in this framework for the specific choice of isotropic kernels and trivial edge transporters , given by identity maps . In Sec . 2 , we will give an outline of our method , deferring details to Secs . 3 and 4 . In Sec . 3.2 , we describe how to compute general geometric quantities , not specific to our method , used for the computation of the convolution . In our experiments in Sec . 6.1 , we find that the enhanced expressiveness of Gauge Equivariant Mesh CNNs enables them to outperform conventional GCNs and other prior work in a shape correspondence task . 2 CONVOLUTIONS ON GRAPHS WITH GEOMETRY . We consider the problem of processing signals on discrete 2-dimensional manifolds , or meshes M . Such meshes are described by a set V of vertices in R3 together with a set F of tuples , each consisting of the vertices at the corners of a face . For a mesh to describe a proper manifold , each edge needs to be connected to two faces , and the neighbourhood of each vertex needs to be homeomorphic to a disk . Mesh M induces a graph G by forgetting the coordinates of the vertices while preserving the edges . A conventional graph convolution between kernel K and signal f , evaluated at a vertex p , can be defined by ( K ? f ) p = Kselffp + ∑ q∈Np Kneighfq , ( 1 ) where Np is the set of neighbours of p in G , and Kself ∈ RCin×Cout and Kneigh ∈ RCin×Cout are two linear maps which model a self interaction and the neighbour contribution , respectively . Importantly , graph convolution does not distinguish different neighbours , because each feature vector fq is multiplied by the same matrix Kneigh and then summed . For this reason we say the kernel is isotropic . Consider the example in figure 1 , where on the left and right , the neighbourhood of one vertex p , containing neighbours q ∈ Np , is visualized . An isotropic kernel would propagate the signal from the neighbours to p in exactly the same way in both neighbourhoods , even though the neighbourhoods are geometrically distinct . For this reason , our method uses direction sensitive ( anisotropic ) kernels instead of isotropic kernels . Anisotropic kernels are inherently more expressive than isotropic ones which is why they are used universally in conventional planar CNNs . Algorithm 1 Gauge Equivariant Mesh CNN layer Input : mesh M , input/output feature types ρin , ρout , reference neighbours ( qp0 ∈ Np ) p∈M . Compute basis kernels Kiself , K i neigh ( θ ) B Sec . 3 Initialise weights wiself and w i neigh . For each neighbour pair , p ∈M , q ∈ Np : B App . A. compute neighbor angles θpq relative to reference neighbor compute parallel transporters gq→p Forward ( input features ( fp ) p∈M , weights wiself , w i neigh ) : f ′p ← ∑ i w i selfK i selffp + ∑ i , q∈Npw i neighK i neigh ( θpq ) ρin ( gq→p ) fq We propose the Gauge Equivariant Mesh Convolution , a minimal modification of graph convolution that allows for anisotropic kernels K ( θ ) whose value depends on an orientation θ ∈ [ 0 , 2π ) .1 To define the orientations θpq of neighbouring vertices q ∈ Np of p , we first map them to the tangent plane TpM at p , as visualized in figure 1 . We then pick an arbitrary reference neighbour q p 0 to determine a reference orientation2 θpqp0 : = 0 , marked orange in figure 1 . This induces a basis on the tangent plane , which , when expressed in polar coordinates , defines the angles θpq of the other neighbours . As we will motivate in the next section , features in a Gauge Equivariant CNN are coefficients of geometric quantities . For example , a tangent vector at vertex p can be described either geometrically by a 3 dimensional vector orthogonal to the normal at p or by two coefficients in the basis on the tangent plane . In order to perform convolution , geometric features at different vertices need to be linearly combined , for which it is required to first “ parallel transport ” the features to the same vertex . This is done by applying a matrix ρ ( gq→p ) ∈ RCin×Cin to the coefficients of the feature at q , in order to obtain the coefficients of the feature vector transported to p , which can be used for the convolution at p. The transporter depends on the geometric type ( group representation ) of the feature , denoted by ρ and described in more detail below . Details of how the tangent space is defined , how to compute the map to the tangent space , angles θpq , and the parallel transporter are given in Appendix A . In combination , this leads to the GEM-CNN convolution ( K ? f ) p = Kselffp + ∑ q∈Np Kneigh ( θpq ) ρ ( gq→p ) fq ( 2 ) which differs from the conventional graph convolution , defined in Eq . 1 only by the use of an anisotropic kernel and the parallel transport message passing . We require the outcome of the convolution to be equivalent for any choice of reference orientation . This is not the case for any anisotropic kernel but only for those which are equivariant under changes of reference orientations ( gauge transformations ) . Equivariance imposes a linear constraint on the kernels . We therefore solve for complete sets of “ basis-kernels ” Kiself and K i neigh satisfying this constraint and linearly combine them with parameters wiself and w i neigh such that Kself = ∑ i w i selfK i self and Kneigh = ∑ i w i neighK i neigh . Details on the computation of basis kernels are given in section 3 . The full algorithm for initialisation and forward pass , which is of time and space complexity linear in the number of vertices , for a GEM-CNN layer are listed in algorithm 1 . Gradients can be computed by automatic differentiation . The GEM-CNN is gauge equivariant , but furthermore satisfies two important properties . Firstly , it depends only on the intrinsic shape of the 2D mesh , not on the embedding of the mesh in R3 . Secondly , whenever a map from the mesh to itself exists that preserves distances and orientation , the convolution is equivariant to moving the signal along such transformations . These properties are proven in Appendix D and empirically shown in Appendix F.2 . 1In principle , the kernel could be made dependent on the radial distance of neighboring nodes , by Kneigh ( r , θ ) = F ( r ) Kneigh ( θ ) , where F ( r ) is unconstrained and Kneigh ( θ ) as presented in this paper . As this dependency did not improve the performance in our empirical evaluation , we omit it . 2Mathematically , this corresponds to a choice of local reference frame or gauge . 3 GAUGE EQUIVARIANCE & GEOMETRIC FEATURES . On a general mesh , the choice of the reference neighbour , or gauge , which defines the orientation of the kernel , can only be made arbitrarily . However , this choice should not arbitrarily affect the outcome of the convolution , as this would impede the generalization between different locations and different meshes . Instead , Gauge Equivariant Mesh CNNs have the property that their output transforms according to a known rule as the gauge changes . Consider the left hand side of figure 2 ( a ) . Given a neighbourhood of vertex p , we want to express each neighbour q in terms of its polar coordinates ( rq , θq ) on the tangent plane , so that the kernel value at that neighbour Kneigh ( θq ) is well defined . This requires choosing a basis on the tangent plane , determined by picking a neighbour as reference neighbour ( denoted q0 ) , which has the zero angle θq0 = 0 . In the top path , we pick qA as reference neighbour . Let us call this gauge A , in which neighbours have angles θAq . In the bottom path , we instead pick neighbour qB as reference point and are in gauge B . We get a different basis for the tangent plane and different angles θBq for each neighbour . Comparing the two gauges , we see that they are related by a rotation , so that θBq = θ A q − θAqB . This change of gauge is called a gauge transformation of angle g : = θ A qB . In figure 2 ( a ) , we illustrate a gauge equivariant convolution that takes input and output features such as gray scale image values on the mesh , which are called scalar features . The top path represents the convolution in gauge A , the bottom path in gauge B . In either case , the convolution can be interpreted as consisting of three steps . First , for each vertex p , the value of the scalar features on the mesh at each neighbouring vertex q , represented by colors , is mapped to the tangent plane at p at angle θq defined by the gauge . Subsequently , the convolutional kernel sums for each neighbour q , the product of the feature at q and kernel K ( θq ) . Finally the output is mapped back to the mesh . These three steps can be composed into a single step , which we could call a geometric convolution , mapping from input features on the mesh to output features on the mesh . The convolution is gauge equivariant if this geometric convolution does not depend on the gauge we pick in the interim , so in figure 2 ( a ) , if the convolution in the top path in gauge A has same result the convolution in the bottom path in gauge B , making the diagram commute . In this case , however , we see that the convolution output needs to be the same in both gauges , for the convolution to be equivariant . Hence , we must have that K ( θq ) = K ( θq − g ) , as the orientations of the neighbours differ by some angle g , and the kernel must be isotropic . As we aim to design an anisotropic convolution , the output feature of the convolution at p can , instead of a scalar , be two numbers v ∈ R2 , which can be interpreted as coefficients of a tangent feature vector in the tangent space at p , visualized in figure 2 ( b ) . As shown on the right hand side , different gauges induce a different basis of the tangent plane , so that the same tangent vector ( shown on the middle right on the mesh ) , is represented by different coefficients in the gauge ( shown on the top and bottom on the right ) . This gauge equivariant convolution must be anisotropic : going from the top row to the bottom row , if we change orientations of the neighbours by −g , the coefficients of the output vector v ∈ R2 of the kernel must be also rotated by −g . This is written as R ( −g ) v , where R ( −g ) ∈ R2×2 is the matrix that rotates by angle −g . Vectors and scalars are not the only type of geometric features that can be inputs and outputs of a GEM-CNN layer . In general , the coefficients of a geometric feature of C dimensions changes by an invertible linear transformation ρ ( −g ) ∈ RC×C if the gauge is rotated by angle g. The map ρ : [ 0 , 2π ) → RC×C is called the type of the geometric quantity and is formally known as a group representation of the planar rotation group SO ( 2 ) . Group representations have the property that ρ ( g + h ) = ρ ( g ) ρ ( h ) ( they are group homomorphisms ) , which implies in particular that ρ ( 0 ) = 1 and ρ ( −g ) = ρ ( g ) −1 . For more background on group representation theory , we refer the reader to ( Serre , 1977 ) and , specifically in the context of equivariant deep learning , to ( Lang & Weiler , 2020 ) . From the theory of group representations , we know that any feature type can be composed from “ irreducible representations ” ( irreps ) . For SO ( 2 ) , these are the one dimensional invariant scalar representation ρ0 and for all n ∈ N > 0 , a two dimensional representation ρn , ρ0 ( g ) = 1 , ρn ( g ) = ( cosng 9 sinng sinng cosng ) . where we write , for example , ρ = ρ0 ⊕ ρ1 ⊕ ρ1 to denote that representation ρ ( g ) is the direct sum ( i.e . block-diagonal stacking ) of the matrices ρ0 ( g ) , ρ1 ( g ) , ρ1 ( g ) . Scalars and tangent vector features correspond to ρ0 and ρ1 respectively and we have R ( g ) = ρ1 ( g ) . The type of the feature at each layer in the network can thus be fully specified ( up to a change of basis ) by the number of copies of each irrep . Similar to the dimensionality in a conventional CNN , the choice of type is a hyperparameter that can be freely chosen to optimize performance .

Summary: The work presents a novel message passing GNN operator for meshes that is equivariant under gauge transformations. It achieves that by parallel transporting features along edges and spanning a space of gauge equivariant kernels. Further, a DFT-based non-linearity is proposed, which preserves the equivariance in the limit of sampling density. The method is evaluated on an MNIST toy experiment and the Faust shape correspondence task.