H OW D
O V ISION
T RANSFORMERS W
ORK
Anonymous authors
Paper under double-blind review
A BSTRACT
The success of multi-head self-attentions (MSAs) for computer vision is nowindisputable.
However, little is known about how MSAs work.
We present various explanations to help better understand the nature of MSAs.
In particular, wedemonstrate the following properties of MSAs and Vision Transformers (ViTs):1 MSAs improve not only accuracy but also generalization by ﬂattening the losslandscapes.
Such improvement is primarily attributable to their data speciﬁcity, notlong-range dependency.
On the other hand, ViTs suffer from non-convex losses.
Large datasets and loss landscape smoothing methods alleviate this problem.
MSAs and Convs exhibit opposite behaviors.
For example, MSAs are low-passﬁlters, but Convs are high-pass ﬁlters.
Therefore, MSAs and Convs are complementary.
3 Multi-stage neural networks behave like a series connection of smallindividual models.
In addition, MSAs at the end of a stage play a key role inprediction.
Based on these insights, we propose AlterNet, a model in which Convblocks at the end of a stage are replaced with MSA blocks.
AlterNet outperforms
CNNs not only in large data regimes but also in small data regimes.
I NTRODUCTION
There is limited understanding of multi-head self-attentions (MSAs), although they are now ubiquitousin computer vision.
The most widely accepted explanation for the success of MSAs is their weakinductive bias and capture of long-range dependencies (See, e.g., (Dosovitskiy et al., 2021; Naseeret al., 2021;
Tuli et al., 2021; Yu et al., 2021; Mao et al., 2021; Chu et al., 2021)).
Yet because of theirover-ﬂexibility, Vision Transformers (ViTs)—neural networks (NNs) consisting of MSAs—havebeen known to have a tendency to overﬁt training datasets, consequently leading to poor predictiveperformance in small data regimes, e.g., image classiﬁcation on CIFAR (Krizhevsky et al., 2009).
However, we show that the explanation is poorly supported.
1. R ELATED W ORK
Self-attentions (Vaswani et al., 2017; Dosovitskiy et al., 2021) aggregate (spatial) tokens with normalized importances:
[Equation]
where Q , K , and V are query, key, and value, respectively.
d is the dimension of query and key, andz j is the j -th output token.
From the perspective of convolutional neural networks (CNNs), MSAsare a transformation of all feature map points with large-sized data-speciﬁc kernels.
Therefore, MSAsare at least as expressive as convolutional layers (Convs) (Cordonnier et al., 2020), although this doesnot guarantee that MSAs will behave like Convs.
Is the weak inductive bias of MSA, such as modeling long-range dependencies, beneﬁcial for thepredictive performance?
To the contrary, appropriate constraints may actually help a model learnstrong representations.
For example, local MSAs (Yang et al., 2019; Liu et al., 2021; Chu et al., 2021),which calculate self-attention only within small windows, achieve better performance than global
MSAs not only on small datasets but also on large datasets, such as ImageNet-21K (Russakovskyet al., 2015).
In addition, prior works observed that MSAs have the following intriguing properties: 1 MSAsimprove the predictive performance of CNNs (Dai et al., 2021;
Guo et al., 2021; Srinivas et al., 2021),
[Figure]
[Figure] [Figure]
and ViTs predict well calibrated
(Minderer et al againstdata corruptions, image occlusions (Naseer et al., 2021), and adversarial attacks (Shao et al., 2021;
Bhojanapalli et al., 2021; Paul & Chen, 2021; Mao et al., 2021).
They are particularly robust againsthigh-frequency noises (Shao et al., 2021).
3 MSAs closer to the last layer signiﬁcantly improvepredictive performance (Graham et al., 2021; Dai et al., 2021; Xiao et al., 2021).
These empirical observations raise immediate questions: 1 What properties of MSAs do we needto better optimize NNs?
Do the long-range dependencies of MSAs help NNs learn?
2 Do MSAsact like Convs?
If not, how are they different?
How can we harmonize MSAs with Convs?
Canwe just leverage their advantages?
We provide an explanation of how MSAs work by addressing them as a trainable spatial smoothing offeature maps , because Eq.
(1) also suggests that MSAs average feature map values with the positiveimportance-weights.
Even non-trainable spatial smoothings, such as a small 2 × 2 box blur, help
CNNs see better (Zhang, 2019; Anonymous, 2022).
These simple spatial smoothings not only improveaccuracy but also robustness by ensembling feature map points and ﬂattening the loss landscapes(Anonymous, 2022).
Remarkably, spatial smoothings have the properties of MSAs 1 – 3 .
See
Appendix B for theoretical and empirical explanations of MSAs as a spatial smoothing.
1. C ONTRIBUTION
We address the three key questions:
What
properties of MSAs do we need to improve optimization?
We present
various evidences to support that MSA is generalized spatial smoothing.
It means that MSAs improveperformance because their formulation—Eq.
(1)—is an appropriate inductive bias.
Their weak inductive bias disrupts NN training.
In particular, a key feature of MSAs is their data speciﬁcity , not long-range dependency.
As an extreme example, local MSAs with a 3 × 3 receptive ﬁeld outperformsglobal MSA because they reduce unnecessary degrees of freedom
[Equation]
ResNet 0.
Frequency -6.
L o g a m p l i t u d e
Stem
ViT 0.
Frequency -2.
L o g a m p l i t u d e
Emb 0.
N o r m a l
i z e d d e p t h
[Figure]
Fourier transformed feature maps.
Noise frequency -15.
A cc u r a c y ( % )
ResNet ViT
[Figure]
How do MSAs improve performance?
MSAs have their advantages and disadvantages.
On the onehand, they ﬂatten loss landscapes as shown in Fig.
The ﬂatter the loss landscape, the better theperformance and generalization (Li et al., 2018; Keskar et al., 2017; Santurkar et al., 2018; Foretet al., 2021; Chen et al., 2021).
Thus, they improve not only accuracy but also robustness in largedata regimes.
On the other hand, MSAs allow negative Hessian eigenvalues in small data regimes.
This means that the loss landscapes of MSAs are non-convex, and this non-convexity disturbs NNoptimization (Dauphin et al., 2014).
Large amounts of training data suppress negative eigenvaluesand convexify losses.2 Do MSAs act like Convs?
We show that MSAs and Convs exhibit opposite behaviors.
MSAsaggregate feature maps, but Convs disperse them.
Moreover, as shown in Fig.
2a, the Fourier analysisof feature maps shows that MSAs reduce high-frequency signals, while Convs, conversely, ampliﬁeshigh-frequency components.
In other words, MSAs are low-pass ﬁlters, but Convs are high-passﬁlters .
In addition, Fig. 2b indicates that Convs are vulnerable to high-frequency noise but that MSAsare not.
Therefore, MSAs and Convs are complementary.
How can we harmonize MSAs with Convs?
We reveal that multi-stage NNs behave like aseries connection of small individual models.
Thus, applying spatial smoothing at the end of a stageimproves accuracy by ensembling feature map outputs from each stage (Anonymous, 2022) as shownin Fig. 3a.
Based on this ﬁnding, we propose an alternating pattern of Convs and MSAs .
NN stagesusing this design pattern consists of a number of CNN blocks and one (or a few) MSA block as shownin Fig.
3c.
The design pattern naturally derives the structure of canonical Transformer, which has one
MSA block per MLP block as shown in Fig.
3b.
It also provides an explanation of how adding Convsto Transformer’s MLP block improves accuracy and robustness (Yuan et al., 2021;
Guo et al., 2021;
Mao et al., 2021).
Surprisingly, models using this alternating pattern of Convs and MSAs outperform CNNs not only onlarge datasets but also on small datasets, such as CIFAR.
This contrasts with canonical ViTs, modelsthat perform poorly on small amount of data.
It implies that MSAs are generalized spatial smoothingsthat complement Convs, not simply generalized Convs.
[Equation]
P ROPERTIES OF MSA S D O W E N EED T O I MPROVE O PTIMIZATION ?
To understand the underlying nature of MSAs, we investigate the properties of the ViT family: e.g.,vanilla ViT (Dosovitskiy et al., 2021); PiT (Heo et al., 2021), which is “ViT + multi-stage”; and Swin(Liu et al., 2021), which is “ViT + multi-stage + local MSA”.
This section shows that these additionalinductive biases enable ViTs to learn strong representations.
We also use ResNet (He et al., 2016a) forcomparison.
NNs are trained from scratch with DeiT-style data augmentation (Touvron et al., 2021)
[Figure]
2048, 1 × 1,512, 3 × 3,512, 3 × 3,
AvgPool
ResNet Blocks × L
Subsampling
[Figure]
512, 1 × 1,2048, 1 × 1,512, MSA, D×H
D×H , 1 × 1,
MLP Block ×
A " ention Block
[Figure]
Transformer 512, MSA, D×H
D×H , 1 × 1,2048, 1 × 1,512, 3 × 3,512, 1 × 1,2048, 1 × 1,
ResNet Blocks × L
A " ention Block
Subsampling
[Figure]
for 300 epochs.
The NN training begins with a gradual warmup (Goyal et al., 2017) for 5 epochs.
Formore detailed conﬁgurations, see Appendix A.
The stronger the inductive biases, the stronger the representations ( not regularizations ).
Domodels with weak inductive biases overﬁt training datasets?
To address this question, we providetwo criteria on CIFAR-100: the error of the test dataset and the cross-entropy, or the negativelog-likelihood, of the training dataset (NLL train , the lower the better).
See Fig.
5a for the results.
Contrary to our expectations, experimental results show that the stronger the inductive bias, thelower the training NLL and test error.
This indicates that ViT does not overﬁt training datasets .
Inaddition, appropriate inductive biases, such as locality constraints for MSAs, helps NNs learn strongrepresentations.
We also observe these phenomena on CIFAR-10 and ImageNet as shown in Fig.
C.1.
Figure C.3 also supports that weak inductive biases disrupt NN training.
In this experiment, anextremely small patch sizes for the embedding hurts the predictive performance of ViT.
ViT does not overﬁt small training datasets.
We observe that ViT does not overﬁt even on smallerdatasets.
Figure 5b shows the test error and the training NLL of ViT on subsampled datasets.
In this
D e n s i t y
ResNet ViT
ViT (6%)
D e n s i t
y warmup
th experiment, as the size of the dataset decreases, the error increases as expected, but surprisingly, NLL train alsoincreases.
Thanks to the strong data augmentation, ViTdoes not overﬁt even on a dataset size of 2%.
This suggests that ViT’s poor performance in small data regimesis not due to overﬁtting.
ViT’s non-convex losses lead to poor performance.
How do weak inductive biases of MSAs disturb theoptimization?
A loss landscape perspective providesan explanation: the loss function of ViT is non-convex,while that of ResNet is strongly (near-)convex .
This poorloss disrupts NN training (Dauphin et al., 2014), especially in the early phase of training (Jastrzebski et al.,2021).
Figure 1c and Fig.
4 provide top-5 largest Hessianeigenvalue densities (
[Figure]
Figure 4 also shows that large datasets suppress negative Hessian eigenvalues in the early phase oftraining.
Therefore, large datasets tend to help ViT learn strong representations by convexifying theloss.
ResNet enjoys little beneﬁt from large datasets because its loss is convex even on small datasets.
[Figure]
Loss landscape smoothing methods aids in ViT training.
Loss landscape smoothing methods can alsohelp ViT learn strong representations.
In classiﬁcationtasks, global average pooling (GAP) smoothens theloss landscape by strongly ensembling feature mappoints (Anonymous, 2022).
We demonstrate how theloss smoothing method can help ViT improve performance by analyzing ViT with GAP classiﬁer instead of
CLS token on CIFAR-100.
Figure 6 shows the Hessian max eigenvalue spectrumof the ViT with GAP.
As expected, the result shows that
GAP classiﬁer suppresses negative Hessian max eigenvalues, suggesting that GAP convexify the loss.
Sincenegative eigenvalues disturb NN optimization, GAP classiﬁer improve the accuracy by +2 .
7 percent point.
Likewise, Sharpness-Aware Minimization (SAM) (Foret et al., 2021), an optimizer that relies on thelocal smoothness of the loss function, also helps NNs seek out smooth minima.
Chen et al.
(2021)showed that SAM improves the predictive performance of ViT.
MSAs ﬂatten the loss landscape.
Another property of MSAs is that they reduces the magnitude of
Hessian eigenvalues.
Figure 1c and Fig.
4 shows that the eigenvalues of ViT are signiﬁcantly smallerthan that of CNNs.
While large eigenvalues impede NN training (Ghorbani et al., 2019), MSAs canhelp NNs learn better representations by suppressing large Hessian eigenvalues .
Figure 1a and Fig.
1balso support this claim.
In Fig.
1a, we visualize the loss landscapes by using ﬁlter normalization (Liet al., 2018).
The loss landscape of ViT is ﬂatter than that of ResNet, and this trend is noticeableat the boundary.
Similarly, Fig.
1b shows that ResNet follows an irregular trajectory, especially inthe early phase of training; ViT converges to the optimum along a smooth trajectory.
In large dataregimes, the negative Hessian eigenvalues—the disadvantage of MSAs—disappears, and only theiradvantages remain.
As a result, ViTs outperform CNNs on large datasets, such as ImageNet and JFT(Sun et al., 2017).
PiT and Swin also ﬂatten the loss landscapes.
See Fig.
C.4.
A key feature of MSAs is data speciﬁcity (not long-range dependency) .
The two distinguishingfeatures of MSAs are long-range dependency and data speciﬁcity.
Contrary to popular belief, thelong-range dependency hinders NN optimization.
To demonstrate this, we analyze convolutional ViT ,which consists of two-dimensional convolutional MSAs (Yang et al., 2019) instead of global MSAs.
Convolutional MSAs calculates self-attention only between feature map points in convolutionalreceptive ﬁelds after unfolding the feature maps in the same way as convolutions.
[Figure]
Figure 7a shows the error and NLL train of convolutional ViTs with kernel sizes of 3 × 3 , 5 × 5 , and8 × 8 (global MSA) on CIFAR-100.
In this experiment, 5 × 5 kernel outperforms 8 × 8 kernel on boththe training and the test datasets.
NLL train of 3 × 3 kernel is worse than that of 5 × 5 kernel, but betterthan that of global MSA.
Although the test accuracies of 3 × 3 and 5 × 5 kernels are comparable, therobustness of 5 × 5 kernel is signiﬁcantly better than that of 3 × 3 kernel on CIFAR-100-C (Hendrycks& Dietterich, 2019).
Figure 7b shows that the strong locality inductive bias not only reduce computational complexity asoriginally proposed (Liu et al., 2021), but also aid in optimization by convexifying the loss landscape.5 × 5 kernel has fewer negative eigenvalues than global MSA because it restricts unnecessary degreesof freedom.
5 × 5 kernel also has fewer negative eigenvalues than 3 × 3 kernel because it ensemblesa larger number of feature map points (See also Fig. 6).
The amount of negative eigenvalues isminimized when the two effects are balanced.
It is clear that data speciﬁcity improves NNs.
MLP-Mixer (Tolstikhin et al., 2021; Yu et al., 2021),a model with an MLP kernel that does not depend on input data, underperforms compared to ViTs.
Data speciﬁcity without self-attention (Bello, 2021) improves performance.
D O MSA S A CT L IKE C ONVS ?
Convs are data-agnostic and channel-speciﬁc.
In contrast, MSAs are data-speciﬁc and channelagnostic.
This section shows that these differences lead to large behavioral differences.
It suggeststhat MSAs and Convs are complementary.
[Figure]
MSAs are low-pass ﬁlters, but Convs are highpass ﬁlters.
As explained in Section 1.1, MSAsspatially smoothen feature maps with self-attentionimportances.
Therefore, we expect that MSAs willtend to reduce high-frequency signals.
See Proposition B.1 for a more detailed discussion.
Figure 8 shows the relative log amplitude ( ∆ log amplitude) of ViT’s
Fourier transformed feature map athigh-frequency ( 1 .
π ) on ImageNet.
In this ﬁgure,
MSAs almost always decrease the high-frequencyamplitude, and MLPs—corresponding to Convs—increase it.
The only exception is in the early stages ofthe model.
In these stages, MSAs behave like Convs,i.e., they increase the amplitude.
This could serveas an evidence for a hybrid model that uses Convsin early stages and MSAs in late stages (Guo et al.
Graham et al., 2021; Dai et al., 2021;
Xiao et al., 2021; Srinivas et al., 2021).
[Figure]
a
[Figure]
u r
e
m a p v a r i a n c
[Figure]
Based on this, we can infer that low-frequency signals and high-frequency signals are informativeto MSAs and Convs, respectively .
In support of this argument, we report the robustness of ViTand ResNet against frequency-based random noise.
Following Shao et al.
(2021) and Anonymous(2022), we measure the decrease in accuracy with respect to data with frequency-based random noisex noise = x 0 + F − 1 ( F ( δ ) (cid:12) M f ) , where x 0 is clean data, F ( · ) and F − 1 ( · ) are Fourier transformand inverse Fourier transform, δ is Gaussian random noise, and M f is frequency mask.
Figure 2b shows the results.
As expected, the results reveals that ViT and ResNet are vulnerable tolow-frequency noise and high-frequency noise, respectively.
Low-frequency signals and the highfrequency signals each correspond to the shape and the texture of images.
The results thus suggeststhat MSAs are shape-biased (Naseer et al., 2021), whereas Convs are texture-biased (Geirhos et al.,2019).
MSAs aggregate feature maps, but Convs do not.
Since MSAs average feature maps, they willreduce variance of feature map points.
This suggests that MSAs ensemble feature maps (Anonymous,2022).
To demonstrate this claim, we measure the variance of feature maps from NN layers.
Figure 9 shows the experimental results of ResNet and ViT.
This ﬁgure indicates that MSAs in ViTtend to reduce the variance; conversely, Convs in ResNet and MLPs in ViT increase it.
In conclusion,
MSAs ensemble feature maps, but Convs do not .
As Anonymous (2022) ﬁgured out, reducing thefeature map uncertainty helps optimization by ensembling and stabilizing the feature maps.
See
Fig.
D.1 for more results on PiT and Swin.
We observe two additional patterns for feature map variance.
First, the variance accumulates in every
NN layer and tends to increase as the depth increases.
Second, the feature map variance in ResNetpeaks at the ends of each stage.
Therefore, we can improve the predictive performance of ResNet byinserting MSAs at the end of each stage.
Furthermore, we also can improve the performance by using
MSAs with a large number of heads in late stages.
H OW C AN W E H ARMONIZE MSA S W ITH C ONVS ?
Since MSAs and Convs are complementary, this section seeks to design a model that leverages onlythe advantages of the two modules.
To this end, we propose the design rules described in Fig.
3c, anddemonstrate that the models using these rules outperforms CNNs, not only in the large data regimesbut also in the small data regimes, such as CIFAR.
4. D ESIGNING A RCHITECTURE
We ﬁrst investigate the properties of multi-stage NN architectures.
Based on this investigation, wecome to propose an alternating pattern, i.e., a principle for stacking MSAs based on CNNs.
Multi-stage NNs behave like individual models.
In Fig.
9, we observe that the pattern of featuremap variance repeats itself at every stages.
This behavior is also observed in feature map similaritiesand lesion studies.
[Figure]
P
P P
[Figure]
E
P
[Figure]
E
P P
P (
[Figure]
) Feature map similarity
[Figure]
Figure 10a shows the representational similarities of ResNet and Swin on CIFAR-100.
In thisexperiment, we use minibatch CKA (Nguyen et al., 2021) to measure the similarities.
As Nguyenet al.
(2021) ﬁgured out, the feature map similarities of CNNs have a block structure.
Likewise, weobserve that the feature map similarities of multi-stage ViTs, such as PiT and Swin, also have a blockstructure.
Since vanilla ViT does not have this structure (Bhojanapalli et al., 2021; Raghu et al.
, 2021),the structure is an intrinsic characteristic of multi-stage architectures.
See Fig.
D.3 for more detailedresults of PiT and Swin.
Figure 10b shows the results of lesion study (Bhojanapalli et al., 2021), where one NN unit is removedfrom already trained ResNet and Swin during the testing phase.
In this experiment, we remove one3 × 3 Conv layer from the bottleneck block of ResNet, and one MSA or MLP block from Swin.
In
ResNet, removing an early stage layers hurts accuracy more than removing a late stage layers.
Moreimportantly, removing a layer at the beginning of a stage impairs accuracy more than removing alayer at the end of a stage.
The case of Swin is even more interesting.
At the beginning of a stage,removing an MLP hurts accuracy.
At the end of a stage, removing an MSA seriously impairs theaccuracy.
These results are consistent with Fig.
See Fig.
D.4 for the results on ViT and PiT.
Based on these ﬁndings, we expect MSAs closer to the end of a stage to signiﬁcantly improve theperformance.
This is contrary to the popular belief that MSAs closer to the end of a model improvethe performance (Srinivas et al., 2021; d’Ascoli et al., 2021; Graham et al., 2021; Dai et al., 2021).
Build-up rule.
Considering all the insights, we propose the following design rules:
• Alternately replace Conv blocks with MSA blocks from the end of a baseline CNN model.
• If the added MSA block does not improve predictive performance, replace a Conv block
located at the end of an earlier stage with an MSA block .
• Use more heads and higher hidden dimensions for MSA blocks in late stages.
We call the model that follows these rules AlterNet .
AlterNet uniﬁes ViTs and CNNs by adjustingthe ratio of MSAs and Convs as shown in Fig.
Figure 11 shows AlterNet based on pre-activation
ResNet-50 (He et al., 2016b) for CIFAR-100 as an example.
Figure D.5 shows AlterNet for ImageNet.
Figure 12a reports the accuracy of Alter-ResNet-50 , which replaces the Conv blocks in ResNet-50with local MSAs (Liu et al., 2021) according to the aforementioned rules, on CIFAR-100.
As expected,
MSAs in the last stage ( c4 ) signiﬁcantly improve the accuracy.
Surprisingly, an MSA in 2 nd stage( c2 ) improves the accuracy, while two or more MSAs in the 3 rd stage ( c3 ) reduce it.
In conclusion,
MSAs at the end of a stage play an important role in prediction.
Conv block
Conv block
MSA block
MSA block
Conv block
Conv block MSA block
Conv block Conv block
Conv block
MSA block
Conv block Conv block
Conv block Conv block
Conv block Conv block
[Figure] [Figure]
Figure 12c demonstrates that MSAs suppress large eigenvalues while allowing only a few negativeeigenvalues.
As explained in Fig.
4, large datasets compensate for the shortcomings of MSAs.
Therefore, more data allows more MSAs for a models.
4. P ERFORMANCE
Figure 12b shows the accuracy and corruption robustness of Alter-ResNet-50 and other baselineson CIFAR-100 and CIFAR-100-C. Since CIFAR is a small dataset, CNNs outperforms canonical
ViTs.
Surprisingly, Alter-ResNet—a model with MSAs following the appropriate build-up rule—outperforms CNNs even in the small data regimes.
This suggests that MSAs complement Convs.
Inthe same manner, this simple modiﬁcation shows competitive performance on larger datasets, such as
ImageNet.
See Fig.
E.1 for more details.
D ISCUSSION
Our present work demonstrates that MSAs are not merely generalized Convs, but rather generalizedspatial smoothings that complement Convs.
MSAs help NNs learn strong representations by ensembling feature map points and ﬂattening the loss landscape.
Since the main objective of this work isto investigate the nature of MSA for computer vision, we preserve the architectures of Conv and
MSA blocks in AlterNet.
Thus, AlterNet has a strong potential for future improvements.
In addition,
AlterNet can conveniently replace the backbone for other vision tasks such as dense prediction(Carion et al., 2020).
As Anonymous (2022) pointed out, global average pooling (GAP) for simpleclassiﬁcation tasks has a strong tendency to ensemble feature maps, but NNs for dense prediction donot use GAP.
Therefore, we believe that MSA to be able to signiﬁcantly improve the results in denseprediction tasks by ensembling feature maps.
Lastly, strong data augmentation for MSA trainingharms uncertainty calibration as shown in Fig.
F.1a.
We leave a detailed investigation for future work.
R EPRODUCIBILITY S TATEMENT
To ensure reproducibility, we provide comprehensive resources, such as code and experimental details.
The codebase will be released as open source under the Apache License 2.0.
See the supplementalmaterial for the code.
Appendix A provides the speciﬁcations of all models used in this work.
Detailedexperimental setup including hyperparameters and the structure of AlterNet are also available in
Appendix A and Appendix E. De-facto image datasets are used for all experiments.
A E XPERIMENTAL S ETUP
We obtain the main experimental results with two kind of machines for CIFAR (Krizhevsky et al.,2009).
The ﬁrst one is the Intel Xeon W-2123 Processor, 32GB memory, and a single GeForce RTX2080 Ti, and the other is four Intel Intel Broadwell CPUs, 15GB memory, and a single NVIDIA
T4.
For ImageNet (Russakovsky et al., 2015), we use AMD Ryzen Threadripper 3960X 24-Core
Processor, 256GB memory, and four GeForce RTX 2080 Ti.
Models are implemented in PyTorch(Paszke et al., 2019).
NNs are trained using categorical cross-entropy (NLL) loss and AdamW optimizer (Loshchilov &
Hutter, 2018) with initial learning rate of 1 .
25 × 10 − 4 and weight decay of 5 × 10 − 2 .
We also usecosine annealing scheduler (Loshchilov & Hutter, 2016).
We train NNs for 300 epochs with batch sizeof 96 on CIFAR, and with batch size of 128 on ImageNet.
The learning rate is gradually increased(Goyal et al., 2017) for 5 epochs.
Following Touvron et al.
(2021), strong data augmentations, suchas RandAugment (Cubuk et al., 2020), Random Erasing (Zhong et al., 2020), label smoothing(Szegedy et al., 2016), mixup (Zhang et al., 2018), and CutMix (Yun et al., 2019), are used fortraining.
Stochastic depth (Huang et al., 2016) is also used to regularize NNs.
Since this DeiT-styleconﬁguration (Touvron et al., 2021) is the de facto standard in ViT training (e.g., (Heo et al., 2021;
Liu et al., 2021)), we believe the insights presented in this paper can be widely useful.
See sourcecode for detailed conﬁgurations.
We mainly report the performances of ResNet-50, ViT-Ti, PiT-Ti, and Swin-Ti.
Their trainingthroughputs on CIFAR-100 are 320 , 434 , 364 , and 469 image/sec respectively, which are comparableto each other.
In Figs.
5a and C.1a, we also report the predictive performance of ResNeXt-50 (Xieet al., 2017),
Twins-S
(Chu et al., 2021), and MLP-Mixer-Ti (Tolstikhin et al., 2021).
In Fig.
E.1, weadditionally report the performance of ConViT-Ti (d’Ascoli et al., 2021), LeViT-128S (Graham et al.
CoaT-Lite-Ti (Xu et al., 2021), and MLP-Mixer-B (Tolstikhin et al., 2021).
We use patch sizeof 2 × 2 for ViT and PiT on CIFAR.
For Swin, we use patch size of 1 × 1 and window size of 4 × 4 .
We use patch size of 4 × 4 for ViT only in Fig.
We halve the depth of the ViT in Fig.
C.5 and
Fig.
C.6 due to the memory limitation.
All models for CIFAR, and ResNet, ViT, and AlterNet for ImageNet are trained from scratch.
Weuse pertained PiT and Swin from Wightman (2019) for ImageNet.
The implementations of Vision
Transformers are based on Wightman (2019) and Wang (2021).
For Hessian max eigenvalue spectrum (Anonymous, 2022), we use power iteration with batch sizeof 16 to produce the top-5 greatest eigenvalues.
To this end, Yao et al.
(2020)’s
implementation isused.
We modify the algorithm to calculate the eigenvalues with respect to (cid:96) 2 regularized NLL onaugmented training dataset—note that measuring Hessian eigenvalues on clean dataset would giveincorrect results.
In the strict sense, the weight decay is not (cid:96) 2 regularization, but we neglect thedifference.
B MSA I S S PATIAL S MOOTHING
As mentioned in Section 1.1, spatial smoothings before subsampling layers help CNNs see better (Zhang, 2019; Anonymous, 2022).
Anonymous (2022) showed that such improvement in performance is possible due to (Bayesian) feature map ensembles.
To this end, they used data-complemented
Bayesian NN (BNN) (Park et al., 2021) which exploits data uncertainty:
[Equation]
where π ( x i | x j ) is normalized importance weight of a feature map point x i with respect to anotherfeature map x j , i.e., (cid:80)
i π
( x i | x j ) = 1 .
The importance is deﬁned as a similarity between x i and x j .p
( z j | x i , w i ) and p ( z j | x j , D ) are one NN prediction and output predictive distribution, respectively.w i is a NN weight sample from posterior p ( w |D ) with respect to training dataset D .
Put shortly,
Eq.
(2) complements a prediction with other predictions based on similarities.
For instance, a 2 × 2box blur ensembles four neighboring feature map points, each with 1 ⁄ 4 of the same importance.
We note that the formulations for self-attention and data-complemented BNN are identical.
The
Softmax term and V in Eq.
(1) exactly correspond to π ( x i | x j ) and p ( z j | x i , w i ) in Eq.
Theweight samples in Eq.
(2) is correspond to the multi-heads of MSAs (See also (Hron et al., 2020)).
[Figure]
Proposition B.1.
Self-attention is spatial smoothing of the values, i.e., the weighted average withself-attention importances is low-pass ﬁlter.
Proof.
Self-attention is convolution of the importance weights A ( Softmax term in Eq.
and the value predictions ( V term in Eq.
A ∗ V (3)
Consider V averaged N times with A :
[Equation]
Since A is probability, A ∗ · · · ∗ A is the probability for the sum of N random variables from A , i.e,a + · · · + a ∼ A ∗· · ·∗ A where a ∼ A .
By deﬁnition, an operator is low-pass ﬁlter if and only if thehigh frequency component vanishes when the operator is applied inﬁnitely.
Therefore, A is low-passﬁlter because Var ( a + · · · + a ) =
N Var ( a ) and F [ A ∗ · · · ∗ A ∗ V ] =
F [ A ∗ · · · ∗ A ] F [ V ]where F is Fourier transform.
Likewise, the properties of spatial smoothing are the same as those of MSAs (Anonymous, 2022):
Spatial smoothing improves the accuracy of CNNs.
In addition, spatial smoothing predictswell-calibrated uncertainty.
2 Spatial smoothing is robust against MC dropout (which is equivalentto image occlusion), data corruption, and adversarial attacks, and particularly robust against highfrequency noise.
3 Spatial smoothing layers closer to the output layer signiﬁcantly improves thepredictive performance.
Taking all these observations together, we provide an explanation of how MSAs work by addressingthem as a general form of spatial smoothing or an implementation of data-complemented BNN .
Spatial smoothing improves performance in the following ways (Anonymous, 2022): 1 Spatialsmoothing helps in NN optimization by ﬂattening the loss landscapes.
Even a small 2 × 2 blur ﬁltersigniﬁcantly improves performance.
2 Spatial smoothing is a low-pass ﬁlter.
CNNs are vulnerableto high-frequency noise, but spatial smoothing improves the robustness against it by reducing highfrequency noise effectively.
3 Spatial smoothing is effective when applied at the end of a stagebecause it aggregates all the feature maps that are transformed in various ways.
The main paper showsthat these mechanisms are maintained for MSAs.
C V I T S F ROM A L OSS L ANDSCAPE P ERSPECTIVE
This section provides further explanations of the analysis in Section 2.
The lower the NLL on training dataset, the lower the error on test dataset.
Figure 5a demonstrates that low training NLL results in low test error on CIFAR-100.
We observe the same pattern on
CIFAR-10 and ImageNet.
Figure C.1 shows the result.
[Figure]
In small data regimes, such as CIFAR-10 (Fig. C.1a) and CIFAR-100 (Fig. 5a), the error and the of ViTs are inferior to those
CNNs.
It suggests that the cons
MSA outweigh the
As discussed in Fig. 4, ViTs suffers from the non-convex losses, and the convex loss
optimization
In large data regimes, such as ImageNet (Fig. C.1b), both the error and the NLL train of ViTs with local
MSAs are superior to those of CNNs.
Since large datasets convexify the loss function as discussed in
Fig.
4, the pros of MSAs outweigh the cons.
Therefore, MSAs help NNs learn strong representationsby ﬂattening the loss landscapes.
Rigorous discussion on regularization of CNN’s inductive bias.
In Fig.
5a, we compare similarlysized models, such as ResNet-50 and ViT-Ti.
Through the comparison, we show that a weak inductivebias hinders NN training, and CNN’s inductive biases—Conv’s inductive bias and multi-stagearchitecture—help NN learn strong representations.
However, the inductive biases of CNN givesbetter test accuracy for the same training NLL, i.e., Convs somewhat regularize NNs.
We analyzetwo comparable models in terms of NLL train on CIFAR-100.
The NLL train of ResNet-18, a modelsmaller than ResNet-50, is 2.31 and the error is 22.0%.
The NLL train of ViT-S, a larger than ViT-Ti, is2.17 and the error is 30.4%.
In summary, the inductive biases of CNNs improve accuracy for similartraining NLL.
Most of the improvement come from multi-stage architecture, not the inductive bias of Conv.
The
NLL train of the PiT-Ti, multi-stage ViT-Ti, is 2.29 and the error is 24.1 %.
The accuracy of PiT is only
[Figure]
receptive ﬁeld also regularize ViT.
ViT does not overﬁt small training dataset even with a largenumber of epochs.
Figure 5b shows that ViT does not overﬁt small training datasets such as CIFAR.
We observe the samephenomenon in ViT training with a large number of epochs.
In Fig.
C.2, we train ViT and ResNet for 75, 150, 300, 600, and1200 epochs.
The experimental results show that the NLL train andthe error decrease as the epoch increases.
However, the predictiveperformances of ViT is inferior to those of ResNet across a wholerange of epochs.
A smaller patch size does not always imply better results.
ViT splits the image into multiple patches.
The smaller the patchsize, the greater the ﬂexibility of expression and the weaker the in-
[Figure]
ductive bias.
By analyzing ViT with three patch sizes— 2 × 2 , 4 × 4 , and 8 × 8 —we again demonstratethat weak inductive bias disturbs NN optimization.
Figure C.3a shows the error on test dataset and NLL on the training dataset of CIFAR-100.
Asexpected, large patch size harms the performance on both test and training datasets.
Surprisingly,a small patch size also harms the performance.
An appropriate patch size helps ViT learn strongrepresentations, not regularizes ViT.
The Hessian max eigenvalue spectra in Fig.
C.3b explain this observation.
The results reveal that asmall patch size allows negative Hessian eigenvalues.
In other words, weak inductive bias makes theloss landscape non-convex.
A large patch size suppresses negative eigenvalues.
On the other hand, itlimits the model expression and sharpens the loss landscape.
The performance is optimized whenthese two effects are balanced.
[Figure] [Figure]
of eigenvalues as in Fig.
C.5.
Multi-stage architecture in PiT and local MSA in Swin also ﬂatten the loss landscapes.
Asexplained in Fig.
1, MSA smoothens the loss landscape.
Similarly, multi-stage architecture in PiTand local MSA in Swin also help NN learn strong representations by smoothing the loss landscapes.
Figure C.4 provides loss landscape visualizations, trajectories in training phase, and Hessian eigenvalue spectra of ResNet, ViT, PiT, and Swin.
Figure C.4a visualizes the global geometries of the lossfunctions.
The loss landscape of PiT is ﬂatter than that of ViT near the optimum.
Since Swin has moreparameters than ViT and PiT, (cid:96) 2 regularization leads the loss landscape.
Figure C.4b demonstratesthat all the ViTs converge to the optima with smooth trajectories.
These support that MSA smoothensthe loss landscape.
Figure C.4c shows the local geometries of the loss function by representing
Hessian eigenvalues.
In the early phase of training, multi-stage in PiT helps training by suppressingnegative Hessian eigenvalues.
Local MSA in Swin allows the negative eigenvalues, but signiﬁcantlyreduces the magnitude of eigenvalues.
Moreover, the magnitude of Swin’s Hessian eigenvalue doesnot signiﬁcantly increases in the late phase of learning.
Lack of heads may lead to non-convex loss.
Neural tangent kernel (NTK)
(Jacot et al., 2018) theoretically implies that the loss landscape of ViT is convex and ﬂat when the number of heads or the number of embedding dimensions per head goes to inﬁnity (Hron et al., 2020; Liu et al., 2020).
In particular, from Liu et al.
(2020), we easily derive the following proposition:
[Equation]
where || H || is the Hessian spectral norm, m is the number of heads or the number of emgeddingdimensions per head, and c ∞ is a small constant.
Therefore, in practical situations, insufﬁcient headsmay cause non-convex and sharp loss.
[Figure]
In Fig.
C.5, we empirically show that a lot of heads in MSA convexify and ﬂatten the loss landscape(c.f.
Michel et al.
In this experiment, we introduce negative max eigenvalue proportion (NEP,the lower the better) and average of positive max eigenvalue (APE, the lower the better) to measurethe non-convexity and the sharpness, respectively.
For a Hessian max eigenvalue spectrum p ( λ ) , NEPis the proportion of negative eigenvalues (cid:82)
−∞ p ( λ ) dλ , and APE is the expected value of positiveeigenvalues (cid:82) ∞
λ p ( λ ) dλ / (cid:82) ∞
p ( λ ) dλ .
The results show that both NEP and APE decrease as thenumber of heads increases.
Likewise, Fig.
C.6 shows that high embedding dimensions per head alsoconvexify and ﬂatten the loss.
The exponents of APE are − 0 .
562 and − 0 .
796 for the number ofheads and the embedding dimensions, which are in close agreement with the value predicted by thetheory of − 1 / 2 .
A large model has a ﬂat loss in an early phase of training.
Figure C.7 analyze the loss landscapesof large models, such as ResNet-101 and ViT-S. As shown in Fig.
C.7a, large models explore low
NLLs.
These can be a surprising result, because the loss landscapes of the large model is globallysharp as shown in Fig.
C.7b.
ResNet 0.
Normalized depth 0.
F e a t u r e m a p v a r i a n c e
ViT 0.
Normalized depth 0.
F e a t u r e m a p v a r i a n c e
PiT 0.
Normalized depth 0.
F e a t u r e m a p v a r i a n c e
Swin 0.
Normalized depth 0.
F e a t u r e m a p v a r i a n c e
Hessian eigenvalue spectra in Fig.
C.7c solve the problem: The Hessian eigenvalues of large modelsare smaller than that of small models in early phase of training.
It indicates that large models haveﬂat loss functions locally.
D V I T S F ROM A F EATURE M AP P ERSPECTIVE
This section provides further explanations of the analysis in Section 3 and Section 4.1.
MSAs in PiT and Swin also ensemble feature maps.
In Fig.
9, we show that MSAs in ViT reducefeature map variance.
We observe same pattern in PiT and Swin.
Figure D.1 demonstrates that MSAsin PiT and Swin also reduce the feature map variance, suggesting that they also ensemble featuremaps.
One exception is the 3 rd stage in Swin.
In the early of this stage, MSAs suppresses the increasein variance, but at the end of the stage it does not suppress it.
MSAs in PiT and Swin are also low-pass ﬁlters.
As discussed in Fig.
8, MSAs in ViTs are lowpass ﬁlters, while MLPs in ViT and Convs in ResNet are high-pass ﬁlters.
Likewise, we demonstratethat MSAs in PiT and Swin are also low-pass ﬁlters.
Figure D.2 shows the relative log amplitude of Fourier transformed feature maps.
As in the case of
ViT, MSAs in PiT and Swin generally decrease the amplitude of high-frequency signals; in contrast,
MLPs increases the amplitude.
Multi-stage ViTs have the block structures.
The feature map similarities of CNNs shows a blockstructure (Nguyen et al., 2021).
As Raghu et al.
(2021) ﬁgured out, ViT has uniform representationsacross all layers.
By investigating multi-stage ViTs, we demonstrate that subsampling layers createthe characteristic block structure of the representation.
See Fig.
D.3.
Conv in early phase of stages and MSA in late phase of stages play an important role.
Figure D.4 shows the results of lesion study for ResNet and ViTs.
In this experiment, we remove one 3 ×
Conv layer from bottleneck block of ResNet, and one MSA or MLP block from Swin.
We observe
ResNet 0.
Frequency -6.
L o g a m p l i t u d e
Stem 0.
Normalized depth -6.
L o g a m p l i t u d e
ViT 0.
Frequency -2.
L o g a m p l i t u d e
Emb 0.
Normalized depth -2.
L o g a m p l i t u d e
PiT 0.
Frequency -3.
L o g a m p l i t u d e
Emb 0.
Normalized depth -3.
L o g a m p l i t u d e
Swin 0.
Frequency -4.
L o g a m p l i t u d e
Emb 0.
Normalized depth -5.
L o g a m p l i t u d e 0.
1.0 Normalized depth
ResNet
E
P P
P E
P
P P
[Figure]
E
P P
[Figure]
P P
Swin
E
P P
P E
P
P P
[Figure]
Conv block Conv block
MSA block
MSA block MSA block
Conv block
MSA block
Conv block Conv block
MSA block MSA block
Conv block
Conv block Conv block
Conv block Conv block
Conv
[Figure]
consistent results for all models: Removing Convs in early phase of stages and removing MSAs inlate phase of stages signiﬁcantly harm the accuracy.
As a result, the accuracy varies periodically.
E E XTENDED I NFORMATION OF A LTER N ET
This section provides further informations on AlterNet.
Detailed architecture of AlterNet.
Section 4 introduces AlterNet to harmonize Convs with MSAs.
Since most MSAs take pre-activation arrangement, pre-activation ResNet is used as a baseline forconsistency.
We add one CNN block to the last stage of ResNet to make the number of blocks even.
Alocal MSA with relative positional encoding from Swin is used for AlterNet.
However, for simplicityof implementation, we do not implement detailed techniques, such as a cyclic shift and layer-speciﬁcinitialization.
For CIFAR, the patch size of the MSA is 1 × 1 and the window size is 4 × 4 .
If all
Conv blocks are alternately replaced with MSA, AlterNet becomes a Swin-like model.
In order to achieve better performance, NNs should strongly aggregate feature maps at the end ofmodels as discussed in Section 3 and Section 4.
To this end, AlterNet use 3, 6, 12, 24 heads for MSAin each stage.
[Figure]
The computational costs of Conv block and MSA block are almost the same.
The training throughput of Alter-ResNet-50 is473 image/sec on CIFAR-100, which is 20% faster than that ofpre-activation ResNet-50.
The optimal number of MSAs depends on the model and dataset,so we empirically determine the number of MSAs as shown in
Fig. 12a.
A large dataset allows a large number of MSAs.
For
ImageNet, we use 6 MSAs as in Fig.
D.5, because large datasetsalleviate the shortcoming of MSA.
MSAs improve the performance of CNN on ImageNet.
Since
MSA complements Conv, MSA improves the predictive performance of CNN with the appropriate build-up rules as shown in
Section 4.1.
Figure E.1 reports the accuracy and robustness—mean accuracy on ImageNet-C—of CNNs and ViTs on ImageNet1K. Since ImageNet is large dataset, a number of ViTs outperform CNNs.
MSAs with the build-up rules signiﬁcantly improves
ResNet, and the predictive performance of AlterNet is on par with that of Swin without heavymodiﬁcations, e.g., the shifted windowing scheme (Liu et al., 2021).
AlterNet is easy-to-implementand has a strong potential for future improvements.
In addition, the build-up rules not only improve
ResNet, but also other NNs.
We observe that the build-up rule also improved the performance ofvanilla post-activation ResNet and ResNeXt; however, we do not report the results because of thevisualization simplicity.
F D ISTINCTIVE P ROPERTIES OF D ATA A UGMENTATION
This section empirically demonstrates that NN training with data augmentation is different from training on large datasets.
We compare DeiT-style strong data augmentation with weak data augmentation,i.e., resize and crop.
F. D ATA A UGMENTATION C AN H ARM U NCERTAINTY C ALIBRATION
Figure F.1a shows the reliability diagrams of NNs with and without strong augment on CIFAR-100.
In the results, both ResNet and ViT without data augmentation predict overconﬁdent results.
We showthat data augmentation makes the predictive results under-conﬁdent (cf.
Wen et al.
Theseare unexpected results because the predictions are not under-conﬁdent without data augmentation onlarge datasets such as ImageNet.
We leave a detailed investigation for future work.
[Figure]
We observe same
Hessian means the data helps NNs converge to better optima by ﬂattening the loss landscapes.
the other hand, the data augmentation allows a lot of negative Hessian eigenvalues i.e , it makes the loss non-
F. D ATA A UGMENTATION R EDUCES THE M AGNITUDE OF H ESSIAN E IGENVALUES
How does data augmentation help MSA avoid overﬁtting on training dataset and achieve betteraccuracy on test dataset?
Figure F.1b shows the Hessian max eigenvalue spectra of NNs with andwithout strong data augmentation.
First of all, data augmentation reduces the magnitude of Hessianeigenvalues, i.e., data augmentation ﬂattens the loss landscapes in an early phase of training.
Theﬂat loss leads to better generalization.
On the other hand, data augmentation allows a lot of negative
Hessian eigenvalues, i.e., data augmentation makes the loss non-convex.
This prevents NN fromconverging to a low loss on training datasets.
It is clearly different from the effect of large datasetsdiscussed in Fig.
4—large datasets convexify the loss landscapes.
We leave a detailed investigationfor future work.