Rethinking Vision Transformers for MobileNet Size and Speed

In plain words

Phone apps need image models that are small enough to store and fast enough to run in real time. Small convolutional networks such as MobileNetV2 have been the usual choice. Vision transformers do well on many vision tasks, but on phones they tend to be slower and to have more parameters. The paper asks whether a transformer can be as small and as fast as MobileNet.

The authors start from EfficientFormer, an already fast transformer, and change its design. They put a small convolution inside the feed-forward blocks, keep the expensive attention only in the last two stages, and use a cheap form of attention on a higher-resolution stage. Then they search the network's depth, width and feed-forward sizes while scoring each candidate on both parameter count and phone latency, not on latency alone. The resulting EfficientFormerV2 models are reported to match MobileNet-sized models in size and latency while being more accurate.

TL;DR

EfficientFormerV2 is a family of four hybrid vision transformers (S0, S1, S2, L) with a four-stage layout. The first two stages use only FFN blocks; the last two add self-attention. Starting from EfficientFormer-L1, the authors add a depthwise 3x3 conv inside the feed-forward blocks, a "stride attention" for the attention on the 1/16-resolution stage, and dual-path attention downsampling. A joint search over depth, width and per-block expansion ratio maximizes a Mobile Efficiency Score (MES) that combines parameters and iPhone 12 latency. Author-reported: S0 reaches 75.7% top-1 at 3.5M parameters and 0.9 ms (300 epochs); S1 reaches 79.0% at 6.1M and 1.1 ms; S0 is 3.5 points above a MobileNetV2 retrained with the same recipe.

Problem

Vision transformers have succeeded on most computer vision benchmarks §1. Their attention cost grows quadratically with the number of tokens, so they are usually slower than lightweight CNNs on mobile hardware §1. Window attention and hybrid CNN-transformer designs reduce this cost, but the authors report that they are still slower or larger than lightweight CNNs. MobileViT is at least 5x slower than MobileNetV2 on iPhone 12 §1. EfficientFormer-L1 is about as fast as MobileNetV2x1.4 but about 2x larger §1.

Earlier search methods optimize one metric. Latency-only search leaves parameters redundant, and size-only search gives slow models §4 fig. 1. Size here means parameter count: the paper defines model size as the number of parameters §4.2.1. Latency is measured on an iPhone 12 neural engine with CoreML at batch size 1 §5.1.

Method

The method has two parts. §3 revisits EfficientFormer-L1 and changes its blocks one at a time, reporting each change on ImageNet-1K tab. 1. §4 defines the EfficientFormerV2 supernet and a search that picks sub-networks under a joint size-speed score.

3.1 Unified FFN. EfficientFormer uses a 3x3 average-pooling token mixer with a residual connection. The authors replace it with a 3x3 depthwise conv and then move that conv into the feed-forward network (FFN). Every block becomes one of two types: a local FFN or a global attention block §3.1 fig. 2(a,b). The pooling-to-depthwise swap gives 79.8% (+0.6). The unified FFN gives 80.3% at +0.1M parameters and the same latency tab. 1.

3.2 Depth and width. A deeper, narrower network gives 80.5% with 12.24M parameters and 1.3 ms tab. 1 §3.2. The authors add a fifth stage at 1/64 resolution and find it parameter-heavy: the best five-stage variant reaches 80.3% at 12.63M parameters and 1.5 ms, against 80.5% at 12.24M and 1.3 ms for the four-stage model. They keep four stages §3.2 tab. 1.

3.3 MHSA improvements. Adding a depthwise 3x3 conv on the value projection, and talking-head mixing across heads, raises accuracy to 80.8% at similar parameters and latency §3.3 tab. 1.

3.4 Stride attention. Attention on the 1/16-resolution stage gives +0.9 points but is about 2.7x slower. The authors downsample queries, keys and values to a fixed 1/32 resolution, run attention there, and interpolate the output back to the stage's resolution. This gives 81.5% at 1.5 ms, against 81.7% at 3.5 ms without it §3.4 tab. 1. Window attention was set aside as hard to accelerate on mobile. Key/value downsampling from Next-ViT was measured at 2.8 ms, still 2x the baseline, so it was not used §3.4.

3.5 Dual-path attention downsampling. Attention is used to halve the number of tokens between stages. The query branch combines a pooling path (static local downsampling) and a learned 3x3 depthwise conv path, then projects them to the query width. The result is added to a strided conv, so the block mixes local and global information §3.5 fig. 2(f). Attention downsampling alone gives 81.4% at 1.6 ms; the dual-path version gives 81.8% at 1.6 ms tab. 1.

4.1 EfficientFormerV2 network. A stem of two 3x3 stride-2 convs embeds the image to 1/4 resolution. Four stages follow at 1/4, 1/8, 1/16 and 1/32. The first two stages contain only FFN blocks. Stages 3 and 4 contain FFN blocks followed by global MHSA blocks, and the number of MHSA blocks is set by the search §4.1 app. architecture table. Each block is residual and carries a learnable layer scale §4.1 eq. local FFN eq. global.

# FFN block with a local depthwise conv  [§3.1] [§4.1]
# x: (B, C, H, W); expansion: per-block ratio E from the config
def ffn_block(x, expansion):
    # 1x1 conv to C * expansion, then BN and activation
    # [code: models/efficientformer_v2.py#L356-L359]
    hidden = gelu(batch_norm(conv_1x1(x, C * expansion)))
    # depthwise 3x3 conv inside the FFN  [§3.1]
    # [code: models/efficientformer_v2.py#L362-L366]
    hidden = gelu(batch_norm(conv_depthwise_3x3(hidden)))
    # 1x1 conv back to C, then BN  [code: models/efficientformer_v2.py#L367-L368]
    out = batch_norm(conv_1x1(hidden, C))
    # learnable layer scale, residual add  [§4.1]
    # [code: models/efficientformer_v2.py#L411-L420]
    return x + layer_scale * out

# Global MHSA block (stages 3 and 4)  [§4.1] [§3.3]
# x: (B, C, H, W); heads and key/value widths per code, not stated in paper
def mhsa_block(x):
    # queries and keys: 8 heads x 32 channels
    # [code: models/efficientformer_v2.py#L66-L70]
    queries = batch_norm(conv_1x1(x, heads * key_dim))
    keys = batch_norm(conv_1x1(x, heads * key_dim))
    # values: 8 heads x 128 channels, i.e. 4x the key width
    # [code: models/efficientformer_v2.py#L89]
    values = batch_norm(conv_1x1(x, heads * value_dim))
    # scaled dot product plus a learned per-head position bias  [§4.1]
    logits = scale * (queries @ keys) + position_bias
    # talking heads: 1x1 mixing across heads, before and after softmax  [§3.3]
    weights = softmax(mix_heads(logits))
    weights = mix_heads(weights)
    out = weights @ values
    # depthwise 3x3 on the values, added as a local term  [§3.3]
    out = out + batch_norm(conv_depthwise_3x3(values))
    out = batch_norm(conv_1x1(gelu(out), C))
    return x + layer_scale * out

4.2 Joint size-speed search. The search space has three parts: depth (blocks per stage), width (channels per stage), and the expansion ratio of each FFN §4.2.2. Expansion ratios are per FFN, not shared within a stage, which the authors say is a finer-grained space than earlier ViT NAS §4.2.

Objective. The score is

\text{MES} = \text{Score} \cdot \prod_i \left(\frac{M_i}{U_i}\right)^{-\alpha_i} §4.2.1

with Score = 100, U_size = 3M parameters, U_latency = 1 ms on iPhone 12 (iOS 16, CoreML), alpha_latency = 1.0 and alpha_size = 0.5 §4.2.1. MobileNetV2 is the unit reference §4.2.1.

Supernet. A slimmable-style supernet with switchable layers. Depth is elastic through stochastic drop path. Each FFN's expansion ratio is set by the channels of its depthwise conv. Width is set by the output channels of each FFN's and MHSA block's last projection. Widths come from a predefined set of multiples of 16 or 32 §4.2.2. Training samples four sub-networks per iteration (smallest, two random, largest), the "sandwich rule" §4.2.2 alg. 1.

Search. Start from the largest sub-network. At each step, evaluate every frontier action: the accuracy drop on a validation split of ImageNet-1K (about 4 GPU-minutes), the parameter drop from layer shapes, and the latency drop from an iPhone 12 lookup table. Take the action with the smallest accuracy drop per unit of MES gain. Repeat until the target is met, then train the chosen architecture from scratch §4.2.3 alg. 1.

# Joint size-speed search  [§4.2.3] [alg. 1]
# supernet: trained with the sandwich rule  [§4.2.2]
# latency_table: iPhone 12 latency per FFN and MHSA setting  [§4.2.3]
def joint_search(supernet, latency_table, target_mes):
    # start from the largest sub-network  [§4.2.3]
    state = largest_subnet(supernet)
    while mes(state, latency_table) < target_mes:
        best_action = None
        best_ratio = None
        # every action on the frontier of the current state  [§4.2.3]
        for action in frontier_actions(state):
            # accuracy drop on an ImageNet-1K validation split  [§4.2.3]
            delta_acc = (accuracy(supernet, state)
                         - accuracy(supernet, apply(state, action)))
            # MES gain from parameters and the latency table  [§4.2.1]
            delta_mes = (mes(apply(state, action), latency_table)
                         - mes(state, latency_table))
            # pick the smallest accuracy loss per unit of MES gained  [§4.2.3]
            ratio = delta_acc / delta_mes
            if best_ratio is None or ratio < best_ratio:
                best_action = action
                best_ratio = ratio
        state = apply(state, best_action)
    # the searched architecture is then trained from scratch  [alg. 1]
    return state

The released EfficientFormerV2 models are fixed networks. The search itself is not in the repository (see Paper vs code).

Hyperparameters

Architecture (from the appendix table; each stage's MHSA count is its last blocks):

Per-block expansion ratios are in the appendix table and in code efficientformer_v2.py.

Results

Author-reported. ImageNet-1K top-1, models trained with the recipe above tab. 2. Latencies are iPhone 12 (NPU) unless stated. The MES column is as printed; for the rows checked it matches the formula in §4.2.1 with iPhone 12 latency and parameter count (e.g. S2: 100 x (12.6/3)^-0.5 x 1.6^-1 = 30.5).

The retrained MobileNet rows (marked with a dagger in the paper) use the same recipe as the authors' models. The 3.5-point headline is S0 at 300 epochs (75.7) against that retrained MobileNetV2 (72.2). The paper's own table lists the published 71.8 too, which gives 3.9.

Downstream tab. 3. Mask R-CNN on COCO 2017 and Semantic FPN on ADE20K, ImageNet-pretrained backbones. Backbone parameter counts as listed.

Search and ablation results §5.3 app.:

• Search vs random and vs EfficientFormer's search, at 3.5M parameters (iPhone 12): random 1 74.7% and random 2 75.0% at 1.0 ms, this search 75.7% at 0.9 ms. At 3.1M parameters and 0.9 ms: EfficientFormer's search 74.2%, this search 75.0% §5.3.
• At 7M parameters: EfficientFormer's search 79.2%, the same plus magnitude pruning for expansion ratios 79.2%, this search 79.4% app. search analysis.
• Expansion ratio at a fixed 13.4M parameters: ratio 4 gives 81.8%, ratio 2 gives 81.6%, ratio 1 gives 81.1% app. expansion analysis.
• Attention bias removed (S1): 79.0 to 78.8 at 300 epochs, 79.7 to 79.5 at 450 epochs app. attention-bias analysis.
• Without distillation (300 epochs): S0 73.7%, S1 77.9%, S2 80.4% app. no-distillation table.
• Score weights, at 3.5M parameters: (alpha_latency, alpha_size) = (1.0, 1.0) gives 1.1 ms and 77.0%; (0.5, 1.0) gives 1.3 ms and 77.3%; (1.0, 0.5) gives 0.9 ms and 75.7%. The authors use (1.0, 0.5) app. alpha analysis.

Paper vs code

Checked against snap-research/EfficientFormer at d01da42, the repository the paper links in its abstract. The repository contains the four EfficientFormerV2 networks, detection and segmentation configs, and a training script. It does not contain the search.

• Search. The paper's supernet, sandwich-rule training, MES search and latency lookup table §4.2 alg. 1 are not in the repository. The only V2 model code is the fixed networks efficientformer_v2.py. The README describes "supernet" and "joint search" only in prose README.md.
• Stride attention. The paper downsamples queries, keys and values to 1/32 and interpolates the output back §3.4. The code applies a learned depthwise 3x3 stride-2 conv (with BN) to the input feature map before the Q/K/V projections, and bilinearly upsamples the output by 2x efficientformer_v2.py efficientformer_v2.py efficientformer_v2.py. The code's stride-2 MHSA is used only in the third stage efficientformer_v2.py.
• Dual-path downsampling. The paper describes the pooling path as static local downsampling §3.5. In the code, the "pool" is AvgPool2d(1, 2, 0): a kernel-1, stride-2 subsample with no averaging. The local path is a depthwise 3x3 stride-2 conv with no normalization or activation. The two are summed, then projected by a 1x1 conv and BN efficientformer_v2.py.
• Heads and widths. The paper does not state head count or key width. The code uses 8 heads, key width 32 and value width 128 (attention ratio 4) in every MHSA block, regardless of stage width efficientformer_v2.py efficientformer_v2.py.
• Expansion-ratio order. The appendix gives counts in compressed form, for example "4(x6), 3(x6)" for S2 stage 3. The code gives an explicit per-block order, which interleaves the values, for example [4,4,3,3,3,3,3,3,4,4,4,4] efficientformer_v2.py.
• Distillation. The paper names the teacher RegNetY-16GF §5.1. The code's default teacher is regnety_160 with a DeiT checkpoint URL main.py. Whether these are the same checkpoint is not stated. Hard distillation is the default main.py.
• Inference output. The code averages the class-head and distillation-head logits at evaluation time efficientformer_v2.py. The paper does not describe this.
• Downstream learning rate. Detection and segmentation settings match the paper mask_rcnn_efficientformerv2_s2_fpn_1x_coco.py fpn_efficientformerv2_s2_ade20k_40k.py.
• Launch script. The only launch script, slurm_train.sh, names an EfficientFormer-L1-style job, uses batch 128 per GPU on 8 GPUs, and does not show the V2 recipe slurm_train.sh. Whether V2 used this is not stated.
• README. The README says "about 4% higher top-1 accuracy than MobileNetV2" README.md. The paper's headline is 3.5 points.

Pitfalls

• The search cannot be reproduced from this release. Implement the search from §4.2 and Algorithm 1 yourself. You will need a latency table for your target device, since the iPhone 12 table is not released. The paper gives the score and weights §4.2.1 but not the search budget.
• Stride attention is easy to misimplement. Follow the code if you want the released numbers: the learned strided depthwise conv comes before Q/K/V, and the output upsampling is bilinear efficientformer_v2.py efficientformer_v2.py.
• The "pool" in dual-path downsampling is a subsample. A 3x3 average pool as the text suggests would change both results and cost efficientformer_v2.py.
• Unstated but needed: normalization and activation. The code uses BatchNorm throughout and GELU in the FFN, the stem, and the MHSA projections efficientformer_v2.py efficientformer_v2.py efficientformer_v2.py. The paper states neither. The Pixel 6 latencies were measured with ReLU in every compared model §5.1. To reproduce those, swap GELU for ReLU.
• Head count and key width are unstated. Use 8 heads, key width 32, value width 128 efficientformer_v2.py. With S0 stage-3 width 96, the value projection expands to 1024 channels and projects back efficientformer_v2.py.
• Layer scale init is unstated. The code uses 1e-5 efficientformer_v2.py.
• Position bias is tied to the training grid. The bias index table is built from the construction resolution efficientformer_v2.py. The released code therefore fixes the MHSA token grid. The paper notes that attention bias is resolution-sensitive and hurts transfer to downstream tasks app. attention-bias analysis.
• Batch size is not given. The learning rate scales with total batch size, 1e-3 per 1024 §5.1. The paper states 16 A100 GPUs but not the per-GPU batch.
• Repeated augmentation is on by default in code main.py but not in the paper's training table. Turn it off to match the table.
• Distillation settings are partly unstated. The paper gives hard distillation with a RegNetY-16GF teacher §5.1. Alpha and tau are not in the paper; code defaults are 0.5 and 1.0 main.py.
• Epoch counts change the ranking. The 450-epoch numbers are better by 0.2 to 0.7 points across the four models tab. 2. Compare like with like.
• Latency units differ across tables. Detection-from-scratch latency (187.9 ms) names no device and sits far from the iPhone 12 numbers [app. design generalization table]. Do not compare it to the rest of the card.
• Equation numbers are unverified. The supplement refers to "Eqn. 1" for the score app. alpha analysis, but the source does not let me confirm the printed equation numbers. This card cites sections instead.
• Attention cost at high resolution. The paper reports 2.7x slowdown from attention at the 1/16 stage without stride attention §3.4. Budget for it.

Lineage

• Builds on: EfficientFormer (NeurIPS 2022), which the paper uses as the baseline; the supernet follows slimmable networks (arXiv:1812.08928) and the universally slimmable variant §4.2.2; talking-head attention (arXiv:2003.02436) §3.3.
• Compared against: MobileViT (arXiv:2110.02178), EdgeViT, UniNet, MobileNetV2, MobileFormer and PoolFormer tab. 2 tab. 3. Next-ViT (arXiv:2207.05501) is cited for key/value downsampling, which the paper rejects for mobile §3.4.
• Superseded by: not stated in this card.