Bottleneck Transformers for Visual Recognition

Introduction

Deep convolutional backbone architectures have enabled significant progress in image classification , object detection , instance segmentation . Most landmark backbone architectures use multiple layers of $3\times 3$ convolutions.

While the convolution operation can effectively capture local information, vision tasks such as object detection, instance segmentation, keypoint detection require modeling long range dependencies. For example, in instance segmentation, being able to collect and associate scene information from a large neighborhood can be useful in learning relationships across objects . In order to globally aggregate the locally captured filter responses, convolution based architectures require stacking multiple layers . Although stacking more layers indeed improves the performance of these backbones , an explicit mechanism to model global (non-local) dependencies could be a more powerful and scalable solution without requiring as many layers.

Modeling long-range dependencies is critical to natural language processing (NLP) tasks as well. Self-attention is a computational primitive that implements pairwise entity interactions with a content-based addressing mechanism, thereby learning a rich hierarchy of associative features across long sequences. This has now become a standard tool in the form of Transformer blocks in NLP with prominent examples being GPT and BERT models.

A simple approach to using self-attention in vision is to replace spatial convolutional layers with the multi-head self-attention (MHSA) layer proposed in the Transformer (Figure 1). This approach has seen progress on two seemingly different approaches in the recent past. On the one hand, we have models such as SASA , AACN , SANet , Axial-SASA , etc that propose to replace spatial convolutions in ResNet botleneck blocks with different forms of self-attention (local, global, vector, axial, etc). On the other hand, we have the Vision Transformer (ViT) , that proposes to stack Transformer blocks operating on linear projections of non-overlapping patches. It may appear that these approaches present two different classes of architectures. We point out that it is not the case. Rather, ResNet botteneck blocks with the MHSA layer can be viewed as Transformer blocks with a bottleneck structure, modulo minor differences such as the residual connections, choice of normalization layers, etc. (Figure 3). Given this equivalence, we call ResNet bottleneck blocks with the MHSA layer as Bottleneck Transformer (BoT) blocks.

Here are a few challenges when using self-attention in vision: (1) Image sizes are much larger ($1024\times 1024$) in object detection and instance segmentation compared to image classification ($224\times 224$). (2) The memory and computation for self-attention scale quadratically with spatial dimensions , causing overheads for training and inference.

To overcome these challenges, we consider the following design: (1) Use convolutions to efficiently learn abstract and low resolution featuremaps from large images; (2) Use global (all2all) self-attention to process and aggregate the information contained in the featuremaps captured by convolutions. Such a hybrid design (1) uses existing and well optimized primitives for both convolutions and all2all self-attention; (2) can deal with large images efficiently by having convolutions do the spatial downsampling and letting attention work on smaller resolutions. Here is a simple practical instantiation of this hybrid design: Replace only the final three bottleneck blocks of a ResNet with BoT blocks without any other changes. Or in other words, take a ResNet and only replace the final three $3\times 3$ convolutions with MHSA layers (Fig 1, Table 1). This simple change improves the mask AP by 1.2% on the COCO instance segmentation benchmark over our canonical baseline that uses ResNet-50 in the Mask R-CNN framework with no hyperparameter differences and minimal overheads for training and inference. Moving forward, we call this simple instantiation as BoTNet given its connections to the Transformer through the BoT blocks. While we note that there is no novelty in its construction, we believe the simplicity and performance make it a useful reference backbone architecture that is worth studying.

Using BoTNet, we demonstrate significantly improved results on instance segmentation without any bells and whistles such as Cascade R-CNN , FPN changes , hyperparameter changes , etc. A few key results from BoTNet are: (1) Performance gains across various training configurations (Section 4.1), data augmentations (Section 4.2) and ResNet family backbones (Section 4.4); (2) Significant boost from BoTNet on small objects (+2.4 Mask AP and +2.6 Box AP) (Appendix); (3) Performance gains over Non-Local layers (Section 4.6); (4) Gains that scale well with larger images resulting in 44.4% mask AP, competitive with state-of-the-art performance among entries that only study backbone architectures with modest training schedules (up to 72 epochs) and no extra data or augmentations.SoTA is based on https://paperswithcode.com/sota/instance-segmentation-on-coco-minival..

Lastly, we scale BoTNets, taking inspiration from the training and scaling strategies in , after noting that BoTNets do not provide substantial gains in a smaller scale training regime. We design a family of BoTNet models that achieve up to 84.7% top-1 accuracy on the ImageNet validation set, while being upto 1.64x faster than the popular EfficientNet models in terms of compute time on TPU-v3 hardware. By providing strong results through BoTNet, we hope that self-attention becomes a widely used primitive in future vision architectures.

Related Work

A taxonomy of deep learning architectures that employ self-attention for vision is presented in Figure 2. In this section, we focus on: (1) Transformer vs BoTNet; (2) DETR vs BoTNet; (3) Non-Local vs BoTNet.

Connection to the Transformer: As the title of the paper suggests, one key message in this paper is that ResNet bottleneck blocks with Multi-Head Self-Attention (MHSA) layers can be viewed as Transformer blocks with a bottleneck structure. This is visually explained in Figure 3 and we name this block as Bottleneck Transformer (BoT). We note that the architectural design of the BoT block is not our contribution. Rather, we point out the relationship between MHSA ResNet bottleneck blocks and the Transformer with the hope that it improves our understanding of architecture design spaces for self-attention in computer vision. There are still a few differences aside from the ones already visible in the figure (residual connections and block boundaries): (1) Normalization: Transformers use Layer Normalization while BoT blocks use Batch Normalization as is typical in ResNet bottleneck blocks ; (2) Non-Linearities: Transformers use one non-linearity in the FFN block, while the ResNet structure allows BoT block to use three non-linearities; (3) Output projections: The MHSA block in a Transformer contains an output projection while the MHSA layer (Fig 4) in a BoT block (Fig 1) does not; (4) We use the SGD with momentum optimizer typically used in computer vision while Transformers are generally trained with the Adam optimizer .

Connection to DETR: Detection Transformer (DETR) is a detection framework that uses a Transformer to implicitly perform region proposals and localization of objects instead of using an R-CNN . Both DETR and BoTNet attempt to use self-attention to improve the performance on object detection and instance (or panoptic) segmentation. The difference lies in the fact that DETR uses Transformer blocks outside the backbone architecture with the motivation to get rid of region proposals and non-maximal suppression for simplicity. On the other hand, the goal in BoTNet is to provide a backbone architecture that uses Transformer-like blocks for detection and instance segmentation. We are agnostic to the detection framework (be it DETR or R-CNN). We perform our experiments with the Mask and Faster R-CNN systems and leave it for future work to integrate BoTNet as the backbone in the DETR framework. With visibly good gains on small objects in BoTNet, we believe there maybe an opportunity to address the lack of gain on small objects found in DETR, in future (refer to Appendix).

Connection to Non-Local Neural Nets:The replacement vs insertion contrast has previously been pointed out in AA-ResNet (Bello et. al) . The difference in our work is the complete replacement as opposed to fractional replacement in Bello et al. Non-Local (NL) Nets make a connection between the Transformer and the Non-Local-Means algorithm . They insert NL blocks into the final one (or) two blockgroups (c4,c5) in a ResNet and improve the performance on video recognition and instance segmentation. Like NL-Nets , BoTNet is a hybrid design using convolutions and global self-attention. (1) Three differences between a NL layer and a MHSA layer (illustrated in Figure 4): use of multiple heads, value projection and position encodings in MHSA; (2) NL blocks use a bottleneck with channel factor reduction of $2$ (instead of $4$ in BoT blocks which adopt the ResNet structure); (3) NL blocks are inserted as additional blocks into a ResNet backbone as opposed to replacing existing convolutional blocks as done by BoTNet. Section 4.6 offers a comparison between BoTNet, NLNet as well as a NL-like version of BoTNet where we insert BoT blocks in the same manner as NL blocks instead of replacing.

Method

BoTNet by design is simple: replace the final three spatial ($3\times 3$) convolutions in a ResNet with Multi-Head Self-Attention (MHSA) layers that implement global (all2all) self-attention over a 2D featuremap (Fig 4). A ResNet typically has 4 stages (or blockgroups) commonly referred to as [c2,c3,c4,c5] with strides relative to the input image, respectively. Stacks [c2,c3,c4,c5] consist of multiple bottleneck blocks with residual connections (e.g, R50 has bottleneck blocks).

Approaches that use self-attention throughout the backbone are feasible for input resolutions ($224\times 224$ (for classification) and $640\times 640$ (for detection experiments in SASA )) considered in these papers. Our goal is to use attention in more realistic settings of high performance instance segmentation models, where typically images of larger resolution ($1024\times 1024$) are used. Considering that self-attention when performed globally across $n$ entities requires $O(n^{2}d)$ memory and computation , we believe that the simplest setting that adheres to the above factors would be to incorporate self-attention at the lowest resolution featuremaps in the backbone, ie, the residual blocks in the c5 stack. The c5 stack in a ResNet backbone typically uses 3 blocks with one spatial $3\times 3$ convolution in each. Replacing them with MHSA layers forms the basis of the BoTNet architecture. The first block in c5 uses a $3\times 3$ convolution of stride $2$ while the other two use a stride of $1$. Since all2all attention is not a strided operation, we use a $2\times 2$ average-pooling with a stride $2$ for the first BoT block. The BoTNet architecture is described in Table 1 and the MHSA layer is presented in Figure 4. The strided version of the BoT block is presented in the Appendix.

In order to make the attention operation position aware, Transformer based architectures typically make use of a position encoding . It has been observed lately that relative-distance-aware position encodings are better suited for vision tasks . This can be attributed to attention not only taking into account the content information but also relative distances between features at different locations, thereby, being able to effectively associate information across objects with positional awareness. In BoTNet, we adopt the 2D relative position self-attention implementation from .

Experiments

We study the benefits of BoTNet for instance segmentation and object detection. We perform a thorough ablation study of various design choices through experiments on the COCO dataset . We report the standard COCO metrics including the AP${}^{\text{bb}}_{\text{}}$ (averaged over IoU thresholds), AP${}^{\text{bb}}_{\text{50}}$, AP${}^{\text{bb}}_{\text{75}}$, AP${}^{\text{mk}}_{\text{}}$; AP${}^{\text{mk}}_{\text{50}}$, AP${}^{\text{mk}}_{\text{75}}$ for box and mask respectively. As is common practice these days, we train using the COCO train set and report results on the COCO val (or minival) set as followed in Detectron train - 118K images, val - 5K images. Our experiments are based on the Google Cloud TPU detection codebase https://github.com/tensorflow/tpu/tree/master/models/official/detection. We run all the baselines and ablations with the same codebase. Unless explicitly specified, our training infrastructure uses v3-8 Cloud-TPU which contains 8 cores with 16 GB memory per core. We train with the bfloat16 precision and cross-replica batch normalization using a batch size of 64.

We consider the simplest and most widely used setting: ResNet-50We use the ResNet backbones pre-trained on ImageNet classification as is common practice. For BoTNet, the replacement layers are not pre-trained but randomly initialized for simplicity; the remaining layers are initialized from a pre-trained ResNet. backbone with FPNFPN refers to Feature Pyramid Network . We use it in every experiment we report results on, and our FPN levels from 2 to 6 (p2 to p6) similar to Detectron .. We use images of resolution $1024\times 1024$ with a multi-scale jitter of $[0.8,1.25]$ (scaling the image dimension between $820$ and $1280$, in order to be consistent with the Detectron setting of using $800\times 1300$). In this setting, we benchmark both the ResNet-50 (R50) and BoT ResNet-50 (BoT50) as the backbone architectures for multiple training schedules: 1x: 12 epochs, 2x: 24 epochs, 3x: 36 epochs, 6x: 72 epochs1x, 2x, 3x and 6x convention is adopted from MoCo ., all using the same hyperparameters for both the backbones across all the training schedules (Table 2). We clearly see that BoT50 is a significant improvement on top of R50 barring the 1x schedule (12 epochs). This suggests that BoT50 warrants longer training in order to show significant improvement over R50. We also see that the improvement from BoT50 in the 6x schedule (72 epochs) is worse than its improvement in the 3x schedule (32 epochs). This suggests that training much longer with the default scale jitter hurts. We address this by using a more aggressive scale jitter (Section 4.2).

2 Scale Jitter helps BoTNet more than ResNet

In Section 4.1, we saw that training much longer (72 epochs) reduced the gains for BoT50. One way to address this is to increase the amount of multi-scale jitter which has been known to improve the performance of detection and segmentation systems . Table 3 shows that BoT50 is significantly better than R50 ( + 2.1% on AP${}^{\text{bb}}_{\text{}}$ and + 1.7% on AP${}^{\text{mk}}_{\text{}}$) for multi-scale jitter of $[0.5,2.0]$, while also showing significant gains ( + 2.2% on AP${}^{\text{bb}}_{\text{}}$ and + 1.6% on AP${}^{\text{mk}}_{\text{}}$) for scale jitter of $[0.1,2.0]$, suggesting that BoTNet (self-attention) benefits more from extra augmentations such as multi-scale jitter compared to ResNet (pure convolutions).

3 Relative Position Encodings Boost Performance

BoTNet uses relative position encodings . We present an ablation for the use of relative position encodings by benchmarking the individual gains from content-content interaction ($qk^{T}$) and content-position interaction ($qr^{T}$) where $q,k,r$ represent the query, key and relative position encodings respectively. The ablations (Table 4) are performed with the canonical settingres:1024x1024, 36 epochs (3x schedule), multi-scale jitter:[0.8, 1.25]. We see that the gains from $qr^{T}$ and $qk^{T}$ are complementary with $qr^{T}$ more important, ie, $qk^{T}$ standalone contributes to 0.6% AP${}^{\text{bb}}_{\text{}}$ and 0.6% AP${}^{\text{mk}}_{\text{}}$ improvement over the R50 baseline, while $qr^{T}$ standalone contributes to 1.0% AP${}^{\text{bb}}_{\text{}}$ and 0.7 % AP${}^{\text{mk}}_{\text{}}$ improvement. When combined together ($qk^{T}+qr^{T}$), the gains on both AP${}^{\text{bb}}_{\text{}}$ and AP${}^{\text{mk}}_{\text{}}$ are additive ( 1.5% and 1.2% respectively). We also see that using absolute position encodings ($qr_{\textrm{abs}}^{T}$) does not provide as much gain as relative. This suggests that introducing relative position encodings into architectures like DETR is an interesting direction for future work.

4 BoTNet improves backbones in ResNet Family

How well does the replacement setup of BoTNet work for other backbones in the ResNet family? Table 5 presents the results for BoTNet with R50, R101, and R152. All these experiments use the canonical training setting (refer to footnote in 4.3). These results demonstrate that BoTNet is applicable as a drop-in replacement for any ResNet backbone. Note that BoT50 is better than R101 (+ 0.3% AP${}^{\text{bb}}_{\text{}}$, + 0.5% AP${}^{\text{mk}}_{\text{}}$) while it is competitive with R152 on AP${}^{\text{mk}}_{\text{}}$. Replacing 3 spatial convolutions with all2all attention gives more improvement in the metrics compared to stacking 50 more layers of convolutions (R101), and is competitive with stacking 100 more layers (R152), supporting our initial hypothesis that long-range dependencies are better captured through attention than stacking convolution layers.Note that while one may argue that the improvements of BoT50 over R50 could be attributed to having 1.2x more M. Adds, BoT50 ($121\times 10^{9}$ M.Adds) is also better than R101 ($162.99\times 10^{9}$ B M. Adds and is competitive with R152 ($240.56\times 10^{9}$ M. Adds) despite performing significantly less computation.

5 BoTNet scales well with larger images

We benchmark BoTNet as well as baseline ResNet when trained on $1280\times 1280$ images in comparison to $1024\times 1024$ using the best config: multi-scale jitter of $[0.1,2.0]$ and training for 72 epochs. Results are presented in Tables 6 and 8. Results in Table 6 suggest that BoTNet benefits from training on larger images for all of R50, R101 and R152. BoTNet trained on $1024\times 1024$ (leave alone $1280\times 1280$) is significantly better than baseline ResNet trained on $1280\times 1280$. Further, BoT200 trained with $1280\times 1280$ achieves a AP${}^{\text{bb}}_{\text{}}$ of 49.7% and AP${}^{\text{mk}}_{\text{}}$ of 44.4%. We believe this result highlights the power of self-attention, in particular, because it has been achieved without any bells and whistles such as modified FPN , cascade RCNN , etc. This result surpasses the previous best published single model single scale instance segmentation result from ResNeSt evaluated on the COCO minival (44.2% AP${}^{\text{mk}}_{\text{}}$).

6 Comparison with Non-Local Neural Networks

How does BoTNet compare to Non-Local Neural Networks? NL ops are inserted into the c4 stack of a ResNet backbone between the pre-final and final bottleneck blocks. This adds more parameters to the model, whereas BoTNet ends up reducing the model parameters (Table 5). In the NL mould, we add ablations where we introduce BoT block in the exact same manner as the NL block. We also run an ablation with the insertion of two BoT blocks, one each in the c4,c5 stacks. Results are presented in Table 7. Adding a NL improves AP${}^{\text{bb}}_{\text{}}$ by 1.0 and AP${}^{\text{bb}}_{\text{}}$ by 0.7, while adding a BoT block gives +1.6 AP${}^{\text{bb}}_{\text{}}$ and +1.2 AP${}^{\text{mk}}_{\text{}}$ showing that BoT block design is better than NL. Further, BoT-R50 (which replaces instead of adding new blocks) provides +1.5 AP${}^{\text{bb}}_{\text{}}$ and + 1.2 AP${}^{\text{mk}}_{\text{}}$, as good as adding another BoT block and better than adding one additional NL block.

7 Image Classification on ImageNet

While we motivated the design of BoTNet for detection and segmentation, it is a natural question to ask whether the BoTNet architecture design also helps improve the image classification performance on the ImageNet benchmark. Prior work has shown that adding Non-Local blocks to ResNets and training them using canonical settings does not provide substantial gains. We observe a similar finding for BoTNet-50 when contrasted with ResNet-50, with both models trained with the canonical hyperparameters for ImageNet : 100 epochs, batch size 1024, weight decay 1e-4, standard ResNet data augmentation, cosine learning rate schedule (Table 9). BoT50 does not provide significant gains over R50 on ImageNet though it does provide the benefit of reducing the parameters while maintaining comparable computation (M.Adds).

A simple method to fix this lack of gain is to take advantage of the image sizes typically used for image classification. In image classification, we often deal with much smaller image sizes ($224\times 224$) compared to those used in object detection and segmentation ($1024\times 1024)$. The featuremaps on which the BoT blocks operate are hence much smaller (e.g $14\times 14$, $7\times 7$) compared to those in instance segmentation and detection (e.g $64\times 64$, $32\times 32$). With the same number of parameters, and, without a significant increase in computation, the BoTNet design in the c5 blockgroup can be changed to uniformly use a stride of $1$ in all the final MHSA layers. We call this design as BoTNet-S1 (S1 to depict stride 1 in the final blockgroup). We note that this architecture is similar in design to the hybrid models explored in Vision Transformer (ViT) that use a ResNet up to stage c4 prior to stacking Transformer blocks. The main difference between BoTNet-S1 and the hybrid ViT models lies in the use of BoT blocks as opposed to regular Transformer blocks (other differences being normalization layer, optimizer, etc as mentioned in the contrast to Transformer in Related Work (Sec. 2). The architectural distinction amongst ResNet, BoTNet and BoTNet-S1, in the final blockgroup, is visually explained in the Appendix). The strided BoT block is visually explained in the Appendix.

7.2 Evaluation in the standard training setting

We first evaluate this design for the 100 epoch setting along with R50 and BoT50. We see that BoT-S1-50 improves on top of R50 by 0.9% in the regular setting (Table 9). This improvement does however come at the cost of more computation (m.adds). Nevertheless, the improvement is a promising signal for us to design models that scale well with larger images and improved training conditions that have become more commonly used since EfficientNets .

7.3 Effect of data augmentation and longer training

We saw from our instance segmentation experiments that BoTNet and self-attention benefit more from regularization such as data augmentation (in the case of segmentation, increased multi-scale jitter) and longer training. It is natural to expect that the gains from BoT and BoT-S1 could improve when training under an improved setting: 200 epochs, batch size 4096, weight decay 8e-5, RandAugment (2 layers, magnitude 10), and label smoothing of 0.1. In line with our intuition, the gains are much more significant in this setting for both BoT50 (+ 0.6%) and BoT-S1-50 (+ 1.4%) compared to the baseline R50 (Table 10).

7.4 Scaling BoTNets

The previous ablations show the BoNets performance with a ResNet-50 backbone and $224\times 224$ image resolution. Here we study BoTNets when scaling up the model capacity and image resolution. There have been several works improving the performance of ConvNets on ImageNet . Bello et al. recently propose scaling strategies that mainly increase model depths and increase the image resolutions much slower compared to the compound scaling rule proposed in EfficientNets . We use similar scaling rules and design a family of BoTNets. The details of model depth and image resolutions are in the Appendix. We compare to the SENets baseline to understand the impact of the BoT blocks. The BoTNets and SENets experiments are performed under the same training settings (\eg, regularization and data augmentation). We additionally show EfficientNet and DeiT (regularized version of ViT )ViT refers to Vision Transformer , while DeiT refers to Data-Efficient Image Transformer . DeiT can be viewed as a regularized version of ViT with augmentations, better training hyperparameters tuned for ImageNet, and knowledge distillation . We do not compare to the distilled version of DeiT since it’s an orthogonal axis of improvement applicable to all models. to understand the performance of BoTNets compared with popular ConvNets and Transformer models. EfficientNets and DeiT are trained under strong data augmentation, model regularization, and long training schedules, similar to the training settings of BoTNets in the experiments.

ResNets and SENets are strong baselines until 83% top-1 accuracy. ResNets and SENets achieve strong performance in the improved EfficientNet training setting. BoTNets T3 and T4 do not outperform SENets, while T5 does perform on par with S4. This suggests that pure convolutional models such as ResNets and SENets are still the best performing models until an accuracy regime of 83%. BoTNets scale better beyond 83% top-1 accuracy. While SENets are a powerful model class outperforms BoTNets (up to T4), we found gains to diminish beyond SE-350 (350 layer SENet described in Appendix) trained with image size 384. This model is referred to as S5 and achieves 83.8% top-1 accuracy. On the other hand, BoTNets scale well to larger image sizes (corroborating with our results in instance segmentation when the gains from self-attention were much more visible for larger images). In particular, T7 achieves 84.7% top-1 acc., matching the accuracy of B7-RA, with a 1.64x speedup in efficiency. BoTNets perform better than ViT-regularized (DeiT-384), showing the power of hybrid models that make use of both convolutions and self-attention compared to pure attention models on ImageNet-1K.

Conclusion

The design of vision backbone architectures that use self-attention is an exciting topic. We hope that our work helps in improving the understanding of architecture design in this space. Incorporating self-attention for other computer vision tasks such as keypoint detection and 3D shape prediction ; studying self-attention architectures for self-supervised learning in computer vision ; and scaling to much larger datasets such as JFT, YFCC and Instagram, are ripe avenues for future research. Comparing to, and incorporating alternatives to self-attention such as lambda-layers is an important future direction as well.

Acknowledgements

We thank Ilija Radosavovic for several useful discussions; Pengchong Jin and Xianzhi Du for help with the TF Detection codebase; Irwan Bello, Barret Zoph, Neil Houlsby, Alexey Dosovitskiy for feedback. We thank Zak Stone for extensive compute support throughout this project the through TFRC program providing Google Cloud TPUs (https://www.tensorflow.org/tfrc).

References