Gauge Equivariant Transformer
Year: 2022

Abstract:
Attention mechanism has shown great performance and efficiency in a lot of deep learning models, in which relative position encoding plays a crucial role. However, when introducing attention to manifolds, there is no canonical local coordinate system to parameterize neighborhoods. To address this issue, we propose an equivariant transformer to make our model agnostic to the orientation of local coordinate systems (i.e., gauge equivariant), which employs multi-head selfattention to jointly incorporate both position-based and content-based information. To enhance expressive ability, we adopt regular field of cyclic groups as feature fields in intermediate layers, and propose a novel method to parallel transport the feature vectors in these fields. In addition, we project the position vector of each point onto its local coordinate system to disentangle the orientation of the coordinate system in ambient space (i.e., global coordinate system), achieving rotation invariance. To the best of our knowledge, we are the first to introduce gauge equivariance to self-attention, thus name our model Gauge Equivariant Transformer (GET), which can be efficiently implemented on triangle meshes. Extensive experiments show that GET achieves state-of-the-art performance on two common recognition tasks.

heading: 1 Introduction
text: 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.

heading: 2 Related Work
text: 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.

heading: 3 Preliminaries
text: 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.

heading: 3.1 Basic Definitions
text: 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.

heading: 3.2 Gauge Equivariance
text: 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.

heading: 3.3 Riemannian Exponential Map
text: Transformers require encoding the relative position to propagate information. Note that in images, there is still a local point parameterization, which is so natural that one even does not realize it. For general manifolds, it is non-trivial to establish a parameterization criterion, at least in the local frame. Among many charting-based methods, the mostly used one is the Riemannian exponential map expp : TpM →M at point p, which is a mapping from the tangent plane to the surface. For a coordinate vector v ∈ TpM , the output of the Riemannian exponential map is obtained by moving the point p in the direction v along the geodesic curve with a distance of ‖v‖. Denoting the arrival point as q, we have expp(v) = q. Figure 1 visualizes the exponential map as well as some basic definitions introduced in Section 3.1. According to the inverse function theorem, expp is a local diffeomorphism so can avoid metric distortion at the point p. The inverse of Riemannian exponential map is the logarithmic map logp : M → TpM . Under the gauge wp, every point q in the neighborhood of p is associated with coordinate w−1p · logp(q).

heading: 3.4 Parallel Transport
text: The self-attention operation is essentially an aggregation of local neighboring features. However, the feature vectors of different points are from different spaces, thus they need to be parallel transported to the same feature space before being processed. For a tangent vector s at point q, we parallel transport it along the geodesic curve to another point p with respect to Levi-Civita connection [8], which preserves the norm of the vector. Levi-Civita connection is an isometry from TqM to TpM and determines the parallel transport of s, see Figure 2. In a gauge w, the parallel transport of tangent vector corresponds to a 2D rotation gwq→p ∈ SO(2) which contains the relative orientation of gauges in the neighborhood. For a general feature vector of ρ type, parallel transport can be expressed as s′w = ρ(g w q→p)sw.

heading: 3.5 Self-attention
text: Attention enables the model to selectively concentrate on the most relevant parts based on their content information [48, 53, 4, 22]. Consider a set of tokens t = {t1, t2, . . . , tT }, where ti ∈ RF . Attention is composed of three parts, namely query, key and value, denoted by Q : RF → RFQ , K : RF → RFK , and V : RF → RFV , respectively. When Q, K and V are from the same source, it is called self-attention. When there are multiple sets of Q, K and V ’s, it becomes multi-head attention.
The output of a multi-head self-attention transformer at node i is the linear transformation of the concatenation of the outputs at all the heads:
MHSA(t)i = WM (∥∥ h SA(t) (h) i ) , (2)
where ∥∥ is the vector concatenation operator. The single head attention output at head h is
SA(t) (h) i = T∑ j=1 α (h) ij V (h)(tj), (3)
where V (h) is the value function at the head h, and α(h)ij is attention score computed by
α (h) ij = S(K(h)(ti), Q (h)(tj))∑T
j′=1 S(K (h)(ti), Q(h)(tj′))
, (4)
where K(h), Q(h) and S are the key function, query function and score function, respectively.

heading: 4 The Proposed GET

heading: 4.1 Gauge Equivariant Self-Attention Layers
text: Suppose that the dimensions of input feature field f and output feature field f̃ are Cin and Cout, respectively. We define the gauge equivariant multi-head self-attention output at point p under the gauge w as
f̃w(p) = MHSA(f)w(p) = WM (∥∥ h SA(f)(h)w (p) ) , (5)
where WM is the linear transformation matrix. At the head h, the output is defined as
SA(f)(h)w (p) = ∫ ‖u‖<σ α(f)(h)p,quV (h) u (f ′ w(qu))du, (6)
where u = (u1, u2)T ∈ R2, qu = expp wp(u), f ′w(qu) is the numerical value of parallel transported feature vector from point qu to point p under the gauge w, and Vu is the value function incorporating the position information u through an encoder matrix WV (u) ∈ RCout×Cin , i.e.
f ′w(qu) = ρin(g w qu→p)fw(qu), Vu(f ′ w(qu)) = WV (u)f ′ w(qu). (7)
α is the attention score incorporating the content information, and is computed as:
α(f)(h)p,qu = S(K(h)(fw(p)), Q (h)(f ′w(qu)))∫ ‖v‖<σ S(K (h)(fw(p)), Q(h)(f ′w(qv)))dv . (8)
We propose to enforce the attention score to be gauge invariant and the value function to be gauge equivariant, to make the attention layer gauge equivariant. The details of constructing them are presented in Sections 4.3 and 4.4, respectively.

heading: 4.2 Extension of Regular Representation
text: In our model, the feature fields in the intermediate layers are all regular fields (i.e., whose type is regular representation). Regular representation is a special type of group representation of CN . If we use Θk to denote the rotation matrix with angle of k · 2π/N , then CN can be expressed as CN = {Θ0,Θ1, · · · ,ΘN−1}. For k = 0, 1, · · · , N − 1, the regular representation ρCNreg(Θk) is an N ×N cyclic permutation matrix which shifts the coordinates of feature vectors by k steps. Regular representation provides transformation matrices when rotating by angles of multiples of 2π/N , but feature vectors can go through any rotation in SO(2) during parallel transport. Figure 3 illustrates this issue by giving an example in R5 with respect to ρC5reg. We propose to extend the regular representation of CN by finding an orthogonal representation ρ̃N of SO(2), such that it behaves the same as regular representation for any element in CN , i.e.
∀Θ ∈ CN , ρ̃N (Θ) = ρCNreg(Θ). (9)
As ρCNreg takes different forms between odd values and even values of N , Theorem 1 shows that only odd N ’s are vaild in our model.
Theorem 1 (i) If N is even, there is no such real representation ρ̃N of SO(2) that satisfies Eqn. (9). (ii) If N is odd, there is a unique representation ρ̃N of SO(2) that satisfies Eqn. (9). (iii) The representation ρ̃N in (ii) is an orthogonal representation.
Here we only show our method for constructing ρ̃N in Theorem 1. According to group representation theory, regular representation ρCNreg can be decomposed into irreducible representations (irrep for short), i.e.,
ρCNreg(Θ) = Adiag ( ϕ0(Θ), ϕ1(Θ), · · · ,ϕN−1 2 (Θ) ) A−1, (10)
where ϕ0, · · · , ϕ(N−1)/2 are the irreps of CN , and A ∈ GL(N,R). The irreps of CN take the following form for odd N :
∀Θ ∈ CN , ϕ0(Θ) = 1, ϕk(Θ) = [ cos(kθ) − sin(kθ) sin(kθ) cos(kθ) ] , (11)
where θ ∈ [0, 2π) is the rotation angle of the matrix Θ, i.e.
Θ = [ cos θ − sin θ sin θ cos θ ] , (12)
and k = 1, · · · , N−12 . We extend the irreps to SO(2) as
∀Θ ∈ SO(2), ϕ̃0(Θ) = 1, ϕ̃k(Θ) = [ cos(kθ) − sin(kθ) sin(kθ) cos(kθ) ] , (13)
where k = 1, · · · , N−12 . By substituting the ϕ’s in Eqn. (10) with ϕ̃’s, we get that for ∀θ ∈ SO(2),
ρ̃N (Θ) = Adiag ( ϕ̃0(Θ), ϕ̃1(Θ), · · · , ϕ̃N−1 2 (Θ) ) A−1. (14)
Obviously the representation ρ̃N satisfies the condition Eqn. (9). In this way, one can apply ρ̃N (gwq→p) to feature vector of regular field during parallel transport.

heading: 4.3 Gauge Equivariant Value Function
text: Inspired by [11], we choose the value function to be the numerical value of the parallel transported feature vector multiplied by the value encoding matrix. For the value function to be gauge equivariant, the necessary and sufficient condition is that Eqn. (15) always holds for any Θ ∈ SO(2):
WV (Θ −1u) = ρout(Θ −1)WV (u)ρin(Θ). (15)
We propose a practical method to solve Eqn. (15). We first expand the Wv into taylor series:
WV (u) = W0 +W1u1 +W2u2 +W3u 2 1 +W4u1u2 +W5u 2 2 + · · · , (16)
whereWi ∈ RCout×Cin(i = 0, 1, · · · ) is the Taylor coefficient. Since we adopt regular representation in this paper, Eqn. (15) only needs to hold for Θ ∈ CN . Plugging Eqn. (16) into Eqn. (15), by comparing the coefficients, Wi’s need to satisfy that for any Θ ∈ CN ,
W0 = ρout(Θ −1)W0ρin(Θ), (17a)
cos(θ)W1 − sin(θ)W2 = ρout(Θ−1)W1ρin(Θ), (17b) sin(θ)W1 + cos(θ)W2 = ρout(Θ
−1)W2ρin(Θ), (17c) · · · .
To deal with the issue of having infinite terms in Eqn. (16), we may bypass it by simply truncating the Taylor series. We use the second order Taylor expansion and omit higher order terms, i.e.,
WV (u) ,W0 +W1u1 +W2u2 +W3u 2 1 +W4u1u2 +W5u 2 2. (18)
It is worth emphasizing that making truncations does not affect the equivariance property in the slightest, as the equations in (17) show the coupling characteristics.
Eqn. (17a) is the constraint on W0 in the order 0, Eqn. (17b) and Eqn. (17c) are the constraints on W1 and W2 in order 1, and there are three more equations in Eqn. (17) constraining on W3, W4 and W5 in the order 2. We can see that only the Wi’s in the same order are coupled together. This coupling property allows us not only to solve the equations in (17) in separate groups, but also to make truncations in Eqn. (16) without affecting the equivariance property.
After truncation, we can get a set of solution bases of Taylor coefficients {W̃ (1), · · · , W̃ (m)} by solving the first six linear equations in (17) which are separated into three independent groups, where m is the dimension of solution space. Each W̃ (i) is a tuple consisting of six components, W̃
(i) 0 , · · · , W̃ (i) 5 . The details in solving linear equations are provided in supplementary materials.
Then, the equivariant matrix basis W (i) has the following form:
W (i)(u) = W̃ (i) 0 + W̃ (i) 1 u1 + W̃ (i) 2 u2 + W̃ (i) 3 u 2 1 + W̃ (i) 4 u1u2 + W̃ (i) 5 u 2 2, (19)
which satisfies Eqn. (15) for all u. Their linear combination, ∑ ciW
(i), still meets Eqn. (15) and ci’s can be set as learnable parameters during training. With WV = ∑ ciW
(i), the value function in Eqn. (7) is exactly equivariant to gauge transformations at multiples of 2π/N .
It is remarkable that our method of solving the equivariance constraint Eqn. (15) is very general, as the solution process can be applied to any groups, including ρin and ρout. Especially, it can avoid solving analytic solutions when the group is very complex, like when the case is that the group is a high dimensional orthogonal group. In addition, compared to Fourier series used in [54], Taylor series is a better approximation in local neighborhoods. The omitted Taylor terms in Eqn. (18) is O(σ3), which is negligible when the radius σ is small enough. So GET could achieve the same performance with fewer parameters. In addition, we can avoid selecting radial profiles that introduce extra hyperparameters.

heading: 4.4 Gauge Invariant Attention Score
text: In implementation, the manifold is discretized to mesh for computer processing. The discretization details are provided in supplementary materials. Here we set the key and query function to be structurally the same as Graph Attention Network [49], i.e., K(h)(fw(p)) = W (h) K fw(p),
Figure 3: Illustration for the reason of extension. f(q) is a feature vector of type ρC5reg , which takes numerical value fw(q) ∈ R5 under gauge wq. Applying a gauge transformation with angle 2π/5 to w′q , f(q) takes another value fw′(q), which is a permutation of fw(q). The problem here is what value does f(q) takes after it is parallel transported to point p.
Figure 4: Local coordinate projection. xp is the position vector in the global coordinate system marked in red. For better illustration it is moved to the local coordinate system, marked in blue. In the local coordinate system xp is projected onto the directions of up, vp and np, respectively, and the lengths of three directed line segments (in green) form the input Xp.
Q(h)(f ′w(qu)) = W (h) Q f ′ w(qu), where W (h) K ∈ RN×Cin , W (h) Q ∈ RN×Cin . The score function is structurally similar to [49], which takes the following form:
S(K(·), Q(·)) = P (ReLU(K(·) +Q(·))). (20) Here, ReLU is the Nonlinear Rectified Unit acting on each element of the N dimensions, and P : RN → R is the average pooling function. The linear transformation matrices WK and WQ are required to satisfy the constraint in Eqn. (17a) on CN for K and Q to be gauge equivariant. After activation and pooling, the final attention score is gauge invariant.
With the gauge invariant attention score and gauge equivariant value function, the single head attention Eqn. (6) is gauge equivariant. For the multi-head attention to be gauge equivariant, the transformation matrix WM also needs to satisfy Eqn. (17a).

heading: 4.5 Rotation Invariance
text: The rotation invariance property of GET is accomplished by constructing a local coordinate system for every point and making use of the gauge equivariance property. As is shown in Figure 4, assuming that xp is the coordinate vector of p ∈ M in the global coordinate system, np is the corresponding normal vector, and the gauge wp is ascertained by principal axes up and vp. By projecting the raw data xp onto the local coordinate system, we get the local coordinate of point p: Xp = (〈xp, up〉 , 〈xp, vp〉 , 〈xp, np〉), which relies on wp but is invariant to the choice of global coordinate system. The insight is that X is actually a feature map whose corresponding feature field is associated with representation ρlocal as:
ρlocal(Θ) =
[ cos θ − sin θ 0 sin θ cos θ 0
0 0 1
] . (21)
If we feed the local coordinates into an SO(2) gauge equivariant model whose outputs are scalar fields, the result will be SO(3) rotation invariant.

heading: 4.6 Error Analysis
text: Following the conventions, GET stacks multiple self-attention layers with ReLU activation functions. Even if discretized on triangle meshes, GET is still exactly equivariant to gauge transformations at angles that are multiples of 2π/N . Theorem 2 Assume a GET ψ, whose types of input, intermediate, and output feature fields are ρlocal, kiρ CN reg and ρ0, respectively, where ki is the number of regular fields in the i th intermediate feature
field. Denote f as the input feature field on triangle mesh M , and the norm of the feature map is bounded by constant C. Gauges w and w′ are linked by transformation g. Further suppose that ψ is Lipschitz continuous with constant L, then we have:
(i) If gp ∈ CN for every mesh vertex p ∈M , then ψ(fw) = ψ(fw′).
(ii) For general gp ∈ SO(2), we have ‖ψ(fw)− ψ(fw′)‖ ≤ πLN C.
Theorem 2 provides a bound for gauge transformation with respect to any angles. Compared to non-equivariant models, GET decreases the equivariance error by a factor of 1/N . In experiments, we empirically show that the performance of our model increases as N increases.

heading: 5 Experiments
text: We conduct extensive experiments to evaluate the effectiveness of our model. We test the performance of our model on two deformable domain tasks, and conduct parameter sensitivity analysis and several ablation studies to make a comprehensive evaluation. Note that we use data preprocessing to precompute some useful preliminary values in order to save training time. The details of preprocessing can be found in supplementary materials.

heading: 5.1 Shape Classification
text: Our model used here is lightweight but powerful. The details of the architecture and training settings are provided in supplementary materials. Under the same setting, we compare our model with HSN [56], MeshCNN [24], GWCNN [20], GI [45] and MDGCNN [35], whose results are cited in [56]. As is shown in Table 5.2, our model achieves state-of-the-art performance on this dataset. GET significantly improves the previous state-of-the-art model HSN by 3.1% in classification accuracy. This may attribute to the attention mechanism and the intrinsic rotation invariance of our model, while all other models are CNNs and directly accepts the raw coordinates xyz as input. Also, HSN is the most parameter efficient model among the models we compared with. Our model consumes only 1/7 parameters of HSN (11K vs. 78K).

heading: 5.2 Shape Segmentation
text: A widely used task in 3D shape segmentation is Human Body Segmentation [30], in which the model is to predict body-part annotation for each sampled point. The dataset consists of 370 training models from MIT [50], FAUST [5], Adobe Fuse [1] and SCAPE [3] and 18 test models from SHREC07 [23]. The readers may refer to supplementary materials for details of neural network architecture and hyperparameters.
Table 2 reports the percentage of correctly classified vertices across all samples in the test set. The results of comparing models are cited from [56], [60] and [35]. Our model outperforms all these models in the segmentation task. GET consumes only about 1/15 parameters compared with MeshCNN (148K vs. 2.28M) but achieves higher performance.

heading: 5.3 Parameter Sensitivity
text: Order of the Group CN . The hyperparameter N is a key factor to the model equivariance since it controls both the dimension of regular field and the number of angles at which the our model is exactly equivariant. Also, Theorem 2 asserts that the equivariance error is bounded by a factor of 1/N compared to non-equivariant models. Here we study the effect of N on model accuracy while keeping parameter numbers roughly the same. The results of the Human Body Segmentation dataset with different N ’s are shown in Table 3. We can see that the model performance improves considerably as N increases and stabilizes finally.

heading: 5.4 Ablation Study
text: In this section, we perform a series of ablation studies to analyze individual parts of our model. All the experiments are carried out on the Human Body Segmentation dataset under the same setting as in Section 5.2. We evaluate the effectiveness of gauge equivariance, attention, local coordinate and parallel transport method, with the latter two experiments provided in supplementary materials.
Gauge Equivariance and Attention. To confirm the effectiveness of gauge equivariance property and attention mechanism, we design two baseline models with one not equivariant and the other based on convolution. For the non-equivariant baseline, we use Graph Attention Networks [49]. For the convolution-based model, we adopt a similar architecture as GET.
Table 4 shows that GET both benefits from the power of gauge equivariance and attention. We can see that both properties do contribute to the superiority of the model performance.

heading: 6 Conclusion
text: We propose GET, which firstly incorporates attention in gauge equivariance. GET introduces a new input, which is invariant to rotation of the global coordinate system. GET employs a new parallel transport approach, which is plausible for parallel transport between any two points. GET utilizes Taylor expansion in solving equivariant constraints, achieving better approximation ability. GET achieves state-of-the-art performances on several tasks and is efficient among the baselines.

heading: Acknowledgment
text: Zhouchen Lin was supported by the NSF China (No.s 61625301 and 61731018), NSFC Tianyuan Fund for Mathematics (No. 12026606) and Project 2020BD006 supported by PKU-Baidu Fund. Yisen Wang is partially supported by the National Natural Science Foundation of China under Grant 62006153, and Project 2020BD006 supported by PKU-Baidu Fund.

Review summaries:
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.

This paper addresses the main challenge in attention to manifolds, namely, that there are no coordinates to use as input in positional encoding. The network uses regular fields of cyclic groups and a parallel transport method for transporting features to these fields. Invariance to global rotations is achieved by taking only local position vectors. Experiments are performed on SHREC and Human Segmentation using a fraction of the usual transformer capacity.   

This paper considers the problem of introducing attention to data lying on manifolds (without a global symmetry). Since there is no canonical coordinate system available in such cases to parameterize neighborhoods, the authors instead use the notion of gauge equivariance, proposing a so-called gauge equivariant transformer. Some additional innovations in the intermediate layers are introduced to improve expressivity, and equivariance error is characterized. Experiments are reported on SHREC and Human Body Segmentation tasks, showing good performance with significant parameter efficiency compared to prior art. 

This paper proposes a gauge-equivariant transformer (GET) for manifolds. The challenging issue for self-attention on a manifold is the lack of canonical coordinate systems on a surface, which results in rotation ambiguity. To address it, this work introduces the regular field of cyclic groups as feature fields, proposes a novel parallel transport approach that puts the feature vectors in these fields, and develops the gauge-equivariant self-attention by adopting Taylor expansion in solving equivariant constraints. The proposed method is validated on the tasks of 3D shape classification and segmentation, achieving SOTA on SHREC dataset and the Human Body Segmentation dataset.

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.