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.

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

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

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

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

abstractText: 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.

authors:
Brandon Anderson
Truong-Son Hy
Risi Kondor

title: Cormorant: Covariant molecular neural networks

venue: NeurIPS,

authors:
Dragomir Anguelov
Praveen Srinivasan
Daphne Koller
Sebastian Thrun
Jim Rodgers
James Davis

title: SCAPE: Shape Completion and Animation of People

venue: In SIGGRAPH

authors:
Dzmitry Bahdanau
Kyunghyun Cho
Yoshua Bengio

title: Neural machine translation by jointly learning to align and translate

venue: In ICLR,

authors:
Federica Bogo
Javier Romero
Matthew Loper
Michael J Black

title: FAUST: Dataset and Evaluation for 3D Mesh Registration

venue: In CVPR,

authors:
Davide Boscaini
Jonathan Masci
Emanuele Rodoià
Michael Bronstein

title: Learning Shape Correspondence with Anisotropic Convolutional Neural Networks

venue: In NeurIPS,

authors:
Michael M Bronstein
Joan Bruna
Taco S Cohen
Petar Veličković

title: Geometric deep learning: Grids, groups, graphs, geodesics, and gauges

venue: arXiv preprint arXiv:2104.13478,

authors:
Taco S Cohen
Maurice Weiler
Berkay Kicanaoglu
Max Welling

title: Gauge Equivariant Convolutional Networks and the Icosahedral CNN

authors:
Pim de Haan
Maurice Weiler
Taco S Cohen
Max Welling

title: Gauge Equivariant Mesh CNNs: Anisotropic convolutions on geometric graphs

authors:
Sander Dieleman
Jeffrey De Fauw
Koray Kavukcuoglu

title: Exploiting Cyclic Symmetry in Convolutional Neural Networks

venue: In ICML,

authors:
Alexey Dosovitskiy
Lucas Beyer
Alexander Kolesnikov
Dirk Weissenborn
Xiaohua Zhai
Thomas Unterthiner
Mostafa Dehghani
Matthias Minderer
Georg Heigold
Sylvain Gelly

title: An image is worth 16x16 words: Transformers for image recognition at scale

authors:
Carlos Esteves
Christine Allen-Blanchette
Ameesh Makadia
Kostas Daniilidis

title: Learning so (3) equivariant representations with spherical cnns

venue: In ECCV,

authors:
Carlos Esteves
Ameesh Makadia
Kostas Daniilidis

title: Spin-weighted spherical cnns

authors:
Carlos Esteves
Yinshuang Xu
Christine Allen-Blanchette
Kostas Daniilidis

title: Equivariant Multi-view Networks

venue: In ICCV,

authors:
Danielle Ezuz
Justin Solomon
Vladimir G Kim
Mirela Ben-Chen

title: GWCNN: A Metric Alignment Layer for Deep Shape Analysis

venue: In Computer Graphics Forum,

authors:
Fabian B Fuchs
Daniel E Worrall
Volker Fischer
Max Welling

title: SE (3)-transformers: 3D Roto-translation Equivariant Attention Networks

authors:
Zhengyang Geng
Meng-Hao Guo
Hongxu Chen
Xia Li
Ke Wei
Zhouchen Lin

title: Is attention better than matrix decomposition

authors:
Daniela Giorgi
Silvia Biasotti
Laura Paraboschi

title: Shape Retrieval Contest 2007: Watertight Models Track

venue: SHREC competition,

authors:
Rana Hanocka
Amir Hertz
Noa Fish
Raja Giryes
Shachar Fleishman
Daniel Cohen-Or

title: MeshCNN: a Network with an Edge

authors:
Lingshen He
Yuxuan Chen
Zhengyang Shen
Yiming Dong
Yisen Wang
Zhouchen Lin

title: Efficient equivariant network

authors:
Michael Hutchinson
Charline Le Lan
Sheheryar Zaidi
Emilien Dupont
Yee Whye Teh
Hyunjik Kim

title: Lietransformer: Equivariant self-attention for lie groups

authors:
Risi Kondor
Zhen Lin
Shubhendu Trivedi

title: Clebsch–Gordan Nets: a fully fourier space spherical convolutional neural network

authors:
Risi Kondor
Shubhendu Trivedi

title: On the Generalization of Equivariance and Convolution in Neural Networks to the Action of Compact Groups

authors:
Haggai Maron
Meirav Galun
Noam Aigerman
Miri Trope
Nadav Dym
Ersin Yumer
Vladimir G Kim
Yaron Lipman

title: Convolutional Neural Networks on Surfaces via Seamless Toric Covers

authors:
Jonathan Masci
Davide Boscaini
Michael Bronstein
Pierre Vandergheynst

title: Geodesic Convolutional Neural Networks on Riemannian Manifolds

venue: In ICCV Workshops,

authors:
Daniel Maturana
Sebastian Scherer

title: Voxnet: A 3D Convolutional Neural Network for Real-time Object Recognition

venue: In IROS,

authors:
Federico Monti
Davide Boscaini
Jonathan Masci
Emanuele Rodola
Jan Svoboda
Michael M Bronstein

title: Geometric Deep Learning on Graphs and Manifolds using Mixture Model CNNs

authors:
Nathanaël Perraudin
Michaël Defferrard
Tomasz Kacprzak
Raphael Sgier

title: Deepsphere: Efficient spherical convolutional neural network with healpix sampling for cosmological applications

venue: Astronomy and Computing,

authors:
Adrien Poulenard
Maks Ovsjanikov

title: Multi-directional Geodesic Neural Networks via Equivariant Convolution

authors:
Adrien Poulenard
Marie-Julie Rakotosaona
Yann Ponty
Maks Ovsjanikov

title: Effective Rotation-invariant Point CNN with Spherical Harmonics Kernels

venue: In 3DV,

authors:
Charles R Qi
Hao Su
Kaichun Mo
Leonidas J Guibas

title: Pointnet: Deep Learning on Point Sets for 3D Classification and Segmentation

authors:
Charles R Qi
Li Yi
Hao Su
Leonidas J Guibas

title: Pointnet++: Deep Hierarchical Feature Learning on Point Sets in a Metric Space

authors:
Yongming Rao
Jiwen Lu
Jie Zhou

title: Spherical fractal convolutional neural networks for point cloud recognition

venue: In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition,

authors:
Gernot Riegler
Ali Osman Ulusoy
Andreas Geiger

title: Octnet: Learning Deep 3D Representations at High Resolutions

authors:
David Romero
Erik Bekkers
Jakub Tomczak
Mark Hoogendoorn

title: Attentive group equivariant convolutional networks

authors:
David W Romero
Jean-Baptiste Cordonnier

title: Group Equivariant Stand-Alone Self-Attention

venue: For Vision. ICML,

authors:
David W Romero
Mark Hoogendoorn

title: Co-attentive equivariant neural networks: Focusing equivariance on transformations co-occurring in data

authors:
Zhengyang Shen
Lingshen He
Zhouchen Lin
Jinwen Ma

title: PDO-eConvs: Partial differential operator based equivariant convolutions

authors:
Ayan Sinha
Jing Bai
Karthik Ramani

title: Deep Learning 3D Shape Surfaces using Geometry Images

authors:
Hang Su
Subhransu Maji
Evangelos Kalogerakis
Erik Learned-Miller

title: Multi-view Convolutional Neural Networks for 3D Sape Recognition

authors:
Nathaniel Thomas
Tess Smidt
Steven Kearnes
Lusann Yang
Li Li
Kai Kohlhoff
Patrick Riley

title: Tensor field networks: Rotation-and translation-equivariant neural networks for 3D point clouds

authors:
Ashish Vaswani
Noam Shazeer
Niki Parmar
Jakob Uszkoreit
Llion Jones
Aidan N Gomez
Lukasz Kaiser
Illia Polosukhin

title: Attention is all you need

authors:
Daniel Vlasic
Ilya Baran
Wojciech Matusik
Jovan Popović

title: Articulated Mesh Animation from Multi-view Silhouettes

authors:
Chu Wang
Marcello Pelillo
Kaleem Siddiqi

title: Dominant Set Clustering and Pooling for Multi-view 3D Object Recognition

authors:
Peng-Shuai Wang
Yang Liu
Yu-Xiao Guo
Chun-Yu Sun
Xin Tong

title: O-cnn: Octree-based convolutional neural networks for 3D shape analysis

authors:
Xiaolong Wang
Ross Girshick
Abhinav Gupta
Kaiming He

title: Non-local neural networks

authors:
Maurice Weiler
Gabriele Cesa

title: General E (2)-equivariant Steerable CNNs

authors:
Maurice Weiler
Mario Geiger
Max Welling
Wouter Boomsma
Taco S Cohen

title: 3D Steerable CNNS: Learning Rotationally Equivariant Features in Volumetric Data

authors:
Ruben Wiersma
Elmar Eisemann
Klaus Hildebrandt

title: CNNs on Surfaces using Rotation

venue: Equivariant Features. TOG,

authors:
Daniel Worrall
Gabriel Brostow

title: Cubenet: Equivariance to 3D Rotation and Translation

authors:
Daniel E Worrall
Stephan J Garbin
Daniyar Turmukhambetov
Gabriel J Brostow

title: Harmonic Networks: Deep Translation and Rotation equivariance

authors:
Yufei Xu
Qiming Zhang
Jing Zhang
Dacheng Tao

title: Vitae: Vision transformer advanced by exploring intrinsic inductive bias

venue: arXiv preprint arXiv:2106.03348,

authors:
Yuqi Yang
Shilin Liu
Hao Pan
Yang Liu
Xin Tong

title: PFCNN: Convolutional Neural Networks on 3D Surfaces using Parallel Frames

authors:
Tan Yu
Jingjing Meng
Junsong Yuan

title: Multi-view Harmonized Bilinear Network for 3D Object Recognition

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.

title: Gauge Equivariant Transformer