Mobile Hybrid Backbones: MobileViT & FastViT

Chapter 1: Local vs Global

To design a hybrid, you need a sharp picture of what each operation actually gives you. Let's make "local" and "global" precise.

A convolution with a k×k kernel computes each output pixel from a k×k window of the input. Its receptive field — the set of input pixels that can influence one output — is exactly k×k after one layer. Stack L convolution layers (stride 1, kernel k) and the receptive field grows to roughly 1 + L·(k−1) along each axis. With 3×3 kernels (k = 3) you gain 2 pixels of reach per layer. To connect two pixels 64 apart you need about 32 layers. That is the structural reason CNNs are "local": global mixing is possible but expensive in depth.

Worked example. You have a 32×32 feature map and want every pixel to be able to see every other pixel. With 3×3 convolutions, k−1 = 2, so to reach across 31 pixels you need L ≈ 31/2 ≈ 16 layers — and that is just for the receptive field to touch the far corner, not to mix information richly. Each of those 16 layers costs compute and adds latency. A single self-attention layer connects all 1,024 pixels in one shot.

Self-attention, by contrast, is global by construction. For N tokens it computes an N×N matrix of pairwise scores, so the receptive field is the entire input after one layer — RF₁ = N. The price is the O(N²) score matrix and the O(N²·d) compute. On a 32×32 map flattened to 1,024 tokens, that's a 1,024×1,024 = ~1M-entry attention matrix per head. On a phone, prohibitive.

So the design question crystallizes: how do we get attention's instant global reach without paying its O(N²) bill across the whole image, and without discarding convolution's free locality bias? MobileViT's answer is delightfully literal. It runs convolution to do the cheap local work and gather neighborhoods, then runs attention over a small, coarse set of positions to do the global work, then folds the result back. FastViT's answer is to do almost everything with reparameterizable convolutions and use attention sparingly, only in the deepest, lowest-resolution stage where N is tiny.

Chapter 2: The MobileViT Block

MobileViT's central contribution is a single block that gives you "transformers as convolutions" — it learns global representations with the spatial inductive bias of convolutions, while keeping the data flow that mobile hardware likes. Let's trace it tensor by tensor.

The block takes a feature map X of shape [H, W, C] (height, width, channels). It produces an output of the same spatial size and channel count, so it drops into a CNN like any other block. Inside, it does five steps:

The genius is in steps 2–4, the unfold → attend → fold sandwich, so let's slow down on what "P pixel-positions" means, because it is the whole trick.

Take the feature map and tile it into non-overlapping patches of size w×h (the paper uses 2×2). A 2×2 patch has P = 4 pixels. If the map is H×W, the number of patches is N = (H·W)/P. Now here is the unusual reshape: instead of flattening each patch into one token (as a plain ViT would), MobileViT keeps the P positions inside the patch separate. It groups all the "top-left" pixels of every patch into one sequence of length N, all the "top-right" pixels into another sequence of length N, and so on — P sequences in total.

Then it runs the transformer along the N axis, separately for each of the P positions. So the top-left pixel of patch 5 attends to the top-left pixels of all other patches; it does not attend to the bottom-right pixel of its own patch. Why this strange choice? Because every pixel inside a patch already got mixed with its in-patch neighbors during the local rep convolution in step 1. The transformer's job is only to provide the inter-patch, global communication that convolution can't cheaply give. By attending position-wise, MobileViT avoids redundant work and — crucially — preserves the spatial order of pixels, so the fold in step 4 is a clean, exact inverse of the unfold.

python
# MobileViT block — shape-faithful skeleton (PyTorch-style pseudocode)
import torch, torch.nn as nn

class MobileViTBlock(nn.Module):
    def __init__(self, C, d, n=3, patch=(2, 2), L=2):
        super().__init__()
        self.ph, self.pw = patch                 # 2, 2  -> P = 4
        self.local = nn.Sequential(             # step 1: local rep
            nn.Conv2d(C, C, n, padding=n//2, groups=C),  # depthwise n x n
            nn.Conv2d(C, d, 1))                 # 1x1 -> project to d
        self.transformer = nn.Sequential(*[TransformerLayer(d) for _ in range(L)])
        self.proj = nn.Conv2d(d, C, 1)        # back to C
        self.fuse = nn.Conv2d(2*C, C, n, padding=n//2)

    def forward(self, x):                     # x: [B, C, H, W]
        y = self.local(x)                     # [B, d, H, W]
        B, d, H, W = y.shape
        ph, pw = self.ph, self.pw
        # step 2: unfold -> [B, d, P, N] with P = ph*pw, N = (H//ph)*(W//pw)
        y = y.reshape(B, d, H//ph, ph, W//pw, pw)
        y = y.permute(0, 1, 3, 5, 2, 4).reshape(B, d, ph*pw, -1)
        y = y.permute(0, 2, 3, 1)   # [B, P, N, d]
        # step 3: attend across N, independently for each of the P positions
        y = self.transformer(y)             # [B, P, N, d]
        # step 4: fold (exact inverse of the unfold)
        y = y.permute(0, 3, 1, 2).reshape(B, d, ph, pw, H//ph, W//pw)
        y = y.permute(0, 1, 4, 2, 5, 3).reshape(B, d, H, W)
        # step 5: project, concat with input, fuse
        y = self.proj(y)                      # [B, C, H, W]
        return self.fuse(torch.cat([x, y], dim=1))

Chapter 3: Unfold → Attend → Fold

This is the mechanical heart of MobileViT, and it is worth doing with actual numbers so the reshape stops being mysterious. We'll use a tiny example you can verify by hand.

Let the feature map be 4×4 with a single channel (d = 1 for readability), and let the patch be 2×2, so P = 4 and N = (4·4)/4 = 4. Label the 16 pixels by their value:

The four 2×2 patches are: patch 0 = {a,b,c,d} (top-left), patch 1 = {e,f,g,h} (top-right), patch 2 = {i,j,k,l} (bottom-left), patch 3 = {m,n,o,p} (bottom-right). Within each patch, the four positions are top-left (TL), top-right (TR), bottom-left (BL), bottom-right (BR).

The unfold groups by position, not by patch. Position TL collects the TL pixel from every patch; position TR collects every TR; and so on. The result is a [P=4, N=4] table:

Now the attend step runs a transformer along each row independently. The TL row [a, e, i, m] becomes a length-4 sequence; self-attention lets a (TL of patch 0) mix with e, i, m (TLs of the other patches). The four rows are processed in parallel and never talk to each other in this step. After attention each entry is a new mixed value, but the table keeps its [4, 4] shape and its layout.

Finally the fold reverses the grouping exactly: TL of patch 0 goes back to grid cell (0,0), TR of patch 0 to (0,1), and so on. Because each value carries a unique (position, patch) address, the fold is a bijection — no information is lost, no pixel is misplaced. The output is a 4×4 map, same shape as X, but now every pixel has absorbed global context from the same-position pixels across the whole image.

Chapter 4: The Full MobileViT

A block is not a network. MobileViT the architecture interleaves MobileNetv2 inverted-residual blocks (cheap, local, downsampling) with a few MobileViT blocks (global) at the lower-resolution stages, where N is small enough that attention is affordable.

The stem is a strided 3×3 convolution that takes the 256×256×3 image down to 128×128. Then the body alternates: several MV2 blocks downsample and build local features, and at three points — once the spatial size has shrunk to 32×32, 16×16, and 8×8 — a MobileViT block injects a global view. Placing attention only at coarse resolutions is the budget trick from Chapter 1: at 8×8 with a 2×2 patch, N = 16 — a 16×16 attention matrix, trivial.

MobileViT ships in three sizes — XXS (~1.3M params), XS (~2.3M), and S (~5.6M) — by scaling the channel widths and the transformer dimension d. Even the largest, MobileViT-S, is smaller than MobileNetv3-Large in parameters yet more accurate on ImageNet-1k, and it does so without the giant pre-training datasets ViTs usually demand. That last point is the headline: because the convolutional spine supplies the spatial inductive bias, MobileViT trains to strong accuracy on plain ImageNet-1k, where a comparable pure ViT badly underperforms.

Let's quantify the attention savings. A pure ViT operating at 32×32 = 1,024 tokens would build a 1,024×1,024 ≈ 1.05M-entry attention matrix. MobileViT at 32×32 with 2×2 patches has N = 256 patches and runs attention over N for each of P = 4 positions: 4 × (256×256) = 4 × 65,536 = 262,144 entries — a 4× reduction at that stage, and far more once you account for only running attention at three coarse stages instead of every layer. The deeper the stage, the cheaper: at 8×8, N = 16, so the matrix is a negligible 16×16.

One more design choice worth calling out: MobileViT uses no positional embeddings in its transformer. A plain ViT needs them because flattening patches destroys spatial order. MobileViT doesn't, because the convolutions before and after the transformer, plus the position-preserving unfold/fold, already encode where everything is. The spatial structure is carried by the data layout, not by added position vectors — one fewer thing to learn, one fewer thing to get wrong at a new resolution.

Chapter 5: FastViT & Reparameterization

Fast-forward to 2023. MobileViT proved hybrids work; FastViT (Vasu et al.) attacks the next bottleneck — raw on-device latency — with a different and beautiful idea: build a network that is heavy and accurate during training, then collapse it into a lean network for inference, with identical outputs. The collapse is exact algebra, not approximation. FastViT-T8 runs a forward pass in roughly 0.8 ms on a recent mobile GPU while matching the accuracy of models several times slower.

The trick is structural reparameterization, introduced for CNNs by RepVGG (Ding et al., 2021). The insight: a sum of linear operations is itself a single linear operation. During training you give the network several parallel branches — say a 3×3 conv, a 1×1 conv, and an identity skip — which together are easier to optimize and reach higher accuracy. But each branch is linear (a convolution is a linear map), and they're summed, so at inference you can fold all three into one equivalent 3×3 conv whose weights are the sum of the branch weights (after padding the smaller kernels). One kernel, one memory read, one fused operation — and bit-for-bit the same result.

Why is this a giant deal for phones? Because each branch you keep at inference is a separate kernel launch, a separate read of the input tensor from memory, and a separate write — and on mobile accelerators, memory bandwidth, not arithmetic, is the wall. Three branches mean three trips to memory; one fused conv means one. The reparameterization converts training-time accuracy (lots of branches) into inference-time speed (one branch) for free.

Folding BatchNorm into the conv too

FastViT folds more than parallel convs. After training, a convolution followed by BatchNorm can be merged into a single convolution. BatchNorm at inference is an affine map per channel: it subtracts the running mean μ, divides by √(σ² + ε), scales by γ, and shifts by β. Because all of that is linear in the conv output, you can bake it directly into the conv's weights and bias.

Worked example, one channel. Suppose a conv produces a pre-BN value u = W·x + b with W = 2, b = 0. The BN stats are μ = 1, σ² = 4 (so √σ² = 2, take ε ≈ 0), γ = 3, β = 5. BN output is 3·(u − 1)/2 + 5. Fold it: the new weight is W′ = 2 · 3/2 = 3, the new bias is b′ = 5 − 1·3/2 = 3.5. Check at x = 2: original path gives u = 4, BN gives 3·(4−1)/2 + 5 = 4.5 + 5 = 9.5; folded path gives 3·2 + 3.5 = 6 + 3.5 = 9.5. Identical, and now it's one conv with no BatchNorm layer at all.

python
# Fold a Conv2d -> BatchNorm pair into one equivalent Conv2d
import torch

def fuse_conv_bn(W, b, gamma, beta, mu, var, eps=1e-5):
    std = (var + eps).sqrt()           # per-channel
    scale = gamma / std                  # [C_out]
    W_fused = W * scale.view(-1, 1, 1, 1)  # scale every kernel
    b_fused = beta + (b - mu) * scale       # fold the shift
    return W_fused, b_fused                # one conv, BN gone

Chapter 6: RepMixer

FastViT's signature block is the RepMixer, and it applies the reparameterization idea to token mixing — the job attention normally does. The realization is sharp: most of a transformer block's cost is the attention's token mixer, but at the resolutions FastViT cares about, a reparameterizable convolution can mix tokens almost as well, far faster, and then fold itself away.

Recall the standard token-mixer in efficient transformers (e.g. the MetaFormer view): x = x + TokenMixer(Norm(x)). A common cheap choice (PoolFormer) is TokenMixer = Pool. FastViT instead defines the RepMixer's training-time mixing as a depthwise convolution wrapped in a residual, with the normalization folded into the conv:

Read that carefully. The block adds a residual X, then a mixed term BN(DWConv(X)), and subtracts an X. The subtraction is deliberate: it makes the whole token-mixer a single residual around a depthwise conv, structured so that all the linear pieces — the depthwise convolution, the BatchNorm, the identity skip, and the subtraction — can be folded into one depthwise convolution at inference. After reparameterization:

That's the punchline. At training time the RepMixer is a multi-term residual structure (good gradients, good accuracy). At inference time it is a single depthwise convolution — one of the cheapest operations that exists on mobile hardware. No skip connection to read twice, no separate BatchNorm pass, no attention matrix. Token mixing for the price of one depthwise conv.

A full FastViT stage stacks RepMixer blocks (token mixing) with a convolutional FFN (channel mixing), and — only in the last stage, where the feature map is tiny — a few genuine self-attention blocks for the strongest global reasoning. Everything reparameterizable folds; the handful of attention layers at the bottom stay, but N is so small there that they cost almost nothing. FastViT also uses large-kernel depthwise convs and a conv stem that are all reparameterized, so the entire backbone, save those few attention layers, is a stack of single fused convolutions at inference.

Chapter 7: Results & Ablations

Numbers are where the philosophy meets the device. Both papers measure latency on real hardware, not FLOPs, because — as Chapter 0 warned — FLOPs are a poor proxy for mobile speed.

MobileViT's headline

MobileViT-S reaches roughly 78% top-1 on ImageNet-1k with about 5.6M parameters, trained on ImageNet-1k alone (no giant pre-training corpus). At comparable parameter budgets it beats MobileNetv3, DeIT-style ViTs, and earlier lightweight transformers. The ablations show the hybrid structure is load-bearing: remove the convolutional local-rep and accuracy drops, because the transformer then has to learn locality from scratch; remove the transformer global-rep and you lose the long-range reasoning that lifts it above a pure CNN. The two halves are complementary, not redundant.

FastViT's headline

FastViT's family spans roughly 75% to 84% top-1. The smallest, FastViT-T8, hits about 75.6% at around 0.8 ms mobile-GPU latency; larger variants climb toward 84% while staying markedly faster than comparable-accuracy competitors (EfficientFormer, MobileViT itself, ConvNeXt-class models). Crucially, the reparameterization gives an accuracy bump and a latency reduction at once: the train-time branches lift accuracy by a point or so, and folding them away cuts inference latency, so the model lands strictly up-and-left on the accuracy–latency plot.

The FLOPs-vs-latency lesson

Here is the ablation that reshapes how you think about efficiency. FastViT's authors point out that two blocks with the same FLOP count can have wildly different latencies. A multi-branch RepVGG-style block and its fused single-conv twin compute the same number of multiply-adds, yet the fused version is much faster because it touches memory once instead of three times. Likewise, attention's reshapes and softmax add latency that FLOP counters miss. This is why both papers optimize the measured wall-clock on a device — and why a model can have more FLOPs yet run faster.

Chapter 8: Latency Explorer

This is the showcase. You're the architect now. Build a mobile backbone by choosing how much attention to use, at what resolution, and whether to reparameterize — and watch the estimated accuracy, latency, and parameter count respond in real time. The model behind it captures the two papers' central trade-offs: attention buys accuracy but costs latency that explodes with token count N = (size/patch)²; reparameterization folds branches to cut latency for free; deeper attention stages are cheaper because N has collapsed.

Drive it to the conclusions: push the attention stage to high resolution and latency spikes (N² cost); move it deep and it's nearly free; toggle reparameterization and watch latency drop with accuracy held. Find the sub-millisecond, high-accuracy corner — that's where FastViT-T8 lives.

What you should feel through the controls: there is no single "best" model — only points on a frontier. MobileViT chose to push accuracy with affordable mid-depth attention; FastViT chose to push latency by reparameterizing almost everything and rationing attention to the deepest stage. Both are valid answers to "put a transformer on a phone," and the explorer lets you walk the curve between them.

Chapter 9: Limitations & Connections

These two papers are triumphs of pragmatism, but they are not the end of the story. Let's be honest about where they strain, then connect them to the wider map.

Limitations

MobileViT. The unfold/fold reshapes, while mathematically clean, are not always friendly to mobile compilers — some accelerators handle the permute-heavy data movement poorly, which is part of why later work (MobileViTv2) replaced the multi-head attention with a cheaper separable attention. Its latency, while better than a pure ViT, is still higher than the fastest pure CNNs of its era. And the global rep is confined to coarse stages, so very fine-grained long-range structure can be missed.

FastViT. Reparameterization only folds linear branches; any nonlinearity inside a branch blocks the fold, constraining the block design space. The training-time model is heavier (more memory, slower training) even though inference is lean. And the approach assumes you control the full train→fuse→deploy pipeline; if you only have a pretrained checkpoint without the multi-branch structure, you can't recover the training-time gains.

Both. They target classification backbones; transferring the latency wins to dense tasks (detection, segmentation) where high resolution must be preserved is harder, because that's exactly where attention's N² cost reappears. And both are tuned to specific device classes — a kernel fusion that helps one accelerator may not help another.

Connections

The ideas here sit in a rich neighborhood. Follow these threads:

Cheat sheet