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 authors address an important gap in existing literature at the intersection of self-attention networks (transformers) and the equivariant neural networks -- how to apply self-attention to manifolds or meshes (discretized versions of manifolds). This has immediate applications in the area of 3D shapes which are 2D manifolds embedded in 3D. The key challenge is that of gauge-equivariance as there is no canonical coordinate system and we would like the self-attention blocks to be equivariant to the choice of the gauge. To this end, the authors build a provably (approximate) gauge-invariant score function and a gauge-equivariant value function. The experiments show the importance of the proposed methods.