RepViT: Revisiting Mobile CNN From ViT Perspective

In plain words

Small image models for phones come in two families: convolutional networks (CNNs) and Vision Transformers (ViTs). Recently, small ViTs have matched or beaten small CNNs on accuracy and speed on phones. The authors ask whether the design choices behind those ViTs can be brought back into a plain CNN.

They start from MobileNetV3, a well-known small CNN, and change it step by step. Each change borrows a layout choice from the small ViTs: the spatial mixing and the per-pixel mixing become separate steps, and the network gets a different stem, a deeper downsampling step and a simpler classifier. The result, RepViT, is a family of pure CNNs that the authors report beat small ViTs on latency and accuracy on a phone.

The main point for an engineer: how many operations a model has does not predict its speed on a phone well, so the paper measures latency on the device.

TL;DR

• RepViT is a family of pure CNNs built from a MobileNetV3-L starting point by adopting lightweight-ViT layout choices §3.
• The core block separates the spatial token mixer (a re-parameterised depthwise conv) from the channel mixer (an FFN with expansion 2) §3.2.
• Author-reported ImageNet-1K top-1 with distillation: 78.7% (M0.9, 0.9 ms) to 83.7% (M2.3, 2.3 ms) on iPhone 12 with Core ML tab. 1.
• Dense prediction and SAM results are also reported, with latencies on phone and laptop hardware §4.2 §4.3.
• Several architecture details exist only in the supplementary material (not part of our source); the released code fills most of them in (see Paper vs code).

Problem

Lightweight ViTs beat lightweight CNNs on accuracy and latency on mobile devices, and the usual explanation is global self-attention §1 §2. The architectural differences between lightweight ViTs and lightweight CNNs, in the block structure and in the macro and micro design, were not examined closely §1. FLOPs and model size correlate poorly with measured on-device latency §3.1, so the paper measures latency directly.

Assumptions: all models are trained on ImageNet-1K at 224x224 §4.1; latency is measured on an iPhone 12 (iOS 16) compiled with Core ML Tools at batch size 1 §3.1 §4.1.

Method

The paper starts from MobileNetV3-L and changes it in five groups, each cited to the section where it is introduced §3.2–§3.5:

1. Training recipe: align the baseline with the lightweight-ViT recipe (DeiT style, with distillation excluded for now) §3.1.
2. Block design: split token mixer and channel mixer, re-parameterise the depthwise conv, expansion ratio 2, widths doubling per stage §3.2.
3. Macro design: two-conv stem, deeper downsampling, a simple classifier, stage ratio change §3.3.
4. Micro design: 3x3 kernels only, cross-block SE placement §3.4.
5. Network family: RepViT-M0.9, M1.0, M1.1, M1.5, M2.3, where the suffix is the latency in ms on iPhone 12 §3.5.

The paper gives no numbered equations and no numbered algorithms. The pseudocode below is our reading of the layout in the text, checked against the released code repvit.py.

# RepViT block, stride 1  [§3.2]
# Input x: (B, C, H, W); channel mixer expansion 2  [§3.2]
def repvit_block_stride1(x, use_se):
    # Token mixer: depthwise 3x3 and depthwise 1x1 branches plus the
    # identity, summed then batch-normalised  [§3.2]
    # [code: model/repvit.py#L83-L92]
    mixed = batch_norm(dw_3x3(x) + dw_1x1(x) + x)
    # Optional squeeze-and-excitation, placed after the token mixer  [§3.4]
    # [code: model/repvit.py#L149]
    if use_se:
        mixed = squeeze_excite(mixed)
    # Channel mixer: pointwise expand to 2C, GELU, pointwise back to C,
    # residual add  [§3.2]  [code: model/repvit.py#L151-L157]
    hidden = gelu(conv1x1_bn(mixed, out_channels=2 * C))
    return mixed + conv1x1_bn(hidden, out_channels=C, init_gamma=0)

# RepViT block, stride 2 (downsampling)  [§3.3]
# Depthwise stride-2 conv changes resolution, then a 1x1 conv changes channels
def repvit_block_stride2(x, out_channels, use_se):
    # Depthwise 3x3, stride 2, with BN  [§3.3]
    # [code: model/repvit.py#L133-L136]
    x = batch_norm(dw_3x3_stride2(x))
    # Optional squeeze-and-excitation  [code: model/repvit.py#L135]
    if use_se:
        x = squeeze_excite(x)
    # Pointwise 1x1 conv sets the new channel count  [§3.3]
    x = conv1x1_bn(x, out_channels=out_channels)
    # FFN with expansion 2 after the 1x1, residual  [§3.3]
    # [code: model/repvit.py#L138-L143]
    hidden = gelu(conv1x1_bn(x, out_channels=2 * out_channels))
    return x + conv1x1_bn(hidden, out_channels=out_channels, init_gamma=0)

# RepViT network  [§3.3] [§3.5]  [code: model/repvit.py#L218-L243]
# image: (B, 3, 224, 224); configs: one row per block (kernel, expansion,
# channels, SE flag, activation flag, stride)
def repvit(image, configs, channels):
    # Stem: two stride-2 3x3 convs with GELU between  [§3.3]
    # first conv has channels/2 filters (24 for M0.9)  [§3.3]
    x = conv3x3_bn_stride2(image, channels // 2)
    x = gelu(x)
    x = conv3x3_bn_stride2(x, channels)
    # Stacked blocks, stride-1 blocks and stride-2 blocks  [§3.2] [§3.3]
    for block_config in configs:
        x = repvit_block(x, block_config)
    # Head: global average pool, then batch norm and linear  [§3.3]
    # (the paper says linear only; the code adds a BN first)
    # [code: model/repvit.py#L163-L186]
    pooled = global_average_pool(x)
    return bn_linear(pooled)

Re-parameterisation: during training the depthwise branches and the identity are separate; at inference they fold into one 3x3 depthwise conv §3.2 repvit.py. The paper says the skip connection cost is removed at inference §3.2.

Distillation (used for the headline numbers): the student gets two heads, one for the labels and one for the teacher. The loss mixes both with weight alpha losses.py. The paper states the teacher only §4.1.

L = (1 - \alpha)\, L_{\text{base}} + \alpha\, L_{\text{kd}} losses.py, with alpha = 0.5 and, for the hard variant, L_kd = cross-entropy against the teacher's argmax losses.py main.py. This formula is not in the paper.

Hyperparameters

Values marked "paper" are stated in the paper. Values marked "code" are the defaults of the released training script; the paper defers the full recipe to its supplementary material, which is not in our source. Code values were not confirmed against the paper.

Model configs (per block: kernel, expansion, channels, SE, activation, stride) are in repvit.py for M0.9, repvit.py for M1.0, repvit.py for M1.1, repvit.py for M1.5 and repvit.py for M2.3. The paper refers to its supplementary for these tables §3.5; the supplementary is not in our source.

Results

Author-reported. ImageNet-1K, 224 x 224, latency on iPhone 12 with Core ML §4.1. Table 1 reports the models trained with distillation; Table 2 reports them without §4.1.

Without distillation (300 epochs) tab. 2: RepViT-M0.9 77.4% at 0.9 ms, M1.0 78.6% at 1.0 ms, M1.1 79.4% at 1.1 ms, M1.5 81.2% at 1.5 ms, M2.3 82.5% at 2.3 ms. The paper reports, at similar latency, MobileOne-S1 75.9% (0.9 ms), FastViT-T8 75.6% (0.9 ms) and EfficientFormerV2-S0 73.7% (0.9 ms) tab. 2.

Step-by-step ablation on ImageNet-1K, each row adding one change to the previous one. Values are as stated in the text; the row for the classifier gives only the drop §3.1–§3.4 fig. 2:

Ablations on the RepViT block tab. 6 tab. 7:

Downstream, backbone latency at 512 x 512 on iPhone 12 tab. 5: RepViT-M1.1 (4.9 ms) gives 39.8 AP-box, 37.2 AP-mask and 40.6 mIoU, against EfficientFormer-L1 (5.4 ms) at 37.9, 35.4 and 38.9. RepViT-M2.3 (9.9 ms, with 450-epoch ImageNet init) gives 44.6, 40.8 and 46.1, against EfficientFormerV2-L (18.2 ms) at 44.7, 40.4 and 45.2 tab. 5.

RepViT-SAM (RepViT-M2.3 as the SAM image encoder) §4.2 tab. 3 tab. 4:

The paper calls RepViT-SAM "nearly 10x faster" than MobileSAM on MacBook M1 Pro §4.2; the table ratio is 482.2 / 44.8 = 10.8.

Paper vs code

Repository: the link is in the abstract, so no further search was needed. The owner is THU-MIG, and the README lists the paper's five authors. Code read at commit 298f42075eda5d2e6102559fad260c970769d34e. The main code paths are model/repvit.py, main.py, losses.py, logs/ and README.md. The detection, segmentation and SAM sub-projects were not checked against the paper in this card.

• Stage depth. Paper: stage ratio 1:1:7:1, depth 2:2:14:2 §3.3. Code: the M0.9 config has 3, 4, 16 and 3 blocks per stage, where each stage after the first starts with its stride-2 block repvit.py. How the paper counts depth is not stated, so this is unresolved.
• Head. Paper: global average pooling and a linear layer §3.3. Code: a BatchNorm1d before the linear layer repvit.py. For the distillation model, the two heads' logits are averaged at evaluation repvit.py.
• SE placement. Paper: SE in the 1st, 3rd, 5th... block of each stage §3.4. Code: the M1.5 first stage has SE on blocks 1 and 3 [code: model/repvit.py#L391-L395], which matches; the M0.9 first stage has SE on block 1 only repvit.py, which does not.
• Activation flag. The config's HS column selects between two identical nn.GELU() calls, so it has no effect repvit.py [code: model/repvit.py#L154]. The paper says GELU is used throughout §3.1.
• Distillation. Paper: teacher RegNetY-16GF §4.1. Code default: regnety_160 main.py. Code also uses hard-label distillation [code: main.py#L137-L140] losses.py; the paper does not say which.
• Learning rate. Paper: not stated. Code: 1e-3 scaled by batch x GPUs / 512 main.py. The shipped 300-epoch M0.9 log peaks at about 3.0e-3 repvit_m0_9_distill_300e.txt. Under the scaling rule that corresponds to 6 GPUs, not the 8 in the README README.md — the reason is not stated.
• Reported accuracy. Table 1 gives 78.7% for M0.9 at 300 epochs tab. 1. The shipped log's final evaluation (epoch 299) gives 78.58% [code: logs/repvit_m0_9_distill_300e.txt#L300].
• Parameter count. Table 1 gives 5.1M for M0.9 tab. 1. The shipped log reports n_parameters = 5,489,328 for the training model, which includes the distillation head repvit_m0_9_distill_300e.txt. The 5.1M figure is presumably the fused inference model; the paper does not say.
• Extra variant. The code defines RepViT-M0.6, which is not in this paper repvit.py. The README reports it separately README.md.

Pitfalls

• Missing supplementary. The architecture tables and full training details are in the supplementary material §3.5, which is not in our source. The configs and recipe here come from the code, so they carry the code's status, not the paper's.
• Latency is device- and toolchain-specific. Numbers come from Core ML on iPhone 12 with iOS 16 and Xcode's tool §3.1 README.md. Other devices, runtimes or Core ML versions may rank models differently.
• Throughput is a separate metric. Table 1 throughput uses an RTX 3090 with the largest power-of-two batch that fits tab. 1 caption. It is not latency.
• Distillation changes the headline. Table 1 numbers use distillation from a RegNet teacher. Table 2 gives the no-distillation numbers, which are 1 to 2 points lower for the same model tab. 2.
• Re-parameterisation is a deployment step. The training form has separate depthwise branches and BN; the inference form is a single depthwise conv. The latency numbers assume the inference form §3.2. The code has fuse() methods for it repvit.py repvit.py.
• Last FFN norm starts at zero. The last BN of each FFN is initialised to zero, so each residual branch starts as identity repvit.py repvit.py. The paper does not mention this.
• Ambiguous SE pattern. "1st, 3rd, 5th" is the paper's wording §3.4, but the released configs vary by variant and stage. Copy the config table from the code, not the rule in the text.
• Speed claims are relative. The "10x" SAM claim compares against MobileSAM on a laptop §4.2. On a phone, only RepViT-SAM ran at all, since the other two returned OOM tab. 3.
• Ablation noise. Re-parameterisation gains are 0.1 to 0.3 points on single runs tab. 6. No variance across seeds is reported.

Lineage

• Builds on: MobileNetV3-L as the starting point §3; the separated token and channel mixers from MetaFormer-style designs §3.2; structural re-parameterisation from RepVGG §3.2. None of these has a card in this wiki.
• Compared against: EfficientFormerV2 (in this wiki) and EfficientFormer, FastViT, MobileOne and MobileSAM, which have no cards tab. 1 tab. 3.
• Superseded by: none known from content in this wiki.