Gauge Equivariant Transformer

1 Introduction . Recently , Transformer has dominated the area of Natural Language Processing [ 48 ] . Its key advantage over previous methods is its ability to attend to the most relevant part in a given context . This is largely attributed to its self-attention operator , which computes the similarity between representations of words in sequences in the form of attention scores . Because of the superiority , researchers start to apply Transformer to other learning areas , including Computer Vision [ 26 , 53 , 16 , 59 ] and Graphs [ 49 ] . In this work , we aim at applying Transformer to manifolds . Unlike regular data , such as images , where each neighbor owns a clearly quantified relative position to its center in a canonical coordinate system , irregular data do not have a uniquely defined local coordinate system for the neighbors , resulting in the problem of orientation ambiguity , which directly obstructs the Transformer to numerically intake the relative position information . Several works have been proposed to deal with the rotation ambiguity problem , in which a promising way is to exploit gauge equivariance . While most of them are not rotation invariant to global ∗Work was done during an internship at JD Explore Academy . ∗Equal contribution . Sorted by tossing the coin . †Corresponding author . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . coordinate system , all of them are established on convolution , i.e. , equal attention to neighboring points and neglect to content-based information . So it is desirable to propose a gauge equivariant transformer with the support of rotation invariance . In this paper , we propose Gauge Equivariant Transformer , named GET for short , which employs multi-head self-attention to simultaneously utilize position-based and content-based information , and is both gauge equivariant and rotation invariant . To achieve rotation invariance , we first project xyz coordinates in a global coordinate system onto a local coordinate frame , and then design equivariant transformers to overcome the orientation ambiguity problem of local coordinate systems . We adopt the regular field proposed in [ 13 ] as feature fields of intermediate layers , since the representation of regular field commutes with element-wise activation functions . After that , we propose a novel method to accommodate parallel transport of feature vectors in regular field with any rotation angles . Since we adopt regular fields in intermediate layers , we make a relaxation such that they are equivariant only for gauge transformations of angles that are multiples of 2π/N . Exact equivariance can be guaranteed for gauge transformations at multiples of 2π/N , and an equivariance error bound can be obtained for all other angles . In experiments , our model shows better performance and greater parameter efficiency than all baseline methods . Our contributions can be summarized as follows : • We propose GET , which incorporates attention and achieves both gauge equivariance and rotation invariance with superior expressive power . GET is mathematically proven to be exactly equivariant on angles that are multiples of 2π/N ( N ∈ N∗ ) , and an equivariance error bound is derived for other angles to guarantee the overall approximate equivariance property . • We carefully design the model input to ensure that it is irrelevant to the global coordinate system , only depending on the choice of gauge . Our model achieves rotation invariance with the assistance of gauge equivariance . • We propose a novel method to parallel transport the feature vectors in the regular field by extending the representation of a cyclic group to any angle rotation group . Compared to previous methods using truncation or interpolation , our extension can preserve more geometric information . • We elevate the model performance by designing a new approach which incorporates Taylor expansion in solving the equivariance constraint , which has a better approximation ability in local neighborhoods . • We confirm the superiority of our model via extensive experiments . Our model outperforms the HSN model on the SHREC dataset by 3.1 % accuracy , and outperforms the MeshCNN model on the Human Body Segmentation dataset by 0.3 % accuracy with much fewer parameters , presenting state-of-the-art performance . 2 Related Work . Geometric Deep Learning . Geometric deep learning is an emerging field concerning with adapting neural networks to various data types [ 7 ] , especially on irregular data . For research works on modeling curved surfaces , common methods include view-based methods [ 46 , 61 , 51 ] and volumetric methods [ 32 , 40 , 52 ] . To boost efficiency , some works define convolution on point clouds directly [ 37 , 38 ] , but they are vulnerable to pose change since the coordinate inputs are dependent of the global coordinate system . So it is highly desired to develop models that solely intake geometric information of surfaces . Approaches that merely utilize intrinsic information of surfaces are called intrinsic methods . They use local parameterization to assign each neighboring point with a coordinate for information aggregation . A seminal work is Geodesic CNNs [ 31 ] , which uses an exponential map to parameterize each local neighborhood and takes the maximum response across multiple choices of local coordinate orientation . While taking the maximum response direction discards the orientation information of feature maps , as an alternative , aligning local coordinate with principle curvature direction is another approach to deal with the ambiguity problem [ 33 , 6 ] . But this approach can only be applied in limited cases as the curvature direction may be ill-defined at some points or even areas of curved surfaces . MDGCNN [ 35 ] and PFCNN [ 60 ] describes features by the so called directional functions . But both of them adopt scalar equivariant kernels , resulting in limited expressive power . Equivariant Deep Learning . Success of CNNs has been attributed to translation equivariance , which inspires researchers to implement more powerful equivariant models , including equivariance of planar rotation [ 12 , 15 , 13 , 58 , 54 , 44 , 25 ] , 3D space rotation [ 57 , 19 , 47 , 36 , 55 , 35 , 28 , 2 , 39 ] , sphere rotation [ 9 , 17 , 34 , 18 ] , and so on . All above works are about equivariance on homogeneous space [ 29 , 10 ] . Cohen et al . [ 11 ] further extend equivariance to manifolds , in which they identify a new type of equivariance called gauge equivariance . The models in [ 56 , 14 ] are successful extensions of gauge equivariant CNNs on mesh surfaces . However , their model suffers from changes in the orientation of global coordinate system . Also , there are works proposed for equivariant attention . Romero et al . [ 43 ] propose co-attentive equivariant networks , which effectively attends to co-occurring transformations . Romero et al . [ 41 ] further propose attentive group equivariant convolutional networks . Besides this , transformers have also been applied to group equivariant networks , where Fuchs et al . [ 21 ] do so via irreducible representations , Hutchinson et al . [ 27 ] via Lie algebra , and Romero et al . [ 42 ] via generalization of position encodings . All the models above are equivariant to symmetric groups , while currently gauge equivariant attention is still lacking . 3 Preliminaries . Unlike regular data , in which coordinates ( or pixels ) are aligned in a global frame , there is no such specific frame on general manifolds . To begin with , we briefly review and define some mathematical concepts . 3.1 Basic Definitions . We restrict our attention to 2D manifolds in 3D Euclidean space . Consider a 2D smooth orientable manifold M . For a point p in M , denote its tangent plane as TpM . Each point in TpM can be associated with a coordinate by specifying a coordinate system . Namely , we can parameterize the tangent plane TpM with a pointwise linear mapping wp : R2 → TpM , which is defined as the gauge w at point p [ 11 ] . The gauge of manifold M is the set containing gauges at every point in M . For planar data , a feature map is the set of features located at different positions on a plane . Similarly , a feature field on a surface is a set of geometric quantities at different positions of the surface . Note that these two concepts are similar but not the same . From the perspective of geometric deep learning , a feature map is defined as numerical values of geometric quantities that may be gauge dependent , while a feature field refers to geometric quantities themselves that are gauge independent . For example , each point of the surface can be assigned with a tangent vector as its feature vector , all of which form a feature field . As is shown in Figure 1 , the tangent vector v itself is a geometric quantity , which stays the same regardless of arbitrary gauge selection but takes different numerical values in different gauges following an underlying rule . We use f to denote the feature field of a manifold , fw : M → Rn denotes the feature map under the gauge w and fw ( p ) denotes the feature map evaluated at point p. Different gauges can be linked by gauge transformations . The gauge transformation at point p is a frame transformation : gp ∈ SO ( 2 ) , where SO ( 2 ) is the special orthogonal group consisting of all 2D rotation transformation matrices . A new gauge w′p can be produced by applying gauge transformation gp to the original gauge wp , i.e. , w′p = gp · wp . Gauge transformation is usually characterized by group representations . Group representation is a mapping ρ : G→ GL ( n , R ) where GL ( n , R ) is the group of invertible n×n matrices , and ρ meets the condition ρ ( g1 ) ρ ( g2 ) = ρ ( g1g2 ) , where g1 , g2 ∈ G are the elements of the group , g1g2 are element product defined on the group , and ρ ( g1 ) ρ ( g2 ) is matrix multiplication . Therefore , after applying the gauge transformation gp , the feature vector value fw ( p ) transforms to fw′ ( p ) = ρ ( g−1p ) fw ( p ) . Here ρ is a group representation of SO ( 2 ) which is called the type of the feature vector . If all the feature vectors share the same type ρ , the feature field is called a ρ-field and ρ is called the representation type of the field . The above definitions can also be at the manifold level , i.e. , fw′ = ρ ( g−1 ) fw . The notation kρ , where k is a positive integer , refers to the group representation whose output is k-blocks diagonal matrix with each block equals to ρ . In particular , if the representation of a feature field is ρ ( g ) = 1 , then the feature field becomes scalar field , denoted as ρ0 . 3.2 Gauge Equivariance . For a function φ , its input is a feature map fw , where f is a ρin-field , in order to make φ gauge equivariant , and its output f̃w should be a feature map , where f̃ is a ρout-field . When φ is a layer of a neural network , gauge equivariance implies that φ does not rely on the gauge in the forward process . Suppose that there are two gauges w and w′ linked by a gauge transformation g : w′ = g · w , we have fw′ = ρin ( g−1 ) fw since f is a ρin-field . Gauge equivariance means that the outputs f̃w = φ [ fw ] and f̃w′ = φ [ fw′ ] are linked by the ρout representation of the same transformation g , i.e . f̃w′ = ρout ( g −1 ) f̃w . Finally , we get : ρout ( g −1 ) φ [ fw ] = φ [ ρin ( g −1 ) fw ] . ( 1 ) To sum up , a function φ is gauge equivariant if the above equation always holds for any feature field f , gauge w and transformation g .

Summary: The paper presents a gauge equivariant transformer for geometric learning. The gauge equivariance aims to solve the orientation ambiguity when doing the convolution or local aggregation operations within a neighborhood. Specifically, this work achieves the gauge equivariant self-attention within the neighborhood defined by mesh data. Different from existing gauge equivariant layers, the paper proposes to accommodate the parallel transport in the regular field. As a result, the gauge equivariance is up to a certain error bound. Further, this work validates the effectiveness on two datasets including shape classification (SHREC) and human body segmentation.