CMT: Convolutional Neural Networks Meet Vision Transformers

Introduction

The past decades have witnessed the extraordinary contribution of CNNs in the field of computer vision due to its ability of extracting deep discriminative features. Meanwhile, self-attention based transformers has become the de facto most popular models for natural language processing (NLP) tasks, and shown excellent capability of capturing long-distance relationships. Recently, many researchers attempt to apply the transformer-based architectures to vision domains, and achieve promising results in various tasks such as image classification , object detection , and semantic segmentation . Vision transformer (ViT) is the first work to replace the conventional CNN backbone with a pure transformer. Input images (224$\times$224$\times$3) are first split into 196 non-overlapping patches (with a fixed size of 16$\times$16$\times$3 per patch), which are analogous to the word tokens in NLP. The patches are then fed into stacked standard transformer blocks to model global relations and extract feature for classification. The design paradigm of ViT has heavily inspired the following transformer based models for computer vision, such as IPT for low-level vision and SETR for semantic segmentation.

Despite that transformers have demonstrated excellent capabilities when migrated to vision tasks, their performances are still far inferior to similar-sized convolutional neural network counterparts, e.g., EfficientNets . We believe the reason of such weakness is threefold. Firstly, images are split into patches in ViT and other transformer-based models such as IPT and SETR . Doing so can greatly simplify the process of applying transformer to image-based tasks. And the sequence of patches can be directly fed into a standard transformer where long-range dependencies between patches can be well captured. However, it ignores the fundamental difference between sequence-based NLP tasks and image-based vision tasks, e.g., the 2D structure and spatial local information within each patch. Secondly, transformer is difficult to explicitly extract low-resolution and multi-scale features due to the fixed patch size, which poses a big challenge to dense prediction tasks such as detection and segmentation. Thirdly, the computational and memory cost of self-attention modules in transformers are quadratic ($\mathcal{O}(N^{2}C)$) to the resolution of inputs, compared to $\mathcal{O}(NC^{2})$ of convolution-based CNNs. High resolution images are prevalent and common, e.g., 1333$\times$800 in COCO and 2048$\times$1024 in Cityscapes . Using transformers to process such images would inevitably cause the problem of insufficient GPU memory and low computation efficiency.

In this paper, we stand upon the intersection of CNNs and transformers, and propose a novel CMT (CNNs meet transformers) architecture for visual recognition. The proposed CMT takes the advantages of CNNs to compensate for the aforementioned limitations when utilizing pure transformers. As shown in Figure 2(c), input images first go through the convolution stem for fine-grained feature extraction, and are then fed into a stack of CMT blocks for representation learning. Specifically, the introduced CMT block is an improved variant of transformer block whose local information is enhanced by depth-wise convolution. Compared to ViT , the features generated from the first stage of CMT can maintain higher resolution, i.e., $H/4\times W/4$ against $H/16\times W/16$ in ViT, which are essential for other dense prediction tasks. Furthermore, we adopt the stage-wise architecture design similar to CNNs by using four convolutional layer with stride $2$, to gradually reduce the resolution (sequence length) and increase the dimension flexibly. The stage-wise design helps to extract multi-scale features and alleviate the computation burden caused by high resolution. The local perception unit (LPU) and inverted residual feed-forward network (IRFFN) in CMT block can help capture both local and global structure information within the intermediate features and promote the representation ability of the network. Finally, the average pooling is used to replace the class token in ViT for better classification results. In addition, we propose a simple scaling strategy to obtain a family of CMT variants. Extensive experiments on ImageNet and other downstream tasks demonstrate the superiority of our CMT in terms of accuracy and FLOPs. For example, our CMT-S achieves 83.5% ImageNet top-1 with only 4.0B FLOPs, while being 14x and 2x less than the best existing DeiT and EfficientNet , respectively. In addition to image classification, CMT can also be easily transferred to other vision tasks and serve as a versatile backbone. Using CMT-S as the backbone, RetinaNet can achieve 44.3% mAP on COCO val2017, outperforming the PVT-based RetinaNet by 3.9% with less computational cost.

Related Work

The computer vision community prospered in past decades riding the wave of deep learning, and the most popular deep neural networks are often built upon basic blocks, in which a series of convolutional layers are stacked sequentially to capture local information within intermediate features. However, the limited receptive field of small convolutional kernels makes it difficult to obtain global information, withholding the networks of high performance on challenging tasks such as classification, object detection, and semantic segmentation. Therefore, many researchers start to dig deeper into self-attention based transformers which have the ability to capture long-range information. Here we briefly review the conventional CNNs and recently proposed vision transformers.

Convolutional neural networks. The first standard CNN was proposed by LeCun et al. for handwritten character recognition, and the past decades have witnessed that many powerful networks achieved unprecedented success on large scale image classification task . AlexNet and VGGNet showed that a deep neural network composed of convolutional layers and pooling layers can obtain adequate results in recognition. GoogleNet and InceptionNet demonstrated the effectiveness of multiple paths within a basic block. ResNet showed better generalization by adding shortcut connections every two layers to the base network. To alleviate the limited receptive fields in prior research, some researches incorporated attention mechanisms as an operator for adaptation between modalities. Wang et al. proposed to stack attention modules sequentially between the intermediate stages of deep residual networks. SENet and GENet adaptively recalibrated channel-wise feature responses by modeling interdependencies between channels. NLNet incorporated the self-attention mechanism into neural networks, providing pairwise interactions across all spatial positions to augment the long-range dependencies. In addition to above architectural advances, there has also been works focusing on improving over-parameterized deep neural networks by trading accuracy for efficiency. For example, MobileNets and EfficientNets both leveraged neural architecture search (NAS) to design efficient mobile-size network and achieved new state-of-the-art results.

Vision transformers. Since transformers achieved remarkable success in natural language processing (NLP) , many attempts have been made to introduce transformer-like architectures to vision tasks. The pioneering work ViT directly applied the transformer architecture inherited from NLP to classification with image patches as input. While ViT required a large private dataset JFT-300M to achieve promising result, DeiT introduced a new training paradigm to extend ViT to a data-efficient transformer directly trained on ImageNet-1K. T2T-ViT proposed to embed visual tokens by recursively aggregating neighboring tokens into one token. TNT proposed to model both patch-level and pixel-level representation by the inner and outer transformer block, respectively. PVT introduced the pyramid structure into ViT, which can generate multi-scale feature maps for various pixel-level dense prediction tasks. CPVT and CvT are the most related to our work which leverage a convolutional projection into conventional transformer block, but we carefully investigate how to maximize the advantage of utilizing both CNNs and transformers by studying the different components including shortcut and normalization functions and successfully obtain a more superior result. Besides, transformers are also used to solve other vision tasks such as object detection , semantic segmentation , image retrieval , and low-level vision task .

Although there are many works successfully applying transformers for vision tasks, they have not shown satisfactory results compared to conventional CNNs, which are still the primary architectures for vision applications. Transformers are especially good at modeling long-range dependencies necessary for downstream vision tasks. However, locality should also be maintained for visual perception. In this paper, we demonstrate the potential of combining the transformer based network together with convolutional layer, the overall architecture follows the elaborated prior convolutional neural networks such as ResNet and EfficientNet .

Approach

Our intention is to build a hybrid network taking the advantages of both CNNs and transformers. An overview of ResNet-50 , DeiT , and the proposed small version (CMT-S) of CMT architectures are presented in Figure 2. As shown in Figure 2(b), DeiT directly splits an input image into non-overlapping patches, however, the in-patch structure information can only be poorly modeled with linear projections. To overcome this limitation, we utilize the stem architecture which has a $3\times 3$ convolution with a stride of $2$ and an output channel of $32$ to reduce the size of input images, followed by another two $3\times 3$ convolutions with stride $1$ for better local information extraction. Following the design in modern CNNs (e.g., ResNet ), our model has four stages to generate feature maps of different scales which are important for dense prediction tasks. To produce the hierarchical representation, a patch embedding layer consisting of a convolution and a layer normalization (LN) is applied before each stage to reduce the size of intermediate feature (2x downsampling of resolution), and project it to a larger dimension (2x enlargement of dimension). In each stage, several CMT blocks are stacked sequentially for feature transformation while retaining the same resolution of the input. For example, the “Stage 3” of CMT-S contains $16$ CMT blocks as illustrated in Figure 2(c). The CMT block is able to capture both local and long-range dependencies, and we will describe it in Sec. 3.2 in details. The model ends with a global average pooling layer, a projection layer, and a 1000-way classification layer with softmax.

Given an input image, we can obtain four hierarchical feature maps with different resolutions, similar to typical CNNs such as ResNet and EfficientNet . With the above feature maps whose strides are $4$, $8$, $16$, and $32$ with respect to the input, our CMT can obtain multi-scale representations of input images and can be easily applied to downstream tasks such as object detection and semantic segmentation.

2 CMT Block

The proposed CMT block consists of a local perception unit (LPU), a lightweight multi-head self-attention (LMHSA) module, and an inverted residual feed-forward network (IRFFN), as illustrated in Figure 2(c). We will describe these three parts in the following.

Local Perception Unit. Rotation and shift are two commonly used data augmentation manners in vision tasks, and these operations should not alter the final results of the model. In other words, we expect translation-invariance in those tasks. However, the absolute positional encoding used in previous transformers, initially designed to leverage the order of tokens, damages such invariance because it adds unique positional encoding to each patch . Besides, vision transformers ignore the local relation and the structure information inside the patch. To alleviate the limitations, we propose the local perception unit (LPU) to extract local information, which is defined as:

Inverted Residual Feed-forward Network. The original FFN proposed in ViT is composed of two linear layers separated by a GELU activation . The first layer expands the dimension by a factor of $4$, and the second layer reduces the dimension by the same ratio:

where the activation layer is omitted. We also include the batch normalization after the activation layer and the last linear layer according to . The depth-wise convolution is used to extract local information with negligible extra computational cost. The motivation for inserting shortcut is similar to that of classic residual networks, which can promote the propagation ability of gradient across layers. We show that such shortcut helps the network achieve better results in our experiments.

With the aforementioned three components, the CMT block can be formulated as:

where ${\bf Y}_{i}$ and ${\bf Z}_{i}$ denote the output features of LPU and LMHSA module for the $i$-th block, respectively. LN denotes the layer normalization . We stack several CMT blocks in each stage for feature transformation and aggregation.

3 Complexity Analysis

We analyze the computational cost between standard ViT and our CMT in this section. A standard transformer block consists of a MHSA module and a FFN. Given an input feature of size $n\times d$, the computational complexity (FLOPs) can be calculated as:

where $r$ is the expansion ratio of FFN, $d_{k}$ and $d_{v}$ are dimensions of key and value, respectively. More specifically, ViT sets $d=d_{k}=d_{v}$ and $r=4$, the cost can be simplified as:

Under above setting, the FLOPs of CMT block is as follows:

where $k\geq 1$ is the reduction ratio in LMHSA. Compared to standard transformer block, the CMT block is more friendly to computational cost, and is easier to process the feature map under higher resolution (larger $n$).

4 Scaling Strategy

Inspired by , we propose a new compound scaling strategy suitable for transformer-based networks, which uses a compound coefficient $\phi$ to uniformly scale the number of layers (depth), dimensions, and input resolution in a principled way:

where $\alpha$, $\beta$, and $\gamma$ are constants determined by grid search to decide how to assign resources to network depth, dimension and input resolution, respectively. Intuitively, $\phi$ is the coefficient that controls how many more ($\phi\geq 1$) or less ($\phi\leq-1$) resources are available for model scaling. Notably, the FLOPs of the proposed CMT block is approximately proportionalThe precious proportion is associated with $n$ and $d$. For example, CMT-S has $n$=$3136\gg d$=$64$ in “stage 1” and $n$=$49\ll d$=$512$ in “stage 4”. The above proportion can already generate good variants for CMT. to $\alpha$, $\beta^{1.5}$, and $\gamma^{2}$ according to E.q. 16. And we constraint $\alpha\cdot\beta^{1.5}\cdot\gamma^{2}\approx 2.5$ so that for a given new $\phi$, the total FLOPS will approximately increase by $2.5^{\phi}$. This will strike a balance between the increase of computational cost and performance gain. In our experiments, we empirically set $\alpha$=1.2, $\beta$=1.3, and $\gamma$=1.15.

We build our model CMT-S to have similar model size and computation complexity with DeiT-S (ViT-S) and EfficientNet-B4. We also introduce CMT-Ti, CMT-XS and CMT-B according to the proposed scaling strategy. The input resolutions are $160^{2}$, $192^{2}$, $224^{2}$, and $256^{2}$ for all four models, respectively. The detailed architecture hyper-parameters are shown in Table 1.

Experiments

In this section, we investigate the effectiveness of CMT architecture by conducting experiments on several tasks including image classification, object detection, and instance segmentation. We first compare the proposed CMT with previous state-of-the-art models on above tasks, and then ablate the important elements of CMT.

Experimental Settings. ImageNet is a image classification benchmark which contains 1.28M training images and 50K validation images of 1000 classes. For fair comparisons with recent works, we adopt the same training and augmentation strategy as that in DeiT , i.e., models are trained for 300 epochs (800 for CMT-Ti that requires more epochs to converge) using the AdamW optimizer. All models are trained on 8 NVIDIA Tesla V100 GPUs.

Results of CMT. Table 2 shows the performances of the proposed CMTs that are scaled from the CMT-S according to E.q. 17. Our models achieve better accuracy with fewer parameters and FLOPs compared to other convolution-based and transformer-based counterparts. In particular, our CMT-S achieves 83.5% top-1 accuracy with 4.0B FLOPs, which is 3.7% higher than the baseline model DeiT-S and 2.0% higher than CPVT , indicating the benefit of CMT block for capturing both local and global information. Note that all previous transformer-based models are still inferior to EfficientNet which is obtained via a thorough architecture search, however, our CMT-S is 0.6% higher than EfficientNet-B4 with less computational cost, which demonstrates the efficacy of the proposed hybrid structure and show strong potential for further improvement. We also plot the accuracy-FLOPs curve in Figure 1(a) to have an intuitive comparison between these models. We can see that CMTs consistently outperform other models by a large margin.

2 Ablation Study

Stage-wise architecture. Transformer-based ViT/DeiT can only generate single-scale feature map, losing a lot of multi-scale information crucial for dense prediction tasks. We change the columnar DeiT-S to hierarchical DeiT-S-4Stage, which has 4 stages like CMT-S in Table 1, but maintains the original FFN. We also change MHSA to LMHSA to reduce computational cost. As shown in Table 3, DeiT-S-4Stage outperforms DeiT-S by 1.6% with less FLOPs, demonstrating that the widely-adopted stage-wise design in CNNs is a better choice for promoting transformer-based architecture.

CMT block. Ablations on different modules in CMT are shown in Table 5. DeiT-S-4Stage has 4 patch embedding layers (the first is a 4$\times$4 convolution with stride 4). “+ Stem” indicates that we add the CMT stem into the network and replace the first patch embedding layer with a 2$\times$2 convolution with stride 2. The improvement shows the benefit of the convolution-based stem. Besides, the proposed LPU and IRFFN can further boost the network by 0.8% and 0.6%, respectively. It is worth noticing that the shortcut connections in LPU and IRFFN are also crucial for the final performance.

Normalization function. Transformer-based models usually use LN inherited from NLP. However, convolution-based models usually utilize batch normalization (BN) to stabilize the training. CMT maintains the LN before LMHSA and IRFFN, and inserts BN after the convolutional layer. If all LNs are replaced by BNs, the model cannot converge during training. If all BNs are replaced by LNs, the performance of CMT-S drops to 83.0%, indicating that proper application of normalization functions can improve the final performance.

Scaling strategy. Table 4 shows the ImageNet results of CMT architecture under different scaling strategies. Unidimensional scaling strategies are significantly inferior to the proposed compound scaling strategy, especially for depth-only scaling strategy which leads to a even worse result of $83.4\%$ against $83.8\%$ of the original CMT-S, when the network is scaled up.

3 Transfer Learning

Experimental Settings. The experiments are conducted on COCO , which contains 118K training images and 5K validation images of 80 classes. We evaluate the proposed CMT-S using two typical framework: RetinaNet and Mask R-CNN for object detection and instance segmentation, respectively. Specifically, we replace the original backbone with our CMT-S to build new detectors. All models are trained under standard single-scale and “1x” schedule (12 epochs) following PVT .

Results of CMT. We report the performance comparison results of object detection task and instance segmentation task in Table 6 and Table 7 respectively. For object detection with RetinaNet as basic framework, CMT-S outperforms Twins-PCPVT-S with 1.3% mAP and Twins-SVT-S with 2.0% mAP. For instance segmentation with Mask R-CNN as basic framework, CMT-S surpasses Twins-PCPVT-S with 1.7% AP and Twins-SVT-S with 1.9% AP. We also report the inference speed on COCO val2017 with 1280×800 input, CMT-S based RetinaNet and Mask R-CNN achieve 14.8 FPS and 11.2 FPS, respectively.

3.2 Other Vision Tasks

We also evaluate the proposed CMT on five commonly used transfer learning datasets, including CIFAR10 , CIFAR100 , Standford Cars , Flowers , and Oxford-IIIT Pets (see Appendix for more details). We fine-tune the ImageNet pretrained models on new datasets following . Table 8 shows the corresponding results. CMT-S outperforms other transformer-based models in all datasets with less FLOPs, and achieves comparable performance against EfficientNet-B7 with 9x less FLOPs, which demonstrates the superiority of CMT architecture.

Conclusion

This paper proposes a novel hybrid architecture named CMT for visual recognition and other downstream computer vision tasks such as object detection and instance segmentation, and addresses the limitations of utilizing transformers in a brutal force manner in the field of computer vision. The proposed CMT architectures take advantages of both CNNs and transformers to capture local and global information, promoting the representation ability of the network. In addition, a scaling strategy is proposed to generate a family of CMT variants for different resource constraints. Extensive experiments on ImageNet and other downstream vision tasks demonstrate the effectiveness and superiority of the proposed CMT architecture.

Acknowledgment Chang Xu was supported by the Australian Research Council under Project DP210101859 and the University of Sydney SOAR Prize.

References