EfficientViT: Memory Efficient Vision Transformer with Cascaded Group Attention

In plain words

Vision Transformers are accurate image models, but they can be slow on real hardware even when their arithmetic count looks small. Much of the time goes to moving data around memory: reshaping tensors and running small element-wise operations inside attention. The authors argue that speed should be measured on the actual device rather than estimated from operation counts.

EfficientViT changes the layout of each block. A block contains a single attention layer, placed between cheap feed-forward layers, so most of the work is channel mixing that runs fast on hardware. The attention is split across heads so that each head sees a different slice of the features, and each head builds on the one before it. That removes duplicate work between heads and lets the network go deeper at little extra cost. The width of each part is also rebalanced: the value projection gets more channels and the feed-forward layers get fewer.

The result is a family of small models (M0 to M5) that the authors report as fast on a GPU, a CPU and in ONNX format, with accuracy that is competitive with efficient convolutional networks on ImageNet-1K.

TL;DR

• EfficientViT is a family of hierarchical ViT-style classifiers (M0 to M5) designed to maximise measured throughput §1 tab. 2.
• Each block is a sandwich: one self-attention layer between feed-forward layers, with depthwise convolutions before each FFN §3.1 eq. 1.
• Cascaded group attention (CGA) splits the channels across heads and adds each head's output to the next head's input §3.1 eq. 2 eq. 3.
• Author-reported: M5 reaches 77.1% ImageNet-1K top-1 at 10,621 images/s on an Nvidia V100, against 75.2% and 7,560 images/s for MobileNetV3-Large tab. 2.
• Read Pitfalls before implementing. The learning rate, the window partition in stage 1, and the attention bias are not what the paper text alone implies, and several speed ratios quoted in the text do not reproduce from Table 2.

Problem

Efficient CNNs and ViTs are usually compared by FLOPs and parameters, and these correlate weakly with wall-clock speed. For example, MobileViT-XS uses about 700M FLOPs yet runs much slower than DeiT-T with 1,220M FLOPs on a V100 §1. Methods that cut attention cost with sparse attention or low-rank approximation often lose accuracy and give little speed-up §2.1.

The authors profiled Swin-T and DeiT-T and found that memory-bound operations (reshaping, element-wise addition, copying, normalisation) take a large share of runtime. In a Swin-T variant accelerated 1.25× with 20% MHSA layers, these operations account for 44.26% of total runtime §2.1. Subnetworks with 20 to 40% MHSA layers reached better accuracy than the usual 50% §2.1 fig. 3.

Assumed setting: ImageNet-1K training from scratch, 300 epochs, 224×224 inputs, 8 V100 GPUs §4.1.

Method

The design rests on three observations from Section 2 §2:

1. Use fewer MHSA layers and more FFN layers, because MHSA is memory-bound §2.1 fig. 2 fig. 3.
2. Feed each attention head a different channel split, because heads in small ViTs learn similar projections §2.2 fig. 4.
3. Reallocate parameters. Taylor structured pruning §2.3 on Swin-T and DeiT-T kept most of the value channels, cut the Q and K channels and the FFN hidden width. The FFN expansion ratio drops from 4 to 2 §3.1 fig. 5.

Sandwich block. A block applies one self-attention layer for spatial mixing, with FFN layers before and after it. The paper writes it with N FFNs on each side; the final models use N = 1 eq. 1 §4.4 tab. 6.

X_{i+1} = \Phi^{\rm F}_i\big(\Phi^{\rm A}_i\big(\Phi^{\rm F}_i(X_i)\big)\big) eq. 1

where X_i is the input of block i, \Phi^{\rm A}_i is the attention layer and \Phi^{\rm F}_i the FFN. Each FFN has hidden width 2C §3.1. A 3×3 depthwise convolution comes before each FFN to add local structure §3.1.

# EfficientViT block, sandwich layout  [§3.1] [eq. 1]
# Input x: (B, C, H, W); output: (B, C, H, W)

def efficientvit_block(x):
    # Depthwise token mixing before the first FFN  [§3.1]
    # Residual adds around each sub-layer: see Paper vs code
    x = x + depthwise_conv_0(x)
    # First FFN: hidden width 2C, ReLU in between  [§3.1] [§3.2]
    x = x + ffn_0(x)
    # The only self-attention layer in the block  [§3.1] [eq. 2]
    x = x + cascaded_group_attention(x)
    # Depthwise token mixing before the second FFN  [§3.1]
    x = x + depthwise_conv_1(x)
    # Second FFN, same shape as the first  [§3.1] [eq. 1]
    x = x + ffn_1(x)
    return x

Cascaded group attention. Each head gets its own split of the input channels, X_{ij} for head j of h heads. Head outputs are concatenated and projected back to the input width eq. 2 §3.1. The cascade adds the previous head's output to the current split, and the first head uses its split alone eq. 3 §3.1. Because each head's value width is C/h, the outputs can be added to the next split efficientvit.py. A depthwise convolution is applied to the queries after the Q projection §3.1.

\widetilde{X}_{ij} = \mathrm{Attn}\big(X_{ij}W^{\rm Q}_{ij},\ X_{ij}W^{\rm K}_{ij},\ X_{ij}W^{\rm V}_{ij}\big), \qquad \widetilde{X}_{i+1} = \mathrm{Concat}_{j=1..h}\big[\widetilde{X}_{ij}\big]\,W^{\rm P}_i eq. 2

X'_{ij} = X_{ij} + \widetilde{X}_{i(j-1)}, \qquad 1 < j \le h eq. 3

The paper's cascade equation is unlabelled in the LaTeX source; it is numbered eq. 3 here by counting display equations in order.

# Cascaded group attention  [§3.1] [eq. 2] [eq. 3]
# Input x: (B, C, H, W); h heads; per-head value width d = C / h
# Output: (B, C, H, W)

def cascaded_group_attention(x):
    # Split the channels into h groups, one per head  [§3.1] [eq. 2]
    splits = chunk(x, h, dim=1)                # each (B, C/h, H, W)
    head_outputs = []
    for j in range(h):
        # Cascade: add the previous head's output to split j  [§3.1] [eq. 3]
        feat = splits[j]
        if j > 0:
            feat = feat + head_outputs[j - 1]
        # Queries, keys and values for this head  [§3.1] [eq. 2]
        queries, keys, values = qkv_proj_j(feat)
        # Depthwise token interaction on the queries  [§3.1]
        queries = depthwise_conv_j(queries)
        # Scaled dot-product scores plus a learned relative bias.
        # The scale and bias are not in the paper, see Paper vs code.
        scores = scale * dot(queries, keys) + bias_j
        weights = softmax(scores, dim=-1)
        head_outputs.append(matmul(values, transpose(weights)))
    # Concatenate heads, apply ReLU, project back to C channels
    # [§3.1] [eq. 2]; the ReLU is not in eq. 2, see Paper vs code
    return proj(relu(concat(head_outputs, dim=1)))

Network. Three stages, each made of EfficientViT blocks. Between stages a subsample block halves the resolution §3.2. Its self-attention is replaced by an inverted residual block, which limits information loss during downsampling §3.2. The 16×16 patch embedding is an overlapping patch embedding §3.2; the code builds it as four stride-2 3×3 convolutions with BatchNorm and ReLU efficientvit.py. The head uses BatchNorm, not LayerNorm, and ReLU, not GELU or HardSwish, for speed on deployment platforms §3.2 tab. 6 #10, #11.

# EfficientViT forward pass  [§3.2] [tab. 1] [fig. 6]
# Input images: (B, 3, 224, 224); output: (B, num_classes)

def efficientvit(images):
    # Overlapping patch embedding to C1 channels, 16x downsampling  [§3.2]
    x = patch_embed(images)                    # (B, C1, 14, 14)
    # Stage 1: L1 blocks; attention runs in 7x7 windows  [§3.2] [code]
    x = stage_1(x)
    # Subsample block: 2x down, channels C1 -> C2  [§3.2]
    x = subsample_block(x)                     # (B, C2, 7, 7)
    # Stage 2: L2 blocks, global attention over 7x7  [§3.2] [code]
    x = stage_2(x)
    # Subsample block: 2x down, channels C2 -> C3  [§3.2]
    x = subsample_block(x)                     # (B, C3, 4, 4)
    # Stage 3: L3 blocks  [§3.2] [tab. 1]
    x = stage_3(x)
    # Global average pool, BatchNorm, linear classifier  [code]
    return classifier(global_avg_pool(x))

Stage resolutions (14, 7, 4) and window behaviour are from the code, not the paper text efficientvit.py.

Hyperparameters

Architecture variants tab. 1. Widths {C1, C2, C3}, depths {L1, L2, L3} and heads {H1, H2, H3} per stage. The code builds the same values build.py.

Fixed across variants (code: EfficientViT/classification/model/efficientvit.py and build.py; paper: §3.1, §3.2).

Training (image classification; §4.1 unless marked as code-only).

Downstream detection (RetinaNet 1×, COCO; §4.3).

Downstream classification fine-tuning. Paper: 300 epochs, AdamW, batch 256, lr 1e-3, weight decay 1e-8 §4.3. No matching code is in the repository; the downstream/ folder covers detection only.

Throughput protocol. Paper: GPU runs use the largest power-of-two batch that fits in memory; CPU runs use one thread and batch 16 §4.1. The repository measures GPU throughput at a fixed batch of 2048 in fp16 autocast on a traced model with BatchNorm fused speed_test.py speed_test.py speed_test.py. CPU runs use batch 16 and one thread speed_test.py speed_test.py. There is no ONNX or CoreML benchmark script; ONNX appears only in the README links.

Results

Author-reported. ImageNet-1K, trained from scratch for 300 epochs at 224×224 §4.1. Throughput: V100 GPU, Intel Xeon E5-2690 v4 CPU at 2.60 GHz, ONNX on the same CPU tab. 2.

Headline (EfficientViT-M5 against MobileNetV3-Large) tab. 2:

Full ImageNet-1K table tab. 2: Top-1 and Top-5 in %; throughput in img/s on GPU, CPU and ONNX; FLOPs in M; Params in M; Input in px; Epochs. ↑ denotes fine-tuning at higher resolution. The "M4↑384" label is Table 2's; the text calls this row M5↑384 (see Pitfalls).

Tiny variants of larger ViTs tab. 3. Top-1 %, throughput in img/s (GPU, CPU, ONNX), FLOPs (G), Params (M).

Transfer to downstream classification tab. 4. Top-1 %, throughput in img/s on GPU.

COCO object detection, RetinaNet 1× tab. 5. AP metrics in %; FLOPs in M; Params in M.

Ablations on ImageNet-1K, 100 epochs, EfficientViT-M4 baseline tab. 6. Top-1 %, throughput in img/s (GPU, ONNX).

Comparison with 1000-epoch distillation runs tab. 7. Top-1 % on ImageNet-1K and ImageNet-ReaL; † = 1000 epochs with distillation from RegNetY-16GF §4.4.

Text-level results §4.2 §4.3 §4.4:

• M5 beats MobileNetV3-Large by 1.9 points top-1 and is faster on GPU and CPU, but slower in ONNX §4.2 tab. 2.
• M4 surpasses MobileNetV2 by 4.4 AP on COCO and uses 18.1% fewer FLOPs than SPOS with 2.0 AP more §4.3 tab. 5.
• Replacing the sandwich block with a Swin block costs 3.0 points top-1 at similar speed; CGA gives about 1.1 points over MHSA §4.4 tab. 6.
• Going from 1 to 2 or 3 FFNs does not help; N = 1 is the best trade-off §4.4 tab. 6.
• Higher resolution helps: M5↑512 reaches 80.8% top-1 §4.2 tab. 2.

Paper vs code

Official code: microsoft/Cream, directory EfficientViT/, commit 4a13c4091e78f9abd2160e7e01c02e48c1cf8fb9, MIT [paper link in abstract]. The cloned tree is at raw/code/2305.07027/. Citations below use paths from its root.

Differences between paper and code:

1. Learning rate. The paper states an initial lr of 1e-3 at batch 2048 §4.1. The code's --lr default is 1e-3 and it is scaled by batch × GPUs / 512, which is 4e-3 for 8 × 256 main.py main.py. The paper does not say whether 1e-3 is before or after scaling.
2. Attention window. The paper's attention (eq. 2) runs over all tokens and gives no window size §3.1 eq. 2. The code partitions stage 1 (14×14) into 7×7 windows and runs stages 2 and 3 globally efficientvit.py efficientvit.py.
3. Learned attention bias and scale. The code adds a learned bias per head and per relative offset, and scales the dot product by key_dim^-0.5 efficientvit.py efficientvit.py efficientvit.py. Eq. 2 has neither eq. 2.
4. ReLU before the output projection. The code applies ReLU to the concatenated heads before the projection efficientvit.py efficientvit.py. Eq. 2 has no activation eq. 2.
5. Residual connections. Each sub-layer is wrapped in a residual add efficientvit.py efficientvit.py. Eq. 1 is written without them eq. 1.
6. Per-head kernel sizes. The depthwise kernels on the queries vary by head and by variant build.py build.py. The paper gives no kernel sizes §3.1.
7. Subsample block. The code uses a 4× expansion and squeeze-excite efficientvit.py. The paper describes an inverted residual block without these details §3.2.
8. Teacher model. The paper uses RegNetY-16GF §4.4 tab. 7. The code's default teacher name is regnety_160 main.py. The repository does not show the exact command used for the distilled runs (the README has none); only the argument defaults are verified.
9. EMA. The code creates an EMA copy by default main.py, but evaluation calls use the raw model main.py. The paper does not mention EMA.
10. Downstream classification. The paper describes a fine-tuning recipe §4.3; no code for it is in the repository.
11. Detection learning rate. The paper defers to Swin's settings §4.3; the code uses AdamW lr 1e-4, weight decay 0.05 retinanet_efficientvit_m4_fpn_1x_coco.py.

No discrepancy found in the architecture variant table: Table 1 matches the code widths, depths and heads tab. 1 build.py.

Pitfalls

• Learning rate depends on what "1e-3" means. Implement the code's rule (lr × total batch / 512) or the paper's literal 1e-3 at 2048, and say which in your report. The two differ by 4× §4.1 main.py.
• Speed-up ratios in the text do not reproduce from Table 2. The abstract and describe M2 as 5.8× / 3.7× faster than MobileViT-XXS on GPU / CPU and 7.4× in ONNX (abstract; §1). Table 2 gives 18218/4456 ≈ 4.1× (GPU), 121.2/29.4 ≈ 4.1× (CPU) and 158.7/41.7 ≈ 3.8× (ONNX). The paper does not say how the quoted figures were computed. Use the table values.
• Other text/table mismatches. §4.2 gives M4 vs EdgeViT-XXS as 4.4× (CPU) and 3.0× (GPU); Table 2 gives 15914/3638 ≈ 4.4× (GPU) and 88.5/28.2 ≈ 3.1× (CPU). §4.2 gives M2 vs MobileViTV2-0.5 as 3.4× (GPU) and 3.5× (CPU); Table 2 gives ≈ 3.5× and ≈ 3.5×. The abstract's 40.4% GPU advantage of M5 over MobileNetV3-Large differs from §4.2's 40.5%; Table 2 gives 40.5%.
• Row name. §4.2 calls the 384-pixel fine-tune "M5↑384"; Table 2 calls the same numbers (79.8%, 3,986 img/s) "M4↑384" §4.2 tab. 2.
• Throughput is not reproducible from the stated protocol alone. The paper states a "largest power-of-two batch" rule for GPU; the repository uses a fixed batch of 2048, fp16 autocast and a traced model with BatchNorm fused speed_test.py speed_test.py speed_test.py. Expect numbers to differ with other batch sizes, precisions and hardware.
• ONNX and CoreML numbers have no benchmark code here. The repository has only the PyTorch GPU/CPU script. The ONNX numbers in Table 2 are not reproducible from this repository.
• Percentages in §4.4 have no stated base. For example, the GPU throughput gain from QKV reallocation is quoted as 4.9%; from Table 6, (15914−15132)/15914 = 4.9% but (15914−15132)/15132 = 5.2%. Check the base before quoting.
• Ablations are 100-epoch runs with single results. The paper reports no seeds or variance §4.4. Differences of 0.1 to 0.5 points are within the noise you should expect, but the paper does not quantify that.
• HardSwish is more accurate but slower. HardSwish gains 0.9 points top-1 but costs about 20% ONNX speed §4.4 tab. 6 #11. The paper keeps ReLU.
• Model size. The authors list the larger model size (from the extra FFNs) as a limitation compared with MobileNetV3 §6.
• Cooldown epochs. The code adds a 10-epoch cooldown at minimum LR main.py. I did not check how timm counts these against the 300 epochs, so the total schedule length is unverified.
• Cascade equation numbering. The cascade equation has no label in the source. It is numbered eq. 3 here by counting.
• Q/K width. The code fixes the Q/K width at 16 per head for every variant efficientvit.py. The paper's ablation shows 16 as the best value of the swept range §4.4. Keep 16 unless you are re-running the sweep.
• Augmentation beyond the paper. The paper says "same as DeiT" §4.1. The code enables more by default: label smoothing, AGC clipping, repeated augmentation, warmup and EMA. Implement from the code if you want the reported numbers.

Lineage

• Builds on: Swin Transformer (the stage layout and the window idea are the starting point §2), and DeiT, which the paper treats as a baseline and a training-recipe source §2 §4.1. The paper also uses LeViT's throughput protocol and its distillation recipe §4.1 §4.4.
• Compared against: MobileNetV2, MobileNetV3 and EfficientNet-B0 (CNNs); MobileViT, MobileViTV2, EdgeViT, Mobile-Former and GLiT (efficient ViTs); PVTV2, T2T-ViT, PoolFormer, EfficientFormer and Swin-T (tiny ViTs) tab. 2 tab. 3.
• Superseded by: not known from wiki content.