Global Context Vision Transformers

We propose global context vision transformer (GC ViT), a novel architecture that enhances parameter and compute utilization for computer vision tasks. The core of the novel model are global context self-attention modules, joint with standard local self-attention, to effectively yet efficiently model both long and short-range spatial interactions, as an alternative to complex operations such as an attention masks or local windows shifting. While the local self-attention modules are responsible for modeling short-range information, the global query tokens are shared across all global self-attention modules to interact with local key and values. In addition, we address the lack of inductive bias in ViTs and improve the modeling of inter-channel dependencies by proposing a novel downsampler which leverages a parameter-efficient fused inverted residual block. The proposed GC ViT achieves new state-of-the-art performance across image classification, object detection and semantic segmentation tasks. On ImageNet-1K dataset for classification, the tiny, small and base variants of GC ViT with 28M, 51M and 90M parameters achieve 83.4%, 83.9% and 84.4% Top-1 accuracy, respectively, surpassing comparably-sized prior art such as CNN-based ConvNeXt and ViT-based Swin Transformer. Pre-trained GC ViT backbones in downstream tasks of object detection, instance segmentation, and semantic segmentation on MS COCO and ADE20K datasets outperform prior work consistently, sometimes by large margins. Code and pre-trained models are available at https://github.com/NVlabs/GCViT.


INTRODUCTION
During the recent years, Transformers (Vaswani et al., 2017) have achieved State-Of-The-Art (SOTA) performance in Natural Language Processing (NLP) benchmarks and became the de facto model for various tasks. A key element in the success of Transformers is the self-attention mechanism which allows for capturing contextual representations via attending to both distant and nearby tokens (Yin et al., 2021). Following this trend, Vision Transformer (ViT)  proposed to utilize image patches as tokens in a monolithic architecture with minor differences comparing to encoder of the original Transformer. Despite the historic dominance of Convolutional Neural Network (CNN) in computer vision, ViT-based models have achieved SOTA or competitive performance in various computer vision tasks. In essence, the self-attention mechanism in ViT allows for learning more uniform short and long-range information (Raghu et al., 2021) in comparison to CNN. However, the monolithic architecture of ViT and quadratic computational complexity of self-attention baffle their swift application to high resolution images (Yang et al., 2021a) in which capturing multi-scale long-range information is crucial for accurate representation modeling.
Several efforts Tu et al., 2022), most notably Swin Transformer , have attempted to address the balance between short-and longrange spatial dependencies by proposing multi-resolution architectures in which the self-attention is computed in local windows. In this paradigm, cross-window connections such as window shifting are used for modeling the interactions across different regions. Despite the progress, the limited receptive field of local windows challenges the capability of self-attention to capture long-range information, and window-connection schemes such as shifting only cover a small neighborhood in the vicinity of each window. Subsequent efforts such as Focal Transformer  attempted to address this issue by designing highly sophisticated self-attention modules with increased model complexity. In this work, we introduce the Global Context (GC) ViT network to address these limitations. Specifically, we propose a hierarchical ViT architecture consisting of local and global self-attention modules. At each stage, we compute global query tokens, using a novel fused inverted residual blocks, which we refer to as modified Fused-MBConv blocks, that encompass global contextual information from different image regions. While the local self-attention modules are responsible for modeling short-range information, the global query tokens are shared across all global self-attention modules to interact with local key and value representations. The design of our proposed framework for global query generator and self-attention is intuitive and simple and can be efficiently implemented using major deep learning framework. Hence, it eliminates sophisticated and computationally expensive operations and ensures the effectiveness of self-attention when applied to high-resolution images. In addition, we propose a novel downsampling block with a parameter-efficient fused-MBConv layer to address the lack of inductive bias in ViTs and enhancing the modeling of inter-channel dependencies.
We have extensively validated the effectiveness of the proposed GC ViT using three publicly available datasets for various computer vision tasks. For image classification using ImageNet-1K dataset, GC ViT with 28M, 51M, 90M and 201M parameters, referred to as tiny, small, base and large variants, achieve new SOTA benchmarks of 83.4%, 83.9%, 84.4% and 84.6% Top-1 accuracy. Hence, GC ViT consistently outperforms both ConvNeXt (Liu et al., 2022) andSwin Transformer (Liu et al., 2021) models by a significant margin (see Fig. 1). Using a pre-trained GC ViT base backbone with a Cascade Mask RCNN  head, our model achieves a box mAP of 52.9 for object detection and a mask mAP of 45.8 for instance segmentation on the MS COCO dataset. In addition, using an UPerNet (Xiao et al., 2018) head, our model achieves a mIoU of 49.0 on ADE20K for semantic segmentation. Other variants of GC ViT with different learning capacities also demonstrate SOTA results when compared to similarly-sized models on both MS COCO and ADE20K datasets. Hence, GC ViT demonstrates great scalability for high-resolution images on various downstream tasks, validating the effectiveness of the proposed framework in capturing both short and long-range information.
The main contributions of our work are summarized as follows: • We introduce a compute and parameter-optimized hierarchical ViT with reparametrization of the design space (e.g., embedding dimension, number of heads, MLP ratio). • We design an efficient CNN-like token generator that encodes spatial features at different resolutions for global query representations. • We propose global query tokens that can effectively capture contextual information in an efficient manner and model both local and global interactions. • We introduce a parameter-efficient downsampling module with modified Fused MB-Conv blocks that not only integrates inductive bias but also enables the modeling of inter-channel dependencies. • We demonstrate new SOTA benchmarks for : (1) ImageNet classification with Pareto fronts on ImageNet-1K for model size and FLOPs (see Fig. 1), and (2) downstream tasks such as detection, instance segmentation and semantic segmentation on MS COCO and ADE20K, respectively.

GC VIT ARCHITECTURE
Architecture. Fig. 2 depicts the architecture of GC ViT. We propose a hierarchical framework to obtain feature representations at several resolutions (called stages) by decreasing the spatial dimensions while expanding the embedding dimension, both by factors of 2. At first, given an input image with resolution of x ∈ R H×W ×3 , we obtain overlapping patches by applying a 3 × 3 convolutional layer with a stride of 2 and appropriate padding. Then patches are projected into a C-dimensional embedding space with another 3 × 3 convolutional layer with stride 2. Every GC ViT stage is composed of alternating local and global self-attention modules to extract spatial features. Both operate in local windows like Swin Transformer , however, the global self-attention has access to global features extracted by the global token generator. The token generator is a CNN-like module that extracts features from the entire image only once at every stage. After each stage, the spatial resolution is decreased by 2 while the number of channels is increased by 2 via a downsampling block. Resulting features are passed through average pooling and linear layers to create an embedding for a downstream task.
The GC ViT architecture benefits from novel blocks such as a downsampling operator, a global query generator and a global self-attention module described in the next sections.  Downsampler. We leverage an idea of spatial feature contraction from CNN models that imposes locality bias and cross channel interaction while reducing dimensions. We utilize a modified Fused-MBConv block, followed by a max pooling layer with a kernel size of 3 and stride of 2 as a downsampling operator, see Fig 3. The Fused-MBConv block in our work is similar to the one in EfficientNetV2  with modifications as inx where SE, GELU and DW-Conv 3×3 denote Squeeze and Excitation block (Hu et al., 2018), Gaussian Error Linear Unit (Hendrycks & Gimpel, 2016) and 3×3 depth-wise convolution, respectively. In our proposed architecture, the Fused-MBConv blocks provide desirable properties such as inductive bias and modeling of inter-channel dependencies. It is ablated in Table 5.

Attention.
Multi-head self-attention is the the core computational operator in the GC ViT architecture, it extracts semantic information from the image. GC ViT is composed of local and global selfattention modules illustrated in Fig 4. Similar to Swin Transformer , we benefit from splitting the image into windows and performing local self-attention within them, this leads to linear complexity scaling with image size. The local self-attention extracts local, short-range, information and in order to facilitate long range dependencies we propose to use a novel global self attention to allow cross-patch communication with those far beyond the local window. Global self-attention attends other regions in the image via global query token that represents image embedding extracted with CNN-like module. On the other hand, the global features are extracted from the entire input features and then repeated to form global query tokens. The global query is interacted with local key and value tokens, hence allowing to capture long-range information via cross-region interaction. Best viewed in color.

Fused MBConv
Max pool 2x2 Extracted global features reshape repeat

Stage-wise global tokens
Spatially matched with local tokens Tokenized features times for input-to-stage dimension matching Figure 5 -Global query generator schematic diagram. It is designed to (i) transform an input feature map to the current stage of dimension H, W, C denoting height, width, and channel respectively, (ii) extract features via repeating the modified Fused MBConv block, joint with down-sampling, log 2 H h times for dimension matching to local window size h (iii) output is reshaped and repeated to ( H h ) 2 number of local tokens that can attend to global contextual information. denotes merged dimensions during reshaping.

GLOBAL QUERY GENERATOR
We propose to generate global query tokens that encompass information across the entire input feature maps for interaction with local key and value feature pairs. Specifically, as shown in Fig. 5, a layer f in the the generator consists of a Fused-MBConv block followed by a max pooling layer, similar to the one described in Sec. 2, and the final global query q g,i at stage i (i ∈ {1, 2, 3, 4}) of GC ViT is computed according to These query tokens are computed once at every stage of the model and shared across all global attention blocks, hence decreasing number of parameters and FLOPs and improving the generalizability. In addition, the global attention layers only learn local key and value features which will be used for interaction with global query tokens. both local and global spatial information. Fig. 6 illustrates the difference between local and global self-attention. The global attention query q g has a size of B × C × h × w, wherein B, C, h and w denote batch size, embedding dimension, local window height and width, respectively. Moreover, q g is repeated along the batch dimension to compensate for the overall number of windows and batch size B * = B × N where N is the number of local windows. q g is further reshaped into multiple heads. The value and key are computed within each local window using a linear layer. The global self-attention query, key and value features are computed as follows

GLOBAL SELF-ATTENTION
) attn = softmax(qk) return matmul(attn, v).reshape (B * , N, C) Since the partitioned windows only contain local information, interaction with rich contextual information embedded in the global query tokens provides an effective way of enlarging the receptive field and attending to various regions in the input feature maps. The self-attention module is computed as in where d is scaling factor and b is a learnable relative position bias term. Assuming position change between [−p + 1, p − 1] along horizontal and vertical axes, b is sampled from the gridb ∈ R (2p−1)×(2p−1) . As shown in Sec. 5, relative position bias improves the performance, especially for dense prediction downstream tasks. In Algorithm 1, we present a PyTorch-like pseudocode for computing global selfattention in GC ViT. A complexity analysis of the global self-attention is presented in the supplementary materials.

RELATED WORK
ViT. The ViT  was first proposed as an alternative to CNNs with the advantage of enlarged receptive field, due to its self-attention layers. However, it lacked desirable properties of CNNs such as inductive biases and translation invariance and required large-scale training datasets to achieve competitive performance. Data-efficient Image Transformers (DeiT)  introduced a distillation-based training strategy which significantly improved the classification accuracy. LeViT (Graham et al., 2021) proposed a hybrid model with re-designed multi-layer perceptron (MLP) and self-attention modules that are highly-optimized for fast inference. Cross-covariance Image Transformer ( Following previous efforts, we used a random crop size of 512 × 512 for the input images. For fair assessment, we only compare against models with a pre-trained ImageNet-1K backbone.

CLASSIFICATION
We present the ImageNet-1K classification benchmarks in Table 1 Fig. 1, GC ViT models have better or comparable computational efficiency in terms of number FLOPs over the competing counterpart models.   (Xiao et al., 2018) and pretrained ImageNet-1K backbone. All models use a crop size of 512 × 512.
In Table 2, we present object detection  and instance segmentation benchmarks

ADE20K SEMANTIC SEGMENTATION RESULTS
We present semantic segmentation benchmarks on ADE20K dataset in Table 3  As shown in Table 4, we study the role of each component in GC ViT model for classification, detection, instance and semantic segmentation. For simplicity, we start with Swin Transformer as the base model and progressively re-design the components to demonstrate their importance in improving the performance. Firstly, we remove the window shifting and predictably observe significant performance degradation across all tasks. Changing distribution of parameters to our design improves the performance by +1.7, +2.8, +2.2 and +1.7 in terms of accuracy, box AP, mask AP and mIoU. Such reparametrization includes changing the window size, MLP ratio, number of layers to name but a few. Adding the CNN-based stem of GC ViT to the model provides additional improvements of +0.3, +0.2, +0.2 and +0.2 in terms of accuracy, box AP, mask AP and mIoU. In addition, using our proposed downsampler further improves the accuracy, box AP, mask AP and mIoU by +0.4, +0.1, +0.1 and +0.3, respectively. The last two changes demonstrate the importance of convolutional inductive bias and capturing the inter-channel dependencies in our model. Finally, leveraging the proposed global self-attention improves the performance by by +0.8, +0.8, +0.6 and +1.2 in terms of accuracy, box AP, mask AP and mIoU. Hence, this validates the effectiveness of the  proposed global self-attention, in particular for downstream tasks with high resolution images such as semantic segmentation in which modeling long-range spatial dependencies is critical.

DOWNSAMPLER DESIGN
We studied the effectiveness of various downsampler blocks in Table 5. The simplest alternative to our design is a pair of convolutional and maxpooling layers. However, it results in a reduction of ImageNet Top-1 accuracy by -0.7. Patch merging is another variant which was introduced in Swin Transformers . However, it reduces the accuracy by -0.5. Finally, our down-sampler which consists of a modified Fused-MBConv block and strided convolution and shows the best result. Importance of the former component is explained by the SE operation which boosts cross channel interaction while keeping number of parameters and FLOPs low. We conclude that our proposed down-sampler is essential to achieve high accuracy as it introduces convolutional inductive bias.

INTERPRETABILITY
To provide further insights on interpretability of the proposed global self-attention and query tokens, we demonstrate visualization of the learned attention and Grad-CAM (Selvaraju et al., 2017) maps in Fig. 7  C.2 EMA AND BATCH SIZE We also used used Exponential Moving Averages (EMA) and observed slight improvement in terms of ImageNet TOp-1 accuracy. Furthermore, the performance of the model across different batch sizes were stable as we did not observe significant changes. Table S.3 demonstrates the effect of EMA and batch size on the accuracy of a GCViT-T model.

E TRAINING DETAILS
For image classification, GC ViT models were trained using four computational nodes with 32 NVIDIA A100 GPUs. The total training batch size is 1024 (32 per GPU) for GC ViT-S, GC ViT-B, GC ViT-L and 4096 (128 per GPU) for GC ViT-XXT, GC ViT-XT and GC ViT-T. On average, each model required 32 hours of training with the specified hyper-parameters as indicated in the paper. All classification models were trained using the timm package (Wightman, 2019). Object detection and instance segmentation models as well as semantic segmentation models were trained using one computational node with 8 NVIDIA A40 GPUs using a total batch size of 16, hence a batch size of 2 per GPU. Detection and instance segmentation models were trained using mmdetection (Chen et al., 2019) package and on average required 56 hours of training. Semantic segmentation models were trained using mmsegmentation (Contributors, 2020) package, and on average required 34 hours of training.

F COMPLEXITY ANALYSIS
Given an input feature map of x ∈ R H×W ×C at each stage with a window size of h × w, the computational complexity of GC ViT is as follows The efficient design of global query token generator and other components allows to maintain a similar computational complexity in comparison to Swin Transformer Liu et al. (2021) while being able to capture long-range information and achieve better higher accuracy for classification and downstream tasks such as detection and segmentation.

G IMAGENET CLASSIFICATION BENCHMARKS
In Table S.4, we provide a comprehensive benchmark in terms of Top-1 accuracy for the models that are only trained on ImageNet-1K ) dataset, and without additional data.