EdgeViTs: Competing Light-weight CNNs on Mobile Devices with Vision Transformers

In plain words

Vision transformers recognise images well but are heavy to run on a phone. This paper builds a small one, EdgeViT, aimed at running about as fast as the best small convolutional networks. The main idea is to split the expensive attention step into parts. Nearby pixels are mixed with cheap convolutions. Attention runs across the whole image, but only over one representative point per small window. The result is then spread back to the neighbouring points.

The authors also measure real speed and energy on a phone instead of counting operations, which they argue is a poor proxy for speed. Their models are reported to sit on the best trade-off between accuracy and speed or energy among the models they compare.

TL;DR

EdgeViT is a family of light hierarchical vision transformers. Each layer uses a local-global-local block: depthwise and pointwise convolutions mix neighbours, self-attention runs over one sampled token per window, and a transposed convolution propagates that context back to every token §3.2. The authors report 74.4%, 77.5% and 81.0% ImageNet-1K top-1 for their three sizes, on a Galaxy S21 CPU at 32.8, 54.1 and 85.3 ms tab. 2. Code covers ImageNet classification only.

Problem

Self-attention costs grow quadratically with the number of tokens, so vision transformers are costly for mobile and edge devices §1. MobileViT, the earlier transformer designed for mobile devices, still lags behind efficient CNNs such as MobileNetV2 §1. The authors argue that FLOPs and parameter counts are indirect proxies for on-device latency and energy §1. They also argue that the parameter count is a weak selection criterion on modern phones, and that operations such as the cyclic shift in Swin are hard to deploy §1.

Method

Overview §3.1. The network is a four-stage pyramid. Spatial resolution falls and channel width grows from stage to stage, as in other hierarchical vision transformers §3.1 fig. 2.

Local-global-local (LGL) bottleneck §3.2. Each block replaces standard multi-head self-attention with three steps fig. 3:

1. Local aggregation. Depthwise and pointwise convolutions mix information within a k×k neighbourhood of every token §3.2.
2. Global sparse attention. One delegate token is sampled per r×r window, so the delegates are spread evenly across the feature map. Standard self-attention runs over the delegates only §3.2. The paper says this differs from PVT and Twins, where every token is a query §3.2.
3. Local propagation. A transposed convolution with kernel and stride equal to r spreads each delegate's context back to its window §3.2.

The paper states the block as four residual updates eq. 1. The block uses layer normalisation before each sub-block, a two-layer feed-forward network with GELU after aggregation and after propagation, and omits positional encoding from the equations eq. 1 §3.2.

# LGL block, one per layer  [§3.2] [eq. 1]
# Input and output: x (B, C, H, W); r = sample rate of the stage

def lgl_block(x_in, r):
    # 1. Local aggregation over k x k windows  [§3.2] [eq. 1]
    x = x_in + local_aggregation(norm(x_in))
    # 2. Feed-forward network, residual  [§3.3] [eq. 1]
    y = x + ffn(norm(x))
    # 3. Attention on one delegate per r x r window, then spread back  [§3.2]
    z = y + local_propagation(sparse_attention(norm(y), r), r)
    # 4. Second feed-forward network, residual  [§3.3] [eq. 1]
    x_out = z + ffn(norm(z))
    return x_out

Architecture §3.3. Stage 1 downsamples by 4 with a 4×4 convolution with stride 4. Stages 2 to 4 downsample by 2 with a 2×2 convolution with stride 2 §3.3. Positional information comes from a conditional positional encoding, implemented as a 3×3 depthwise convolution with zero padding and a residual connection, placed before local aggregation and before global attention §3.3. The feed-forward network is two linear layers with GELU §3.3. The sample rates are 4, 2, 2 and 1 for the four stages §3.3 app. A.2. The variants differ in channels, blocks and heads tab. 1. The paper points to its supplementary material for the exact block §3.3.

Complexity app. A.1. The paper counts the spatial mixing cost only. For local aggregation it gives O(k² h w c). For sparse global attention it gives O(h² w² c / r⁴), which assumes one delegate per r² tokens. For local propagation it gives O(r² h w c). The total is compared with O(h² w² c) for full self-attention. The paper sets k = 3 for this count app. A.1.

Hyperparameters

Values are from the paper unless the notes say otherwise. Code values are from the official repository at code_commit.

Results

Author-reported. Setup: ImageNet-1K classification with 224×224 inputs §4.1. MobileViTs are tested at 256×256 tab. 2. The baselines are limited to under 2 GFLOPs §4.1.

ImageNet-1K, latency tab. 2 (top-1 on the validation set; CPU latency on Galaxy S21):

Caption of tab. 2: * marks architectures scaled down by the authors; † is LeViT-384 retrained for 300 epochs without distillation.

ImageNet-1K, energy tab. 3 (energy from the Snapdragon HDK; latency column repeats the Galaxy S21 CPU latency). The efficiency column is top-1 divided by energy in mJ, see Pitfalls:

Bold in this copy marks the EdgeViT rows only; the source uses bold for its Pareto-optimal rows. The text says all three EdgeViT sizes are Pareto-optimal for accuracy against latency and against energy §1 app. A.3. The text says LeViT-384† lies between EdgeViT-S and EdgeViT-XS on efficiency §4.1.

Ablations on ImageNet-1K, base EdgeViT-XXS rescaled to about 0.5 GFLOPs §4.1.1 tab. 4 (CPU ms and top-1 %):

COCO detection and instance segmentation §4.2.1 tab. 5 (val2017; AP values in %; #Par. in millions; mAP@100 as stated in §4.2.1):

ADE20K semantic segmentation, Semantic FPN §4.2.2 tab. 6 (validation mIoU):

Dense-task efficiency app. A.4 tab. 7 tab. 8. Detection (input 800×800) and segmentation (input 512×512) CPU time and energy, with the same efficiency formula. Detection: PVTv2-B0 37.2 AP, 1.34 s, 3.50 J, efficiency 0.011; EdgeViT-XXS 38.7 AP, 0.59 s, 2.02 J, 0.019; EdgeViT-XS 40.6 AP, 0.90 s, 2.89 J, 0.014; EdgeViT-S 43.4 AP, 1.88 s, 4.36 J, 0.010. Segmentation: PVTv2-B0 37.2 mIoU, 0.35 s, 1.10 J, 0.034; EdgeViT-XXS 39.7 mIoU, 0.19 s, 0.71 J, 0.056; EdgeViT-XS 41.4 mIoU, 0.31 s, 1.11 J, 0.037; EdgeViT-S 45.9 mIoU, 0.52 s, 1.73 J, 0.027. The source tables give ResNet18 and PVTv1-Tiny rows too. The protocol for these tables is not stated.

Paper vs code

Compared with src/ at commit df3d070addc3dc761173429426eab5edeb0d6343 (Apache-2.0). The repository holds ImageNet classification only.

• Depth per stage. Paper: XS [1,1,2,2], S [1,2,3,2] tab. 1. Code: XS depth=[1, 1, 3, 1], S depth=[1, 2, 5, 3] edgevit.py edgevit.py. XXS matches: [1, 1, 3, 2] edgevit.py.
• Heads. Paper: XS and S use [1,2,4,8] tab. 1. Code computes heads as embed_dim / head_dim with head_dim 48 for XS and S edgevit.py edgevit.py edgevit.py. That gives [1,2,5,8] for XS and S (240/48 = 5).
• Local aggregation kernel. Paper: k = 3 in the complexity count app. A.1. Code uses a 5×5 depthwise convolution in the middle of the block edgevit.py.
• Local aggregation contents. The paper's equation shows only aggregation plus a feed-forward step eq. 1. The code's local block also has a 3×3 positional convolution, batch normalisation, two pointwise convolutions and a convolutional MLP with ratio 4 edgevit.py edgevit.py.
• Last stage. The paper applies the bottleneck in all four stages with rate 1 in the last one §3.3. The code adds the local aggregation block only when the sample rate is above 1, so stage 4 has global attention without local aggregation edgevit.py edgevit.py. With rate 1 the code also skips the sampler and propagation edgevit.py.
• Local propagation. Text says depthwise separable transposed convolution §3.2. Code has a depthwise transposed convolution only, with kernel and stride r edgevit.py. The code then applies a LayerNorm and an output linear projection edgevit.py edgevit.py. Neither appears in eq. 1.
• Sampler. The paper's default is "center" §4.1.1. The code uses AvgPool2d(1, stride=r), which keeps the first position of each r×r window rather than its middle edgevit.py. With kernel 1 the average is the identity, so this is stride subsampling.
• Attention bias. Code uses bias on the QKV projection by default edgevit.py. The paper does not state it.
• FFN width. Code uses ratio 4 edgevit.py. The paper does not state a ratio.
• Final normalisation. Code applies batch norm to the last feature map before pooling and the head edgevit.py edgevit.py. The paper does not mention it.
• Augmentation. Paper lists crop, flip, mixup, random erasing and label smoothing §4.1. Code defaults also include RandAugment rand-m9, colour jitter 0.4, CutMix with alpha 1.0 and mixup 0.8, and repeated augmentation through RASampler main.py main.py main.py main.py.
• Stochastic depth and smoothing. Code defaults stochastic depth to 0.1 and label smoothing to 0.1 main.py main.py. The paper states neither drop-path nor a smoothing value.
• EMA. Code enables a model EMA by default with decay 0.99996 main.py. The paper does not mention it. Validation in main.py uses the raw model, not the EMA main.py.
• Learning rate and batch. Code scales the lr linearly with batch size over 512, from a default of 5×10⁻⁴ main.py main.py. The paper's base lr 1×10⁻³ corresponds to a total batch of 1024. The README example uses --batch-size 256 on 8 GPUs, which is a total of 2048 if the flag is per process README.md. That gives lr 2×10⁻³ by the code's rule.
• Warm-up, minimum lr, cooldown. Code defaults: warm-up lr 1×10⁻⁶, minimum lr 1×10⁻⁵, cooldown 10 epochs main.py main.py main.py. The paper gives only the 5-epoch warm-up. I did not check how timm 0.3.2 applies the cooldown, so its effect on the schedule is unverified.
• Pretrained flag. The model builders accept pretrained=True but do not load weights edgevit.py. I did not read hubconf.py, which may do so.
• Detection and segmentation. No detection or segmentation code is in the repository. Those numbers have no code here §4.2.
• Latency and energy. No benchmarking code is in the repository. The protocol in §4.1 cannot be checked against code.

Pitfalls

• Figure 1 and latency. The Figure 1 caption text gives 33.3 ms for the "lightest" model against 52.1 ms for EfficientNet-B0 fig. 1. tab. 2 gives 32.8 ms for EdgeViT-XXS and 33.3 ms for MobileNet-v2. Read the figure text against the table.
• Detection gain in the text. The text gives +1.2 AP^m for EdgeViT-S over PVTv2-B1 on Mask R-CNN §4.2.1. tab. 5 gives 41.0 against 38.8, which is +2.2.
• Efficiency formula. The text calls the efficiency "the average gain in top-1 accuracy from each 1 mJ" §4.1 tab. 3. The values match top-1 divided by energy in mJ: 74.4 / 127.4 = 0.584 for EdgeViT-XXS. The dense tables print energy in J, and their efficiency is top-1 / (1000 × J) tab. 7. The metric is not an accuracy gain, and it favours less accurate models, as the paper notes §4.1. That is why the paper compares only Pareto-optimal points.
• Power column. Latency is measured on a Galaxy S21 and energy on a separate Snapdragon 888 HDK §4.1. Energy ÷ latency does not reproduce the Power column: EdgeViT-XXS gives 127.4 mJ / 32.8 ms = 3.88 W, not the listed 4.27 W. Do not derive energy from latency.
• Batch size. The paper's base lr is for batch 1024 §4.1. The code scales the lr from the per-process batch main.py. Use the total batch when reproducing.
• Sampler "center". The default sampler takes index 0 of each window, not its middle, see Paper vs code. The "center" name and the "avg" and "max" alternatives (kernel r) are not equivalent in position.
• Benchmark conditions. CPU, FP32, batch 1, and TorchScript lite §4.1. The authors chose this because it was the only setup that ran all their models; NPU and other hardware paths were not used §4.1. Results on other runtimes or hardware may rank models differently.
• Baselines. MobileViT is run at 256×256 while the others use 224×224 tab. 2. LeViT-384† is retrained for this paper, so it is not the released LeViT tab. 2. Only baselines under 2 GFLOPs are compared §4.1.
• Ablation rescaling. Ablation variants are rescaled to about 0.5 GFLOPs by changing layer counts, so the ablation models are not the released ones app. A.2.
• Dense tasks. Detection and segmentation are initialised from the ImageNet-1K weights of this paper's own training §4.2. Their schedules are given only as the values in the Hyperparameters table, so warm-up, decay and weight decay for them are not stated §4.2.1 §4.2.2.
• Missing details. Not stated: warm-up start lr, minimum lr, label smoothing, stochastic depth, EMA, RandAugment, and the dense-task schedule details. The supplementary block description is pointed to but not reproduced in the paper text §3.3.
• Normalisation after downsampling. The paper does not say which norm follows each downsampling convolution. The code applies LayerNorm after the patch embedding convolution edgevit.py edgevit.py.

Lineage

• Builds on: hierarchical pyramid ViTs such as PVT and PVTv2 and Twins §1 §3.1; conditional positional encodings §3.3; DeiT training recipe §4.1. Its local aggregation is compared with the locally-grouped self-attention of Twins and Swin §4.1.1.
• Compared against: MobileNetV2, MobileNetV3 and EfficientNet-B0 §4.1; PVT-v2, MobileViT, DeiT, LeViT (retrained), Uniformer, T2T-ViT, TNT and Twins-SVT tab. 2.
• Superseded by: not known from content in this wiki.