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