Setting the Record Straight on Transformer Oversmoothing

[..4. ...The authors analyse a simplified Transformer architecture...Pre-LN architectures [2] have been shown to effectively counteract some of the oversmoothing...]

[..Questions:..]

[..1. ...how your findings will change in the presence of different weights per layer and in the presence of pre-ln layers?...]

We were also interested in understanding how our theoretical results generalize to a more general scenario with changing parameters and Pre-LN architectures. To test this, in our experiments we allow both attention matrices $\mathbf{A}$ and weights $\mathbf{H}$ to change every layer, and we use the same architecture as ViT-Ti (a Pre-LN architecture). We were excited to see that even though these models don’t follow the assumptions of our theoretical results, we can still use them to control the smoothing/sharpening behavior of these models. Thank you for brining this up. We have clarified these details in the paper.

[..2. If oversmoothing is the problem, what about other ways to mitigate it? I am talking about scaling the residuals (Noci et al, [1, 2]) or initializing the attention layers differently, e.g. [3] or different initializations per layer, e.g. [4]...]

Thanks for bringing this up. One of the strengths of our approach is that it is complementary to these approaches (scaling residuals [1 ,2], different initializations [3, 4]) and so we can apply them at the same time as our approach. Because of this we opted to compare against an approach that isn’t complementary, specifically FeatScale by Wang et al., 2022. Alongside the above increased depth experiment we also compare its data efficiency and corruption robustness below:

Data Efficiency (Table 2)

Corruption Robustness (Table 4)

[..3. Before section 5, should the superscript $\mbox{}^+$ model be initialized as $\mbox{diag}(\Lambda_H) = +(\psi^2)$?..]

This is correct! Thank you for catching this typo!

[..4 . Can you comment on the stability of your new parametrization? Especially what (if any) are the differences in the early stage of training..]

This is also interesting! To better understand what is happening early in training we have added a plot of HFC/LFC for the final layer against training iterations in Figure 4 in Section 5. Interestingly, we see that the last layers of both ViT-Ti$^*$ and ViT-Ti have the same filtering properties from the beginning to the end of training, despite the fact that the first $L/2$ layers of ViT-Ti$^*$ are sharpening. Further at the beginning of training, the last layer of ViT-Ti$^-$ becomes less sharp, but by the end of training it becomes sharper than at initialization. Thank you for asking about this and for these citations, we will add them to the paper.

Thank you again for your time. Let us know if you have any other questions and we will respond as soon as possible.

Dear Reviewer yd6A,

Thank you so much for your time and many comments. We respond to each question below.

[..Weaknesses:..]

[..1. Overall I find the main contributions/results hard to digest. Most of the results seem to be extensions from Wang et al...]

Thank you for asking about this. We completely agree that Wang et al., 2022 provides big inspiration for this work. Their analysis is extremely insightful and our goal is to further extend this analysis to uncover the relationship between the eigenspectrum of the weights and the suppression (or not) of high frequencies when the weights H and residual connection are taken into account. This turns out to be very important as it gives us tools to control the eigenspectrum of the weights in order to induce smoothing or sharpening.

Beyond the nice work of Wang et al., 2022, we were also heavily inspired by the work of [Dong et al., Attention is not all you need: Pure attention loses rank doubly exponentially with depth (ICML 2021)] who show that feature representations in Transformers can rank collapse to rank 1, but that there exist infinitely many weights that do not rank collapse with the residual connection. We extend this to show exactly which weights do not rank collapse: only when multiple eigenvalues $(1 + \lambda^H_j \lambda^A_i)$ simultaneously have equivalent dominant magnitudes. However, this is a rare case: because $\mathbf{A}$ and $\mathbf{H}$ are learned, it is unlikely the magnitude $|1 + \lambda^H_j \lambda^A_i|$ of any two eigenvalues will be identical.

We were also influenced by the extensive analysis of [Park & Kim How do vision transformers work? (ICLR 2022)]: we base our experiment setup off of theirs. Finally, our inspiration to look deeper into how the eigenspectrum can be connected to smoothing, while including the weights and residual connection is also inspired by the related line of work on oversmoothing in graph neural networks e.g., [Di Giovanni et al., Understanding convolution on graphs via energies. (TMLR 2023)].

[..2. show if/how their new parametrization enables training in severe cases of oversmoothing, as e.g. when the networks become much deeper...]

Thank you for raising this point, we agree. To answer this we have added an experiment where we double the depth of the ViT-Ti models on CIFAR100. We have also added a parameterization variant where the first $L/2$ layers sharpen (as in ViT-Ti$^-$) and the remaining $L/2$ layers are left normal (as in ViT-Ti), which we call ViT-Ti*. Finally we have added a comparison FeatScale by Wang et al., 2022. All results are averaged over 3 trials with standard deviations (we bold the best result in each column and any method whose mean plus standard deviation is greater than or equal to mean of the best result):

[new] Increased Depth (Table 3)

Here ViT-Ti$^+$ and ViT-Ti$^-$ both underperform ViT-Ti, and we suspect this is because ViT-T$^-$ sharpens too much, amplifying noise in the data (whereas ViT-Ti$^+$ oversmooths the data). On the other hand ViT-Ti* significantly outperforms all other models not just for larger depths in Table 3, but also in corruption robustness (Table 4), and data efficiency (Table 2). This suggests that both oversmoothing and oversharpening are detrimental to performance, and that a careful balance is key to an accurate model. We have updated the paper to describe these results and look forward to investigating this further in future work.

[..3. ...a task in vision or language that requires a higher rank output, requiring a different prediction per token, could be more desirable...]

This is a very good point. We are interested to try out the approach in a more label-dense vision task such as segmentation, where oversmoothing poses much larger problems than image classification as you suggest.

(response continues below)

Dear Reviewer yd6A,

Thank you again for your time on the paper. Do you have any extra questions about the paper or our response?

Let us know and we will try to response as soon as possible before the discussion period ends on November 22.

Thank you,

The authors

[..3. ...redundancy in the cases in Theorem 1..]

Thank you for pointing out this redundancy! We have fixed this.

[..4. ...the authors show the distribution of the dominating eigenvalues but do not show whether the features oversmooth when $\lambda_1^A$ dominates...]

We believe there is a small confusion here. For ViT-Ti$^+$ we have that (a) Figure 3 shows that the features do oversmooth more and more, and (b) Table 1 shows that $\lambda_n^A$ dominates. On the other hand, for ViT-Ti$^-$ we have that (a) Figure 3 indicates that the features do not oversmooth, and (b) Table 1 shows that $\lambda_1^A$ dominates. This shows that even when the attention matrix $\mathbf{A}$ is not fixed across layers the results agree with the Theorem statements. Thank you for letting us clarify this, we will add these details to the text to make this clearer.

[..Questions..]

[..1. The precise definition of dominating eigenvalues should be given...]

We agree. We have pulled this definition out of the proof and placed it into the main paper.

[..2. I think the authors mean $\mbox{diag}(\Lambda_H ) := +(\psi^2)$..]

Thank you for catching this! We have fixed this.

[..Minor Comments..]

[..1. compare the runtime of their ViT-Ti$^+$ with other methods...]

Thanks. The training time of ViT-Ti versus our original proposed methods is as follows (averaged over 3 runs):

The throughput of each method (in images/second, averaged over 10 runs) is:

We will add these timings to the text.

[..2. The authors showed the distribution of dominating eigenvalues of their - version on CIFAR10, but not on ImageNet...]

We think there is a small confusion. The dominating eigenvalue distribution of our $\mbox{}^-$ version is always 0%: $(1 + \lambda^H_j \lambda^A_n)$ and 100%: $(1 + \lambda^H_j \lambda^A_1)$, regardless of the data, as we control the eigenvalues via bounding $\Lambda^H$ in our parameterization. We will clarify this in the paper.

Thank you again for your time. Let us know if you have any other questions and we will respond as soon as possible.

Dear Reviewer QJ8g,

Thank you so much for your time and questions. We respond to each of them below.

[..Weaknesses:..]

[..1. ...parameterizing H as a symmetric matrix might restrict the expressive power..]

Thanks for this comment. Our original motivation for using a symmetric parameterization was inspired by [Hu et al., 2019. Exploring weight symmetry in deep neural networks.] who show that weight symmetry in NNs does not affect their universal approximation capabilities. However, we agree that it would be interesting to test a non-symmetric parameterization. Based on your suggestion we have implemented a new parameterization for $\mathbf{H}$ as $\mathbf{H} = \mathbf{V}_H \Lambda \mathbf{V}_H^{-1}$. Specifically, we take gradients to $\mathbf{V}_H$ in the backwards pass, and compute $\mathbf{H}$ using the above equation in the forwards pass. For both the symmetric and non-symmetric versions we test out 2 variants: 1. (-) i.e., the sharpening model: $\mbox{diag}(\Lambda_H ) := −(\psi^2)$ and a new variant 2. (*) where the first L/2 layers are parameterized as the sharpening model above, and the remaining L/2 layers are left normal (as in ViT-Ti). We use the suffix “-NS” for the non-symmetric parameterizations. We have also included a new baseline, FeatScale [Wang et al., 2022], as well as a new experiment where we double the depth of the network. Here are the new accuracy results, averaged over 3 trials with standard deviations (we bold the best result in each column and any method whose mean plus standard deviation is greater than or equal to mean of the best result):

Data Efficiency (Table 2)

[new] Increased Depth (Table 3)

Corruption Robustness (Table 4)

We have added these results to the paper. The non-symmetric parameterization does improve upon the symmetric parameterization in some cases (particularly for increased depth and corrupted data). We appreciate your suggestion to investigate this.

[..2. ...robustness and data efficiency...]

Thanks. This seems to be a very deep question. We originally suspected a link between sharpening and data efficiency because of the eigenvalue distribution of the Data Efficient Image Transformer (DeiT-B/16) in Table 1: compared to all other baseline models, it has the highest distribution of sharpening eigenvalues. This is what led us to test data efficiency. We also hypothesized that it may help with input corruption robustness as many of the corruption techniques in Hendrycks & Dietterich, 2019 blur or reduce the resolution of the image: any further smoothing would likely make classification even harder. As far as we are aware we are the first to investigate the relationship between sharpening and data efficiency and corruption robustness. We agree this is interesting to further investigate. If you have any suggestions for additional experiments we will try to implement them as quickly as possible.

(response continues below)

Dear Reviewer QJ8g,

Thank you again for your time on the paper. Do you have any extra questions about the paper or our response?

Let us know and we will try to response as soon as possible before the discussion period ends on November 22.

Thank you,

The authors

Data Efficiency (Table 2)

[new] Increased Depth (Table 3)

Corruption Robustness (Table 4)

We have added these results to the paper. The non-symmetric parameterization does improve upon the symmetric parameterization in some cases (particularly for increased depth and corrupted data). We appreciate your suggestion to investigate this.

[..Additionally, can you connect it to the condition number?...How should I interpret the asymmetry metric? How does it relate to condition number?..]

This is an interesting question. To investigate this we computed the average condition number for all H matrices used in ViT-Ti for every epoch in training on CIFAR100. We have added the results to Appendix C. The condition number fluctuates during training but gradually increases throughout training. We do not know of a connection between the condition number and the asymmetry of H, but agree it would be an interesting direction for future work.

[..3. how do you update $\Theta$ in $QR(\Theta) = [V_H, R]$?...Can you give more details on updating QR decomposition?..]

We backpropagate through the QR decomposition, which requires a QR decomposition for each gradient step. This is because we require $V_H$ to be orthogonal. If we instead parameterized $V_H = \Theta R^{-1}$ and used gradient descent to update $\Theta$ then this could break the orthogonality of $V_H$. Thank you for this, we will add this detail to make this clearer in the paper.

[..4. How does the HFC/LFC evolve in training steps?..]

Good question. To answer this we plot the HFC/LFC of the final layer for every epoch during training for the best performing symmetric models (ViT-Ti$^-$ and ViT-Ti$^*$), as well as ViT-Ti, and present the results in Figure 4 in Section 5. Interestingly, we see that the last layers of both ViT-Ti$^*$ and ViT-Ti have the same filtering properties throughout training, despite the fact that the first $L/2$ layers of ViT-Ti$^*$ are sharpening. Further at the beginning of training, the last layer of ViT-Ti$^-$ becomes less sharp, but by the end of training it becomes sharper than at initialization. Thank you for this suggestion!

[..5. why do you have $V_H^{-1}$ after clip?..]

[..6. I think it should be $\mbox{diag}(\Lambda_H) = \Phi^2$..]

[..8. "||" is missing in definition of asymmetry..]

Thank you for catching these! There should be no $V_H^{-1}$, no negative sign, and additional ||.

[..7. Please define $Q$ in the main statement of Theorem-2..]

Good call, we have moved the definition from the proof to the theorem statement.

[..9. Please describe the metrics for Table-2 and Table-3..]

Sorry, these are test accuracies, thank you for asking!

Thank you again for your time. Let us know if you have any other questions and we will respond as soon as possible.

Dear Reviewer rqLo,

Thank you so much for your time and detailed response. We respond to each question below.

[..Weaknesses:..]

[..1. ...a more detailed comparison to Wang et al is needed..]

Thanks for this! We agree that it would be enlightening to compare with Wang et al., 2022 and so we have implemented their FeatScale method. Here are the new accuracy results. We have added another a new experiment where we double the model depth (Table 3) and a parameterization variant where the first $L/2$ layers sharpen (as in ViT-Ti$^-$) and the remaining $L/2$ layers are left normal (as in ViT-Ti), which we call ViT-Ti*. All results are averaged over 3 trials with standard deviations (we bold the best result in each column and any method whose mean plus standard deviation is greater than or equal to mean of the best result):

Data Efficiency (Table 2)

[new] Increased Depth (Table 3)

Corruption Robustness (Table 4)

We have added these results to the paper. We suspect the reason that FeatScale underperforms is because it is still motivated assuming high frequencies are always suppressed, specifically they (Wang et al., 2022) say “According to our analysis in Section 2.2, MSA module will indiscriminately suppress high-frequency signals, which leads to severe information loss. Even though residual connection can retrieve lost information through the skip path, the high-frequency portion will be inevitably diluted (Theorem1). To this end, we propose another scaling technique that operates on feature maps, named Feature Scaling (FeatScale).” We agree with you that their analysis suggests that residual connections might prevent high frequencies from collapsing to zero, but Wang et al., 2022 disagree: “Although multi-head, FFN, and skip connection all help preserve the high-frequency signals, none would change the fact that MSA block as a whole only possesses the representational power of low-pass filters”. We believe the reason they argue this is because when they analyze the update equations that include residual connections (Proposition 5) they do not account for the eigenspectrum of those equations. Our goal with this work is to set the record straight with a new analysis that uncovers the relationship between the eigenspectrum of the weights and the suppression (or not) of high frequencies in the update equations.

[..2. is 0.7 small enough to suggest symmetry, given that it starts at 0.95?..]

Thank you for bringing this up. It’s true that Figure 2 does not suggest perfect symmetry. Our original motivation for using a symmetric parameterization was inspired by [Hu et al., 2019. Exploring weight symmetry in deep neural networks.] who show that weight symmetry in NNs does not affect their universal approximation capabilities. However, we agree that it would be interesting to test a non-symmetric parameterization. Based on your suggestion we have implemented a new parameterization for $\mathbf{H}$ as $\mathbf{H} = \mathbf{V}_H \Lambda \mathbf{V}_H^{-1}$. Specifically, we take gradients to $\mathbf{V}_H$ in the backwards pass, and compute $\mathbf{H}$ using the above equation in the forwards pass. We use the suffix “-NS” for the non-symmetric parameterizations. Here are the new accuracy results:

(response continues below)

Dear Reviewer rqLo,

Thank you again for your time on the paper. Do you have any extra questions about the paper or our response?

Let us know and we will try to response as soon as possible before the discussion period ends on November 22.

Thank you,

The authors

Dear Reviewer m7iA,

Thank you so much for your time and feedback. We respond to each question below.

[..Weaknesses:..]

[..an expectation for a broader and more diverse range of tasks...what is the performance impact of this approach on NLP tasks?..]

We definitely aim to investigate this in future work. In large part, most research we are aware of on Transformer oversmoothing divides into either vision-focused or NLP-focused work, both of which are usually entirely empirical (with the rare exceptions of Dong et al., 2021 and Wang et al., 2022). Because our theory extends the theoretical work of Dong et al., 2021 and Wang et al., 2022 who restrict their focus to vision Transformers, we originally opted to focus solely on vision as well.

This said, we agree that it would strengthen the paper to extend the evaluation to another domain. We are currently running an NLP experiment that should be ready on Monday. We will get back to you with these results as soon as possible!

[..Questions..]

[..Can the proposed approach effectively mitigate oversmoothing even when employing an exceptionally high number of layers?..]

This is a good question! To answer this we have added an experiment where we increase the depth of the ViT-Ti models on CIFAR100. We have also added a parameterization variant where the first L/2 layers sharpen (as in ViT-Ti$^-$) and the remaining $L/2$ layers are left normal (as in ViT-Ti), which we call ViT-Ti*. Finally we have added a comparison FeatScale by Wang et al., 2022. All results are averaged over 3 trials with standard deviations (we bold the best result in each column and any method whose mean plus standard deviation is greater than or equal to mean of the best result):

[new] Increased Depth (Table 3)

Here ViT-Ti$^+$ and ViT-Ti$^-$ both underperform ViT-Ti, and we suspect this is because ViT-Ti$^-$ sharpens too much, amplifying noise in the data (whereas ViT-Ti$^+$ oversmooths the data). On the other hand ViT-Ti* significantly outperforms all other models not just for larger depths in Table 3, but also in corruption robustness (Table 4), and data efficiency (Table 2). This suggests that both oversmoothing and oversharpening are detrimental to performance, and that a careful balance is key to an accurate model. We have updated the paper to describe these results and look forward to investigating this further in future work.

Thank you again for your time. Let us know if you have any other questions and we will respond as soon as possible.

Dear Reviewer m7iA,

Thank you again for your time on the paper. Do you have any extra questions about the paper or our response?

Let us know and we will try to response as soon as possible before the discussion period ends on November 22.

Thank you,

The authors

[..an expectation for a broader and more diverse range of tasks...what is the performance impact of this approach on NLP tasks?..]

We are still running NLP experiments and hope to have it by tomorrow