Invariant Attention: Provable Clustering Under Transformations

Dear Reviewer 4tjm, thank you for your valuable feedback. Below we respond to your comments in detail.

1. We thank you for this important remark. We have improved the text to make the motivation and approach more clear, while acknowledging that incorporating IA into a ViT is yet to be shown. However, we believe that the following intuitions hold promise. One of the reasons that transformers require large amounts of data compared to CNNs to achieve the same accuracy is that transformers ignore structural properties of images, especially geometric and spatial properties. Images naturally contain repeated motifs that are some geometric transformation of each other due to factors such as symmetry, camera motion/perspective, object orientation, etc. These geometric transformations are combined with statistical variability to produce intra-class variability in these motifs. CNNs are equipped with the geometric property of translation-equivariance. Since attention mechanisms in transformers offer no such benefits, in our work, we propose Invariant Attention (IA) as a potential replacement to the kernel (self) attention sublayer in a ViT layer while keeping the feedforward sublayer intact. We believe that, enhanced with IA, the transformer architecture would be invariant to geometrical transformations in data, thus eliminating the need to learn invariance from scratch and reducing the total amount of data needed to train a Vision Transformer. Attention mechanisms possess intrinsic clustering properties, as elucidated in the Introduction and visually represented in Figure 2. This is particularly evident with image data that contain motifs subject to geometric transformations. In such cases, the Euclidean distance between transformed images belonging to the same class would not be small, thereby incorrectly clustering upon repeated application of kernel attention. It is this clustering nature of attention that we have tried to address through a generalized optimization framework, and provided a proof for one group of transformations. We believe this inherent clustering property of Invariant Attention will introduce benefits of transformation-invariance to transformers.

2. Thank you for pointing out these excellent references. These works propose valuable ideas for endowing neural networks with invariance [1] and/or locality [2, 3]. For example, the authors of Spatial Transformers [1] present a learnable module that achieves invariance to transformations via a so-called localization network (a fully-connected network or a CNN). The authors in [2] present LocalViT, where locality is added to vision transformers via adding depth-wise convolution into the feed-forward network. Both works use approaches that take advantage of the 2D structure of images (e.g. adding convolutions). While these works provide substantial empirical results, they remain a black-box approach, and do not aim to provide theoretical analysis that would provide further insight or performance guarantees. Additionally, none of them propose an attention-mechanism based method. In contrast, the proposed Invariant Attention mainly provides theoretical results of provable clustering and simple proof-of-concept experiments that do not involve training. We agree that [1-3] are notable works in the field. Lastly, our proposed kernel is not based on averaging but on maximum similarity. The authors in [3] also propose such a kernel and use probabilistic analysis to show its strengths and valuable properties, e.g. positive definiteness with high probability. In our proposed Invariant Attention, however, we utilize this kernel to calculate invariant weights.

3. Thank you for the constructive feedback on the limitations of empirical analysis. We acknowledge this need and are working to provide more quantitative experiments over MNIST. However, to run experiments on CIFAR100 or other datasets of natural images where statistical variability takes many forms other than geometry, one would need to incorporate a more complex model (e.g. a network with nonlinearities), which we leave for future directions of this work.

Responses to Weaknesses:

1. Thanks for the comment! Clustering is an important problem in unsupervised machine learning. Currently there is no provable clustering technique for data that contains motifs subject to transformations. We have formulated a framework useful for this task for any general transformation group and have proved clustering for a certain class of transformations, the SE(2) group. Application to ViTs: A reason for transformers needing more data than CNNs to train up to the same level of accuracy is that transformers ignore structural properties of image data, especially geometric and spatial properties. Image patches have repeated motifs that are some geometric transformations of each other due to factors such as symmetry, camera motion and perspective, object orientation, etc. These geometric transformations are combined with statistical variability to produce intra-class variability in these motifs. CNNs are equipped with the geometrical property of translation equivariance. Since attention mechanisms in transformers offer no such benefits, in our work, we propose Invariant Attention (IA) as a potential replacement to the kernel attention (self-attention) sublayer in a ViT layer while keeping the feedforward sublayer intact. We believe that, enhanced with IA, the transformer architecture would be invariant to geometrical transformations in data and additional data that was otherwise needed by the transformer to learn this invariance will no longer be required. This, we believe, would reduce the total amount of data needed to train a Vision Transformer. Attention mechanisms have inherent clustering properties as we have explained in the introduction Section and through experiment in Figure 2. when our data contains images/image patches containing motifs subject to geometric transformations since Euclidean distance between transformed images of the same class is not small would not be clustered together on repeated application of Kernel attention. It is this clustering nature of attention that we have tried to address through a generalized optimization framework and proved theoretically for one group of transformations. We believe this inherent clustering property of Invariant Attention will introduce benefits of transformation-invariance to transformers.

2. Thanks for the suggestion. The proof involves many geometric ideas, and we could not sketch them in the main body due to space limitations. However, we have improved the Section on Uniqueness, Section 3 to better motivate some ideas leading up to the main results.

3. Thanks for the comment. From the formulation in equation 2.3, we have used 1. closure properties of a transformation group and norm preserving property of Euclidean transforms to remark that the optimization problem in Equation 2.3 contains solutions that are transformations of each other. This is also not readily evident for other groups of transformations. The closed form expression of the invariant mean also requires additional geometric properties of the data as required by Danskin's theorem and we have highlighted them in the updated version to signify the importance of these observations.

4. Thank you for the feedback. We have added a new Section, Section 2.1 that details the optimization of transformations and discusses the generalization ability of this framework to other transformation groups We can impart invariance to the group of transformations we desire. Some examples of transformation groups that are geometric transformation groups are Affine group, Similarity group, etc. When implemented for these groups, invariant attention enforces the Invariance property to the domain of the selected transformation group. We have added to text to distinguish the generalizability of the framework to other transformation groups and the transformation group we have chosen for our proof.

5. Thank you for the feedback. We have added two Sections: Section 2.1 and 2.4 that detail the optimization procedure. In these Sections we have included an algorithm box, optimization procedure and parallelizable nature of our computations.

6. Thanks for the feedback. When compared to the current Kernel Attention mechanism, Invariant Attention does add computational overhead. However, as described in Section 2.4, most of this overhead is parallelizable and the benefits of invariance properties to transformations can be imparted for clustering or into ViT architectures.

Response to Questions:

2. Thank you for the question. We have added an new Section, Section 2.1 that details the optimization of transformations and discussues the generalizability of this framework. Some examples of transformation groups that are geometric transformation groups are affine group, similarity group, etc.

3. Invariant attention will be helpful for problems in classification where the class label is invariant to transformations. For problems in which orientation is of interest, one solution is to extend the output of IA to output not only the invariant mean, but also the optimal transformations which generated the Invariant Mean. We note that these transformations and the Invariant Mean are jointly optimized in the computation.

Responses to Weaknesses:

1. Thank you for the feedback. We have added two Sections: Section 2.1 and 2.4 that detail the optimization procedure.
2. Thanks for the comment! Clustering is an important problem in unsupervised machine learning. Currently there is no provable clustering technique for data that contains motifs subject to transformations. We have formulated a framework useful for this task for any general transformation group and have proved clustering for a certain class of transformations, the SE (2) group. We have added details about the optimization procedure for the invariant mean in Section 2.4.
3. Thank you for the constructive feedback on the limitations of empirical analysis. We acknowledge this need and are working to provide more quantitative experiments over MNIST. However, to run experiments on CIFAR100 or other datasets of natural iamges where statistical variability takes many forms other than geometry, one would need to incorporate a more complex model (e.g. a network with nonlinearities), which we leave for future directions of this work.
4. Thank you for the question. We have added a new Section, Section 2.1 that details the optimization of transformations and discusses the generalizability of this framework. Some examples of transformation groups that are geometric transformation groups are affine group, similarity group, etc.

Response to Weaknesses:

Thank you for your constructive feedback. We acknowledge the need for quantitative results, and are currently performing a quantitative comparison of kernel attention and Invariant Attention to measure clustering accuracy. This experiment will be added to the next version of the paper.