Mitigating Over-smoothing in Transformers via Regularized Nonlocal Functionals
Authors: Tam Minh Nguyen
Tan Minh Nguyen
Richard Baraniuk
Venue: NeurIPS 2023 poster (2023)
Keywords: transformers
self-attention
total variation
nonlocal functionals
over-smoothing
TL;DR: We develop a new regularized nonlocal functional for self-attention to overcome the over-smoothing issue in transformers.

Abstract: Transformers have achieved remarkable success in a wide range of natural language processing and computer vision applications. However, the representation capacity of a deep transformer model is degraded due to the over-smoothing issue in which the token representations become identical when the model's depth grows. In this work, we show that self-attention layers in transformers minimize a functional which promotes smoothness, thereby causing token uniformity. We then propose a novel regularizer that penalizes the norm of the difference between the smooth output tokens from self-attention and the input tokens to preserve the fidelity of the tokens. Minimizing the resulting regularized energy functional, we derive the Neural Transformer with a Regularized Nonlocal Functional (NeuTRENO), a novel class of transformer models that can mitigate the over-smoothing issue. We empirically demonstrate the advantages of NeuTRENO over the baseline transformers and state-of-the-art methods in reducing the over-smoothing of token representations on various practical tasks, including object classification, image segmentation, and language modeling.

Mitigating Over-smoothing in Transformers via Regularized Nonlocal Functionals
Tam Nguyen
Tan M. Nguyen
Department of Electrical & Computer Engineering
Department of Mathematics
Rice University
National University of Singapore
Houston, USA
Singapore
mn72@rice.ede
tanmn edu...ss.dde... edu.
Richard G. Baraniuk
Department of Electrical & Computer Engineering
Rice University
Houston, USA
richberice. edu
Abstract
Transformers have achieved remarkable success in a wide range of natural language processing and computer vision applications. However, the representation capacity a deep transformer model is degraded due to the over-smoothing issue in which the token representations become identical when the model's depth grows. In this work, we show that self-attention layers in transformers minimize a functional which promotes smoothness, thereby causing token uniformity. We then propose a novel regularizer that penalizes the norm of the difference between the smooth output tokens from self-attention and the input tokens to preserve the fidelity of the tokens. Minimizing the resulting regularized energy functional, we derive the Neural Transformer with Regularized Nonlocal Functional (NENTRENO), a novel class of transformer models that can mitigate the over-smoothing issue. We empirically demonstrate the advantages NeuTRENO over the baseline transformers and state-of -fe---fftthtttee art methods in reducing the over-smoothing of token representations on various practical tasks, including object classification, image segmentation, and language modeling.
1 Introduction
Transformer models [62] have achieved substantial success in natural language processing [16, 2, 13,
10, 47, 4, 6, 14], reinforcement learning [9, 32), computer vision [19, 40, 59, 49, 44, 3, 41, 71, 27), and other practical applications [50, 33, 70, 26, 66). Transformers also excel at transferring knowledge from pre-trained models to new tasks, even when limited supervision available [45, 46, 16, 69, 39]. At the heart of transformers lies the self-attention mechanism, which computes weighted average of token representations within sequence. These weights are determined based on the similarity scores between pairs of tokens, determining their relative importance in the sequence [11, 43, 38]. This flexibility in capturing diverse syntactic and semantic relationships has been identified as a crucial factor contributing to the success of transformers [57, 63, 12, 64, 31].
1.1 Background: Self-Attention
For given input sequence X : [=x11, [[N] Nx of N feature vectors, self-attention transforms X into the output sequence H in the following two steps:
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

DeiT
NeuTRENO DelT
ImageNet Classification
ADE20K Segmentation
0.9
0.5
Cosine Similarity
between tokens
0.
0.6
0.5
0.4
0.3
0.3
10
12
10
Layer
Layer
Figure 1: The cosine similarity between tokens representations across layers of NeuTRENO DeiT vs. the baseline DeiT models on Imagenet classification and ADE20K image segmentation tasks. In both tasks, DeiT baseline suffers from over-smoothing as tokens become similar to identical when the model gets deeper. In contrast, tokens in TRENO models are significantly more diverse, suggesting reduction in over-smoothing. Further details regarding this analysis can be found Appendix E.
Step The input sequence X is projected into the query matrix Q, the key matrix K, and the value matrix V via three linear transformations
Q = 2: XW Ki XWT.
where We, WK K E D D2 and Wv E RxDx are the weight matrices. We denote Q [g(1)..... (N) K == ((),, (i), and vi). for I,..., N are the query, key, and value vectors, respectively.
[(N) and V := [v(1),..., v(N)] where the vectors
Step 2. The output sequence U := [u(1). u(N) E RNN is then computed as follows
U softmax
QK /TTT/DDs V AV,
where the softmax function applied each row of the matrix QK T /v Dyk- The matrix
softmax VD Duk
E NxN and its component oij for i, j = 1, N are called the attention
matrix and attention scores, respectively. For each query vector q(i) for i 1, ,N, an equivalent form of Equ. (2) to compute the output vector u(i) is given by
softmax (q()) ((())((()))))))))))))))))))))))) Dan v(i).
(3)
The self-attention computed by Eqn. (2) and (3) is refered softmax attention. In our work, we refer to transformer that uses softmax attention as a softmax transformer.
1.2 Over-smoothing in Transformers
Despite their remarkable success, deep transformer-basee models have been observed to suffer from the over-smoothing issue, in which all token representations become identical when more layers added to the models [55, 65, 18). This over-smoothing phenomenon, also known the "token uniformity" problem, significantly limits the representation capacity of transformers. To illustrate this phenomenon, examine the average cosine similarity between pairs of token representations across different layers in softmax transformer trained for the Imagenet object classification and ADK20 image segmentation tasks [73]. As depicted in Fig. in both tasks, this cosine similarity between tokens increases as the models become deeper. Particularly, in the last two layers, the cosine similarity scores approximately 0.9, indicating high degree of similarity among tokens.
1.3
Contribution
We develop nonlocal variational denoising framework for self self-attention, providing insights into the over-smoothing phenomenon transformers. particular, by viewing self-attention as a gradient descent step toward minimizing nonlocal functional that penalizes high-frequency noise in the signal, we uncover the diffusive nature of self-artention, which explains the over-smoothing issue of transformers. Motivated by this understanding, propose the Neural Transformer with Regularized

---

Nonlocal Functional (NeuTRENO), novel class of transformers designed mitigate over-smoothing. NeuTRENO is derived by optimizing regularized nonlocal functional, which includes an additional convex fidelity term. This fidelity term penalizes the norm of the difference between the smooth output tokens from self-attention and the input tokens, thereby reducing the over-smoothing effect. Our contribution is three-fold.
1. We develop nonlocal variational denoising framework for self-attention and shed light on the over-smoothing issue that hampers the representation capacity of transformers.
2. We develop NeuTRENO, a novel class of transformers that are capable of alleviating the over-smoothing issue.
3. We theoretically prove that transformers with softmax self-attention are prone to over- smoothing while NeuTRENO avoid this issue.
We empirically demonstrate the benefits of NeuTRENO on various large-scale applications, including the ImageNet object classification, ADE20K image segmentation, and WikiTexx-103 language modeling tasks.
We organize our paper as follows: in Section 2, we develop a nonlocal variational denoising framework for self-attention and provide explanation for the over-smoothing issue in transformer-basee models. In section 3, we propose NeuTRENO, and present a theoretical result that guarantees NeET'SNO's capability of mitigating over-smoothing. In Section 4, empirically validate the benefits of NeuTRENO. We discuss the related work in Section Finally, we conclude main contributions and remarks. Further results, details, and proofs are provided in the Appendix.
Organization:
2 A Nonlocal Variational Denoising Framework for Self-attention
We first consider the output matrix U == [u(1), u(N]] E RND in self-attention as given by Eqn. 2 in Section 1.1. Let Ω R, Ω2 and u(x) = (z) "D(x))² be a real rector-valued function, u: : Ω + RD, u E (2). The output matrix U in self-attention discretizes the function 1-D grid. In the context of signal/image denoising. U can be considered as the desired clean signal, and u((z) (xx) is its corresponding intensity function denoting the signal values at the position Ω. We further let the observed intensity function fx) denote the values of the observed noisy signal at x E Ω f Ω + RP, f (I()). For example, fxx can be given as f(x) u(x) + n(x),
(x))
where is the additive noise. We wish to reconstruct uxx) from f(z). Following the variational denoising method proposed in [23] and [24], the denoised image uxx can be obtained by minimizing the following regularized functional with respect to u:
E(u, f) = J(u)= + G(u, f)
(5)
À 1 { [1((x) u(y)l1k(x, yddxd + + f [|(()) f(x) [1
2
- ( y) dxdy is a nonlocal functional of weighted differences. The weights k(xr, 3) represent the affinity between signal values at positions x and For example, for images, k(x,, 3) captures the proximity between pixels x and y the image. ((u) works regularizer. Minimizing ((u) promotes the smoothness of u and penalizes high-frequency noise in
Here, J(u) = Saxa (lu(x)-u(v))
the signal. Adding the convex fidelity term G(u, f) = aे Jo [() - ( to the functional J(u) allows the denoised signal u(x) to preserve relevant information the observed noisy signal
fx).. The regularized functional E(u,, f) can be considered
energy functional.
2.1 Self-attention
Gradient Descent Step to Minimize the Nonlocal Functional J
We show that self-attention is equivalent to taking a gradient descent step toward minimizing the functional (u) in the energy functional Elu, f). We expand J (u) as follows
1 / 1 2 AxR 1=1
J(u)
Σ
u. (,())) (xx()))
dxdd
(6)
The gradient of J with respect is then given by
OJ aJ
0J
T
V VuJ(u)= =
(7)
Fu
Ouz
Oun
3

---

The partial derivative 0J 0J/Ou,. = 1,2,.... ..., D, is defined through its dot product with an arbitrary function h E (92) as follows
DJ
d hy(z)= (uy + Thy)l,=
Juj
( d
dr Jaxa (u,(z) uy(x) + (h)(2) (h) k(x, yddrdy
==
(n (u)(x) uy(y) rh)(2)
- - Thy - (v))(((a y) y)dzyy |
I (u,(x) u,(v))(h) (,(1)) (2f((x.))dd) g)ddy
Jaxa
Jaxa (1z(x) u,(y))) (xx yddxdy J. (uy(z) us()) y)ddd
Applying a change of variables (I, v) - (y. .) to the second term of the above integral, we have
OJ
hg(x) I Jan (u)(x) - u (x) (x yddxdy - Saxp (u,(w) (x))) x) ddddy
0u,
- I (u5(x) - -u,(y) v) kyy x) (d)bh
Thus, the Frechet derivative of J with respect to uj is given by
Ou, OJ [ (u,)(x) - (,(0)(( + 4(y,x)))dg
(8)
Substituting the formula for 0J/ /Ju in Eqn. into Eqn. for vu) we obtain the following gradient flow
du(x, 1) [()) -v - ulx, (u(x,())u(x;()))) (M(2,, + k(y,:))dd
V
(9)
dt
where is the time variable we introduce to capture the dynamics of u when gradient descent is applied to minimize J(u). Let vxx) == [v) (ar) (cr),.... un (x))² be a real vector-valued function, v Ω RDD. E (22).. We discretize v(x)
a 1-D grid to attain the value vectors v(1).... v(N) E R
which form the value matrix V := [v(1).
v(N) E NxD in self-attention as defined in Eqn. 2. We initialize u at t 0 with v(z). i... u(z, 0) = u(x).
Self-attention an Euler Discretization of the Gradient Flow Given in 9. We discretize the gradi- ent flow in Eqn. 9 using the Euler method [21] with step size At(x) 1/52 (k(2, v) + k(y, and obtain the following update
(t()) = u(x, + At(x) (u(v. 0) uxx 0)) (k(r, v) + k(y. xy
1 (k(2, v) + k(y, (k(k(z (1) + kly x)) dy dy I K(I,u)o(v) In ((.,8')d
dy. (10)
Here, K(x, y) == k(x, y) k(y, x) is a symmetric kernel and u 0) = v(y) since u is initialized at t = 0 with aforrementioeed. Let k(x) : [Ka (x)...,, KD., (z))² be real vector-valued function, : Ω. + Pt &s E (22( Similar (x)))) and v(z)), we discretize k(x)) on 1-D grid attain the key vectors k(1)..... A<N) E Pev, which form the key matrix K = [A(()) [(N) e xx in self-attention as defined in Eqn. 2. We choose K K(x,1) v) = exp Ar x k(w) Dat) and
rewrite Eqn. 10 as follows
u At(z))= / (()(K())"())/////
v(y)dy.
(11)
(xp(kx(x)TT(()))) /////D/ /Dxx)

---

Estimating the integrals in Eqn. 11 via Monte-Carke approximation using the key vectors E Po and and value vectors v( v(N) E RD,, we obtain
N
exp (
u(r, At(x)) R
vii. (12)
2j= exp(k(x)) /Dxk)
Discretizing uar, At)))) another 1-D grid, we attain
u exp exp(k(i) ((())))())//// /Dak)
v(j)
j'=l exp 1(1) (())"
softmax k(i) k()/ V D v(j), i= 1, I,...,N. N.
(13)
==1
Comparing Eqn. 13 and Eqn. 3, we observe that Eqn. 13 implement a symmetric self-attention, in which the query matrix Q and the key matrix K are the same, i.e. We "W where We and into Q and K as given Eqn. 1. This
WK are the linear projections that map the input sequence symmetry of the attention scores is desirable in some image processing tasks due the symmetric similarities between pixels, but can be relaxed for other tasks. To break the symmetry of attention scores in Eqn. 13, we replace the key vectors (<i) by the query vectors g i 1,..., N. obtain the exact formula of self-attention given by Eqn. 3. The following theorem summarizes our results:
Theorem (Self-attention as a Gradient Descent Step to Minimize a Nonlocal Functional). Given the nonlocul functional J (w) = Saxn [|((z) u (vv) y)dzdy of a vector-valued function u : + >R,, u E L (S2), and let K ((x y) == k(x, y) + k(y. x) =
(kyy)///////)))) Dyk), where k : Ω + R k E (2))). Then, taking a gradient descent step on u at time t = 0, where to minimize J is
(x, v(x), with adaptive step size A1(x) Sa(k(x, k(y, x)
equivalent to updating via symmetric self-attention
u(r, A((x)) Σ softmax (M(x) k(j)/ V DDt vo(),
which results in
N
u11)= = Σ softmax (/(T)(/////// Dak o(1).
i i= 1...,.,,
(14)
Here, u(n)) v(n). and un)). = 1. ..... N, are the key, value, and output vectors in self-attention, respectively. Breaking the symmetry of the attention scores by replacing k(i) with q(),, 1,..., N, in Eqn. 14, we obtain the exact formula of self-attention
u(i) = softmax (9(1) kT) ( (D w(f), i i= 1
Remark 1. In Eqn. 9, the change in position is proportional to the sum of differences between (x)) and u at other position in the domain Q. In particular, when u(x) is smaller or larger than the diffusion
values at other positions, it will increase decrease, respectively. This is analogous process in which particles or substances move from high-concentration to cov-concentration regions. It has been proved that diffusion process converges saturating state in which the concentrations at all positions are the same. This suggests that u(x) tends to suffer from the over-smoothing issue.
2.2 Random Walk Analysis of Over-smoothing
The diffusion process and random walk are closely related concepts, as diffusion can be seen as collective behavior of numerous random walks performed by individual particles molecules. Inspired by the analogy between the dynamics of u in Eqn 9 and a diffusion process, as well as the relationship between diffusion process and random walk, in this section, show the connection
5

---

between the evolution of and random walk. By adopting random walk perspective graph neural network [58], we demonstrate that ux)) under the dynamics given in Eqn 9 suffers from over-smoothing.
Recall from the gradient flow in Eqn by using Euler method discretization, after update steps starting from the initial u(x 0) = (xx) with adaptive stepsize AL - 1//o (A((z, v) + k(y, x)) dy.
obtain the following
(a, kat())) / (x, y)u (A. 1))t(a)) dy.
(15)
Ja K(x, d'dd
Discretizing uax, AA)))) and using Monte-Care approximation for the integrals in 15, we attain
N u (1) Σ A
(16)
where A, is computed using the keys and queries as either softmax (<(1) ((j)) /Dgx or
softmax ((()axx(()))a) (g))))(()) D Let ( be random walk on {((())}= as defined:
B ( v(i) P(B (1) e(j)
(17)
v(i)) = Au
where ((x)(n) is the random value of a k-step walk, starts at node n. and v(n) is the initial value node n, respectively, for 1, 1, ... N. The transition probability A is defined as above. To investigate the connection between the update process of u and the random walk defined in 17, we show that, for é = 1, 1.2...., NN, after update steps as in 16, with initial value u (0) = vi). (1) equals to the expected value of the k-step walk, starting at node i:
Lemma 1. Let u (i) defined in 16 and (1)) the random walk defined by 17. Then
(1) E ( (i)).
(18)
We next present the Lemma 2 which is necessary to show the convergence of (k) (i).
Lemma 2. The random walk (i) 17 with the transition matrix either be Aij softmax ((i))(()/////) Dakk) or Aij softmax (B)])(()//// has a unique stationary distribution π = [TI, T2....,mT such that Ti = P(B()) = v(i)), for E
i,j 1.2...., N, i==1 Te 1, and n- = T.
If A. softmax (A(iT (())///////// /(t)/ the stationary distribution is:
di
d2 Ej
dn
π E d
(19)
Σ
d
where d; = EE exp k(i) ((())
are the key vectos.
In general. can be found by finding the left eigenvector of A corresponding to the domi- nant eigenvalue 1.
From the Lemma and Lemma 2. we see that, for all = 1,2, N,
N
(+) (i) = E (x) = Σ (B)P(B( lk (1) v()) Σ T=(())=0
(20)
as 1: - 0o. This shows that when & increases, u(i) (+) converges to a constant vector, indicating that uu22 (z). under the dynamic in 9, suffers from over-smoothing.

---

42
Ko
Self-Attention
FFN
k1
Self Attention
FFN
Kg
Self-Attention
V0
1/2
a(Vo v1)
a
V22
Figure 2: Our proposed NeuTRENO model adds proportion of the difference between the values of the first and that of the current layer the self-attention's 's output each layer.
NeuTRENO: Mitigating the Over-smoothing Transformers via Minimizing Regularized Functional
In Section 2.1, we have shown that self-attention implicitly performs a gradient descent step to minimize the nonlocal functional J(u) in Eqn. 5, which results the diffusive characteristics of and causes the over-smoothing phenomenon in transformers, proved in Section 2.2. Fortunately, our objective is not to minimize J(()) but the energy/regularized functional E(u, f) defined by Eqn. 5. This regularized functional consists of not only J(u) but also the convex fidelity term G(u, f) Ja [|x(x)) f(x) This fidelity term aims to preserve the relevant information in the observed noisy signal f(x) by penalizing solution uxx) that deviates significantly from f(z), thereby mitigating the effects of over-smoothing caused by minimizing J(u).
In this section, we will derive our Neural Transformer with a Regularized Nonlocal Functional (Neu'REEN by minimizing the regularized functional E(u, f). We then provide theoretical result to prove that NeuTRENO does not suffer from over-smoothing. Recall from Eqn. 5 that E(u, f) is given by
D J (1y(z)
E(u. f) = J(u) + G(u, = J(u) +
Following similar derivation in Section 2.1 (see Appendix for the detailed derivation), we obtain the following gradient flow when minimizing E(u, f) using gradient descent
du(x,, t) VV, vu) -(u(z)-f(a)) f(x)),
(21)
dl
Neu TRENO-attention is an Euler Discretization of the Gradient Flow Given in 21. Following the similar derivation in Section 2.1, we discretize the gradient flow in Eqn. 21 using the Fuler method [21] with step size At(x) 1/Jg (A:( 8) + k(y, x))x and initializing u at with (2x), i.e., uxx 0) v(x). Choosing > = \/A((1)) we obtain the following update
u(a, Al(x)) = u (r 0) A.. (u(x)) f(x))
[ K K(z,g)o(v) dy +
(22)
Jn K(x d'ydy
We choose the observed noisy signal f(x) = (x) where v²(x) v(r) at the first layer in the transformer model. The update in Eqn. 22 becomes
u(x, At(z)) [
K K(I,9)o(V) Jo K(x, d'y dy X v(x)).
(23)
Applying the Monte-Carlo method to approximate the integrals in Eqn. 23 and discretizing (xx At(x)), v(z), and (xr) on a 1-D grid, we attain the following new formula for calculat- ing symmetric self-attention:
N
(1i)= Σ
softmax (b(1)T(()) h(t) DDz v(j) + ((o w()),, i= 1, 1,...,N N. (24)
j=1

---

Table 1:
Top-1 1 and Top-S accuracy (%) of NeuTRENO D vs. DeiT on the ImageNet benchmark. We also present the performance of adapting NeuTRENO the pre- rre-trained DeiTT baseline. NeuTRENO Adaptation. In addition, we compare NeuTRENO with FeatScale [65] and incorporate our method with FeatScale model.
Model/Metric
Top-I Ace (%)
Top-5 Acc (%)
Softmax DeiT
72.17
91.02
NeuTRENO
73.01
91.56
NeUTRENO Adaptation
72.63
91.38
DeiT FeatScale
72.346 73.23
91.22
NeuTRENO DeiT FeatScale
91.73
Its corresponding asymmetric self-attention is obtained by replacing the key vectors k(i) with the query vectors q(i), i- 1, ,, NN and given by
N
uli)= Σ softmax ((())(()//// Das v(1) ((x v()), i= I....,,, N. (25)
Leveraging Eqn. 25, define the Neural Transformer with Regularized Nonlocal Functional (NeuTRENO) as follows.
Definition (Neural Transformer with a Regularized Nonlocal Functional (NeuTRENO). Given set of key and value vectors {k)(), of (G));= in each layer f, e= =......, L, for each query vector (i), I....,N, in the same layer, the self-attention unit layer a Neural Transformer with Regularized Nonlocal Functional (NeuTRENO) computes the corresponding output vector ui (i) of the query (i) by the following attention formula:
N
ul (1) = Σ
soft max ( + o i 1= N. (26)
Ø=1
where 00 (1), (N) E P are the value vectors in the first layer of NeuTRENO.
Fig. 2 illustrates the architecture of NeuTRENO.
Proposition 1. The evolution of u(x) under the dynamic in 21 does not converge to a constant vector.
Proposition indicates that our NeuTRENO mitigates the over-smoothing issue, suggesting the benefit of our method. The proof for Proposition is given in Appendix B.3.
4 Experimental Results
In this section, we empirically demonstrate the advantages of our proposed NeuTRENO approach across various tasks, including ImageNet classification [15], ADE20K image segmentation [73], and language modeling on the WikiText-1003 [42]. Our aim to show: (i) NeuTRENO significantly outperforms the transformer baseline with softmax-attentie moreover, TeUTRENO surpass FeatScale, vision transformer that addresses wwerssmothhing combining NeuTRENO with FeatScale beneficial; (ii) the advantages of incorporating our proposed method with trainee models. We also demonstrate the benefits of our NeuTRENO in the symmetry setting and we point to Appendix for the results. Throughout our experiments, we compare the performance of our proposed models with baselines of the same configuration. For additional details regarding datasets, models, and training procedures, please refer to Appendix A.
defined in across various tasks;
Object classification ImageNet. To demonstrate the advantage of NeuTRENO method, compare it with the DeiT baseline [59] on the ImageNet image classification task. Our NeUTOENO DeiT surpasses the DeiT baseline, as shown in Table 1. Notably, our NeuTRENO DeïT achieves significantly higher performance in terms of both Top-1 Accuracy and Top-5 Accuracy. We also compare method with FeatScale [65], vision transformer model addressing over-smoothing (see Table 1). Our NeuTRENO significantly outperforms FeatScale, and combining NeuTRENO with FeatScale leads substantial improvements. These results confirm the benefits of our model.
Image Segmentation on ADF20K dataset. To further validate the advantages of our proposed methods, we compare the performance of the Segmenter models [56] using the NeuTRENO DeiT

---

Table 2: Single-scale (SS) MIoU and multi-scale MIoU (MS) of the NeuTRENO DeiT vs. the DeiT on the ADE20K image segmentation.
Model/Metric
SS MloU
MS MloU (9)
Softmax DeiT
35.72
36.68
NeuTRENO DeiT
37.24
38.06
Table 3: Test and valid perplexity (Test PPL and Valid PPL) WikiText- (03TTxxt-133-o of NeuTRENO compared to the softmax transformer. Our proposed method achieves significantly better performance PPL than the baseline.
Method/Metric
Valid PPL
Test PPL
Softmax Trunsformer
33.15
34.29
NeuTRENO
32.60
33.70
Test
16
109
30
J(u)
10"
207
10'
10
Layer
Layer
Figure 3: The average value of functional J(u) over 1000 training (Left) samples and test (Right) samples. When softmax attention is applied. the functional decreases the depth of the trained DeiT increases. and DeiT backbones the ADE20K image segmentation task [72], as shown in Table 2. The results demonstrate the substantial performance improvements achieved by utilizing the NeuTRENO DeiT backbone over the DeiT backbone, in terms of both single-scale (SS) MIoU and multi-scale (MS) MIoU metrics. These results strongly emphasize the effectiveness of our NeuTRENO approach in enhancing image segmentation performance.
Language Model on WikiText-103. addition to computer vision tasks, we also evaluate the effec- tiveness of our model on large-scale natural language processing application, specifically language Our Neu language model demonstrates better performance in terms of both test perplexity and valid perplexity when compared to the softmax transformer language model [68]. These findings, combined with the results obtained across various tasks, empirically confirm the significant benefits of our NeuTRENO models.
modeling on WikiText-
Combine with pre-trained models. Furthermore, our proposed method also beneficial combine with pre-trained models. To empirically demonstrate that we incorporate NeuTRENO with pre-trained DeiT and fine-tune on the ImageNet dataset with one-third number of epochs that are used in training. The result is presented in Table showing that combined with our method improves both the Top-1 and Top-5 accuracies of the pre- pre-trained models.
Empirical Analysis
Applying Neltmax-Attentioo Reduces the functional (u). We present evidence supporting that the employment of softmax attention minimizes the functional J(u). Initially, we observe that the average cosine similarity between the numerical approximation of V.J(u)) using symmetric or asymmetric kernel K (a, y) for both the trained Sym-DeiT (using symmetric self-attention 14) and DeiT models, closed as shown in Table 4. This suggests that reversing the direction of the asymmetric approximation effectively decreases (u)) Considering that softmax attention takes steps in this reversed direction numerically, its application leads reduction in J (u) This is further substantiated by Fig. 3, which demonstrates a decrease in J(u) as the depth of the trained DeiT increases when softmax attention is employed. More details of this analysis in Appendix E
Over-smoothing Analysis. We empirically illustrate the effectiveness of NeuTRENOs in mitigating the over- smoothin problem in transformers. Fig. I compares the cosine similarity between token representations across layers for both NeuTRENO and softmax baseline models, specifically focusing on the Imagenet classification task (Left) and ADE20K image segmentation (Right). The token

---

Table 4: The average cosine similarity between the numerical approximation of (1)))) using symmetric asymmetric kernel K(2; 2). for the trained Sym-DeiT on 1000 training and 1000 test data samples. The average score close to shows a strong alignment between symmetric and asymmetric gradient approximations, suggesting that reversing the direction the asymmetric approximation effectively reduces the functional J(u).
DeiT and softmax DeiT models. The metric evaluated
Model
Training data
Test data
Sym-DeiT
0.982
0.976
Softmax DeiT
0.973
0.964
features extracted by NeuTRENOs exhibit significantly lower similarity, particularly in the final layers. This finding highlights the ability of NeuTRENOs to address the over-smoothing issue and improve the diversity of token representations. We provide more details of this analysis in Appendix E.
6 Related Work
Over-smoothing in Transformers. Over-smoothing in deep transformers has been observed in various domains and applications from natural language processing [55] computer vision [65, 18]. In vision tasks, [74] observes that the performance of the vision transformer (ViT [20]) quickly saturates more layers are added to the model. Moreover, experiments in [74] show that the 32-layer ViT underperforms the 24-layer VIT, indicating the difficulty of ViTs in gaining benefits from deeper architectures. The authors point out that over-smoothing results in this phenomenon by causing the token representations to become identical when the model grows deeper. Based on this observation, the authors propose a cross- head communication method that helps enhance the diversity of both token representations and attention matrices. Furthermore, it has been shown in [60] that the training of VIT models encounters instability with greater depths. [25] proposes that this instability arises from the over-smoothing, where token representation for patches within an image becomes progressively alike as the model's depth increases. To explain this issue, [65] finds out that self-attention acts as a low-pass filter and smoothens the token representations in ViTs. This leads to the proposal of the FeatScale method [65], which regulates feature frequencies, whether low or high, to counteract the consequences of over-smo.
[55] observes the phenomenon in BERT [16], deep language model, and explores over-smoothing through the graph perspective. The work utilizes hierarchical fusion strategies by preserving the output of self-attention through all layers, which is memory-costly. On the other hand, [65, 18] investigate over wwer-smoothing in the image domain through the lens of Fourier spectrum, showing that self-attentions are low-pass filters, retaining only low-frequency, causing over-smoothed outputs. Our work is an orthogonal explanation of the previous work. We focus on developing variational denoising framework to understand the self-attention of transformers as a gradient descent approximation of functional. Our new finding explains the over-smoothing issue of transformers due to self-attention minimizing a functional and inspires us to derive the novel NeuTRENO method to overcome over-smoothing.
In addition,
Nonlocal Functionals for Image Processing. Total variation [51] is well-known as an image- denoising technique. It denoises noisy image by solving a constraint optimization problem. The method is also related to SoDE-floo-basee image-
techniques [24], namely isotropic and anisotropic diffusion [67] models. The method is edge preserving, meaning to avoid over-blurring edges' information [7]. Nonlocal functionals [35, 24] is considered as an extension of total variation to nonlocal scale. Nonlocal functional and edge preservation properties are the motivation of our work to explain and overcome over-smoothing in transformers.
7 Concluding Remarks
In this paper, we establish a nonlocal variational denoising framework for self-attention. From this variational perspective, explain over-smoothing in self-attention, which hinders the representation capacity of transformer models. We also derive the novel Neural Transformer with a Regularized Nonlocal Functional (NeuTRENO) alleviate the over-smoothing. We empirically verify the benefits of NeuTRENO with wide range of large-scale applications including ImageNet object classification, ADE20K object segmentation, and WikiText- 103 language modeling. A limitation of our paper is that the regularized nonlocal functional can also help improve the prreec-pres of transformer models. We leave this exciting research idea as future work.
of TRENO has not been addressed. It is interesting to explore if
10

---

Acknowledgments and Disclosure of Funding
RGB acknowledges support from the NSF grants CCF-1111094, IIS- 1838177 and IIS-1730574.ONR grants N00014-18-1257, N00014-20-1-2334 and MURI 900014-20-1-2758 AFOSR grant FA9550- 22-1-0060; and a Vannevar Bush Faculty Fellowship, ONR grant N00014-18-1-2047, TMN acknowl- edges support from his start-up grant the National University of Singapore (Grant Number: 4-0009807-00-00),
References
[1] Eneko Agirre, Daniel Cer, Mona Diab, and Aitor Gonzalez-Agirre SemEval-2012 task 6: A pilot on semantic textual similarity. In #SEM 2012: The First Joinl Conference on Lexical and Computational Semantics - Volume 1: Proceedings of the main conference and the shared task, and Volume 2: Proceedings of the Sixth International Workshop on Semantic Evaluation (Se- mEval 2012), pages 385-393, Montréal, Canada, 7-8 June 2012. Association for Computational Linguistics.
[2] Rami AI-Rfou, Dokook Choe, Noah Constant, Mandy Guo, and Llion Jones. Character-level language modeling with deeper self-attention. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pages 3159-3166, 2019.
[3] Anurag Amna, Mostafa Dehghani, Georg Heigold, Chen Sun, Mario Lučić, and Cordelia Schmid. Vivit: A video vision transformer. 2021 IEEE/CVF International Conference on Computer Vision (ICCV), pages 68266, 2021.
[4] Alexei Baevski and Michael Auli. Adaptive input representations for neural language modeling. In International Conference on Learning Representations, 2019.
[5] Alfonso S. Bandeira, Amit Singer, and Thomas Strohmer. Mathematics of Data Science. 2020.
[6] Tom B Brown, Benjamin Mann, Nick Ryder, Melamie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. arXiv preprint arXiv:2005. 2020.
[7] A. Buades, B. Coll, and J. M. Morel. A review of image denoising algorithms, with a new one. Multiscale Modeling & Simulation, 4(2):490-530, 2005.
[8] Kung-Ching Chang, Kelly Pearson, and Tan Zhang. Perron-frobenius theorem for nonnegative tensors. Communications in Mathematical Sciences, 6(2):507-520. 2008
[9] 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:1 15044- 2021.
[10] Rewon Child, Scott Gray, Alec Radford, and llya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv: 190,:099044, 2019.
[11] Kyunghyun Cho, Bart van Merriënboer, Caglar Gulcehre, Dzmitry Bahdanau, Fethi Bougares, Holger Schwenk, and Yoshua Bengio. Learning phrase representations using RNN encoder- decoder for statistical machine translation. In Proceedings of the 2014 Conference Empirical Methods in Natural Language Processing (EMNLP), pages 1724-1734, Doha, Qatar, October 2014. Association for Computational Linguistics.
[12] Kevin Clark, Urvashi Khandelwal, Omer Levy, and Christopher D. Manning. What does BERT look at? analysis of BERT's attention. In Proceedings of the 2019 ACL Workshop BlackboxNLP: Analyzing and Interpreting Neural Nenworks for NLP, pages 276-286, Florence, Italy, August 2019. Association for Computational Linguistics.
[13] Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc Le, and Ruslan Salakhutdinov. Transformer- XL: Attentive language models beyond fixed-length context. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 2978-2988. Florence, Italy, July 2019. Association for Computational Linguistics.

---

[14] Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Uni- versal transformers. arXiv preprint /
2018.
[15] Jia Deng. Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large- scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pages 248-254 255. leee, 2009.
[16] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv: 1810.04880 2018.
[17] Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Confer- ence of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies. Volume (Long and Short Papers), pages 4171-4186, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics.
[18] Yihe Dong, Jean- -aanttBpeiste Cordonnier, and Andreas Loukas. Attention is not all you need: Pure attention loses rank doubly exponentially with depth. In International Conference on Machine Learning, pages 2793-2803. PMLR, 2021.
[19] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021.
[20] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021.
[21] Leonhard Euler. Institutiones calculi integralis, volume 1. impensis Academiae imperialis scientiarum, 1792.
[22] Tianyu Gao, Xingcheng Yao, and Danqi Chen. SimCSE:
Simple contrastive learning of sentence embeddings. In Proceedings of the 2021 Conference on Empirical Methods Natural Language Processing, pages 6894-6916 10, Online and Punta Cana, Dominican Republic, November 2021. Association for Computational Linguistics.
[23] Guy Gilboa and S. Osher. Nonlocal linear image regularization and supervised segmentation. Multiscale Model. Simul., 6:595-630, 660, 2007.
[24] Guy Gilboa and S. Osher. Nonlocal operators with applications image processing. Multiscale Model. Simul., 7:1005-1028, 2008
[25] Chengyue Gong, Dilin Wang, Meng Li, Vikas Chandra, and Qiang Liu. Vision transformers with patch diversification. arXiv preprint arXiv:2104. 12753,, 2021.
[26] Anmol Gulati, James Qin, Chung-Cheng Chiu, Niki Parmar, Yu Zhang, Jiahui Yu, Wei Han, Shibo Wang, Zhengdong Zhang, Yonghui Wu, et al. Conformer: Convolution-augmented transformer for speech recognition. arXiv preprint 3XXv:2005.081000 2020.
[27] Meng-Hao Guo, Jun-Xiong Cai, Zheng-Ning Liu, Tai-Jiang Mu, Ralph R Martin, and Shi-Min Hu. Pet: Point cloud transformer. Computational Visual Media, 187) 2021.
[28) Dan Hendrycks, Steven Basart, Norman Mu, Saurav Kadavath, Frank Wang, Evan Dorundo, Rahul Desai, Tyler Zhu, Samyak Parajuli, Mike Guo, et al. The many faces of robustness: A critical analysis of out-of-distribution generalization. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 8340-8349, 2021.
[29] Dan Hendryeks and Thomas Dietterich. Benchmarking neural network robustness to common corruptions and perturbations. arXiv preprint 19Xhv:1903.122611 2019.
12

---

[30] Dan Hendrycks, Kevin Zhao, Steven Basart, Jacob Steinhardt, and Dawn Song. Natural adversarial examples. In Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition, pages 15262-15271, 2021.
[31] John Hewitt and Percy Liang. Designing and interpreting probes with control tasks. In Proceed- ings of the 2019 Conference on Empirical Methods Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNN) 2733 2743, Hong Kong, China, November 2019. Association for Computational Linguistics.
pages
[32] 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.
[33] John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ron- neberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Zidek. Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. Nature, 596(7873):583--55) 2021.
[34] Patrick Kahardipraja, Brielen Madureira, and David Schlangen. Towards incremental trans- formers: An empirical analysis of transformer models for incremental NLU. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages 1178-1189, Online and Punta Cana, Dominican Republic, November 2021. Association for Computational Linguistics.
[35] Stefan Kindermann, S. Osher, and Peter W. Jones. Deblurring and denoising of images by nonlocal functionals, Multiscale Model. Simul., 4:1091-1115, 2005.
[36] Dimitrios Kotzias, Misha Denil, Nando de Freitas, and Padhraic Smyth. From group to individual labels using deep features. Proceedings of the 2/th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, 2015.
[37] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
[38] Zhouhan Lin, Minwei Feng, Cícero Nogueira dos Santos, Mo Yu, Bing Xiang, Bowen Zhou, and Yoshua Bengio. A structured self-attentive sentence embedding. CoRR, 2017.
1330,
[39] Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: robustly optimized bert pretraining approach. arXiv preprint arXiv: 1907. / 1692, 2019.
[40] 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 on Computer Vision, pages 10012-10022, 2021.
[41] Ze Liu, Jia Ning, Yue Cao, Yixuan Wei, Zheng Zhang, Stephen Lin, and Han Hu. Video swin transformer. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2022.
[42] Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. In 5th International Conference on Learning Representationss ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net, 2017.
[43] Ankur Parikh, Oscar Täckström, Dipanjan Das, and Jakob Uszkoreit. A decomposable attention model for natural language inference. In Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing, pages 2249-2255, Austin, Texas, November 2016. Association for Computational Linguistics.
[44] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International Conference on Machine Learning, pages 8748-8763 PMLR, 2021.
13

---

[45] Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. Improving language understanding by generative pre-training. OpenAl report, 2018.
[46] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and llya Sutskever. Language models are unsupervised multitask learners. OpenAl blog, 18):9, 2019.
[47] 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 unified text-to-text transformer. Journal of Machine Learning Research, 21 (1(11)): 2020.
[48] Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. SQuAD: 100,000+ ques- tions for machine comprehension of text. In Proceedings of the 2016 Conference on Empirical Methods in Naturul Language Processing, pages 2383-2392. Austin, Texas, November 2016. 6. Association for Computational Linguistics.
[49] Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever. Zero-shot text-to-image generation. In International Conference on Machine Learning, pages 8821-8831. PMLR, 2021.
[50] Alexander Rives, Joshua Meier, Tom Sercu, Siddharth Goyal, Zeming Lin, Jason Liu, Demi Guo, Myle Ott, C Lawrence Zitnick, Jerry Ma, et al. Biological structure and function emerge from scaling unsupervised learning to 250 million protein sequences. Proceedings of the National Academy of Sciences, I 118(15), 2021.
[51] Leonid I. Rudin, Stanley Osher, and Emad Fatemi. Nonlinear total variation based noise removal algorithms. Physica D: Nonlinear Phenomena, 60(1):259-268, 1992.
[52] Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, et al. Imagenet large scale visual recognition challenge. International Journal of Computer Vision, 115(3):211-252, 2015.
[53] Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear transformers are secretly fast weight programmers. In International Conference Machine Learning, pages PMLR, 2021.
[54] Matthias Seeger. Gaussian processes for machine learning. International journal of neural systems, 14(02):69-10) 106, 2004.
[55] Han Shi, JIAUII GAO, Hang Xu, Xiaodan Liang, Zhenguo Li, Lingpeng Kong, Stephen M. S. Lee, and James Kwok. Revisiting over-smoothing in BERT from the perspective of graph. In International Conference on Learning Representations, 2022.
(56) Robin Strudel, Ricardo Garcia, Ivan Laptev, and Cordelia Schmid. Segmenter: Transformer for semantic segmentation. Proceedings of the IEEF/CVF international conference on computer vision, pages 7262-7272, 2021.
[57] Ian Tenney, Dipanjan Das, and Ellie Pavlick. BERT rediscovers the classical NLP pipeline. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 3593-4601 Florence, Italy, July 2019. Association for Computational Linguisties.
[58] Matthew Thorpe, Tan Minh Nguyen, Hedi Xia, Thomas Strohmer, A. Bertozzi, Stanley J. Osher, and Bao Wang. Grand++: Graph neural diffusion with a source term. In International Conference on Learning Representations, 2022.
[59] Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Herve Jegou. Training data-efficient image transformers distillation through attention. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 10347-10357. PMLR, 18-24 Jul 2021.
[60] Hugo Touvron, Matthieu Cord, Alexandre Sablayrolles, Gabriel Synnaeve, and Hervé Jé- gou. Going deeper with image transformers. Proceedings of the IEEF/CVE International Conference Computer Vision (ICCV), pages 32-42, October 2021.
14

---

[61] Yao-Hung Hubert Tsai, Shaojie Bai, Makoto Yamada, Louis-Philippe Morency, and Ruslan Salakhutdinov. Transformer dissection: An unified understanding for transformer's 8 attention via the lens of kernel. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Process- ing (EMNLP-JJCNLP)) pages 4344-4353, Hong Kong, China, November 2019, Association for Computational Linguistics.
[62] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pages 5998-6008, 2017.
[63] Jesse Vig and Yonatan Belinkov. Analyzing the structure of attention in a transformer language model. In Proceedings of the 2019 ACL Workshop BlackboxNLP: Analyzing and Interpret- ing Neural Networks for NLP, pages 63-76, Florence, Italy, August 2019, Association for Computational Linguistics.
[64] Elena Voita, David Talbot, Fedor Moiseev, Rico Sennrich, and Ivan Titov. Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 5797-5808, Florence, Italy, July 2019. Association for Computational Linguistics.
[65] Peihao Wang, Wenqing Zheng, Tianlong Chen, and Zhangyang Wang. Anti-oversmoothing in deep vision transformers via the fourier domain analysis: From theory practice. In International Conference on Learning Representations, 2022.
[66] Zifeng Wang and Jimeng Sun. TransTab: Learning Transferable Tabular Transformers Across Tables. In Advances in Neural Information Processing Systems (NeurIPS 2022), 2022.
[67] Joachim Weickert, Wissenschaftlicher Werdegang, Steven Zucker, Allan Dobbins, Lee Iverson, Benjamin Kimia, and Allen Tannenbaum. Anisotropic diffusion in image processing. 01 1996.
[68] Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty. Mingxing Tan, Glenn Fung, Yin Li, and Vikas Singh. Nyströmformer: A Nyströmm basedd Algorithm for Approximating Aolf-Atttntion. 2021.
[69] Zhilin Yang, Zihang Dai, Yiming Yang. Jaime Carbonell, Ruslan Salakhutdinov, and Quoe Le. Xlnet: Generalized autoregressive pretraining for language understanding. arXiv preprint arXiv:1906.08237, 2019.
[70] Shuai Zhang, Lina Yao, Aixin Sun, and Yi Tay. Deep learning based recommender system: A survey and new perspectives. ACM Computing Surveys (CSUR), 52(1):1-38, 2019.
[71] Hengshuang Zhao, Li Jiang, Jiaya Jia, Philip HS Torr, and Vladlen Koltun. Point transformer. In Proceedings of the IEEF/CVÉ International Conference on Computer Vision, pages 16259- 16268, 2021.
[72] Bolei Zhou, Hang Zhao, Xavier Puig, Sanja Fidler, Adela Barriuso, and Antonio Torralba. Scene parsing through ade20k dataset. In Proceedings of the IEFE conference computer vision and pattern recognition, pages 633-641, 2017.
[73] Bolei Zhou, Hang Zhao, Xavier Puig, Tete Xiao, Sanja Fidler, Adela Barriuso, and Antonio Torralba. Semantic understanding of scenes through the ade20k dataset, 2018
[74] Daquan Zhou, Bingyi Kang, Xiaojie Jin, Linjie Yang, Xiaochen Lian, Zihang Jiang, Qibin Hou, and Jiashi Feng. Deepvit: Towards deeper vision transformer, 2021.
15

---

Supplement to "Mitigating Over-smoothing in Transformers via Regularized Nonlocal Functionals"
Table of Contents
A Additional Details the Experiments in Section 4
16
A.1 Image Classification on Imagenet
17
A.2 Image Segmentation ADK20 dataset
17
A.3 Language Modeling on WikiText-
17
B Technical Proofs
17
B. Proof of Lemma
17
B.2 Proof of Lemma 2
18
B.3 Proof of Proposition
18
C Derivation of Gradient of E as Given in Eqn. 21
19
Results of Symmetric Setting
19
Additional Details on the Empirical Analysis Section
20
E.1 Average Cosine Similarity between Gradient Approximations
20
E,2 Average Value Function
20
E.3
Over-smoothing Analysis
20
Additional Experimental Results
F11
Object classification on Imagenet with DeiT-small baseline
21
F.2
Beyond Softmax-Attention
Additional Empirical Analysis Results
21
G.I
Visualizing Attention Matrices
21
G.2 Head Redundancy between Layers
21
G.3 NeuTRENO Inherently Mitigates Over-smoothing, even without Training the Models
21
G.4 Efficiency Analysis
G.5 Stability and Significance of NeuTRENO
22
G.6 Robustness NeuTRENO
23
G.7 NeuTRENO in Incremental Learning.
23
G.8 Ablatation study the choice of Á
23
G.9 Scalability of NeuTRENO
A
Additional Details on the Experiments in Section
This section provides datasets, models, and training details experiments Section 4. The code to reproduce our experimental results is included in our Supplementary Material submission.
16

---

A.1 Image Classification Imagenet
Datasets and Metrics. The ImageNet dataset [15, 52] comprises 1.28 million training images and 50, 000 validation images, encompassing the classification of 1000 categories. The evaluation metrics used for performance assessment are the top-1 and top-5 accuracies.
Models and Baselines. Our baseline model is the DeiT-tiny model [59], which consists of 12 transformer layers, 3 attention heads per layer, and model dimension of 192. For model setting and setting and configuration, we follow [59]. Their implementation is available at
https//gttbbb.com/ffccoorrrrree
The used for our NeuTRENO method is 0.6.
A.2 Image Segmentation on ADK20 dataset
Datasets and Metrics. The ADE20K dataset is recognized for its inclusion of challenging scenes with fine-grained labels, making it one of the most demanding semantic segmentation datasets. The training set consists of 20,210 images encompassing 150 semantic classes. Additionally, there 2,000 images in the validation set and 3,352 images in the test set. This in task the Single-scale mean Intersection over Union (SS mloU) and the Multi-scale (MS mloU).
Models and baselines. The training configuration and setting for our models are followed by [56]. The baseline model is finetuned with the pretrained DeiT-tiny backbone while segmenter model used the pretrained NeuTRENO DeiT-tiny, with À = 0.6.
A.3 Language Modeling on WikiText-103
Datasets and Metrics. The WikiText-103 dataset consists of articles extracted from Wikipedia and is specifically designed to capture long contextual dependencies. The training set comprises approximately 28, 000 articles, totaling 103 million running words. Each article contains text blocks consisting of approximately 3, 600 words. The validation and test sets contain 218, 000 and 246, 000 running words, respectively, with each set consisting of 60 articles and approximately 268, 000 words. Our experiment follows the standard setting [42, 53], which involves dividing the training data into independent long segments of L words. For evaluation, we employ a batch size of and process the text sequence using sliding window of size L. When computing perplexity (PPL), we consider only the last position, except for the first segment where all positions are evaluated, following the approach in [2, 53].
Models and baselines. For our language modeling implementation, we rely on the publicly available code http://gtuub.cmm/DSSSA////
developed by [53]. In our experiments, we set the dimen- sions of keys, values, and queries to 128, while the training and evaluation context length set to 256. In this experiment, À = 0.4 yields the best performance of NeuTRENO language model.
B Technical Proofs
B.1 Proof of Lemma
For all I....,A N, we have E (B) v().. Assume that (i) (1), then
n
(+))
((-))() v())
N
Σ
v(j)
P =
v(j)B( v())P(B()( v())
N
N
Σ
((B (1) = v(j) = (B) (0) oi))
1
1=1
N
Σ
v(j) Σ AaP(B ( = v()))
Σ
ABE (4) (1) Σ A (1)
u
(+)))
17

---

Thus, by induction, obtain the conclusion of the lemma.
B.2 Proof of Lemma
Since the transition matrix A N is right- its largest eigenvalue is 1 (see Theorem 4.1 in [5]). Also, A is regular positive matrix since its elements are positive. Thus, the Perron- Frebenius theorem [8] implies the existence of unique probability distribution x, which a positive left eigenvector of the transition matrix A associated with its largest eigenvalue 1. In particular, in the case of symmetricity constraint, æ can be chosen as follows
di
d2
dn Σ5=1 d2
Eラ= dj
Σラ= dj
where d; = En exp ((i)) ((((iTT()) kkj)/ DDyk)
is easy to see that
N
N
di
exp
Dat
Aij
am1
2i=1
di
d2
> i=1
exp k(i) k(j)/ /DR
Σ di
d,
Tj..
E1=1 di
As consequence, æ must be the unique stationary distribution the random walk (k())() (1)))) This concludes the proof.
B.3 Proof of Proposition
Recall from the gradient flow in Equ 21, by using the method of Euler discretization, af- ter & update steps starting from the initial u(x, 0) v(x) with adaptive stepsize At
1/ (/Ka(k(z))) k(y, x)) dy and by choosing \/))(r), we obtain the following
(At()) u(x, (k. 1)At(=)) At(:) AAt(z) (k. 1)At(=)) f(x))
f K K(x, y)uly, (k- 1)A((1)) dy xff((z) u(z, (k ((_,((=1))At())
(27)
Jo K (x,
Discretizing uxx, ((x))) and using Monte-Carlo approximation for the integrals in 27 obtain
(x) Σ A (K- (2) - 2 (1)))
(28)
where Aú is computed using the keys and queries as either softmax (())///)/// /D(x) or softmax (())))(())))))))))))((((())))))) k()) DDxb)
Suppose that (K)) defined as Eqn. 28, converges to constant vector á as A: - We have
18

---

u (1)
u (f)
N
N
(1) - Ayu
(8) (1) + (() (1)+ s(i)
1=1
(k) -- (i) Σ Aa) 1 ()) An)
(1)
(29)
1=1
1=1
+
((((())))) - u (1)) + (()(i) - f(())
N
N
(j) 2x Σ (n) (1) - E - =x( + +A - 11)(2())(()
+(ff(i) f())
(ulk+1) (1) - -u (x(+t(())) +
(ulk) (1) -u(² +0
Since u (1) + u. for i 1,2, ....,........................... N, A oc, we have
(ulk) (1) +
(u(k) - (x) +
as k DC. This contradiction since while the LHS of 29 approaches O. its RHS approaches
((f(1) 100--foo f()), which is not in general. Thus, we obtain the conclusion Proposition
C Derivation of Gradient of E as Given in Eqn. 21
Taking the gradient of Elu, f) with respect to u, we obtain
8G DG
8G
V E = VuJ J+ +
(30)
0ur Oug
Ōup!
The partial derivative aG/Ou,.J = 1,2,
D, is defined through its dot product with an arbitrary
function hy E (2)) as follows
0G
d G(u thy))
Ou, hy(x)
dr
A d
dr / (1z(x)- - (x) (h)()))
x / (u,(z) h
Thus, the Frechet derivative of with respect to uj is given by
aG
X(u)(x) fz(1)) ou,
(31)
Substituting the formula for 0G/0u, in Eqn. 31 into Eqn. 30 for V.E(u, f), obtain the following gradient flow
du(x,t) - -V - (v)(((( + +ff(()) u(x)),
(32)
dt
where is dummy time variable and (u) is defined as in 9.
D Results of Symmetric Setting
In this section, we show that NeuTRENO significantly improves the performance of a symmetric transformer baseline, which utilizes symmetric self-attention. We refer to the DeiT with symmetric attention, defined in 14, as Sym-DeiT and the Sym-DeiT combined with our NeuTRENO method as Sym-NeuTRENO DeiT.
Object classification on Imagenet To further illustrate the advantage of our NeuTRENO method, compare Sym-NeuTRENO DeiT with the Sym-DeiT baseline the ImageNet image
19

---

Table 5: Top-1 and Top-5 accuracy (%) of Nem-NeeTRENO DeiT vs. Sym DeiT on the ImageNet classification task. The Sym-NeuTRENO DeïT models significandl outperform the Sym-DeiT DeiT in terms of accuracy, indicating the benefit NeuTRENO method.
Model/Metric
Top-I Acc (%)
Top-5 Acc (%)
Sym-DeiT
71.14
90.54
SyR-NeuTRENO DeiT
72.07
91.22
Table 6: Single-scale (SS) MIoU and multi-scale (MS) MloU of the Sym-NeuTRENO DeiT Sym-DeiT. The Syn-NeuTRENO DeiT model beneficial since they significantly outperform the Sym-DeiT.
Model/Metric
SS MloU
MS MIoU (%)
Sym-Deil
35.18 35.68
36,00
DeiT
36.39
classification task. Our Cemn-NuuRRENT TRENO DeiT outperforms the Sym-Deil baseline, as shown in Table 5. Notably, the Sym-NeuTRENT DeiT achieves higher performance in terms of both top- accuracy and top-5 accuracy than Sym-DeiT baseline. These results further confirm the benefits of our proposed NeuTRENO model.
Image Segmentation on ADE20K dataset We also compare the performance of the Segmenter models [56] using the SyN-NeuTRENNO DeiT backbone with models using the Sym-DeiT backbone on ADE20K image segmentation [72], as shown in Table 6. The results demonstrate the substantial performance improvements achieved by utilizing the SyN-NeuTRENO DeiT backbone compared to the Sym-DeiT backbone in terms of both single-scale (SS) MIoU and multi-scale (MS) MloU metries. This result further validates the advantages of our NeuTRENO models in enhancing image segmentation performance in the symmetric setting.
E Additional Details on the Empirical Analysis in Section 5
In this section, we provide the details for the empirical analysis in Section
E.1 Average Cosine Similarity between Gradient Approximations
To produce the results in Table 4, we derive the approximation for the gradient V. J(u),,,, from Eqn at time f = 0:
VuJ(u) = (u(x, 0) uly 0)) K() (z, - ((v) K(x, ddd..
where K K(x,8) 9) = k(z,1)= + k(y, x). Using Monte-Carlo approximation for the integral and choosing K(I, 1) = exp (((x))))² k(y)/ VDa),, the symmetric approximation of the gradient is derived as ()=(((ii - v(i)) exp (ki)"
(())))(())))))))))))()(())))(( k(j)/ / /Dgk). Otherwise, by choosing K K(x, v) exp(g(x)T ((())/ V DDn). the assymmetric approximation of the gradient is derived as E v(i)) exp A())) yD,x). In this analysis, we take the dot product between the symmetric and asymmetric approximation of the gradient (uu) and average these dot products over positions. We finally report the average cosine similarity over 1000 training data and 1000 test data, as shown in Table
E.2 Average Value of Function
In order to report the average value of function J(u) in Fig. 3, we follow the process of computing (u) for 1000 data points for each transformer block. Subsequently, the average value reported for each layer. This procedure is carried out for both the training and test datasets.
E.3
Over-smoothing Analysis
The average cosine similarity between all pairs of token's 's representations (æ,, xx)) in a sequence is computed
Σ 141
N(N 1)
x
The result then averaged over 1000 randomly chosen test data ImageNet and ADF20K. The result is then reported for each layer, as in Fig. 1.
20

---

Table 7: Top-1 and T accuracy (%) of NeuTRENO DeiT-small vs. DeiT-small on the ImageNet benchmark. The NeuTRENO small significantly outperform the DeiT-small in terms of accuracy. We also compare NeuTRENO Deil-small with DeiT plus FeatScale. a vision transformer model that addresses over-smoothing. showing the advantage NeuTRENO. The accuracies reported [59] for DeiT-small and [65] for DeiT-small plus FcatScale, respectively, are in parentheses.
Model/Metric
Top-1 Acc (%)
Top-5 Ace (%)
DeiT-small
79.97 (79.9)
95.05 (95.0)
Deil-small + FeatScale
79.96 (80.9)
95.06
NeuTRENO DeiT-small
80.68
95.30
Table 8: Accuracy of NeuTRENO vs. Kemel Transformerr on the CIFAR-10 dataset [37]. The NeuTRENO model significantly outperforms the terms of accuracy.
Model/Metric
Accuracy (%)
Kernel Transformer
75.89
NeuTRENO
76.75
F Additional Experimental Results
F1
Object classification on Imagenet with DeiT-small baseline
In this section, show the advantages of method when we further scale up the model by doubling the model dimension and the number of heads compared to that of the Dei'T-tiyy. In particular, the NeuTRENO DeiT-small achieves better results in both Top-1 Accuracy and Top-5 Accuracy, as shown in Table 7. Our method also outperforms DeiT plus FeatScale. Here, we did our best to reproduce the results of DeiT-small plus FeatScale [65]. In Table 7, we include our reproduced results and the results reported in [59] for DeiT--mmal and [65] for DeiT-small plus FeatScale, respectively.
F.2 Beyond Softmax-Attention
We show that NeuTRENO can be combined with other baseline attention mechanisms other than softmax attention. In particular, our NeuTRENO significantly improves transformer-basee models with kernel attention [54, 61], on the CIFAR-10 image classification task [37], as shown in Table This further confirms the benefits of our model. Here, both models share the same configuration regarding training, the model's size, and the model's depth (12 layers).
G Additional Empirical Analysis Results
This section provides extra empirical analysis to further demonstrate the benefits of NeuTRENO models in mitigating over-smoothing
G.1 Visualizing Attention Matrices
Fig. 4 displays the 3-head attention matrices obtained from layer [1, 6, 12] of both the pre-trained NeuTRENO DeiT-tiny and the DeiT-tiny baseline models, using a random sample from the ImageNet dataset.
G.2 Head Redundancy between Layers
NeuTRENO mitigates head redundancy between layers, particularly in the final transformer layers where over-smoone wer-shoo is most pronounced. Fig. 5 shows the average cosine similarity of attention matrices between two successive layers, over 1000 randomly sampled data. The trained NeuTRENO DeiT obtains lower cosine similarity than that of the trained DeiT as the model depth increases.
G.3 NeuTRENO Inherently Mitigates Over-smoothing, even without Training the Models
mandomly-initializn NeuTRENO DeiT-tiny significantly reduces the average cosine similarity DeiT-tiny model, as shown in Fig, 6, on the Imagenet classification task. This observation highlights the ability of our NeuTRENO models in mitigating over-smoothing.
between token representations of 12-layer
G.4 Efficiency Analysis
We report the ratios of the floating-poit operations per second (FLOPs), the inference memory, and the inference real-time running of NeuTRENO DeiT vs. DeiT per sample on the ImageNet dataset, which are 1.00005, 1. 1.00002,, 1.00013, respectively. This indicates that the significant gain in the performance of Neu does not come with the cost of efficiency.
21

---

DeiT
NeuTRENO DeiT
50
50
20
30
Head
100
100
103
100
100
100
150
150
150
130
150
150
100
100
Head 2
100
100
100
150
150
150
100
1.00
100
100
Head 3
50
50
50
50
100
100
100
100
100
100
150
150
130
130
150
100
100
Layer
Layer 6
Layer 12
Layer 1
Layer
Layer 12
Figure 4: Plot of attention matrices attained from layer [1, 6, 12] of both the pretrained DeiT-tiny baseline (Left)
and the NeuTRENO DeiT-tiny (Right) models, for each head, using a random sample from the Imagenet dataset.
DeiT
NeuTRENO DeiT
Cosine Similarity of Attemtion Matrices between layers
0.7
0.6
0.5
0.3
6
8
9
10
11
12
Layer
Figure 5: The average cosine similarity of attention matrices between two successive layers, over 1000 randomly sampled data, of the trained NeuTRENO DeiT and trained DeiT models on the Imagenet classification task.
DeiT
NeuTRENO DeiT
Cosine Similarity between tokens
0.7
0.6
0.5
0.4
Layer
Figure The average cosine similarity between token representations of 12-layer randomly-initialin NeuTRENO DeiT and DeiT models, the Imagenet classification task. Here, 1000 data are randomly sampled for the analysis.
G.5 Stability and Significance of NeuTRENO
To further confirm the stability and significance of NeuTRENO's performance, we provide the
standard deviations from five runs for both the NeuTRENO and baseline models for each experiment (in the main text) in Tables 9, 10, 11.
22

---

Table 9: Means and standard deviations over five runs with different random seeds of models trained on the Imagenet Classification task.
Modrl/Metric
Top-1 Ace (%)
Top-5 Acc (%)
Softmax DeiT-Tiny
72.17 + 0.07
91.02 0.04
NeuTRENO DeñT-Tiny
73.01 +
91.56
NeuTRENO Adaptation
72.63 + 0.07
91.38
DeiT-Tiny FeatScale
72.346 0.06
91.22 0.04
NeuTRENO DeiT-Tiny FeatScale
73.23 +: 0.08
91.73 0.05
Table 10: Means and standard deviations over five runs with different random seeds of models trained on the
ADE20K image segmentation task
Metric/Model
Pretrained Softmax Deit Deit-Tiny
Pretrained NeuTRENO DeiT-Tiny
SS MloU
35.72 0.57
37.24 0.62
MS MIoU
36.68 0.42
38.06 0.54
Table 11: Means and standard deviations over five runs with different random seeds of models trained on the WikiText-103 language model task.
Metrio/Model
Softmax Transformer
NeuTRENO
Valid PPL
33.15 = 0.07
32.60 + 0.08
Test PPL
34.29 1 0.09
33.70 1 0.07
Table 12: Evaluation of NeuTRENO DeiT-Tiny vs. Softmax DeiT-Tiny on the ImageNet-C (mean corruption error mCE), Imagenet-A (Accuracy), and Imagenet-R (Aecuracy) datasets.
Model/Dataset (Metric)
Imagenet-C (mCE)
Imagenet-A (Accuracy %) Imagenet-R R (Accuracy %)
Softmax DeiT-Tiny NeuTRENO-Deii-Tiny
71.6
6.9
32.83
70.1
8.2
33.82
Table 13: Ablation study of different values hyperparameter Ā of NeuTRENO DeiT-Tiny on the ADE20K Image Segmentation task.
Metric/Model
Baseline
À 0.2
0.5
0.6
À 2.0
SS MIoU
35.72
34.94
35.60
37.54
37.38
37.24
36.71
36.37
26.25
MS MIoU
36.68
35.53
36.45
38.62
38.22
38.06
37.82
37.26
27.2
G.6 Robustness of TRETOEN
In addition the standard metrics, we evaluate the robustness of our NeuTRENO model compared to the baseline transformer model, particularly under adversarial examples and for -f-oofiittriuttion generalization. Table 12 demonstrates that NeuTRENO DeiT-Tiny is consistently more robust than the DeiT-Tiny baseline on the Imagenet-C (сommon data corruption and perturbations, such as adding noise and blurring the images) [29], Imagenet-A (adversarial examples) [30], and Imagenet-R (out of distribution generalization) [28] datasets, which are widely used to test the model's robustness.
G.7 NeuTRENO in Incremental Learning
In an incremental learning setting [34], our 8-layer NeuTRENO achieves 1.97% higher accuracy on the sentiment classification task [36] than the 8-layer baseline transformer.
G.8 Ablatation study on the choice
We also conduct an ablation study on the impact of the hyperparameter. In particular, on the ADE20K image segmentation task, we train NeuTRENO with different values. We summarize our results in Table 13. Our findings reveal that within the range of [0.2, 1). NeuTRENO consistently outperforms the softmax baseline. However, when values become small or big (below 0.2 above 1, respectively). NeuTRENO's performance declines.
G.9 Scalability of NeuTRENO
To demonstrate the scalability of our proposed model, we conduct additional experiments show that our NeuTRENO method can effectively mitigate the oversmoothing issue in the BERT-base model. In particular, in Figure 7 (Left), we plot the cosine similarity between token representations across layers of pre-trained BERT-base model [17] on the SQuAD v question answering task [48] and observe
23

---

SQuAD Question Answering
STS-12 Semantic Textual Similarity
Pre-trained
Randomized
Pre-trained
Cosine Similarity
0.8
10
0.8
belween takens
0.9
0.7
0.6
0.5
12
12
Layer
Layer
Layer
Pre-trained SimCSE
Pre-trained BERT
NeuTRENO Pre-trained BERT
NeuTRENO Pre-trained SimCSE
Figure 7: The average cosine similarity between token representations of 12-layer trained (Left) and randomly- initialized (Middle) BERT base and NeuTRENO BERT-base the SQuAD question answering task. We also plot same cosine similarity scores for trained SimCSE and NeuTRENO SimCSE models (Right) on STS-12 semantic textual similarity task. Here, 1000 and 500 are randomly sampled the analysis on SQuAD and STS-12 datasets, respectively.
the presence of the oversmoothing issue as the model gets deeper, causing tokens become identical. We then apply NeETOENO on the same pre-trained BERT model, and without any fine-tuning, we observe significant reduction the cosine similarity between token embeddings in each layer (see Figure 7 (Left)), indicating that NeuTRENO effectively mitigates the oversmoothing problem in BERT. Additionally, our NeuTRENO BERT finetuned on the task yields better accuracy than the finetuned BERT (81.39 exact match score and 88.62 Fl-score vs. 80.77 exact match score and 88. 88.12 F1-score). Moreover, we have conducted the same analysis for a randomized BERT-base model and randomized NeuTRENO BERT-base model and obtained the same encouraging results (see Figure 7 (Middle)). These results further suggest that NeuTRENO helps alleviate the over-smoothing issue in large-scale transformer models.
We also obtain additional results and show that our NeuTRENO SimCSE, after fine-tuned the STS- 12 semantic textual similarity task [1], gains a significant improvement over the baseline SimCSE [22], which is also fine-tuned on the same task (77.32% vs. 75.29% Spearman's correlation. Here, the higher correlation, the better). This additional result further verifies that decreasing the cosine dissimilarity between tokens within trained transformer-but based models leads improved empirical performance.
24


Reviews:
[{"review_id": "HQ4pSjW6Fr", "reviewer": "Reviewer_qvvE", "summary": "This paper analyzes the reason of over-smoothing in Transformer based on Self-attention as a Gradient Descent Step to Minimize a Nonlocal Function and random walk analysis. The paper then proposes a novel regularizer that penalizes the norm of the difference between the output tokens from self-attention and the input tokens to preserve the fidelity of  the tokens.  Experimental results on ImageNet classification, image segmentation and LM tasks show that the proposed approach achieves better performance than the baseline vanilla Transformer.", "questions": "Please address the points under weaknesses.\n\nAlso, there are some presentation errors in the paper. For example,\n(1)\tThe word “functional” is an adjective, so using it as a noun in the paper is ungrammatical, for example, in the title and phrases such as “minimize a functional which…” “Minimizing the resulting regularized energy functional, …” throughout the paper.\n\n(2)\tSome math symbols have not been defined, for example, D_x , D_{qk} in Section 1.1.\n\n(3)\tLine 89 defines the weights k(x,y), but the definition is kind of vague. Line 117 provides a form of K(x,y) :=k(x,y)+k(y,x), which needs to be explained clearer about this choice of K(x,y).\n", "limitations": "This is N/A for this work.", "rating": 4, "confidence": 5, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "Strengths:\n\n(1)\tThe paper provides interesting aspects to analyze the reason of over-smoothing in Transformers based on Self-attention as a Gradient Descent Step to Minimize a Nonlocal Function and random walk analysis.\n \n(2)\tThe theoretical proofs are detailed and clear.\n", "weaknesses": "Weaknesses:\n\n(1)\tAlthough the paper provides detailed theoretical proofs of the nonlocal variational denoising framework for self-attention and provides an explanation for the over-smoothing issue in transformer-based models, the empirical evaluations missed some important details and analyses.  In Table 1, the configuration of NeuTRENO Adaptation is not explained. \n\n(2)\tExcept the results in Table 2 on image segmentation, the gains from the proposed approach in Table 1 ImageNet classification and WikiText-103 LM are all quite small. In Table 1, the gains from the proposed approach over the baseline are 0.84 on Top-1 acc and 0.54 on Top-5 acc. In Table 3, the gains from the proposed approach over baseline are 0.55 on validation PPL and 0.59 on test PPL. It is not clear how stable the proposed model is on these datasets, since there is no reporting of standard deviations from multiple runs with different random seeds, so it is not clear whether the gains are by chance. Also, it is not clear whether these gains are statistically significant. \n\n\n(3)\tThere have been prior works analyzing the uniformity and alignment problem of BERT on sentence representations and proposed postprocessing solutions, for example, the flow https://arxiv.org/pdf/2011.05864.pdf and BERT whitening method https://arxiv.org/abs/2103.15316 and supervised method such as SimCSE https://arxiv.org/abs/2104.08821 and enhancement over SimCSE on sentence representations.\n\nThese works are not compared to or cited in the paper.\n\n(4)\tAn important point is missing. Based on the results from models addressing uniformity in BERT and improving sentence representations, these models alleviate uniformity but cause degradations in transfer learning capability, e.g., SimCSE. This issue has not been discussed in the paper since the paper did not investigate the proposed approach in pre-training and fine-tuning paradigm, to investigate the impact of the proposed approach on transferability.\n\n(5)\tThere are some presentation problems. Some math symbols and functions are not defined or clearly defined, as listed under Questions.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful review and valuable feedback. Below we address your concerns.\n\n**Q1**. The empirical evaluations missed some important details and analyses. In Table 1, the configuration of NeuTRENO Adaptation is not explained.\n\n**Reply:**\nThank you for your comment. In our NeuTRENO Adaptation experiment, we integrate the NeuTRENO architecture into a pre-trained DeiT-Tiny model with $\\tilde\\lambda = 0.6$ (determines how much the regularization impacts the solution, a small value makes the solution smoother and vice versa). This combined model is then finetuned for 100 additional epochs. The experiment and its configuration are detailed from lines 250 to 254 (main text) and lines 520 to 523 (Appendix A). For comparisons with related work (e.g., FeatScale method), please refer to lines 234 to 236 (main text). For evaluation metrics, we provide thorough descriptions in Appendix A, under \"Datasets and Metrics\" for each task.\n\nAs for the analysis, we provide numerous empirical analyses of our models in Section 5 (main text) and in Appendix G. In particular, Fig. 1 shows that the cosine similarity between token representations across layers in a trained NeuTRENO is significantly reduced compared to the baseline model, demonstrating its advantage in mitigating oversmoothing. Furthermore, Fig. 6 presents similar results for a randomly initialized NeuTRENO, i.e., before training. Additionally, Fig. 5 (Appendix G2) demonstrates NeuTRENO's advantage in reducing head redundancy compared to the baseline. To address the efficiency of NeuTRENO, an efficiency analysis is provided in Appendix G4. Finally, to verify our theoretical results in Theorem 1, in Figure 3, we show that softmax attention indeed minimizes the functional $J(\\bf{u})$ in Eqn. 6 in our paper.\n\n**Q2**. Question on the stability and significance of NeuTRENO's performance.\n\n**Reply:** Thank you for your question. Please kindly refer to points 1, 2, and 3 in the general response for our answers to your questions.\n\n**Q3**. Reviewer's concerns that prior work on the uniformity and alignment problem of BERT on sentence representations are overlooked. \n\n**Reply:** Thanks for bringing up these prior works. However, we respectfully disagree that our paper disregards these works. \n\nWe would like to emphasize that our work focuses on developing a variational denoising framework to understand the self-attention of transformers as a gradient descent approximation of a functional. From this new finding, we explain the oversmoothing issue of transformers as a result of self-attention minimizing a functional. **It is important to note that we primarily aim at explaining and resolving the collapse of representations, i.e., oversmoothing, at the token level (tokens representations become identical) rather than at the sentence level (the model gives high similarity scores for semantically different sentences)**. Although the merits of our method can be generalized to enhance sentence-level embeddings, it is out of the scope of our work, and we will explore it in future work.\n\nIn addition, our new empirical study shows that SimCSE, proposed to enhance sentence representations via contrastive learning, still suffers from oversmoothing at the token level. Fig. 7 in the attached PDF shows that the cosine similarity between token representations for each layer of the SimCSE, which is trained on the STS-12 dataset, increases as the model depth grows. In contrast, our NeuTRENO is shown to mitigate the issue. When integrating NeuTRENO with the pre-trained SimCSE, without any fine-tuning, we observe the cosine similarity scores reduce significantly (Fig. 7)\n\n**References**\n\n[1] Gao, T., Yao, X., and Chen, D. \"SimCSE: Simple Contrastive Learning of Sentence Embeddings\". EMNLP, 2021.\n\n**Q4**. The paper did not investigate the transferability of the proposed approach.\n\n**Reply:** In our paper, we empirically study the transferability of NeuTRENO and summarize the results in Table 2 (Section 4). In particular, after pretraining DeiT and NeuTRENO DeiT on the ImageNet classification task, we finetune both models for the ADE20k image segmentation task. The results demonstrate that the finetuned NeuTRENO DeiT yields significantly better performance than the finetuned DeiT. Further details regarding this transfer learning experiment are provided in Section 4, from lines 238 to 244, and also in Appendix A2. Note that we provide the accuracies of the pre-trained models on the ImageNet classification task in Table 1.\n\n**Q5**. Some math symbols have not been or are vaguely defined, for example, D_x , D_{qk}, k(x,y), and K(x, y).\n\n**Reply:** Thank you for your comment. In our paper, $D_{x}$ represents the feature dimension of the token $x_{i}$, $i=1,\\dots,N$, i.e. $X \\in \\mathbb{R}^{N \\times D_{x}}$. $D_{qk}$ represents the feature dimension of $q_{i}$ and $k_{i}$, $i=1,\\dots,N$, i.e., $Q, K \\in \\mathbb{R}^{N \\times D_{qk}}$.\n\nThe function $k(x, y)$ captures the similarity between signal values at positions x and y. For example, $k(x, y)$ can be a radial basis function (RBF) kernel $e^{-\\gamma||\\Phi(x) - \\Phi(y)||^2}$ for a feature function $\\Phi$ [1]. $K(x, y)$ is indeed just a shorthand notation for $k(x, y) + k(y, x)$. We use this notation to simplify the expression.\n\n**References**\n\n[1] Bishop, Christopher M. \"Pattern Recognition and Machine Learning\". New York. Springer, 2006.\n\n **Q6**. The word “functional” is an adjective, so using it as a noun in the paper is ungrammatical.\n\n**Reply:** In the context of variational calculus, a functional is a mathematical concept that maps from a function space to the real numbers. A more formal definition can be found in [1].\n\n**References**\n\n[1] A. N. Kolmogorov and S. V. Fomin, “Elements of the Theory of Functions and Functional Analysis”, Graylock Press, 1957.\n\n-----\nWe hope we have cleared your concerns about our work. We would appreciate it if we could get your further feedback at your earliest convenience.\n"}, {"review_id": "Dy5kCUUuq2", "reviewer": "Reviewer_YNWk", "summary": "This paper studies the oversmoothing problem in transformers. Roughly speaking, it was observed that embeddings start to converge when the network is deeper. The authors built a model to explain this phenomenon. Roughly speaking, they relate having deeper architecture with making progress towards minimizing a function. Then they built some regularizers out of this intuition.  ", "questions": "It was proven that an update of a self-attention is equivalent to taking a gradient over a function. But will the function have many local optimal so that even we have the main theorem, it does not quite say all the embeddings will eventually be the same? Or maybe the function is a convex function. ", "limitations": "NA", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "\nI think these analysis and regularization tools are useful addition to the exploding area of transformers. \n", "weaknesses": "\nI feel I was not able to properly parse some math texdt. My most unsure part is whether the amount of hand-waiving is too much even under today’s standard. For example, the authors feel very comfortable in exchanging integration operator and summation. It is true in general the approximation errors are usually inconsequential. But it is still different from pretending that they are the same in theorem statements/proofs. Also, I dont understand how K(x, y), k(x, y) and boldk(x) and boldk(y) are related. First k(x, y) are defined, then there will always exist boldk? It also looks that sometimes k/K are treated as kernels whereas sometimes it is treated as a kernel function. Probability functions and kernel functions have different properties? \n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful review and valuable feedback. Below we address your concerns.\n\n**Q1**. I feel I was not able to properly parse some math texdt. My most unsure part is whether the amount of hand-waiving is too much even under today’s standard. For example, the authors feel very comfortable in exchanging integration operator and summation. It is true in general the approximation errors are usually inconsequential. But it is still different from pretending that they are the same in theorem statements/proofs.\n\n**Reply:** In the main text, we employ the Monte Carlo approximation method to estimate integrals in Eqns. 11, 15, 23, and 27. We choose the Monte Carlo integration method because it is an unbiased estimator. Additionally,  when the sample size $N$ used in Monte Carlo is sufficiently large, the estimation converges to the correct value. This tends to hold true for transformer models, which are often applied on long sequences with many token samples.\n\n**Q2**. Also, I don't understand how K(x, y), k(x, y) and boldk(x) and boldk(y) are related. First k(x, y) are defined, then there will always exist boldk? It also looks that sometimes k/K are treated as kernels whereas sometimes it is treated as a kernel function. Probability functions and kernel functions have different properties?\n\n**Reply:** The function $k(x, y)$ captures the similarity between signal values at positions x and y. For example, $k(x, y)$ can be a radial basis function (RBF) kernel $e^{-\\gamma||\\Phi(x) - \\Phi(y)||^2}$ for a feature function $\\Phi$ [1]. $K(x, y)$ is indeed just a shorthand notation for $k(x, y) + k(y, x)$. We use this notation to simplify the expression.\n\n$\\bf{k}$$(\\cdot)$ is a vector-valued function which maps positions ($x$ or $y$) to feature vectors. In transformer models, ${\\bf k}(i)$, $i=1,\\dots,N$, is the key vector of token $i$-th in self-attention. For images, $\\bf{k}$$(x)$ indicates the feature vector at pixel $x$.\n\n**References**\n\n[1] Bishop, Christopher M. \"Pattern Recognition and Machine Learning\". New York. Springer, 2006.\n\n**Q3. It was proven that an update of a self-attention is equivalent to taking a gradient over a function. But will the function have many local optimal so that even we have the main theorem, it does not quite say all the embeddings will eventually be the same? Or maybe the function is a convex function.**\n\n**Reply:** Yes, the functional $E$ is a convex function of $\\bf{u}$ since $J(\\bf{u})$ is a sum of quadratic functionals of $u_j$ for $j = 1,\\dots, D$, and $G(\\bf{u}, \\bf{f})$ is a quadratic functional of $\\bf{u}$. Hence, $E(\\bf{u})$, the sum of two convex functional, is itself a convex functional.\n\n-----\nWe hope we have cleared your concerns about our work. We would appreciate it if we could get your further feedback at your earliest convenience."}, {"review_id": "tmy5lhfY8H", "reviewer": "Reviewer_h3ak", "summary": "This paper analyses the over-smoothing problem of transformer architecture by showing that self-attention layers minimize a functional that causes over-smooth. To address this problem, the authors introduce a regularizer that penalizes the norm of the difference between the smooth output tokens and input tokens. In the experimental section, the authors demonstrate that their proposed solution, NeuTRENO, outperforms baseline models on visual and language-related downstream tasks.", "questions": "please refer to weaknesses", "limitations": "The generalizability and scalability of the proposed solution to larger transformer models have not been investigated. ", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The issue of over-smoothing in transformers is a fascinating topic. While previous research has explored this problem, this paper offers a unique perspective by examining self-attention as a gradient descent step from a variational standpoint. This novel approach sheds new light on the over-smoothing problem and contributes to a deeper understanding of the issue within the research community.\n\nThe effectiveness of the proposed solution, NeoTRENO, in mitigating the problem of over-smoothing is demonstrated both theoretically and empirically. The experiments show NeoTRENO's effectiveness in both vision and language transformers. ", "weaknesses": "Current experiments are mostly conducted on small transformer models( e.g., DeiT-tiny).  Given that the transformer architecture is crucial for large pre-trained language models, it remains unclear whether the proposed solution can effectively alleviate the over-smoothing problem when the model size is increased, and whether can be combined with large pre-trained language models (e.g.,  BERT). \n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful review and valuable feedback. Below we address your concerns.\n\n**Q1.** Current experiments are mostly conducted on small transformer models( e.g., DeiT-tiny). Given that the transformer architecture is crucial for large pre-trained language models, it remains unclear whether the proposed solution can effectively alleviate the over-smoothing problem when the model size is increased, and whether can be combined with large pre-trained language models (e.g., BERT).\n\n**Reply:** Thanks for your comments. The DeiT-tiny model used in our experiments has 5M parameters and is trained on the large-scale Imagenet dataset for classification task. To examine the scalability of our NeuTRENO model, we have applied NeuTRENO to the DeiT-small, which is 4 times larger than the DeiT-tiny and has 22M parameters. We summarize the results in Table 7 in Appendix F.1. The improvements achieved by our NeuTRENO DeiT-small over the baseline model demonstrate the scalability of NeuTRENO when applied to very large transformer architectures. \n\nFurthermore, we have conducted additional experiments to show that our NeuTRENO method can effectively mitigate the oversmoothing issue in the BERT-base model. In particular, in Figure 7 (Left) in the attached PDF, we plot the cosine similarity between token representations across layers of a pre-trained BERT-base model [1] on the SQuAD v1.1 question answering task [2] and observe the presence of the oversmoothing issue as the model gets deeper, causing tokens to become identical. We then apply NeuTRENO on the same pre-trained BERT model, and without any fine-tuning, we observe a significant reduction in the cosine similarity between token embeddings in each layer (see Figure 7 (Left) in the attached PDF), indicating that NeuTRENO effectively mitigates the oversmoothing problem in BERT. Moreover, we have conducted the same analysis for a randomized BERT-base model and a randomized NeuTRENO BERT-base model and obtained the same encouraging results (see Figure 7 (Middle) in the attached PDF). These results further suggest that NeuTRENO helps alleviate the oversmoothing issue in large-scale transformer models. \n\n**References**\n\n[1] Jacob Devlin, et al. “BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding”. NAACL, 2019.\n\n[2] Pranav Rajpurkar, et al. “SQuAD: 100,000+ Questions for Machine Comprehension of Text”. EMNLP, 2016.\n\n-----\nWe hope we have cleared your concerns about our work. We would appreciate it if we could get your further feedback at your earliest convenience.\n"}, {"review_id": "WPstSyHuha", "reviewer": "Reviewer_qcqa", "summary": "This work shows that self-attention layers in transformers minimize a funcitonal which promotes smoothneess, thereby casuing token uniformity. The work also proposes a novel regularizer to preserve fidelity of the tokens. The work empirically shows that NeuTRENO outperforms baseline transformers in reducing the over-smoothing of token representation. ", "questions": "- Although the authors motivates this work from the point of over-smoothing, the paper lacks discussion on the actual cause, impact and consequence of over-smoothing. For example, from Figure.1, we observe that NeuTRENO has a smaller cosine similarity comparing to naive DeiT. One question would be, what should be the desired cosine sim? Is smaller cosine sim siginificantly better than say sim over 0.8? ", "limitations": "Please see the previous two sections", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- I like the flow of the paper. The authors dive straight into the key issue of the work, without much retrospect on the previous known matters.\n- The paper rewrites the self-attention as a Gradient Descent Step to minimizie a nonlocal functional.  \n- The paper shows that as k (number of layers) increases, the model is more likely to suffer from over-smoothing. \n- I think viewing each self-attention layer in the transformer as an impliciit gradient ascent in very interesting. It echos with some literature in the past, providing more insights on the interpretability of deep neural networks. ", "weaknesses": "- Some use of the wording could be improved, for example in line 220s, the authors like to use words like \"significantly\", \"addresses\" etc. In my opinion, some of the experimental results are not strong enough to support author's claim. It could create misleading for the audiences. \n- Experiments looks very interesting. I think the paper might benefit more by discussing broader impact of \"alleviating over-smoothness\". Say for example. would be help in OOD setting? incremental setting? or other interesting applications. In my opinion, 1% increase on standard metric should not be the best selling point for this interesting work.   ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful review and valuable feedback. Below we address your concerns.\n\n**Q1**. Some use of the wording could be improved, for example in line 220s, the authors like to use words like \"significantly\", \"addresses\" etc. In my opinion, some of the experimental results are not strong enough to support author's claim. It could create misleading for the audiences. \n\n**Reply:** Thank you for your feedback. We have removed those words in our revised manuscript. We have also included the standard deviations from 5 runs for each experiment (in the main text) in Tables 9, 10, and 11 in the attached PDF. The standard deviations for both the NeuTRENO and baseline models in these tasks are small compared to the gains achieved by the NeuTRENO method. This observation suggests that NeuTRENO’s  improvements over the baseline are significant and not by chance. \n\n**Q2**. Experiments looks very interesting. I think the paper might benefit more by discussing broader impact of \"alleviating over-smoothness\". Say for example. would be help in OOD setting? incremental setting? or other interesting applications. In my opinion, 1% increase on standard metric should not be the best selling point for this interesting work.\n\n**Reply:** Thank you for your valuable suggestions.  We have evaluated the robustness of our NeuTRENO model compared to the baseline transformer model, particularly under adversarial examples and for out-of-distribution generalization. Table 12 in the attached PDF demonstrates that NeuTRENO DeiT-Tiny is consistently more robust than the DeiT-Tiny baseline on the Imagenet-C ( corruption and perturbations data, such as adding noise and blurring the images) [1], Imagenet-A (adversarial examples) [2], and Imagenet-R (for out of distribution generalization) [3], which are widely used to test the model’s robustness. \n\nFurthermore, in an incremental learning setting [4], our 8-layer NeuTRENO achieves 1.97% higher accuracy on the sentiment classification task [5] than the 8-layer baseline transformer. We hope that these additional analyses and results address your concerns and provide further evidence for the stability, significance, and robustness of our proposed NeuTRENO approach. \n\n**References**\n\n[1] Hendrycks, Dan, et al. \"Benchmarking neural network robustness to common corruptions and perturbations.\" arXiv, 2019.\n\n[2] Hendrycks, Dan, et al. \"Natural adversarial examples.\" CVPR, 2021.\n\n[3] Hendrycks, Dan, et al. \"The many faces of robustness: A critical analysis of out-of-distribution generalization.\" ICCV, 2021.\n\n[4] Kahardipraja, et al. \"Towards incremental transformers: An empirical analysis of transformer models for incremental NLU.\"  EMNLP, 2021.\n\n[5] Dimitrios Kotzias, et al. \"From group to individual labels using deep features\". KDD, 2015\n\n**Q3**. Although the authors motivates this work from the point of over-smoothing, the paper lacks discussion on the actual cause, impact and consequence of over-smoothing. For example, from Figure.1, we observe that NeuTRENO has a smaller cosine similarity comparing to naive DeiT. One question would be, what should be the desired cosine sim? Is smaller cosine sim siginificantly better than say sim over 0.8? \n\n**Reply:** Thanks for your comments. We believe there is a misunderstanding of the contributions of our paper. Please allow us to clear this misunderstanding by clarifying that our work first focuses on developing a variational denoising framework to understand the self-attention of transformers as a gradient descent approximation of a functional. Using this new finding, we provide an explanation for the oversmoothing issue in transformers as a result of self-attention minimizing a functional, leading to the smoothing effect on the input sequence, analogous to a diffusion process  (see Remark 1 in our main text). Then, we  rigorously prove this observation on the existence of oversmoothing in transformers using a random walk analysis in Section 2.2. Thus, in our paper, we not only discuss the actual cause of oversmoothing but also prove its existence theoretically.\n\nRegarding the impact and consequence of oversmoothing, oversmoothing leads to a loss of diversity among token representations, hindering the model's capacity to capture diverse features. As a result, the model's performance can be adversely affected. Additionally, oversmoothing often constrains the ability of transformer models to be scaled up in depth. For instance, more-layer transformer models might underperform compared to those with fewer layers due to this issue [1]. \n\nThe cosine similarity between token representations of each layer is an indicator of the oversmoothing issue. It shows how similar a token is to the other, on average, at a layer. When cosine similarity at a layer is high, it indicates that the representations of tokens are too similar, diminishing the expressive power of the model. Consequently, smaller cosine similarity scores are preferable.\n\n**References**\n\n[1] Daquan Zhou, et al. “Deepvit: Towards deeper vision transformer”. arXiv, 2021.\n\n-----\nWe hope we have cleared your concerns about our work. We would appreciate it if we could get your further feedback at your earliest convenience.\n"}, {"review_id": "OtSkbDFQtd", "reviewer": "Reviewer_k2Mg", "summary": "Recent research works have revealed that the over-smoothing issue, a prevalent challenge in Graph Neural Networks, similarly plagues Transformers. Contrary to expectations, the performance of a Transformer model does not invariably improve with increased depth of the self-attention layers. In fact, a deeply layered Transformer may not necessarily outperform a shallow one. This paper elucidates that the over-smoothing phenomenon occurs due to a consistent decrease in the non-local function of weighted differences in each self-attention layer as depth increases. Drawing on this finding, the paper puts forward a solution featuring modified skip connections. While similar solutions have been proposed in prior research, this paper sets itself apart by providing an in-depth theoretical analysis from a gradient flow perspective, which previous works have not addressed. Furthermore, the experimental results corroborate the effectiveness of the proposed solution in mitigating over-smoothing in Transformers.", "questions": "1. This paper introduces a concept of symmetric self-attention, i.e., the query matrix is equal to the key matrix. The experimental results demonstrate an unexpectedly high performance from a Transformer predicated on symmetric self-attentions. This performance seems to contravene the conventional understanding that self-attention should capture the bi-directional correlation among tokens. Can the authors provide a logical elucidation to account for these particular experimental outcomes?", "limitations": "N/A", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "1. The theoretical investigation conducted in this paper is rigor and in-depth.", "weaknesses": "1. The theoretical explication put forth in this paper applies accurately only to Transformers with single-head self-attention mechanisms. However, given that the prevailing Transformers utilize multi-head self-attention mechanisms and feed-forward networks, it becomes prudent to critically evaluate the efficacy of the resolution proposed herein.\n2. This paper lacks a comprehensive set of comparative experiments to analyze the influence of diverse initial values of the parameter $\\widetilde{\\lambda}$ on the model's final performance.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "\nThank you for your thoughtful review and valuable feedback. Below we address your concerns.\n\n**Q1**. The theoretical explication put forth in this paper applies accurately only to Transformers with single-head self-attention mechanisms. However, given that the prevailing Transformers utilize multi-head self-attention mechanisms and feed-forward networks, it becomes prudent to critically evaluate the efficacy of the resolution proposed herein.\n\n**Reply:** Thanks for your comments. The variational denoising framework we develop for the single-head self-attention mechanism can be extended to derive the multi-head self-attention mechanism. Please allow us to explain this derivation below.\n\nLet us consider one of the common implementations of the multi-head self-attention, in which the input sequence $X \\in R^{N \\times D_{x}}$ is truncated into $H$ pieces $X_d \\in R^{N \\times D_{x}/H}$, for $d = 1, ..., H$. Here, we assume that $D_{x}$ is divisible by $H$.  In order to derive the multi-head attention, we apply our variational denoising framework proposed in Section 2 and Theorem 1 in our paper on each of these $X_d$ signals. Note that the vector $\\bf{q}$$(i)$, $\\bf{k}$$(i)$, and $\\bf{v}$$(i)$, for $i = 1,...,N$,  in Theorem 1 are replaced by the vectors $\\bf{q}$$_d(i)$, $\\bf{k}$$_d(i)$, and $\\bf{v}$$_d(i)$ which are linear transformations of the input vector $X_d(i)$, $i = 1,...,N$. The feed-forward network then combines the output sequences of all of these $H$ heads. \n\nGiven our derivation above, it follows that the oversmoothing issue still persists in transformers with multi-head self-attention, and the proposed NeuTRENO still helps mitigate it. Indeed, our experiments in the paper are conducted with multi-head self-attention and validate the advantage of NeuTRENO over the baseline multi-head self-attention in resolving oversmoothing.\n\n**Q2**. This paper lacks a comprehensive set of comparative experiments to analyze the influence of diverse initial values of the parameter on the model's final performance. \n\n**Reply:** Thank you for your suggestion. We have conducted an ablation study on the impact of the hyperparameter $\\tilde\\lambda$. In particular, on the ADE20K image segmentation task, we train NeuTRENO with different $\\tilde\\lambda$ values. We summarize our results in Table 13 in the attached PDF. Our findings reveals that within the range of [0.2, 1], NeuTRENO consistently outperforms the softmax baseline. However, when $\\tilde\\lambda$ values become small or big (below 0.2 or above 1, respectively), NeuTRENO's performance declines. \n\n**Q3**. This paper introduces a concept of symmetric self-attention, i.e., the query matrix is equal to the key matrix. The experimental results demonstrate an unexpectedly high performance from a Transformer predicated on symmetric self-attentions. This performance seems to contravene the conventional understanding that self-attention should capture the bi-directional correlation among tokens. Can the authors provide a logical elucidation to account for these particular experimental outcomes? \n\n**Reply:** Thank you for your question. First, we would like to clarify that transformer models can be bi-directional or uni-directional, depending on the task or model design. For instance, BERT [1], a bi-directional model, allows tokens to attend to both preceding and succeeding tokens. On the other hand, models like GPT [2], trained for autoregressive language modeling, is a uni-directional model as tokens can only attend to their previous counterparts. Second, regarding the symmetric attention, its competitive empirical performances have been demonstrated in existing works [3, 4].\n\n**References**\n\n[1] Jacob Devlin, et al.. “BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding”. NAACL, 2019.\n\n[2] Tom Brown, et al. \"Language models are few-shot learners.\" NEURIPS, 2020.\n\n[3] Yao-Hung Hubert Tsai,et al. “Transformer Dissection: An Unified Understanding for Transformer’s Attention via the Lens of Kernel”. EMNLP-IJCNLP, 2019.\n\n[4] Wenlong Chen, et al. \"Calibrating Transformers via Sparse Gaussian Processes.\" ICLR, 2022.\n\n-----\nWe hope we have cleared your concerns about our work. We would appreciate it if we could get your further feedback at your earliest convenience."}]

Author rebuttal:
Dear AC and reviewers,

Thanks for your thoughtful reviews and valuable comments, which have helped us improve the paper significantly. We are encouraged by the endorsements that: 1) Our paper's variational denoising framework for self-attention is novel and interesting  (Reviewer qvvE, h3ak, qcqa); 2) The derivation of the oversmoothing in transformers from our variational denoising framework is useful (Reviewer YNWk) and contributes to the understanding of the (oversmoothing) issue (Reviewer h3ak); 3) The theory in our paper is rigorous, in-depth (Reviewer k2Mg), clear, and detailed (Reviewer qvvE); 4) The effectiveness of our NeuTRENO model is demonstrated both theoretically and empirically (Reviewer h3ak). 

Among the concerns of the reviewers is the significance of our NeuTRENO's advantages over the baseline transformer model. We address this concern here.
 
1. We provide the standard deviations from 5 runs for each experiment (in the main text) in Tables 9, 10, and 11 in the attached PDF. The standard deviations for both the NeuTRENO and baseline models in these tasks are small compared to the gains achieved by the NeuTRENO method. This observation suggests that NeuTRENO’s improvements over the baseline are significant and not by chance.

2. In addition to the standard metrics, we have evaluated the robustness of our NeuTRENO model compared to the baseline transformer model, particularly under adversarial examples and for out-of-distribution generalization. Table 12 in the attached PDF demonstrates that NeuTRENO DeiT-Tiny is consistently more robust than the DeiT-Tiny baseline on the Imagenet-C (common data corruption and perturbations, such as adding noise and blurring the images) [1], Imagenet-A (adversarial examples) [2], and Imagenet-R (out of distribution generalization) [3] datasets, which are widely used to test the model’s robustness. 

3. Furthermore, in an incremental learning setting [4], our 8-layer NeuTRENO achieves 1.97% higher accuracy on the sentiment classification task [5] than the 8-layer baseline transformer.

4. To demonstrate the scalability of our proposed model, we have conducted additional experiments to show that our NeuTRENO method can effectively mitigate the oversmoothing issue in the BERT-base model. In particular, in Figure 7 (Left) in the attached PDF, we plot the cosine similarity between token representations across layers of a pre-trained BERT-base model [6] on the SQuAD v1.1 question answering task [7] and observe the presence of the oversmoothing issue as the model gets deeper, causing tokens to become identical. We then apply NeuTRENO on the same pre-trained BERT model, and without any fine-tuning, we observe a significant reduction in the cosine similarity between token embeddings in each layer (see Figure 7 (Left) in the attached PDF), indicating that NeuTRENO effectively mitigates the oversmoothing problem in BERT. Moreover, we have conducted the same analysis for a randomized BERT-base model and a randomized NeuTRENO BERT-base model and obtained the same encouraging results (see Figure 7 (Middle) in the attached PDF). These results further suggest that NeuTRENO helps alleviate the oversmoothing issue in large-scale transformer models.

We would also like to summarize the main contributions of our paper here. Our work first focuses on **developing a variational denoising framework to understand the self-attention of transformers as a gradient descent approximation of a functional**. Using this new finding, we **provide an explanation for the oversmoothing issue in transformers** as a result of self-attention minimizing a functional, leading to the smoothing effect on the input sequence, analogous to a diffusion process  (see Remark 1 in our main text). Then, we **rigorously prove this observation on the existence of oversmoothing in transformers using a random walk analysis** in Section 2.2. Thus, in our paper, we not only discuss the cause of oversmoothing but also prove its existence theoretically. Finally, we **propose the Neural Transformer with a Regularized Nonlocal Functional (NeuTRENO), a novel class of transformers designed to mitigate oversmoothing**. NeuTRENO is derived by optimizing a regularized nonlocal functional, which includes an additional convex fidelity term. This fidelity term penalizes the norm of the difference between the smooth output tokens from self-attention and the input tokens, thereby reducing the over-smoothing effect. It is important to note that in our analysis, we primarily aim at explaining and resolving the collapse of representations, i.e., oversmoothing, at the token level (tokens representations become identical) rather than at the sentence level (the models give high similarity scores for semantically different sentences).

**References**

[1] Hendrycks, Dan, et al. "Benchmarking neural network robustness to common corruptions and perturbations." arXiv, 2019.

[2] Hendrycks, Dan, et al. "Natural adversarial examples." CVPR, 2021.

[3] Hendrycks, Dan, et al. "The many faces of robustness: A critical analysis of out-of-distribution generalization." ICCV, 2021.

[4] Kahardipraja, et al. "Towards incremental transformers: An empirical analysis of transformer models for incremental NLU."  EMNLP, 2021.

[5] Dimitrios Kotzias, et al. "From group to individual labels using deep features". KDD, 2015

[6] Jacob Devlin, et. al. “BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding”. NAACL, 2019.

[7] Pranav Rajpurkar, et al. “SQuAD: 100,000+ Questions for Machine Comprehension of Text”. EMNLP, 2016.

-----

We are glad to answer any further questions you have on our submission.

Meta-review: This paper studies how self-attention minimizes a functional which promotes smoothness, which had an effect on deeper transformer networks leading to token uniformity. The paper also proposes a novel regularizer term (NeuTRENO), based on the proportion of the difference between the first and current layers' self-attention activations. The experimental validation shows improvements (with NeuTRENO) compared to benchmark/standard very competitive transformers on ImageNet, ADE20k, WikiText103. Overall, all the reviewers who participated in the rebuttal agree that this paper is NeurIPS material. The rebuttal discussion/answers improved the clarity of the paper, which we hope to see in the camera ready.
Decision: Accept (poster)