S CALING THE D EPTH OF V ISION T RANSFORMERS VIA THE F OURIER D OMAIN A NALYSIS
Anonymous authors
Paper under double-blind review
A BSTRACT
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 ﬁrst establish an analytic framework to investigate ViT from the spectrum domain.
We show that the self-attention mechanism inherently amounts to a low-pass ﬁlter, which indicates when ViT scales up its depth, excessive lowpass ﬁltering 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 ﬁrst technique, termed AttnScale , decomposes a self-attention block into low-pass and high-pass components, then rescales and combines these two ﬁlters 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 efﬁcient, hyperparameter-free, and can effectively avoid attention collapse and patch uniformity caused by low-pass ﬁltering.
Our experiments demonstrate our proposed methods consistently help ViT beneﬁt from deeper architectures, bringing ≥ 1.0% performance gain with little parameter overhead.
In addition to the baseline model, our techniques are also successfully applied to ViT variants.
All our codes and pre-trained models will be publicly released upon acceptance.
I NTRODUCTION
Transformers have achieved a 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 classiﬁcation 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 signiﬁcantly 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 diversiﬁcation (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 ﬁrst 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 ﬁlter, 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 ﬁnding 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 ﬁltering (Section 2.3).
Built on the aforementioned analysis framework in the Fourier domain, we propose two novel techniques, to mitigate the low-pass ﬁltering effect of self-attention and effectively scale up the depth of ViTs.
The ﬁrst technique, termed Attention Scaling ( AttnScale ), directly manipulates on the calculated attention map to enforce an all-pass ﬁlter (Section 3.1).
It decomposes the self-attention matrix into a low-pass ﬁlter plus a high-pass ﬁlter, then adopts a learnable weight to adaptively amplify the effect of high-pass ﬁlter.
The second technique, termed Feature Scaling ( FeatScale ), hinges on feature maps to re-weight different frequency bands separately (Section 3.2).
It employs trainable coefﬁcients 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 ﬁrst rigorous theoretical analysis of ViT from the spectral domain.
We characterize the low-pass ﬁltering effect of cascading MSAs, which connects to the recent empirical ﬁndings 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 efﬁcient, 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 and CaiT with 1 . 0% and 0 .
3% performance gain without whistles and bells.
W HY V I T C ANNOT G O D EEPER ? 2. N OTATION AND P RELIMINARIES
We begin by introducing our notations.
Let X ∈ R n × d denote the feature matrix, where n is the number of samples, and d is the feature dimension.
Let x i ∈ R d , ∀ i = 1 , · · · , n , the i -th row of X , denote the feature vector of the i -th sample, and z j ∈ R n , ∀ 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, x i denotes the ﬂatten version of the i -th patch embedding ( d = patch width × patch height ) , and z j 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 classiﬁcation.
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):
[Equation]
where W K ∈ R d × d k , W Q ∈ R d × d q , W V ∈ R
d × d are the key, query, and value weight matrices, respectively, √ d here denotes a scaling factor, and softmax( · ) operates on X row-wisely.
Multi-Head
Self-Attention (MSA) involves a group of SA heads and combines their outputs through a linear
[Figure]
L o g R e l a t i v e
D i s t a n c e
MHA
Actual Theorem  12 Layer Index
MHA + FFN Actual
Theorem  12 Layer Index
MHA + FFN
[Figure]
Res
[Figure]
Theorem projection (Vaswani et al., 2017):
[Equation]
where the subscripts denote the SA head number, H is the total number of SA heads, and W O ∈ R Hd ×
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:
[Equation]
Fourier Analysis The main analytic tool in this paper is Fourier transform.
Denote F : R n → C n be the Discrete Fourier Transform (DFT) with the Inverse Discrete Fourier Transform (IFT)
F − 1 : C n → R n .
Applying F to a ﬂatten image signal x is equivalent to left multiplying a DFT matrix, whose rows are the Fourier basis f k = (cid:2) e 2 π j( k − 1) ·
[Equation]
where k denotes the k -th row of DFT matrix, and j is the imaginary unit.
Let ˜ z
= F z be the spectrum of z , and ˜ z dc ∈ C , ˜ z hc ∈ C n − 1 take the ﬁrst element and the rest elements of ˜ z , respectively.
Deﬁne DC [ z ] = ˜
z dc f 1 ∈ C n as the Direct-Current (DC) component of signal z , and
HC [ z ] =
[ f· · · f n ] ˜ z hc ∈ C n the complementary high-frequency component.
In signal processing, a low-pass ﬁlter 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 ﬁlter as a particular type of ﬁlters that only preserve the DC component, while diminishing the remaining high-frequency component.
To be more precise, we deﬁne low-pass ﬁlters in Deﬁnition 1.
Deﬁnition 1.
Given an endomorphism f : R n → R n with f t denoting applying f for t times, f is a low-pass ﬁlter if and only if for all z ∈ R n :
[Equation]
Deﬁnition 1 reveals the nature of low-pass ﬁlters: 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. S ELF -A TTENTION I S L OW -P ASS F ILTER
In this subsection, we will give theoretical justiﬁcation on self-attention in terms of its spectral-domain effect.
Our main result is that self-attention is constantly a low-pass ﬁlter, 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 ﬁlter independent of the input token features or key/query matrices.
Theorem 1.
(SA matrix is a low-pass ﬁlter)
Let A = softmax( P ) , where P ∈ R n × n .
Then A must be a low-pass ﬁlter.
For all z ∈ R n , lim t →∞ (cid:107)HC [ A t z ] (cid:107) 2 / (cid:107)DC [ A t z ] (cid:107) 2 = 0 .
[Figure]
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 (Veliˇckovi´c et al., 2017), or L distance (Kim et al., 2021), the resulting attention matrix is always a low-pass ﬁlter.
One can see consecutively applying self-attention matrix simulates the process of
ViT’s forward propagation.
As the layer number increases inﬁnitely, the ﬁnal output will only keep the DC bias, and
ViT loses all the feature expressive power.
Corollary 2.
Let P 1 , P 2 , · · · , P n be a sequence of matrix in R n × n , and each P k , ∀ k
= 1 , · · · , L has A k = softmax( P k ) .
Then (cid:81)
Lk
=
A k is also a low-pass ﬁlter.
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 ﬁlter 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 ﬁlters, 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 α = max i,j | P ij | , where P ∈ R n × n .
Deﬁne SA( X ) =
AXW
V as the output of a self-attention module, then
[Equation]
In particular, when P = XW Q ( XW K ) T / √ d , and assume tokens are distributed inside a ball with radius γ > 0 , i.e., (cid:107) x
i (cid:107) 2 ≤ γ, ∀ i = 1 , · · · , n , then α ≤ γ 2 (cid:107) W Q W TK (cid:107) 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 (cid:107) W V (cid:107) 2 (cid:113) ne 2 α e 2 α + n − 1 .
When (cid:107) W V (cid:107) 2 (cid:113) ne 2 α e 2 α + n − 1 < 1 , HC [SA( X )] converges to zero exponentially.
We note that, no matter how attention is computed or signals are initialized, since (cid:113) ne 2 α e 2 α + n − 1 is bounded by √ n , (cid:107) W V (cid:107) 2 < 1 / √ n will deﬁnitely cause a monotonically decreasing high-frequency component.
When dot-product attention is adopted, a sufﬁcient condition that (cid:107)HC [ X ] (cid:107) F decreases to zero within logarithmic time is γ 2 (cid:107) W Q W TK (cid:107) 2 / √ d + log (cid:107) W V (cid:107) 2 ≤ log n − 1 n / 2 .
2. E XISTING M ECHANISMS THAT C OUNTERACT L OW -P ASS F ILTERING
In this section, we take other building blocks of ViT into consideration.
We will justify whether
Multi-Head Self-Attention (MSA), Feed-Forward Network (FFN), and residual connections can effectively alleviate the low-pass ﬁltering drawbacks.
All the derivations follow from Theorem 3, and some proof ideas are borrowed from Dong et al.
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 ) =
(cid:80) Hh =1 SA( X ) W hO .
We show by Proposition 4 in
Appendix C.1 that the convergence rate turns to σ 1 σ 2 H (cid:113) ne 2 α e 2 α + n − 1 , where H is the number of heads, σ 1 = max
Hh =1 (cid:107) W hV (cid:107) 2 and σ 2 = max
Hh =1 (cid:107) W hO (cid:107) 2 .
One can see MSA can only slow down the convergence up to a constant σ 2 H , which does not root out the problem.
Does residual connection beneﬁt?
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 σ 2 H (cid:113) ne 2 α e 2 α + n − 1 to 1 + σ 1 σ 2 H (cid:113) ne 2 α e 2 α + n − 1 > 1 .
Refer to Proposition 5 in Appendix C.2.
Does FFN make a change?
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 σcontributes a σ 3 (cid:16) 1 + σ 1 σ 2 H (cid:113) ne
α e 2 α + n − 1 (cid:17) 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 help preserve the high-frequency signals, it does not change the fact that MSA block as a whole only possesses the representational power of low-pass ﬁlters.
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.
2. C ONNECTIONS TO E XISTING T HEORETIC U NDERSTANDING W ORKS
It is known that Graph Convolutional Networks (GCN) are not more than low-pass ﬁlters (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 ﬁlter.
Combining with our theoretical derivation, one can see the root reason is that both graph Laplacian matrices and selfattention matrices consistently own a ﬁxed 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 ﬁltering.
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 speciﬁes 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 deﬁcient, 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 (Veliˇckovi´c et al., 2017) and L2 distance (Kim et al., 2021).
See Appendix D.3.
A TTN S CALE & F EAT S CALE : S CALING FROM THE F OURIER -D OMAIN 3. A TTN S CALE : M AKE A TTENTION AN A LL -P ASS F ILTER
As we discussed in Section 2.2, self-attention matrix can only perform low-pass ﬁltering, which narrows the ﬁlter 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 ﬁltering and produce all-pass ﬁlters.
AttnScale decomposes the self-attention matrix to a low-pass ﬁlter plus a high-pass ﬁlter, and introduces a trainable parameter to rescale the high-pass ﬁlter to match the magnitude with the low-pass component.
Formally, let A denote a self-attention matrix.
To decompose a low-pass ﬁlter from A , we ﬁnd the largest possible low-pass ﬁlter 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 = 11 T /n from A and take the complementary part as the high-pass ﬁlter.
Afterward, we can rescale the
[Figure]
+ + ! [Figure] [Figure] +
[Figure]
+
[Figure]
+ + ! "
[Figure] [Figure] [Figure] [Figure]
Refer to Appendix F.3 for details.
high-pass component of the ﬁlter, 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:
[Equation]
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 ﬁlters: 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.
3. F EAT S CALE : R EWEIGHT H IGH -F REQUENCY S IGNALS
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 ﬁrst 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
[Equation] [Equation]
[Equation]
(12) where s l ∈ R d and t l ∈ R
d are learnable parameters to perform channel-wise re-weighting.
We initialize s l and t l 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 matrix X , and HC [ X ] can be efﬁciently computed by X − DC
[ X ] .
3. D ISCUSSION
We have proposed two methods to facilitate the deeper stacking of ViT MSA modules, and discussed their motivations and strengths from our perspective of ﬁltering 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 ﬁlter, 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, modiﬁed attention maps aggregate high-pass components from other heads and is endowed with richer ﬁltering 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 ( 11 T /n ) with the complementary residual map that encodes diverse patterns in a self-attention matrix.
By re-weighting the residual map, the diversiﬁed patterns can be ampliﬁed, 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-ﬁgure 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 ﬁltering, and the remaining DC bias signiﬁes 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.
R ELATED W ORK
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.
ﬁrst adopts a pure transformer architecture (ViT) for image classiﬁcation.
The follow-up works 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., 2021; 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 efﬁciently extract multi-scale features.
To dismiss ViT from the heavy reliance on large-scale dataset pre-training, Touvron et al.
(2021a)
; Yuan et al.
propose knowledge distillation and progressive tokenization for data-efﬁcient training.
Despite impressive effectiveness, most of these model are only based on relatively shallow ViT backbones with a dozen of MSA blocks or so.
Advances in deep ViTs.
Building deeper ViTs has arisen many interests.
Zhou et al.
(2021a)
ﬁrst 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.
[Table]
Role of depth in NNs.
Discussing the importance of deep structures in Neural Networks (NNs) is an over 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 beneﬁts 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 ﬁltering 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).
E XPERIMENTS
In this section, we report experiment results to validate our proposed methods.
First, we testify 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.
5. H OW C AN A TTN S CALE & F EAT S CALE B ENEFIT D EEP V I T?
Experiment Settings.
In this subsection, we intend to testify our models are beneﬁcial to various
ViT backbones with different depth settings and training modes.
We choose DeiT (Touvron et al., 2021a) as our ﬁrst 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.
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.
In addition to training from scratch, we also investigate the ﬁne-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.
Speciﬁcally, AttnScale brings less than 100/ [Table] 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 ﬁne-tuning setting, as we will show in Appendix G.1, only tens of epoch’s ﬁne-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 ﬁltering properties (compare Fig. 9 with Fig. 10), and 2) more high-frequency data can be preserved (refer to Fig. 11).
5. C OMPARED WITH SOTA M ODELS
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.
C ONCLUSION
In this paper, we investigate the scalability issue with ViT and propose two practical solutions via Fourier domain analysis.
Our theoretical ﬁndings indicate Multi-Head Self-Attention (MSA) inherently performs low-pass ﬁltering 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 ﬁltering bottleneck by adaptively scaling high-pass ﬁlter component and high-frequency signals, respectively.
Our experiments also validate the effectiveness of our methods.
Both techniques can boost various ViT backbones by a signiﬁcant performance gain.
Grounded with our theoretical framework, interesting directions for further work include designing parameter regularizations and spectrum-speciﬁc normalization layers.
R EPRODUCIBILITY S TATEMENT
To ensure reproducibility, we use the Machine Learning Reproducibility Checklist v2.0, Apr. 7(Pineau et al., 2021).
An earlier version of this checklist (v1.2) was used for NeurIPS 2019.
• For all models and algorithms presented, – A clear description of the mathematical settings, algorithm, and/or model.
We clearly describe all of the settings, formulations, and algorithms in Section 2 and 3.
– A clear explanation of any assumptions.
We do not make assumptions.
– An analysis of the complexity (time, space, sample size) of any algorithm.
We provide a brief piece of analysis at the end of Section 3.1.
• For any theoretical claim , – A clear statement of the claim.
We make clear and rigorous statements in Section 2.2 and Appendix C, E. – A complete proof of the claim.
Complete proofs can be found in Appendix B, C, E. • For all datasets used, check if you include: – The relevant statistics, such as number of examples.
We only use widely adopted dataset ImageNet.
All the information can be found at its ofﬁcial website.
– The details of train/validation/test splits.
We follow the standard splits.
– An explanation of any data that were excluded, and all pre-processing step.
We did not exclude any data.
We ported the data loading pipeline from DeiT repository.
– A link to a downloadable version of the dataset or simulation environment.
We have attached our code within the supplementary materials.
– For new data collected,a complete description of the data collection process, such as instructions to annotators and methods for quality control.
We do not collect or release new datasets.
• For all reported experimental results , check if you include: – The range of hyper-parameters considered, method to select the best hyperparameter conﬁguration, and speciﬁcation of all hyper-parameters used to generate results.
Most of our hyper-parameters inherit from DeiT repository.
We provide hyper-parameter and training recipe differences in Section 5.1.
– The exact number of training and evaluation runs.
Same with DeiT, CaiT, etc., we only train and evaluate once with a ﬁxed seed.
– A clear deﬁnition of the speciﬁc measure or statistics used to report results.
We use standard image classiﬁcation loss and metric, same with DeiT, etc.
– A description of results with central tendency (e.g. mean) & variation (e.g. error bars).
We only report the test accuracies run by once.
– The average runtime for each result, or estimated energy cost.
We do not report the running time or energy cost.
– A description of the computing infrastructure used.
All of our experiments are conducted on NVIDIA RTX A6000.
A M ORE P RELIMINARIES ON F OURIER A NALYSIS
In this appendix, we provide more preliminary knowledge about Fourier analysis.
Here, we only consider discrete Fourier transform on real-value domain F : R n → C n .
A Discrete Fourier transform (DFT) can be written in a matrix form as below:
[Equation]
 
[Equation]
(13) and its inverse discrete Fourier transform is DF T − 1 = DF T .
In this paper, we regard matrices as multi-channel signals.
For example, X ∈ R n × 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 ( x 1 ) · · · F ( x d )]
= DF T · X .
Hereby, we can simply operators DC [ · ] and HC [ · ] using the matrices in Eqn. 13.
By deﬁnition, we can write DC [ · ] as below:
[Equation] [Equation]
namely DC [ · ] = 11 T /n .
Conversely, we can write HC [ · ] as:
[Equation]
[Equation] [Equation]
which indicates HC [ · ] = I − 11 T /n .
We will frequently use these derivations later in the proofs.
B D EFERRED P ROOFS B. P ROOF OF T HEOREM
Theorem 1.
(SA matrix is a low-pass ﬁlter)
Let A = softmax( P ) , where P ∈ R n × n .
Then A must be a low-pass ﬁlter.
For all z ∈ R n ,
[Equation]
t
[Equation]
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 ( A ij > 0 , ∀ i, j ).
Besides, for all i = 1 , · · · , n , (cid:80) nj =1 A
i,j = 1 .
Therefore, A 1 = 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 = P JP − 1 :
A = [ v· · · v n ] (cid:124) (cid:123)(cid:122) (cid:125) P   λ
J ( λ 2 ) ...
J ( λ s ) 
 (cid:124) (cid:123)(cid:122) (cid:125)
J   u T...
u Tn 
 (cid:124) (cid:123)(cid:122)
[Equation]
where the Jordan block J ( λ ) can be written as
[Equation]
Applying A to z for t times can be written as A t z which is equivalent to:
[Equation]
Let f ( x ) = x
t , then A t =
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
[Equation]
Therefore, on the diagonal number m ≤ min( t, k − 1) above the main diagonal stands:
[Equation]
[Equation]
For arbitrary m ≤ k − 1 and | λ | < 1 ,
[Equation]
t
[Equation] [Equation]
Recall that λ 1 = 1 and according to the deﬁnition of Jordan canonical form, v 1 = 1 .
By Eqn. 25:
[Equation]
Plug the result from Eqn. 29 into the original limit:
[Equation]
where Eqn. 31 is due to the orthogonality of DC and HC terms.
B. P ROOF OF C OROLLARY
Corollary 2.
Let P 1 , P 2 , · · · , P n be a sequence of matrix in R n × n , and each P k , ∀ k
= 1 , · · · , L has A k = softmax( P k ) .
Then (cid:81)
Lk
=
A k is also a low-pass ﬁlter.
Proof.
Let A = (cid:81)
Lk
=
A k , then we show A satisﬁes the following conditions, so that A can be regarded as another self-attention matrix.
Then we can conclude the proof by Theorem 1.
For every i = 1 , · · · , n , (cid:80) nj =1 A ij = 1 .
Suppose (cid:80) nj =1 B ij = 1 for every i = 1 , · · · , n , then for every k = 1 , · · · , L and i = 1 , · · · , n , n (cid:88)
j =
( A k B ) ij = n (cid:88)
j = n (cid:88) m =
A k,im B mj (36) = n (cid:88) m =  
A k,im   n (cid:88) j =
B mj   
 (37) = n (cid:88) m =
A k,im = 1 .
By induction, for every i , (cid:80) nj =
A 1 ,ij = 1 ⇒ (cid:80) nj =1 ( A 2 A 1 ) ij = 1 ⇒ · · · ⇒ (cid:80) nj
=1 A ij = 1 .
For every i
, j = 1 , · · · , n , A ij > 0 .
Suppose B ij > 0 , ∀ i, j , then for every k = 1 , · · · , L , ( A k B ) ij = (cid:80) nk =
1 A k,im B mj .
Since
A k,im > 0 , B mj > 0 , then ( A k B ) ij > 0 .
By induction, for every i, j , A 1 ,ij > 0 ⇒ ( A 2 A 1 ) ij > 0 ⇒ · · · ⇒ A ij > 0 .
B. P ROOF OF T HEOREM
Theorem 3.
(convergence rate of SA) Let A = softmax( P ) and α = max i,j | P ij | , where P ∈ R n × n .
Deﬁne SA( X ) =
AXW
V as the output of a self-attention module, then
[Equation]
In particular, when P = XW Q ( XW K ) T / √ d , and assume tokens are distributed inside a ball with radius γ > 0 , i.e., (cid:107) x
i (cid:107) 2 ≤ γ, ∀ i = 1 , · · · , n , then α ≤ γ 2 (cid:107) W Q W T K (cid:107) 2 / √ d .
Proof.
First
[Equation]
we write X =
DC
[ X ] + HC [ X ] = 1 T z + H , where DC [ X ] = 1 z T equals to the orthogonal projection of X to subspace span( 1 ) , and H =
HC [ X ] represents the remaining part of the original signals.
[Equation]
Therefore,
[Equation]
The Eqn.
45 leverages a special case of Hölder’s inequality, and the Eqn.
46 can be yielded from (cid:107) softmax( P ) (cid:107) ∞
= 1 .
Now we need to upper bound (cid:107) softmax( P ) (cid:107) 1 .
Suppose α = max i,j | P ij | , then for each i = 1 , · · · , n , we have the following inequality for the element with the largest value (say the j -th column):
[Equation]
[Equation]
Hence, we have (cid:107) softmax( P ) (cid:107) 1 ≤
(cid:80) i max j A ij ≤ ne
[Equation]
e 2 α +( n − 1) .
Insert this result to Eqn.
46, we can conclude the proof.
In particular, when P = XW Q ( XW K ) T / √ d ,
[Equation]
(cid:12)(cid:12)(cid:12)(cid:12)
[Equation]
Since (cid:107) x i (cid:107) 2 , (cid:107) x j (cid:107) 2 ≤ γ, ∀ i, j , α ≤ max i,j (cid:107) x i (cid:107) 2 (cid:107) W Q W TK (cid:107) 2 (cid:107) x j (cid:107) 2 / √ d ≤ γ 2 (cid:107) W Q W TK (cid:107) 2 / √ d .
C E XTENSION OF T HEOREMC. M ULTI -H EAD A TTENTION
Proposition 4.
(convergence rate with MSA) Let A h = softmax( P h ) , where P h ∈ R n × n with h = 1 , · · · , H .
Let α = max Hh =1 max i,j | P hij | .
Deﬁne MSA( X ) =
(cid:80)
Hh =1 A h XW hV W hO as the output of a multi-head self-attention module, then
[Equation]
[Equation]
where H is the number of heads, σ 1 = max H h =1 (cid:107) W h V (cid:107) 2 and σ 2 =
max H h
=1 (cid:107)
W h O (cid:107) 2 .
Proof.
For the h -th head, according to Theorem 3:
[Equation] [Equation]
where α h = max i,j | P hij | .
Then we have:
[Equation]
where Eqn. 51 follows from the linearity of HC [ · ] and triangle inequality.
Eqn. 54 can be obtained by relaxing α h , (cid:107) W hV (cid:107) 2 , and (cid:107) W hO (cid:107) 2 to α , σ 1 and σ 2 (Note that (cid:113) ne
α e 2 α + n − 1 is monotonically increasing with α ).
C. R ESIDUAL C ONNECTION
Proposition 5.
(convergence rate with skip connection) Let A h = softmax( P h ) , where P h ∈ R n × n with h = 1 , · · · , H .
Let α = max Hh =1 max i,j | P hij | .
Deﬁne X (cid:48)
=
MSA( X ) + X as the output of a multi-head self-attention module with skip connection, then
[Equation]
where H is the number of heads, σ 1 = max Hh =1 (cid:107) W hV (cid:107) 2 and σ 2 = max
Hh =1 (cid:107) W hO (cid:107) 2 .
Proof.
By Proposition 4,
[Equation] [Equation]
Again, Eqn. 56 follows from the linearity of HC [ · ] and triangle inequality.
C. F EED -F ORWARD N ETWORK
Proposition 6.
(convergence rate with FFN) Let A h = softmax( P h ) , where P h ∈ R n × n with h = 1 , · · · , H .
Let α = max Hh =1 max i,j | P hij | .
Deﬁne Y = FFN(MSA( X ) + X ) as the output of a transformer block, then
[Equation]
where FFN : R d → R d represents a feed-forward network, H is the number of heads, σ 1 = max
Hh =1 (cid:107) W hV (cid:107) 2 , σ 2 = max Hh =1 (cid:107) W hO (cid:107) 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.
Let X (cid:48)
=
MSA( X ) + X and z (cid:48)
= 1 T ( X (cid:48) /n )
∈ R d , then we have
[Equation]
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 deﬁnition of Lipschitz constant.
Finally, Eqn. 63 is yielded from
Proposition 5.
Lemma 7.
Given X ∈ R n × d , (cid:107)HC [ X ] (cid:107) F ≤ (cid:107) X −
[Equation]
z T
[Equation]
F for all z ∈
[Equation]
d
[Equation]
Proof.
We prove the Lemma by showing that z ∗ =
X T 1 /n achieves the minimum of the optimization problem
arg min z (cid:107) X − 1 z T (cid:107) 2 F .
[Equation]
It is easy to show the derivative in terms of z :
[Equation]
Therefore, z ∗ = 1 n X T 1 achieves the minimum.
D D EFERRED R EMARKS ON S ECTION 2. D. C ONNECTION WITH R ANDOM W ALK T HEORY
We add that the asymptotic evolution of feature representations can be interpreted through the lens of random walk theory.
We can regard self-attention map
A 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 a i denote the i -th row of A , then for all i = 1 , · · · , n we have (cid:107) ( a i ) T A − π T (cid:107) ≤ λ (cid:107) a i − π (cid:107) , where λ ∈ (0 , 1) is the mixing rate of the transition matrix A .
As a consequence, lim l →∞ A l = 1 π T yields a pure low-pass ﬁlter.
When repeatedly applying this self-attention matrix to feature maps, lim l →∞ A l X → 1 π T X 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).
D. R EMARKS ON D ONG ET AL . (2021)
Here we respectfully elaborate on the hidden assumptions in the proof of the current preprint of Dong et al.
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 satisﬁed, we conjectured that the authors implicitly assumed E ij − E ij (cid:48) is bounded around zero.
After directly communicating with the authors, they conﬁrmed that a missed assumption here is max i,j ( E ij − E ij (cid:48) ) ≤ 1 . 2)
In the proof of Lemma A.1, to let Eqn.
(8)-(9) hold, the authors may have assumed R , W V ≥ 0 , where ≥ denotes entry-wise inequality.
As the authors suggested, an entry-wise absolute value can be imposed to R and W V as a simple ﬁx, without inﬂuencing their (cid:96) 1 and (cid:96) ∞ norm.
However,
even after those changes, we still have difﬁculty walking through Eqn.
and we are currently communicating with the authors on this matter.
In the proof of Lemma A.1, we ﬁnd Eqn.
(12) may not be satisﬁed in general.
We can raise the following counterexample: Since E , r , R , W V can be any matrices, we simply let D = diag(2 , 3) , softmax( r ) =
2] T , R = W V
= I .
Then (cid:107) D 1 softmax( r ) T RW V (cid:107)
=while (cid:107) D 1 (cid:107) ∞ (cid:107) R (cid:107) 1 (cid:107) W V (cid:107) 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.
D. G ENERALIZE TO O THER A TTENTION M ECHANISM
Our theorizing can be smoothly generalized to other attention mechanisms because our Theoremand 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 Veliˇckovi´c et al.
(2017); Verma et al.
(2018), where attention is calculated via a linear combination:
[Equation]
(68) where u Q and u K are query/key parameters, b is the bias term.
With the same condition in Theorem 3, we can upper bound α by | ( (cid:107) u K (cid:107)
+ (cid:107) u Q (cid:107) 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:
[Equation]
69) where W Q and W k are query/key weights, and τ is a scaling factor.
Similar to Theorem 3, we can upper bound α by ( (cid:107) W K (cid:107) 2 + (cid:107) W Q (cid:107) 2 ) 2 γ 2 /τ .
E A N A UXILIARY L EMMA FOR A TTN S CALE
Lemma 8.
Let ˜
A = F A F − 1 be the spectral response of attention matrix A , and parameterize a low-ﬁlter by L = F − 1 diag( β, 0 , · · · , 0) F .
Then β ∗ = 1 is the optimal solution of the following optimization problem: arg min β (cid:107)
A − L (cid:107) F .
Proof.
First we make simpliﬁcation L = F − 1 diag( β, 0 , · · · , 0) F = β 11 T (refer to Appendix A).
Then we have:
[Equation]
From Eqn.
74, β ∗ = 1 achieves the minimum of the objective function.
F M ORE ON V ISUALIZATION F. D ETAILS ON F IGURE
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 ofﬁcial 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.
Let X l denote the output of the l -th layer, and Xbe the initial inputs.
For red line, we directly calculate log( (cid:107)HC [ X l ] (cid:107) F / (cid:107) X 0 (cid:107) F ) at each layer.
For blue line, we ﬁrst obtain the coefﬁcient γ l in Section 2.2 and 2.3 with respect to network parameters (e.g., we can compute γ l = (cid:113) ne 2 α e 2 α + n − 1 (cid:107) W V (cid:107) 2 for attention only architecture).
Then we estimate the upper bound by γ l (cid:107)HC [ X l − 1 ] (cid:107) F and apply the logarithm by log( γ l (cid:107)HC [ X l − 1 ] (cid:107) F / (cid:107) X 0 (cid:107) F ) .
To summarize, one can see without residual connection, the ﬁrst two sub-ﬁgures imply an exponential convergence rate, which is consistent with our Theorem 3.
F. M ORE V ISUALIZATION ON S PECTRUM
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 regarding A as a linear ﬁlter, its Fourier-domain response is another linear kernel Λ = F A F − 1 .
When Λ is applied to a spectrum ˜ x = F x of signals x , the i -th frequency response will be Λ i ˜ x , where Λ i is the i -th row of Λ .
Hence, we can use (cid:107) Λ i (cid:107) 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.
[Equation]
[Table] [Table]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure]
g n i t u d e
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 deﬁnition 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 maps A ( l,h ) ∈ R n × n , the cosine similarity can be computed by:
[Equation]
where A ( l,h ) : ,i denotes the i -th column of A ( 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 output X ( l ) ∈ R n × d , the cosine similarity is estimated by:
[Equation]
where X ( l ) i, : denotes the i -th row of X ( 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.
We additionally provide a visualization of these two metrics for 12-layer DeiT in Fig.
G D EFERRED E XPERIMENTS AND M ODEL I NTERPRETATION G. F INE -T UNING E XPERIMENTS
Our deferred ﬁne-tuning experiment with CaiT (Touvron et al., 2021b) results are presented in
Table 3.
Different from trining scratch, we ﬁne-tune CaiT with AttnScale and FeatScale parameters from the pre-trained models for 60 epochs following Gong et al.
For a fair comparison, we simultaneously train a plain CaiT for another 60 epochs.
During ﬁne-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)
[Table]
G. V ISUALIZATION AND I NTERPRETATION OF A TTN S CALE
In this subsection [Table] we provide visualization to interpret our AttnScale and further support our experiments.
(cid:202)
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 ﬁnd when layer index goes larger, the scaling weights turns larger to prevent attention collapse at deeper layer.
(cid:203)
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.
(cid:204) 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 ﬁnd that attention maps from AttnScale enjoy richer ﬁltering diversity, capable of performing high-pass (row 2, column 3) and band-pass (row 12, column 5) ﬁltering, instead of only low-pass ﬁltering (see Fig. 9).
[Figure]
[Figure] [Figure] [Figure] [Figure]
[Figure]
t
[Figure]
{+t 0 
 [Figure] [Figure] t 0 
+}
[Figure]
[Figure] [Figure] t
[Figure] [Figure] [Figure] [Figure]
[Figure] [Figure] [Figure] [Figure]
t
[Figure]
G. V ISUALIZATION AND I NTERPRETATION OF F EAT S CALE
In this subsection, we provide visualization to interpret our FeatScale.
(cid:202) We plot the scaling weights of FeatScale in Fig. 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.
(cid:203) 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 (cid:107)HC [ X ] (cid:107) F / (cid:107) X (cid:107) F . We ﬁnd 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. D e i T -
[Figure]
Block
Block
Block
Block
Block
Block
[Figure] [Figure] [Figure] [Figure]
[Figure]
[Figure]
[Figure]
M a g n i t u d e
[Table]
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
M a g n i t u d e
[Figure]
[Figure]