Scaling the Depth of Vision Transformers via the Fourier Domain Analysis

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 . 1 INTRODUCTION . 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 . 2 WHY VIT CAN NOT GO DEEPER ? . 2.1 NOTATION AND PRELIMINARIES . 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 . 2.2 SELF-ATTENTION IS A LOW-PASS FILTER . 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 .

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