[Paper Review] Anti-Oversmoothing in Deep Vision Transformers via the Fourier Domain Analysis: From Theory to Practice

Peihao Wang, Wenqing Zheng+2 more|arXiv (Cornell University)|Mar 9, 2022

CCD and CMOS Imaging Sensors45 citations

TL;DR

The paper analyzes Vision Transformers in the Fourier domain, showing self-attention acts as a low-pass filter that hampers depth, and introduces AttnScale and FeatScale to recover high-frequency information, yielding consistent accuracy gains across ViT variants.

ABSTRACT

Vision Transformer (ViT) has recently demonstrated promise in computer vision problems. However, unlike Convolutional Neural Networks (CNN), it is known that the performance of ViT saturates quickly with depth increasing, due to the observed attention collapse or patch uniformity. Despite a couple of empirical solutions, a rigorous framework studying on this scalability issue remains elusive. In this paper, we first establish a rigorous theory framework to analyze ViT features from the Fourier spectrum domain. We show that the self-attention mechanism inherently amounts to a low-pass filter, which indicates when ViT scales up its depth, excessive low-pass filtering will cause feature maps to only preserve their Direct-Current (DC) component. We then propose two straightforward yet effective techniques to mitigate the undesirable low-pass limitation. The first technique, termed AttnScale, decomposes a self-attention block into low-pass and high-pass components, then rescales and combines these two filters to produce an all-pass self-attention matrix. The second technique, termed FeatScale, re-weights feature maps on separate frequency bands to amplify the high-frequency signals. Both techniques are efficient and hyperparameter-free, while effectively overcoming relevant ViT training artifacts such as attention collapse and patch uniformity. By seamlessly plugging in our techniques to multiple ViT variants, we demonstrate that they consistently help ViTs benefit from deeper architectures, bringing up to 1.1% performance gains "for free" (e.g., with little parameter overhead). We publicly release our codes and pre-trained models at https://github.com/VITA-Group/ViT-Anti-Oversmoothing.

Motivation & Objective

• Motivate and understand why ViTs saturate or underperform as depth increases.
• Formulate a spectral-domain framework to analyze ViT features and self-attention.
• Propose two efficient, hyperparameter-free techniques to counteract over-smoothing in ViTs.
• Demonstrate that deeper ViT models can benefit from the proposed methods across multiple backbones.

Proposed method

• Formulate ViT as a sequence of self-attention blocks and analyze in the Fourier domain.
• Prove that self-attention acts as a low-pass filter independent of input tokens or query/key matrices (Theorem 1).
• Introduce AttnScale to decompose attention into low- and high-pass components and reweight to achieve all-pass behavior (Eq. 7-9).
• Introduce FeatScale to reweight DC and high-frequency components of feature maps without explicit Fourier transforms (Eq. 10-12).
• Argue that AttnScale and FeatScale are lightweight, hyperparameter-free, and plug into various ViT variants.
• Provide empirical validation showing consistent gains across DeiT, CaiT, and Swin backbones.

Experimental results

Research questions

• RQ1 Why does ViT performance saturate or degrade with increasing depth?
• RQ2 Can Fourier-domain analysis reveal mechanisms behind attention-driven smoothing in ViTs?
• RQ3 Do lightweight scaling techniques in attention maps or feature maps mitigate over-smoothing and enable deeper ViTs?
• RQ4 Are AttnScale and FeatScale transferable across different ViT architectures without extra hyperparameters?
• RQ5 What are the practical accuracy gains when applying these methods to standard Vision Transformer variants?

Key findings

• Self-attention cascades act as a low-pass filter, eroding high-frequency information as depth increases.
• AttnScale can convert the effective attention into an all-pass filter by decomposing into low- and high-frequency parts and reweighting the high-frequency residual.
• FeatScale reweights DC and high-frequency components at the feature map level to preserve high-frequency information.
• Both methods are lightweight, hyperparameter-free, and compatible with DeiT, CaiT, and Swin backbones.
• Empirical results show up to 1.1% accuracy gains on ImageNet across evaluated models without additional parameter overhead.
• Compared to several prior techniques, AttnScale and FeatScale outperform methods like re-attention and LayerScale in reported gains.