Transformer Meets Twicing: Harnessing Unattended Residual Information
Authors: Laziz Abdullaev
Tan Minh Nguyen
Venue: ICLR 2025 Poster (2025)
Keywords: transformers
self-attention
oversmoothing
nonlocal smoothing
nonparametric regression
Primary area: foundation or frontier models, including LLMs
TL;DR: We propose a novel attention mechanism to enhance expressive power of transformers by leveraging residual information.

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.

Published as a conference paper at ICLR 2025
TRANSFORMER MEETS TWICING:
HARNESSING
UNATTENDED RESIDUAL INFORMATION
Laziz U. Abdullaev
Tan M. Nguyen
Department of Mathematics National University of Singapore a z 2. abdul laev@u u. nus edu
Department of Mathematics
National University of Singapore t Canmn& nus. edu sg
ABSTRACT
Transformer-based deep learning models have achieved state-of-the-art perfor- mance across numerous language and vision tasks. While the self-attention mech- anism, core component of transformers, has proven capable of handling com- plex 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 behav- ior of associated NLM smoothing with compelling theoretical guarantees and en- hanced adversarial robustness. This approach enables the extraction and reuse of meaningful information retained in the residuals following the imperfect smooth- ing 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 robustness and 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. The code is publicly avail- able https: http//gtttt thub. com/ az zcodes twi twiiiiiiiin _g tent on.
INTRODUCTION
Attention mechanisms and transformers (Vaswani al., 2017) have achieved state of the art perfor- mance across a wide variety of tasks machine learning (Khan et al., 2022; Lin et al., 2022; Tay al., 2022) and, in particular, within natural language processing (Al-Rfou et al., 2019; Baevski & Auli, 2018; Dehghani et al., 2018; Raffel et al., 2020; Dai et al., 2019), computer vision (Liu et al., 2021; Touvron et al., 2021; Radford et al., 2021), and reinforcement learning (Janner et al., 2021; Chen et al., 2021). They have also demonstrated strong performance in knowledge transfer from pretraining tasks various downstream tasks with weak or no supervision (Radford et al., 2018; 2019: Devlin et al., 2018). At the core of these models is the dot-product self-attention mechanism, which learns self-alignment between tokens an input sequence by estimating the relative impor- tance of each token with respect to all others. The mechanism then transforms each token into weighted average of the feature representations the other tokens with weights proportional to the leamned importance scores. The relative importance scores capture contextual information among tokens and are key to the success of the transformer architecture (Vig & Belinkov, 2019; Tenney et al., 2019; Cho et al., 2014; Tran al., 2025; Parikh et al., 2016; Lin et al., 2017; Nguyen et al., 2021).
Even though deep transformer-basee models have achieved notable success. they are prone to the representation collapse issue, where all token representations become nearly identical as more layers are added. This phenomenon, often referred to the "over-smoothing" prob- lem, substantially reduces the transformers' ability to represent diverse features (Shi et al., 2022; Wang al., 2022; Nguyen et al., 2023a; Devlin et al., 2018; Nguyen al., 2024a). Correspondence laziz, abdul Laevvu. eds

---

Published as a conference paper at ICLR 2025
DeiT
DefffWwcing
DelT
cei
Dei'T (Touvron et al., 2021) and DeiT- Twicing (ours) attention heatmaps. Our model shows better representational capacity compared to the baseline by paying attention to more meaningful parts of objects while DeiT attention scores are collapsed to one or few points.
Figure 2:
To demonstrate this phenomenon, we analyze the average cosine similarity between token pairs across layers a softmax transformer trained for the Imagenet classification tasks. As shown Figure the cosine similarity be- tween tokens increases with depth. In the fi- nal layers, the cosine similarity scores are just under 0.9, suggesting high level of similarity among token representations.
ImageNet Classification
Token Cosine Similarity
DeiT
0.8
DelT-Twicing
0.5
3
10
12
Layer
A prior line of research explores representation collapse in transformers through the lens of im- age denoising, showing that self-attention com- putation is equivalent to a gradient descent step towards minimizing energy functional that pro- motes smoothness in the input image (Nguyen al., 2023b; Gilboa & Osher, 2007). Additionally, investigating the over-smoothing phenomenon from graph-based perspective has gained significant attention in recent studies (Wu et al., 2023; Shi et al., 2022; Nguyen et al., 2023a; 2024b).
Figure I: Average token cosine similarities across layers of DeiT and DeiT-Twicing over 100 dom samples. Our model retains better token di- versity compared to the baseline.
Contribution. In this work, we take the connection between the self-attention mechanism and the monss image smoothing filter (Buades al., 2005) further, and show that rapidly vanish- ing eigenvalues of associated NLM filter across iterations major cause of representation collapse transformers. The NLM similarity matrix, the heart of NLM smoothing procedure, computes pairwise similarities between image patches based intensity differences, effectively serving weight matrix in the smoothing process. We then propose the Twicing Attention, novel attention mechanism, redesigned from the modified NLM smoothing operation that is tailored decrease the rate of decay of the eigenvalues of the NLM similarity matrix and thereby offering advantages over the standard NLM based self-attention. In particular, establish a connection between mod- ification technique and the twicing kernels in nonparametric regression (Stuetzle & Mittal, 1979; Newey al., 2004; Abdous, 1995), uncovering the modified NLM filter's ability exploit mean- ingful information in the residuals of each transformer layer after applying smoothing operation. In summary, our contributions are three-fold:
1. We develop the novel Twicing Attention mechanism, self-attention mechanism variant that promotes better token diversity across transformer layers which also enjoys enhanced robustness.
2. We develop theoretical framework highlighting the effectiveness of Twicing Attention in mitigating representational collapse by decelerating the rate of eigenvalue vanishing phe- nomenon.
3. We show, through the lens of twicing kernels in nonparametric regression, how unattended but useful residual information between self-attention input and output can be used self-correction at each transformer layer.
Moreover, we empirically validate the performance improvements of Twicing Attention over stan- dard self-attention in large-scale tasks such ImageNet-1K classification (Touvron et al., 2021),
2

---

Published as a conference paper at ICLR 2025
ADE20K image segmentation (Strudel et al., 2021) and WikiText-IO3 103 language modelling (Merity et al., 2016), and offer additional insights into its implementation with minimal additional compu- tational overhead. We also assess its robustness against adversarial attacks, data contamination, and various distribution shifts.
Organization. The paper is written in the following structure: In Section 2, we introduce the reader with some background context on self-attention mechanism and its connection to image smoothing operation as warm up to achieve better readability overall. In Section 3, we leverage the connec- tion between self-attention mechanism and nonlocal-means (NLM) smoothing filters to show that representation collapse phenomenon is particularly caused by low-pass behaviour of such filtering procedure. Then, we propose novel technique to alleviate the low-pass behaviour of associated NL.M smoothing, thereby enabling a redesign of the standard self-attention mechanism with better expressive power across the transformer layers. In Section 4, we present our experimental results using Twicing Attention while Section 6 contains a brief overview of related work in the literature. Finally, we end with concluding remarks in Section 7 and defer most of the technical proofs and derivations as well as extra experimental observations to appendix.
BACKGROUND
2.1
SELF-ATTENTION MECHANISM
Given an input sequence X = [01...., ØNJ" IR NxD, of N feature vectors, the self-attention mechanism transforms the input to U= [u1
uN I NxDs as follows:
x W. K
Wo
uui)
softmax
Wvæj
VD
MM2 softmax
9 VVD 1i
for = NN, where softmax(a)) (a,) == softmax (a), for [01,... ......... aN] is abuse of notation
for convenience. The vectors q. ky. and Vy, j = 1, I....,,N, N. are the query, key, and value vectors,
respectively. They are computed as Q = [q....., qv = XW XWT RNND and V == [V,,..., ] XW, R.NxD where Wq. W k DxD. Wy DDDxD are the weight matrices. Eqn. be expressed in matrix form as:
o N K = [ki,..., ky]
QK
U softmax
where the softmax function is applied row-wise to the matrix QK//VD. We refer to transformers built with Eqn. 2 as standard transformers or just transformers.
2.2 NONLOCAL VARIATIONAL DENOISING FRAMEWORK FOR SELF-ATTENTION
Based on the framework established by (Nguyen et al., 2023b), first consider the output matrix "N e RND in self-attention as given by Eqn. in Section 2.1. Let Q U := [u(1),..., ...,, u
R, E Ω2 and uxx) == [uj (x), ( be a real vector-valued function, u Ω + RD, u e (2)). The output matrix U in self-attention discretizes the function u(x) with respect to In the context of signal/image denoising, U can be considered as the desired clean signal, and u(x) is its corresponding intensity function denoting the signal values at the position E Q.. We further let the observed intensity function f(x) denote the values of the observed noisy signal at æ E 2, f IR, o e (2). For example, f (x) can be given as = (())=((>> n
where n is the additive noise (see Eqn. I of (Buades et al., 2005)). We wish to reconstruct u(x) from fx)) Following the variational denoising method proposed in (Gilboa & Osher, 2007), the denoised image u(r) can be obtained by minimizing the following regularized functional with respect to w:
E(u, f) (w) + G(u. f) = Jaxa [lu(z)) u(v) 12, v)dddy 2/ X f |u(z) f(x)
S
3

---

Published as a conference paper at ICLR 2025
Here, = 5ax |u(x) (()))>u(ylāăud (1y)l y)dzdy is nonlocal functional of weighted differ- ences. The weights w(x,V) v) represent the affinity between signal values at positions and y- For example, for images, w(x,)) y) captures the proximity between pixels a and in the image. (u) works regularizer. Minimizing Ju) promotes the smoothness of u and penalizes high- frequency noise in the signal as discussed in the next section. Adding the convex fidelity term G(u, 2/2 la [u()) (x) with the regularization parameter X, to the functional J(u) allows the denoised signal (()r) to preserve relevant information the observed noisy signal f()). In the following section, we show that NLM algorithm for image filtering corresponds to fixed point iteration step to solve the stationary point equation of J.
=
2.3 TRANSFORMERS IMPLEMENT ITERATIVE SMOOTHING
Note that the functional Jur imposes stronger penalty discontinuities sharp transitions in the input signal u, thereby promoting smoothness throughout the signal. To get the minimizer of (u, f), we consider the following system of equation:
f(()(z).f(z)) aJ=(uxx)) M(u(x)- f(x)) =0, Vor Ω. @u(x)
Ou(x)
Direct gradient calculation, detailed in Appendix A.3, then yields
So (((())) u())))((z,))) M(u(x) f(x)) =0 Va Q.
(6)
Rearranging the terms in Eqn. 6, we obtain
u( x Af(x) So w( () A Sa w (r.y)dd
Vx f
worth noting that Eqn. 7 becomes NLM filter with weights w(x, y) when À = 0 (see Eqn. of (Buades et al., 2005)). In order to establish connection between NLM and self-attention, let k(x) := [Ai (z),.., ., (=)]² be a real vector-valued function, k: Ω + RR, k ε L L'()). Similar to u ((x) and v(z), we can discretize k(x) on I-D grid to attain the key vectors k(1)...., ((N) E ER, which form the key matrix K [k(1),... [(N]]T RRN in self-attention as defined in Eqn. 2. Neglecting the symmetry of the kernel, we choose w(x,y) y) = cxp(a(=)\<<v//V> and rewrite Eqn. (7) with À 0 follows:
u(x) (exp(a(a))
E
(8)
In line with the methodology proposed by (Nguyen et al., 2023b), the Monte-Carle discretization of the above expression with respect to x, y Ω yields
ΣE ( R(IND))()
u ii
(9)
ΣN exp( (OP(())))))
for which the following iterative solver a natural choice:
ΣN (OP()()))))/ VeeN,
u + (1)
EN exp(p(p)) (
n (i)
=((),
where é is an iteration step. It can be seen that setting À = 0 and u u'(i) = v'(), one iteration step becomes
ΣN (xp(a())())//)))
softmax ((i)T((()
" which is equivalent the self-attention computation given by Eqn. (1).
1. ('())
(10)
Σ (xp(((())((()/)))))))))))))))))(
VD
4

---

Published as a conference paper at ICLR 2025
(*(x
(F(0)
677(0)
0.5
0.4
0.4
0.4
0.4
0.2
02-
0.
Figure 3: Dynamics of p" (2)) xx" and (2) (2x xx2 ) for 1,2, 6, 12.
HARNESSING UNATTENDED RESIDUAL INFORMATION VIA TWICING ATTENTION
In this section, we shall associate the representation collapse phenomenon with the rapidly vanish- ing spectrum of NLM similarity matrix during the iterative process, and then propose method to alleviate this issue. Then, we will give deeper theoretical support for the modification method before proposing our Twicing Attention associated with this modified NLM filter.
3.1 VANISHING EIGENVALUES IN ITERATIVE NLM FILTERING
The denoising iteration can be written the matrix vector multiplication
u D D-DD--WW -W''""
(11)
where W is an N x N matrix given by Wij = w(,,j), and D is a diagonal matrix with D N W ig- Introducing the averaging operator A D- the denoising iteration Eqn. 11 becomes conjugate the positive definite matrix D DDD 1/2 This implies that A has complete set of right eigenvectors ( and positive eigenvalues A1 X 2 2 AN O. The largest eigenvalue is λ 1, corresponding to the trivial all-ones right eigenvector (A1 = 1). We expand the signal vector u the eigenbasis as
UA Au. The matrix
D DDD 1/2 via A
Σ CjEy, where ej (E,, u). Applying one step of NLM gives
Au Σ
G AE
Σ 3-1
NgosEt.
Iteratively applying NLM n times, however, yields
A u Σ
(12)
by the same argument. Eqn. 12 reveals that denoising accomplished by projecting the image onto the basis (E and attenuating the contributions of the eigenvectors associated with smaller cigenvalues. Observing that in Eqn. 12, the dynamics of eigenvalues are represented as
N x" p"(A)u= Σ P'))))
(13)
where p(A) X, an identity polynomial whose iterations exhibit steep inclines near À and declines sharply towards zero elsewhere. This results in the iterations converging rapidly toward constant degenerate solution loosing salient information in the input.
3.2 LEVERAGING QUADRATIC KERNEL TOWARDS BETTER INFORMATION CAPACITY
In this section, we revisit the eigenvector expansion of the matrix A as indicated in Eqn. 12. Al- though, in theory, figh-frequency noise is effectively captured by the eigenvectors corresponding to the smallest cigenvalues, in practice, the iterative denoising process can also suppress the contribu- tions of eigenvectors with larger eigenvalues, leading to potential information loss.
5

---

leading to Published as a conference paper at ICLR 2025
To address this issue, we work out alternative polynomial dynamics Pr. which aims to: (eigen- value enhancement) maximally enhance eigenvalues such that Pr. (A) 2 Pa (A) for all A [0. 1], and (0-1 boundedness) ensure that values remain within the range [0. 1] to prevent any eigenvalue from exploding limits as n - oo, therehy ensuring -oo < O s lim Pr. (A) S < co. Owing to the computational overhead associated with higher-degree polynomials, we limit our focus quadratic polynomials. By setting p(0) O, general form of such polynomials is given by p()) = aA + 6X².
Conditions (eigenvalue enhancement) and (0-1 boundedness) imply 1 ((1) 2 p( b= 1 - a. This constraint reformulates p as:
(()) ad + (1-a)1²
(14)
Basic analysis reveals that ş attains its maximum at Xa 2(a-1) which is feasible for all & (0, 2). To satisfy condition (0-1 boundedness), we determine by solving {(aa = 1, yielding a unique solution of 2. This confirms that the optimal quadratic polynomial fulfilling both conditions (eigenvalue enhancement) and (0-1 boundedness) is ((A) 2A 22. Figure 3 illustrates that (") and (2) perform similarly in discarding small eigenvalues near O, essential for effective noise removal. However, 8n (A) remains significantly larger 1, thus better retaining the salient information captured by the input. This observation suggests using 2A- A as a candidate similarity matrix for smoothing the given input without drastically loosing mid mid-ranged eigenvalues and, thus, being capable of capturing more salient information.
3.3 WHY 2A A HELPS: THEORETICAL GROUNDING
Now we shall attempt to provide deeper theoretical insights on the benefits of employing 2A A as a step denoiser, or a similarity matrix in general. First, we demonstrate in Proposition 1 that it achieves substantially slower decay rate of representational capacity the long run. The connection to twicing kernels, from which the paper title originates, is established and Proposition 2 is presented to demonstrate how these kernels effectively reduce estimation bias in nonparametric regression, another smoothing procedure associated with self-attention.
Mitigating representation collapse. To rigorously analyze the differences in denoising dynamics between the kemels p(A)= A and p(A) 2A A A², we define eigencapacity, which correlates with the model's information representation capacity, in Definition 1. Then, we demonstrate in Proposi- tion that the eigencapacity of the former kernel decays at significantly faster rate compared to that of the latter.
Definition 1 (Eigencapacity). Let E C[O,1] and p(A) represent the filter kernel applied during the th denoising step, as specified by Eqn. 13. The eigencapacity of this step, denoted by Ko,, (p), is
defined by the integral
Kon kn) == So (²())do.
(15)
Note that Kn, (p), which represents the area under the curve p"(x) over the interval æ e [0, 1]. exhibits strong correlation with (the sum of) the well-preserved magnitudes of the eigenvalues of p"(A) at iteration step This correlation arises because the integral of p²(x) over this range provides an effective approximation of this sum, particularly for matrices of considerable size since mean value theorem for definite integrals implies that
1 N
S/'' (''()))(s)der = p(c)
E nh (A.)
for some c E [0, 1]. where is a PDF of eigenvalue distribution. This observation underscores the integral's utility in approximating eigenvalue-related characteristics of the filter dyncamics repre- sented by (A).. In the following Proposition 1, we show that the eigencapacity of 2A - 222 decays at significantly slower rate than A.
Proposition (Representational capacity decay rates). Consider denoising process employing the filter kernels p(A)= = A and p( 2A 2 The eigencapacity kn( decays at a rate of Onn 1/2), in contrast to (n² for kn(p), Specifically, the behavior of these eigencapacities as

---

Published as a conference paper at ICLR 2025
oo is given by:
Kon (p) -
(16)
VT
Kon (p)
(17)
2
Remark 1. Due to the equivalence that has been established between NLM smoothing and self- attention computation in Section 2.3, Proposition demonstrates that 2f A was used as a similarity matrix in self-attention mechanism, the output would correspond to a nonlocal smooth- ing operation for which the convergence to a degenerate solution is significantly slower and, thus, capable of maintaining representational capacity for more iterations.
We refer the reader to Appendix A.1 for the proof of Proposition 1.
Relation twicing kernels in nonparametric regression. Equivalence between standard self- attention computation and Nadaraya- estimator using isotropic Gaussian kernels in nonpara- metric regression has been established and used in numerous recent works (Nguyen al., 2022c; Han et al., 2023; Nielsen et al., 2025). In particular, it has been shown that the output of a self- attention block is a discrete form of convolution of Gaussian kernel with bandwidth vD and the value function (detailed in Appendix A.4). We reinterpret attention computation as 2A A rather than A in the nonparametric regression setting. If multiplying by the attention matrix A is equivalent to using some kernel K for NW estimation, then using A 2 is equivalent to applying the convolved kernel K K instead of K (see Appendix A.5).
Therefore, while standard self-attention computation implicitly performs Nadaraya- Watton estima- tor by employing the kernel K, attention computation with 2A AA is equivalent to employing the modified kernel 2K K K, which is exactly the same as applying the kernel twicing procedure to the original regression kernel (Stuetzle & Mittal, 1979; Newey al., 2004; Abdous, 1995). This constructs higher order kernels with small bias property (SBP) which refers to a kernel's 's ability to reduce the leading-order term in the bias of the estimator as demonstrated in Proposition 2 below. Proposition 2 (Twieing kernels reduce the estimator bias). Let K (u) be a symmetric kernel function used in the Nadaraya- Watson estimator with bandwidth h. Define new kernel k(u) as
k(u)= = 2K (w) - (K ())),,
where (K ()) denotes the convolution of K(u) with itself. Then, the kernel K(u) yields Nadaraya- Watson estimator with a smaller bias than that using K ((u).
Remark 2. Due to the relation that has been established between 2A - A and 2K - K K, Proposition 2 implies that 2f - A2 was used as a similarity matrix in self-attention mechanism, the output would correspond to a Nadaraya- Watson estimator with lower bias and arguably less sensitive to bandwidth selection (Newey et al., 2004). This reduced sensitivity mitigates the bias fluctuations often introduced by slight adjustments, making the attention mechanism inherently more resilient to adversarial perturbations and improve model's 's robustness (Chernozhukov et al., 2022).
The proof of Proposition 2 is provided in Appendix A.2. For comprehensive statistical discussion on the topic, we direct the reader to (Newey et al.. 2004) and the references therein.
Twicing kernels benefit from residuals. Recall that we have established connection between self- attention matrices A and 2A A to regression kernels K and 2K - K * K, respectively. Therefore, we use smoothing and computing the attention output interchangeably. Now we provide a core constructive difference between the two kernel computations. Given a kernel-type smoother A and observations V'(r) at iteration f, twicing procedure takes the following three steps:
1. Smooth V'(2) and obtain AV
2. Smooth the residual Vé(x) (V and obtain the correction A (VV==))AV'(
(A (x).
3. Combine Step 1 and Step 2 and define (2A A2)V((z>) the new estimator.
Note that the final estimator actually consists of two terms: the first term corresponds to the denoised image via the filter A, and the second term is the residual - A which is also smoothed with A.
7

---

Published as a conference paper at ICLR 2025
Table 1: Top-1 and Top-5 Test Accuracy on ImageNet corrupted by projected gradient descent (PGD), fast gradient sign method (FGSM), and simultaneous perturbation stochastic approximation (SPSA).
ImageNet
PGD
FGSM
SPSA
Model
Top
Top 5
Top
Top 5
Top
Top 5
Top
Top 5
DeiT (Touvron al., 2021)
72.00
91.14
8.16
22.37
29.88
63.26
66.41
90.29
NeuTRENO (Nguyen et al., 2023b)
72.44
91.39
8.85
23.83
31.43
65.96
66.98
90.48
DeiT-Twicing [10-12]
72.31
91.24
8.66
22.58
31.63
64.74
66.47
90.49
DeiT-TTweing
72.60
91.33
9.15
24.10
32.28
65.67
67.12
90.53
FAN (Zhou et al., 2022) FAN-Twicing
77.09
93.72
11.91
24.11
33.81
65,25
67.15
92.14
77.18
94.02
12.80
28.86
35.52
67.23
68.89
93.75
Table Evaluation of the performance of DeiT and DeiT-Twicing in ImageNet classification under the presence of different corruptions, using appropriate evaluation metries for each.
Dataset Metric
IetgeNet-R Top
ImageNct-A Top
ImageNet-C mCE (1)
ImageNet-C (Extra) mCE (J)
DeiT (Touvron al., 2021)
32.22
6.97
72.21
63.68
Deif-Twwiing [10-12]
32.31
8.14
70.25
62.63
DeiT-Twwcing
32.74
7.66
70.33
62,46
FAN (Zhou al., 2022)
42.24
12.33
60.71
52.70
FAN-Twicing
42.36
12.30
60.48
52.21
Therefore, denoising with kernel p(A) is equivalent denoising with kernel p(A) and subsequently feeding the smoothed method noise of this denoising step back into the output of the current iteration to effectively extracts salient information remaining in the residual and reincorporates it into the denoising output.
3.4 TWICING ATTENTION: FULL TECHNICAL FORMULATION
Stemming from the theoretical benefits discussed in the previous sections, we formulate Twicing Attention as follows:
Definition 2 (Twicing Attention). Given query Qf = [g [gf....,.kI] I RND, key Kt Ikf...., NxD, and value V = [ff...., "y é RND matrices as in Section 2.1 at 8 layer of transformer, the output of Twicing Attention mechanism is computed as:
Uº (2A - A
(18)
where == softmax (Q'K*T) IVD) and the softmax function is applied row-wise.
Remark Even though Definition? gives Twicing Attention computation as Eqn. 18, we use the following equivalent, twicing procedure-inspien form in practice:
U²= ===AV AV= + -VV")
(19)
In other words, instead of computing the square of attention matrix A we decompose Eqn. 18 into regular self-attention output and smoothed residual parts AV + A(Vé -VV).. This allows us to compute AV* once and reuse it in the residual to replace attention squaring operation, which is O(N',, with cheaper matrix multiplication of O N2 N)D) runtime complexity matching the standard self-attenticn computation.
4 EXPERIMENTAL RESULTS
In this section, we empirically justify the advantage of Twicing Attention over baseline transformers with standard self-attention mechanism. Whenever employ Twicing Attention replace the standard one in a given model, we append a Twicing suffix to imply this in the reports. Moreover, if Twicing Attention is inserted in specific transformer layers only, specify the layer indices in square brackets ([10-12] for Twicing Attention in layers 10, 11 and 12, etc.). We evaluate our method
8

---

Published as a conference paper at ICLR 2025
Table 3: Image segmentation on ADE20K.
Table 4: Test PPL on WikiText-103.
Model
Pix. Acc.
Mean Ace.
Mean IoU
Model
Test PPI.
Attacked PPI.
DeiT DeiT-Twieing 77.51
77.25
44,48 45.53
34.73
Transformer
37,51 36.69
55.17 54.46
35.12
Tr. Tr.-Twicing
on Wikitext-103 modeling both under clean and Word Swap contamination (Merity et al., 2016), 6), and ImageNet- IK classification under wide range of attacks (Deng al., 2009; Russakovsky et al., 2015) as described in detail in the following paragraphs.
4.1 IMAGE CLASSIFICATION AND SEGMENTATION
Object classification on ImageNet-1K. To demonstrate the advantage of our model, compare with the DeiT baseline (Touvron et al., 2021) and NeuTRENO (Nguyen et al., 2023b) on the Im- a image classification task (Deng et al., 2009). Our model surpasses the DeiT' baseline, as shown in Table 1 in the clean data setting as well as under adversarial attacks such as fast gradient sign method (FGSM) (Goodfellow et al., 2014), projected gradient descent (PGD) (Madry et al., 2017) with perturbation budget 4/255 and provide comparison of results for different values of perturbations in appendix, and simultaneous perturbation stochastic approximation (SPSA) (Uesato et al., 2018) () with perturbation budget 1/255. Furthermore, Table 2 shows Twicing to be con- sistently more robust than the DeiT baseline across various testing conditions, including adversarial examples and cut-of-distributie which involves common data corruptions and perturbations such as noise addition and image blur- ring, well as on Imagenet-A and Imagenet-R R datasets, which assess adversarial example handling and outto--distribution generalization, respectively (Hendrycks et al., 2021). ImageNet-C (Extra) contains four extra image corruption types: spatter, gaussian blur, saturate, and speckle noise.
datasets. This includes its performance on the Imagenet-C dataset,
When combining with a state-of- offte-of-th-arr robust transformer backbone, Fully Attentional Network (FAN) (Zhou et al., 2022), Twicing Attention is able to improve performance in terms of clean accu- racy as well as its robustness against adversarial attacks such as PGD and FGSM (with perturbation budget 4/255) as well as SPSA (with perturbation budget 1/255) substantially included in Table 1. We also find better out-of-disrrbuution generalization in FAN-Twicing over standard FAN except for ImageNet-A benchmark where FAN-Twicing is still highly competitive (see Table 2).
Image segmentation on ADE20K. On top of the classification task, we compare the performance of the Segmenter models using DeiT and Dwii-T backbones the ADE20K (Zhou al., 2019) image segmentation task to further validate the advantages of our proposed method by adopting the experimental setup of (Strudel et al., 2021). In Table we report the key metrics: pixel accuracy, mean accuracy, and mean intersection over union (IOU). We observe performance boost across all 3 metrics with Twicing over the DeiT baseline (Touvron et al., 2021).
4.2 LANGUAGE MODELING ON WIKITEXT-103
In addition to computer vision tasks, we also evaluate the effectiveness of our model on a large-scale natural language processing application, specifically language modeling on WikiText- 103 (Merity et al., 2016). Our language model demonstrates better performance in terms of both test perplexity (PPL) and valid perplexity when compared to the standard transformer language model (Vaswani et al.. 2017) as shown in Table 4. We also show that test PPL on WikiTeex- 103 contaminated by Word Swap Attack, where words are randomly swapped with a generic token AAA', We follow the setup of (Han et al., 2023; Teo & Nguyen, 2025) and assess models by training them on clean data before attacking only the test data using an attack rate of 4%.
5 EMPIRICAL ANALYSIS
Representation collapse analysis. We empirically demonstrate in Figure that Twicing Attention mechanism promotes better token diversity and, therefore, it is able to slow down representation collapse phenomenon in transformers. We observe that, in DeiT, the average cosine similarity score between tokens quickly exceeds three-quarters, whereas in model, it remains consistently below this threshold. Additionally, we demonstrate in Figure 2 that our model is indeed capable of retaining better expressive power by being able to pay attention to notably more and important parts of objects in images while DeiT indicates collapsed behaviour. See Appendix D.2 for a dozen of additional attention heatmaps supporting the comparative argument.

---

Published as a conference paper at ICLR 2025
Table 5: Efficiency comparison between DeiT and Dei'--Twiiin models.
Model
Avg. Compute Speed (ms/it)
GFLOPs sample
Param. Count (M)
DeiT (Touvron et al., 2021)
8,58
1.25
5.7
Dei'--Twicing
9.14
1.33
5.7
DeiT-Twicing [10-12]
8.72
1.27
5.7
Efficiency analysis. As stated in Remark 3, our Twicing Attention mechanism can be implemented with O( runtime complexity, which is on par with the standard self-attention mechanism. Table 5 compares the prediction average compute speed per iteration over 1000 runs as well as the gooating-poitt operations (FLOPs) per sample, which is a measure of trade-off between efficiency (lower FLOPs) and accuracy (higher FLOPs) of models. We observe that employing Twicing Atten- tion only in last 3 layers, the model can still enjoy performance gains over the baseline while seeing almost negligible increase in average compute speed and FLOPSs per sample.
6 RELATED WORK
Theoretical frameworks for attention. Attention mechanisms have been studied from a range of perspectives. (Tsai et al., 2019) shows that attention can be derived from kernel similarity functions. (Nguyen et al., 2023d; Tarzanagh et al., 2023) relate attention to support vector regression/maahine, while (Tao et al., 2023) explains attention through nonlinear singular value decomposition of asym- metric kernels. Furthermore, (Teo & Nguyen, 2025) establishes a connection between self-attention and kernel principal component analysis, demonstrating that self-attention maps query vectors the principal component axes of the key matrix in a feature space. Attention has also been explained through Gaussian mixture models, ordinary/partial differential equations, optimization algorithms, and graph-structured learning (Nguyen et al., 2022a;b; 2023c; Tang & Matteson, 2021; Gabbur et al., 2021; Lu et al., 2019; Sander al., 2022; Nguyen et al., 2022d; Kreuzer al., 2021: Zhang & Feng, 2021) or an energy functional minimization associated with a variational image denoising frame- work (Nguyen et al., 2023b). (Nguyen et al., 2022c; Han et al., 2023; Nielsen et al., 2025) show that self-attention performs Nadaraya-Watte regression with Gaussian isotropic kernels.
Representation collapse in transformers. Representation collapse or over-smoothing in transform- has been observed across various domains, including NLP (Shi et al., 2022) and computer vision (Wang et al., 2022). (Shi et al., 2022) analyzes this issue in BERT (Devlin et al., 2018) using a graph- based approach, employing hierarchical fusion techniques to retain self-attention outputs, though a high memory cost. (Dong et al., 2021) is among the first to explore oversmoothing in transformers through rank collapse, while (Caron al., 2021) examines self-supervised ViTs, revealing their em- bedded semantic segmentation information. Additionally, (Darcet et al., 2024) identifies high-norm token artifacts in these feature maps. Furthermore, (Nguyen al., 2024a) explores the link between self-attention and the state-space model, showing that attention output collapses into steady-state solution.
Robust transformers. For transformers, robust strategies include an ensemble defense against ad- versarial attacks (Mahmood et al., 2021), position-aware attention scaling with patch- wise aug- mentation (Mao et al., 2022), and a fully-attentional network for of-te-of-the-art performance on corrupted images (Zhou et al., 2022). Efficiency usually refers to optimal performance under ideal conditions, while robustness describes maintaining strong performance under less-than-ideal cir- cumstances. A common trend among robust models, such as (Mao et al., 2022; Han et al., 2023; Zhou et al., 2022), is their reliance additional computational overhead, often matching exceeding that of our proposed model.
CONCLUDING REMARKS
In this paper, we introduced the Twicing Attention mechanism, enhancing the transformer's rep- resentational capacity by utilizing residuals between self-attention inputs and outputs. This novel self-attention variant improves token diversity and mitigates representational collapse by leveraging useful residual information as a form of self-correction. We empirically demonstrated performance gains on ImageNet- 1k, ADE20K, tational overhead by trying selective layer placement for Twicing Attention. However, limitations include its efficient application across transformer all layers with negligible additional compu- tation. Ongoing work explores approximation techniques and sparsity to improve efficiency, while extending the theoretical framework to even more practical scenarios remains an open challenge.
MikiText-103, and robustness benchmarks with minimal compu-
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
This research project is supported by the National Research Foundation Singapore under the AI Singapore Programme (AISG Award No: AISGG-TC-2203--122--010
This research project is supported by the Ministry of Education, Singapore, under the Academic Research Fund Tier This research / project is also supported by the NUS Presidential Young Professorship Award (A-0009807-01-00) and the NUS Artificial Intelli- gence Institute-Seed Funding (A-8003062-00-00),
(FY2023) (A-8002040-00-00) A 1-8002039-00-000
Reproducibility Statement. We have made efforts to ensure the reproducibility of our work through several measures. Source codes for our experiments are provided in the supplementary materials of the paper. The details of our experimental settings and computational infrastructure are given in Section 4 and the Appendix. All datasets that we used in the paper are published, and they are easy to find in the Internet. These resources and explanations should allow others to replicate our results with relative ease.
Ethics Statement. Given the nature our work and contributions, we do not foresce any negative societal and ethical impacts of our work.
REFERENCES
Belaid Abdous. Computationally efficient classes of higher-order kernel functions. The Canadian
Journal of Statistics / La Revue Canadienne de Statistique, 23(1):21-27, 1995. doi: 10.2307/
3315548.
Rami AI-Rfou, Dokook Choe, Noah Constant, Mandy Guo, and Llion Jones. Character-level lan- guage modeling with deeper self-attention. In Proceedings of the AAAI conference on artificial intelligence, volume 33, pp. 3159-3166, 2019.
Alexei Baevski and Michael Auli. Adaptive input representations for neural language modeling. arXiv preprint arXiv: /0853, 2018
A. Buades, B. Coll, and J.-M. Morel. non-local algorithm for image denoising. In 2005 IEEE Computer Society Conference on Computer Vision and Pattern Recognition (CVPR 05), vol- ume 2. pp. 60-65 vol. 2, 2005, doi: 10.1109/CVPR.2005.33
Mathilde Caron, Hugo Touvron, Ishan Misra, Herv'e J' J'egou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in self-supervised vision transformers. 2021 IEEE/CVF International Conference on Computer Vision (ICCV), pp. 9630-9640, 2021.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 34: 15041-1509,97 2021.
Victor Cheroozuukov, Juan Carlos Escanciano, Hidehiko Ichimura, Whitney K. Newey, and James M. Robins. Locally robust semiparametric estimation. Econometrica, 9014)1150 July 2022. doi: 10.3982/ECTA16299 URL /tttps:://oi..ooo 111.:/1...9922////.// 6294.
Kyunghyun Cho, Bart Van Merriënboer, Caglar Gulcehre, Dzmitry Bahdanau, Fethi Bougares, Hol- ger Schwenk, and Yoshua Bengio. Learning phrase representations using rnn encoder-decoder for statistical machine translation. arXiv preprint arXiv: 1406. 1078, 2014.
Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc V Le, and Ruslan Salakhutdi- nov. Transformer-xl: Attentive language models beyond fixed-length context. arXiv preprint arXiv:1901.088 02860, 2019.
Timothée Darcet, Maxime Oquab, Julien Mairal, and Piotr Bojanowski. Vision transformers need registers. In The Twelfth International Conference on Learning Representations, 2024. URL. https://openreve/
net wwtp:/eeeee/tendeeeeeeeeee 2ttm:mm2---2
Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Universal transformers. arXiv preprint arXiv: 2018.

---

Published as a conference paper at ICLR 2025
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A scage-scale hi- erarchical image database. In 2009 JEEE conference on computer vision and pattern recognition, pp. 248-255. Icee, 2009.
Jacob Devlin, Ming-Wer Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint 1110::v810.048000 2018.
Yihe Dong, Jean- SannBBapisste Cordonnier, and Andreas Loukas. Attention is not all you need: pure attention loses rank doubly exponentially with depth. Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference Machine Learning, volume 139 of Proceed- ings of Muchine Learning Rexearch, pp. 2793-2803, PMLR, 18-24 Jul 2021. URL https:
//proceedings mll.. wallrpressse / dong2la html.
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23(120):1339
2022.
Prasad Gabbur, Manjot Bilkhu, and Javier Movellan. Probabilistic attention for interactive segmen- tation. Advances in Neural Information Processing Systems, 34:4448 4460, 2021.
Guy Gilboa and Stanley Osher. Nonlocal linear image regularization and supervised segmentation. Multiscale Modeling Simulation, 6(2):595-630, 2007. doi: 10.1137/000669358.
lan J Goodfellow, Jonathon Shlens, and Christian Szegedy. Explaining and harnessing adversarial examples. arXiv preprint arXiv: 14iiv::412.66 42.65722 2014.
Xing Han, Tongzheng Ren, Tan Nguyen, Khai Nguyen, Joydeep Ghosh, and Nhat Ho. Design- ing robust transformers using robust kernel density estimation. Advances in Neural Information Processing Systems, 36:53362-53384, 2023.
Dan Hendrycks, Kevin Zhao, Steven Basart, Jacob Steinhardt, and Dawn Song. Natural adversarial examples. In Proceedings of the IEEE/CVF conference computer vision and pattern recogni- tion, pp. 15262- 15271, 2021.
Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. Advances in neural information processing systems, 34:1273-1286, 2021.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are mnss: Fast autoregressive transformers with linear attention. In International conference on ma- chine learning, pp. 5156-5165, PMLR, 2020.
Salman Khan, Muzammal Naseer, Munawar Hayat, Syed Waqas Zamir, Fahad Shahbaz Khan, and Mubarak Shah. Transformers in vision: A survey. ACM computing surveys (CSUR), 54(10s):
2022.
2022.
Devin Kreuzer, Dominique Beaini, Will Hamilton, Vincent Létourneau, and Prudencio Tossou. Re- thinking graph transformers with spectral attention. Advances in Neural Information Processing Systems, 14-21618-21621 2021.
Tianyang Lin, Yuxin Wang, Xiangyang Liu, and Xipeng Qiu. A survey of transformers. Al open, 3: 111-132, 2022.
Zhouhan Lin, Minwei Feng, Cicero Nogueira dos Santos, Mo Yu, Bing Xiang, Bowen Zhou, and Yoshua Bengio. A structured self-attentive sentence embedding. arXiv preprint 1703.07 2017.
Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF international conference computer vision, pp. 10012-10022, 2021.
Yiping Lu, Zhuohan Li, Di He, Zhiqing Sun, Bin Dong, Tao Qin, Liwei Wang, and Tie-Yan Liu. Understanding and improving transformer from a multi-partich. dynamic system point of view. 2019.
arXiv preprint arXiv: /
12

---

Published as a conference paper at ICLR 2025
Aleksander Madry, Aleksandar Makelov, Ludwig Schmidt, Dimitris Tsipras, and Adrian Vladu. Towards deep learning models resistant to adversarial attacks. arXiv preprint rrXiv:1706.066 2017.
Kaleel Mahmood, Rigel Mahmood, and Marten Van Dijk. On the robustness of vision transformers to adversarial examples. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 7838-7847, 2021.
Xiaofeng Mao, Gege Qi, Yuefeng Chen, Xiaodan Li, Ranjie Duan, Shaokai Ye, Yuan He, and Hui Xue. Towards robust vision transformer. In Proceedings of the IEEE/CVF conference on Com- puter Vision and Pattern Recognition, pp. 12042-12051, 2022.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. arXiv preprint arXiv: 1609000800 2016.
Whitney K. Newey, Fushing Hsich, James M. Robins. Twicing kernels and a small bias property of semiparametric estimators. Econometrica, 72(3):47--62,, 2004. ISSN 001 29682, 14680262. URL hwtp::ttp////. //tp::://wwwww// wwwtt..org/staabeeeeeeeee
Khang Nguyen, Nong Minh Hieu, Vinh Duc Nguyen, Nhat Ho, Stanley Osher, and Tan Minh Nguyen. Revisiting over-smoothing and over-squashing using ollivier-ricci curvature. In In- lernational Conference on Machine Learning, pp. 25956-25979, PMLR, 2023a.
Tam Nguyen, Tan Nguyen, and Richard Baraniuk. Mitigating over-smoothing in transformers via regularized nonlocal functionals, Advances in Neural Information Processing Systems, 36:80233- 80256, 2023b.
Tam Minh Nguyen, Tan Minh Nguyen, Dung DD Le, Duy Khuong Nguyen, Viet-Anh Tran, Richard Baraniuk, Nhat Ho, and Stanley Osher. Improving transformers with probabilistic attention keys. In International Conference on Machine Learning, pp. 16595-16621. PMLR, 2022a.
Tam Minh Nguyen. César A Uribe, Tan Minh Nguyen, and Richard Baraniuk. Pidformer: Trans- former meets control theory. In Forty-first International Conference on Machine Learning, 2024a.
Tan Nguyen, Vai Suliafu, Stanley Osher, Long Chen, and Bao Wang. Fmmformer: Efficient and flexible transformer via decomposed near-field and far- field attention. Advances in neural infor- mation processing systems, 34:29449 -94633, 2021.
Tan Nguyen, Tam Nguyen, Hai Do Khai Nguyen, Vishwanath Saragadam, Minh Pham, Khuong Duy Nguyen, Nhat Ho, and Stanley Osher. Improving transformer with an admixture of attention heads. Advances in neural information processing systems, 35:27937-27952 2022b.
Tan Nguyen, Minh Pham, Tam Nguyen, Khai Nguyen, Stanley Osher, and Nhat Ho. Fourierformer: Transformer meets generalized fourier integral theorem. Advances in Neural Information Pro- cessing Systems, 35:29319 29335,, 2022c
Tan M Nguyen, Tam Nguyen. Long Bui, Hai Do, Duy Khuong Nguyen, Dung D Le, Hung Tran- The, Nhat Ho, Stan J Osher, and Richard G Baraniuk. A probabilistic framework for pruning transformers via a finite admixture of keys. ICASSP 2023-2023 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 1-5. IEEE, 2023c.
Tan Minh Nguyen, Richard Baraniuk, Robert Kirby, Stanley Osher, and Bao Wang. Momentum transformer: Closing the performance gap between self-attention and its linearization. In Mathe- matical and Scientific Machine Learning, pp. 189-204. PMLR, 2022d.
Tan Minh Nguyen, Tam Minh Nguyen, Nhat Ho, Andrea L. Bertozzi, Richard Baraniuk, and Stanley Osher. A primal- framework for transformers and neural networks. In The Eleventh Interna- tional Conference Learning Representations, 2023d. URL pttttps://pperevebee .net/ forum?id=U T8- 5hC1V.
Tuan Nguyen, Hirotada Honda, Takashi Sano, Vinh Nguyen, Shugo Nakamura, and Tan Minh Nguyen. From coupled oscillators to graph neural networks: Reducing over-smoothing via ku- ramoto model-based approach. In International Conference on Artificial Intelligence and Statis- fics, pp. 2710-2718. PMLR, 2024b.
13

---

Published as a conference paper at ICLR 2025
Stefan Nielsen, Laziz Abdullaev, Rachel SY Teo, and Tan Nguyen. Elliptical attention. Advances in Neural Information Processing Systems, 37:109748-109777 2025.
Ankur P Parikh, Oscar Täckström, Dipanjan Das, and Jakob Uszkoreit, A decomposable attention model for natural language inference. arXiv preprint 1606.0660..1933 2016.
Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, Improving language under- standing by generative pre-training, 2018.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Dly Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, al. Learning transferable visual models from natural language supervision. In International conference on machine learning, pp- 8748-8763. PMLR, 2021.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research, 21 (1(10)):165 2020.
Olga Russakovsky, Jia Deng. Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang. Andrej Karpathy, Aditya Khosla, Michael Bernstein, al. Imagenet large scale visual recognition challenge. International journal of computer vision, 15:2111222, 2015.
Michael E Sander, Pierre Ablin, Mathieu Blondel, and Gabriel Peyré. Sinkformers: Transform- ers with doubly stochastic attention. In International Conference on Artificial Intelligence and Statistics, pp. 3515-3525 35300. PMLR, 2022.
Han Shi, JIAHUI GAO, Hang Xu, Xiaodan Liang. Zhenguo Li, Lingpeng Kong, Stephen M. S. Lee, and James Kwok. Revisiting over-smoothing BERT from the perspective of graph. In International Conference Learning Representations, 2022. URL https: /ttp:////eeeeeee....e. net/forum: ddddUV99uuaam
Robin Strudel, Ricardo Garcia, Ivan Laptev, and Cordelia Schmid. Segmenter: Transformer for semantic segmentation. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 7262-7272, 2021.
W. Stuetzle and Y. Mittal. Some comments on the asymptotic behavior of robust smoothers. In T. Gasser and M. Rosenblatt (eds.), Smoothing Techniques for Curve Estimation, volume 757 of Lecture Notes, pp. 191-195. Springer- Verlag, New York, 1979.
Binh Tang and David S Matteson. Probabilistic transformer for time series analysis. Advances in Neural Information Processing Systems, 14.23592-23608, 2021.
Qinghua Tao, Francesco Tonin, Panagiotis Patrinos, and Johan AK Suykens. Nonlinear svd with asymmetric kernels: feature learning and asymmetric nystr)" om method. arXiv preprint arXiv:2306.077 07040, 2023.
Davoud Ataee Tarzanagh, Yingcong Li, Christos Thrampoulidis, and Samet Oymak. Transformers as support vector machines. In NeurlPS 2023 Workshop on Mathematics of Modern Machine Learning. 2023. URL https: / /ttp///ppeeeeeeeeeeeeeeeeeeeeeee .nt/ www.:oodduuuwwaaaaaaao 9K.
Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey. ACM Computing Surveys, (5(6):1-28) 2022.
Ian Tenney, Dipanjan Das, and Ellie Pavlick. Bert rediscovers the classical nlp pipeline. arXiv preprint wwx.vv1905.05990 2019.
Rachel SY Teo and Tan Nguyen. Unveiling the hidden structure of self-attention via kernel principal component analysis. Advances in Neural Information Processing Systems, 17:101.93-101127, 2025.
14

---

Published as a conference paper at ICLR 2025
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. In International conference machine learning, pp. 10347-10357. PMLR, 2021.
Hoang V. Tran, Thieu Vo, An Nguyen The, Tho Tran Huu, Minh-Khoi Nguyen-Nhat, Thanh Tran, Duy-Tung Pham, and Tan Minh Nguyen. Equivariant neural functional networks for transformers. In The Thirteenth International Conference on Learning Representations, 2025. URL https:
//wpnneevee..eeee wwt.:/rumtttttmemooooooo wtttt//ooommiiiiieeeee a OukstY.
Yao-Hung Hubert Tsai, Shaojie Bai, Makoto Yamada, Louis-Philippe Morency, and Ruslan Salakhutdinov. Transformer dissection: unified understanding of transformer's attention via the lens of kernel. arXiv preprint arXiv:1908. 11775, 2019.
Jonathan Uesato, Brendan O' d'donoguue, Pushmeet Kohli, and Aaron Oord. Adversarial risk and the dangers of evaluating against weak attacks. In International Conference on Machine Learning, pp. 5025-5034. PMLR, 2018.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Eukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural informa- tion processing systems, 30, 2017.
Jesse Vig and Yonatan Belinkov. Analyzing the structure of attention in a transformer language model. arXiv preprint 1rXi::1906.044 2019.
M. P. Wand and M. C. Jones. Kernel Smoothing. CRC Press, Chapman and Hall, 1995.
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 Reprexentations, 2022. URL TtLhttps: /wwp//evieeeeeeeeee//ee..e net/ forum? id=0476 60/mm NNp.
Xinyi Wu, Amir Ajorlou, Zihui Wu, and Ali Jadbabaie. Demystifying oversmoothing in attention- based graph neural networks. seirth- Conference on Neural Information Processing Systems, 2023. URL htttps://ppereete
net forum? d= 65q ei.B.
Shaolei Zhang and Yang Feng. Modeling concentrated crttttttioonn for neural machine translation with gaussian mixture model. arXiv preprint arXiv:2/09.05244, 2021.
Bolei Zhou, Hang Zhao, Xavier Puig, Tete Xiao, Sanja Fidler, Adela Barriuso, and Antonio Torralba. Semantic understanding of scenes through the ade20k dataset. International Journal of Computer Vision, 127:302-321, 2019.
Daquan Zhou, Zhiding Yu, Enze Xie, Chaowei Xiao, Animashree Anandkumar, Jiashi Feng, and Jose M Alvarez. Understanding the robustness vision transformers. In International Conference PMLR, 2022.
on Machine Learning, pp.
15

---

Published as a conference paper at ICLR 2025
Supplement to "Transformer Meets Twicing: Harnessing Unattended Residual Information"
Table of Contents
A Technical Proofs and Derivations
16
A.I Proof of Proposition
16
A.2 Proof of Proposition? 2
17
A.3 Derivation of Gradient of Ju
18
A.4 Equivalence of Self-attention and Nadaraya- Watson Estimator
19
A.5 Equivalence between Sclf-convolution and Square of Attention Matrix
20
B Experimental Details and Additional Results
20
B.1 Wikitext-103 Language Modelling
20
B.2 ImageNet Image Classification and Adversarial Attack
22
B.3 Out-of-Distribution Robustness and Data Corruption on ImageNet-A,R,C....
B.4 ADE20K Image Segmentation
23
Compute Resources
23
Additional Empirical Analysis
23
D.1 Extra Over-smoothing Analysis
23
D.2 Extra Attention Heatmap Analysis
23
E Ablation Studies
24
A TECHNICAL PROOFS AND DERIVATIONS
A. PROOF OF PROPOSITION
The equivalence in Eqn. 16 is straightforward to obtain since kx) can be calculated as
(p) S.' (p'(x)dx= S. x=dx= dx
n +
To prove the equivalence given by Eqn. 17, we first observe that
kn)(p) = S. (=)()d f((2 22) = ((² (2x - x² x²)" dx,
where the last equality is due the symmetry of 2x-x² = 1- 1-11-11111 x along 1. Now, employing a variable change = 2y yields
Kin (n) 1 16² (2x -2)-2 [ (4y "
1
4" S: (1- 4)" 4" B(n n l,n+
(20)
4" [rInn+ 1)2 4t (n!)
(21)
T(2n+2) (2n 1)!
where B(x,y) y) and I(x) denote the Euler Beta function and Gamma function, respectively, and used the identity B(x, y) (())T())T((r y) to transform Eqn. 20 into Eqn.21. Now using
16

---

Published as a conference paper at ICLR 2025
Stirling's approximation n! V2Tn(n/e)" for Eqn. 21, obtain
4th 2mm 2e2"
(0)
2m 2n 1)
1)2m-1 /
2n
2n-1
22 +
12n+1
CV
(1- 2n +
Vn+ + 1/2
V
(22)
2/n²
where we used the fact that e-1 10 (1
2n+1
to derive Eqn. 22.
2n=
A.2 PROOF OF PROPOSITION 2
Proof of Proposition 2. To compare the biases of the estimators using kernels K(u) and R(u), we analyze the moments of these kernels, they determine the bias in kernel estimators.
We begin with showing that K has valid kernel propoerties.
Normalization. Since K (u) is a valid kernel, we have:
Se ((u) du 1.
The convolution of K (uu) with itself satisfies:
J (K K)(u) du (5s ((w))du)) =1.
Therefore,
J k(u)du= du f K(u) du [Kx R)(u)) du =
Thus, ((u)) is normalized.
Symmetry, If K (u) is symmetric, i.e., K(u) = K(-u), then (K * (u) also symmetric. There- k--u)= 2K(-u) (K (- 2K( -u)))))) (K* (u)
fore,
(u).
Thus, K(u) is symmetric.
Zero First Moment. The first moment of kernel should be zero:
f (u((u) du f (K(uu) du f uu(K (KK+ ()(u) du.
Since K(u) is symmetric, f (u)(u) du O, and the convolution (K K)(u) also symmetric, so f u(KK K * (1) du = O. Therefore,
fx (u) 0.
This confirms that (u) has a zero first moment.
Next, note that the second moment of a kernel function, #2, influences the leading term in the bias of the kernel estimator. For Ru)). have:
(2( = S u² (2((u) du =2 f u² (() du - f u²(K ()) du.
(23)
17

---

Published as a conference paper at ICLR 2025
We know that f (u) du (2(K). The tenm u²(K K)) du can be evaluated follows:
ʃ ())
S1 12 (J K(t)K(u-o)) dr du
K(V) (Ja [[2KK(u (oK(u-1) - du du
JR
J K(W) (f( (f((+t)((())) ds de
((O)( (ʃ (s) 20 f (K(s)ds+O) f K(s) ds)
Since (ss) symmetric:
f sk()) f ( ʃ {())dd==1.
Thus, the expression simplifies to:
J ((t) + 0 (2) x2()) / K(u)du+ f (² 2())(())))))
Thus,
f (K ())
du (x)) H2(K) 2K)(K).
Finally, returning to (2(K)
(K) in Eqn. 23:
(2(K)= 2 f u² K (u) du f xK K)) du
2uz ()) (K)(K)
Implications for the bias. classical result in statistics imply that the leading bias term of the Nadaraya-Wateee estimator using kernel K (u) is proportional to v2(K)
h Bias mi (x)] 222 (K)m"(x)
where (x) the second derivative of the true regression function point æ (see, for example, (Wand& & Jones, 1995)).
For the estimator using (u), since 222 R) = O, the leading bias term of order h2 disappears. The next non-zero term in the bias expansion involves the fourth moment sa ()), resulting in a bias of order h::
1ea
Bias & (z)] 24 "1(( (r))mm(())))) (m))
This demonstrates that the estimator using (u) reduces leading order bias terms that appear when
K ((u) is used.
A.3 DERIVATION OF GRADIENT OF J..
Expand the functional (u) as follows:
Ju (u) f E
(1(())))(())? w1 ddddy
(24)
The gradient of J., with respect to is then given by:
[J_ aJm
aJ.
(u)
(25)
Ou," Oug" Oup
18

---

Published as a conference paper at ICLR 2025
The partial derivative aJm//uu, for j 1.2....,D. D, is defined through its dot product with an arbitrary function h1 (S2) as follows:
aJw,
d hz(2) Julu Th₂)
du,
dr
1-0
1 1 {
+h)((() wo,, y)
2 dT Janse
==0
+ +h) Thy(w)) (v))² wa y) drdy 2 dr Sann (v))
1
T=0
= Sma (u,(2) - uj (1)) (h) (v)) w( 4) dxdy
Applying a change of variables (x, - (v, x) to the second term of the above integral, we have:
0Jw h2(2)) J (u)(2) u (v)) (x) (w(x v) + w(y dy
Oug
Thus, the Frechet derivative of Jw with respect uj is given by: 0Jw
So (us(tt) uy (v)) (k(2, v) k(y, dy:
(26)
du,
which then gives the desired gradient with wxx v) w(x 2v, (Nguyen al., 2023b).
A.4 EQUIVALENCE OF
AND NADARAYA- SATARAYY-WATSON ESTIMATOR
We establish the relationship between self-attention, as defined in Eqn. 1, and non-parametric re- gression following the approaches of (Nguyen al., 2022e; Han al., 2023; Nielsen et al., 2025).
To begin, let us assume that the key and value vectors {ky, 22 je[N] are generated by the following data process:
f(k) E
(27)
where e represents random noise with zero mean, i.e., E[<] 0, and f is the unknown function we aim estimate. In this setup, the keys [ are independent and identically distributed (i.id.) samples drawn from the marginal distribution p(k), characterizing the random design setting. We use p(u u. k) to denote the joint distribution of the pairs (v, k) generated by the process described in Eqn. 27, For new query q, our goal is to estimate the function f(q).
Recall NW estimator is non- parametric estimator of the unknown f at any given query described by
f(k) = E[ k] - v.v(v p(u k))vu= f v v.p(vuk) dv,
p(k)
where the first equality comes from the noise being zero mean, the second equality comes from the definition of conditional expectation and the final equality comes from the definition of conditional density. Eqn. 27 implies that if we can just obtain good estimates of the joint density p(v, k) and marginal density p(k) then we can estimate the required f()). The Gaussian isotropic kernels with bandwidth are given by
1 N
1
(o k))
Σ (o (o-v-v;)ee v.) k2), Po (k)) = [C[N]
Σ (k) N [c[N]
(28)
where fo is the multivariate Gaussian density function with diagonal covariance matrix o²Ip Given the kernel density estimators in Eqn. 28, the unknown function can be estimated as
(k) = S.P U. (,,( du= ffo V. V-Eje[N] Po v;) (k k;)
du
(k)
EjeNT [ILN)&o(A) Po - kj)
ΣjEN Po k,)) V. p0 (v-v,)de v, = Ej(N] Ujso (k- kk)) EjeNNi k;))
EjeNN 40 -k;))
19

---

Published as a conference paper at ICLR 2025
85
De Tlwwiinnng
Corruption Errer (CE)
75-
70-
63
50
Brigtnees
Contranst
Defocus Blar
Elastic Tansform
Fag
Gaussion Maise
Frost
Glass Blur
impulse Noise
JPEG compression
Nation Blar
Ruslate
Shot Noisee
Snow
Zoom Blar
Compation mCE
Corruption Types
Figure 4: ImageNet-C corruption error (CE) ()) and mean CE (mCE) (L) comparison of our model
and DeiT across all corruption types. Our model consistently outperforms DeïT.
Then, using the definition of the Gaussian isotropic kernel and evaluating the estimated function at q we have
Σ vy exp
- - 13/022)
i(q:)
Σÿ exp (- (-1q:-1) ky 1
Σ By exp [1 + |k; 2=000 exp( q: k/
Σ exp
(-q1; 1 k, 2)/ exp(1 q: /0)
Σ Uy exp(q, kz/
N
Σ softmax k )
Σj exp(q k2)
as desired.
A.5 EQUIVALENCE BETWEEN SELF-CONVOLUTION AND SQUARE OF ATTENTION MATRIX
Let K denote the isotropic Gaussian kernel with bandwidth h. Then,
(K vx)) I ((x-((((E))) - () v)( dt S 1(x--o - 1) J K(t- v) dy dt
( K(x-t)K(( 3) dt v(y) dy f K(x (xxxt))(( - (1))())))))))))))))))))))))) dy
Σ ((==))((--))(
(30)
Taking A to be the NLM matrix whose entries are given by A.= = w((,j) = K(i - j), it becomes evident that Eqn. 30 can be represented as
* K E(A)),
(31)
B EXPERIMENTAL DETAILS AND ADDITIONAL RESULTS
B. WIKITEXT- 103 LANGUAGE MODELLING
Dataset. The WikiText- 103' (WT-103) dataset contains around 268K words and its training set consists of about 28K articles with 103M tokens. This corresponds to text blocks of about 3600 words. The validation set and test sets consist of 60 articles with 218K and 246K tokens respectively. hag.y///aanggggmmmoliinnngnaaaa.
20

---

Published as a conference paper at ICLR 2025
Tog Accuracy
Ep Acturaty
Accuraty
Accuraty
Accurcy (86)
Accurasy (16)
75
30
10-
Perturbation
Perturbation
Perurbattion Budget
Perturbation
Figure 5: Top-1 and Top-5 accuracies on FGSM attack with 6 increasingly different perturbation budgets (x (x255).
Figure 6: Top-1 1 and Top-5 accuracies on PGD attack with 6 increasingly different perturbation budgets (x25)).
Figure 7: Validation PPL (L) training curves for baseline Transformer (higher) and Transformer- Twicing (lower). Left: small models (9.4M); Right: medium models (21M). We observe relatively faster convergence for Twicing Attention compared to standard self-attention.
Table 6: Valid/Test PPL on WT-103.
Corruption. Word Swap Text Attackè corrupts the data by substituting random words with generic token "AAA". We follow the setup of (Han et al., 2023) and assess models by training them on clean data before attacing only the evaluation set using a substitution rate of 4%.
Model
Valid PPL
Test PPI.
Transformer (small)
38.11
37.51
Tr.- Tr.-Twicing (small)
37.12
36.69
Transformer (med)
31.98
26.17
Tr. Ft.-Twicing (med)
30.90
25.65
Model, Optimizer & Train Specification. We
Linear Trans.
40.00
41.26
adopt the training regime of (Nguyen et al., 2022c).
Linear-Twicing
39.45
40.61
To this end, the small backbone uses 16 layers, 8
heads of dimension 16, feedforward layer of size 2048 and embedding dimension of 128. We use a dropout rate of 0.1. We trained with Adam using starting learning rate of 0.00025 and cosine scheduling under default Py PyTorch settings. We used a batch size of 96 and trained for 120 epochs and 2000 warmup steps. The train and evaluation target lengths were set to 256.
Larger Language Modeling. To verify if Twicing Attention scales, we conduct extra language mod- eling on Wikitext-103 with medium sized models (21M parameters) on top of the small models (9.4M parameters) reported in the main text. The results in Figure and Table 6 imply a positive answer to this matter.
Table 7: Fine-tuning Switch Transformer on SST-2.
Attention Type
Test Acc.
#Params
Self-attention Twicing Attention
77.78 78.34
33M
33M
Fine-tuning pre-trained language model. To
Table 8: Attacked Test PPI. on WT-103.
show how Twicing Attention can offer improve- ments to the pre-trained models, pre-train medium sized (33M parameters) Switch Trans-
Model
Clean PPL
Attacked PPL
Linear Trans. 41.26
72.13
former (Fedus et al., 2022), a Mixture of Experts architecture, with the standard self-attention on 1iki Then we finetune this pretrained language model Stanford Sentiment Treebank 2 (SST-2) dataset using standard self-attention (baseline) as well as Twicing Attention (ours) for 8 epochs. Table 7 compares Top 1 fine-tune test accuracies for both cases and we find that fine-tuning with Twicing Attention achieves higher ac- curacy, provided that the fine-tuning is long enough for the model to adapt to the new attention mechanism.
Linear- Twicing
40.61
70.40
"mplemmntation available at github.com/QDDataeeeettttaa
21

---

Published as a conference paper at ICLR 2025
Table 9: Top-1 and Top-5 Test Accuracy on ImageNet corrupted by projected gradient descent (PGD), fast gradient sign method (FGSM), and simultaneous perturbation stochastic approximation (SPSA).
ImageNet
PGD
FGSM
SPSA
Model
Top 1
Top 5
Top 1
Top 5
Top 1
Top 5
Top 1
Top 5
DeiT (Touvron al., 2021)
72.00
91.14
8.16
22.37
29.88
63.26
66.41
90.29
DeiT-Twicing [1-12]
72.60
91.33
9.15
24.10
32.28
65.67
67.12
90.53
DeiT-Twicing [7-12]
72.45
91.35
8.67
22.90
31.60
64.79
66.48
90.52
Dwii-T [10-12]
72.31
91.24
8.66
22.58
31.63
64.74
66.47
90.49
Table 10: Evaluation of the performance of DeiT and DeiT-Twicing in ImageNet classification under the presence of different corruptions, using appropriate evaluation metrics for each.
Dataset Metric
ImageNet-R Top 1
ImageNet-A
ImageNet-C mCE (1)
ImageNet-C (Extra) mCE (4)
Top 1
DeiT (Touvron et al., 2021)
32.22
6.97
72.21
63.68
NeuTRENO (Nguyen et al., 2023h)
31.65
8.36
70.51
63.56
Deif-Twicing
32.74
7.66
70.33
62.46
DeiT-Twicing [7-12]
32.68
8.10
69.98
62.35
DeiT-Twicing [10-12)
32.31
8.14
70.25
62.63
Non-conventional attention mechanisms. To further validate the broad applicability of Twic- ing Attention (or twicing procedure general), we conduct additional experiments following the theoretical setup of Linear Transformers (Katharopoulos et al., 2020). To be more precise, Ta- ble (last 2 rows) compares the perplexities recorded for Linear Transformers with feature map ((x) elu(x) )++ matching their original choice, and Linear- Twiciin Transformers for which we apply the twicing transformation 2A A where
normalize (Q))(K))) trained for 75 epochs, while Table 8 shows that Linear-Twicin model has superior robustness against the Word Swap Attack. Note that we explicitly construct the similarity matrix A for both of the models for framework to work. The positive results indicate that the applicability of Twicing Attention is not limited to standard softmax self-attention but any reasonable similarity matrix can be covered.
B.2 IMAGENET IMAGE CLASSIFICATION AND ADVERSARIAL ATTACK
Dataset. We use the full ImageNet dataset that contains 1.28M training images and 50K validation images. The model learns to predict the class of the input image among 1000 categories. We report the top-1 and top-5 accuracy on all experiments.
Corruption. We use attacks FGSM (Goodfellow et al., 2014) and PGD (Madry et al., 2017) with perturbation budget 4/255 while SPSA (Uesato et al., 2018) uses a perturbation budget 1/255. All attacks perturb under L. norm. PGD attack uses 20 steps with step size of 0.15.
Model, Optimizer & Train Specification. The configuration follows the default DeiT tiny config- uration (Touvron al., 2021). In particular, we follow the experimental setup of (Han al., 2023; Nguyen et al., 2022c). To this end, the DeiT backbone uses 12 layers, 3 heads of dimension 64, patch size 16, feedforward layer of size 768 and embedding dimension of 192. We train using Adam with starting learning rate of 0.0005 using cosine scheduling under default Py PoTorc settings, momen- tum of 0.9, batch size of 256, 5 warmup epochs starting from 0.000001 and 10 cooldown epochs, for overall train run of 300 epochs. The input size is 224 and we follow the default Augment policy and color jätter 0.4.
Extra results. In Figure 5 and Figure 6, report that our model consistently outperforms the baseline across increas- ing six levels of severity under FGSM and PGD attacks. Ad- ditionally, Table 11 compares our model versus FeatScale (Wang al., 2022), vision transformer addressing over- smoothing. Our model outperforms in both metrics on clean ImageNet classification.
Table 11: ImageNet Classification.
Model
Top 1
Top 5
DeiT
72.00
91.14
DeiT-FeatScale
72.35
91.23
Dwii-nn
72.60
91.33
22

---

Published as a conference paper at ICLR 2025
ImageNet Classification
ADE20K Segmentation
Token Cosine Similerity
Token Cosine Similarity
0.9
0.8
Dei -lwicing
0.8
Deil-lwicing
0.7
0.7
0.0
0.6
0.5
0.4
0.5
0.3
Layer
Layer
Figure 8: Comparison of average token cosine similarities across layers for DeiT and DeiT-Twicing models. Subfigure (a) uses ImageNet, while subfigure (b) evaluates ADE20K segmentation.
B.3 OUT-OF-DISTRIBUTION ROBUSTNESS AND DATA CORRUPTION IMAGENET-A.R.C
ImageNet-A,R,C are benchmarks capturing a range of out-of-distribution and corrupted samples (Hendrycks et al., 2021). ImageNet-A contains real world adversarially filtered images that fool current ImageNet classifiers. ImageNet-R contains various artistic renditions of object classes from the original ImageNet. ImageNet-C consists of 15 types of algorithmically generated corruptions with 5 levels of severity. (e.g blurring, pixelation, speckle noise etc). Figure 4 shows that DeiT- Twicing (our) model outperforms DieTT baseline all 15 main types of ImageNet-C corruptions.
B.4 ADE20K IMAGE SEGMENTATION
Experimental setup. We adopt the setup in Strudel al. (2021). The encoder is pretrained on ImageNet-1K following the same specification described in Appendix B.2. In particular, the encoder DeïT-tiny backbone of 5.7M parameters pretrained for 300 epochs. After pretraining, we then attach a decoder that contains 2-layer masked transformer and finetune the full encoder-decoder model for 64 epochs on the ADE20K Zhou al. (2019) image segmentation dataset.
C COMPUTE RESOURCES
Training. All models are trained using four NVIDIA A100 SXM4 40GB GPUs including both language and vision models.
Imagenet Classification under adversarial attacks are evaluated using two NVIDIA A A00 SXM4 40GB GPUs while only one of such GPUs was used evaluate against ImageNet- A,R,C and Word Swap Attack for language modelling.
Evaluation.
D ADDITIONAL EMPIRICAL ANALYSIS
D.1 EXTRA OVER-SMOOTHINS ANALYSIS
We show the layer-wise over-smogthing analysis on both ImageNet classification and ADE20K image segmenta- tion in Figure 8. We observe that in both cases, aver- age token cosine similarities with Twicing Attention grow slower than those with standard self attentoten,on, once again validating our theoretical findings.
D.2 EXTRA ATTENTION HEATMAP ANALYSIS
Extending the visualizations in the main text, Figure 9 illustrates how attention heatmaps evolve from early to late layers for DeiT and Twii-T models given the input images. In Figure 10, we provide 12 more examples show that when employed Twicing Attention, DeiT- Twicing is usually more accurate to recognize objects in images and shows substantially better expressive power by capturing more meaningful parts of objects without missing target.
Figure 9: Evolution of attention heatmaps from early late layers. Odd rows: DeiT-Twicing: Even rows: DeÎT.
23

---

Published as a conference paper at ICLR 2025
DelT
DeffTluccig
DelT
DeifTwichng
DaT
DeflTlichng
DalT
DefT.wichhg
DeiT
DeffTluccig
DeffTwichng
DaIT
DeftTWichig
DalT
DefTTwichng
Figure 10: More examples showing how Twicing Attention is capable of retaining model's expres- sive power. The DeiT baseline model frequently collapses the entire image into the background, particularly when the background occupies a significant portion of the image, making it challenging distinguish object details. Only few cases, such as the example in bottom right, trying cap- ture more information was not as successful while still being highly competitive.
ABLATION STUDIES
Since our proposed method does require any additional parameters learnable. neither tunable, we only study the layer placement for Twicing Attention. Table 9 demonstrates 3 different layer (full),
placements of Twicing Attention find that as long as Twicing Attention is placed starting from later layer till the end, the performance improvements are almost always proportional to the number of layers employing Twicing Attention. In Table 10, however, we observe that this proportionality does not happen in general since all three types layer placements lead to good results in different categories, but all beating the baseline by notable margins. We also find that putting Twicing Attention only in few inital layers may not offer significant overall improvements.
12 (last six), and 10 to 12 (last three). We
24


Reviews:
[{"review_id": "NJhy156OC1", "reviewer": "Reviewer_pR7y", "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$.", "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.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper is relatively easy to follow and well-written.\n2. The proposed \"Twicing Attention\" is simple and easy to implement.\n3. 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.\n2. 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "kif2jxPPDU", "reviewer": "Reviewer_3adS", "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.", "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.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. Novelty: The authors introduce the Twicing Attention mechanism to slow down the eigenvalue decay associated with representational collapse.\n2. Theoretical Contribution: the authors provide mathematical validation for the Twicing Attention’s capability to retain information across layers.\n3. 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.\n2. Lack of comparison: the work does not compare its method with alternative solutions that address oversmoothing, such as regularization strategies.\n3. Lack of ablations: the authors are suggested to consider applying the proposed method at different layer depths or intervals and evaluate their difference.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Jz8IIqDvMI", "reviewer": "Reviewer_CqxC", "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.", "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.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well-written, with clear expression of formulae and symbols, and will be highly readable.\n2. 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.\n3. 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.\n2. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "km3dlA1vSG", "reviewer": "Reviewer_7QSP", "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.", "questions": "- Visualizations on earlier layers and more heads of the transformers would help to strengthen your claim.\n- Please refer to the weakness.\n- I am open to increase my score if you alleviate my concerns.", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- Theoretical foundation: The paper provides analysis connecting self-attention to NLM filters. \n- Fluent presentation flow: Messages of this paper are presented well, with well-demonstrated background knowledge and motivation.\n- 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.\n- **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.\n- **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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)