Transformer Meets Twicing: Harnessing Unattended Residual Information

Authors: Laziz Abdullaev, Tan Minh Nguyen
Primary area: foundation or frontier models, including LLMs
Keywords: transformers, self-attention, oversmoothing, nonlocal smoothing, nonparametric regression

Abstract:
Transformer-based deep learning models have achieved state-of-the-art performance across numerous language and vision tasks. While the self-attention mechanism, a core component of transformers, has proven capable of handling complex data patterns, it has been observed that the representational capacity of the attention matrix degrades significantly across transformer layers, thereby hurting its overall performance. In this work, we leverage the connection between self-attention computations and low-pass non-local means  (NLM) smoothing filters and propose the Twicing Attention, a novel attention mechanism that uses *kernel twicing procedure* in nonparametric regression to alleviate the low-pass behavior of associated NLM smoothing with compelling theoretical guarantees. This approach enables the extraction and reuse of meaningful information retained in the residuals following the imperfect smoothing operation at each layer. Our proposed method offers two key advantages over standard self-attention: 1) a provably slower decay of representational capacity and 2) improved accuracy across various data modalities and tasks. We empirically demonstrate the performance gains of our model over baseline transformers on multiple tasks and benchmarks, including image classification and language modeling, on both clean and corrupted data.

Decision: accept (poster)
Reviewer scores: 6, 6, 8, 5

Reviews and discussion:
<thread>
creation date:
1730480016960

thread_type:
official_review

<thread_object>
creator: reviewer_7qsp

<content>
summary:
This paper introduces the Twicing Attention mechanism, drawing inspiration from established connections between self-attention and low-pass non-local means (NLM) smoothing filters. The authors demonstrate two key advantages of their proposed method: 1) a theoretically proven slower decay of representational capacity across transformer layers, and 2) improved performance on both vision and language tasks across multiple datasets. The paper's primary contribution lies in its theoretical framework. It first establishes that representation collapse in transformers stems from the inherent low-pass characteristics of NLM filters. The authors then provide proof showing that the twicing formulation ($2A^2-A$) offers superior theoretical properties compared to standard attention ($A$), particularly in preserving token diversity and meaningful feature representations.

strengths:
- Theoretical foundation: The paper provides analysis connecting self-attention to NLM filters. 
- Fluent presentation flow: Messages of this paper are presented well, with well-demonstrated background knowledge and motivation.
- Empirical validation: The authors provide visualizations of attention heatmaps, which validates their claim that their method preserve the diversity of token representations

weaknesses:
- **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.
- **Additional computational cost and marginal empirical improvements**: Performance increase in Table 4 is in trade of computational cost. Hardly can engineers be convinced to substitute the original attention with the proposed one.
- **Limited Evaluation Scope**: The authors report the empirical performance on classification tasks for vision models. Yet dense tasks such as segmentation are more direct and effective in evaluating the structure of patch representations produced by the method.

questions:
- Visualizations on earlier layers and more heads of the transformers would help to strengthen your claim.
- Please refer to the weakness.
- I am open to increase my score if you alleviate my concerns.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730627058486

thread_type:
official_review

<thread_object>
creator: reviewer_cqxc

<content>
summary:
This paper propose the Twicing Attention, a novel attention mechanism that uses kernel twicing procedure in nonparametric regression to achieve slower decay of representational capacity and improved accuracy across various data modalities and tasks. And the design of this module builds on the study of the connection between self-attention and NLM smoothing filters. The method was tested on a public dataset, yielding promising results.

strengths:
1. The paper is well-written, with clear expression of formulae and symbols, and will be highly readable.
2. The authors discuss the recurring problem of decay of representational capacity in Transformer, which has also been recognized as a manifestation of rank collapse in other studies. Instead of directly continuing the study of related work on rank collapse, the authors start with the NLM and try to gradually restore the cause of this phenomenon and again based on the proposed method that can alleviate this problem, the research angle is more interesting and also quite theoretical significance.
3. The author's description of the solution is complete and accompanied by thorough proofs, the process is clear and easy to understand, and the work done is very informative.

weaknesses:
1. 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.
2. 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.

questions:
1. 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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730638997089

thread_type:
official_review

<thread_object>
creator: reviewer_3ads

<content>
summary:
The self-attention mechanism's representational capacity diminishes significantly across layers, and this oversmoothing effect is reducing overall performance. This paper introduces Twicing Attention, a novel mechanism that connects self-attention computations with low-pass non-local means smoothing filters. By employing a kernel twicing procedure, it alleviates the low-pass effects of NLM smoothing while preserving meaningful information from residuals. Twicing Attention offers slower decay of representational capacity and improved accuracy across different data modalities. Significant performance improvement brought by Twicing attention is observed in multiple tasks.

strengths:
1. Novelty: The authors introduce the Twicing Attention mechanism to slow down the eigenvalue decay associated with representational collapse.
2. Theoretical Contribution: the authors provide mathematical validation for the Twicing Attention’s capability to retain information across layers.
3. Experiments: the authors evaluate their proposed methods on both language models and vision models.

weaknesses:
1. Limited improvement: The gains in clean data settings (such as on ImageNet in Tab. 1) are modest.
2. Lack of comparison: the work does not compare its method with alternative solutions that address oversmoothing, such as regularization strategies.
3. Lack of ablations: the authors are suggested to consider applying the proposed method at different layer depths or intervals and evaluate their difference.

questions:
My question lies in the efficiency comparison (Tab. 4). Despite the fact that Twicing has the same complexity of $O(N^2 d)$ 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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730686016849

thread_type:
official_review

<thread_object>
creator: reviewer_pr7y

<content>
summary:
The over-smoothing problem in Transformers is a well-known phenomenon, where the outputs of different attention layers in a Transformer model are highly similar. This paper introduces Twicing Attention to address this problem, which uses low-pass NLM smoothing filters to tackle this problem. The core idea can be phrased as, instead of using the standard attention matrix $A$, to use $2A - A^2$.

strengths:
1. The paper is relatively easy to follow and well-written.
2. The proposed "Twicing Attention" is simple and easy to implement.
3. Theoretical motivation and the mathematical details behind their motivation and choices have been provided.

weaknesses:
1. 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.
2. 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

questions:
I would be grateful if the authors could respond and address the weaknesses. I am willing to increase my score if the authors could address the weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732067679789

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Global Response

comment:
Dear AC and reviewers,

First of all, we thank all reviewers for their endorsements as well as valuable feedback on our work. In particular, reviewers' positive comments on the clarity of our presentation (pR7y, CqxC, 7QSP), significance of our theoretical contribution (all 4 reviewers), and informativeness of the paper (CqxC, pR7y, 7QSP) have been encouraging for us.

In this global response, we would like to address some of the shared concerns among the reviewers as well as reiterate and clarify what major benefits Twicing Attention can offer, especially other than merely alleviating representation collapse.

We respond to each common concern as follows:

1. **Limited accuracy improvement on clean data.** We agree that Twicing Attention offers relatively modest accuracy improvements in clean data settings. However, the clean data performance is not the only claim that we make about our model but improved overall accuracy (both under clean and corrupted data settings). Rather, we believe that the complementary robustness comparisons make the Twicing model stand out as a substantially better model overall. In particular, Twicing models show capability to offer up to a significant ~19\% improvement (FAN, PGD) with average of about ~8\% performance gains across all adversarial 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). 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.
2. **Additional computational cost for modest clean accuracy gain.** As mentioned in Point 1 above, the additional computation is also serving to obtain a significantly more robust model. In particular, notice how DeiT-Twicing is comparable to FAN against adversarial attacks while FAN introduces a more sophisticated architecture to achieve that. Additionally, refer to the relative improvements (\%) provided in Point 1. It is worth noting that most tailored robust models available also introduce similar (or sometimes more) computational complexity compared to Twicing (added a new paragraph in Related Works for this comparison). This is also sometimes known as the robustness-efficiency trade-off (RETO) which is hardly avoidable.
4. **Narrow problem formulation **(respresntation collapse)**.** While we genuinely understand why some reviewers tend to think that the paper only deals with representation collapse due to our problem introduction style, we would like to reiterate an almost equally important subject of our paper--improving the underlying theoretical denoiser/estimator framework through the twicing procedure [3, 4]--which also ensures more robustness [1, 2, 4] as it helps the model learn more stable representations. Furthermore, another importance of such a theoretical observation along with empirical justification is that it could foster interesting future research to explore more similar frameworks to improve deep learning models in various aspects. In light of this concern, we have adjusted our introduction and the following sections to give a little more importance to the robustness of Twicing Attention both theoretically and empirically.

### 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.

[5]: Victor Chernozhukov, Juan Carlos Escanciano, Hidehiko Ichimura, Whitney K. Newey, and James M. Robins (2022): "Locally robust semiparametric estimation". Econometrica, 90(4):1501–1535

[6]: Caron, M., Touvron, H., Misra, I., Jégou, H., Mairal, J., Bojanowski, P., & Joulin, A. (2021). Emerging properties in self-supervised vision transformers. Proceedings of the International Conference on Computer Vision (ICCV).

[7]: Darcet, T., Oquab, M., Mairal, J., & Bojanowski, P. (2024). Vision transformers need registers. Published as a conference paper at ICLR 2024.

-----

We are glad to answer any further questions you have on our submission.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732067808605

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Summary of Revisions

comment:
According to the comments and suggestions from reviewers, we have applied the following changes to the updated version of the paper:

1. General edit: We have cut the text on pages 7 and 8 to make the presentation more compact than before as suggested. To fill the created space, we have added more experiments, insights and theoretical explanations for the robustness of Twicing Attention, an aspect of our model which we seem not to have emphasized enough before. In particular, we associate the small bias property of the underlying twicing kernels to the reduction of bandwidth sensitivity [4, 1] and robustness to input perturbations through [5].
2. Extra experiments: We have conducted additional experiments on the image segmentation task on ADE20K dataset, and presented the comparison results between DeiT and DeiT-Twicing in Table 3 of the main text evaluated across three key metrics. We observe performance improvements over all metrics as reported. We have also provided the necessary experimental details in Appendix B.5. Besides, we have added a new model NeuTRENO (Nguyen et al, 2023) as a comparison model as requested. Also, we have trained a larger language modeling on Wikitext-103 to verify the scaling potential of Twicing Attention when implemented inside LLMs and obtained a positive answer as reported in Figure 7 and Table 6 of Appndix B.1.
3. Extra empirical analysis: As suggested by the reviewer 7QSP, we have provided the evolution of attention heatmaps for DeiT and DeiT-Twicing from early to late layers together with dozen of extra last layer heatmaps for more input images to strengthen our claims in Appendix D.2. We have also extended oversmoothing analysis in Figure 2 by conducting a similar experiment on ADE20K image segmentation task, and the results are positive and shown in Figure 8 in the appendix. In both cases, token smoothing is slower with Twicing Attention, validating our theoretical results.
4. Related works: We have added a discussion on the two papers [6, 7] studying the feature maps of Vision Transformers as suggested by the reviewer 7QSP since we found them indeed relevant. We have also added a new paragraph to the section dedicated to the research on robust transformer models building upon Point 1 of our Summary of Revisions.

### 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.

[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.

[6]: Caron, M., Touvron, H., Misra, I., Jégou, H., Mairal, J., Bojanowski, P., & Joulin, A. (2021). Emerging properties in self-supervised vision transformers. Proceedings of the International Conference on Computer Vision (ICCV).

[7]: Darcet, T., Oquab, M., Mairal, J., & Bojanowski, P. (2024). Vision transformers need registers. Published as a conference paper at ICLR 2024.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732286297019

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Additional Experimental Results

comment:
Dear reviewers,

We would like to thank all reviewers again for your thoughtful reviews and feedback. We have obtained additional experimental result that validates the concept of Twicing Attention is not tied to the exact form of standard softmax self-attention, but it offers improvements for any reasonable similarity matrices including different types and forms of attention mechanisms as described below.

We have conducted additional experiments with Linear Transformers [9] as described in our previous comment. Table A1 below compares the perplexities recorded for Linear Transformers with feature map $\phi(x) = \text{elu}(x)+1$ matching their original choice, and Linear-Twicing Transformers for which we apply the twicing transformation $2A-A^2$ where $A = \text{normalize}(\phi(Q)\phi(K)^\top)$. Note that we explicitly construct the similarity matrix $A$ for both of the models for our framework to work. On top of Table A1 results, we also observe relatively faster convergence for Linear-Twicing, very similar trend to what is illustrated in Figure 7 in the revised appendix of the paper. The positive results indicate that the applicability of Twicing Attention is not limited to standard softmax self-attention, but it is compatible with any reasonable similarity matrix. We have appended this result to Appendix B.1 and Table 6 of the revised document (highlighted by blue color).

**Table A1:** Validation/Test PPL on Wikitext-103 trained for 75 epochs.
| Model | Validation PPL | Test PPL |
|-------|----------------|----------|
| Linear Trans. | 40.00 | 41.26 |
| Linear-Twicing Trans. | **39.45** | **40.61** 

We would be happy to engage in any follow-up discussion or address any additional comments by the reviewers.

**References:**

[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.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper receives ratings of 8, 5, 6, 6, where the reviewers generally provided a positive assessment of this manuscript. In this paper, the authors propose Twicing Attention, a new self-attention mechanism designed to address representational collapse in transformers caused by over-smoothing. The authors established a connection between self-attention and nonparametric regression using twicing kernels. The proposed Twicing Attention mechanism leverages residuals between attention input and output, providing improvements in representational capacity and robustness. Experimental results also show taht Twicing Attention demonstrates consistent improvements across multiple tasks and benchmarks.

Strengths:
- The proposed Twicing Attention mechanism introduces a new perspective by leveraging residual information through twicing kernels, offering a novel solution to the representational collapse issue in transformers.
- Theoretical Rigor: this paper provides a strong theoretical foundation, connecting self-attention mechanisms to nonparametric regression techniques and demonstrating slower decay of representational capacity.
- Comprehensive experiments on vision (ImageNet, ADE20K) and language tasks (WikiText-103) showcase consistent performance improvements over baseline models. Enhanced robustness is also demonstrated under adversarial attacks and distribution shifts, indicating the practicality of the method.
- The mechanism is computationally efficient, requiring minimal additional overhead when selectively applied to specific layers of transformers.
- By addressing the limitation in transformer models (over-smoothing), the method has the potential to influence various applications in NLP, computer vision, and beyond.

Areas for Improvement:
- While the empirical results are compelling, comparisons with more diverse state-of-the-art methods could strengthen the claims o superiority and broader applicability.
- Some aspects of the theoretical framework, while rigorous, could be made more accessible to practitioners through intuitive explanations/visualizations.
- While the method is validated on several tasks, demonstrating its effectiveness rigorously across more diverse domains and larger datasets could enhance its impact.
- The paper briefly acknowledges limitations but could provide a more detailed discussion of areas for improvement and specific directions for future research.

The reviewers praised the paper’s strong theoretical underpinnings and empirical results, particularly its contributions to addressing over-smoothing in transformers. The paper also introduces a novel and elegant solution that aligns well with recent trends in enhancing transformer architectures. And therefore we recommend acceptance of this paper.