Transformers Learn Low Sensitivity Functions: Investigations and Implications

Questions:

Q1: Please see the response to weakness 2a.

Q2: How does the variance $\sigma$ impact the results?

Thank you for the question. We evaluate the sensitivity values of the five models compared on the CIFAR10 dataset at the end of training using different values of $\sigma$. We added the results in Appendix A.3 and we see that even though the sensitivity values are different for different $\sigma$ (as expected), the conclusion that transformers learn lower sensitivity functions than CNNs is robust to the value of $\sigma$. Appendix A.4 also has results for the QQP dataset with a different value of $\sigma$ than considered in the paper.

Q3: Using fixed noise level for different tokens.

We believe that token norms are comparable and thus, it makes sense to use a fixed noise level. Specifically, noise is added at patch level for images, and as mentioned in line 377 in the paper, for language tasks, noise is added after layer normalization.

Q4: Attention architecture considered in Section 3.1

The model considered in Section 3.1 is a standard attention layer with key, query, and value weights $W_K, W_Q, W_V$, composed with a linear decoder $U$. Please also see the response to weakness 1a.

Q5: Does test accuracy correlate with sensitivity?

Fig 13 in the Appendix shows the test accuracies for the models considered in Fig. 4 on the CIFAR-10 dataset. Although the ViTs have a slightly higher test accuracy, it is comparable across the five models.

Q6: How does model scale affect the difference between sensitivity values?

This is an interesting question. However, as mentioned in the response to weakness 1b, comparing larger models is challenging because we also have to control the pretraining data and strategies for a fair comparison.

Q7: “Section 5 only considers the RoBERTa model. Do the claims hold for more recent language models?”

As mentioned in line 401 in the paper, we include the results with GPT-2 in the Appendix which lead to the same conclusions. Please also see the response to weakness 1b.

Thank you for the detailed comments and feedback to help improve our work. We hope that the following responses will address the reviewer’s concerns and would be happy to address further questions the reviewer may have.

Weaknesses:

1: Regarding the scale of experiments

a. “Results in Section 3 use a single attention layer model”

We emphasize that the goal of Section 3 is to show that even in a very simple setting with a single-layer self-attention model, we can see the low sensitivity simplicity bias. While the dataset can be considered for analyzing more complex model architectures as well, we choose not to do so because a) it is not necessary to use a more complex model for this data and b) it is more interesting to analyze larger models on more realistic settings. Hence, we compare larger models on real-world vision and language datasets in Sections 4 and 5. That being said, the experiments in Section 3 provide some insight into what part of the model gives rise to the low sensitivity bias.

b. “Experiments in Sections 4 and 5 use small-scale datasets and models”

We emphasize three points to address this concern.

First, we respectfully disagree with the statement that “bold conclusions are extracted from a single model/dataset”. We compare several datasets and models to validate our claims. Specifically, to show that transformers learn lower sensitivity functions than CNNs, we consider two datasets as follows.

• We consider the CIFAR10 dataset and compare two CNN-based models, namely ResNet18 and DenseNet12, with a ConvMixer model and and two ViT models, and
• We compare a CNN (ResNet18) and a ViT on the SVHN dataset.

Next, to show that transformers learn lower sensitivity functions than MLPs,

• we consider the FashionMNIST dataset and compare a ViT, a simpler 3-layer CNN and MLP-based models with two activation functions.
• we also consider a binary classification task with the MNIST dataset and compare a ViT and an MLP.

We consider simpler models and datasets for comparison with MLPs because we expect these models to perform reasonably well on these datasets.

Similarly, for comparisons with LSTMs, we consider two datasets MRPC and QQP, and compare with two language models, namely RoBERTa (in the main body) and GPT-2 (in Fig. 20 the Appendix, as mentioned in line 401 in the paper).

Second, we note that we focus on relatively simpler datasets and models because we can train the models from scratch in these settings to get reasonably good performance. This is important because we compare the sensitivity of models that have comparable and reasonably good train accuracy. We want to single out the effect of the architecture, without any confounders such as the choice of the optimization algorithm or data augmentation or use of any pretraining strategies, which might vary for different architectures to get good accuracy.

That being said, we compare (pre-trained) ConvNeXT and VIT-B/16 models on ImageNet-1K dataset to show that the conclusions are relatively robust to pretraining and increasing scale. However, a systematic and fairer comparison for large-scale models would warrant better control over the pretraining strategies, which is beyond the scope of this work.

Third, we acknowledge the reviewer’s concern about the claims not directly transferring to LLMs. However, we emphasize that our work aims to give insights about the inductive biases and other properties of the transformer architecture in general, and we don’t claim that these insights transfer directly to LLMs or about the interactions with pretraining and finetuning. That being said, comparisons with GPT-2 model, which is a causal model and more similar to recently used language models compared to RoBERTa, lead to the same conclusions and indicate that we can expect them to also hold for other language models.

2: Clarity

a. “It is not clear how noising is applied to image data”

Fig.1 is exactly what is done for images – adding noise to one patch at a time. Since noise is at the patch level, the process is the same for CNNs and ViTs.

b. “It is hard to understand the importance of Proposition 2.1”

As stated in lines 154-157 in the paper, “Larger eigenvalues for lower-order monomials indicate that simpler features are learned faster. Since low sensitivity implies learning low-degree polynomials, Proposition 2.1 also implies a weak form of low sensitivity bias.”

Thank you for the positive feedback and the helpful comments.

W1, Q1: How do sensitivity and robustness relate to the tradeoff between robustness and accuracy observed in different contexts?

Thank you for the question. We note that these tradeoffs occur when measuring robustness to adversarial perturbations or out-of-distribution (OOD) data. In this paper, we focus on benign shifts to measure robustness and hence, don’t observe this type of a tradeoff.

W2-4: Thank you for the suggestions to improve the presentation of the paper. We have made some edits to the introduction as suggested. We will also incorporate the other suggestions in the final version.

Q2: Relation between generalization in representation learning to generalization in sensitivity

This is an interesting question. We believe that under benign shifts, such as on new samples from the same distribution as the train set, or under small random corruptions, most properties should generalize. At the very least, the conclusions drawn based on the metrics should generalize, even if the values change. Exploring the generalization of sensitivity systematically can be an interesting direction for future work.

Q3: Comparing LSTM and RoBERTa with the same number of parameters.

As mentioned in the paper, we compare models that have the same accuracy to ensure the comparison of sensitivity values is fair. Since both the models attain very similar accuracy, they can potentially learn similar functions that could attain similar sensitivity values. However, we observe that they learn functions that have similar accuracy but differ significantly in terms of sensitivity.

Thank you for the helpful comments and feedback.

W1-2: Regarding the claim that low sensitivity explains the better robustness of transformers.

Following the reviewer’s suggestion, we have rephrased the aforementioned statement in the paper as follows: “As encouraging lower sensitivity improves robustness, the inductive bias of transformers to learn functions of lower sensitivity could explain their better robustness (to common corruptions) compared to CNNs.”

As the reviewer mentioned, prior works have shown that transformers are more robust in CNNs. We note that we discuss related work on robustness and data augmentation in lines 1467-1479 and lines 1489-1499, respectively in the Appendix. However, discussing the difference in the robustness of transformers and CNNs in more detail helps us elucidate the connection between lower sensitivity and better robustness.

We summarize our results from Section 6.1 below and hope our response and the revised statement address the reviewer’s concern about the section.

In Sections 4 and 5, we showed that transformers have lower sensitivity compared to other architectures, and in Section 6.1, we investigate the role of lower sensitivity in the improved robustness of transformers. Using the CIFAR-10-C dataset, we (a) observe that transformers have better robustness compared to CNNs, and (b) show that encouraging lower sensitivity while training the transformer further improves the robustness. Here, we emphasize three points.

First, transformers exhibit a bias towards learning low-sensitivity functions even when trained without explicit regularization, as shown in the experiments in Sections 4 and 5.

Second, our goal while training with data augmentation (with Gaussian noise added to the images randomly) and sensitivity regularization (with patch-wise Gaussian noise added to the images) is to see if we can show that reducing sensitivity leads to improved robustness. These methods seem like the simplest ways to encourage low sensitivity. We welcome any suggestions the reviewer may have about other ways we can encourage lower sensitivity, and will try our best to test those.

Third, while one may expect that encouraging lower sensitivity while training would improve robustness to noise corruption, we observe improved robustness to other corruptions as well. For instance, corruptions from the blur, weather and digital transform categories are significantly different from the noise corruptions as they cannot be implemented by making small pixel-wise changes to the image. The improved robustness across these perturbations suggests that it could be a consequence of the lower sensitivity.

In other words, while the definition of sensitivity is slightly similar to the corruptions from the noise category, it is quite different from the other corruptions and still correlates well with the improved performance on those.

W3, Q1: Comparing (train) sensitivity to perturbation robustness at test time.

As suggested by the reviewer, there could be connections between sensitivity and robustness to random perturbations, however, sensitivity as a metric not only serves as a measure of inductive bias that distinguishes transformers from other architectures but also has important implications, as a measure of robustness and flatness of the minimum and as a progress measure for grokking.

For instance, in App. A.2, we evaluate sensitivity to Gaussian noise added across the input and observe that while transformers have lower sensitivity (or better robustness), this metric does not distinguish transformers from other architectures as clearly as the sensitivity to token-wise Gaussian perturbations. Similarly, some of the noise corruptions considered in the experiments in Section 6.1 also indicate that lower sensitivity is correlated with better robustness.

That being said since we consider sensitivity a notion of inductive bias, it makes sense to measure it on the train test, analogous to other metrics of inductive bias such as maximum $\ell_2$-margin for linear predictors on separable data. This allows leveraging a property of the train data to predict things like generalization on the test data.

Thank you for the positive feedback and the helpful comments.

W1: Thank you for the suggestion, we have revised the sentence.

Q1: How does sensitivity relate to robustness to adversarial/random perturbations?

Thank you for the question. While low sensitivity (to token-wise perturbations) correlates with better robustness, they are not equivalent. Specifically, in App. A.2, we evaluate sensitivity to Gaussian noise added across the input and observe that this metric does not distinguish transformers from other architectures as clearly as the sensitivity to token-wise Gaussian perturbations.