Paper: Invariant Attention: Provable Clustering Under Transformations (ICLR-2024)

Reviewer comment (Novelty):
The proposed method is still based on dynamic kernels [1, 2, 3]. Why the kernels based on averaging are better than previous attempts?

Author response:
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.

Reviewer comment (Experiments):
Current empirical validation is limited - more quantitative experiments ($e.g.,$ overall accuracy over MNIST) would strengthen the claims. In addition, more qualitative results on complex image datasets ($e.g.,$ CIFAR100) or the impact of downstream tasks would be useful.

Author response:
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.

Reviewer comment (Writing):
The problem of the paper is not clear to me. What are the problems where the Invariant Attention is needed? Will it increase the overall data efficiency of ViT-type models? How is the clustering property helpful in this case?

Author response:
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.

Reviewer comment (Theory):
There are no proofs of the theorems and claims, and not even a sketch of the proof or an idea is provided. Though the authors claim to put it in the appendix, it is not possible right now to review the correctness. The definition of “infimal convex combination reach”, which is an important part of the theory, is not provided even on an intuitive level.

Author response:
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.

Reviewer comment (Theory):
It is not clearly stated which transformations are admissible. Authors claim that their method works for any transformation set, but provide theoretical guarantees only in the case of the T=SE(2). At the same time, they state that Invariant Attention enforces invariance under unknown transformations of the domain, which is clearly misleading, and the transformation set should be known beforehand to construct the kernel.

Author response:
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.

Reviewer comment (Reproducibility):
The optimization procedure for finding invariant mean is not fully described. How are the transformation vectors \tau_i parametrized in each experiment? Exactly, what parameters are we optimizing, and how? It would be nice to have a clearly described algorithm in the form of pseudo-code or something like that. Also, no information about the time complexity or needed resources is provided in the experiment part.

Author response:
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.

Reviewer comment (Theory):
Not a learnable algorithm. Though the authors claim that they are currently working at implementing learnable weights inside Invariant Attention, its real applicability to modern visual transformer models in the presented form is questionable. It requires running an optimization procedure for each attention head and each pair of image patches only to calculate the kernel weights. This is also dependent on the dimension and complexity of the transformation set and will require training separate models for different symmetry groups. The issue is not addressed in the paper.

Author response:
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.

Reviewer comment (Theory):
Your results (Theorem 3.1) are formulated for SE(2) explicitly. How is that transferable to other groups of transformations?

Author response:
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.

Reviewer comment (Theory):
What kind of structural properties are preserved or exploited by Invariant Attention? How will it help in prediction? May invariance to transformations actually harm the prediction quality, when the focus is on orientation, for example? Like classifying the right arrow and left arrow, for example?

Author response:
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.

Reviewer comment (Reproducibility):
There are no details on how to obtain the optimal transformations from 2.6. How to obtain these transformations is crucial. Without an efficient way to obtain them, the proposed method cannot be applied in practice.

Author response:
Thank you for the feedback. We have added two Sections: Section 2.1 and 2.4 that detail the optimization procedure.

Reviewer comment (Novelty):
The method does not involve actual feature learning. It's hard to argue the importance of the method for machine learning methods when there is no actual representation learning happening.

Author response:
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.

Reviewer comment (Experiments):
The experiments are extremely simple: 6 transformations of the same image or 10 MNIST samples. These might be good for a first step to see that the method/implementation is sound, but more validations are needed for a novel machine learning method.

Author response:
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.

Reviewer comment (Theory):
“Invariant Attention, enforces invariance under unknown transformations of the domain by optimizing over these transformations” How general are the transformations that Invariant Attention can optimize over? What kind of transformations can be optimized in practice?

Author response:
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.

Reviewer comment (Evaluation):
There is no quantitative result.

Author response:
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.

Reviewer comment (Experiments):
Without comparison to previous works or quantitative results, it is not possible to judge if the proposed technic is useful or not and the significance of the proposed method.

Author response:
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.