Venue: NeurIPS 2021 Conference

Reviews:

final_score: rating: 8: Top 50% of accepted NeurIPS papers, clear accept
review_content: summary: 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.



main_review: 

This is the first gauge equivariant self-attention on manifolds.
Several contributions of the paper are novel: the parallel transport, the discretization of local rotations and their interpolation, and the approach for solving equivariant constraints. 

The key in the authors' contribution are the cyclic fields obtained by the discretization of the local mesh angles. 

Figure 3 is very illuminating in the parallel transport of feature vectors.

The computation of the value function follows the icosahedral CNN paper by Cohen.

Theorem 2 provides a bound with respect to any angle (not just the discretized ones). 

The appendix is as useful to read as the paper and the derivations are a delight to follow. 

The main puzzle for the reader is in the ablation study. Against the trend, it seems that it is mainly the gauge equivariance in human segmentation that matters rather than the self-attention (performance difference 0.3\%). The explanation though 
"Without attention, the convolution-based baseline performs as well as the second best baseline (MeshCNN)" does not reflect the table. Please elaborate. 

In a similar tone, please explain the challenges of gauge equivariance in self-attention rather than just gauge equivariance in meshes. After following the math and the lack of any difference in performance I am not sure whether the self-attention really matters and whether it is worth writing a paper about it.

POST-REBUTTAL: My main question was answered here and in the rest of the rebuttal. Stay with initial mark.
limitations_and_societal_impact: See above.
review_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: summary: 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. 
main_review: This paper is concerned with the introduction of attention for operating on data lying on manifolds. General manifolds lack a canonical coordinate system to parameterize local neighborhoods. Since attention-based methods rely on such information for positional encodings, the paper considers a model that is gauge-equivariant (and thus agnostic to the orientation of the local coordinate system), introducing what the authors dub as a "Gauge Equivariant Transformer." 

The starting point is the influential work on gauge-equivariant CNNs, at the core of which is a convolution like operation that is equivariant to local gauge transformations, thus providing principled manifold CNNs. However, the use of convolution also implies that all points in a neighborhood are accorded equal importance. This paper deals with the extension of this machinery to attention mechanisms -- such that the aggregation only depends on the intrinsic geometry. Further, an additional requirement pursued is global rotational invariance (which might not always be the case in gauge-equivariant CNNs for general manifolds). This work can also be considered an addition to a growing body of work on group-equivariant transformer models that have been proposed recently. 

Attention is restricted to 2-manifolds (embedded in 3D). The basic definitions and preliminaries are introduced in section 3.1, where quantities such as gauges, tangent vectors, feature fields are defined, followed by explication on gauge transformations. The condition for gauge equivariance is specified in section 3.2, while the exponential and log map are covered in 3.3. Since self-attention amounts to aggregation of neighborhood features, in a general manifold, feature vectors of different points reside in different spaces, and thus have to be parallel transported to the same feature space where they can be processed. This process is described in section 3.4. The usual self-attention formulation is described in 3.5.

Following the above background, a gauge equivariant self-attention layer is defined where the value function is required to be gauge-equivariant and the attention score gauge-invariant. The value function is taken to be the numerical outcome of the parallel transported feature vector times the value matrix, subject to the usual gauge equivariant condition. The resulting equation is then reduced to a system of linear equations using a Taylor expansion. Then, using a similar structure as the graph attention network, the key and query are combined to give a gauge-invariant attention score. Finally experimental results are shown on SHREC and Human Body Segmentation datasets (the manifolds are represented as triangular meshes), the results show that the proposed model gives superior performance while using a fraction of parameters (compared to the baseline models). Results are also shown for the choice of N (the order of the cyclic group, which is used for equivariance in the intermediate layers). 

While I have tried to verify all the details in the paper, my chief concern is that the paper can sometimes be hard to parse (which is also a reason for my lower confidence score). So, even with quite some familiarity with the related literature, I don't feel fully confident about some aspects of the paper (such as solving for the gauge equivariant condition). Nonetheless, the overall formulation seems correct, and the appendix fills in a lot of the missing details. Despite these concerns, I think the paper presents good results on two benchmarks and makes a step towards more principled self-attention networks for data on manifolds. 


Minor comments:

- The paper keeps using the phrase "orientation ambiguity problem", but doesn't quite explain what they mean by it, perhaps due to the presumption that since the paper is on gauge equivariance, it will be clear. However, ii would improve readability if this is clarified concretely in the context of attention right in the beginning. 
- The paper has many grammatical errors and typos, which the authors might want to iron out. I am listing some such examples below.
- Line 34: "neglection to content-based information" --> Neglect of content-based information?
- Line 49: "which initiatively" is this a typo? 
- Line 63: "emerging field concerning on adapting neural networks on various data types" --> "emerging field concerned with adapting neural networks to various data types"
- Line 64: "For researches" --> For research on modeling/for research works on modeling
- Line 160: "selectively concentrating on the most relevant parts" --> "selectively concentrate on the most relevant parts"
- Line 211: "feature vector multiply by the value encoding matrix" --> "feature vector multiplied by the value encoding matrix"
- Line 261: "the resulted one will be" --> result will be
limitations_and_societal_impact: Limitations are not discussed.
review_title: null

final_score: rating: 7: Good paper, accept
review_content: summary: 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.
main_review: Strength
+ This work is the first for gauge-equivariant self-attention (transformer). 
+ The paper is well written with a solid background and the related properties being proven. 
+ The technical contribution is high; a new parallel transport method by extending regular representation,  gauge equivariant value function by Taylor approximation.
+ The proposed method achieves SOTA on two popular benchmarks. 

Weakness
- Unclear parts. The design of the gauge-invariant attention (Eq.19) needs to be clarified. Why not use the conventional product-based score?  
- Lack of comparison to group-equivariant transformers.
While group-equivariant transformers [17, 21, 34] are mentioned in the related work, any of them are not compared in the experiments. The benefit of gauge equivariance would be better shown in comparison to the work of group equivariance. I think at least the work of [17] can be compared since they are developed for 3D data. 
- Lack of in-depth comparison with gauge-equivariant CNNs.
Table 4 shows that the convolution-based gauge-equivariant network performs only slightly worse than GET. Does it imply the marginal gain of the proposed self-attention layer compared to gauge-equivariant convolution? In this sense, a more in-depth comparison with gauge-equivariant CNNs [10, 46, 12] needs to be done to show the advantage of GET. 


limitations_and_societal_impact: N/A
review_title: null

final_score: rating: 8: Top 50% of accepted NeurIPS papers, clear accept
review_content: 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. 
main_review: There are many strengths in this paper. 
- The authors address an important and challenging problem and effectively advance research in applying self-attention networks on manifolds. The novelty in terms of ideas as applied to this problem appears to be high. 
- They provide an intuitive generalization of the multi-head self-attention (MHSA) for feature-fields on manifolds. 
- MHSA involves computing SA for multiple heads and then linearly combining them. However, the numerical values of the features depend on the gauge which in turn depends on the point on the manifold. Thus in order to compute pairwise self-attention scores at a point p, we need to first transform the numerical values features at a point q to the same gauge as at p and then use parallel transport to move them to p. The various weight matrices themselves need to satisfy constraints such that the score function is gauge-invariant and the value function is gauge-equivariant. This is challenging and the authors do a good job explaining and designing these modules. 
- Theoretical guarantees ensure the desired gauge-equivariance of MHSA.
- The experiments clearly show the usefulness of the proposed approach. The number of trainable parameters to achieve the same performance as methods like MeshCNNs for shape classification and segmentation is reduced by an order of magnitude. 

Points of clarification:
- I don't fully follow some of the notation and ideas in Section 4.1:
   - Why is the subscript $u$ needed in $q_u$? 
   - And shouldn't the value function and the corresponding weight matrix be $V_q$ and $W_q$ instead of $V_u$ and $W_u$? This is because $u$ refers to a point using the gauge at point $p$, but the the weight matrix for the value function acts on the parallel transported point $q_u$. 

The authors should discuss possible limitations of this work. 

Table 4 suggests that the attention mechanism itself is not crucial for good performance. Can the authors comment on this?
Also, what is this baseline exactly which uses equivariance and not attention? The authors say that this is a convolution based model created by removing the attention scores. How does removing the attention scores make it convolutional. And is the resulting model similar to gauge-equivariant mesh CNNs?
limitations_and_societal_impact: The authors have not explicitly addressed the limitations of the work and they say that it's because of space limitations.
review_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: 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.
main_review: From my viewpoint, this work will have a certain impact -- incorporating the gauge equivariance into the recently popular transformer. The main concern I have is about the comparison with other gauge equivariant CNN, especially given that the paper shows the negligible contribution of attention to the performance improvement in ablation study. 

Originality:
Gauge equivariant self-attention is novel. Although it’s a combination between existing gauge equivariance and self-attention, this work closes the gap in the transformer. In addition, the work presents a new way to accommodate parallel transport, which is interesting to me.

Quality:
The paper is of high quality in terms of theoretical analysis. But I’m a bit concerned about the experimental validation. 
- Missing baselines related to gauge equivariance. For example, the geodesic CNN [25] also tried to solve the orientation ambiguity problem. Also, there is no other method for comparison in the tables that also incorporate gauge equivariance. 
- I would appreciate it more if the paper could discuss and compare with gauge equivariant mesh CNN [46, 12]. 

Clarity:
Overall, the paper is well-written. But I still have several comments. Hope the author could address them during the rebuttal.
- What’s the motivation for using a regular field for the intermediate features?
- The gauge equivariance is not strict (up to error bound). I am not sure if the title will be misleading.
- What’s the efficiency of the proposed method?
- Adding citations into table 1 and 2 would be more readable.


Significance:
This paper injects the gauge equivariance into the recent popular transformer. However, due to the missing comparison with other gauge equivariant CNN, I am not sure if the result is significant. Further, the ablation study shows that the key to the improvement is due to the gauge equivariance. Does this mean that convolution, as used in [46, 12], would be better? I would like to hear more from the authors during the rebuttal. 


### Post rebuttal
I have the feedback from the authors and other reviews. Since the feedback addressed my concerns, I would like to improve my rating. 
limitations_and_societal_impact: From my viewpoint, I suggest discussing the limitation in terms of efficiency.
review_title: null

Meta-review: This paper develops the tools necessary to construct Gauge-equivariant transformers -- and to apply self-attention to manifolds or meshes with a focus on 2D smooth and orientable manifolds.
The presented methods can be applied to a wide variety of problems.


This paper is clearly written and technically sound. The experiments clearly show the strength of the proposed method (e.g. in comparison to MeshCNNs for shape classification).

Experimental results (Tables 1, 2) also support the idea that directly incorporating equivariances and invariances into the model's architecture is more efficient than using data-augmentation.
Decision: Accept (Poster)