Transformer Meets Twicing: Harnessing Unattended Residual Information

Thank you for your thoughtful review and valuable feedback. Below we address your concerns.

Q1. Narrow Problem Framing: The paper's central premise regarding "representation collapse" in transformers warrants closer scrutiny. Recent research has demonstrated that this phenomenon is not inherent to transformer architectures. For instance, DINO(Caron et al., 2021) demonstrates that self-supervised training can produce well-structured, diverse token representations in Vision Transformers. Furthermore, Darcet et al. (2024) provide evidence that apparent "collapse" may actually reflect a more nuanced information distribution pattern, where artifacts in attention heatmaps encode global information while non-artifact tokens maintain distinct representations, albeit with lower similarity to the CLS token.

Answer: We find the two papers that the reviewer brings to our attention interesting in terms of characterizing and understanding the emergent artifacts in the feature maps. Intriguingly, interpreting artifacts as locations where the model stores global input information—elegantly handled with extra register tokens in Darcet et al. (2024)—aligns (in a weak sense) with an image denoising perspective as well. When weighted averaging (blurring) is repeatedly applied, sharp edges are smoothed out, letting global information coming from large background sections dominate the output image. We note that the twicing procedure [3, 4] is tailored to guide the model to benefit from those local information and details before proceeding with another blurring iteration to accomodate both local and global information flow.

On the other hand, there are at least a few fundamental scope differences between the cited papers and ours, and our subject of study is not limited to representation collapse: (i) we mainly focus on the interpretation of our method through the perspective of twicing procedure and its analytical and statistical properties; (ii) while slowing down the decay of representational capacity is one of our contribution, it is not the only one. We believe the theoretical relation to twicing kernels with small bias property and its implications on learning more stable and robust representations is equally important matter of our paper; (iii) Unlike some prior works trying to mitigate over-smoothing completely by constantly fusing with initial layer tokens, we merely aim to slow it down to balance the mitigation of this phenomenon and largely deviating from the native behaviour of transformers to benefit from both worlds. All that being said, it is interesting to note how Twicing Attention heatmaps are usually concentrated over the body of the object while reducing the abovementioned artifacts as shown in a dozen of more sample images in Figure 11 in the appendix. Lastly, attention heatmaps are not the only way we illustrate "collapse", but we observe that, with Twicing Attention, average token similarities across layers indeed increase slower than the baseline as shown in Figure 2, which complements the other visualizations to validate slower collapse. Also, please see newly added Figure 8 in the appendix for both ImageNet and ADE20K oversmoothing analysis as a validation of our theoretical results on slower collapse.

References

[3]: Tukey, J.W. (1977). "Exploratory Data Analysis". Reading, MA: Addison-Wesley.

[4]: Stuetzle, W., and Y. Mittal (1979): "Some Comments on the Asymptotic Behavior of Robust Smoothers", in Smoothing Techniques for Curve Estimation, Lecture Notes, 757. New York: Springer-Verlag, 191–195.

We would like to thank the reviewer again for your thoughtful reviews and valuable feedback.

We would appreciate it if you could let us know if our responses have addressed your concerns and whether you still have any other questions about our rebuttal.

We would be happy to do any follow-up discussion or address any additional comments.

Thank you for your thoughtful review and valuable feedback. Below we address your concerns.

Q1. Limited improvement: The gains in clean data settings (such as on ImageNet in Tab. 1) are modest.

Answer: We agree that Twicing Attention offers relatively modest accuracy improvements in clean data settings in Table 1. However, the clean data performance is not the main/only claim that we make about our model but improved overall accuracy (both under clean and corrupted data settings). Rather, we believe that the complimentary robustness comparisons make the Twicing model stand out as a substantially better model overall. In particular, Twicing models show capability to offer up to ~19% improvement (FAN, PGD) with average of about ~8% performance gains across all attacks. Besides, Figure 4 in the appendix shows that Twicing Attention can notably and consistently outperform the baseline across all 15 types of natural corruption types (about ~10% improvement on "contrast", "gaussian noise", and "impulse noise" to name but a few). Zooming out a little bit, this observation is interesting and consistent with the theory in the following sense: it suggests that Twicing models are inherently more stable across both clean and corrupted data settings by prioritizing stable representations over overtuned accuracy on clean accuracy—an aspect that can make models more susceptible to small perturbations, such as common adversarial attacks. Additionally, drawing on the small-bias property of Twicing kernels in nonparametric regression, one can argue that the resulting estimator is relatively less sensitive to bandwidth selection [1]. This reduced sensitivity mitigates the bias fluctuations often introduced by slight adjustments, making the estimator inherently more resilient to minor perturbations and improve model's robustness in general. Our experimental results under 3 widely adopted adversarial attacks validate that Twicing Attention is indeed significantly more robust compared to the baseline self-attention. We also refer to [2] for a more detailed robustness of twicing kernels in regression compared to the Nadaraya-Watson estimator kernel before twicing. At the same time, nonetheless, we also see in Table 4 of the revised document that improvements on clean and contaminated data for language modeling are comparable.

[1]: Newey, W.K., F. Hsieh, and J.M. Robins (2004). "Twicing Kernels and a Small Bias Property of Semiparametric Estimators." Econometrica, Vol. 72, No. 3, pp. 947–962.

[2]: Chernozhukov, V., Escanciano, J. C., Ichimura, H., Newey, W. K., & Robins, J. M. (2022). Locally robust semiparametric estimation. Econometrica: Journal of the Econometric Society.

Q2. Lack of comparison: the work does not compare its method with alternative solutions that address oversmoothing, such as regularization strategies.

Answer: We have conducted additional experiments comparing our method with an alternative model, NeuTRENO [8], that uses a nonlocal functional regularization to mitigate oversmoothing by constantly fusing with initial layer tokens. In the table below, we report the Top 1/Top 5 accuracy on ImageNet, as well as their robustness against PGD, FGSM and SPSA adversarial attacks. We observe that while both models outperform the baseline DeiT, DeiT-Twicing offers relatively larger improvements in almost all metrics.

Table C1: ImageNet classification under clean and adversarially attacked settings.

[8]: Nguyen, T. M., Nguyen, T. M., & Baraniuk, R. (2023). Mitigating over-smoothing in transformers via regularized nonlocal functionals. NeurIPS 2023.

Q3. Lack of ablations: the authors are suggested to consider applying the proposed method at different layer depths or intervals and evaluate their difference.

Answer: In Appendix E, we compare 3 different choices of layer placements for Twicing Attention [1 to 12, 7 to 12, and 10 to 12]. As a result, we observe in particular that overall performance is roughly proportional to the number of Twicing layers in terms of clean data and under adversarial attacks. In Table C2 below, we report 2 more models--twicing at even layers, and using previous layer residual for twicing procedure for efficiency. We observe that even though DeiT-Twicing [even layers] has 6 Twicing Attention layers just as DeiT-Twicing [7-12], the latter model does better than the former. This validates that if one has capability to implement $n$ layers (out of $L > n$ total layers) with Twicing Attention, it is better to place them as the latest n contiguous layers of the transformer.

Table C2: Ablation of Twicing Attention placed at different layers.

Q4. My question lies in the efficiency comparison (Tab. 4). Despite the fact that Twicing has the same complexity of as claimed in the paper, it still increases the overhead by an additional 50% due to the extra matrix multiplication in line 7, Alg. 1. However, Tab. 4 indicates that implementing Twicing or not will not incur big difference on both speed and GFLOPs. What is the reason behind that? I would appreciate a more detailed efficiency analysis & comparison in the rebuttal phase if possible.

Answer: We appreciate the reviewer’s attention to a potential source of confusion regarding Table 4. We elaborate on the details of that efficiency analysis as follows. Our model does not add 50% more computational cost as reported is the efficiency statistics considering the end-to-end flow of an input through the transformer. In fact, the additional computation--specifically calculating $A(V - A V)$ (with the pre-computed $AV$ as in Algorithm 1)--only marginally increases the total workload when considering the entire Transformer architecture. While this operation does add extra steps to the attention mechanism, the overall computational cost is dominated by other components, such as the feed-forward networks and linear transformations. These components combined require significantly more computation than the attention mechanism alone. Furthermore, attention layer itself is not doubled in terms of computational complexity since Twicing only adds an extra attention-weighted averaging while the basement of standard self-attention already consists of computing $QK^T$ and $\text{softmax}(\cdot)$ (which we do not repeat for Twicing) other than $AV$ matrix operation. As a result, theoretically, the added computation increases the total computational cost by only roughly about 7% (which is approximately consistent with Table 4 results) if we analyze a modestly simplified Transformer architecture in terms of component-wise runtime complexities and their contributions to the overall computational overhead. Considering the partial model, Twicing [10-12], we see that the additional overall computation is virtually negligible while offering a decent relative performance. It is also worth noting that since Twicing does not introduce any learnable parameters, its contribution to complexity of backward passes is minimal during pre-training.

We hope we have cleared your concerns about our work. We have also revised our manuscript according to your comments, and we would appreciate it if we can get your further feedback at your earliest convenience.

We would like to thank the reviewer again for your thoughtful reviews and valuable feedback.

We would appreciate it if you could let us know if our responses have addressed your concerns and whether you still have any other questions about our rebuttal.

We would be happy to do any follow-up discussion or address any additional comments.

On top of DeiT-NeuTRENO model compared in our previous response, we conducted an additional experiment with the FeatScale [10], another state-of-the-art vision transformer variant that tries to mitigate representation collapse. As shown in Table A2 below, our model DeiT-Twicing outperforms DeiT+FeatScale in both metrics on ImageNet classification. We report the same results in Table 9 of Appendix B.2 of our revision.

Table A2: Top 1/Top 5 accuracies on clean ImageNet classification.

[10]: Peihao Wang, Wenqing Zheng, Tianlong Chen, and Zhangyang Wang. Anti-oversmoothing in deep vision transformers via the fourier domain analysis: From theory to practice. In International Conference on Learning Representations, 2022.

We would appreciate it if you could let us know if our responses have addressed your concerns and whether you still have any other questions about our rebuttal. We would be happy to do any follow-up discussion or address any additional comments.

If you agree that our responses to your reviews have addressed the concerns you listed, we kindly ask that you consider whether raising your score would more accurately reflect your updated evaluation of our paper. Thank you again for your time and thoughtful comments!

We would like to thank the reviewer again for your valuable initial reviews and feedback.

For additional robustness comparison, we benchmarked DeiT-NeuTRENO against ImageNet out-of-distribution and natural image corruptions, and found our model DeiT-Twicing being better in 3 out of 4 tests except ImageNet-A. However, NeuTRENO deteriorates the performance of the baseline DeiT in ImageNet-R, an observation which supports our model DeiT-Twicing to be more stable with no extra hyper-parameters introduced. This additional result has been included in Table 9 of the revised document appendix.

We hope this additional results complement our previous response to your question and clears your related concerns. We would be glad to hear your futher feedback on our work and rebuttal at your earliest convenience.

We sincerely appreciate the reviewer’s explanation of your score and your positive assessment of our paper's quality. We recognize that the applicability of any method to LLMs is a significant factor. While we agree with the reviewer that pretraining a large model with Twicing Attention applied at all layers may not be suitable for low-budget scenarios, we would like to share the following results and insights as a potential use case for Twicing Attention, both with off-the-shelf pretrained models and in full pretraining scenarios with lower budgets.

Fine-tuning a pretrained Switch Transformer. To show how Twicing Attention can offer improvements to the pretrained models, we pretrain a medium sized (33M params) Switch Transformer [11], a Mixture of Experts architecture, with the standard self-attention on WikiText-103. Then we finetune this pretrained language model on Stanford Sentiment Treebank 2 (SST-2) dataset using standard self-attention (baseline) as well as Twicing Attention (ours) for 8 epochs. Table L1 compares Top 1 finetune test accuracies for both cases and we find that finetuning with Twicing Attention achieves higher accuracy, provided that the fine-tuning is long enough (usually a few more epochs than usual) for the model to adapt to the new attention mechanism.

Table L1: Switch Transformer Pretrained on WikiText-103 and Finetuned on SST-2.

Partial Twicing model. Additionally, we would like to highlight how DeiT-Twicing [10-12] (last 3 layers only) increases the FLOPs by just over 1% while improving robustness by 14.3% (ImageNet-A), 2.7% (ImageNet-C) [Table 2 of the paper] even surpassing the full model, and 5.5% (FGSM) [Table 1 of the paper]. We believe such a partially deployed Twicing Attention allows its application for almost negligible extra cost in practice.

Thank you once again for your time and thoughtful feedback!

References:

[11]: Fedus, W., Zoph, B., & Shazeer, N. (2022). Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research.

Thank you for your thoughtful review and valuable feedback. Below we address your concerns.

Q1. Admittedly, the authors' work is very rich and makes a very profound contribution at the theoretical level, but in my humble opinion, the authors' approach serves as a skillful level of reconciliation that moderates the rank collapse in depth, whereas a similar reconciliation skill is actually not uncommon in rank collapse-related research directions. I am not accusing the authors of not being innovative enough, but I hope that the authors can go further at the theoretical level and expand the sequence model that can really reduce the loss of information unlike the classical Transformer.

Answer: We thank the reviewer for endorsing the theoretical contributions of the paper. While we agree that manipulating spectrum in different ways is not utterly uncommon in literature, we believe that a modification--in a way that it further connects the attention mechanism to the well-established Twicing procedure in image processing [3] and nonparametric regression [4] with small bias property [1]--is a novel theoretical perspective. Unlike pure engineering heuristics or linear algebraic approaches to moderate the rank collapse, our work offers broader interpretability of Transformers along with the modified attention mechanism. Additionally, we believe that this interpretation can foster further research to study the potential benefits that such traditional statistical frameworks still has to offer for the modern deep learning theory.

Q2. The author's research is more profound, but the experiments are less adequate, with too few test datasets and too few comparison methods. I tend to think that this is the result of too much time constraints, and I hope that the author will add more datasets as well as other experiments on the Transformer if there is enough time.

Answer: We appreciate the reviewer's understanding of time constrains. We took our chance to carry out extra image segmentation experiments on another widely adopted dataset ADE20K and report the pixel accuracy, mean accuracy, and mean intersection over union (IoU) metrics to compare against the baseline in Table D below. We find that Twicing Attention offers improvements across all three metrics evaluated.

Table D: Image segmentation on ADE20K.

Furthermore, we have done experiments with an additional competetitor model, NeuTRENO (Nguyen et al, 2023), that uses a nonlocal functional regularization to mitigate oversmoothing by constantly fusing with initial layer tokens. In the Table E below, we report the Top 1/Top 5 accuracy on ImageNet as well as their robustness against PGD, FGSM and SPSA adversarial attacks as in the paper. We observe that while both models outperform the baseline DeiT, our DeiT-Twicing offers relatively more improvements in almost all metrics.

Table E: Image classification on ImageNet-1K.

Q3. For pure curiosity, I would like to ask what the authors think the performance of this method would be in more extreme cases, which in this case refers to two main scenarios: first, the performance on LLMs with a very large number of parameters. Second, on non-classical Transformer structures, such as Linear Transformer and other analogs.

Performance on LLMs: To answer the question about the potential performance on LLMs with larger number of parameters, we trained a medium-sized model with 20.97M parameters compared to the small-model with 9.43M parameters. As a result, we observe that Transformer-Twicing still offers improvements across both validation and test perplexities indicating scaling properties about as good as the baseline Transformer. Also, in Figure 7 of the appendix, we show the training curves for the language models of both sizes, and observe that Twicing Attention helps the models converge relatively faster as well.

Table F: Language modeling on Wikitext-103.

Extreme Unconventional Transformers: Since the Twicing Attention's theoretical framework does not depend on how exactly the weight matrix is built, we believe that as long as any Transformer architecture-based model leverages an attention mechanism with a concrete similarity (attention) matrix that can be connected to either NLM denoising or nonparametric Nadaraya-Watson estimation as in the paper, Twicing is highly likely to offer extra representation capacity and robustness. In particular, as transformers with linear attention [9] are concerned, the implementation steps would involve denoting their separable similarity matrix in Eqn. (4) of [9] as $A$, and replacing it with the corresponding twicing weight matrix $2A-A^2$.

References:

[1]: Newey, W.K., F. Hsieh, and J.M. Robins (2004). "Twicing Kernels and a Small Bias Property of Semiparametric Estimators." Econometrica, Vol. 72, No. 3, pp. 947–962.

[2]: Chernozhukov, V., Escanciano, J. C., Ichimura, H., Newey, W. K., & Robins, J. M. (2022). Locally robust semiparametric estimation. Econometrica: Journal of the Econometric Society.

[3]: Tukey, J.W. (1977). "Exploratory Data Analysis". Reading, MA: Addison-Wesley.

[4]: Stuetzle, W., and Y. Mittal (1979): "Some Comments on the Asymptotic Behavior of Robust Smoothers", in Smoothing Techniques for Curve Estimation, Lecture Notes, 757. New York: Springer-Verlag, 191–195.

[9]: Katharopoulos, A., Vyas, A., Pappas, N., & Fleuret, F. (2020). Transformers are RNNs: Fast autoregressive transformers with linear attention. In Proceedings of the International Conference on Machine Learning (ICML). PMLR.

We hope we have cleared your concerns about our work. We have also revised our manuscript according to your comments, and we would appreciate it if we can get your further feedback at your earliest convenience.

We would like to thank the reviewer again for your thoughtful reviews and valuable feedback.

We would appreciate it if you could let us know if our responses have addressed your concerns and whether you still have any other questions about our rebuttal.

We would be happy to do any follow-up discussion or address any additional comments.

We would like to thank the reviewer again for your valuable initial reviews and feedback.

We took our time to also test Linear-Twicing model's robustness againt Word Swap Attack using the same experimental setting as in Section 4.2 of the main text. As we show in Table B1 below, Linear-Twicing Transformer offers a significant improvement of almost 2 PPL points over the baseline Linear Transformer. This further validates that Twicing Attention can indeed enhance robustness of the model even with unconventional attention mechanisms. This result has been included in Table 7 of the revised document appendix.

Table B1: Clean/Attacked Test PPL on Wikitext-103 under Word Swap attack.

We hope this additional result provides additional justification of our insights provided in the previous replies and further addresses your question.

Thank you for your thoughtful review and valuable feedback. Below we address your concerns.

Q1. The paper compensates for the simplicity of the core idea by over-explaining and being overly verbose. For example, most of the material on pages 7-8 can be summarised in 2-3 paragraphs. Even Algorithm 1 on page 8 is redundant and too verbose. The algorithm's objective is clear and simple: to compute $2A-A^2$. I don't think one needs 12 lines to explain that.

Answer: While we intended to make our narrative of NLM denoising and nonparametric regression perspective more comprehensible to the readers, we agree with the reviewer on the point that the content on the page 7, in particular, can be compressed into a more compact text. We have editted that section in our revision to achieve a concise alternative accordingly. We use the created space for extra insights on the robustness of Twicing Attention (Section 3.3), as well as extra experimental results (Section 4.1).

Q2. Instead, the paper could have added to its contribution through a more thorough study. E.g., one avenue for improvement would be to consider other candidates besides the $2A-A^2$ and then compare them in the considered benchmarks.

Answer: Even though we agree that there is still room for further empirical studies, we would argue that considering other "candidates" besides $2A-A^2$ is actually a little bit controversial since $2A-A^2$ is an only theoretically justified choice enabling us to study Twicing Attention through the lens of the well-established twicing procedure (which actually sets the paper title) in image reconstruction theory and nonparametric regression regime [3, 4]. The identity $(2A_\ell-A_\ell^2)V_\ell = A_\ell V_\ell + A_\ell(V_\ell - A_\ell V_\ell) = A_\ell V_\ell + A_\ell\cdot \text{r}_{\ell}$ is a quick way of reiterating the core motivation behind this very choice.

For the sake of comparison and further insights, however, we have conducted additional experiments to study other candidates that are intended to approximate the twicing procedure without compromising the baseline efficiency. We report the results in Table A below. Note that each model in Table A is trained for 200 epochs. The compared models in this study are inspired by the core idea of twicing procedure--adding back the smoothed residual to the estimation. We observe that efficient approximations often exhibit faster initial convergence rates; however, they are less stable tend to fall behind the full model in later stages of training, as they struggle to capture and learn the more complex patterns which models are expected to learn in later stages. We still believe that such efficient versions can be made work well, yet we leave it for future research.

Table A: Comparison of DeiT-Twicing and its efficient approximations as explained.

For further comparison with different candidates, we introduced a hyper-parameter into Twicing Attention as $AV + \lambda A(V-AV) = [(1+\lambda)A - \lambda A^2]V$. Then, we train this model on ImageNet classification with $\lambda = 1/2$, which lies right in the middle of baseline self-attention and our Twicing attention, so that it can capture the general effect of such scaling. We present our results in Table B below. While we find that it still offers improvements over the baseline, it falls behind the original Twicing Attention, and justifies the use of the proposed model with theoretical support.

Table B: Comparison of DeiT-Twicing models using $2A-A^2$ and $(1+\lambda) A - \lambda A^2$ as a similarity matrix.

We hope we have cleared your concerns about our work in this response and revised document. We would appreciate it if we can get your further feedback at your earliest convenience.

We would like to thank the reviewer again for your thoughtful reviews and valuable feedback.

We would appreciate it if you could let us know if our responses have addressed your concerns and whether you still have any other questions about our rebuttal.

We would be happy to do any follow-up discussion or address any additional comments.