Paper: Fine-grained Local Sensitivity Analysis of Standard Dot-Product Self-Attention (ICLR-2024)

Reviewer comment (Reproducibility):
Solving Eq. (10) requires SVD for the n-by-d matrix. How to ensure the computational efficiency?

Author response:
To be clear, our approach only requires computing the spectral norm, and SVD is not needed for solving the spectral norm. In our paper, we just use the power iteration method which is known to be very efficient in computing the spectral norm and has been used many times in deep learning (e.g. Tsuzuku'18,  Miyato'18, Meunier'22).

Reviewer comment (Theory):
Before Proposition 1, the authors mention the robustness under l_2 perturbation. How about using l_\inf perturbations for robustness when compared to the adversarially chosen l_2 perturbations? In this case, Eq. (5) will be changed to the l_\inf norm but I’m not sure the used techniques are still valid.

Author response:
Extending our fine-grained $\ell_2$ analysis to the $\ell_\infty$ case is definitely non-trivial and would require some major changes, since we are using the fact that softmax is $1$-Lipschitz with respect to the $\ell_2$ norm. We want to emphasize that it is totally reasonable for our current paper to focus on the $\ell_2$ perturbation cases, and there are many papers published in top machine learning venues (e.g Singla'21, Trockman'21, Meunier'22, Singla'22, Prach'22, Araujo'23, Wang'23,Hu'23) that only focus on robustness and sensitivity of neural networks under the $\ell_2$ setting. Our paper is the first obtaining a non-vacuous certified robustness result of standard dot-product self-attention networks on CIFAR.  

[Singla'21]. Sahil Singla and Soheil Feizi. Skew orthogonal convolutions. ICML  

[Trockman'21] Asher Trockman and J Zico Kolter. Orthogonalizing convolutional layers with the Cayley transform. ICLR  

[Meunier'22] Laurent Meunier, Blaise J Delattre, Alexandre Araujo, and Alexandre Allauzen. A dynamical system
perspective for lipschitz neural networks. ICML  

[Singla'22] Sahil Singla, Surbhi Singla, and Soheil Feizi. Improved deterministic l2 robustness on CIFAR-10 and CIFAR-100. ICLR  

[Prach'22] Bernd Prach and Christoph H Lampert. Almost-orthogonal layers for efficient general-purpose lipschitz networks. ECCV  

[Araujo'23] Alexandre Araujo, Aaron J Havens, Blaise Delattre, Alexandre Allauzen, and Bin Hu. A unified algebraic perspective on lipschitz neural networks. ICLR  

[Wang'23] Ruigang Wang and Ian Manchester. Direct parameterization of lipschitz-bounded deep networks. ICML  

[Hu'23] Kai Hu, Andy Zou, Zifan Wang, Klas Leino, and Matt Fredrikson. Unlocking deterministic robustness certification on imagenet. NeurIPS

Reviewer comment (Novelty):
My main concern is that this work does not provide enough contribution. In Section 4, the gap caused by perturbation is derived. However, these results are not novel, in fact, they are easy to derive. The main idea of Section 4 is just finding a Lipschitz constant to bound the gap when perturbation is added to input. This can be easily done if we take derivative over input X and find an upper bound for the Frobenius norm of the gradient over X. In some other works, the closed form gradients (maybe over $W^Q,W^K$, but similar to gradient over X) are easily derived, e.g, Tian, Yuandong, et al. "Scan and Snap: Understanding Training Dynamics and Token Composition in 1-layer Transformer." arXiv preprint arXiv:2305.16380 (2023). Thus, I do not think the theory has much contribution.

Author response:
> However, these results are not novel, in fact, they are easy to derive. This can be easily done if we take derivative over input X and find an upper bound for the Frobenius norm of the gradient over X. 

We think our results are novel for the following reason. Our local sensitivity analysis is not a local Lipschitz bound result (see the explanation in Appendix A of our paper), and **our analysis is really needed to achieve our main goal that is to produce tight local bounds which can lead to non-vacuous certified robustness of standard dot-product self-attention networks on realistic data sets such as CIFAR.** We emphasize that getting **non-vacuous** certified robust accuracy is an important topic that is being studied in the deep learning field (e.g. see Singla'21, Trockman'21, Meunier'22, Singla'22, Prach'22, Araujo'23, Wang'23,Hu'23), and this is very different from the cases when one studies generalization bounds which are typically used to provide design guidelines and hence not required to be non-vacuous. If one follows the suggestions in this comment (e.g. finding upper bound for the Frobenius norm of the gradient), then the resultant bounds (e.g. the local Lipschitz bound in a concurrent submission available at https://openreview.net/forum?id=mivL0akE5E is exactly derived using this gradient norm idea) is too loose and can only produce 0 certified robust accuracy for self-attention networks on CIFAR.  Now we elaborate as follows.

- Non-vacuous certified robustness: Given a data point $X$, a classifier $F$ is said to be certifiably robust at radius $\epsilon > 0$ at this data point with label $y$ if for all $\tau$ such that $||\tau|| \leq \epsilon$, we have $arg max_j [F (X+\delta)]_j = y$. Given a perturbation level $\epsilon$, the certified robust accuracy is defined to be the percentage of the data points (in the testing test) where the classifier $F$ is certifiably robust at radius $\epsilon$, and this notion of certified robustness has been adopted in many recent works (e.g. Singla'21, Trockman'21, Meunier'22, Singla'22, Prach'22, Araujo'23, Wang'23,Hu'23). Based on our Proposition 1, our local sensitivity bound can be combined with the prediction margin to calculate the certified robust accuracy. As commented by Reviewer d9ub, achieving non-vacuous certified accuracy for dot-product self-attention on CIFAR10 task is one of our main contributions.  If ones tries to develop a local Lipschitz bound and uses it for certified robustness, the bound is usually too loose and the robust accuracy will be 0. For example, the local Lipschitz bound in a concurrent submission (Specformer) available at https://openreview.net/forum?id=mivL0akE5E is exactly derived using this gradient norm idea (see Theorem 4.3 of that paper). If we compare Theorem 4.3 of the SpecFormer paper with our bound (see table below), we can see our bounds are better by magnitude. The consequence is that Theorem 4.3 of the SpecFormer paper can only achieve 0 certified accuracy on CIFAR. **Our key sensitivity metric (Lemma 1 in our paper) is novel and crucial for obtaining a non-vacuous certified accuracy result.**

| $\ell_2$ perturbation bound $\epsilon$ | .01| .02|.03|.04|.05|.06| .07 |.08 | .09 | .10 |
|----|---|--------|--------|--------|--------|---------|---------|---------|---------|---------|
| PGD Lower Bound | 0.0286 | 0.0573 | 0.0860 | 0.1147 | 0.1427 | 0.1719  | 0.2008  | 0.2295  | 0.2582  | 0.2865  |
| Local Bound Theorem 1. (ours) | 0.0291 | 0.0582 | 0.0875 | 0.1168 | 0.1462 | 0.1757  | 0.2052  | 0.2348  | 0.2646| 0.2943|
| SpecFormer Local Lipschitz Bound (Gradient-based)| 16.901 | 34.403 | 52.515 | 71.245 | 90.600 | 110.589 | 131.219 | 152.499 | 174.436 | 197.036 |

[Figure Comparing to the Gradient-based Bound](https://drive.google.com/file/d/1riLDjkjXMbf0bmi3Ux6t0puZ_z60mSFq/view?usp=sharing)


- Our analysis is not a local Lipschitz constant analysis: To reduce conservatism, the local bounds derived in our paper are actually quite different from a local Lipschitz constant bound. Our local bound $||X-X'||\leq \epsilon \implies ||F(X)-F(X')|| \leq \delta(X, \epsilon)$ fixes $X$ and only allows $X'$ to vary. This is sufficient for calculating certified robust accuracy. The Lipschitz constant bound you describe $||F(X'')-F(X')||\leq L(X, \epsilon)||X''-X'||$ actually holds for any two points $(X',X'')$ in the $\epsilon$-ball around $X$.  This introduces unnecessary conservatism when applied to compute the certified robustness.

Reviewer comment (Theory):
How to theoretically guarantee that a model can both have good expressivity and stability?

Author response:
> How to theoretically guarantee that a model can both have good expressivity and stability?

As mentioned before, our focus is more on certified robustness rather than stability (which is typically used for generalization). Now we address the above comment via emphasizing the trade-off between certified robustness and clean performance. Our results and other papers on certifiably robust networks (e.g Singla'21, Trockman'21, Meunier'22, Prach'22, Araujo'23, Wang'23) all demonstrate that there is a trade-off between deterministic $\ell_2$ certified robustness and the clean performance (expressivity). Currently, improving certified robustness typically leads to degraded performance. One just navigates the optimal Pareto trade-off between robustness and expressivity depending on the possible level of perturbations for the tasks at hand.  Continuing to boost the clean performance of certifiably robust models is an ongoing research effort in the robust learning community. Our main contribution is on characterizing the robustness/performance trade-off for dot-product self-attention, which is a very reasonable self-contained topic. How to improve the structures of transformers to obtain better Pareto trade-off curves for robustness/robustness is beyond the scope of our paper, and should be investigated in the future. Our work serves as a foundation for such future study.

Reviewer comment (Theory):
When weights $W^Q,W^K,W^V$ follows some specific distribution, can the sensitivity bound be improved? Or the bound is only related to the magnitude of weights?

Author response:
> When weights follows some specific distribution, can the sensitivity bound be improved? 

We follow the standard setup in the certified robustness literature  and consider the certified robustness of a trained model with deterministic weight (Singla'21, Trockman'21, Meunier'22, Prach'22, Araujo'23, Wang'23). Evaluating the robust accuracy for a fixed trained model makes sense, since robustness certification is typically applied to a trained model. To the best of our knowledge, certified robustness under distributions of the network weight has not been extensively studied in the robust learning field. It is unclear how to extend our analysis or any other existing deterministic $\ell_2$ certified robustness analysis to such a setting considering a weight distribution. 

>Or the bound is only related to the magnitude of weights?

Our bound in Theorem 1 not only depends on the spectral norm of the weights, but is highly dependent on the actual matrix elements and how they interact with the input and other attention heads. This is especially true in the "key sensitivity metric" of eq. (12) since we are looking at the spectral norm after summing many attention weights together, making our bound more fine-grained and less conservative.
Here's a simple example: Suppose we have an attention layer with two attention heads. We set $H=0$, $W^Q_1 = W^Q_2$, $W^K_1=W^K_2$ ($P(X)_1=P(X)_2$), and $W^V_1 = W^V_2$. Then if $W^O_2 = -1/2*W^O_1$, the sensitivity metric of Eq (12) is: $||\sum_i P(X)_i \otimes (W^V_i W^O_i)^T|| = 1/2 ||P(X)_1 \otimes (W^V_1W^O_1)^T||$. However, if $W^O_2= 1/2 W^O_1$, then we have $ 3/2 ||P(X)_1 \otimes (W^V_1W^O_1)^T||$ which is a bound 3 times larger. In this case, we have not changed the spectral norm of any weight matrix, only the sign direction. But our key sensitivity metric changes significantly. To summarize, our bound is mainly related to our key sensitivity metric (Lemma 1).  Decreasing the spectral norms of $(W^Q, W^K, W^V, W^O)$ can eventually lead to a decreased value of the key sensitivity metric. However, that is not the only way to decrease the key sensitivity metric.

Reviewer comment (Experiments):
There are multiple naive bounds described in theory but not evaluated in practice. For example, equation 9 for bounding $\Delta_1$ and $||W_l^V W_l^O||(||X|| + \epsilon)$ for bounding $||X' W_l^V W_l^O||$ should also be evaluated in Figure 1 to make the conservatism argument strong.

Author response:
We agree that evaluating these naive bounds, especially the naive estimate of eq (9), would make the conservatism argument strong.  We have revised the experimental section (Section 5) and Figure 1 with the additional naive bounds. The results further show that our choice of the key-sensitivity metric is crucial for getting tight local sensitivity bounds of dot-product self-attention. We have also included an additional study on the tightness of our bound similar to Figure 1 for several input norm values found in Appendix C.2, Figure 4. We hope this provides a more complete picture of how our bound and architecture deal with conservatism.

Reviewer comment (Experiments):
Is it possible to perform another trial of robustness experiments on NLP tasks (text classification, entailment, etc.)? The analysis of this paper is applied to general self-attention and it is definitely great to see a practical evaluation on vision tasks. But it would be even better to see if the same robustness argument is applicable across domains.

Author response:
Thanks for this comment. We followed your instruction and produced a robustness study of dot-product self-attention on the Stanford Sentiment Tree-bank (SST) dataset. We also revised our paper to include these new results in Appendix C.3 and Figure 5. We follow the setup in Xu’23, Wang’20, Zhu’20, and Li’21, and consider the deterministic $\ell_2$ certified robustness on the word embedding space (we use the BERT embedding from Devlin'19). Our analysis here only serves as a proof of concept for using our bounds to NLP tasks, and we acknowledge that considering perturbations on the word embedding space (instead of the tokens directly) is a common limitation of applying sensitivity analysis to NLP benchmarks, as shared in many other papers (Hou’23). The certified robustness radii that we obtained is measured using $\ell_2$ norm, and this is similar to those in prior work *without* using dot-product attention (Xu’23). Our results demonstrate that our sensitivity analysis bounds lead to non-vacuous robustness results on SST. For more details, see Appendix C of our revised paper. If the reviewer further wants us to work on other specific tasks, please let us know and we are willing to provide more results.

[Xu’23] Xu et al. Certifiably Robust Transformers with 1-Lipschitz Self-Attention. 
[Hou’23] Hou et al. TextGrad: Advancing Robustness Evaluation in NLP by Gradient-Driven Optimization. ICLR  
[Wang’20] Wang et al. InfoBERT: Improving Robustness of Language Models from An Information Theoretic Perspective. ICLR  
[Zhu’20] Zhu et al. FreeLB: Enhanced Adversarial Training for Natural Language Understanding. ICLR  
[Li’21] Li et al. TAVAT: Token-Aware Virtual Adversarial Training for Language Understanding. AAAI  
[Devlin’19] Devlin et al. Bert: Pre-training of deep bidirectional transformers for language understanding. NAACL-HLT