[Paper Review] LocalViT: Analyzing Locality in Vision Transformers

Yawei Li, Kai Zhang+5 more|arXiv (Cornell University)|Apr 12, 2021

Advanced Neural Network Applications44 references283 citations

TL;DR

LocalViT adds a locality mechanism to vision transformers by inserting depth-wise convolutions into the feed-forward network, improving ImageNet accuracy with minimal extra cost and showing generalization across multiple transformer architectures.

ABSTRACT

The aim of this paper is to study the influence of locality mechanisms in vision transformers. Transformers originated from machine translation and are particularly good at modelling long-range dependencies within a long sequence. Although the global interaction between the token embeddings could be well modelled by the self-attention mechanism of transformers, what is lacking is a locality mechanism for information exchange within a local region. In this paper, locality mechanism is systematically investigated by carefully designed controlled experiments. We add locality to vision transformers into the feed-forward network. This seemingly simple solution is inspired by the comparison between feed-forward networks and inverted residual blocks. The importance of locality mechanisms is validated in two ways: 1) A wide range of design choices (activation function, layer placement, expansion ratio) are available for incorporating locality mechanisms and proper choices can lead to a performance gain over the baseline, and 2) The same locality mechanism is successfully applied to vision transformers with different architecture designs, which shows the generalization of the locality concept. For ImageNet2012 classification, the locality-enhanced transformers outperform the baselines Swin-T, DeiT-T, and PVT-T by 1.0%, 2.6% and 3.1% with a negligible increase in the number of parameters and computational effort. Code is available at https://github.com/ofsoundof/LocalViT.

Motivation & Objective

• Motivate the integration of locality mechanisms into vision transformers to capture local image structures.
• Propose a locality-enhanced transformer by inserting depth-wise convolutions into the feed-forward network.
• Analyze how locality, activation functions, and expansion ratio affect performance.
• demonstrate the approach across multiple vision transformer architectures to show generality.

Proposed method

• Interpret the input as a sequence of token embeddings and rearrange to a 2D lattice (Seq2Img).
• Replace the feed-forward network with a module inspired by inverted residuals that includes 1x1 convolutions and a depth-wise 2D convolution.
• Use a non-linear activation after the depth-wise convolution (e.g., ReLU6, h-swish) and optional attention modules (ECA/SE).
• Split the class token before the feed-forward network and concatenate it back after processing image tokens to preserve classification behavior.
• Apply locality to selected transformer layers and analyze the placement and expansion ratio (gamma) effects on performance.

Experimental results

Research questions

• RQ1 Does injecting locality into the feed-forward network improve vision transformer accuracy without substantial parameter or FLOPs increases?
• RQ2 How do activation functions, layer placement, and hidden dimension expansion affect the benefits of locality?
• RQ3 How well does the locality mechanism generalize across different vision transformer architectures (e.g., DeiT, T2T-ViT, PVT, TNT)?

Key findings

• Depth-wise convolution alone improves baseline transformers.
• Activation choice after depth-wise convolution significantly impacts gains (e.g., h-swish with SE/ECA yields larger gains).
• Locality is more beneficial in lower transformer layers than higher ones.
• Increasing the hidden dimension expansion ratio (gamma) provides larger capacity and accuracy gains.
• The locality mechanism generalizes across DeiT, T2T-ViT, PVT, and TNT, with notable improvements over baselines in several cases.
• On ImageNet, LocalViT variants outperform DeiT-T and PVT-T baselines by up to several percentage points with minimal parameter/compute overhead.