Year: 2022

Abstract:
Vision Transformer (ViT) has recently demonstrated promise in computer vision problems. However, unlike Convolutional Neural Networks (CNN), it is known that the performance of ViT saturates quickly with depth increasing, due to the observed attention collapse or patch uniformity. Despite a couple of empirical solutions, a rigorous framework studying on this scalability issue remains elusive. In this paper, we first establish a rigorous theory framework to analyze ViT features from the Fourier spectrum domain. We show that the self-attention mechanism inherently amounts to a low-pass filter, which indicates when ViT scales up its depth, excessive low-pass filtering will cause feature maps to only preserve their Direct-Current (DC) component. We then propose two straightforward yet effective techniques to mitigate the undesirable low-pass limitation. The first technique, termed AttnScale, decomposes a self-attention block into low-pass and high-pass components, then rescales and combines these two filters to produce an all-pass self-attention matrix. The second technique, termed FeatScale, re-weights feature maps on separate frequency bands to amplify the high-frequency signals. Both techniques are efficient and hyperparameter-free, while effectively overcoming relevant ViT training artifacts such as attention collapse and patch uniformity. By seamlessly plugging in our techniques to multiple ViT variants, we demonstrate that they consistently help ViTs benefit from deeper architectures, bringing up to 1.1% performance gains “for free” (e.g., with little parameter overhead). We publicly release our codes and pre-trained models at https://github.com/VITA-Group/ViT-Anti-Oversmoothing.

text: Vision Transformer (ViT) has recently demonstrated promise in computer vision problems. However, unlike Convolutional Neural Networks (CNN), it is known that the performance of ViT saturates quickly with depth increasing, due to the observed attention collapse or patch uniformity. Despite a couple of empirical solutions, a rigorous framework studying on this scalability issue remains elusive. In this paper, we first establish a rigorous theory framework to analyze ViT features from the Fourier spectrum domain. We show that the self-attention mechanism inherently amounts to a low-pass filter, which indicates when ViT scales up its depth, excessive low-pass filtering will cause feature maps to only preserve their Direct-Current (DC) component. We then propose two straightforward yet effective techniques to mitigate the undesirable low-pass limitation. The first technique, termed AttnScale, decomposes a self-attention block into low-pass and high-pass components, then rescales and combines these two filters to produce an all-pass self-attention matrix. The second technique, termed FeatScale, re-weights feature maps on separate frequency bands to amplify the high-frequency signals. Both techniques are efficient and hyperparameter-free, while effectively overcoming relevant ViT training artifacts such as attention collapse and patch uniformity. By seamlessly plugging in our techniques to multiple ViT variants, we demonstrate that they consistently help ViTs benefit from deeper architectures, bringing up to 1.1% performance gains “for free” (e.g., with little parameter overhead). We publicly release our codes and pre-trained models at https://github.com/VITA-Group/ViT-Anti-Oversmoothing.

heading: 1 INTRODUCTION
text: Transformers have achieved phenomenal success in Natural Language Processing (NLP) (Vaswani et al., 2017; Devlin et al., 2018; Dai et al., 2019; Brown et al., 2020), and recently in a wide range of computer vision applications too (Dosovitskiy et al., 2020; Liu et al., 2021; Arnab et al., 2021; Carion et al., 2020; Jiang et al., 2021a). One representative advance, the Vision Transformer (ViT) (Dosovitskiy et al., 2020), stacks Multi-head Self-Attention (MSA) blocks, by treating each local image patch as semantic tokens and modeling their interactions globally. Unlike Convolutional Neural Networks (CNNs) that hierarchically enlarge the receptive from local to global, even a shallow ViT is able to effectively capture the global contexts, leading to their very competitive performance on image classification and other tasks (Liu et al., 2021; Jiang et al., 2021a).
Going deep has always been a trend in deep learning (LeCun et al., 2015; Krizhevsky et al., 2012), and ViT was expected to make no exception. One might reasonably conjecture that a deeper ViT with more MSA blocks significantly outperform its shallower baseline. Unfortunately, building deeper ViTs face practical challenges. Empirically, Zhou et al. (2021a) shows a vanilla ViT of 32 layers under-performs the 24-layer one. Gong et al. (2021) demonstrates a downgraded patch diversity in deeper layers, and Dong et al. (2021) mathematically reveals the rank collapse phenomenon when Transformer goes deeper. Despite efforts towards deep ViT through patch diversification (Gong et al., 2021; Zhou et al., 2021b), rank collapse alleviation (Zhou et al., 2021a; Zhang et al.), and training
stabilization (Touvron et al., 2021b; Zhang et al., 2019), most of them are restricted to empirical studies. Rethinking the problem with deep ViT from a more principled angle pends further efforts.
In this paper, we present the first rigorous analysis of stacking self-attention mechanism in the Fourier space. We mathematically show that cascading self-attention blocks is equivalent to repeatedly applying a low-pass filter, regardless of the input key or query tokens (Section 2.2). As a consequence, going deeper with vanilla ViT blocks only preserves Direct Component (DC) of the signal at the output layer. This theoretical finding explains the observations of patch uniformity and rank collapse, and is also inherently related to the over-smoothing phenomenon in Graph Convolutional Networks (GCNs) (Kipf & Welling, 2016; Nt & Maehara, 2019; Oono & Suzuki, 2019; Cai & Wang, 2020). Moreover, we also reveal the role of other transformer modules (e.g., MLP and residual connection) in preventing this undesirable low-pass filtering (Section 2.3).
Built on the aforementioned analysis framework in the Fourier domain, we propose two novel techniques, to mitigate the low-pass filtering effect of self-attention and effectively scale up the depth of ViTs. The first technique, termed Attention Scaling (AttnScale), directly manipulates on the calculated attention map to enforce an all-pass filter (Section 3.1). It decomposes the self-attention matrix into a low-pass filter plus a high-pass filter, then adopts a learnable weight to adaptively amplify the effect of high-pass filter. The second technique, termed Feature Scaling (FeatScale), hinges on feature maps to re-weight different frequency bands separately (Section 3.2). It employs trainable coefficients to re-mix the DC and high-frequency components, hence selectively enhancing the high-frequency portion of the MSA output. Both AttnScale and FeatScale are extremely memory and computationally friendly. Neither runs Fourier transformation explicitly, bringing little extra complexity to the original ViTs.
Our contributions can be summarized as follows:
• We establish the first rigorous theoretical analysis of ViT from the spectral domain. We characterize the low-pass filtering effect of cascading MSAs, which connects to the recent empirical findings of ViT patch diversity loss or rank collapse.
• We present two theoretically grounded Fourier-domain scaling techniques, named AttnScale and FeatScale. They operating on re-adjusting the low- and high-frequency components of the attention maps and feature maps, respectively. Both are efficient, hyperparameter-free, easy-to-use, and able to generalize across different ViT variants.
• We conduct extensive experiments by integrating AttnScale and FeatScale with different ViT backbones. Both of our approaches substantially boost DeiT, CaiT, and Swin-Transformer with up to 1.1%, 0.6% and 0.5% performance gains, without whistles and bells.

heading: 2 WHY VIT CANNOT GO DEEPER?

heading: 2.1 NOTATION AND PRELIMINARIES
text: We begin by introducing our notations. Let X ∈ Rn×d denote the feature matrix, where n is the number of samples, and d is the feature dimension. Let xi ∈ Rd,∀i = 1, · · · , n, the i-th row ofX , denote the feature vector of the i-th sample, and zj ∈ Rn,∀j = 1, · · · , d, the j-th column of X , represent signals of the j-th channel. In the context of ViT, X denotes a set (sequence) of image patches, xi denotes the flatten version of the i-th patch embedding (d = patch width× patch height), and zj denotes a whole image signal of the j-th channel.
Transformer Architecture Vision Transformer (ViT) consists of three main components: a patch embedding and position encoding part, a stack of transformer encoder block with Multi-Head Self-Attention (MSA) and Feed-Forward Network (FFN), and a score readout function for image classification. We depict a transformer block in Fig. 3a. The key ingredient here is the Self-Attention (SA) module, which takes in the token representation of the last layer, and encodes each image token by aggregating information from other patches with respect to the computed attention value, formulated as below (Vaswani et al., 2017):
SA(X) = softmax
( XWQ(XWK) T
√ d
) XWV , (1)
where WK ∈ Rd×dk ,WQ ∈ Rd×dq ,WV ∈ Rd×d are the key, query, and value weight matrices, respectively, √ d here denotes a scaling factor, and softmax(·) operates onX row-wisely. Multi-Head Self-Attention (MSA) involves a group of SA heads and combines their outputs through a linear projection (Vaswani et al., 2017):
MSA(X) = [SA1(X) · · · SAH(X)]WO, (2) where the subscripts denote the SA head number, H is the total number of SA heads, and WO ∈ RHd×d projects multi-head outputs to the hidden dimension. Besides MSA module, each transformer block is equipped with a normalization layer, feed-forward network, and skip connections to cooperate with MSA. Formally, a transformer block can be written as follows:
X ′ = MSA(LayerNorm(X)) +X, (3)
Y = FFN(LayerNorm(X ′)) +X ′. (4)
Fourier Analysis The main analytic tool in this paper is Fourier transform. Denote F : Rn → Cn be the Discrete Fourier Transform (DFT) with the Inverse Discrete Fourier Transform (IFT) F−1 : Cn → Rn. Applying F to a flatten image signal x is equivalent to left multiplying a DFT matrix, whose rows are the Fourier basis fk = [ e2πj(k−1)·0 · · · e2πj(k−1)·(n−1) ]T /√ n ∈ Rn, where k denotes the k-th row of DFT matrix, and j is the imaginary unit. Let z̃ = Fz be the spectrum of z, and z̃dc ∈ C, z̃hc ∈ Cn−1 take the first element and the rest elements of z̃, respectively. Define DC [z] = z̃dcf1 ∈ Cn as the Direct-Current (DC) component of signal z, and HC [z] = [f2 · · · fn] z̃hc ∈ Cn the complementary high-frequency component. In signal processing, a low-pass filter is a system that suppresses the high-frequency component of signals and retain the low-frequency component. In this paper, we refer to low-pass filter as a particular type of filters that only preserve the DC component, while diminishing the remaining high-frequency component. To be more precise, we define low-pass filters in Definition 1. Definition 1. Given an endomorphism f : Rn → Rn with f t denoting applying f for t times, f is a low-pass filter if and only if for all z ∈ Rn:
lim t→∞ ‖HC [f t(z)]‖2 ‖DC [f t(z)]‖2 = 0. (5)
Definition 1 reveals the nature of low-pass filters: they will produce a dominant response on DC component, while imposing an inhibition effect on the high-frequency band. We refer interested readers to Appendix A for more useful backgrounds.

heading: 2.2 SELF-ATTENTION IS A LOW-PASS FILTER
text: In this subsection, we will give theoretical justification on self-attention in terms of its spectral-domain effect. Our main result is that self-attention is constantly a low-pass filter, which continuously erases high-frequency information, thus causing ViT to lose features expressiveness at deep layers.
Formally, we have the following theorem that shows attention matrix produced by a softmax function (e.g, Eqn. 1) is a low-pass filter independent of the input token features or key/query matrices. Theorem 1. (SA matrix is a low-pass filter) LetA = softmax(P ), where P ∈ Rn×n. ThenA must be a low-pass filter. For all z ∈ Rn, limt→∞‖HC [Atz]‖2/‖DC [Atz]‖2 = 0.
Theorem 1 is a straightforward result of Perron-Frobenius theorem. See Appendix B.1 for a proof. Theorem 1 also reveals that no matter how attention is computed inside the softmax function, including dot product (Vaswani et al., 2017), linear combination (Veličković et al., 2017), or L2 distance (Kim et al., 2021), the resulting attention matrix is always a low-pass filter. One can see consecutively applying self-attention matrix simulates the process of ViT’s forward propagation. As the layer number increases infinitely, the final output will only keep the DC bias, and ViT loses all the feature expressive power. Corollary 2. Let P1,P2, · · · ,Pn be a sequence of matrix in Rn×n, and each Pk,∀k = 1, · · · , L has Ak = softmax(Pk). Then ∏L k=1Ak is also a low-pass filter.
In fact, ViT re-computes self-attention matrices per layer, which seems to avoid the consecutive power of an identical self-attention matrix. However, we also provide Corollary 2, which suggests even the ViT consists of distinctive self-attention matrix at each layer, their composition turns out to act like a low-pass filter as well. We also visualize the spectrum of attention maps (Fig. 2 and more on Appendix F.2) to support our theoretical conclusions.
Knowing that self-attention matrices amount to low-pass filters, we are also interested in to which extent an MSA layer would suppress the high-frequency component. Thereby, we also provide a convergence rate to illustrate this speed that the high-frequency component are being annihilated. Theorem 3. (smoothening rate of SA) Let A = softmax(P ) and α = maxi,j |Pij |, where P ∈ Rn×n. Define SA(X) = AXWV as the output of a self-attention module, then
‖HC [SA(X)]‖F ≤ √
ne2α
e2α + n− 1 ‖WV ‖2‖HC [X]‖F . (6)
In particular, when P =XWQ(XWK)T / √ d, and assume tokens are distributed inside a ball with radius γ > 0, i.e., ‖xi‖2 ≤ γ,∀i = 1, · · · , n, then α ≤ γ2‖WQW TK‖2/ √ d.
The proof of Theorem 3 can be found in Appendix B.3. Theorem 3 says the high-frequency intensity ratio to the pre- and post- attention aggregation is upper bounded by ‖WV ‖2 √ ne2α
e2α+n−1 . When ‖WV ‖2 √ ne2α e2α+n−1 < 1, HC [SA(X)] converges to zero exponentially. We note that, no
matter how attention is computed or signals are initialized, since √ ne2α e2α+n−1 is bounded by √ n, ‖WV ‖2 < 1/ √ n will definitely cause a monotonically decreasing high-frequency component. When dot-product attention is adopted, a sufficient condition that ‖HC [X]‖F decreases to zero within logarithmic time is γ2‖WQW TK‖2/ √ d+ log‖WV ‖2 ≤ log n−1n /2.

heading: 2.3 EXISTING MECHANISMS THAT COUNTERACT LOW-PASS FILTERING
text: In this section, we take other ViT building blocks into consideration. We will justify whether MultiHead Self-Attention (MSA), Feed-Forward Network (FFN), and residual connections can effectively alleviate the low-pass filtering drawbacks. All the derivations follow from Theorem 3, and some proof ideas are borrowed from Dong et al. (2021). We further present Fig. 1 to justify our results.
Does multi-head help? MSA employs weights to combine the results of multiple self-attention blocks. We can rewrite it as MSA(X) = ∑H h=1 SA(X)W h O. We show by Proposition 4 in
Appendix C.1 that the convergence rate turns to σ1σ2H √ ne2α
e2α+n−1 , where H is the number of heads,
σ1 = max H h=1‖W hV ‖2 and σ2 = maxHh=1‖W hO‖2. One can see MSA can only slow down the convergence up to a constant σ2H , which does not root out the problem.
Does residual connection benefit? In addition to MSA, a transformer block also leverages a skip connection, which can be formulated as Res(X) = MSA(X)+X . We show that residual connection
can effectively prevent high-frequency component from diminishing to zero by promoting the rate σ1σ2H √ ne2α e2α+n−1 to 1 + σ1σ2H √ ne2α e2α+n−1 > 1. Refer to Proposition 5 in Appendix C.2.
Does FFN make any difference? A feed-forward network is appended to MSA module. We characterize its effect in Appendix C.3. Our Proposition 6 suggests that a FFN with Lipschitz constant σ3 contributes a σ3 ( 1 + σ1σ2H √ ne2α
e2α+n−1
) convergence rate, which does not improve the original
one. However, if skip connection is adopted over FFN, σ3 > 1 can guarantee the upper bound of the high-frequency component is non-contractive.
Although multi-head, FFN, and skip connection all help preserve the high-frequency signals, none would change the fact that MSA block as a whole only possesses the representational power of low-pass filters. Our Proposition 4, 5, 6 states multi-head, FFN, skip connections can only slow down the convergence by indistinguishably amplifying low- and high-frequency components with the same factor. However, since they are incapable of promoting high-frequency information separately, it is inevitable that high-frequency components are continuously diluted as ViT goes deeper. This restricts the expressiveness of ViT, resulting in the performance saturation in deeper ViT.

heading: 2.4 CONNECTIONS TO EXISTING THEORETIC UNDERSTANDING WORKS
text: It is known that Graph Convolutional Networks (GCN) are not more than low-pass filters (Nt & Maehara, 2019). In the meanwhile, Oono & Suzuki (2019); Cai & Wang (2020) pointed out GCN’s node features will be exponentially trapped into the nullspace of the graph Laplacian matrix. Similarly, our work concludes that self-attention module is yet another low-pass filter. Combining with our theoretical derivation, one can see the root reason is that both graph Laplacian matrices and selfattention matrices consistently own a fixed leading eigenvector, namely the DC basis. This makes aggregating information via such matrices inherently project the token representation onto these invariant eigenspaces. And we note that over-smoothing, rank collapse, and patch uniformity are all the manifestation of excessive low-pass filtering. See Appendix D.1 for more discussion.
In Dong et al. (2021), the authors proved that ViT’s feature maps will doubly exponentially collapses to a rank-1 matrix, which reveals ViT loses feature expressiveness at deep layers. Besides, they also gave a systematic study on other building blocks of transformer. While they share the similar insights with us, our work further specifies which rank-1 matrix the feature activation will converge to, namely the subspace spanned by the DC basis. That makes our theory to be better grounded with signal-processing and geometric interpretations, via directly measuring the intensity of the highfrequency residual, instead of examining a composite norm distance to an agnostic rank-1 matrix. Although Dong et al. (2021) presented a faster convergence speed, we respectfully suggest that the current proof of Dong et al. (2021) might be deficient, or at least incomplete in the assumptions (see Appendix D.2). Moreover, our theory can be generalized to other attention mechanisms such as logistic attention (Veličković et al., 2017) and L2 distance (Kim et al., 2021). See Appendix D.3.

heading: 3 ATTNSCALE & FEATSCALE: SCALING FROM THE FOURIER-DOMAIN

heading: 3.1 ATTNSCALE: MAKE ATTENTION AN ALL-PASS FILTER
text: As we discussed in Section 2.2, self-attention matrix can only perform low-pass filtering, which narrows the filter space ViT can express. Inspired by this, we propose a scaling techniques directly manipulating the attention map, termed Attention Scaling (AttnScale), to balance the effects of lowand high-pass filtering and produce all-pass filters. AttnScale decomposes the self-attention matrix to a low-pass filter plus a high-pass filter, and introduces a trainable parameter to rescale the high-pass filter to match the magnitude with the low-pass component.
Formally, letA denote a self-attention matrix. To decompose a low-pass filter fromA, we find the largest possible low-pass filter that can be extracted from A. We use Lemma 8 in Appendix E to justify our solution. By Lemma 8, we can simply extract L = F−1 diag(1, 0, · · · , 0)F = 11T /n from A and take the complementary part as the high-pass filter. Afterward, we can rescale the high-pass component of the filter, and combine low-pass and high-pass together to form a new self-attention matrix. We illustrate this scaling trick in Fig. 3b. To be more precise, for the l-th layer and h-th head, we recompute the self-attention map as follows:

heading: Multi-Head
text: A (l,h) LP =
1 n 11T , (7)
A (l,h) HP = A (l,h) −A(l,h)LP , (8)
Â(l,h) = A (l,h) LP + (ωl,h + 1)A (l,h) HP , (9)
where ωl,h is a trainable parameter, and different layers and heads adopt separate ωl,h. During training time, ωl,h’s are initialized with 0, and jointly tuned with the other network parameters. By adjusting ωl,h, Â(l,h) can simulate any type of filters: low-pass, high-pass, band-pass, or all-pass. Note that our AttnScale is extremely lightweight, as it only brings O(HL) extra parameters, where H is the number of heads, and L is the number of ViT blocks.

heading: 3.2 FEATSCALE: REWEIGHT HIGH-FREQUENCY SIGNALS
text: According to our analysis in Section 2.2, MSA module will indiscriminately suppress high-frequency signals, which leads to severe information loss. Even though residual connection can retrieve lost information through the skip path, the high-frequency portion will be inevitably diluted (Theorem 1). To this end, we propose another scaling technique that operates on feature maps, named Feature Scaling (FeatScale). FeatScale processes the output of MSA by mixing the information from varying frequency bands discriminatively. FeatScale first decomposes the resultant signals into their DC and high-frequency components. Then it introduces two groups of parameters to re-weight the two components for each channel, respectively. The pipeline of this scaling technique is depicted in Fig. 3c. To be more precise, we re-weight the output of the l-th MSA by
X (l) DC = DC [MSA(X)] (diag(sl) + I), (10)
X (l) HC = HC [MSA(X)] (diag(tl) + I), (11)
X(l) =X (l) DC +X (l) HC , (12)
where sl ∈ Rd and tl ∈ Rd are learnable parameters to perform channel-wise re-weighting. We initialize sl and tl with zeros and tune them with gradient descent. After adjusting the proportion of different frequency signals, FeatScale can prevent the dominance of the DC component. DC [·] and HC [·] are cheap to compute without explicit Fourier transform. Calculating DC [X] is as simple as running the column average of matrixX , andHC [X] can be efficiently computed byX −DC [X].

heading: 3.3 DISCUSSION
text: We have proposed two methods to facilitate the deeper stacking of ViT MSA modules, and discussed their motivations and strengths from our perspective of filtering and signal processing. In this section, we will connect these two techniques with commonly mentioned problems with ViTs.
How does AttnScale prevent attention collapse? Deep ViT suffers from the attention collapse issue (Zhou et al., 2021a). When transformer goes deepr, the attention maps gradually become similar
and even much the same after certain layers. Combining with our thoery, one can see collapsed attention maps turn out to be a pure low-pass filter, which wipes off all the high-frequency information in one shot. Zhou et al. (2021a) proposed the re-attention trick, which blends attention map across different heads. By doing this, modified attention maps aggregate high-pass components from other heads and is endowed with richer filtering property. We note that our AttnScale is akin to a more lightweight re-attention mechanism with better interpretability. We rewrite the attention map as a sum of an already-collapsed attention (11T /n) with the complementary residual map that encodes diverse patterns in a self-attention matrix. By re-weighting the residual map, the diversified patterns can be amplified, which prevents it from degenerating to a rank-1 matrix. We further verify this argument using cosine similarity metric (Zhou et al., 2021a) in the upper sub-figure of Fig. 4.
How does FeatScale conserve patch diversity? Self-attention blocks tend to map different patches into similar latent representations, yielding information loss and performance degradation (Gong et al., 2021). By our theory, this asymptotic smoothness of feature map is caused by excessive low-pass filtering, and the remaining DC bias signifies uniform patch representations. Conventional approaches to addressing this problem include incorporating convolutional layers (Wu et al., 2021; Jiang et al., 2021b) and enforcing patch diversity regularizations (Gong et al., 2021). As diverse features are often characterized by high-frequency signals, our FeatScale instead elevating the high-frequency component via a learnable scaling factor, can be regarded as a more straightforward way to reconstruct the patch richness. Compared with LayerScale (Touvron et al., 2021b), in which each frequency band is equally scaled, our FeatScale treating DC and high-frequency components differently, not only perform a per-channel normalization, but also perform a spectral-domain calibration with highfrequency details and low-frequency characteristics. The lower sub-plot of Fig. 4 shows ViT with our FeatScale has lower feature similarity at deep layer.

heading: 4 RELATED WORK
text: Transformers in Vision. Transformer (Vaswani et al., 2017) entirely relies on self-attention mechanism to capture correlation and exchange information globally among the input. It has achieved a remarkable performance in natural language processing (Devlin et al., 2018; Dai et al., 2019; Brown et al., 2020) and many cross-disciplinary applications (Jumper et al., 2021; Ying et al., 2021). Recent advances have also successfully applied Transformer to computer vision tasks. Dosovitskiy et al. (2020) first adopts a pure transformer architecture (ViT) for image classification. The follow-up works (Chen et al., 2021b) extend ViT to various vision tasks, such as object detection (Carion et al., 2020; Zhu et al., 2020; Zheng et al., 2020; Sun et al., 2020), segmentation (Chen et al., 2021a; Wang et al., 2021), image generation (Parmar et al., 2018; Jiang et al., 2021a), video processing (Zhou et al., 2018; Arnab et al., 2021), and 3D instance processing (Guo et al., 2021; Lin et al., 2021). To capture multi-scale non-local contexts, Zhang et al. (2020) designs transformers in self-level, top-down, and bottom-up interaction fashion. Liu et al. (2021) presents hierarchical ViTs with shifted window based attention that can efficiently extract multi-scale features. To dismiss ViT from the heavy reliance on large-scale dataset pre-training, Touvron et al. (2021a); Yuan et al. (2021) propose knowledge distillation and progressive tokenization for data-efficient training. Despite impressive effectiveness, most of these model are only based on relatively shallow ViT backbones with a dozen of MSA blocks.
Advances in deep ViTs. Building deeper ViTs has arisen many interests. Zhou et al. (2021a) first investigated the depth scalability of ViT. The authors found that the attention collapse hinders ViT from scaling up, and propose two methods to conquer this problem i) increasing the embedding dimension, and ii) a cross-head re-attention trick to regenerate attention map. A concurrent work Touvron et al. (2021b) came up with a LayerScale layer that performs per-channel multiplication for each residual block. More importantly, they make explicit separation of transformer layers involving self-attention between patches, from class-attention layers that are devoted to extract the global content into a single embedding to be decoded. Gong et al. (2021) further proposed a series of losses that can enforce patch diversity in ViT. Such regularizations include penalty on cosine similarity, patch-wise contrastive loss, and mixing loss. Tang et al. (2021) presented a shortcut augmentation scheme with block-circulant projection to improve feature diversity. Although these existing solutions manage to deepen ViTs, most of them are empirical works and bring no principled theory.
Role of depth in NNs. Discussing the importance of deep structures in Neural Networks (NNs) is an overly broad topic. Here we only focus on a subset of works that scaling up a transformer could relate to. For ordinal deep learning models, such as FFNs and CNNs, deep architecture immediately
benefits from the universal approximation power and expressive capacity (Cybenko, 1989; Hornik, 1991; Telgarsky, 2016; Lu et al., 2017; Petersen & Voigtlaender, 2020; Zhou, 2020). In contrast, several studies in graph learning domain have reported severe performance degradation due to oversmoothing when stacking many layers (Kipf & Welling, 2016; Wu et al., 2019; Li et al., 2018). The subsequent studies (Nt & Maehara, 2019; Oono & Suzuki, 2019; Cai & Wang, 2020) gave theoretical explanations of the over-smoothing phenomena from the views of graph signal filtering and feature dynamics. Likewise, ViT have been witnessed performance saturation when going deeper. However, to our best knowledge, Dong et al. (2021) is the sole work in the literature that systematically and rigorously analyzes this issue with deep ViT. The main idea of this work is showing the self-attention block will downgrade the rank of the feature maps. Our work takes one step forward by studying ViT on spectral domain, and manages to reveal the signals will ultimately fall into the one-dimension DC subspace. We see our theory and techniques are also applicable to NLP transformers. However, we only focus on ViT because empirical observations indicate NLP modeling (including Transformer) does not require a deep structure (Vaswani et al., 2017; Brown et al., 2020), while vision tasks always demand one (LeCun et al., 2015).

heading: 5 EXPERIMENTS
text: In this section, we report experiment results to validate our proposed methods. First, we validate the effectiveness of our AttnScale and FeatScale when integrated with different deep ViT backbones (Section 5.1). Second, we compare our best models with state-of-the-art (SOTA) results (Section 5.2). All of our experiments are conducted on the ImageNet dataset (Russakovsky et al., 2015) with around 1.3M images in the training set and 50k images in the validation set. Our implementations are based on Timm (Wightman, 2019) and DeiT (Touvron et al., 2021a) repositories.

heading: 5.1 HOW CAN ATTNSCALE & FEATSCALE BENEFIT DEEP VIT?
text: Experiment Settings. In this subsection, we intend to testify our models are beneficial to various ViT backbones with different depth settings and training modes. We choose DeiT (Touvron et al., 2021a) as our first backbone in order to train from scratch. When training 12-layer DeiT, we follow the same training recipe, hyper-parameters, and data augmentation with Touvron et al. (2021a). When training 24-layer DeiT, we follow the setting in Gong et al. (2021). Specially, we set dropout rate to 0.2 when training 24-layer DeiT (Touvron et al., 2021b). Our second backbone is CaiT (Touvron et al., 2021b). We only apply our techniques to the patch embedding layers. The third backbone is the SOTA model Swin-Transformer (Liu et al., 2021). All experimental settings share the same with Liu et al. (2021). In addition to training from scratch, we also investigate the fine-tuning setting. We defer this part to Appendix G.1.
Results. All of our experimental evaluations are summarized in Table 1. The results suggest our proposed AttnScale and FeatScale successfully facilitate both DeiT, CaiT, Swin-Transformer under different depth settings and training modes. Specifically, AttnScale brings less than 100/150 extra parameters for 12/24-layer DeiT while boosting the top-1 accuracy by 0.9% for 12-layer DeiT and 0.6% for 24-layer DeiT. Our FeatScale substantially improves top-1 accuracy by 1% for 12-layer DeiT and 0.8% for 24-layer DeiT. Compared with existing techniques, the improvements of AttnScale and FeatScale already surpass re-attention (0.6%) (Zhou et al., 2021a), LayerScale
(0.7%) (Touvron et al., 2021b), and late class token insertion (0.6%) (Touvron et al., 2021b). We also observe a consistent 0.6% performance gain when AttnScale and FeatScale plugged into CaiT. Under fine-tuning setting, as we will show in Appendix G.1, only tens of epoch’s fine-tuning can further promote their performance by ≥ 0.2% (see Table 3). On Swin-Transformer, both our AttnScale and FeatScale bring around 0.5% accuracy gain. This makes Swin-S with 50M parameters even comparable to Swin-B (83.5% top-1 accuracy on ImageNet1k) with 88M parameters. We defer more model interpretation and visualization to Appendix G. For a brief summary, we observe that both shallow and deep ViT enjoy from AttnScale and FeatScale that: 1) the attention maps can simulate richer filtering properties (compare Fig. 9 with Fig. 10), and 2) more high-frequency data can be preserved (refer to Fig. 11).

heading: 5.2 COMPARISON WITH SOTA MODELS
text: In this subsection, we compare our best models with state-of-the-art models on ImageNet benchmark. We choose SOTA models from three classes: CNN only, CNN + transformer, and pure transformer. For transformer domain, we only conduct experiments with those lightweight models with comparable number of parameters, such as ViT-S and DeiT-S. All the results are presented in Table 2.
Among all methods, Swin-Transformer combined with our methods achieves the state-of-the-art performance. Our CaiT-S + AttnScale and CaiT-S + FeatScale on 24-layer CaiT-S also attain superior results over all other pure transformers, while keeping low parameter cost. That our performance surpasses some CNN-based models (e.g., ResNet-152 and CVT) indicates by increasing depth, ViT will be endowed with higher potential to surpass CNNs that have been dominating computer vision domain so far. Our DeiT-S+FeatScale result also outperforms ViT-B/16 and DeiT-S Distilled, which suggests deepening network can bring more considerable accuracy gain than increasing model width or employing a teacher model.

heading: 6 CONCLUSION
text: In this paper, we investigate the scalability issue with ViT and propose two practical solutions via Fourier domain analysis. Our theoretical findings indicate Multi-Head Self-Attention (MSA) inherently performs low-pass filtering on image signals, thus causes rank collapse and patch uniformity problems in deep ViT. To this end, we proposed two techniques, AttnScale and FeatScale, that can effectively break such low-pass filtering bottleneck by adaptively scaling high-pass filter component and high-frequency signals, respectively. Our experiments also validate the effectiveness of our methods. Both techniques can boost various ViT backbones by a significant performance gain. Grounded with our theoretical framework, interesting directions for further work include designing parameter regularizations and spectrum-specific normalization layers.

heading: ACKNOWLEDGMENTS
text: Z.W. is in part supported by an NSF SCALE MoDL project (#2133861).

heading: A MORE PRELIMINARIES ON FOURIER ANALYSIS
text: In this appendix, we provide more preliminary knowledge about Fourier analysis. Here, we only consider discrete Fourier transform on real-value domain F : Rn → Cn. A Discrete Fourier transform (DFT) can be written in a matrix form as below:
DFT = 1√ n  1 1 · · · 1 1 e2πj · · · e2πj(n−1) ... ... . . . ... 1 e2πj(k−1)·1 · · · e2πj(k−1)·(n−1) ... ... . . .
... 1 e2πj(n−1) · · · e2πj(n−1)2
 , (13)
and its inverse discrete Fourier transform isDFT−1 =DFT . In this paper, we regard matrices as multi-channel signals. For example,X ∈ Rn×d means d-channel n-length signals. When DFT and inverse DFT are applies to multi-channel signals, each channel is transformed independently, i.e., F(X) = [F(x1) · · · F(xd)] =DFT ·X . Hereby, we can simply operators DC [·] andHC [·] using the matrices in Eqn. 13. By definition, we can write DC [·] as below:
DC [x] =DFT−1 diag(1, 0, · · · , 0)DFTx (14)
= 1
n 11Tx, (15)
namely DC [·] = 11T /n. Conversely, we can writeHC [·] as:
HC [x] =DFT−1 diag(0, 1, · · · , 1)DFTx (16) =DFT−1(I − diag(1, 0, · · · , 0))DFTx (17)
= I − 1 n 11Tx, (18)
which indicatesHC [·] = I − 11T /n. We will frequently use these derivations later in the proofs.

heading: B DEFERRED PROOFS

heading: B.1 PROOF OF THEOREM 1
text: Theorem 1. (SA matrix is a low-pass filter) LetA = softmax(P ), where P ∈ Rn×n. ThenA must be a low-pass filter. For all z ∈ Rn,
lim t→∞ ‖HC [Atz]‖2 ‖DC [Atz]‖2 = 0.
Proof. Let λ1, λ2, · · · , λs ∈ C be the eigenvalues of A with ordering |λ1| ≥ |λ2| ≥ · · · ≥ |λs|. Notice that,A is a positive matrix, each of whose element is strictly greater than zero (Aij > 0,∀i, j). Besides, for all i = 1, · · · , n, ∑n j=1Ai,j = 1. Therefore, A1 = 1 implies A must have an eigenvalue 1 and its corresponding eigenvector is the all-one vector 1.
By Perron-Frobenius Theorem (Meyer, 2000), eigenvalue 1 corresponds to a all-positive eigenvector 1, implies λ1 = 1 should be the largest eigenvalue without multiplicity, and the absolute value of other eigenvalues λ2, · · · , λs must be less than 1. Let us rewrite A in the Jordan canonical form
A = PJP−1:
A = [v1 · · · vn]︸ ︷︷ ︸ P
 λ1 J(λ2) . . .
J(λs)  ︸ ︷︷ ︸
J
u T 1 ... uTn  ︸ ︷︷ ︸ P−1 , (19)
where the Jordan block J(λ) can be written as
J(λ) =  λ 1 λ 1 . . . . . .
λ 1 λ
 . (20)
ApplyingA to z for t times can be written asAtz which is equivalent to:
Atz = PJ tP−1z = P  λ1 J(λ2) . . .
J(λs)
 t P−1z (21)
= P  λt1 J(λ2) t
. . . J(λs) t P−1z (22) Let f(x) = xt, thenAt = P f(J)P−1 = P diag(f(λ1), f(J(λ2)), · · · , f(J(λs)))P−1. Suppose a Jordan block with shape k × k, then
f(J(λ)) =  f(λ) f ′(λ) f ′′(λ) 2! · · ·
f(k−1)(λ) (k−1)!
f(λ) f ′(λ) . . . ... . . . . . . f
′′(λ) 2!
f(λ) f ′(λ) f(λ)
 (23)
Therefore, on the diagonal number m ≤ min(t, k − 1) above the main diagonal stands:
t(t− 1)...(t−m+ 1) m! λt−m (24)
For arbitrary m ≤ k − 1 and |λ| < 1,
lim t→∞ t(t− 1)...(t−m+ 1) m! λt−m = 0 (25)
Recall that λ1 = 1 and according to the definition of Jordan canonical form, v1 = 1. By Eqn. 25:
lim t→∞ Atz = P lim t→∞ diag(f(λ1), f(J(λ2)), · · · , f(J(λs)))P−1z (26)
= P diag(λt1,0, · · · ,0)P−1z (27) = λt1v1u T 1 z (28)
= 1uT1 z (29)
Plug the result from Eqn. 29 into the original limit:
lim t→∞ ‖HC [Atz]‖2 ‖DC [Atz]‖2 = lim t→∞ √ ‖HC [Atz]‖22 ‖z −HC [Atz]‖22
(30)
= lim t→∞
√ ‖HC [Atz]‖22
‖z‖22 − ‖HC [Atz]‖22 (31)
= lim t→∞
√ ‖(I − 1n11T )Atz‖ 2 2
‖z‖22 − ‖(I − 1n11T )Atz‖ 2 2
(32)
=
√ ‖(I − 1n11T )1u T 1 z‖22
‖z‖22 − ‖(I − 1n11T )1u T 1 z‖22
(33)
=
√ ‖(I1uT1 z − 1uT1 z‖22
‖z‖22 − ‖(I1uT1 z − 1uT1 z‖22 (34)
= 0 (35)
where Eqn. 31 is due to the orthogonality of DC and HC terms.

heading: B.2 PROOF OF COROLLARY 2
text: Corollary 2. Let P1,P2, · · · ,Pn be a sequence of matrix in Rn×n, and each Pk,∀k = 1, · · · , L hasAk = softmax(Pk). Then ∏L k=1Ak is also a low-pass filter. Proof. Let A = ∏L k=1Ak, then we show A satisfies the following conditions, so that A can be regarded as another self-attention matrix. Then we can conclude the proof by Theorem 1.
1) For every i = 1, · · · , n, ∑n j=1Aij = 1.
Suppose ∑n j=1Bij = 1 for every i = 1, · · · , n, then for every k = 1, · · · , L and i = 1, · · · , n,
n∑ j=1 (AkB)ij = n∑ j=1 n∑ m=1 Ak,imBmj (36)
= n∑ m=1 Ak,im  n∑ j=1 Bmj  (37) =
n∑ m=1 Ak,im = 1. (38)
By induction, for every i, ∑n j=1A1,ij = 1⇒ ∑n j=1(A2A1)ij = 1⇒ · · · ⇒ ∑n j=1Aij = 1.
2) For every i, j = 1, · · · , n,Aij > 0. Suppose Bij > 0,∀i, j, then for every k = 1, · · · , L, (AkB)ij = ∑n k=1Ak,imBmj . Since Ak,im > 0,Bmj > 0, then (AkB)ij > 0. By induction, for every i, j, A1,ij > 0⇒ (A2A1)ij > 0⇒ · · · ⇒ Aij > 0.

heading: B.3 PROOF OF THEOREM 3
text: Theorem 3. (convergence rate of SA) Let A = softmax(P ) and α = maxi,j |Pij |, where P ∈ Rn×n. Define SA(X) = AXWV as the output of a self-attention module, then
‖HC [SA(X)]‖F ≤ √
ne2α
e2α + n− 1 ‖WV ‖2‖HC [X]‖F .
In particular, when P =XWQ(XWK)T / √ d, and assume tokens are distributed inside a ball with radius γ > 0, i.e., ‖xi‖2 ≤ γ,∀i = 1, · · · , n, then α ≤ γ2‖WQW TK‖2/ √ d.
Proof. First, we write X = DC [X] +HC [X] = 1Tz +H , where DC [X] = 1zT equals to the orthogonal projection ofX to subspace span(1), andH = HC [X] represents the remaining part of the original signals.
HC [SA(X)] = (I − 11T )AXWV (39)
= ( I − 1
n 11T
) A(1zT +H)WV (40)
= ( I − 1
n 11T
) A1zTWV + ( I − 1
n 11T
) AHWV (41)
= ( I − 1
n 11T
) AHWV (42)
Therefore, ‖HC [SA(X)]‖F = ∥∥∥∥(I − 1n11T ) AHWV ∥∥∥∥ F
(43)
≤ ∥∥∥∥I − 1n11T ∥∥∥∥ 2 ‖softmax(P )‖2‖WV ‖2‖H‖F (44)
≤ √ ‖softmax(P )‖1‖softmax(P )‖∞‖WV ‖2‖H‖F (45)
= √ ‖softmax(P )‖1‖WV ‖2‖H‖F (46)
The Eqn. 45 leverages a special case of Hölder’s inequality, and the Eqn. 46 can be yielded from ‖softmax(P )‖∞ = 1. Now we need to upper bound ‖softmax(P )‖1. Suppose α = maxi,j |Pij |, then for each i = 1, · · · , n, we have the following inequality for the element with the largest value (say the j-th column):
Aij = ePij∑n t=1 e Pit ≤ e α eα + ∑ t6=j e −α = e2α e2α + (n− 1) (47)
Hence, we have ‖softmax(P )‖1 ≤ ∑ imaxjAij ≤
ne2α
e2α+(n−1) . Insert this result to Eqn. 46, we
can conclude the proof. In particular, when P =XWQ(XWK)T / √ d,
α = max i,j |Pij | = max i,j ∣∣∣∣xTi WQW TKxj√d ∣∣∣∣ . (48)
Since ‖xi‖2, ‖xj‖2 ≤ γ,∀i, j, α ≤ maxi,j‖xi‖2‖WQW TK‖2‖xj‖2/ √ d ≤ γ2‖WQW TK‖2/ √ d.

heading: C EXTENSION OF THEOREM 3

heading: C.1 MULTI-HEAD ATTENTION
text: Proposition 4. (smoothening rate with MSA) Let Ah = softmax(P h), where P h ∈ Rn×n with h = 1, · · · , H . Let α = maxHh=1 maxi,j |P hij |. Define MSA(X) = ∑H h=1A hXW hVW h O as the output of a multi-head self-attention module, then
‖HC [MSA(X)]‖F ≤ σ1σ2H √
ne2α
e2α + n− 1 ‖HC [X]‖F ,
where H is the number of heads, σ1 = maxHh=1‖W hV ‖2 and σ2 = maxHh=1‖W hO‖2.
Proof. For the h-th head, according to Theorem 3: ‖HC [SAh(X)]‖F ≤ √ ne2αh
e2αh + n− 1 ‖W hV ‖2‖HC [X]‖F , (49)
where αh = maxi,j |P hij |. Then we have:
‖HC [MSAh(X)]‖F = ∥∥∥∥∥HC [ H∑ h=1 SAh(X)W h O ]∥∥∥∥∥ F
(50)
≤ H∑ h=1 ∥∥HC [SAh(X)W hO]∥∥F (51) ≤
H∑ h=1
√ ne2αh
e2αh + n− 1 ‖W hV ‖2‖HC [X]‖2 (52)
≤ H∑ h=1
√ ne2αh
e2αh + n− 1 ‖W hV ‖2‖W hO‖2‖HC [X]‖F (53)
≤ σ1σ2H √
ne2α
e2α + n− 1 ‖HC [X]‖F , (54)
where Eqn. 51 follows from the linearity ofHC [·] and triangle inequality. Eqn. 54 can be obtained by relaxing αh, ‖W hV ‖2, and ‖W hO‖2 to α, σ1 and σ2 (Note that
√ ne2α
e2α+n−1 is monotonically increasing with α).

heading: C.2 RESIDUAL CONNECTION
text: Proposition 5. (smoothening rate with skip connection) Let Ah = softmax(P h), where P h ∈ Rn×n with h = 1, · · · , H . Let α = maxHh=1 maxi,j |P hij |. Define X ′ = MSA(X) +X as the output of a multi-head self-attention module with skip connection, then
‖HC [X ′]‖F ≤ ( 1 + σ1σ2H √ ne2α
e2α + n− 1
) ‖HC [X]‖F
where H is the number of heads, σ1 = maxHh=1‖W hV ‖2 and σ2 = maxHh=1‖W hO‖2.
Proof. By Proposition 4,
‖HC [X ′]‖F = ‖HC [MSA(X) +X]‖F (55) ≤ ‖HC [MSA(X)]‖F + ‖HC [X]‖F (56)
≤ √
ne2α
e2α + n− 1 ‖HC [X]‖F + ‖HC [X]‖F (57)
= ( 1 + σ1σ2H √ ne2α
e2α + n− 1
) ‖HC [X]‖F . (58)
Again, Eqn. 56 follows from the linearity ofHC [·] and triangle inequality.

heading: C.3 FEED-FORWARD NETWORK
text: Proposition 6. (smoothening rate with FFN) Let Ah = softmax(P h), where P h ∈ Rn×n with h = 1, · · · , H . Let α = maxHh=1 maxi,j |P hij |. Define Y = FFN(MSA(X) +X) as the output of a transformer block, then
‖HC [Y ]‖F ≤ σ3 ( 1 + σ1σ2H √ ne2α
e2α + n− 1
) ‖HC [X]‖F
where FFN : Rd → Rd represents a feed-forward network, H is the number of heads, σ1 = maxHh=1‖W hV ‖2, σ2 = maxHh=1‖W hO‖2, and σ3 = Lips(FFN) is the Lipschitz constant of FFN. In particular, σ3 = 1 + Lips(FFN) when residual connection is considered in FFN.
Proof. LetX ′ = MSA(X) +X and z′ = 1T (X ′/n) ∈ Rd, then we have
‖HC [FFN(X ′)]‖F ≤ ‖FFN(X ′)− 1FFN(z)T ‖F (59) = ‖FFN(X ′)− FFN(1zT )‖F (60) ≤ σ3‖X ′ − 1zT ‖F (61)
= σ3 ∥∥∥∥(I − 1n11T ) X ′ ∥∥∥∥ F = σ3‖HC [X ′]‖F (62)
≤ σ3 ( 1 + σ1σ2H √ ne2α
e2α + n− 1
) ‖HC [X]‖F , (63)
where Eqn. 59 follows from Lemma 7, Eqn. 60 holds because FFN operates row-wisely on feature matrix, and Eqn. 61 is due to the definition of Lipschitz constant. Finally, Eqn. 63 is yielded from Proposition 5.
Lemma 7. GivenX ∈ Rn×d, ‖HC [X]‖F ≤ ‖X − 1zT ‖F for all z ∈ Rd.
Proof. We prove the Lemma by showing that z∗ =XT1/n achieves the minimum of the optimization problem argminz‖X − 1zT ‖2F .
‖X − 1zT ‖2F = Tr(XT − z1T )(X − 1zT ) (64) = Tr(XTX)− Tr(z1TX)− Tr(XT1zT ) + Tr(zT1T1z) (65) = nzTz − 2Tr(XT1zT ) + Tr(XTX) (66)
It is easy to show the derivative in terms of z:
∇z‖X − 1zT ‖2F = 2nz − 2XT1. (67)
Therefore, z∗ = 1nX T1 achieves the minimum.

heading: D DEFERRED REMARKS ON SECTION 2.4

heading: D.1 CONNECTION WITH RANDOM WALK THEORY
text: We add that the asymptotic evolution of feature representations can be interpreted through the lens of random walk theory. We can regard self-attention mapA as probability transition matrices for a Markov chain. Since each entry is larger than zero, the Markov chain should be irreducible. This implies the Markov chain will converge into a unique stationary distribution π (Randall, 2006). Let ai denote the i-th row ofA, then for all i = 1, · · · , n we have ‖(ai)TA−πT ‖ ≤ λ‖ai−π‖, where λ ∈ (0, 1) is the mixing rate of the transition matrix A. As a consequence, liml→∞Al = 1πT yields a pure low-pass filter. When repeatedly applying this self-attention matrix to feature maps, liml→∞A
lX → 1πTX only preserves the rank-1/DC portion of the signals, which is consistent with our Theorem 1. Nevertheless, this interpretation does not bring other transformer components into consideration. And our theory further provides a concrete convergence rate with respect to the network parameters (Theorem 3).

heading: D.2 REMARKS ON DONG ET AL. (2021)
text: Here we respectfully elaborate on the hidden assumptions in the proof of the current preprint of Dong et al. (2021).
1) In the proof of Lemma A.3, Taylor expansion was used to approximate and upper bound an exponentiation. However, to let the right-hand side upper bound satisfied, we conjectured that the authors implicitly assumed Eij −Eij′ is bounded around zero. After directly communicating with the authors, they confirmed that a missed assumption here is maxi,j(Eij −Eij′) ≤ 1. 2) In the proof of Lemma A.1, to let Eqn. (8)-(9) hold, the authors may have assumedR,WV ≥ 0, where ≥ denotes entry-wise inequality. As the authors suggested, an entry-wise absolute value can be imposed to R and WV as a simple fix, without influencing their `1 and `∞ norm. However,
even after those changes, we still have difficulty walking through Eqn. (6)-(8), and we are currently communicating with the authors on this matter.
3) In the proof of Lemma A.1, we find Eqn. (12) may not be satisfied in general. We can raise the following counterexample: Since E, r,R,WV can be any matrices, we simply let D = diag(2, 3), softmax(r) = [0.8 0.2] T ,R = WV = I . Then ‖D1softmax(r)TRWV ‖1 = 4 while ‖D1‖∞‖R‖1‖WV ‖1 = 3, which disproves the claim. We conjecture that some additional prerequisite constraints on E, r might be needed here to proceed the derivation, and we are currently communicating with the authors on this matter.

heading: D.3 GENERALIZE TO OTHER ATTENTION MECHANISM
text: Our theorizing can be smoothly generalized to other attention mechanisms because our Theorem 1 and 3 do not require any prior knowledge on pre-softmax pairwise correlation P .
Logistic Attention. We refer logistic attention to the attention mechanism used in Veličković et al. (2017); Verma et al. (2018), where attention is calculated via a linear combination:
Aij = exp
( xTi uQ + x T j uK + b )∑ t exp (xiuQ + xtuK + b)
(68)
where uQ and uK are query/key parameters, b is the bias term. With the same condition in Theorem 3, we can upper bound α by |(‖uK‖2 + ‖uQ‖2)γ + b|.
L2 Distance Attention. L2 distance based attention (Kim et al., 2021) Lipschitz formulation of self-attention. The pair-wise attention can be written as follows:
Aij = exp
( −‖xTi WQ − xTjWK‖22/τ )∑ t exp ( −‖xTi WQ − xTt WK‖22/τ
) (69) whereWQ andWk are query/key weights, and τ is a scaling factor. Similar to Theorem 3, we can upper bound α by (‖WK‖2 + ‖WQ‖2)2γ2/τ .

heading: E AN AUXILIARY LEMMA FOR ATTNSCALE
text: Lemma 8. Let Ã = FAF−1 be the spectral response of attention matrix A, and parameterize a low-filter by L = F−1 diag(β, 0, · · · , 0)F . Then β∗ = 1 is the optimal solution of the following optimization problem: argminβ‖A−L‖F .
Proof. First we make simplification L = F−1 diag(β, 0, · · · , 0)F = β11T (refer to Appendix A). Then we have:
‖A−L‖2F = ‖A− β11T ‖2F = Tr(A− β11T )T (A− β11T ) (70)
= Tr ( ATA− β
n AT11T − β n 11TA+
β2
n 11T
) (71)
= β2 n Tr11T − 2β n Tr11TA+TrATA (72)
= β2 − 2β n n∑ j=1 n∑ i=1 Aij +TrA TA (73)
= β2 − 2β +TrATA (74) From Eqn. 74, β∗ = 1 achieves the minimum of the objective function.

heading: F MORE ON VISUALIZATION

heading: F.1 DETAILS ON FIGURE 1
text: To verify our Theorem 3, we depict the high-frequency intensity of each layer’s output and its theoretical upper bound. Our visualization is based on the official checkpoint of 12-layer DeiT-S. Since
training a ViT without either FFN or residual connection will certainly cause failure, we remove these components directly from the pre-trained model to illustrate the effects of different components. We use logarithmic scale for the purpose of better view. LetXl denote the output of the l-th layer, andX0 be the initial inputs. For red line, we directly calculate log(‖HC [Xl]‖F /‖X0‖F ) at each layer. For blue line, we first obtain the coefficient γl in Section 2.2 and 2.3 with respect to network parameters
(e.g., we can compute γl = √ ne2α
e2α+n−1‖WV ‖2 for attention only architecture). Then we estimate the upper bound by γl‖HC [Xl−1]‖F and apply the logarithm by log(γl‖HC [Xl−1]‖F /‖X0‖F ). To summarize, one can see without residual connection, the first two sub-figures imply an exponential convergence rate, which is consistent with our Theorem 3.

heading: F.2 MORE VISUALIZATION ON SPECTRUM
text: In this appendix, we provide more visualization on the spectrum of attention map to validate our Theorem 1. We compute the spectrum of attention map A for both Fig. 2 and Fig. 5 in the following way. By regardingA as a linear filter, its Fourier-domain response is another linear kernel Λ = FAF−1. When Λ is applied to a spectrum x̃ = Fx of signals x, the i-th frequency response will be Λix̃, where Λi is the i-th row of Λ. Hence, we can use ‖Λi‖2 to evaluate the spectral response intensity of the i-th frequency band. Below we provide a complete spectral visualization of attention maps computed from a random sample in ImageNet validation set.

heading: F.3 DETAILS ON SIMILARITY CURVES (FIGURE 4)
text: In Fig. 4, we visualize the cosine similarity of attention maps and feature maps to show the effectiveness of our AttnScale and FeatScale on 24-layer DeiT, respectively. We follow the definition in Zhou et al. (2021a) to compute the cosine similarity metric for attention maps. Instead of measuring cross-layer similarity, we calculate average cross-patch similarity at the same layer. Given the layer index l and corresponding attention mapsA(l,h) ∈ Rn×n, the cosine similarity can be computed by:
M lattn = 2
n(n− 1)H H∑ h=1 n∑ i=1 n∑ j=i+1 ∣∣∣A(l,h)T:,i A(l,h):,j ∣∣∣∥∥∥A(l,h):,i ∥∥∥ 2 ∥∥∥A(l,h):,j ∥∥∥ 2 , (75)
whereA(l,h):,i denotes the i-th column ofA (l,h), and H is the number of heads. The cosine similarity between i-th and j-th column of A(l,h) measures how the contribution of one token (say the i-th token) varies from the other (say the j-th token). We average the similarity between every pair of tokens’ attention map (excluding the self-to-self similarity) and every attention head. We refer interested readers to Zhou et al. (2021a) for more details.
We use the similar metric to compute similarity for feature maps. Following Gong et al. (2021), we compute pair-wise cosine similarity between every two different tokens. Formally, given the layer index l, and its outputX(l) ∈ Rn×d, the cosine similarity is estimated by:
M lfeat = 2
n(n− 1) n∑ i=1 n∑ j=i+1 ∣∣∣X(l)Ti,: X(l)j,: ∣∣∣∥∥∥X(l)i,: ∥∥∥ 2 ∥∥∥X(l)j,: ∥∥∥ 2 , (76)
whereX(l)i,: denotes the i-th row ofX (l). The cosine similarity between between i-th and j-th row of X(l) measures how similar the feature representations of two tokens are. Likewise, we average the similarity between every pair of tokens’ features except for the self-to-self similarity. More details can found in Gong et al. (2021). We additionally provide a visualization of these two metrics for 12-layer DeiT in Fig. 12.

heading: G DEFERRED EXPERIMENTS AND MODEL INTERPRETATION

heading: G.1 FINE-TUNING EXPERIMENTS
text: Our deferred fine-tuning experiment with CaiT (Touvron et al., 2021b) results are presented in Table 3. Different from trining scratch, we fine-tune CaiT with AttnScale and FeatScale parameters from the pre-trained models for 60 epochs following Gong et al. (2021). For a fair comparison, we simultaneously train a plain CaiT for another 60 epochs. During fine-tuning, we reduce learning rate to 5× 10−5 and weight decay to 5× 10−4. All other hyper-parameters and training recipe are kept consistent with the original paper (Touvron et al., 2021b).
G.2 VISUALIZATION AND INTERPRETATION OF ATTNSCALE
In this subsection, we provide visualization to interpret our AttnScale and further support our experiments. Ê In Fig. 6 we visualize the learned weights of our AttnScale. We observe conclude our AttnScale are successfully trained to amplify the high-pass component. We also find when layer
index goes larger, the scaling weights turns larger to prevent attention collapse at deeper layer. Ë We also compare the attention map produced by AttnScale with those produced by original DeiT. We observe from Fig. 8 that our AttnScale can extract more salient and higher contrastive attention than vanilla DeiT, which indicates our AttnScale possesses higher capability to distinguish tokens from larger variety of attention schemes. Ì To be more objective, we plot the spectrum of a 24-layer DeiT’s attention maps with/without our AttnScale in Fig. 9 and 10. The visualization procedure has been elaborated in Sec. F.2. We find that attention maps from AttnScale enjoy richer filtering diversity, capable of performing high-pass (row 2, column 3) and band-pass (row 12, column 5) filtering, instead of only low-pass filtering (see Fig. 9).
G.3 VISUALIZATION AND INTERPRETATION OF FEATSCALE
In this subsection, we provide visualization to interpret our FeatScale. Ê We plot the scaling weights of FeatScale in Fig. 7. We observe that the re-weighting factors learned for high-frequency components t is consistently larger than the weights for the DC term s, which indicates our FeatScale is successfully trained to elevate high-frequency features against the dominance of DC component. Similarly, the gap between s and t becomes huger when going deeper. Ë We also demonstrate the proportion of feature maps’ high-frequency component in Fig. 11 for both 12 (lower one) / 24(upper one) -layer DeiT. The proportion value is calculated by ‖HC [X]‖F /‖X‖F . We find high-frequency signals diminish quickly at deeper layer, and 24-layer DeiT suffers from a faster pace. Our FeatScale is effective to keep the high-frequency signals stand for both 12-layer and 24-layer DeiT.

Review summaries:
This paper explores the reasons why ViTs cannot go deeper. The article provides clear and solid proofs, clarifying the rank collapse in attention matrix via Fourier analysis. Meanwhile, the authors suggest two types of possible solutions, AttnScale and FeatScale, which are useful when scaling up the ViTs architecture. The experiments on DeiT and CaiT show that the adjustment on attention helps avoiding rank collapse when going deeper, holding high-frequency information.  

This paper provides a viewpoint of spectrum domain to show self-attention amounts to a low-pass filter and will cause feature maps to only preserve direct-current components as depth increases. Based on this viewpoint, it proposes two technique named AttnScale, FeatScale to reweight low-pass, high-pass components with minimal parameters. Experiments are conducted on ImageNet. 

This paper first analysis the difficulty of going deeper using transformer models. The challenge is that the mechanism of self-attention is applying a low-pass filter.  Keep stacking self-attention layers will lose the feature expressive power and only preserve the DC bias.  To alleviate this issue, this paper propose two modules to make the MHA to be an all pass filter. First module AttenScale combines the high-pass component and FeatScale reweights the featuremap. The experimental results demonstrate the improvement on accuracy for DeiT and CaiT when adopting the proposed method.