MobileNetV4: A Practitioner’s Deep-Dive

Why MACs are not enough

Many papers use MAC count as a latency proxy. That implicitly assumes memory is free. On accelerators, compute capability grows much faster than memory bandwidth, so many layers become memory-bound, and a model with fewer MACs can be slower.

The model

Following the classic Roofline model [@williams2009roofline], each layer \(i\) has a compute time and a memory time, and the slower one determines its latency, because compute and memory overlap roughly in parallel:

\[ \text{ModelTime}=\sum_i \max\big(\text{MACTime}_i,\ \text{MemTime}_i\big),\qquad \text{MACTime}_i=\frac{\text{LayerMACs}_i}{\text{PeakMACs}},\quad \text{MemTime}_i=\frac{\text{WeightBytes}_i+\text{ActivationBytes}_i}{\text{PeakMemBW}} \tag{1}\]

A workload’s operational intensity (OI) is \(\text{LayerMACs}/(\text{WeightBytes}+\text{ActivationBytes})\). A hardware target is summarized by one number, its Ridge Point:

\[ \text{RP}=\frac{\text{PeakMACs}}{\text{PeakMemBW}}\quad[\text{MACs/byte}] \]

A layer whose OI is below the RP is memory-bound on that hardware; above it, compute-bound.

The paper’s footnotes add two caveats: the model assumes software implementation does not matter (so tricks with complex memory access, like pruning, look better on paper than on device), and a similar framework was proposed independently by Lavin.

A toy example (Not from the paper)

Let both chips have 3.2 GB/s memory bandwidth, but chip A has 100 GMAC/s (RP ≈ 31, CPU-like) and chip B has 1.39 TMAC/s (RP ≈ 434, accelerator-like). Take two layers: a big conv (200 M MACs, 1 MB traffic, OI = 200) and a final FC layer (1.28 M MACs, about 1.28 MB of weights, OI ≈ 1).

The big conv speeds up about 6.5× on the accelerator while the FC layer does not speed up at all and becomes a larger share. This is the mechanism behind the paper’s design logic below.

Reading Fig. 2 and Fig. 3

The take-aways are:

• Large Conv2D layers are expensive at low RP but add cheap capacity at high RP.
• FC layers and depthwise convs are cheap at low RP and expensive at high RP.
• MobileNetV1-1.5x and MobileNetV3-Large-1.5x illustrate the two extremes: V3 minimizes MACs at the cost of memory complexity (good on CPUs), V1 spends MACs for capacity (fine on accelerators).

What MNv4’s design takes from this

MNv4 puts expensive, MAC-heavy conv layers at the start of the network (they matter for capacity, and cost only at low RP) and keeps the same large final FC layers in every variant (they matter for accuracy, and cost only at high RP). Because the two costs hurt on opposite hardware classes, no single device pays both at once. The consequence is that smaller variants carry a relatively heavier FC cost on high-RP hardware.

The practical reading is that MAC counting is already a decent proxy on CPUs (\(r_s\) of 0.94 to 0.96), but degrades on the DSP and EdgeTPU, where the roofline model stays above 0.96. Note this is a rank correlation across a handful of MobileNet-family models, which says nothing about absolute latency error.

The idea

The MobileNetV2 inverted bottleneck (IB) is [@sandler2018mobilenetv2]: 1×1 pointwise (PW) expand → depthwise (DW) conv → 1×1 PW project. UIB makes one change: it adds an optional DW before the expansion and makes the middle DW optional too. Two on/off switches give four combinations, and NAS picks per block which to use.

The same figure redrawn as a Mermaid diagram (my redraw of the paper’s figure, with the structure read off the image and from the paper’s architecture tables):

flowchart LR
  subgraph UIB["UIB (general form)"]
    direction TB
    u0([input]) --> u1["DW K1 (optional)"] --> u2["PW expand 1x1"] --> u3["DW K2 (optional)"] --> u4["PW project 1x1"] --> u5([output])
  end
  subgraph EDW["ExtraDW (K1 and K2)"]
    direction TB
    e0([input]) --> e1["DW K1"] --> e2["PW expand"] --> e3["DW K2"] --> e4["PW project"] --> e5([output])
  end
  subgraph IB["MobileNet IB (K2 only)"]
    direction TB
    i0([input]) --> i2["PW expand"] --> i3["DW K2"] --> i4["PW project"] --> i5([output])
  end
  subgraph CNX["ConvNeXt-like (K1 only)"]
    direction TB
    c0([input]) --> c1["DW K1"] --> c2["PW expand"] --> c4["PW project"] --> c5([output])
  end
  subgraph FFN["FFN (neither)"]
    direction TB
    f0([input]) --> f2["PW expand"] --> f4["PW project"] --> f5([output])
  end
  subgraph FIB["Fused IB (alternative)"]
    direction TB
    g0([input]) --> g2["kxk Conv2D expand"] --> g4["PW project 1x1"] --> g5([output])
  end
  UIB --> EDW
  UIB --> IB
  UIB --> CNX
  UIB --> FFN

The four instantiations

The paper notes ExtraDW can be viewed as a MobileNetV1-style factorization of two standard convolutional blocks. At every stage UIB lets the network choose an ad-hoc spatial vs channel mixing trade-off, enlarge the receptive field where needed, and maximize compute utilization.

Why this stays NAS-friendly

The PW expansion and projection weights are shared across the candidate instantiations; only the DW ops are the searchable options. The paper reports that a SuperNet-based search then shares over 95% of the parameters between candidates, so the search space is manageable.

Fused IB

A \(k\times k\) FusedIB replaces the expand-PW + DW pair with a single \(k\times k\) full Conv2D followed by a 1×1 PW [@akin2022edgetpu]. All MNv4 model stems use FusedIBs (Appendix D).

Does UIB help? (Table 1)

Where NAS actually puts each block (Appendix D)

The authors summarize what TuNAS chose in the final Conv models:

• ExtraDW tends to appear at the start of each searchable stage, where resolution drops sharply. The two DWs enlarge the receptive field and soften the resolution loss. It is also common in early stages.
• FFN and ConvNeXt-like blocks show up in the final layers, where earlier layers have already mixed spatially and channel mixing gives more marginal gain.

A PyTorch sketch (Not from the paper)

The paper’s reference implementation is in TensorFlow Model Garden (footnote 1 of the paper); a timm port also exists (footnote 7). The sketch below is my own illustration of the block’s logic, not the official code.

import torch
import torch.nn as nn


def conv_bn(c_in, c_out, k, stride=1, groups=1, act=True):
    layers = [
        nn.Conv2d(c_in, c_out, k, stride, padding=k // 2, groups=groups, bias=False),
        nn.BatchNorm2d(c_out),
    ]
    if act:
        layers.append(nn.ReLU(inplace=True))  # paper favours ReLU/BN (hardware support)
    return nn.Sequential(*layers)


class UIB(nn.Module):
    """Universal Inverted Bottleneck.
    start_dw_k: kernel of the optional DW *before* expansion (0 = absent)
    mid_dw_k:   kernel of the optional DW *between* expand and project (0 = absent)
      (0, 3|5) -> IB          (3|5, 0) -> ConvNeXt-like
      (k1, k2) -> ExtraDW     (0, 0)   -> FFN
    """

    def __init__(self, c_in, c_out, expand_ratio=4, start_dw_k=0, mid_dw_k=3, stride=1):
        super().__init__()
        c_mid = int(round(c_in * expand_ratio))
        self.use_res = stride == 1 and c_in == c_out  # assumption: standard IB-style residual
        # Where the stride goes is NOT specified in the paper text: verify against the
        # official TF code / timm before comparing against published latencies.
        s_start = stride if (start_dw_k and not mid_dw_k) else 1
        s_mid = stride if mid_dw_k else 1
        if stride > 1 and not (start_dw_k or mid_dw_k):
            raise ValueError("an FFN-type UIB cannot downsample")
        self.start_dw = conv_bn(c_in, c_in, start_dw_k, s_start, groups=c_in, act=False) if start_dw_k else nn.Identity()
        self.expand = conv_bn(c_in, c_mid, 1)
        self.mid_dw = conv_bn(c_mid, c_mid, mid_dw_k, s_mid, groups=c_mid) if mid_dw_k else nn.Identity()
        self.project = conv_bn(c_mid, c_out, 1, act=False)

    def forward(self, x):
        y = self.project(self.mid_dw(self.expand(self.start_dw(x))))
        return x + y if self.use_res else y


class FusedIB(nn.Module):
    """k x k full conv (expand) -> 1x1 project; used in MNv4 stems."""

    def __init__(self, c_in, c_out, expand_ratio=4, k=3, stride=1):
        super().__init__()
        c_mid = int(round(c_in * expand_ratio))
        self.conv = conv_bn(c_in, c_mid, k, stride)
        self.project = conv_bn(c_mid, c_out, 1, act=False)
        self.use_res = stride == 1 and c_in == c_out

    def forward(self, x):
        y = self.project(self.conv(x))
        return x + y if self.use_res else y

Background: MHSA vs MQA

Multi-head self-attention (MHSA) projects queries, keys, and values separately per head. Multi-query attention (MQA) [@shazeer2019mqa] keeps multiple query heads but shares a single key head and a single value head across all of them. The paper cites language-model work showing that sharing one key/value head does not noticeably cost accuracy.

Why MQA suits mobile vision

The argument is again about operational intensity. When the number of tokens is small relative to the feature dimension, loading separate K/V projections per head is mostly memory traffic. Sharing them reduces bandwidth and raises operational intensity. In mobile hybrids this regime is typical, because attention is used only in the low-resolution, high-channel final stages and batch size is usually 1. To the authors’ knowledge this is the first use of MQA for mobile vision.

Asymmetric spatial downsampling

Borrowing the “asymmetric” spirit of MQA, they add spatial reduction attention (SRA, from Pyramid Vision Transformer [@wang2021pvt]) so keys and values live at lower resolution than queries. Two differences from PVT: the downsampling is a stride-2 3×3 depthwise conv instead of average pooling, which is cheap and adds capacity. The justification is that neighbouring tokens in hybrid models are already correlated thanks to early spatial-mixing convs.

The resulting block (the paper’s Eq. 2):

\[ \text{MobileMQA}(X)=\text{Concat}(\text{attn}_1,\dots,\text{attn}_n)\,W^O,\qquad \text{attn}_j=\text{softmax}\!\left(\frac{(XW^{Q}_j)\,(\text{SR}(X)W^{K})^{\top}}{\sqrt{d_k}}\right)\text{SR}(X)W^{V} \tag{2}\]

Here \(\text{SR}\) is the stride-2 DW conv when spatial reduction is used, or the identity otherwise. Note \(W^K\) and \(W^V\) carry no head index (the sharing), while \(W^Q_j\) does.

The measured gains

The Einsum optimization (Appendix G)

Even plain MQA, though faster than MHSA, was still about 12× slower than UIB for the same MACs. The culprit was not arithmetic but transposes: on-device, an Einsum is lowered to transposes, reshapes, and batched matmuls. Transposes cost zero MACs but need a full read and write of the tensor, so they are pure memory traffic.

Two rules explain when transposes appear:

1. Contracted and non-contracted indices are not contiguous in an input. A matmul needs the contracted indices adjacent so that reshaping to 2-D is free.
2. Interleaved non-contracted indices or a reordered output: the output indices must come out in an order that avoids a post-matmul transpose.

The fix is only a re-ordering of indices in the Einsum strings and in the weight tensor layout:

# Slow: 2 transposes (non-contracted h, k interleaved with n; output reordered)
Q = tf.einsum("bnd,hdk->bhnk", X, P_q)
# Fast: no transpose
Q = tf.einsum("bnd,dhk->bnhk", X, P_q)

# Slow: 2 transposes (contracting h,v are not adjacent in the operands)
Y = tf.einsum("bhnv,hdv->bnd", O, P_o)
# Fast: no transpose
Y = tf.einsum("bnhv,hvd->bnd", O, P_o)

(Fig. 11 of the paper writes the optimized output projection as "bnhv,dhv->bnd", with the weight stored as dhv. That is the same idea as the hvd form in the Appendix G text: the contracted axes h, v are adjacent in both operands, so no transpose is needed. Treat the weight-axis order as a layout choice, not a fixed rule.)

The authors also report almost a 3× training speed-up with tf.keras on Google Tensor v4, so the benefit is not specific to inference.

A PyTorch sketch (Not from the paper)

The paper states that at the time of writing there was no PyTorch implementation of Mobile MQA, which is why hybrids were left out of the CoreML comparison. The following is my own sketch of the math in Equation 2; it does not reproduce the TensorFlow layout optimization, and PyTorch/CoreML lowering will differ.

import torch
import torch.nn as nn


class MobileMQA(nn.Module):
    """Multi-query attention on a (B, C, H, W) feature map, with optional stride-2 DW
    spatial reduction applied to keys/values only (sketch of Eq. 2)."""

    def __init__(self, dim, num_heads=4, key_dim=64, value_dim=64, kv_downsample=True):
        super().__init__()
        self.h, self.kd, self.vd = num_heads, key_dim, value_dim
        self.q = nn.Conv2d(dim, num_heads * key_dim, 1, bias=False)   # per-head queries
        self.sr = (
            nn.Sequential(
                nn.Conv2d(dim, dim, 3, stride=2, padding=1, groups=dim, bias=False),
                nn.BatchNorm2d(dim),
            )
            if kv_downsample
            else nn.Identity()
        )
        self.k = nn.Conv2d(dim, key_dim, 1, bias=False)    # ONE shared key head
        self.v = nn.Conv2d(dim, value_dim, 1, bias=False)  # ONE shared value head
        self.o = nn.Conv2d(num_heads * value_dim, dim, 1, bias=False)

    def forward(self, x):
        B, _, H, W = x.shape
        N = H * W
        q = self.q(x).view(B, self.h, self.kd, N).transpose(2, 3)   # B,h,N,kd
        src = self.sr(x)
        k = self.k(src).flatten(2)                                   # B,kd,M
        v = self.v(src).flatten(2).transpose(1, 2)                   # B,M,vd
        attn = (q @ k.unsqueeze(1)) * self.kd ** -0.5                # B,h,N,M
        attn = attn.softmax(dim=-1)
        out = attn @ v.unsqueeze(1)                                  # B,h,N,vd
        out = out.transpose(1, 2).reshape(B, N, self.h * self.vd)    # B,N,h*vd
        out = out.transpose(1, 2).reshape(B, self.h * self.vd, H, W)
        return self.o(out)

Design philosophy: simplicity

Before searching, the authors ran correlation analyses between cost models and hardware, and kept only components whose costs correlate well across devices and that sit near the Pareto frontier. Findings that shaped the search space:

• Multi-path blocks (group convs, ShuffleNet-style) have low MACs but are inefficient because of memory-access complexity.
• Modules with weak hardware support: Squeeze-and-Excite, GELU, and LayerNorm are poorly supported on DSPs (LayerNorm lags BatchNorm), and SE is slow on accelerators.
• Conventional components win: depthwise and pointwise convs, ReLU, BatchNorm, and plain attention.

From these they derived three principles: use standard widely supported components, use flexible UIB blocks, and use straightforward attention (Mobile MQA).

Neural architecture search

NAS uses TuNAS [@bender2020tunas] with modifications, and a separate search per model size (rather than EfficientNet-style fixed scaling rules [@tan2019efficientnet]).

flowchart TD
  A["Stage 1: coarse-grained search<br/>fixed block: IB, expansion 4, 3x3 DW<br/>search 'filter sizes'"] --> B["Stage 2: fine-grained search<br/>per block: presence of 2 DWs (K1, K2)<br/>kernel 3x3 or 5x5 each<br/>expansion held at 4"]
  B --> C["Final UIB macro-architecture<br/>(MNv4-Conv-S/M/L)"]
  D["Offline JFT-distillation dataset<br/>(reduces reward noise)"] -.-> A
  D -.-> B

Search space and targets (Appendix A)

• Fixed initial layers: a 3×3 stride-2 Conv2D, then a NAS-optimized stride-2 FusedIB stage.
• Searched: the number of UIB blocks and each block’s instantiation for the remaining four stages.
• Fixed head: same as MobileNetV3.
• Pointwise convs have low operational intensity at high resolution, so MAC-dense ops were favoured in the early layers.

Key claim: because the space is restricted to components with well-correlated cost models, optimizing EdgeTPU latency alone produced universally efficient models. This is why a single search target suffices, and it is also the biggest assumption to test if you adapt the recipe to other hardware.

From Conv to Hybrid models

Attention is most effective in the last stages, in line with other hybrids. MNv4-Hybrid models interleave Mobile MQA blocks with UIB blocks in the late stages. The architecture tables show the pattern: in Hybrid-M, MQA blocks alternate with ExtraDW/ConvNeXt/FFN blocks at the 16×16 stage and again at 8×8.

ImageNet-1k classification

Protocol. MNv4 models are trained on the ImageNet-1k training split only (no extra data) and evaluated on its validation split. MobileNetV1–V3 were retrained with modern recipes (+3.4 to 74.0 for V1, +1.4 to 73.4 for V2, +0.3 to 75.5 for V3) so the comparison isolates architecture. Other baselines (MobileOne, FastViT, EfficientViT, ConvNeXt, NextViT) use their reported accuracies with the authors’ own latency measurements.

What the authors conclude

• CPU: roughly 2× faster than MobileNetV3, and well ahead of iso-accuracy competitors.
• EdgeTPU: as accurate as MobileNetV3 at about half the latency. MNv4-Conv-M is more than 50% faster than MobileOne-S4 and FastViT-S12, and 1.5 points above MobileNetV2 at comparable latency.
• S23 GPU and iPhone 13 (Apple Neural Engine): mostly on the Pareto front. The nearest competitor on S23 GPU, EfficientViT, has more than 2× the CoreML latency at equal accuracy; FastViT, which is tuned for the Neural Engine, is second on CoreML but more than 5× slower on S23 GPU.
• DSP: Conv models remain top performers; hybrids are incompatible.
• Some models reach the same accuracy with fewer MACs (EfficientViT is the example), but the paper’s argument is that MACs are the wrong target.

COCO object detection

Backbones are plugged into RetinaNet with a 256-d FPN on P3–P7 and a 256-d head of 4 conv layers, using depthwise-separable convs in decoder and heads (standard for mobile detectors). Training: COCO-17, 600 epochs, 384 px input, random flip/scale and RandAugment without Shear and Rotate (they hurt small-object AP), batch size 2048, Adam, L2 weight decay 3e-5, cosine LR with a 24-epoch warm-up, stochastic depth 0.2 for MNv4 backbones, and a per-model tuned learning rate.

MNv4-Conv-M matches the AP of MobileNet-MultiAvg and V2 (32.6–32.9) while being 12% and 23% faster on Pixel 6 CPU. The Hybrid-M detector is 1.4 AP above Conv-M by the table (the paper’s prose says +1.6; see the check below) at 18% higher latency. MobileNet-V3-Large @ 2.0 scores 0.6 AP above Conv-M with fewer MACs but more parameters and about 17% more latency, so MNv4’s detection wins are on the latency axis, not on raw AP.

Benchmarking methodology (Appendix B)

• Formats: TensorFlow Lite, INT8 for mobile CPUs, Hexagon DSP, and EdgeTPUs; FP16 for mobile GPUs.
• Protocol: about 1000 runs per measurement, mean taken; repeated 5 times; the median of the means is reported.
• CPU: pinned to the fastest core, XNNPACK backend.
• Apple Neural Engine: iPhone 13 (iOS 16.6.1), CoreMLTools 7.1, Xcode 15.0.1; PyTorch → CoreML MLProgram in Float16 with float16 MultiArray inputs to minimize input copying.

Training recipe (Appendix C)

Augmentation: Inception Crop, horizontal flip, RandAugment, Mixup, CutMix. Regularization: weight decay/L2, stochastic depth, with strength scaled by model size.

Dynamic dataset mixing

Instead of one fixed augmentation pipeline, the student trains on a dynamic mixture of datasets produced with different augmentations:

• D1: Inception Crop + RandAugment (l2m9) on replicas of ImageNet-1k.
• D2: Inception Crop + extreme Mixup on replicas of ImageNet-1k (the Patient Teacher setup).
• D3: class-balanced in-domain data from JFT-300M [@sun2017jft]: resampled to 130K images per class (130 M total) using a relevance threshold above 0.3 from an ImageNet-1k-trained EfficientNet-B0 (following Noisy Student [@xie2020noisystudent]); the top 130K for plentiful classes, replication for rare ones; the dataset is replicated 10×, with weaker augmentation (Inception Crop + RandAugment l2m5).

What the ablation shows

Note the +1.8 compares a 2000-epoch run with the 400-epoch D2 baseline, so part of it is simply longer training. Interpretation: mixing augmentations expands the effective sample space and makes the task harder and more diverse; JFT data alone hurts (81.8), but adding it to ImageNet data helps. Note the train accuracy gets lower (82.6) while validation rises, which is consistent with better generalization.

The 2000-epoch Conv-L student reaches 85.9%, only 1.6 below its EfficientNet-L2 teacher with 15× fewer parameters and 48× fewer MACs. The headline 87.0% for Hybrid-L comes from combining this distillation with JFT pretraining (stated in the last sentence of Sec. 8, not in Table 9, which lists 86.6).

Implementation details (Appendix H)

flowchart LR
  A["Candidate image"] --> B["Augment<br/>(crop + RandAug / Mixup)"] --> C["JPEG encode and decode"] --> D["Teacher: EfficientNet-L2<br/>(480 M params, 290 G MACs, 87.5% top-1)"] --> E["Store softmax over 1000 classes<br/>(soft label)"]
  E --> F["Student training<br/>flip only, dropout in final FC<br/>cross-entropy to soft labels<br/>AdamW + cosine LR + warm-up"]

• Because augmentations and teacher outputs are precomputed, training is about 2.5× faster than regular ImageNet training and much less sensitive to hyperparameters.
• Example cost: MNv4-Conv-L, batch size 16 k, 400 epochs takes 2.3 hours on 128 TPU v5e.

Which variant should I pick? (My synthesis, not from the paper)