Title: Mitigating Over-smoothing in Transformers via Regularized Nonlocal Functionals
Venue: NeurIPS 2023
Accepted: yes
Keywords: transformers

self-attention

total variation

nonlocal functionals

over-smoothing
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

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.

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 is available [45, 46, 16, 69, 39]. At the heart of transformers lies the self-attention mechanism, which computes a weighted average of token representations within a 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 a given input sequence X := [ x (1) , · · · , x ( N )] ⊤ ∈ R N × D x of N feature vectors, self-attention transforms X into the output sequence H in the following two steps:

DeiT

NeuTRENO DeiT

Layer

Layer

between tokens

ImageNet Classification

ADE20K Segmentation

0.9

0.9

0.8

0.8

0.7

0.7

0.6

0.6

0.5

0.5

0.4

0.4

0.3

0.3

0.2

0.2

2

4

6

8

10

12

2

4

6

8

10

12

Step 1. The input sequence X is projected into the query matrix Q , the key matrix K , and the value matrix V via three linear transformations

where W Q , W K ∈ R D qk × D x , and W V ∈ R D × D x are the weight matrices. We denote Q := [ q (1) , . . . , q ( N )] ⊤ , K := [ k (1) , . . . , k ( N )] ⊤ , and V := [ v (1) , . . . , v ( N )] ⊤ , where the vectors q ( i ) , k ( i ) , and v ( i ) , for i = 1 , . . . , N are the query, key, and value vectors, respectively.

Step 2. The output sequence U := [ u (1) , . . . , u ( N )] ⊤ ∈ R N × D qk is then computed as follows

where the softmax function is applied to each row of the matrix QK ⊤ / √ D qk . The matrix A := softmax ( QK ⊤ √ D qk ) ∈ R N × N and its component a ij 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 Eqn. (2) to compute the output vector u ( i ) is given by

The self-attention computed by Eqn. (2) and (3) is refered as softmax attention. In our work, we refer to a transformer that uses softmax attention as a softmax transformer.

1.2 Over-smoothing in Transformers

Despite their remarkable success, deep transformer-based models have been observed to suffer from the over-smoothing issue, in which all token representations become identical when more layers are added to the models [55, 65, 18]. This over-smoothing phenomenon, also known as the 'token uniformity' problem, significantly limits the representation capacity of transformers. To illustrate this phenomenon, we examine the average cosine similarity between pairs of token representations across different layers in a softmax transformer trained for the Imagenet object classification and ADK20 image segmentation tasks [73]. As depicted in Fig. 1, in both tasks, this cosine similarity between tokens increases as the models become deeper. Particularly, in the last two layers, the cosine similarity scores are approximately 0.9, indicating a high degree of similarity among tokens.

1.3 Contribution

We develop a nonlocal variational denoising framework for self-attention, providing insights into the over-smoothing phenomenon in transformers. In particular, by viewing self-attention as a gradient descent step toward minimizing a nonlocal functional that penalizes high-frequency noise in the signal, we uncover the diffusive nature of self-attention, which explains the over-smoothing issue of transformers. Motivated by this understanding, we propose the Neural Transformer with a Regularized

Nonlocal Functional (NeuTRENO), a novel class of transformers designed to mitigate over-smoothing. 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. Our contribution is three-fold.

We develop a nonlocal variational denoising framework for self-attention and shed light on the over-smoothing issue that hampers the representation capacity of transformers.

We develop NeuTRENO, a novel class of transformers that are capable of alleviating the over-smoothing issue.

We theoretically prove that transformers with softmax self-attention are prone to oversmoothing while NeuTRENO can avoid this issue.

We empirically demonstrate the benefits of NeuTRENO on various large-scale applications, including the ImageNet object classification, ADE20K image segmentation, and WikiText-103 language modeling tasks.

Organization : We organize our paper as follows: in Section 2, we develop a nonlocal variational denoising framework for self-attention and provide an explanation for the over-smoothing issue in transformer-based models. In section 3, we propose NeuTRENO, and present a theoretical result that guarantees NeuTRENO's capability of mitigating over-smoothing. In Section 4, we empirically validate the benefits of NeuTRENO. We discuss the related work in Section 6. Finally, we conclude our main contributions and remarks. Further results, details, and proofs are provided in the Appendix.

2 A Nonlocal Variational Denoising Framework for Self-attention

We first consider the output matrix U := [ u (1) , · · · , u ( N )] ⊤ ∈ R N × D in self-attention as given by Eqn. 2 in Section 1.1. Let Ω ⊂ R , x ∈ Ω , and u ( x ) := [ u 1 ( x ) , . . . , u D ( x )] T be a real vector-valued function, u : Ω → R D , u ∈ L 2 (Ω) . The output matrix U in self-attention discretizes the function u ( x ) on a 1-D grid. In the context of signal/image denoising, U can be considered as the desired clean signal , and u ( x ) is its corresponding intensity function denoting the signal values at the position x ∈ Ω . We further let the observed intensity function f ( x ) denote the values of the observed noisy signal at x ∈ Ω , f : Ω → R D , f ∈ L 2 (Ω) . For example, f ( x ) can be given as

where n is the additive noise. We wish to reconstruct u ( x ) from f ( x ) . Following the variational denoising method proposed in [23] and [24], the denoised image u ( x ) can be obtained by minimizing the following regularized functional with respect to u :

Here, J ( u ) = 1 2 ∫ Ω × Ω ∥ u ( x ) -u ( y ) ∥ 2 2 k ( x, y ) dxdy is a nonlocal functional of weighted differences. The weights k ( x, y ) represent the affinity between signal values at positions x and y . For example, for images, k ( x, y ) captures the proximity between pixels x and y in the image. J ( u ) works as a regularizer. Minimizing J ( u ) promotes the smoothness of u and penalizes high-frequency noise in the signal. Adding the convex fidelity term G ( u , f ) = λ 2 ∫ Ω ∥ u ( x ) -f ( x ) ∥ 2 2 dx to the functional J ( u ) allows the denoised signal u ( x ) to preserve relevant information in the observed noisy signal f ( x ) . The regularized functional E ( u , f ) can be considered as an energy functional.

2.1 Self-attention as a 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 J ( u ) in the energy functional E ( u , f ) . We expand J ( u ) as follows

The gradient of J with respect to u is then given by

The partial derivative ∂J/∂u j , j = 1 , 2 , . . . , D , is defined through its dot product with an arbitrary function h j ∈ L 2 (Ω) as follows

Applying a change of variables ( x, y ) → ( y, x ) to the second term of the above integral, we have

Thus, the Frechet derivative of J with respect to u j is given by

Substituting the formula for ∂J/∂u j in Eqn. 8 into Eqn. 7 for ∇ u J ( u )( x ) , we obtain the following gradient flow

where t is the time variable we introduce to capture the dynamics of u when gradient descent is applied to minimize J ( u ) . Let v ( x ) := [ v 1 ( x ) , . . . , v D ( x )] T be a real vector-valued function, v : Ω → R D , v ∈ L 2 (Ω) . We discretize v ( x ) on a 1-D grid to attain the value vectors v (1) , . . . , v ( N ) ∈ R D , which form the value matrix V := [ v (1) , · · · , v ( N )] ⊤ ∈ R N × D in self-attention as defined in Eqn. 2. We initialize u at t = 0 with v ( x ) , i.e., u ( x, 0) = v ( x ) .

Self-attention is an Euler Discretization of the Gradient Flow Given in 9. We discretize the gradient flow in Eqn. 9 using the Euler method [21] with step size ∆ t ( x ) = 1 / ∫ Ω ( k ( x, y ) + k ( y, x ) ) dy and obtain the following update

Here, K ( x, y ) := k ( x, y ) + k ( y, x ) is a symmetric kernel and u ( y, 0) = v ( y ) since u is initialized at t = 0 with v as aforementioned. Let k ( x ) := [ k 1 ( x ) , . . . , k D qk ( x )] T be a real vector-valued function, k : Ω → R D qk , k ∈ L 2 (Ω) . Similar to u ( x ) and v ( x ) , we can discretize k ( x ) on a 1-D grid to attain the key vectors k (1) , . . . , k ( N ) ∈ R D qk , which form the key matrix K := [ k (1) , · · · , k ( N )] ⊤ ∈ R N × D qk in self-attention as defined in Eqn. 2. We choose K ( x, y ) = exp ( k ( x ) T k ( y ) / √ D qk ) and rewrite Eqn. 10 as follows

Estimating the integrals in Eqn. 11 via Monte-Carlo approximation using the key vectors k (1) , . . . , k ( N ) ∈ R D qk and and value vectors v (1) , . . . , v ( N ) ∈ R D , we obtain

Discretizing u ( x, ∆ t ( x )) on another 1-D grid, we attain

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. W Q = W K where W Q and W K are the linear projections that map the input sequence X into Q and K as given in Eqn. 1. This symmetry of the attention scores is desirable in some image processing tasks due to 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 k ( i ) by the query vectors q ( i ) , i = 1 , . . . , N , to obtain the exact formula of self-attention given by Eqn. 3. The following theorem summarizes our results:

Theorem 1 (Self-attention as a Gradient Descent Step to Minimize a Nonlocal Functional) . Given the nonlocal functional J ( u ) = 1 2 ∫ Ω × Ω ∥ u ( x ) -u ( y ) ∥ 2 2 k ( x, y ) dxdy of a vector-valued function u : Ω → R D , u ∈ L 2 (Ω) , and let K ( x, y ) := k ( x, y ) + k ( y, x ) = exp ( k ( x ) T k ( y ) / √ D qk ) , where k : Ω → R D qk , k ∈ L 2 (Ω) . Then, taking a gradient descent step on u at time t = 0 , where u ( x, 0) = v ( x ) , with an adaptive step size ∆ t ( x ) := 1 ∫ Ω ( k ( x, y ) + k ( y, x ) ) dy to minimize J is equivalent to updating u via a symmetric self-attention

which results in

Here, u ( n ) , v ( n ) , and u ( n ) , n = 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 ( i ) , i = 1 , . . . , N , in Eqn. 14, we obtain the exact formula of self-attention

Remark 1. In Eqn. 9, the change in u at position x is proportional to the sum of differences between u ( x ) and u at other position in the domain Ω . In particular, when u ( x ) is smaller or larger than the values at other positions, it will increase or decrease, respectively. This is analogous to a diffusion process in which particles or substances move from high-concentration to low-concentration regions. It has been proved that a diffusion process converges to a 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 a collective behavior of numerous random walks performed by individual particles or 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, we show the connection

between the evolution of u and a random walk. By adopting a random walk perspective on graph neural network [58], we demonstrate that u ( x ) under the dynamics given in Eqn 9 suffers from over-smoothing.

Recall from the gradient flow in Eqn 9, by using Euler method discretization, after k update steps starting from the initial u ( x, 0) = v ( x ) , with adaptive stepsize ∆ t = 1 / ∫ Ω ( k ( x, y ) + k ( y, x ) ) dy , we obtain the following

Discretizing u ( x, k ∆ t ( x )) and using Monte-Carlo approximation for the integrals in 15 , we attain

where A ij is computed using the keys and queries as either softmax ( k ( i ) ⊤ k ( j ) / √ D qk ) or softmax ( q ( i ) ⊤ k ( j ) / √ D qk ) . Let { B ( k ) ( i ) } k ∈ K be a random walk on { v ( i ) } N i =1 as defined:

where B ( k ) ( n ) is the random value of a k -step walk, starts at node n , and v ( n ) is the initial value at node n , respectively, for n = 1 , 2 , . . . , 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 i = 1 , 2 , . . . , N , after k update steps as in 16, with initial value u (0) ( i ) = v ( i ) , u ( i ) ( k ) equals to the expected value of the k -step walk, starting at node i :

Lemma 1. Let u ( k ) ( i ) defined in 16 and { B ( k ) ( i ) } k ∈ K is the random walk defined by 17. Then

We next present the Lemma 2 which is necessary to show the convergence of u ( k ) ( i ) .

Lemma 2. The random walk B ( k ) ( i ) in 17 with the transition matrix A either be A ij = softmax ( k ( i ) ⊤ k ( j ) / √ D qk ) or A ij = softmax ( q ( i ) ⊤ k ( j ) / √ D qk ) , has a unique stationary distribution π = [ π 1 , π 2 , . . . , π N ] such that π i := P ( B ( k ) ( j ) = v ( i )) , for i, j = 1 , 2 , . . . , N , ∑ N i =1 π i = 1 , and π T = π T A .

If A ij = softmax ( k ( i ) ⊤ k ( j ) / √ D qk ) , the stationary distribution is:

where d i = ∑ N j =1 exp ( k ( i ) ⊤ k ( j ) / √ D qk ) , k (1) , k (2) , . . . , k ( N ) are the key vectos.

In general, π i can be found by finding the left eigenvector of A corresponding to the dominant eigenvalue 1.

From the Lemma 1 and Lemma 2, we see that, for all i = 1 , 2 , . . . , N ,

as k →∞ . This shows that when k increases, u ( i ) ( k ) converges to a constant vector, indicating that u ( x ) , under the dynamic in 9, suffers from over-smoothing.

𝒙

𝟎

𝛼 𝒗

𝟎

-𝒗

𝟏

𝛼 𝒗

𝟎

-𝒗

𝟐

+

𝒖

𝟏

Self-Attention

FFN

𝒙

𝟏

+

𝒖

𝟑

…

-

+

𝒒

𝟎

𝒌

𝟎

𝒗

𝟎

𝒒

𝟏

𝒌

𝟏

𝒗

𝟏

𝒖

𝟐

Self-Attention

FFN

𝒙

𝟐

𝒒

𝟐

𝒌

𝟐

𝒗

𝟐

Self-Attention

-

+

3 NeuTRENO: Mitigating the Over-smoothing in Transformers via Minimizing a 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 in the diffusive characteristics of u and causes the over-smoothing phenomenon in transformers, as proved in Section 2.2. Fortunately, our objective is not to minimize J ( u ) 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 ) = λ 2 ∫ Ω ∥ u ( x ) -f ( x ) ∥ 2 2 dx . This fidelity term aims to preserve the relevant information in the observed noisy signal f ( x ) by penalizing solution u ( x ) that deviates significantly from f ( x ) , 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 (NeuTRENO) by minimizing the regularized functional E ( u , f ) . We then provide a theoretical result to prove that NeuTRENO does not suffer from over-smoothing. Recall from Eqn. 5 that E ( u , f ) is given by

Following a similar derivation as in Section 2.1 (see Appendix C for the detailed derivation), we obtain the following gradient flow when minimizing E ( u , f ) using gradient descent

NeuTRENO-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 Euler method [21] with step size ∆ t ( x ) = 1 / ∫ Ω ( k ( x, y ) + k ( y, x ) ) dy and initializing u at t = 0 with v ( x ) , i.e., u ( x, 0) = v ( x ) . Choosing λ = ˜ λ/ ∆ t ( x ) , we obtain the following update

We choose the observed noisy signal f ( x ) = v 0 ( x ) where v 0 ( x ) is v ( x ) at the first layer in the transformer model. The update in Eqn. 22 becomes

Applying the Monte-Carlo method to approximate the integrals in Eqn. 23 and discretizing u ( x, ∆ t ( x )) , v ( x ) , and v 0 ( x ) on a 1-D grid, we attain the following new formula for calculating symmetric self-attention:

Its corresponding asymmetric self-attention is obtained by replacing the key vectors k ( i ) with the query vectors q ( i ) , i = 1 , . . . , N , and given by

Leveraging Eqn. 25, we define the Neural Transformer with a Regularized Nonlocal Functional (NeuTRENO) as follows.

Definition 1 (Neural Transformer with a Regularized Nonlocal Functional (NeuTRENO)) . Given a set of key and value vectors { k ℓ ( j ) , v ℓ ( j ) } N j =1 in each layer ℓ , ℓ = 1 , . . . , L , for each query vector q ℓ ( i ) , i = 1 , . . . , N , in the same layer, the self-attention unit at layer ℓ in a Neural Transformer with a Regularized Nonlocal Functional (NeuTRENO) computes the corresponding output vector u ℓ ( i ) of the query q ℓ ( i ) by the following attention formula:

where v 0 (1) , . . . v 0 ( N ) ∈ R D 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 1 indicates that our NeuTRENO mitigates the over-smoothing issue, suggesting the benefit of our method. The proof for Proposition 1 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-103 [42]. Our aim to show: (i) NeuTRENO significantly outperforms the transformer baseline with softmax-attention defined in 2 across various tasks; moreover, NeuTRENO surpass FeatScale, a vision transformer that addresses over-smoothing, combining NeuTRENO with FeatScale is beneficial; (ii) the advantages of incorporating our proposed method with pre-trained models. We also demonstrate the benefits of our NeuTRENO in the symmetry setting and we point to Appendix D 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.

Object classification on ImageNet. To demonstrate the advantage of our NeuTRENO method, we compare it with the DeiT baseline [59] on the ImageNet image classification task. Our NeuTRENO DeiT surpasses the DeiT baseline, as shown in Table 1. Notably, our NeuTRENO DeiT achieves significantly higher performance in terms of both Top-1 Accuracy and Top-5 Accuracy. We also compare our method with FeatScale [65], a vision transformer model addressing over-smoothing (see Table 1). Our NeuTRENO significantly outperforms FeatScale, and combining NeuTRENO with FeatScale leads to substantial improvements. These results confirm the benefits of our model.

Image Segmentation on ADE20K dataset. To further validate the advantages of our proposed methods, we compare the performance of the Segmenter models [56] using the NeuTRENO DeiT

Layer

Layer

Train

Test

1010

1010

109

109

(n)r

108

108

107

107

106

106

2

4

6

8

10

12

2

4

6

8

10

12

Language Model on WikiText-103. In addition to computer vision tasks, we also evaluate the effectiveness of our model on a large-scale natural language processing application, specifically language modeling on WikiText-103. Our NeuTRENO 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.

Combine with pre-trained models. Furthermore, our proposed method is also beneficial to 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 1, showing that combined with our method improves both the Top-1 and Top-5 accuracies of the pre-trained models.

5 Empirical Analysis

Applying Softmax-Attention Reduces the functional J ( 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 ∇ u J ( u ) using symmetric or asymmetric kernel K ( x, y ) for both the trained Sym-DeiT (using symmetric self-attention 14) and DeiT models, closed 1, as shown in Table 4. This suggests that reversing the direction of the asymmetric approximation effectively decreases J ( u ) . Considering that softmax attention takes steps in this reversed direction numerically, its application leads to a 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 are in Appendix E

Over-smoothing Analysis. We empirically illustrate the effectiveness of NeuTRENOs in mitigating the over-smoothing problem in transformers. Fig. 1 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

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] to computer vision [65, 18]. In vision tasks, [74] observes that the performance of the vision transformer (ViT [20]) quickly saturates as 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-smoothing.

In addition, [55] observes the phenomenon in BERT [16], a 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-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 a variational denoising framework to understand the self-attention of transformers as a gradient descent approximation of a 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.

Nonlocal Functionals for Image Processing. Total variation [51] is well-known as an imagedenoising technique. It denoises a noisy image by solving a constraint optimization problem. The method is also related to PDE-flow-based image-denoising 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 a 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, we 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) to alleviate the over-smoothing. We empirically verify the benefits of NeuTRENO with a 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 privacy-preserving of NeuTRENO has not been addressed. It is interesting to explore if regularized nonlocal functional can also help improve the privacy-preserving of transformer models. We leave this exciting research idea as future work.

Acknowledgments and Disclosure of Funding

RGBacknowledges support from the NSF grants CCF-1911094, IIS-1838177, and IIS-1730574;ONR grants N00014-18-12571, N00014-20-1-2534, and MURI N00014-20-1-2787; AFOSR grant FA955022-1-0060; and a Vannevar Bush Faculty Fellowship, ONR grant N00014-18-1-2047. TMN acknowledges support from his start-up grant at the National University of Singapore (Grant Number: A-0009807-00-00).

References

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 Joint 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 (SemEval 2012) , pages 385-393, Montréal, Canada, 7-8 June 2012. Association for Computational Linguistics.

Rami Al-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.

Anurag Arnab, Mostafa Dehghani, Georg Heigold, Chen Sun, Mario Luˇ ci´ c, and Cordelia Schmid. Vivit: A video vision transformer. In 2021 IEEE/CVF International Conference on Computer Vision (ICCV) , pages 6816-6826, 2021.

Alexei Baevski and Michael Auli. Adaptive input representations for neural language modeling. In International Conference on Learning Representations , 2019.

Alfonso S. Bandeira, Amit Singer, and Thomas Strohmer. Mathematics of Data Science . 2020.

Tom B Brown, Benjamin Mann, Nick Ryder, Melanie 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.14165 , 2020.

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.

Kung-Ching Chang, Kelly Pearson, and Tan Zhang. Perron-frobenius theorem for nonnegative tensors. Communications in Mathematical Sciences , 6(2):507-520, 2008.

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:15084-15097, 2021.

Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509 , 2019.

Kyunghyun Cho, Bart van Merriënboer, Caglar Gulcehre, Dzmitry Bahdanau, Fethi Bougares, Holger Schwenk, and Yoshua Bengio. Learning phrase representations using RNN encoderdecoder for statistical machine translation. In Proceedings of the 2014 Conference on Empirical Methods in Natural Language Processing (EMNLP) , pages 1724-1734, Doha, Qatar, October 2014. Association for Computational Linguistics.

Kevin Clark, Urvashi Khandelwal, Omer Levy, and Christopher D. Manning. What does BERT look at? an analysis of BERT's attention. In Proceedings of the 2019 ACL Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP , pages 276-286, Florence, Italy, August 2019. Association for Computational Linguistics.

Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc Le, and Ruslan Salakhutdinov. Transformer-XL: Attentive language models beyond a 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.

Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Universal transformers. arXiv preprint arXiv:1807.03819 , 2018.

Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A largescale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition , pages 248-255. Ieee, 2009.

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805 , 2018.

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 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers) , pages 4171-4186, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics.

Yihe Dong, Jean-Baptiste 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.

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.

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.

Leonhard Euler. Institutiones calculi integralis , volume 1. impensis Academiae imperialis scientiarum, 1792.

Tianyu Gao, Xingcheng Yao, and Danqi Chen. SimCSE: Simple contrastive learning of sentence embeddings. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing , pages 6894-6910, Online and Punta Cana, Dominican Republic, November 2021. Association for Computational Linguistics.

Guy Gilboa and S. Osher. Nonlocal linear image regularization and supervised segmentation. Multiscale Model. Simul. , 6:595-630, 2007.

Guy Gilboa and S. Osher. Nonlocal operators with applications to image processing. Multiscale Model. Simul. , 7:1005-1028, 2008.

Chengyue Gong, Dilin Wang, Meng Li, Vikas Chandra, and Qiang Liu. Vision transformers with patch diversification. arXiv preprint arXiv:2104.12753 , 2021.

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 arXiv:2005.08100 , 2020.

Meng-Hao Guo, Jun-Xiong Cai, Zheng-Ning Liu, Tai-Jiang Mu, Ralph R Martin, and Shi-Min Hu. Pct: Point cloud transformer. Computational Visual Media , 7(2):187-199, 2021.

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.

Dan Hendrycks and Thomas Dietterich. Benchmarking neural network robustness to common corruptions and perturbations. arXiv preprint arXiv:1903.12261 , 2019.

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 Pattern Recognition , pages 15262-15271, 2021.

John Hewitt and Percy Liang. Designing and interpreting probes with control tasks. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP) , pages 2733-2743, Hong Kong, China, November 2019. Association for Computational Linguistics.

Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. Advances in neural information processing systems , 34:12731286, 2021.

John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Žídek, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. Nature , 596(7873):583-589, 2021.

Patrick Kahardipraja, Brielen Madureira, and David Schlangen. Towards incremental transformers: 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.

Stefan Kindermann, S. Osher, and Peter W. Jones. Deblurring and denoising of images by nonlocal functionals. Multiscale Model. Simul. , 4:1091-1115, 2005.

Dimitrios Kotzias, Misha Denil, Nando de Freitas, and Padhraic Smyth. From group to individual labels using deep features. Proceedings of the 21th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining , 2015.

Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.

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 , abs/1703.03130, 2017.

Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv:1907.11692 , 2019.

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.

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.

Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings . OpenReview.net, 2017.

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.

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.

Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. Improving language understanding by generative pre-training. OpenAI report , 2018.

Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. OpenAI blog , 1(8):9, 2019.

Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of Machine Learning Research , 21(140):1-67, 2020.

Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. SQuAD: 100,000+ questions for machine comprehension of text. In Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing , pages 2383-2392, Austin, Texas, November 2016. Association for Computational Linguistics.

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.

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 , 118(15), 2021.

Leonid I. Rudin, Stanley Osher, and Emad Fatemi. Nonlinear total variation based noise removal algorithms. Physica D: Nonlinear Phenomena , 60(1):259-268, 1992.

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.

Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear transformers are secretly fast weight programmers. In International Conference on Machine Learning , pages 9355-9366. PMLR, 2021.

Matthias Seeger. Gaussian processes for machine learning. International journal of neural systems , 14(02):69-106, 2004.

Han Shi, JIAHUI 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.

Robin Strudel, Ricardo Garcia, Ivan Laptev, and Cordelia Schmid. Segmenter: Transformer for semantic segmentation. In Proceedings of the IEEE/CVF international conference on computer vision , pages 7262-7272, 2021.

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 4593-4601, Florence, Italy, July 2019. Association for Computational Linguistics.

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.

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.

Hugo Touvron, Matthieu Cord, Alexandre Sablayrolles, Gabriel Synnaeve, and Hervé Jégou. Going deeper with image transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV) , pages 32-42, October 2021.

Yao-Hung Hubert Tsai, Shaojie Bai, Makoto Yamada, Louis-Philippe Morency, and Ruslan Salakhutdinov. Transformer dissection: An unified understanding for transformer's 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 Processing (EMNLP-IJCNLP) , pages 4344-4353, Hong Kong, China, November 2019. Association for Computational Linguistics.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems , pages 5998-6008, 2017.

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 Interpreting Neural Networks for NLP , pages 63-76, Florence, Italy, August 2019. Association for Computational Linguistics.

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.

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

Zifeng Wang and Jimeng Sun. TransTab: Learning Transferable Tabular Transformers Across Tables. In Advances in Neural Information Processing Systems (NeurIPS 2022) , 2022.

Joachim Weickert, Wissenschaftlicher Werdegang, Steven Zucker, Allan Dobbins, Lee Iverson, Benjamin Kimia, and Allen Tannenbaum. Anisotropic diffusion in image processing. 01 1996.

Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty, Mingxing Tan, Glenn Fung, Yin Li, and Vikas Singh. Nyströmformer: A Nyström-based Algorithm for Approximating Self-Attention. 2021.

Zhilin Yang, Zihang Dai, Yiming Yang, Jaime Carbonell, Ruslan Salakhutdinov, and Quoc V Le. Xlnet: Generalized autoregressive pretraining for language understanding. arXiv preprint arXiv:1906.08237 , 2019.

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.

Hengshuang Zhao, Li Jiang, Jiaya Jia, Philip HS Torr, and Vladlen Koltun. Point transformer. In Proceedings of the IEEE/CVF International Conference on Computer Vision , pages 1625916268, 2021.

Bolei Zhou, Hang Zhao, Xavier Puig, Sanja Fidler, Adela Barriuso, and Antonio Torralba. Scene parsing through ade20k dataset. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 633-641, 2017.

Bolei Zhou, Hang Zhao, Xavier Puig, Tete Xiao, Sanja Fidler, Adela Barriuso, and Antonio Torralba. Semantic understanding of scenes through the ade20k dataset, 2018.

Daquan Zhou, Bingyi Kang, Xiaojie Jin, Linjie Yang, Xiaochen Lian, Zihang Jiang, Qibin Hou, and Jiashi Feng. Deepvit: Towards deeper vision transformer, 2021.

Supplement to 'Mitigating Over-smoothing in Transformers via Regularized Nonlocal Functionals'

Table of Contents

A Additional Details on the Experiments in Section 4

This section provides datasets, models, and training details for experiments in Section 4. The code to reproduce our experimental results is included in our Supplementary Material submission.

A.1 Image Classification on 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 a model dimension of 192. For model setting and setting and configuration, we follow [59]. Their implementation is available at https://github.com/facebookresearch/deit. 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 are 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 mIoU) and the Multi-scale (MS mIoU).

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 our 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 1 and process the text sequence using a 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 https://github.com/IDSIA/lmtool-fwp developed by [53]. In our experiments, we set the dimensions of keys, values, and queries to 128, while the training and evaluation context length is set to 256. In this experiment, ˜ λ = 0 . 4 yields the best performance of NeuTRENO language model.

B Technical Proofs

B.1 Proof of Lemma 1

For all i = 1 , . . . , N , we have E [ B (0) ( i )] = v ( i ) . Assume that E [ B ( k ) ( i )] = u ( k ) ( i ) , then

Thus, by induction, we obtain the conclusion of the lemma.

B.2 Proof of Lemma 2

Since the transition matrix A ∈ R N × N is right-stochastic, its largest eigenvalue is 1 (see Theorem 4.1 in [5]). Also, A is a regular positive matrix since its elements are positive. Thus, the PerronFrebenius theorem [8] implies the existence of a unique probability distribution π , which is 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

where d i = ∑ N j =1 exp ( k ( i ) ⊤ k ( j ) / √ D qk ) . It is easy to see that

As a consequence, π must be the unique stationary distribution of the random walk { B ( k ) ( i ) } k ∈ K . This concludes the proof.

B.3 Proof of Proposition 1

Recall from the gradient flow in Eqn 21, by using the method of Euler discretization, after k update steps starting from the initial u ( x, 0) = v ( x ) with adaptive stepsize ∆ t = 1 / ∫ Ω ( k ( x, y ) + k ( y, x ) ) dy and by choosing λ = ˜ λ/ ∆ t ( x ) , we obtain the following

Discretizing u ( x, k ∆ t ( x )) and using Monte-Carlo approximation for the integrals in 27 , we obtain

where A ij is computed using the keys and queries as either softmax ( k ( i ) ⊤ k ( j ) / √ D qk ) or softmax ( q ( i ) ⊤ k ( j ) / √ D qk ) .

Suppose that u ( k ) ( i ) , defined as Eqn. 28, converges to a constant vector ¯ u as k → ∞ . We have

as k → ∞ . This is a contradiction since while the LHS of 29 approaches 0 , its RHS approaches ˜ λ ( f ( i ) -f ( j )) , which is not 0 in general. Thus, we obtain the conclusion of Proposition 1.

C Derivation of Gradient of E as Given in Eqn. 21

Taking the gradient of E ( u , f ) with respect to u , we obtain

The partial derivative ∂G/∂u j , j = 1 , 2 , . . . , D , is defined through its dot product with an arbitrary function h j ∈ L 2 (Ω) as follows

Thus, the Frechet derivative of F with respect to u j is given by

Substituting the formula for ∂G/∂u j in Eqn. 31 into Eqn. 30 for ∇ u E ( u , f ) , we obtain the following gradient flow

where t is a dummy time variable and -∇ u J ( 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, we compare Sym-NeuTRENO DeiT with the Sym-DeiT baseline on the ImageNet image

classification task. Our Sym-NeuTRENO DeiT outperforms the Sym-DeiT baseline, as shown in Table 5. Notably, the Sym-NeuTRENO DeiT achieves higher performance in terms of both top-1 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 Sym-NeuTRENO 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 Sym-NeuTRENO DeiT backbone compared to the Sym-DeiT backbone in terms of both single-scale (SS) MIoU and multi-scale (MS) MIoU metrics. 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 5.

E.1 Average Cosine Similarity between Gradient Approximations

To produce the results in Table 4, we derive the approximation for the gradient ∇ u J ( u ) , from Eqn 9, at time t = 0 :

where K ( x, y ) := k ( x, y ) + k ( y, x ) . Using Monte-Carlo approximation for the integral and choosing K ( x, y ) = exp ( k ( x ) T k ( y ) / √ D qk ) , the symmetric approximation of the gradient is derived as ∑ N j =1 ( v ( i ) -v ( j ) ) exp ( k ( i ) T k ( j ) / √ D qk ) . Otherwise, by choosing K ( x, y ) = exp ( q ( x ) T k ( y ) / √ D qk ) , the assymmetric approximation of the gradient is derived as ∑ N j =1 ( v ( i ) -v ( j ) ) exp ( q ( i ) T k ( j ) / √ D qk ) . In this analysis, we take the dot product between the symmetric and asymmetric approximation of the gradient ∇ u J ( u ) 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 4.

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 J ( u ) for 1000 data points for each transformer block. Subsequently, the average value is 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 representations ( x i , x j ) in a sequence is computed as

̸

The result is then averaged over 1000 randomly chosen test data in ImageNet and ADE20K. The result is then reported for each layer, as in Fig. 1.

F Additional Experimental Results

F.1 Object classification on Imagenet with DeiT-small baseline

In this section, we show the advantages of our method when we further scale up the model by doubling the model dimension and the number of heads compared to that of the DeiT-tiny. 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-small 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-based models with kernel attention [54, 61], on the CIFAR-10 image classification task [37], as shown in Table 8. 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-smoothing 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

Randomly-initialized NeuTRENO DeiT-tiny significantly reduces the average cosine similarity between token representations of 12-layer randomly-initialized 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.

G.4 Efficiency Analysis

We report the ratios of the floating-point 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 . 000002 , 1 . 00013 , respectively. This indicates that the significant gain in the performance of NeuTRENO does not come with the cost of efficiency.

Layer 1

Layer 6

Layer 12

Layer 1

Layer 6

Layer 12

Head 1

Head 2

Head 3

DeiT

NeuTRENO DeiT

0

0

0

0

0

50

50

50

50

50

50

100

100

100

100

100

100

150

150

150

150

150

150

0

100

0

100

0

100

0

100

0

100

0

100

0

0

01

01

0

01

50

50

50

50

50

50

100

100

100

100

100

100

150

150

150

150

150

150

0

100

0

100

0

100

0

100

0

100

0

100

0

0

01

0

0

0

50

50

50

50

50

50

100

100

100

100

100

100

150

150

150

150

150

150

100

0

100

0

100

100

0

100

0

100

DeiT

NeuTRENO DeiT

Cosine Similarity of Attention

Matrices between layers

Layer

0.7

0.6

0.5

0.4

0.3

23456789101112

DeiT

NeuTRENO DeiT

Layer

Cosine Similarity

between tokens

1.0

0.9

0.8

0.7

0.6

0.5

0.4

0.3

1

2345

56789101112

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.

G.6 Robustness of NeuTRENO

In addition to the standard metrics, we evaluate the robustness of our NeuTRENO model compared to the baseline transformer model, particularly under adversarial examples and for out-of-distribution generalization. Table 12 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) [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 of ˜ λ

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 or above 1 , respectively), NeuTRENO's performance declines.

G.9 Scalability of NeuTRENO

To demonstrate the scalability of our proposed model, we conduct 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), we plot the cosine similarity between token representations across layers of a pre-trained BERT-base model [17] on the SQuAD v1.1 question answering task [48] and observe

Pre-trained BERT

NeuTRENO + Pre-trained BERT

Layer

Layer

Cosine Similarity

between tokens

SQuAD Question Answering

STS-12 Semantic Textual Similarity

Pre-trained

Randomized

Layer

Pre-trained

Pre-trained SimCSE

NeuTRENO + Pre-trained SimCSE

0.8

1.0

0.8

0.7

0.9

0.7

0.6

0.8

0.6

0.7

0.5

0.5

0.6

0.4

0.4

0.5

0.3

0.3

0.4

0.2

123456789101112

3456789101112

123456789101112

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)), 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 F1-score vs. 80 . 77 exact match score and 88 . 12 F1-score). 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)). 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 on the STS12 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-based models leads to improved empirical performance.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nRecent 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.\n\nSTRENGTHS:\n1. The theoretical investigation conducted in this paper is rigor and in-depth.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\n1. 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\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-06-17"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\n- 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\nLIMITATIONS:\nPlease see the previous two sections", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-03"}

comments: 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.

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.

The 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).

QUESTIONS:
please refer to weaknesses

LIMITATIONS:
The generalizability and scalability of the proposed solution to larger transformer models have not been investigated.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-06

comments: 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.

STRENGTHS:
I think these analysis and regularization tools are useful addition to the exploding area of transformers.

WEAKNESSES:
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. 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?

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
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-07

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nStrengths:\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\nWEAKNESSES:\nWeaknesses:\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\nQUESTIONS:\nPlease 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\nLIMITATIONS:\nThis is N/A for this work.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-10"}