EdgeNeXt: Efficiently Amalgamated CNN-Transformer Architecture for Mobile Vision Applications

In plain words

Phone and robot cameras need image models that are small and fast. Today the strongest small models are either convolutional networks, which look at neighbouring pixels well but miss the whole picture, or transformers, which see the whole picture but are slow on weak hardware. This paper builds a model that mixes both.

The model is a stack of convolution blocks, with a few special blocks near the end. Each special block splits its channels into groups. Some groups get small convolutions that look at slightly different neighbourhood sizes. The result is then passed through an attention step, but the attention compares channels to channels rather than every pixel to every pixel. That keeps its cost roughly proportional to the image size, which is what makes it usable on mobile devices.

The authors report that this gives better accuracy for a given speed and size than a comparable hybrid model, MobileViT, on image classification, and that the same backbone also helps object detection and segmentation.

TL;DR

EdgeNeXt is a small hybrid CNN-transformer for mobile vision. Its SDTA block combines multi-scale depth-wise convolutions with transposed (channel-wise) self-attention, which is linear in the number of pixels. The authors report 79.4% ImageNet-1K top-1 at 256px for the 5.6M-parameter model, versus 78.4% for MobileViT-S, with 35% fewer MAdds. The official code is MIT-licensed and matches the paper in most places, but it differs in several details that matter for reimplementation (see Paper vs code).

Problem

Self-attention over all pixels is quadratic in the number of tokens, so transformer-style vision models are slow on edge hardware §1. Pure CNNs are cheap but have a local receptive field and fixed weights at inference §1. Earlier hybrids such as MobileViT keep spatial attention and therefore still spend most of their MAdds on attention; the authors argue this is the main bottleneck for mobile inference §2. The goal is a model that is small in parameters and MAdds and still models global context §1, Contributions.

Method

The network is a four-stage hierarchy at strides 4, 8, 16 and 32 [§3 Overall Architecture]. Stage 1 uses a patchify stem and convolution encoders. Stages 2 to 4 start with a 2×2 strided downsampling convolution. The last block of each of stages 2 to 4 is an SDTA encoder; all other blocks are convolution encoders §3 Overall Architecture fig. 2 tab. 1.

1. Convolution encoder §3 Convolution Encoder eq. 1. A depth-wise k×k convolution mixes space, followed by layer norm, two pointwise (1×1) linear layers and a residual connection. The kernel size grows with depth: k = 3, 5, 7, 9 for stages 1 to 4 §3 Convolution Encoder tab. 1. Eq. 1 writes the second pointwise layer together with a GELU and places the activation after it. The residual adds the block input to that output.

2. SDTA encoder §3 SDTA Encoder eq. 2 eq. 3 eq. 4. Three steps.

• Multi-scale split. The C input channels are split into s subsets of C/s channels each. Subsets are processed in a chain: each depth-wise 3×3 conv gets its own subset plus the previous subset's output eq. 2. The number of subsets s depends on the stage (t ∈ {2, 3, 4}), but the text does not give the values §3 SDTA Encoder.
• Transposed attention. Query, key and value projections of the normalized tensor are computed with three linear layers. Q and K are L2-normalized. The attention matrix is formed across channels, giving a C×C softmax score, which is multiplied by V and added to the input eq. 3 eq. 4.
• Light MLP. Two 1×1 convolutions with layer norm and GELU follow the attention output §3 SDTA Encoder.

3. Positional encoding. Added once, before the SDTA block of stage 2 only; adding it in every stage increased latency and hurt dense prediction §3 Overall Architecture.

4. Head. Global average pooling and a linear classifier tab. 1.

# Convolution encoder  [§3 Convolution Encoder]  [eq. 1]
# Input: x (B, C, H, W); k = kernel size for this stage (3, 5, 7 or 9)
# Output: x (B, C, H, W)

def conv_encoder(x, k):
    # Depth-wise k x k conv for spatial mixing  [eq. 1: Dw]
    mixed = depthwise_conv(x, kernel_size=k)
    # Layer norm, then the first pointwise linear  [eq. 1: LN, Linear]
    hidden = linear_first(layer_norm(mixed))
    # Second pointwise linear followed by GELU  [eq. 1: Linear_G]
    out = gelu(linear_second(hidden))
    # Residual connection  [eq. 1]
    return x + out

# SDTA encoder  [§3 SDTA Encoder]  [eq. 2] [eq. 3] [eq. 4]
# Input: x (B, C, H, W); s = number of subsets for this stage (not stated)
# Output: x_hat (B, C, H, W)

def sdta_encoder(x, s):
    # Split channels into s subsets of C / s channels  [§3 SDTA Encoder]
    subsets = split_channels(x, s)
    # The first subset is kept as is  [eq. 2, i = 1]
    outputs = [subsets[0]]
    for i in range(2, s + 1):
        # Depth-wise 3x3 conv of this subset plus the previous output  [eq. 2]
        outputs.append(dwconv3x3(subsets[i - 1] + outputs[i - 2]))
    # Put the subsets back together along channels  [§3 SDTA Encoder]
    multi_scale = concat_channels(outputs)
    # Flatten to tokens (B, H*W, C); stage 2 adds a positional encoding
    # before this block  [§3 Overall Architecture]
    tokens = flatten(multi_scale)                  # (B, H*W, C)
    # Linear projections of the normalized tokens  [§3 SDTA Encoder]
    queries = linear_q(layer_norm(tokens))         # (B, H*W, C)
    keys = linear_k(layer_norm(tokens))            # (B, H*W, C)
    values = linear_v(layer_norm(tokens))          # (B, H*W, C)
    # L2-normalize queries and keys  [§3 SDTA Encoder]
    queries = l2_normalize(queries)
    keys = l2_normalize(keys)
    # Channel-by-channel score matrix, softmax over channels  [eq. 4]
    scores = softmax(transpose(queries) @ keys)    # (B, C, C)
    # Weight the values by the scores and add the input  [eq. 3] [eq. 4]
    attended = values @ scores                     # (B, H*W, C)
    x = x + unflatten(attended)
    # Two 1x1 convs with layer norm and GELU  [§3 SDTA Encoder]
    return x + mlp(layer_norm(x))

Overall network. Patchify stem (4×4 non-overlapping conv, then layer norm), then convolution encoders in stage 1 and downsampling plus blocks in stages 2 to 4 §3 Overall Architecture tab. 1. The parameter-size variants XXS, XS and S use the channel widths in tab. 1. The block counts in tab. 1 for XS and S match the official code; the XXS counts do not (see Paper vs code).

Variants and training data. Three sizes: XXS (1.3M parameters, 0.3G MAdds), XS (2.3M, 0.5G), S (5.6M, 1.3G) tab. 1. A larger model, EdgeNeXt-B (18.5M parameters, 3.8G MAdds), is trained on an ImageNet-21K subset §4.4 tab. 4; its layout is not in the paper.

Hyperparameters

Values for the ImageNet-1K classification runs unless marked otherwise. "Code" means the official repository at code_commit; code-only values are flagged because the paper does not state them.

Results

Author-reported. ImageNet-1K top-1 on the validation set, no retraining by us. Input size varies by row as published: Table 2 lists 224 and 256 px; the MobileViT comparison in Table 3 is at 256 px §4.3. MAdds and latency use the paper's own measurement.

ImageNet-1K classification, state of the art tab. 2:

Source: tab. 2. A knowledge-distilled S model reaching 81.1% is mentioned in the text §4.3 and in a commented-out row of tab. 2; it is not in the table.

Against MobileViT at matched size tab. 3, all 256 px:

Source: tab. 3 §4.5. The text says the XXS model is about 34% faster than MobileViT-XXS on A100 and about 8% faster on Jetson Nano §4.5; these figures match the table.

ImageNet-21K-P pretraining tab. 4:

Source: tab. 4 §4.4.

COCO object detection, SSDLite at 320×320 tab. 5:

Source: tab. 5 §4.6. The table compares against three MobileNets and MobileViT-S only.

Pascal VOC semantic segmentation, DeepLabv3 at 512×512 tab. 6:

Source: tab. 6 §4.7.

Ablations (all on EdgeNeXt-S, ImageNet-1K top-1, latency as published).

Components tab. 7 §5:

Hybrid placement of SDTA blocks tab. 8 §5:

Where the single SDTA block goes tab. 9 §5:

SDTA at the start versus end of each stage tab. 10 §5: start 79.0, end 79.4, both at 332 µs.

Activation and normalization: replacing GELU with Hard-Swish and LayerNorm with BatchNorm lowers accuracy slightly and cuts latency by a large margin; the numbers exist only in a figure §5 fig. 3 and are not transcribed here. The code's comments for the BN and Hard-Swish models give 70.33, 74.87 and 78.39 top-1 for XXS, XS and S model.py; these are the code's own claims and are not in the paper.

Paper vs code

Code: mmaaz60/EdgeNeXt at commit 388450404fdd6d6e097f1cf94f5afcbbdc352f08, MIT licence (Copyright 2022 Muhammad Maaz) LICENSE. The owner is the first author. The paper's short link t.ly/_Vu9 returned HTTP 403 when we tried to open it, so we could not confirm where it points; the repo was found by GitHub search, and its README names this paper and its arXiv id. Read: models/*.py, the argument list and training loop entry in main.py, the evaluation transform in datasets.py, and the README. Not read: engine.py, sampler.py, utils.py, optim_factory.py. The repo has no detection or segmentation code.

• XXS block counts. Table 1 gives XXS three, three, nine and three blocks per stage tab. 1. The code's XXS config uses two, two, six and two model.py. XS and S match the table [code: models/model.py#L39-L46] model.py.
• GELU placement in the convolution encoder. Eq. 1 applies GELU after the second pointwise layer and not after the first eq. 1. The code applies GELU between the two pointwise layers and none after the second conv_encoder.py.
• Channel split order in SDTA. Eq. 2 keeps the first subset as identity and applies the depth-wise convs to subsets 2 to s eq. 2. The code applies a conv to the first subset and passes the last subset through unchanged, accumulating the outputs in the opposite order sdta_encoder.py.
• Subset counts. The paper gives no values for s §3 SDTA Encoder. The code uses 2, 2, 3, 4 for stages 1–4 in S, XS, XXS and B alike model.py model.py model.py model.py. Stage 1 has no SDTA block in the code, so its value does not apply.
• Attention temperature and heads. Eq. 3 and eq. 4 have no temperature and no head split. The code has a learnable temperature per head, initialized to 1, and splits channels into heads (8 by default; 4 for XXS and XS) sdta_encoder.py. Q and K are L2-normalized sdta_encoder.py.
• Extra normalization. The code adds layer norm before each downsampling convolution and a final layer norm before the classifier edgenext.py edgenext.py. Neither is shown in Table 1 or in the text.
• Positional encoding type. Not stated in the paper. The code uses a 2D Fourier sine/cosine encoding layers.py, added to the XCA input in stage 2 sdta_encoder.py. This matches the stage placement in the paper.
• Augmentation defaults. The paper lists random resized crop, flip and RandAugment (S only) §4.2. The code's defaults also turn on colour jitter at 0.4 and label smoothing at 0.1, and RandAugment with magnitude 9 main.py. Neither colour jitter nor label smoothing is mentioned in the paper.
• EMA. The paper uses EMA with decay 0.9995 §4.2. The code disables EMA by default main.py; the README's training command enables it README.md. Evaluation without the USI flag loads the EMA weights main.py.
• Stochastic depth. The paper uses 0.1 for S only §4.2. The code's default is 0.1 for every variant main.py, and the variant comments say 0.0 for XXS and XS model.py [code: models/model.py#L36].
• Mixup and CutMix in the B model. Not in the paper; the code's B comment lists both as used model.py.
• EdgeNeXt-B architecture. The paper gives only its parameter count and MAdds §4.4 tab. 4. The code's B config uses depths 3, 3, 9, 3 and widths 80, 160, 288, 584 model.py.
• Distillation. The paper says the 81.1% S model was trained with knowledge distillation §4.3. The code README calls the same model a USI training recipe and provides weights README.md. No distillation training code was found in the repo; the teacher is not named.
• Eval preprocessing. Matches: resize to 292 and centre crop 256 §4.2 datasets.py.
• Batch size and optimizer. Matches: 300 epochs, lr 6e-3, weight decay 0.05, 20 warmup epochs, effective batch 4096 (256 × 2 × 8 GPUs) §4.2 main.py README.md.

Pitfalls

• Subset counts are missing from the text. The paper's SDTA equation leaves out the case i = 2 for stages 3 and 4 (eq. 2 covers i = 2 only for t = 2). The code applies the additive rule to every subset after the first. The caption of fig. 2 mentions branches B3 and B4 but does not define B. Use the code's values, and say so.
• Equation 3 is incomplete. It has no temperature, no head count, and no explicit L2 normalization, although the text mentions the normalization. A faithful reimplementation needs all three from the code.
• Table 1's XXS block counts do not match the code's XXS model. The code's comment for the XXS config reports 1.33M parameters and 71.23% top-1 model.py, and that config uses two, two, six, two blocks. The table's three, three, nine, three layout would not give the same model. The depths in the code are the ones tied to the reported XXS number.
• Naming. The abstract and §3 say "STDA"; the contribution list and most of the text say "SDTA". The card uses SDTA. Tables also contain typos ("MobileViT-XSS", "Model Prameters", "3.8MAdds").
• Latency figures are not directly comparable. Jetson figures are TensorRT FP16 at batch 1 §4.2 §4.5. A100 figures are PyTorch at batch 256 and are given in µs without saying whether that is per image or per batch §4.2 tab. 3. The A100 measurements may be sensitive to the software version (PyTorch 1.8.1 is stated §4.2).
• Some comparison claims are stronger than the tables. The text says EdgeNeXt has fewer MAdds and parameters than MobileFormer, ViT-C and CoaT-Lite-T §4.3. In tab. 2, EdgeNeXt-S has more MAdds than MobileFormer (965M or 1.30G against 214M) and more parameters than ViT-C (5.6M against 4.6M). The introduction says EdgeNeXt beats all lightweight models in literature on detection and segmentation with fewer MAdds §1; tab. 5 shows EdgeNeXt using more MAdds than all three MobileNets, and tab. 6 compares only MobileViT and two MobileNets.
• Table 5 ablation rows are ambiguous. "+ w/o adaptive kernels" and "+ w/o PE" are read here as cumulative (on top of the row above) because the reported numbers are consistent with that reading. The paper does not say so.
• Input resolution differs across comparison rows. Table 2 mixes 224 and 256 px inputs. The 224-px numbers are the ones the text uses for the ConvNet and ViT comparisons (for example 78.8 against 74.5 gives the stated 4.3% margin) §4.3.
• Unstated training details. Not stated: cosine schedule details beyond warmup, weight-decay exemptions, drop-path schedule (the code increases it linearly over blocks edgenext.py), layer-scale value, MLP expansion ratio, and the detection and segmentation training schedules. The paper refers to MobileViT for detection and segmentation settings, and to ConvNeXt for the 21K pretraining recipe §4.2 §4.4.
• Hard-Swish and BatchNorm variant. The paper reports only the qualitative trade-off §5 fig. 3. Its numbers are in the code's comments model.py and the README model zoo README.md.
• Attention cost hidden in the d² term. The paper states linear complexity in the number of tokens, O(N d²) §3 Overall Architecture. The C×C attention matrix still grows with the square of the channel count, so wider stages are not free.
• Eval uses EMA weights. A checkpoint evaluated with the default flags uses the EMA state dict main.py. Evaluating the non-EMA weights requires changing the code.

Lineage

• Builds on: ConvNeXt (block design, training recipe) §3 Overall Architecture §4.4; Res2Net (multi-scale splitting) §3 SDTA Encoder; cross-covariance attention from XCiT (transposed attention) §3 SDTA Encoder.
• Compared against: MobileViT (main hybrid comparison) tab. 2 tab. 3 tab. 5 tab. 6; MobileNetV1, V2 and V3; ShuffleNetV2; ViT, DeiT, T2T-ViT, MobileFormer, ViT-C, CoaT-Lite-T tab. 2.
• Superseded by: not known from content in this wiki.