NextViT

Vision Transformer for Efficient Deployment

1. Introduction

This notebook is on the NextViT backbone, which incorporates the use of convolutional and Transformer blocks in an efficient manner.

Many works struggle to free the ViT architectures from high latency. Swin Transformer and PVT try to design more efficient spatial attention mechanisms to alleviate the quadratic-increasing computation which is highly complex from our mHSA (multi head self attention). Other consider combining efficient convolution blocks and powerful Transformer blocks to design a CNN-Transformer hybrid approach. Which should then achieve a good trade-off between accuracy which is boosted via the Transformer architecture, and the latency improvement, which comes from our convolutional blocks.

Many of these hybrid architectures use the convolutional blocks in the more shallow stages (Meaning the early/high resolution stages). And then later stack Transformer blocks in the last few stages. NextViT has noticed that such an approach leads to performance saturation on downstream tasks, such as object segmentation/detection. They have also found that both the convolutional blocks and Transformer blocks in those works, can not possess characteristics of efficiency and performance at the same time.

Even though the latency has improved, the accuracy is far from what can be achieved.

1.1 How the issues are addressed

The platforms we will be benchmarking on is TensorRT and CoreML, which are industrially standard approaches for deployments, server-side and mobile-side respectively.

comparison-among-nextvit-and-efficient-networks-in-terms-of-accuracy-latency-tradeoff

The graphs shows the mixture between latency and accuracy, and how NextViT consistently achieves the "best" symbiosis between both metrics in both TensorRT and CoreML given the 3 major datasets: ImageNet-1K, COCO, ADE20K.

Their main contributions are:

2. Recap: CNN & ViT

Convolutional Networks: CNNs have dominated vision architectures in a variety of cv (computer vision) tasks, including image classification, object detection and semantic segmentation. ResNet uses residual connections to eliminate network degradation (Ground-breaking paper which unlocked several new architectures). DenseNet alternately enhances feature reuse and concatenates feature maps through dense connections. MobileNets introduce depthwise convolution and point-wise convolution to build models that have a smaller memory layer with lower latency.

Convolutional networks essentially downsample their input feature map (usually just an RGB image), and can perform some classification there. If we are working with segmentation/detection then it would upsample back to its original (typically) dimensions. And these downsampling/upsampling paths have weights which are learnt (upsampling typically happens via some sort of interpolation which is actually parameter-free). And the different CNN architectures then have some specifically developed modules which assists this core architecture of downsampling/upsampling.

Vision Transformers: The transformer architecture was first introduced within the field of NLP (Natural Language Processing). Then slowly integrated into CV (Computer Vision). Where ViT splits the image into patches and treats those patches as tokens to perform self attention with, the same way these tokens would represent words in NLP. These ViT based models achieve very good metrics. Where DeiT is a model which introduces a teacher-student strategy. Teacher model's predictions are used as soft training targets for a smaller student model, so the student learns from the teacher's full output distribution rather than only the hard ground-truth labels, in DeiT's case, via an extra distillation token that attends to the teacher's output alongside the class token.

Hybrid Models: Recent papers have shown that the combination of convolutional and transformer blocks as a hybrid approach helps absorb the strengths of both. Which is the global context (ViT) plus the local inductive bias (CNN). As convolutions bring locality and translation equivariance, which helps data efficiency and low-level features, not just speed. While ViT's aren't inherently more accurate, but better at long range dependencies, where patches of pixels far from one another can communicate.

3. Overview:

NextViT follows a hierarchical pyramid architecture that makes use of patch embeddings between every layer and a series of convolutions and/or transformer blocks in each stage. The spatial resolution is reduced by a factor of 32 at every stage, while the channel dimension will be expanded across the stages. I will underline the mechanisms of NCB and NTB to model the short term and long term dependencies in our visual data.

3.1 Next Convolutional Block (NCB)

To underline the improvements to this block, the paper then revisits some of the more classical architectures of convolution and Transformer blocks.

variation-of-convolutional-transformer-based-architectures

BottleNeck block proposed by ResNet had dominance in the visual neural networks for some time, as it was inherently inductive biased and had very friendly deployment features. However, the effectiveness of this BottleNeck block was inadequate compared to the Transformer block. Then ConvNeXt block started modernizing the BottleNeck block by imitating design from the Transformer blocks. While ConvNeXt block somewhat improves the network performance, its inference speed on TensorRT/CoreML is very limited by inefficient components, such as 7 x 7 depthwise convolution, LayerNorm, and GELU.

Transformer based blocks had achieved excellent results in various visual tasks and its superiority is appreciated by the paradigm of MetaFormer, and the attention based token mixer module. However once again, the inference speed of this Transformer block is much slower than the original BottleNeck block due to its complex attention mechanisms. Which is unacceptable in most real world applications.

NextViT overcomes the defects from the blocks shown in the image, by introducing the NCB. Which maintains the deployment advantage of BottleNeck block while obtaining prominent performance as a Transformer block.

As shown in drawing (f) NCB follows the general architecture of MetaFormer, which has been verified to be essential to the Transformer block. NextViT designs a novel Multi-Head Convolutional Attention block, as an efficient token mixer with deployment friendly convolution operation.

They build NCB with MHCA and MLP layer in the paradigm of MetaFormer. The proposed NCB can be formulated as follows:

Where ${z}^{l-1} $ is the output from the block which is used as our input, and are the outputs of MHCA and the NCB.

3.2 Multi-Head Convolutional Attention (MCHA)

To escape the high altency dilemma from the attention-based token mizer, they design a novel attention mechanism with efficient convolutiuon operation, i.e. Convolutiuonal Attention (CA) for fast inference speed. They are inspired by the effective multi head design in MHSA, they then build a convolutional attention iwth multiple heads, which jointly attends to information from different representation spaces, at a different position for effective local representation learning.

The definition of the proposal MHCA can be formulated as:

MHCA captures information from h parallel representation subspaces. Simply put, you split the feature's channels into h groups, and each group is a "subspace", meaning a different learned view of the same pixels.

F.ex: 96 channels, head dim 32 -> 3 heads of 32 channels each. Head 1 gets channels 0-31, head 2 gets 32-63, head 3 gets 64-95. Each has its own 3x3 kernel, which spans 9 neighboring tokens, so the mixing across tokens happens inside each head, while the heads themselves are just channel groups. So one might learn edges, another texture, another color gtradients, independently. In parallel. The (1x1) conv at the end then lets them talk.

The term comes from the original transformer paper, here it's the same idea applied to grouped conv rather than attention.

Think of this as, the ouput at a position is the inner product of a learned kernel with that position's neighbors. Which is the definition of a convolution. The paper frames it as "attention" because W learns the affinity ebtween neighnoring tokens, but unlike real attention, W is fixed after training rather than computed from the input. Which is what makes it fast

  1. Split: carves the channels into h groups of 32 ( in this example)
  2. Independent CA (convolutional attention) per group: each gets its own kernel , so no cross group communicatino. Concatenating these back together is the grouped 3x3 conv
  3. Project: is a conv that mixes the heads, since step 2 left them isolated

So the entire equation is just: grouped conv -> BN -> ReLu -> 1x1 conv, which is exactly the MHCa box in the figure. CA mixes across tokens but not across heads. mixes across heads but not across tokens. When put together, they then cover both axes.

overview-of-mhca

Simply put:

We take our input, split the channel dimension by the number of heads we have. Then per head, we perform groupped convolution throughout the channel subspace we are given in that head. In that sense we have performed attention onto the adjacent tokens, which is what represent. From there we concatenate all our heads, where we are back to input dimensions. We then perform 1x1 with , for every token which then allows the channels of the token to speak to each other, whereas before they were isolated in their separate subspaces.

3.3 Next Transformer Block (NTB)

Now our local "attention" has been performed. NextViT underlines the importance of addressing the global attention. The Transformer block has a strong ability to capture low-frquency signals which provide global information, such as global shapes and structures. Nonetheless relevant studies, have observed that these Trasnformer blocks may detriorate high frquency information, such as local texture sinformation. These signals coming from different frqency segmentas are both important for the human visual system, and must be fused in a specific manner to extract more essential and distainct features.

NextViT was motivated by these findings, and therefore develop NTB which is used to capture multi-frquency signals in a lightweight manner.

nextvit-overview

As shown in the above figure, NTB first captures low frquency signals with an Efficient Multi-Head Self Attention (E-MHSA) block, which can be written as:

where denotes to divide the input feature z into multi head form in channel dimension. SA is a spatial reduction self attention operator which is inspired by Linear SRA, and performing as:

where Attention represents the standard attention calculating as . For recap, are linear layers for context encoding, while is an avg-pool operation with stride s for downsampling the spatial dimension before we perform attention, so we save computation. E-MHSA is heavily affected by the channel dimension, therefore the NTB block also performs a channel dimension reduction before entering the E-MHSA module with point-wise convolutions to further accelerate inference. A ratio is introduced for the channel reduction. Batch Normalization is also used in E-MHSA for more efficient deployment, as BN is "foldable." Which deservers a notebook by itself.

NTB is also using MHCA module that works with the E-MHSA module to capture the multi-frequency signals. At last, the output features from E-MHSA and MHCA are concatenated to mix high and low frquency information. MLP layer is then used at the end to extract more infromation an distinct feature. The order of operations for NTB can be broken down as follows:

In english:

  1. We take our feature map, perform point-wise convolution to reduce channel count given our ratio.
  2. Then we performn E-MHSA on our reduced channel feature map.
  3. We take that output, and project it back to original channel count.
  4. We perform MHCA on the above output with the original channel count
  5. We then concatenate this output with the raw output from E-MHSA as well.
  6. MLP on the final output

3.4 Next Hybrid Strategy (NHS)

The paper mentions how many approaches, have simply relied on convolutional blocks on the earlier stages, then stacking transformer blocks on the later stages. This performs well, and is effective for the classification task, but hinders the performance in tasks such as segmentation and detection. Their guess to this is, that the classification task simply uses the outputs from the last stage for prediction. While other tasks (seg. & det.) usually rely on the features from each stage (looking at different resolutions of the input) to gain better results. And this reliance on convolutional blocks at the earlier stage, ignores capturing global information which the transofmer block is great at.

NHS comes into play, with a different structure, where they stack convolutional blocks (NCB) and Transformer blocks (NTB) with (N + 1) * hybrid paradigm. This significantly promotes model performance in detection & segmentation (or atleast they say it does 0-0, but the metrics in the below table backs it up).

To allow global capturing in the earlier stages, they present a strategy with the following pattern , which sequentially stack N \text{Next-ViT}(X) = \oint_{i} \left\{ \left[ \Gamma\!\left( \Psi(X) \times N_{i} \right) \right] \times L_{i} \right\}

where denotes the stage index. denotes NCB. denotes identity layer when i - 1, otherwise, NTB. indicates the operation of stacking the stages sequentailly.

They then run plenty of ablations which can be seen below in the results section.

4. Recap

NextViT is a hybrid architecture, consisting of convolutional & transformer blocks, where both blocks perform their own attention in an efficient manner, both mathematically, but also in a computational efficient manner which allows for friendly deployment with quick inference, which is something which often lacks with this hybrid models.

5. Results:

Latency measured on TensorRT-8.0.3 (T4 GPU, batch 8) and CoreML (iPhone 12 Pro Max, batch 1).

ImageNet-1K classification

Model · Params · FLOPs · TensorRT · CoreML · Top-1 ResNet101 · 44.6M · 7.9G · 7.8 ms · 4.0 ms · 80.8% Swin-T · 29.0M · 4.5G · 24.5 ms · – · 81.3% ConvNeXt-T · 29.0M · 4.5G · 19.0 ms · 83.8 ms · 82.1% Next-ViT-S · 31.7M · 5.8G · 7.7 ms · 3.5 ms · 82.5% CSWin-T · 23.0M · 4.3G · 29.5 ms · – · 82.7% Swin-S · 50.0M · 8.7G · 20.1 ms · 137.3 ms · 83.0% ConvNeXt-S · 50.0M · 8.7G · 28.1 ms · 159.5 ms · 83.1% Next-ViT-B · 44.8M · 8.3G · 10.5 ms · 4.5 ms · 83.2% EfficientFormer-L7 · 82.0M · 7.9G · 17.4 ms · 6.9 ms · 83.3% TRT-ViT-D · 103.0M · 9.7G · 15.1 ms · 8.3 ms · 83.4% Next-ViT-L · 57.8M · 10.8G · 13.0 ms · 5.5 ms · 83.6% ConvNeXt-B · 88.0M · 15.4G · 37.3 ms · 247.6 ms · 83.9%

With SSLD pretraining, S/B/L reach 84.8 / 85.1 / 85.4%, and 86.4% at 384 input.

ADE20K segmentation (Semantic FPN 80k)

Backbone · TensorRT · CoreML · mIoU ResNet101 · 32.8 ms · 13.2 ms · 38.8% Swin-T · – · – · 41.5% EfficientFormer-L3 · 35.9 ms · 10.6 ms · 43.5% Next-ViT-S · 38.2 ms · 18.1 ms · 46.5% EfficientFormer-L7 · 84.0 ms · 23.0 ms · 45.1% CSWin-T · 182.3 ms · – · 48.2% Next-ViT-B · 51.6 ms · 24.4 ms · 48.6% CSWin-S · 298.2 ms · – · 49.2% Next-ViT-L · 65.3 ms · 30.1 ms · 49.1%

COCO detection (Mask R-CNN 1×)

Backbone · AP<sup>b</sup> · AP<sup>m</sup> ResNet101 · 40.4 · 36.4 Swin-T · 42.2 · 39.1 PVTv2-B2 · 45.3 · 41.2 Next-ViT-S · 45.9 · 41.8 Swin-S · 44.8 · 40.9 CSWin-T · 46.7 · 42.2 Next-ViT-B · 47.2 · 42.8 PVTv2-B4 · 47.5 · 42.7 CSWin-S · 47.9 · 43.2 Next-ViT-L · 48.0 · 43.2

Ablations (matched TensorRT latency)

Ablation · Result Hybrid pattern · C H_N H_N H_N best; a transformer in stage 1 hurts Repetition · (NCB×N + NTB×1)×L beats one large group; N=4 optimal Block type · NCB 82.5% vs BottleNeck 81.9, ConvNeXt 79.6, LSA 78.2 Shrink ratio · r=0.75 beats pure transformer (r=1.0): 82.5% vs 82.1%, and faster Norm / activation · BN+ReLU 7.7 ms / 82.5% vs LN+GELU 9.3 ms / 82.7%

Fourier spectrum analysis shows ResNet captures high-frequency signal, Swin low-frequency, and Next-ViT both.