Scaling the Depth of Vision Transformers via the Fourier Domain Analysis

Vision Transformer (ViT) has recently demonstrated promise in computer vision problems. However, unlike Convolutional Neural Networks (CNN), it is known that the performance of ViT saturates quickly with depth increasing, due to the observed attention collapse or patch uniformity. Despite a couple of empirical solutions, a rigorous framework studying on this scalability issue remains elusive. In this paper, we first establish an analytic framework to investigate ViT from the spectrum domain. We show that the self-attention mechanism inherently amounts to a low-pass filter, which indicates when ViT scales up its depth, excessive low-pass filtering will cause feature maps to only preserve their Direct-Current (DC) component. We then propose two straightforward yet effective techniques to mitigate the undesirable low-pass limitation. The first technique, termed AttnScale, decomposes a self-attention block into low-pass and high-pass components, then rescales and combines these two filters to produce an all-pass self-attention matrix. The second technique, termed FeatScale, re-weights feature maps on separate frequency bands to amplify the high-frequency signals. Both techniques are efficient, hyperparameter-free, and can effectively avoid attention collapse and patch uniformity caused by low-pass filtering. Our experiments demonstrate our proposed methods consistently help ViT benefit from deeper architectures, bringing $\ge$ 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.