Next-ViT: Next Generation Vision Transformer for Efficient Deployment in Realistic Industrial Scenarios

In plain words

Vision Transformers are accurate but slow on the chips that run real products, such as a server GPU with TensorRT or an iPhone with CoreML. Plain convolutional networks are fast on those chips but less accurate. Many recent designs mix the two, and the authors found that the usual mixes stop helping on detection and segmentation.

Next-ViT keeps the speed of convolutions and adds a light form of attention. Its convolution block mixes tokens with a grouped convolution, and uses BatchNorm and ReLU in place of LayerNorm and GELU, which these chips run quickly. Its attention block works on a shrunken, pooled set of tokens, so it costs little. Stages 2 to 4 end with one attention block, after several convolution blocks, so those stages also get global context. The authors report that this gives a better speed-accuracy trade-off than the earlier CNNs, Transformers and hybrids they compare with, on classification, detection and segmentation.

TL;DR

• Next-ViT is a four-stage hybrid backbone. Stages 2 to 4 each stack N convolution blocks (NCB) and end with one Transformer block (NTB), written (NCB×N + NTB); stage 1 is convolution only §3.4 eq. 7.
• Both blocks use BatchNorm and ReLU, not LayerNorm and GELU, and the attention in NTB is pooled to cut its cost §3.2 §3.3.
• Author-reported: 82.5% ImageNet-1K top-1 at 7.7 ms TensorRT for Next-ViT-S at 224 px tab. 4; 45.9 box AP for Next-ViT-S on COCO Mask R-CNN 1x tab. 6.
• Use it when you need a backbone whose latency you can measure on your target chip. Do not expect the speed rankings to transfer to other hardware; the paper measures only TensorRT and CoreML.

Problem

Standard Transformer blocks are slow on TensorRT and CoreML because of multi-head self-attention, LayerNorm and GELU. Those operations are costly on these runtimes §1. The authors also observe that hybrid CNN-Transformer networks that put Transformer blocks only in the last one or two stages reach a performance ceiling on dense tasks such as detection and segmentation, because those tasks need global context in every stage §3.4.

The paper's target is the latency-accuracy trade-off measured on those two runtimes, not FLOPs §1 §4.1.1.

Method

Next-ViT follows a hierarchical pyramid. A stem reduces the image by 4×, then three further stages reduce it by 2× each, giving a 32× total reduction. Each stage starts with a patch embedding and then runs a sequence of blocks §3.1 tab. 3 fig. 2. The blocks are of two kinds.

1. Next Convolution Block (NCB) §3.2. It follows the MetaFormer layout, a token mixer followed by an MLP, both with residual connections eq. 1. The token mixer is Multi-Head Convolutional Attention (MHCA), a grouped convolution with 32 channels per group, followed by a 1×1 projection eq. 2 eq. 3 §3.2.1. Normalisation is BatchNorm and the activation is ReLU.
2. Next Transformer Block (NTB) §3.3. It first reduces the channels by a shrink ratio r = 0.75 with a 1×1 projection §3.5. It then applies Efficient Multi-Head Self-Attention (E-MHSA) to the reduced channels, whose keys and values come from average-pooled tokens eq. 4 eq. 5. The output of E-MHSA is concatenated with a second branch that projects the attended features and applies MHCA to them, and the result goes through an MLP eq. 6.
3. Next Hybrid Strategy (NHS) §3.4. Each stage is a repetition of the pattern (NCB×N + NTB×1). The NTB sits at the end of each stage. The first stage has no NTB, so it is pure NCB eq. 7 tab. 2. The paper uses N = 4 for the basic models §3.4.

Next-ViT-S, -B and -L differ only in the number of repeats in stage 3; widths and the other stages are shared tab. 3. Stage 3 holds 10, 20 or 30 blocks in groups of 5 (four NCB and one NTB), so the model depths are 3-4-10-3, 3-4-20-3 and 3-4-30-3 §3.5 tab. 3.

The network head is a BatchNorm, global average pooling and a linear layer nextvit.py. The paper does not describe the head.

Pseudocode for the three components. Equation and table references are to the paper; the code citations show where each step is implemented in the official classification code.

# Next-ViT forward pass, classification  [tab. 3] [eq. 7]
# Input:  images (B, 3, 224, 224)
# Output: logits (B, num_classes)

def next_vit(images):
    # 1. Convolutional stem, 4x downsampling  [tab. 3]
    #    [code: classification/nextvit.py#L296-L301]
    x = stem(images)                          # (B, 64, H/4, W/4)
    # 2. Four stages; NCB blocks first, NTB last in each stage  [eq. 7]
    #    [code: classification/nextvit.py#L285-L294]
    for stage in stages:
        for block in stage.blocks:
            x = block(x)
    # 3. BatchNorm, global average pooling, linear classifier
    #    [code: classification/nextvit.py#L330-L335]
    x = avg_pool(batch_norm(x))               # (B, C)
    return linear(x)

# Next Convolution Block (NCB)  [eq. 1] [eq. 2] [eq. 3]
# Input:  x (B, C_in, H, W); output (B, C_out, H', W')

def ncb(x):
    # 1. Patch embedding: 2x2 average pool when stride 2, then 1x1 conv
    #    and BatchNorm  [tab. 3] [code: classification/nextvit.py#L46-L67]
    x = patch_embed(x)
    # 2. MHCA residual: 3x3 grouped conv with 32 channels per group,
    #    BatchNorm, ReLU, then 1x1 projection  [eq. 2] [eq. 3]
    #    [code: classification/nextvit.py#L77-L88]
    x = x + mhca(x)
    # 3. BatchNorm, then MLP residual with expansion 3, ReLU  [eq. 1]
    #    [code: classification/nextvit.py#L139-L147]
    x = x + mlp(batch_norm(x))
    return x

# Next Transformer Block (NTB)  [eq. 4] [eq. 5] [eq. 6]
# Input:  x (B, C_in, H, W); C_out is the stage width (256, 512 or 1024)
# C_mhsa = round(0.75 * C_out) to a multiple of 32; C_mhca = C_out - C_mhsa

def ntb(x):
    # 1. Patch embedding to C_mhsa channels  [eq. 6]
    #    [code: classification/nextvit.py#L255-L257]
    x = patch_embed(x)
    # 2. Global branch: E-MHSA on BatchNorm-ed tokens, residual  [eq. 4]
    #    [code: classification/nextvit.py#L258-L264]
    tokens = flatten(batch_norm(x))           # (B, H*W, C_mhsa)
    x = x + reshape(e_mhsa(tokens))           # back to (B, C_mhsa, H, W)
    # 3. Local branch: project to C_mhca channels, then MHCA residual
    #    [eq. 6] [code: classification/nextvit.py#L266-L268]
    local = projection(x)
    local = local + mhca(local)
    # 4. Concatenate both branches, then MLP residual (expansion 2)
    #    [eq. 6] [code: classification/nextvit.py#L270-L275]
    x = concat([x, local], channels)          # (B, C_out, H, W)
    return x + mlp(batch_norm(x))

# Efficient multi-head self-attention (E-MHSA)  [eq. 4] [eq. 5]
# Input:  tokens (B, N, C), head dim 32, reduction r per stage
#         r = 8, 4, 2, 1 for stages 1 to 4  [§3.5]

def e_mhsa(tokens):
    # 1. Queries from every token  [eq. 5]
    #    [code: classification/nextvit.py#L186-L188]
    queries = linear_q(tokens)
    # 2. Keys and values from pooled tokens. The code pools r*r
    #    consecutive tokens in raster order, a 1-D pool, not an
    #    r-by-r spatial block  [code: classification/nextvit.py#L190-L199]
    pooled = avg_pool_1d(tokens, kernel=r * r, stride=r * r)
    if r > 1:
        pooled = batch_norm_1d(pooled)
    keys = linear_k(pooled)
    values = linear_v(pooled)
    # 3. Scaled dot-product attention per head, scale = 32 ** -0.5
    #    [code: classification/nextvit.py#L154-L160] [eq. 5]
    weights = softmax(queries @ keys_transposed * scale)
    return linear_proj(weights @ values)

Hyperparameters

The paper states some training settings in the text §4.1.1 §4.2.1 §4.3.1. Where it does not, the value is taken from the official classification, detection or segmentation code or from the README command, and is marked as such. Values in code columns are the defaults in the file. Where the README command overrides a default, the README value is the one the authors used for their reported numbers.

ImageNet-1K classification

COCO Mask R-CNN (mmdetection)

ADE20K segmentation

Latency measurement

Results

Author-reported. Setup: ImageNet-1K training at 224 px, with 384 px fine-tuning §4.1.1; COCO with Mask R-CNN and ImageNet-1K pretraining §4.3.1; ADE20K with Semantic FPN and UperNet at 512 × 512 §4.2.1. Latency is measured by the authors under one setup §4.1.1. The text does not say whether the baseline accuracies were re-run or copied from their papers.

ImageNet-1K classification tab. 4

Models marked † in the paper use large-scale pretraining. Next-ViT-S†, B† and L† reach 84.8, 85.1 and 85.4 at 224 px tab. 4. Next-ViT-S improves on ResNet101 by 1.7 points at a similar TensorRT latency and runs at 3.5 ms on CoreML against 4.0 ms §4.1.2.

COCO Mask R-CNN tab. 6. Latency is backbone only, at 512 × 512.

ResNet101 reaches 40.4 box AP at 1x with 32.8 ms TensorRT backbone latency tab. 6. Next-ViT-S gains 5.5 AP over it §4.3.2 tab. 6, but its backbone latency in the same table is higher, 38.2 ms against 32.8 ms (see Pitfalls).

ADE20K segmentation tab. 5. Latency is backbone only, at 512 × 512.

Next-ViT-S gains 7.7 mIoU over ResNet101 under Semantic FPN §4.2.2 tab. 5.

Design ablations (small-model variants built at similar TensorRT latency §4.4.1) tab. 1 tab. 2 tab. 7 tab. 8 tab. 9:

• Hybrid strategy tab. 1: C H H H (convolution in stage 1, then NHS in stages 2 to 4) gives 82.3 / 45.5 / 46.0 at 7.7 ms. C C C T gives 82.4 / 44.7 / 45.2 at 7.8 ms (classification, detection, segmentation).
• Stacking the NCB×N + NTB pattern more than once per stage tab. 2: (4+1)×2 gives 82.5 / 45.9 / 46.5 at 7.7 ms. Base (4+1)×4 gives 83.2 / 47.2 / 48.6 at 10.5 ms.
• Convolution block tab. 7: NCB gives 82.5 / 45.9 / 46.5 at 7.7 ms, against 79.6 / 41.4 / 43.7 at 7.8 ms for the ConvNeXt block.
• Shrink ratio r tab. 8: r = 0.75 gives 82.5 / 45.9 / 46.5 at 7.7 ms; r = 1.0 gives 82.1 / 45.4 / 45.5 at 8.2 ms.
• Norm and activation tab. 9: BatchNorm with ReLU gives 82.5 / 45.9 / 46.5 at 7.7 ms; LayerNorm with GELU gives 82.7 / 46.0 / 46.7 at 9.3 ms. The paper calls the accuracy gap negligible.

Paper vs code

Code: the official repository at 922e771ae6028b9d956e2622a984459a50aa63d4, Apache-2.0. Paths below are relative to its root. The repository has classification/, detection/ (mmdetection), segmentation/ and deployment/ directories, plus the README and training logs README.md.

• NCB applies BatchNorm before the MLP. Eq. 1 writes MLP(z̃) directly eq. 1. The code normalises the tensor before the MLP nextvit.py.
• NTB applies BatchNorm before both attention and MLP. Eq. 6 shows no normalisation eq. 6. The text says BatchNorm and ReLU are used §3.3. The code places BatchNorm before E-MHSA nextvit.py and before the MLP nextvit.py.
• MHCA contains BatchNorm and ReLU. Eq. 2 and eq. 3 do not show them. The code places BatchNorm and ReLU between the grouped convolution and the 1×1 projection nextvit.py. The kernel size is 3×3 nextvit.py; the text does not give it.
• E-MHSA pooling is one-dimensional. The paper pools with stride s §3.3. The code pools groups of r² consecutive tokens in raster order nextvit.py nextvit.py. It is not an r-by-r spatial block. Divisibility is not checked; the 224 px and 384 px inputs divide evenly at every stage, but other input sizes can drop trailing tokens. This applies to the detection and segmentation backbones too, which share the same class.
• Head. The network ends with BatchNorm, global average pooling and a linear layer nextvit.py. The paper does not describe the head.
• Learning rate. The paper's 2e-3 is the value after scaling. The CLI value is multiplied by total batch / 512 main.py. The README command passes --lr 5e-4 with 2048 total images, which gives 2e-3 README.md.
• Weight decay and warm-up defaults differ from the paper. The defaults are 0.05 and 5 epochs main.py main.py; the paper and README use 0.1 and 20 epochs §4.1.1 README.md.
• Drop-path flag is ignored. --drop-path is parsed main.py but the model is built with only num_classes main.py. The drop-path rate comes from the model factory: 0.1 for S and 0.2 for B and L nextvit.py nextvit.py nextvit.py.
• Training log. The log for Next-ViT-S at 224 px shows the learning rate rising about 1e-4 per epoch from epoch 2 to 1.9e-3 at epoch 20, then cosine decay nextvit_small_in1k_224.log. This matches the 20-epoch warm-up in the README, not the code default.
• Detection and segmentation backbones are not the classifier. They remove the final BatchNorm, pooling and linear head, and return features from each stage nextvit.py nextvit.py. The paper does not describe this.
• Segmentation backbone references a removed layer. The segmentation configs set with_extra_norm=True fpn_512_nextvit_small_80k.py upernet_512_nextvit_small_160k.py. In that mode the forward pass calls self.norm for stage 4 nextvit.py, but that attribute is commented out in the constructor nextvit.py. As read, this path would raise AttributeError. I did not run it, so treat this as a reading of the code, not a confirmed failure.
• Drop-path differs between detection and segmentation. The detection backbone uses 0.1 for S nextvit.py; the segmentation backbone uses 0.2 nextvit.py. The paper does not state drop-path for detection.
• 3x detection augmentation. The 3x config uses an AutoAugment policy with a resize to 400 to 600, a random crop to 384 × 600 and a second resize mask_rcnn_nextvit_small_3x.py. The paper describes only the multi-scale range 480 to 800 §4.3.1.
• FPN 80k runs 40k iterations. The config sets max_iters=80000//gpu_multiples with gpu_multiples=2 fpn_512_nextvit_small_80k.py. The paper's text says 40k §4.2.1, and its Table 5 caption calls the 80k run the doubled schedule tab. 5.
• Table 4 has an implausible cell. ConvNeXt-T's CoreML entry is printed as 83.8 ms, a value that sits in the column next to its top-1 of 82.1 tab. 4. It looks like a misplaced number. I did not check it against the ConvNeXt paper, so treat it as unverified and do not use it.

Pitfalls

• Scaled learning rate. The paper's 2e-3 is post-scaling. Passing 2e-3 to the CLI gives an effective rate of 8e-3 at 2048 images. Passing the code default 5e-4 at batch 64 on one GPU gives about 6e-5, which is far too small for this recipe.
• Weight decay and warm-up are not in the code defaults. Running main.py with defaults trains with 0.05 weight decay and a 5-epoch warm-up. Use the README command to get the paper's settings.
• Unstated training recipe. Mixup, CutMix, RandAugment, random erasing, colour jitter, label smoothing and gradient clipping are set by code defaults and not stated in the paper main.py. The paper says it follows the training settings of several recent Transformers (PVT, Twins and others) "with minor changes" §4.1.1.
• Stochastic depth is linear across blocks. torch.linspace gives each block a rate from 0 up to the maximum nextvit.py. The paper says "increasing" and does not give the rule.
• E-MHSA pooling sizes. The pooling kernel r² must divide H × W at each stage. For 224 px and 384 px inputs, the token counts divide evenly: 3136/64, 784/16, 196/4 and 9216/64, 2304/16, 576/4. For other sizes, check the divisibility before using the model.
• E-MHSA normalisation. BatchNorm is applied to pooled tokens, not to queries nextvit.py nextvit.py. This follows from the code; the paper does not show it.
• Dense-task latency. Detection and segmentation latencies are backbone only at 512 × 512 §4.2.1. Do not compare them with the full-model numbers in the paper's text. On those tables, Next-ViT-S is slower than ResNet101 on both runtimes: 38.2 ms against 32.8 ms on TensorRT and 18.1 ms against 13.2 ms on CoreML tab. 5 tab. 6. The accuracy gains on dense tasks come with that extra cost. The classification advantage at similar latency §4.1.2 does not carry over to these tables.
• Latency is hardware-specific. The paper measures only a T4 with TensorRT 8.0.3 and an iPhone 12 Pro Max with CoreML. The rankings may change on other devices, as the paper itself notes for other models.
• Table numbers are derived. Table, figure and equation numbers follow the order of the LaTeX source, which is the version fetched. The compiled PDF was not read. If a number looks off, check the section.
• No appendix in the source. The LaTeX source contains no appendix, so all hyperparameters are in the main text, tables, or the code.
• Speed claims. The abstract's 3.6× speed-up over CSWin does not match the rows I checked. For example, Next-ViT-L's 65.3 ms against CSWin-S's 298.2 ms in the ADE20K table is about 4.6× tab. 5. Check the tables before quoting a speed-up factor.
• Config warnings. The detection and segmentation configs use resume= to load a checkpoint from ../checkpoints/, which the repository does not include mask_rcnn_nextvit_small_1x.py.

Lineage

• Builds on: MetaFormer for the block layout (NCB follows its token-mixer plus MLP design) §3.2; ConvNeXt, whose block the paper ablates against NCB §3.2 tab. 7; Linear SRA from PVTv2 for the spatially reduced attention in E-MHSA §3.3; the Swin and CSWin attention work §3.2; BottleNeck from ResNet as a reference for deployment-friendly convolution §3.2.
• Compared against: ResNet, ResNeXt, ConvNeXt, Swin, CSWin, Twins, PVTv2, UniFormer, EfficientFormer, CMT, TRT-ViT and others, as cited in its tables tab. 4 tab. 5 tab. 6.
• Superseded by: not known from content already in this wiki.