6 MobileNetV2

A standard conv does two jobs at once

A standard 3\(\times \)3 convolution from \(C_\text {in}\) channels to \(C_\text {out}\) channels computes, for every output cell:

That sum is doing two things. The \(k_h, k_w\) part is spatial mixing: it combines a pixel with its neighbors. The \(c_\text {in}\) part is cross-channel mixing: it combines the \(C_\text {in}\) channels into each of \(C_\text {out}\) new channels. Both happen inside the same kernel and at the same cost \(C_\text {in} \times C_\text {out} \times k^2\) weights per layer.

Howard et al. 2017 (MobileNetV1, arXiv:1704.04861) and Sandler et al. 2018 (MobileNetV2, arXiv:1801.04381) asked: what if we factored those two responsibilities and only paid the full cost of one of them? The spatial mixing is locally information-rich and benefits from a real \(3 \times 3\) neighborhood. The channel mixing can be approximated by a much cheaper operation. If we separate them, we get parameter efficiency without losing much expressive power.

6.1 Run it first

Before the factorization, train the factorized network. Four commands, and about thirty-five minutes of GPU time:

lake exe cache get                            # Mathlib oleans, ~30 s
./scripts/datasets/download_imagenette.sh     # 330 MB dl, 2.4 GB unpacked
lake build mobilenetv2-verified-adam
./.lake/build/bin/mobilenetv2-verified-adam data

On one RTX 4060 Ti (CUDA 12.9), from runs/2026-08-12-mnv2-imagenette-xla-cuda/. XLA’s startup banner is removed, the network description line is wrapped, and the middle of the run is elided:

[pjrt_ffi] XLA backend: PJRT 0.112, 1 device(s)
[pjrt_ffi] compiled verified_mlir/mobilenetv2_adam_train_step.mlir
             (@mobilenetv2_adam_train_step, 581 outputs, 1 replica) in 8101 ms
[pjrt_ffi] compiled verified_mlir/mobilenetv2_fwd.mlir
             (@mobilenetv2_fwd, 1 outputs, 1 replica) in 665 ms
[pjrt_ffi] compiled verified_mlir/mobilenetv2_fwd_eval.mlir
             (@mobilenetv2_fwd_eval, 1 outputs, 1 replica) in 250 ms
MobileNetV2 on Imagenette 224² (stem-s2 → 17 inverted-residual blocks,
  full-paper [t,c,n,s] config, stride-2 depthwise downsamples 224→7 →
  head conv-BN-relu6 → GAP → dense)
  via the VERIFIED renderer → XLA/PJRT → GPU
  train 9469, val 3925; bs 32, MobileNetV2 adam
    (cosine+warmup 3ep, baseLR 0.001000), He init
  running-stats BN: 52 layers, 34112 stat floats → eval via @mobilenetv2_fwd_eval
Epoch 1/80: loss=2.121050 lr=0.000333
  epoch 1: val_acc = 947/3925 = 24.127389%  top5 = 2392/3925 = 60.942675%
Epoch 2/80: loss=1.784318 lr=0.000667
  epoch 2: val_acc = 680/3925 = 17.324841%  top5 = 2357/3925 = 60.050955%
Epoch 3/80: loss=1.590689 lr=0.001000
  epoch 3: val_acc = 1596/3925 = 40.662420%  top5 = 3423/3925 = 87.210191%
Epoch 4/80: loss=1.446373 lr=0.001000
  epoch 4: val_acc = 1851/3925 = 47.159236%  top5 = 3428/3925 = 87.337580%
Epoch 79/80: loss=0.517682 lr=0.000000
  epoch 79: val_acc = 3506/3925 = 89.324841%  top5 = 3874/3925 = 98.700637%
Epoch 80/80: loss=0.517213 lr=0.000000
  epoch 80: val_acc = 3503/3925 = 89.248408%  top5 = 3873/3925 = 98.675159%
done (trained MobileNetV2 adam + cosine/warmup via packed threading).

Eighty epochs at about 26 seconds each, thirty-five minutes in total, and 89.25% top-1 with 98.68% top-5 on Imagenette’s 3,925-image validation split. The best epoch reached \(89.32\% \). Set that against Chapter 5, which ran the same data, the same recipe and the same number of epochs through a network with \(9.5\times \) the parameters and averages \(89.99 \pm 0.32\) top-1 over five seeds (runs/2026-09-12-r34-ablation-fp32-seeds/). MobileNetV2 gives up three quarters of a point of top-1 — less than one run’s own \(\pm 0.97\) interval — picks up a fifth of a point of top-5, and does it in thirty-five minutes against forty-seven.

Epoch 2 is worth a glance before moving on. Validation accuracy falls, from \(24.1\% \) to \(17.3\% \), while the training loss drops normally from \(2.12\) to \(1.78\). That is the warmup: the learning rate is still climbing toward its peak at epoch 3, and the BatchNorm running statistics are two epochs old and averaging over a network whose weights are moving fast. The evaluation forward is using statistics that no longer describe the thing being evaluated. It corrects itself by epoch 3 and never recurs.

Depthwise separability doesn’t weaken it. That run did not execute a reimplementation of the network this chapter describes. It executed verified_mlir/mobilenetv2_adam_train_step.mlir, which is pretty(provenGraph) off the renderer. Every depthwise convolution in it backpropagates by the rule Theorem 71 proves, every pointwise convolution by Theorem 34, every BatchNorm by Theorem 49, and every residual skip by Theorem 58. Those per-operation theorems are unconditional and hold at any width or depth, so they cover all seventeen blocks.

One scope note, because this is the chapter where the distinction starts to bite. The theorem that folds those rules into a single statement about the whole network, at the batch statistics this artifact runs, is mobilenetv2ForwardBFullHasVJPAt_correct, in Proofs/Nets/MobileNet/MobileNetV2FullBVJP.lean, and it covers the stem, all seventeen bottlenecks and the head. The typed graph the artifact is printed from denotes that forward by mobilenetv2FwdGraphBFull_faithful — with classifier dropout, at every mask, by mobilenetv2FwdGraphBFullDo_faithful; the inference render’s typed graph denotes the frozen-statistics forward of one example by mobilenetv2FwdGraphPaperEval_faithful, and both renders’ block text is pretty of their typed blocks (FwdGraphTextTies) — and the tie from the mobilenetv2Verified layer list to that forward is mobilenetv2VerifiedB_denote_eq, a rfl that stops reducing if any [t,c,n,s] entry of the spec drifts, with mobilenetv2VerifiedB_fwd_faithful above it.

The full-depth fold is pointwise. It holds at a given input, and it carries one side condition per activation site, at every example of the batch, saying that the value entering relu6 is neither \(0\) nor \(6\). There are thirty-five such sites per example, and the reason for every one of them is that relu6 has two kinks. The network is differentiable away from them and is not differentiable at them, so a statement quantified over all inputs would be false. Chapter 7 does fold EfficientNet-B0 over all inputs (Theorem 87), and it can do that because swish is smooth everywhere. The axis that moved here was the depth, from two blocks to seventeen, and not the pointwise-to-global one.

The three compile lines are every BatchNorm net’s three: @mobilenetv2_fwd_eval is a separate forward because BatchNorm evaluates differently than it trains, and running-stats BN: 52 layers, 34112 stat floats is the driver threading the running averages out of the train step and into it. This network has 52 BatchNorms against ResNet-34’s 36, which is what a block built from three convolutions instead of two costs.

Depthwise: one kernel per channel, no cross-channel sum

A depthwise convolution does the spatial half on its own, per channel. Each output channel comes from one input channel, with its own \(k \times k\) kernel:

Only one summation level (over the kernel positions), not two. No mixing between channels at all. Channel \(c\) of the output only ever looks at channel \(c\) of the input. The weight tensor is shape \([C, k, k]\) instead of \([C_\text {out}, C_\text {in}, k, k]\):

For a typical \(C_\text {in} = C_\text {out} = 64\) layer, that is 576 weights against 36,864, a 64\(\times \) reduction. The Jacobian from Chapter 3 simplifies in step: one fewer \(\Sigma \) level means the depthwise input-VJP is a slightly simpler reversed-kernel convolution, the weight VJP is a slightly simpler transpose-trick, and the bias VJP is the same per-channel sum. Theorems 71 through 73 prove the depthwise VJPs.

The depthwise-separable sandwich

A depthwise conv alone never lets channels see each other. That is a problem, because the whole point of stacking convolutions is to build hierarchical features that combine information across channels. The fix is to follow each depthwise with a pointwise \(1 \times 1\) convolution, which is purely cross-channel mixing (no spatial structure, since each output pixel is a learned linear combination of the same input pixel’s channels). That two-step combination,

is called a depthwise-separable convolution, and it is what replaces the standard \(3 \times 3\) conv in MobileNet and everywhere downstream. A standard \(3 \times 3\) over 64\(\to \)64 costs 36,864 weights. The separable version costs 576 (depthwise) plus 4,096 (pointwise), for 4,672. Eight times fewer.

Inverted residual: depthwise at the wide point

MobileNetV2 wraps the separable block in one more idea. ResNet’s bottleneck block is wide \(\to \) narrow \(\to \) wide (project down for the expensive spatial conv, project back up for the residual). MobileNetV2 flips it: narrow \(\to \) wide \(\to \) narrow, with the depthwise spatial conv at the wide expansion. The intuition is that depthwise is so cheap that you can afford to make it wider. The narrow bottleneck channels carry the actual information across the residual.

Concretely, an inverted-residual block with input channels \(C_\text {in}\), output channels \(C_\text {out}\), and expansion ratio \(t\):

1. \(1 \times 1\) conv \(C_\text {in} \to t \cdot C_\text {in}\) (the expand)

2. \(3 \times 3\) depthwise at \(t \cdot C_\text {in}\) channels (the spatial mix at the wide point)

3. \(1 \times 1\) conv \(t \cdot C_\text {in} \to C_\text {out}\) (the project, narrowing back down)

4. Residual skip when \(C_\text {in} = C_\text {out}\) and stride equals 1

This is the .invertedResidualNB ic mid oc stride primitive in the spec, one layer constructor that emits all three convs, the BN and relu6 between them, and the conditional residual add. It takes the expanded width mid directly rather than the ratio \(t\), so a block is written the way it is shaped rather than the way it is parameterized, and the NB suffix says the three convs carry no bias, since the BatchNorm immediately after each one subsumes it. The VJP for the whole block composes the depthwise VJP from this chapter with the conv2d VJP (Theorem 34), the BN VJP (Theorem 49) and the residual’s additive fan-in (Theorem 58). Once again there is no new math beyond the foundation rules, only the depthwise primitive plus composition.

Stacking the blocks

MobileNetV2 stacks 17 inverted-residual blocks across seven stages (at depths 1, 2, 3, 4, 3, 3, 1) with widths progressing from 16 to 320 channels, then a final \(1 \times 1\) to 1280 channels, global average pool, and a 1280-to-10 dense head. That is the same stem-stages-head template as ResNet-34. Total: 2.24M parameters, 9.5\(\times \) fewer than ResNet-34, at a modest accuracy cost (\(89.25\% \) against ResNet-34’s \(89.99\% \) on Imagenette, both measured in this book). For mobile and embedded deployment that trade is the whole reason MobileNet exists.

6.2 The theorems

The ladder of Chapter 5, twice: MobileNetV2 and MobileNetV4-Conv-M take each of the six rungs in turn.

6.3 Example: MobileNet V2 on Imagenette

The architecture

The same vertical column as ResNet-34, now built from inverted-residual blocks: a \(3 \times 3\) stride-2 stem, seven stages of .invertedResidualNB blocks (17 in all, widths \(16\) to \(320\), expansion \(t{=}6\) except the first), a \(1 \times 1\) lift to 1280 channels, global average pool, and a dense head. There are no max-pools, since every downsample is folded into a strided block. The inset shows one inverted-residual block: the cheap depthwise spatial mix happens at the wide point.

The verified trainer

§5.3 is the pattern and this is the first chapter to reuse it, so read the two together. Everything that carries weight there carries the same weight here, and only the layers list and the bnChannels that follows from it are new:

import LeanMlir.Verified.NetsCore
import LeanMlir.Verified.Train

-- 1. The network. `slug` names the committed render, `bnChannels`
--    drives BN threading, `NB` means the convs carry no bias.
def mobilenetv2Verified : VerifiedNetSpec where
  name     := "MobileNetV2"
  slug     := "mobilenetv2"
  inC      := 3
  imageH   := 224
  imageW   := 224
  data     := .imagenette
  layers   := [
    .convBnNB 3 32 3 2,                   -- stem       224->112
    .invertedResidualNB  32   32  16 1,   -- stage 1, t=1   @112
    .invertedResidualNB  16   96  24 2,   -- stage 2    112->56
    .invertedResidualNB  24  144  24 1,
    .invertedResidualNB  24  144  32 2,   -- stage 3     56->28
    .invertedResidualNB  32  192  32 1,   --   x2
    .invertedResidualNB  32  192  64 2,   -- stage 4     28->14
    .invertedResidualNB  64  384  64 1,   --   x3
    .invertedResidualNB  64  384  96 1,   -- stage 5        @14
    .invertedResidualNB  96  576  96 1,   --   x2
    .invertedResidualNB  96  576 160 2,   -- stage 6      14->7
    .invertedResidualNB 160  960 160 1,   --   x2
    .invertedResidualNB 160  960 320 1,   -- stage 7         @7
    .convBnNB 320 1280 1 1,               -- head, 1x1 to 1280
    .globalAvgPool,
    .dense 1280 10 ]
  bnChannels := #[32, 32,16, 96,96,24, ... ]  -- all 52, in order

-- 2. The schedule. Identical to ResNet-34's, deliberately.
def mobilenetv2AdamConfig : VerifiedConfig where
  epochs    := 80
  batchSize := 32

-- 3. The entry point: AdamW, cosine with a 3-epoch warmup, baseLR 1e-3.
--    LEAN_MLIR_VARIANT=rms selects the RMSProp render and its own rate.
def runMobilenetV2Adam (argv : List String) : IO Unit := do
  let variant := (← IO.getEnv "LEAN_MLIR_VARIANT").getD "adam"
  let sched := mnv2RmsSchedule
  let rms := variant.startsWith "rms"
  let bs := mobilenetv2AdamConfig.batchSize
  let baseLR := match (← IO.getEnv "LEAN_MLIR_BASE_LR_U").bind (·.toNat?) with
    | some u => u.toFloat * 1e-6
    | none   => if rms then sched.lr * bs.toFloat / 256.0 else 0.001
  mobilenetv2Verified.toNet.trainAdamSched mobilenetv2AdamConfig
    (argv.head?.getD "data") baseLR 0.9 0.999 3 variant
    (if rms then sched.decayRate else 0.0) sched.decayEpochs
def main (argv : List String) : IO Unit := runMobilenetV2Adam argv

Lines marked x2 or x3 appear that many times in the source. The spec has no repeat count, because one constructor is one block: .invertedResidualNB ic mid oc stride takes the expanded width mid rather than a ratio \(t\), so the widths read straight off the paper’s table and the stage boundaries are visible as the places where stride is 2. Internally each block is a \(1 \times 1\) conv to mid channels, a \(3 \times 3\) depthwise at that width, a \(1 \times 1\) project to oc, and a residual skip when ic == oc and stride == 1. Its VJP is Theorem 71 for the depthwise conv composed with the dense, BN and biPath theorems already proved.

The BatchNorm list has 52 entries, not the 53 a per-block count gives. Every inverted-residual block normalizes after each of its three convs, which would be \(17 \times 3 + 2\) for the stem and head, or 53. The first block has expansion \(t{=}1\), so it has no expand conv at all and contributes two entries rather than three, and the real count is 52. That is the kind of detail bnChannels exists to get right: the running-stat region is positional, and an off-by-one there misaligns every frozen statistic downstream of it at evaluation time.

The toSpecs arm carries the same count. .invertedResidualNB is one more entry on Chapter 1’s toSpecs, and the \(t{=}1\) carve-out above is not prose — it is the conditional:

  | invertedResidualNB ic mid oc _ =>          -- Ch. 6
    (if mid != ic then #[(#[mid,ic,1,1],0),(#[mid],1),(#[mid],2)] else #[])
    ++ #[(#[mid,1,3,3],0),(#[mid],1),(#[mid],2),      -- depthwise 3x3 + BN
         (#[oc,mid,1,1],0),(#[oc],1),(#[oc],2)]       -- project 1x1 + BN

mid != ic is exactly “\(t \neq 1\)”. When it holds the block contributes nine tensors and three BatchNorms; when it fails the expand conv is not emitted at all and the block contributes six and two. Sixteen blocks at three plus one at two, plus the stem and the head, is the 52. The depthwise kernel is [mid,1,3,3] — one filter per channel, which is the whole of Theorem 71’s subject written as a shape.

The recipe is ResNet-34’s, unchanged. Eighty epochs at batch 32, AdamW at \(10^{-3}\), cosine with a three-epoch warmup, label smoothing \(0.1\), the same augmentation. That is on purpose. Holding the recipe fixed is what makes the accuracy in §6.1 a comparison between two architectures rather than between two training setups, and it is the same move Chapter 4 used to compare optimizers against one proven gradient. MobileNetV2’s own optimizer is RMSProp, which §6.5 comes back to, and it is rendered at this shape too: LEAN_MLIR_VARIANT=rms loads mobilenetv2_rms_train_step.mlir instead.

Results

§6.1 has the run: 80 epochs at about 26 seconds each on one RTX 4060 Ti, finishing at 89.25% top-1 and 98.68% top-5 on Imagenette’s 3,925-image validation split. Against Chapter 5’s ResNet-34 at \(89.99 \pm 0.32\) (runs/2026-09-12-r34-ablation-fp32-seeds/) on the same data and the same recipe, that is 0.74 points of top-1 for 9.5\(\times \) fewer parameters (2.24M against 21.3M). That is the depthwise-separable trade, and on this evidence it is a better one than the architecture’s reputation suggests. From the log:

• 2.24M parameters, and 1.07 MB of MLIR (1 070 905 chars) for the AdamW training step. ResNet-34 emitted 754 KB with \(9.5\times \) the parameters, so the smaller network renders the larger program. Emission is per-op, and the separable block trades one convolution for three, each with its own BatchNorm. Parameter count went down and operation count went up, and the text follows the second.

• Per-step time is about 90 ms at batch 32, against ResNet-34’s 120 ms at the same batch, the same 295 steps to the epoch and the same resident parameters, so the whole run takes thirty-five minutes rather than forty-seven. That is a real \(1.3\times \), and it is far short of what the parameter count promises: \(9.5\times \) fewer weights buys \(1.3\times \) the throughput, not \(9.5\times \). Depthwise convolution is memory-bound. It does nine multiply-accumulates per output element where the standard \(3 \times 3\) it replaces does \(9 C_\text {in}\), so the kernel spends its time moving activations rather than multiplying them, and arithmetic was never the whole bottleneck. §6.5 sees the same sub-proportional scaling at ImageNet.

• The net renders as three artifacts. The train step is accompanied by mobilenetv2_fwd (152 793 chars) and mobilenetv2_fwd_eval (120 698 chars), the second existing only because BatchNorm evaluates differently than it trains. It is the concrete cost of the bnChannels line in the spec above.

• 52 BN layers, up from ResNet-34’s 36, for the same reason the MLIR is bigger: more internal convolutions means more BN-after-conv pairs. Each one still proves its VJP via Theorem 49.

• Final loss plateaus near 0.50 again. The label-smoothing floor at \(\varepsilon = 0.1\) over ten classes is about \(0.485\), so this run has essentially converged rather than stalled.

6.4 MLIR: Depthwise Convolution

What is already proven. A depthwise convolution applies one \(k\times k\) kernel per channel with no cross-channel sum, so its reverse-mode derivative (Theorem 71, depthwiseHasVJP3) is a reversed-kernel convolution minus the channel transpose the full-conv VJP (Theorem 34) needs: with no \(\sum _{c_\text {in}}\) to take the adjoint of, there is no in/out channel axis to swap. The inverted-residual block composes this with the pointwise convs, BatchNorm, and relu6, and mobilenetv2ForwardBFullHasVJPAt chains the whole network through vjpCompAt.

The gap and how we close it. The emitted graph is denoted and shown equal to depthwiseHasVJP3’s backward. Here is what the printer emits for the depthwise input gradient (two channels, \(3\times 3\)):

func.func @dw_back(%dy: tensor<1x2x4x4xf32>, %W: tensor<2x3x3xf32>)
    -> tensor<1x2x4x4xf32> {
  %We = stablehlo.reshape %W
          : (tensor<2x3x3xf32>) -> tensor<2x1x3x3xf32>
  %Wr = stablehlo.reverse %We, dims = [2, 3] : tensor<2x1x3x3xf32>
  %dx = stablehlo.convolution(%dy, %Wr)
      dim_numbers = [b, f, 0, 1]x[o, i, 0, 1]->[b, f, 0, 1],
      window = {stride = [1, 1], pad = [[1, 1], [1, 1]],
                lhs_dilate = [1, 1], rhs_dilate = [1, 1]}
      {batch_group_count = 1 : i64, feature_group_count = 2 : i64}
      : (tensor<1x2x4x4xf32>, tensor<2x1x3x3xf32>)
        -> tensor<1x2x4x4xf32>
  return %dx : tensor<1x2x4x4xf32>
}

Read it against Chapter 3’s conv_back. That one began with a transpose %W, dims = [1, 0, 2, 3] to swap the in- and out-channel axes before reversing, and this one has no transpose. The reshape only adds the singleton input-channel axis the grouped form expects (\([c,k,k] \to [c,1,k,k]\)), the reverse flips the two spatial axes, and feature_group_count = 2, equal to the channel count, tells the convolution to run each channel through its own filter with no mixing. That single attribute is the depthwise structure, and the bridge theorem is the claim that this graph computes depthwiseHasVJP3’s backward.

Caveats.

• The depthwise convolution bridge is unconditional, because it is linear.

• relu6 is a smooth-point bridge. MobileNetV2’s clamp to \([0,6]\) has no derivative where a pre-activation sits exactly at \(0\) or \(6\), and the bridge may fail only on that measure-zero set.

• Representative scale (two channels at \(4\times 4\)).

6.5 ImageNet recipe

The Imagenette run above trained on \(\sim \)9.5K images. The same backbone scales to full 1000-class ImageNet exactly the way ResNet-34 did in Chapter 5: the inverted-residual stack is unchanged, only the head widens to 1000 and the dataset and schedule grow. As there, the recipe runs on both paths — the reference first, then the verified path, then the two side by side.

The reference (JAX)

The reference trainer is jax/MainMobilenetV2Imagenet.lean; its spec, and its recipe as the run below trained it (the file today carries TF-slim’s README command, the marker’s rows, for the rerun):

-- Same inverted-residual backbone, 1000 classes instead of 10.
def mobilenetV2Imagenet : NetSpec where
  name   := "MobileNetV2 (ImageNet, bf16)"
  imageH := 224
  imageW := 224
  convBnAct := .relu6               -- ReLU6 throughout, stem and 1x1 head included
  layers := [
    .convBn 3 32 3 2 .same,
    .invertedResidual  32  16 1 1 1,
    .invertedResidual  16  24 6 2 2,
    .invertedResidual  24  32 6 2 3,
    .invertedResidual  32  64 6 2 4,
    .invertedResidual  64  96 6 1 3,
    .invertedResidual  96 160 6 2 3,
    .invertedResidual 160 320 6 1 1,
    .convBn 320 1280 1 1 .same,
    .globalAvgPool,
    .dense 1280 1000 .identity      -- 1000-class head
  ]

-- RMSProp recipe in TF-slim's form: staircase decay from step 0, no warmup,
-- no decay on BN or biases, BN decay 0.997 / eps 1e-3, bf16 conv on.
def mobilenetV2ImagenetConfig : TrainConfig where
  learningRate   := 0.045           -- MobileNetV2-native RMSProp peak
  batchSize      := 256
  epochs         := 90              -- near-paper run
  optimizer      := .rmsprop        -- paper: RMSProp + momentum, not SGD
  momentum       := 0.9             -- mu for the RMSProp momentum buffer
  rmspropDecay   := 0.9             -- rho, the running mean-square decay
  rmspropEps     := 1.0             -- MobileNetV2 uses epsilon = 1.0
  weightDecay    := 4e-5            -- smaller than R34's 1e-4
  wdExcludeNormBias := true         -- no decay on BN gamma/beta or biases (slim's rule)
  cosineDecay      := false         -- replaced by the paper exp-decay
  expLRDecayRate   := 0.98          -- MobileNetV2: x0.98 per epoch
  expLRDecayEpochs := 1.0
  expLRStaircase   := true          -- the paper's staircase, counted from step 0
  dropout          := 0.2           -- MobileNetV2 classifier dropout
  warmupEpochs   := 0               -- the paper has no warmup
  augment        := true            -- random-crop + horizontal flip
  useAutoAugment := false           -- MNv2 paper used crop/flip only, no AA
  labelSmoothing := 0.0             -- the paper names none; TF-slim's default
  bf16           := true
  bf16Conv       := true            -- reaches the inverted-residual blocks
  f32StemHead    := true            -- classifier head in fp32
  runningBN      := true            -- paper-faithful eval: running BN stats
  bnMomentum     := 0.997           -- TF-slim's BN decay (PyTorch momentum 0.003)
  bnEps          := 1e-3            -- TF-slim's BN epsilon

-- The 350-epoch run below is the same recipe, longer.
def mobilenetV2ImagenetConfigFull : TrainConfig :=
  { mobilenetV2ImagenetConfig with epochs := 350 }

The run below was mobilenetV2ImagenetConfigFull as it then stood, launched by scripts/jobs/mnv2-full-jax-4gpu.conf; name and conf now carry the README command at \(420\) epochs. On either path, three rows of the paper’s TF-slim recipe are not reproduced. Sandler et al. trained asynchronously on sixteen workers of \(96\) images each, and this run is four replicas of \(64\), synchronous, at the paper’s peak of \(0.045\). The depthwise kernels are weight-decayed, where slim exempts them. And the learning rate multiplies the momentum buffer on its way into the weights, where TensorFlow’s RMSProp folds it into the buffer; the two agree while the rate holds still and part by a little at each step of the staircase.

The one knob specific to MobileNet is bf16Conv. ResNet’s convolutions cast to bfloat16 the instant the flag is on, but MobileNet’s compute lives inside inverted-residual blocks (a \(1\times 1\) expansion, a depthwise \(3\times 3\), a \(1\times 1\) projection) whose convolutions originally ran in fp32 regardless. Routing them through the same convdt cast as the plain convolutions is what lets bf16 reach the bulk of the network, and the whole block runs \(\sim 2\times \) faster on cuDNN for it: the \(1\times 1\)s are really matmuls and love bfloat16, and the depthwise \(3\times 3\) is a wash. The weight decay also drops to \(4\times 10^{-5}\), because MobileNet’s depthwise filters are tiny at nine weights per channel, and the \(10^{-4}\) that suits ResNet over-regularizes them.

Why 350 epochs. The paper fixes the optimizer, the peak, the batch, the weight decay and the decay rate — \(\times 0.98\) per epoch — and leaves the length of the schedule unstated. Given the rest, the length is not free: under \(\times 0.98\)/epoch the learning rate after \(90\) epochs is still \(0.045 \times 0.98^{85} \approx 8.1\times 10^{-3}\), barely annealed, with none of the end-of-schedule polish a cosine run has already collected by then, and only past \(300\) or so does it reach the \(10^{-5}\) range where the run stops moving (\(0.045 \times 0.98^{345} \approx 4.2\times 10^{-5}\) at \(350\)). So \(350\) is this book’s choice, made the way the decay makes it — it is also the length Tan & Le state for EfficientNet’s version of the same RMSProp lineage in Chapter 7 — and the curve below is the check that it is long enough. A shortened run of this recipe has to substitute cosine to be worth anything, which is a different experiment rather than a cheaper version of the same one.

Global batch 256 as \(4 \times 64\) replicas, scored over all \(50{,}000\) validation images every epoch. The wall-clock win over the \(6.2\times \)-larger ResNet-34 is real but sub-proportional, as §6.1 measured at small scale: depthwise convolution is memory-bound, so an order of magnitude fewer FLOPs buys a factor of two or three on this hardware rather than six. The RMSProp peak of \(0.045\) took off cleanly, because \(\varepsilon =1.0\) keeps the adaptive step well damped — in sharp contrast to EfficientNet’s \(\varepsilon =10^{-3}\) in Chapter 7. The last \(25\) epochs average \(71.54\% \), within a tenth of a point of the endpoint. The paper’s figure is \(72.0\% \), and at \(n=50{,}000\) the Wilson 95% interval on this run is \([71.24, 72.03]\): the paper’s recipe, at a length the paper leaves open, is not distinguishable from the paper’s number, which is the claim here and a weaker one than closing a gap by a measured amount.

The run, sampled every 10 epochs. Epoch \(100\) has \(94\% \) of the final accuracy, the run is still climbing at \(200\), and the last \(50\) epochs move top-1 by \(0.14\), so the schedule is exhausted rather than cut short.

The verified path (PJRT)

The verified trainer is apps/imagenette/MainMobileNetV2Imagenet.lean, built as lake build mobilenetv2-imagenet-verified, as it read for the run below (today \(420\) epochs, the README’s):

-- Same inverted-residual backbone, 1000 classes instead of 10.
def mobilenetv2ImagenetVerified : VerifiedNetSpec where
  name       := "MobileNetV2 (ImageNet-1k)"
  slug       := "mobilenetv2in"
  data       := .imagenet
  shimScript := "generated_mobilenet_v2_imagenet_shim.py"
  layers     := [ ... the Imagenette stack, 1000-class head ... ]

-- 350 epochs at 64 PER DEVICE; four replicas give the global 256.
-- 350 because that is the schedule the reference number above was measured on.
def mobilenetv2ImagenetConfig : VerifiedConfig where
  epochs    := 350
  batchSize := 64

def main (argv : List String) : IO Unit :=
  mobilenetv2ImagenetVerified.toNet.trainAdamSched
    mobilenetv2ImagenetConfig (argv.head?.getD "data") ... variant

slug names verified_mlir/mobilenetv2in_<variant>_train_step.mlir, the rendered graph, and shimScript names this net’s own augmentation shim rather than ResNet-34’s — the shim the reference consumes, which is what makes the two a matched pair (§5.7). The head is the only parameter shape that moves, from \(1280 \times 10\) to \(1280 \times 1000\), which takes the count from 2.24M to 3,504,872 and matches the reference exactly. The BatchNorm layout is required to be byte-identical to the Imagenette net’s, and a #guard enforces it, because the running-stat region is positional and a drift there would misalign every frozen statistic at evaluation.

The variant carries the recipe. rmsdp64wxdols0eps0001bf16 is the artifact the run below used, the reference’s recipe row for row: RMSProp at \(\rho = 0.9\), \(\mu = 0.9\), \(\varepsilon = 1.0\) with coupled weight decay \(4 \times 10^{-5}\) off BatchNorm \(\gamma /\beta \) and biases, peak LR \(0.045\) with the \(\times 0.98\) staircase from step \(0\) and no warmup, classifier dropout \(0.2\), no label smoothing, BatchNorm at TF-slim’s \(0.997\) / \(10^{-3}\) with the running statistics read at that \(\varepsilon \) by its eval graph mobilenetv2in_fwd_eval_eps0001, bf16 convolution under an fp32 classifier head, data-parallel over four replicas of \(64\) with each of the \(52\) BatchNorms all-reducing its mean and variance across them, so the statistic group is the global \(256\). The peak and decay rate live in mnv2RmsSchedule, shared with the Imagenette peer so the two cannot carry different values; mnv2ImagenetRmsSchedule adds the staircase and drops the warmup, and the optimizer constants are mnv2RmsHyper. Theorem 84 ties this render’s bf16 gradient nodes to one device’s step at the global batch, its classifier dropout inside the statement. adam64 is AdamW at \(10^{-3}\) instead, the driver’s default because it matches the other ImageNet drivers. The job, which today launches the rerun’s render named in the marker below:

lake run mnv2-default-4gpu

Global batch 256 as \(4 \times 64\) replicas with synchronised BatchNorm statistics, scored over the same \(50{,}000\) images every epoch.

The two paths, side by side

Both columns are MobileNetV2 on the paper’s recipe at 350 epochs, on the same four RTX 3060 cards, scored over the same \(50{,}000\) validation images. Rows above the first rule differ; rows between the rules are held equal by construction, the BatchNorm group included: the reference normalizes over the global \(256\) under its mesh, and the verified render all-reduces to the same \(256\). [TODO: TF-slim’s own reproduction command (research/slim/nets/mobilenet/README.md), on both paths, in code since 2026-10-12 and not yet run: \(0.045\) per \(96\) images, so \(0.12\) at \(256\), stepped \(\times 0.98\) every \(2.5\) epochs rather than every epoch; RMSProp’s learning rate inside the momentum term; weight decay off the depthwise kernels; slim’s truncated-normal initialization (\(\sigma = 0.09\)) on every kernel; label smoothing \(0.1\); the weight average at \(0.9999\), scored; slim’s colour distortion (brightness, then saturation), its crop ranges, TF 1’s bilinear kernel on every resize, its validation crop (\(0.875\) of both sides, then \(224^2\)) and \((x-0.5)\times 2\) inputs; about \(420\) epochs; and the depthwise fan-in init (\(2/k^2\)) on the verified path, the suspect for the \(-0.51\) above. One rerun of both halves takes all of it — mnv2-full-jax-4gpu / mnv2-default-4gpu, \({\sim }49 + 63\) h. The sixteen asynchronous workers of \(96\) stay four replicas of \(64\).]

The two paths differ by \(0.51\) points on top-1 and \(0.42\) on top-5, outside the 95% Wilson interval (\(\pm 0.40\) at \(71\% \) on \(50{,}000\) images) at one seed each, and the difference is made in training, not at evaluation: the verified path’s training loss runs below the reference’s for the first \(25\) epochs and \(0.02\) above it from epoch \(26\) to the end, while the sampled accuracies stay level through epoch \(100\) and part by half a point from \(150\) on. Every row between the rules is held equal by construction. Two inputs to training are not shared: the classifier head’s precision, and the initialization, where every depthwise kernel starts \(\sqrt{C}\) narrower on the verified path, He fan-out over \(C \cdot k^2\) against the reference’s fan over \(k^2\). Those rows are open; a run with them matched is what closes them.

One recipe, two paths to XLA, sampled every 10 epochs and scored over all \(50{,}000\) validation images.

What the MobileNetV2 recipe changes from the 2018 recipe

The ResNet-34 trainer of Chapter 5 is the 2018-era recipe of the SGD lineage: momentum, cosine polish, label smoothing. MobileNetV2 is a 2018 paper too, but from the other side of a fork. Google’s TensorFlow and Inception training tradition runs RMSProp with a huge \(\varepsilon \), an exponential LR staircase, a tiny weight decay, and nothing beyond crop and flip on the data side. So this diff is not a time jump but two contemporary recipes side by side, and it is the base that the next chapter’s 2019 EfficientNet recipe is built on. The entry point is the same on both paths, and every row below is one argument or one field:

What carries over unchanged: batch 256, train and eval both at \(224\), running-BN evaluation, and bf16 convolution under an fp32 classifier head. R34’s 5-epoch warmup does not: MobileNetV2’s staircase starts at step \(0\), as the paper’s does.

The through-line of this fork is that nothing here regularizes harder. MobileNetV2’s recipe is the SGD recipe’s contemporary, tuned for a network \(1/6\) the size, and its distinctive knobs (adaptive optimizer, staircase decay, tiny selective-in-spirit weight decay) are the ones EfficientNet inherits and ConvNeXt eventually replaces wholesale. The 2019–2022 tables that follow all diff against the same 2018 SGD column. This one shows where their other half came from.

The general pattern is deeper and narrower, with a per-channel conv between surrounding \(1 \times 1\) expansions, and it carries forward to MobileNet V3 and the EfficientNet family. All of those add their own small modifications on top (SE blocks, Swish activations, compound scaling) but the core depthwise-separable scaffolding is exactly what this chapter formalized.

6.6 Side quest: MobileNet V4

Six years on, MobileNetV4 (Qin et al. 2024, arXiv:2404.10518) revisits this chapter’s block. Its central move is the universal inverted bottleneck (UIB): take the inverted residual and make two depthwise positions optional, one before the expansion and one after it. Four instantiations fall out of that one template. Both depthwise convs on is ExtraDW. Only the leading one is ConvNeXt, this chapter’s sandwich rearranged into Chapter 8’s block shape. Neither is a plain FFN of two \(1\times 1\)s. Only the inner one is exactly MobileNetV2’s block. The jax/ folder ships the faithful Conv-Medium as a full ImageNet trainer (jax/MainMobilenetV4Imagenet.lean), transcribed block for block from timm’s mobilenetv4_conv_medium and tied to it on shared weights (scripts/parity/mnv4_timm_parity.py, to \(\sim \! 10^{-5}\) in the training forward): \(9{,}715{,}512\) parameters, a fused-IB stage at high resolution, then 21 UIB blocks whose per-block instantiation is visible right in the NetSpec, driven by the same entry point as every other net in this book. What changed from 2018 to 2024 is at least as much recipe as architecture:

The direction of travel matches the recipe tables of Chapters 7 and 8: by 2024 the mobile lineage has converged on the same transformer-era training recipe as everything else (AdamW at large effective batch, cosine, heavy selective decay, RandAugment, EMA) while the block template generalized enough to absorb its rivals.

One trap in that EMA row. Seeding the shadow at the initialization and decaying it away leaves \(d^{\, t}\) of that initialization behind. A random network and a trained one are not joined by a low-loss path, so even a few percent of it puts the average at chance rather than merely behind. Conv-M trips this where longer runs do not, because gradient accumulation leaves it only \(312\) steps per epoch, so \(31{,}200\) total against a time constant of \(10{,}000\). The fix is TensorFlow’s, which this family’s reference implementations assume: ramp the decay, \(d_t = \min \! \big(d,\ (1+t)/(10+t)\big)\).

mobilenet-v4-imagenet ships the paper recipe as full, 500 epochs, and a 100-epoch tier as default: wd \(0.05\) against the paper’s \(0.1\), dropout \(0.1\) against \(0.2\), RandAugment \(m9\) at \(p\, 0.5\) against \(m15\) at \(p\, 0.7\), no drop-path, and \(224\) throughout, because the paper’s full regularization pack underfits a short schedule. The 100-epoch tier is the one run here, on both paths, and the \(79.9\% \) headline is a 500-epoch number.

The reference (JAX)

The reference trainer is jax/MainMobilenetV4Imagenet.lean, built as lake build mobilenet-v4-imagenet, with .imagenet streamed through tfds; its job config is scripts/jobs/mnv4-default-jax-4gpu.conf. AdamW at peak \(0.004\) for an effective batch of \(4096\), reached as a \(512\)-image micro-batch accumulated eight times, so \(312\) optimizer steps an epoch; bf16 convolutions against fp32 master weights, the classifier head in fp32; the EMA shadow is what gets scored, every five epochs and at the end.

Micro-batch \(512\) as \(4 \times 128\) replicas, accumulated \(8\times \) to the effective \(4096\); EMA weights scored over all \(50{,}000\) validation images. Set against this chapter’s other half that is the whole six-year argument in two rows: \(\mathbf{+4.94}\) points over MobileNetV2’s paper-faithful \(71.63\% \), from \(2.8\times \) the parameters and under a third of the epochs, \(100\) against \(350\) and \(16.2\) wall-clock hours against \(38.8\) on the same four cards. By 2024 the mobile line is not trading accuracy for size so much as reaching the same accuracy with a block that happens to be cheap.

The run, scored every five epochs. The last ten epochs move top-1 by \(0.09\) (\(76.49 \to 76.51 \to 76.57\)): the schedule is exhausted, and the points between this and \(79.9\% \) live in the regularization and the length the full tier restores.

The verified path (PJRT)

The spec follows §6.5 and §5.7. The trunk is copied unchanged from the Imagenette spec, a #guard pins every parameter shape but the head against it, and the head moves \(1280{\times }10 \to 1280{\times }1000\), taking the count from \(8{,}447{,}322\) to \(\mathbf{9{,}715{,}512}\), the reference’s count exactly:

def mnv4ImagenetVerified : VerifiedNetSpec where
  name       := "MobileNetV4-Conv-M (ImageNet-1k)"
  slug       := "mnv4in"
  imageH     := 224
  imageW     := 224
  data       := .imagenet
  shimScript := "generated_mobilenet_v4_imagenet_shim.py"
  layers     := [
    .convBnNB 3 32 3 2,               -- stem                       224->112
    .fusedMbConvNB 32 48 4 3 2,       -- stage 0, fused IB          112->56
    .uib  48  80 4 2 3 5,             -- ... the Imagenette UIB stack, 21 blocks
    .convBnNB 256 960 1 1,
    .globalAvgPool,                   -- timm pools BEFORE conv_head
    .convBnNB 960 1280 1 1,
    .dense 1280 1000 ]                -- 1000-class head
  runningBN   := true
  dropoutKeep := some (0.9, 1280)     -- classifier dropout 0.1, read by the `do` variants

-- The head is the ONLY thing that may differ from the 10-class net.
#guard mnv4ImagenetVerified.toSpecs.size == mobilenetv4Verified.toSpecs.size
#guard mnv4ImagenetVerified.toSpecs.pop.pop == mobilenetv4Verified.toSpecs.pop.pop
#guard mnv4ImagenetVerified.toSpecs.back!  == (#[1000], 2)
#guard mnv4ImagenetVerified.bnChannels == mobilenetv4Verified.bnChannels

That is apps/imagenette/MainMobilenetV4Imagenet.lean. It builds as mobilenetv4-imagenet-verified. The renders come from the same chain and the same block table as the three 10-class artifacts, differing only in the head, the batch and the slug, which is why adding them was three #eval lines rather than a new proof chain. The inference graphs, mnv4in_fwd_eval and mnv4in_fwd_eval_s256, denote the frozen-statistics forward at any final feature size (mnv4FwdGraphBFullEval_faithful), and their block text is pretty of the typed graph at both sizes (FwdGraphTextTies); the classifier-dropout forward is stated at every mask and at either precision (mnv4FwdGraphBFullDo_faithful).

The variant carries the recipe. emaaccdp8x128wxdowd005bf16 is the artifact the run below used: AdamW at \(0.004\), accumulation \(k = 8\) over four replicas of \(128\) to the reference’s effective \(4096\), weight decay \(0.05\) with BN \(\gamma /\beta \) and biases excluded, classifier dropout \(0.1\), the EMA shadow at \(0.9999\) with the warmup correction above, bf16 convolution under an fp32 classifier head. Each of the 77 BatchNorms all-reduces its mean and variance across the four replicas, so the statistic group is the \(512\)-image micro-batch, which is the group the reference’s accumulation gives it. Theorem 85 ties the data-parallel render’s bf16 gradient nodes to one device’s step at the global batch; this render’s EMA, gradient accumulation and classifier dropout sit outside that statement. It runs with:

lake run mnv4-default-4gpu

Micro-batch \(512\) as \(4 \times 128\) replicas, accumulated \(8\times \) as on the reference; EMA weights scored over the same \(50{,}000\) images.

The two paths, side by side

Both columns are MobileNetV4-Conv-M on the 100-epoch tier, scored over the same \(50{,}000\) validation images under the same protocol. Rows above the first rule differ; rows between the rules are held equal by construction, the BatchNorm group included: the reference normalizes over the \(512\)-image shard its accumulation hands each step, and the verified render all-reduces to the same \(512\). [TODO: the paper’s tier — its own training config (tensorflow/models, imagenet_mobilenetv4_conv_medium_tpu.yaml: \(500\) epochs at \(256^2\), wd \(0.1\), dropout \(0.2\), RandAugment \(m15\) at \(p\, 0.7\) under TensorFlow’s mappings, drop-path \(0.075\) over its \(25\)-entry block list, BN \(0.997\) / \(\varepsilon \, 10^{-3}\), AdamW \(\varepsilon \, 10^{-7}\), the cosine from step 0, its crop and resize) — exists on both paths and has not run: mnv4-full-jax-4gpu / mnv4-full-4gpu, the mnv4in256 renders. Disclosed: the init and the head order are timm’s, the decay reaches the classifier bias. Scored at the paper’s \(256\) / \(0.875\) and at timm’s \(320\) / \(1.0\) for its e500_r256 weights; the \(224\) tier above at timm’s \(256\) / \(1.0\).]

The two paths agree to \(0.11\) points on top-1 and \(0.16\) on top-5, inside the 95% Wilson interval (\(\pm 0.37\) at \(77\% \) on \(50{,}000\) images), at one seed each. The BatchNorm group is held equal. Above the rule, beside the lowering, sit the precision of the classifier head and of the evaluation forward, and the depthwise initialization, which starts every depthwise kernel \(\sqrt{C}\) narrower on the verified path.

The curves converge; they do not merely end together. The verified path trails at epoch 5 and leads from epoch 10, by \(4.40\) at the widest, and the lead narrows through the anneal — not monotonically, it widens again at epochs \(25\), \(60\) and \(70\) — to \(0.11\). Over the last three scored epochs the two average within \(0.15\) of each other on top-1 and \(0.18\) on top-5:

One recipe, two paths to XLA, scored every five epochs over all \(50{,}000\) validation images.

The tier above is detuned by construction, at wd \(0.05\) against the paper’s \(0.1\), dropout \(0.1\) against \(0.2\), RandAugment \(m9\) at \(p\, 0.5\) against \(m15\) at \(p\, 0.7\), no drop-path against \(0.075\), \(224\) against \(256\) and \(100\) epochs against \(500\), so its \(76.57\% \) and \(76.68\% \) are not a reproduction of \(79.9\% \). Both curves are flat over their last ten epochs; the remaining points live in the regularization and the length, and the paper’s tier is the marker above. Both paths above score the EMA shadow, which the paper’s Conv-M recipe does not carry; the tie stands, since both columns carry it, and the full pair drops it, scoring the live weights as the paper does.