Verified Deep Learning with Lean 4

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:

\[ y_{c_\text {out},\, i,\, j} \; =\; \sum _{c_\text {in},\, k_h,\, k_w} W_{c_\text {out},\, c_\text {in},\, k_h,\, k_w} \cdot x_{c_\text {in},\, i+k_h,\, j+k_w}. \]

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, then V2 in 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 the same GPU time Chapter 5 asked for:

lake exe cache get                            # Mathlib oleans, ~30 s
./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 finished at \(89.71\% \) top-1 and \(98.27\% \) top-5. MobileNetV2 gives up less than half a point of top-1, picks up four tenths of top-5, and does it in thirty-five minutes against an hour and a half.

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 39 proves, every pointwise convolution by Chapter 3’s, every BatchNorm by Chapter 4’s, and every residual skip by Chapter 5’s additive fan-in. 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 is mobilenetv2_full_has_vjp_at, in MobileNetV2FullVJP.lean, and it covers the stem, all seventeen bottlenecks and the head. What ties the full 17-block spec to the paper’s architecture is a separate denotational result about the forward, mobilenetv2Verified_denote_eq, with mobilenetv2Verified_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 saying that the value entering relu6 is neither \(0\) nor \(6\). There are thirty-five such sites, 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 with efficientnetForwardB_full_has_vjp, 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 the same three Chapter 5 showed, and for the same reason. @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.

To run this net alongside the other six at the same scale and recipe:

lake run imagenette

That is the whole Part-I tier at \(224^2\), and it is an afternoon rather than the minute lake run mnist takes.

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:

\[ y_{c,\, i,\, j} \; =\; \sum _{k_h,\, k_w} W_{c,\, k_h,\, k_w} \cdot x_{c,\, i+k_h,\, j+k_w}. \]

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]\):

\[ \text{depthwise weights: } C \times 3 \times 3 \quad \text{vs}\quad \text{standard weights: } C_\text {out} \times C_\text {in} \times 3 \times 3. \]

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 39 through 41 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,

\[ \text{depthwise}\ 3 \times 3 \quad \longrightarrow \quad \text{pointwise}\ 1 \times 1, \]

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 standard conv2d VJPs from Chapter 3, the BN VJP from Chapter 4, and the additive fan-in for the residual from Chapter 5. 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.71\% \) on Imagenette, both measured in this book). For mobile and embedded deployment that trade is the whole reason MobileNet exists.

6.2 The theorems

Definition 38 Depthwise conv forward
#

Concrete per-channel cross-correlation.

Theorem 39 Depthwise input VJP

prove: \(\mathsf{HasVJP3}\, \bigl(\mathrm{depthwiseConv2d}(W, b)\bigr)\), with the per-channel reversed-kernel backward \(dx_{c, h, w} = \sum _{k_h, k_w} W_{c,\, k_H - 1 - k_h,\, k_W - 1 - k_w} \cdot dy_{c,\, h + k_h - p,\, w + k_w - p}\).

Proof

Sketch: the conv2d input VJP (Theorem 25) with one fewer \(\Sigma \) level — depthwise has no cross-channel mixing, so the channel of every \(v\)-read is forced to equal the output channel.

  1. Define \(B(x, dy)_{c, h_i, w_i} := \sum _{h_o, w_o} [\text{pad-valid}]\; W_{c, \hat{k}_h, \hat{k}_w} \cdot dy_{c, h_o, w_o}\) with reconstructed offsets \(\hat{k}_h = h_i + p_H - h_o\), \(\hat{k}_w = w_i + p_W - w_o\) — same shape as the conv2d formula, minus the channel sum.

  2. suffices: \(B\) matches the \(\operatorname {pdiv}_3\) contraction at every index.
    proof: Definition 24.

  3. The flattened forward is affine in \(v\): constant bias plus \(\sum _{k_h, k_w} W_{o(\mathrm{idx}), k_h, k_w} \cdot (\text{pad-conditional projection of } v)\), where the projection reads channel \(o(\mathrm{idx})\) itself.
    proof: Unfold (Definition 38); steps and rules exactly as in Theorem 25 step 4 (Theorems 3, 6, 8 \(\times 2\), 4, 7), with a double instead of triple sum.

  4. Collapse: the channel component of the indicator forces \(c_o = c_i\) (channel mismatch zeroes every other term); per \((h_o, w_o)\) the remaining two-conjunct indicator (\(k_h + h_o = h_i + p_H\), \(k_w + w_o = w_i + p_W\)) collapses the \((k_h, k_w)\) sum.
    proof: Finset.sum_eq_single on the decoded components.

  5. q.e.d.
    proof: What remains is \(B(x, dy)_{c_i, h_i, w_i}\) (equivalent to the reversed-kernel form under \(k_h \leftrightarrow k_H - 1 - k_h\)); with 2, done.

Theorem 40 Depthwise weight VJP

DepthwiseKernel is definitionally \(\mathsf{Tensor3}\), so \(\mathsf{HasVJP3}\) applies to the kernel directly — no flatten bijection needed, unlike the conv2d weight VJP (Theorem 26). prove: \(\mathsf{HasVJP3}\, \bigl(W \mapsto \mathrm{depthwiseConv2d}(W, b)\, x\bigr)\), with the per-channel transpose-trick backward \(dW_{c, k_h, k_w} = \sum _{h_o, w_o} [\text{pad-valid}]\; x_{c,\, k_h + h_o - p_H,\, k_w + w_o - p_W} \cdot dy_{c, h_o, w_o}\).

Proof
  1. Define \(B(W, dy)_{c, k_h, k_w}\) as the formula above.

  2. suffices: \(B\) matches the \(\operatorname {pdiv}_3\) contraction at every kernel index.
    proof: Definition 24, applicable since the kernel is a \(\mathsf{Tensor3}\).

  3. Per-\((c_o, h_o, w_o)\) Jacobian: \(\operatorname {pdiv}_3 = [c_o = c]\, \cdot (\text{pad-conditional } x \text{ term})\) — a single channel equality where the conv2d weight VJP had a packed \(\varphi \)-comparison.
    proof: At output \((c_o, h_o, w_o)\) the forward is \(b_{c_o} + \sum _{k_h, k_w} W_{c_o, k_h, k_w} \cdot x\text{-pad-term}\): affine decomposition and the same foundation-rule sequence as Theorem 26 step 4 (sum + constant + finite-sum \(\times 2\) + product + reindex rules).

  4. q.e.d.
    proof: Contract 3 with \(dy\): collapse the channel sum at \(c_o = c\) (Finset.sum_eq_single); the spatial sums survive, leaving exactly \(B\); with 2, done.

Theorem 41 Depthwise bias VJP

prove: \(\mathsf{HasVJP}\, \bigl(b \mapsto \mathrm{flatten} (\mathrm{depthwiseConv2d}(W, b)\, x)\bigr)\), with backward \(db_c = \sum _{h_i, w_i} dy_{\mathrm{flat}(c, h_i, w_i)}\).

Proof

Same three steps as the conv2d bias VJP (Theorem 27), with input channel \(=\) output channel throughout:

  1. As a function of \(b\), the flattened forward is \((\text{channel reindex of } b) + (\text{$W\! ,x$ term constant in } b)\), so \(\operatorname {pdiv}= \delta _{c,\, \mathrm{chan}(\mathrm{idx})}\).
    proof: Sum rule (Theorem 3), reindex Jacobian (Theorem 7), constant rule (Theorem 6).

  2. suffices: \(db_c = \sum _{\mathrm{idx}} \operatorname {pdiv}\cdot dy_{\mathrm{idx}}\).
    proof: Definition 9.

  3. q.e.d.
    proof: Contract 1 with \(dy\); the Kronecker keeps channel \(c\)’s spatial positions (Finset.sum_eq_single).

6.3 Example: MobileNet V2 on Imagenette

Depthwise convolution on its own is rarely used directly. You almost always see it wrapped in a depthwise-separable sandwich: \(1 \times 1\) expand, \(3 \times 3\) depthwise, \(1 \times 1\) project, with a residual connection around the whole thing. That sandwich is the inverted-residual block that defines MobileNet V2, and it is how depthwise convolutions got their fame.

The point of depthwise is parameter efficiency. A normal \(3 \times 3\) conv over 64 channels uses \(3 \times 3 \times 64 \times 64 = 36864\) weights. A depthwise version uses \(3 \times 3 \times 64 = 576\) weights. 64\(\times \) fewer. The depthwise-separable block recovers the missing cross-channel expressivity with the surrounding \(1 \times 1\) convs, at a combined cost far below a full \(3 \times 3\). MobileNet V2 is 17 of these blocks together.

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.

\begin{tikzpicture} [
  >={Stealth[length=1.8mm]},
  every node/.style={font=\sffamily\scriptsize},
  col/.style    = {align=center, rounded corners=2pt, inner sep=2pt, minimum height=0.56cm, minimum width=6.3cm},
  io/.style     = {col, draw=blue!55!black,   fill=blue!8},
  convbn/.style = {col, draw=orange!65!black, fill=orange!12},
  invres/.style = {col, draw=teal!65!black,   fill=teal!10, minimum height=0.6cm},
  gap/.style    = {col, draw=purple!60!black, fill=purple!8},
  head/.style   = {col, draw=red!60!black,    fill=red!8},
  logits/.style = {col, draw=green!50!black,  fill=green!14, very thick},
  arr/.style    = {->, thick, gray!60, shorten >=1pt, shorten <=1pt},
  stage/.style  = {font=\sffamily\scriptsize\itshape, gray!55!black, anchor=west},
]
  \node[io]                            (input) {Input \;\; $224\times224\times3$};
  \node[convbn, below=0.16cm of input] (stem) {\textbf{ConvBN} stem $3\to32$, $3\times3$, /2 \;\; $112{\times}112$};
  \node[invres, below=0.16cm of stem]  (s1)   {$1\times$ \textbf{InvRes} $32\to16$, $t{=}1$ \;\; $112{\times}112$};
  \node[invres, below=0.16cm of s1]    (s2)   {$2\times$ \textbf{InvRes} $16\to24$, $t{=}6$, /2 \;\; $56{\times}56$};
  \node[invres, below=0.16cm of s2]    (s3)   {$3\times$ \textbf{InvRes} $24\to32$, $t{=}6$, /2 \;\; $28{\times}28$};
  \node[invres, below=0.16cm of s3]    (s4)   {$4\times$ \textbf{InvRes} $32\to64$, $t{=}6$, /2 \;\; $14{\times}14$};
  \node[invres, below=0.16cm of s4]    (s5)   {$3\times$ \textbf{InvRes} $64\to96$, $t{=}6$ \;\; $14{\times}14$};
  \node[invres, below=0.16cm of s5]    (s6)   {$3\times$ \textbf{InvRes} $96\to160$, $t{=}6$, /2 \;\; $7{\times}7$};
  \node[invres, below=0.16cm of s6]    (s7)   {$1\times$ \textbf{InvRes} $160\to320$, $t{=}6$ \;\; $7{\times}7$};
  \node[convbn, below=0.16cm of s7]    (pw)   {\textbf{ConvBN} $1\times1$ $320\to1280$ \;\; $7{\times}7$};
  \node[gap,    below=0.16cm of pw]    (g)    {global avg pool \;\; $7\times7\times1280 \to 1280$};
  \node[head,   below=0.16cm of g]     (d)    {\textbf{Dense} $1280\to10$ \;(identity)};
  \node[logits, below=0.16cm of d]     (out)  {Logits \;\; 10 classes, softmax-CE};
  \foreach \a/\b in {input/stem, stem/s1, s1/s2, s2/s3, s3/s4, s4/s5, s5/s6,
                     s6/s7, s7/pw, pw/g, g/d, d/out}
     \draw[arr] (\a) -- (\b);
  \foreach \n/\k in {s1/1, s2/2, s3/3, s4/4, s5/5, s6/6, s7/7}
     \node[stage] at ($(\n.east) + (0.30,0)$) {stage \k};
\end{tikzpicture}
\begin{tikzpicture} [
  >={Stealth[length=1.6mm]},
  every node/.style={font=\sffamily\scriptsize},
  box/.style  = {align=center, rounded corners=2pt, inner sep=3pt, minimum height=0.62cm, draw=teal!70!black, fill=teal!10},
  dot/.style  = {circle, draw=teal!70!black, fill=teal!14, inner sep=0pt, minimum size=0.42cm},
  term/.style = {font=\sffamily\scriptsize\itshape, inner sep=1pt},
  arr/.style  = {->, thick, gray!65, shorten >=1pt, shorten <=1pt},
  skip/.style = {->, thick, teal!70!black, shorten >=1pt},
]
  \node[term]                       (x)   {$x$};
  \node[box, right=0.5cm of x]      (e)   {$1\times1$ expand\\$C\!\to\!tC$, BN, ReLU6};
  \node[box, right=0.4cm of e]      (dw)  {$3\times3$ DW\\BN, ReLU6};
  \node[box, right=0.4cm of dw]     (p)   {$1\times1$ project\\BN (linear)};
  \node[dot, right=0.5cm of p]      (sum) {$+$};
  \node[term, right=0.5cm of sum]   (y)   {$y$};
  \draw[arr] (x)  -- (e);
  \draw[arr] (e)  -- (dw);
  \draw[arr] (dw) -- (p);
  \draw[arr] (p)  -- (sum);
  \draw[arr] (sum) -- (y);
  \draw[skip] (x.north) .. controls +(0,0.85) and +(0,0.85) .. (sum.north)
        node[midway, above, term, teal!70!black] {skip if stride 1, $C_{\text{in}}{=}C_{\text{out}}$};
  \node[term, below=0.16cm of dw, gray!55!black] {one inverted-residual block: depthwise spatial mix at the wide point};
\end{tikzpicture}

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

-- 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
  nClasses := 10
  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.
def main (argv : List String) : IO Unit :=
  mobilenetv2Verified.toNet.trainAdamSched
    mobilenetv2AdamConfig (argv.head?.getD "data") 0.001 0.9 0.999 3 "adam"

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 § 39 for the depthwise conv composed with the dense, BN and biPath theorems already proved.

Fifty-two BatchNorms, and the list is not three per block. 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.

And the arm shows you the 52. .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 39’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.71\% \) on the same data and the same recipe, that is 0.46 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. Rather than reprint the log, here is what is worth noticing in it.

- 2.24M parameters, and 1.05 MB of MLIR (1 047 327 chars) for the AdamW training step. ResNet-34 emitted 729 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 220 ms at the same batch and the same 295 steps to the epoch, so the whole run takes thirty-five minutes rather than an hour and a half. That is a real \(2.4\times \), and it is also well short of what the parameter count promises: \(9.5\times \) fewer weights buys \(2.4\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.

- Three artifacts, not one. The train step is accompanied by mobilenetv2_fwd (143 476 chars) and mobilenetv2_fwd_eval (111 150 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 § 36.

- Final loss plateaus near 0.50 again, and for the same reason it did in Chapter 5. 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 (§ 39, depthwise_has_vjp3) is a reversed-kernel convolution minus the channel transpose the full-conv VJP of Chapter 3 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 mobilenetv2_full_has_vjp_at chains the whole network through vjp_comp_at.

The gap and how we close it. The emitted graph is denoted and shown equal to depthwise_has_vjp3’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 depthwise_has_vjp3’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

Why the phase-2 trainer? The ImageNet numbers in this book come from the phase-2 (Lean\(\to \)JAX) trainer: at this scale its job is to validate the framework’s logic end to end, and it is the reference the verified path is measured against. Phase-3 verified-IREE codegen led to phase 4 — PJRT, which is the verified path itself. Where a section has both, it reports both.

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. Everything in this section up to §6.5 was measured on the phase-2 (Lean\(\to \)JAX) trainer, whose spec and recipe are:

-- 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, MobileNet-native: smaller weight decay, bf16 conv on.
def mobilenetV2ImagenetConfig : TrainConfig where
  learningRate   := 0.045           -- MobileNetV2-native RMSProp peak (was 0.1 for SGD)
  batchSize      := 256
  epochs         := 90              -- the real run; validate at 30 first
  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
  cosineDecay      := false         -- replaced by the paper exp-decay (gap B)
  expLRDecayRate   := 0.98          -- MobileNetV2: x0.98 per epoch, after warmup
  expLRDecayEpochs := 1.0
  dropout          := 0.2           -- MobileNetV2 classifier dropout (gap C)
  warmupEpochs   := 5
  augment        := true            -- random-crop + horizontal flip
  useAutoAugment := false           -- MNv2 paper used crop/flip only, no AA
  labelSmoothing := 0.0             -- Sandler 2018 used none
  bf16           := true
  bf16Conv       := true            -- reaches the inverted-residual blocks
  runningBN      := true            -- paper-faithful eval (gap A)

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. The payoff is the whole inverted-residual block running \(\sim 2\times \) faster on cuDNN, where the \(1\times 1\)s (which are really matmuls) 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.

Compute budget. On four RTX 3060 (CUDA, bf16, batch 256, four replicas of \(64\)), steady-state throughput is \(6.6\) minutes per epoch — \(374\) s of training plus \(18\) s of validation, held to within a second an epoch across all \(350\) — so the paper’s schedule costs \(38.4\) wall-clock hours. The wall-clock win over the \(9.5\times \)-larger ResNet-34 is real but sub-proportional, exactly as §6.1 measured at small scale: depthwise convolution is memory-bound and low-arithmetic-intensity, so MobileNet’s order-of-magnitude FLOP saving buys a factor of two or three on this hardware rather than a factor of nine.

GPU

Precision

Per epoch

Epochs

Total

Val top-1

Val top-5

4\(\times \) 3060 (CUDA)

bf16

\(6.6\) min

350 (exp-decay)

\(38.4\) hr

\(\mathbf{71.90\% }\)

\(\mathbf{90.41\% }\)

The RMSProp peak LR of \(0.045\) took off cleanly with no collapse, because the large \(\varepsilon =1.0\) keeps the adaptive step well-damped (in sharp contrast to EfficientNet’s \(\varepsilon =10^{-3}\) in Chapter 7, which is far more LR-sensitive), so the \(\sim \)0.02 fallback was never needed.

Two evaluations get reported for a run like this and they are not the same measurement, so it is worth saying which is which once. The curve below is the in-loop evaluation the trainer runs at the end of an epoch, on the 49,920 validation images that survive batching, using the live weights. The headline is an offline pass over the full 50,000 images using the EMA weights. The second is the number to quote against a paper, and the first is the one to read for shape.

The paper-faithful tier: 350 epochs

The recipe above is the paper’s, and the schedule is the part of it that is easiest to get wrong. Sandler et al. use no label smoothing, an exponential decay of \(\times 0.98\) per epoch, and on the order of 350 epochs, and those three go together rather than being independently adjustable. Shortening the schedule forces the other two.

The reason is that 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. Only over \(350\) epochs does the exponential reach \(0.045 \times 0.98^{345} \approx 4.2\times 10^{-5}\) and behave as intended. So a shortened run of this recipe has to substitute cosine to be worth anything at all, which is a different experiment rather than a cheaper version of the same one.

Run end to end on four RTX 3060 it lands at \(\mathbf{71.90\% }\) top-1 and \(\mathbf{90.41\% }\) top-5 over the full 50,000-image validation split in \(38.4\) hours, one attempt, no restarts and no thermal pauses; the best epoch, \(331\), reaches \(71.99\% \). The paper’s figure is \(72.0\% \), and at \(n=50{,}000\) a Wilson 95% interval on this run is \([71.50, 72.29]\) — so the recipe is not distinguishable from the paper, which is the claim here and a weaker one than closing a gap by a measured amount:

\begin{tikzpicture} 
\begin{axis}[
    width=0.92\linewidth, height=6.5cm,
    xlabel={Epoch}, ylabel={Validation accuracy (\%)},
    xmin=0, xmax=355, ymin=0, ymax=95,
    xtick={0,50,100,150,200,250,300,350}, ytick={20,40,60,80},
    legend pos=south east, legend cell align={left},
    grid=major, grid style={gray!18},
    tick label style={font=\small}, label style={font=\small},
    every axis plot/.append style={line width=1pt, mark size=0.7pt},
]
\addplot[blue, mark=*, mark options={fill=blue}] coordinates {
(1,2.75) (10,46.69) (20,56.27) (30,59.36) (40,60.53) (50,61.71) (60,64.28) (70,64.93) (80,65.86) (90,66.14) (100,67.40) (110,68.37) (120,68.76) (130,69.48) (140,69.94) (150,69.78) (160,69.93) (170,70.36) (180,70.31) (190,70.85) (200,71.23) (210,71.31) (220,71.46) (230,71.52) (240,71.69) (250,71.63) (260,71.75) (270,71.74) (280,71.80) (290,71.81) (300,71.85) (310,71.84) (320,71.85) (330,71.92) (340,71.91) (350,71.90)
};
\addlegendentry{top-1}
\addplot[orange, mark=*, mark options={fill=orange}] coordinates {
(1,8.75) (10,72.81) (20,80.58) (30,82.90) (40,83.60) (50,84.48) (60,86.09) (70,86.38) (80,87.16) (90,87.17) (100,87.85) (110,88.59) (120,88.95) (130,89.12) (140,89.39) (150,89.42) (160,89.50) (170,89.79) (180,89.63) (190,89.98) (200,90.06) (210,90.13) (220,90.18) (230,90.05) (240,90.33) (250,90.31) (260,90.30) (270,90.33) (280,90.35) (290,90.35) (300,90.35) (310,90.38) (320,90.40) (330,90.30) (340,90.42) (350,90.41)
};
\addlegendentry{top-5}
\end{axis}
\end{tikzpicture}

MobileNetV2 / ImageNet-1k, paper-faithful 350-epoch exponential-decay run (bf16, 4\(\times \) RTX 3060, RMSProp), 2026-08-30.

Two thirds of the final accuracy is in place by epoch \(50\), the run is still climbing at \(150\), and it is flat from about epoch \(250\) onward, so the last \(100\) epochs are worth almost nothing and the schedule is genuinely exhausted rather than cut short.

[TODO: re-run.] This result predates four fixes that landed while it was training, so it is a clean run of a recipe that is no longer the current one. The ledger, worst first:

In

Still out, worst first

Antialiased resize on train and eval, and timm’s \(50{,}000\)-image denominator — both halves of that protocol

convBnAct := .relu6 — the stem and the \(1\times 1\) head trained plain relu while every inverted-residual interior used ReLU6. Landed 2026-08-30, one day after this run started, and it is a phase-2 only defect: the verified render was always ReLU6

Label smoothing \(0\), as Sandler et al. used

BN running-stat decay as a per-net knob — still \(0.99\) here where PyTorch’s default works out to \(0.9\)

RMSProp \(\varepsilon \) inside the root, matching TensorFlow

Mixup/CutMix label-smoothing fix — inert for this recipe, which runs neither, but it moves the shared emitter

Exponential decay \(\times 0.98\)/epoch over the full \(350\), dropout \(0.2\), running-BN eval

Nothing re-run: the four above are known, none is measured

\(38.4\) h, one attempt, zero restarts, zero thermal pauses

 

The first row is the one to price. Two clamped activations are not obviously worth much, but this net carried the fleet’s largest paper gap for months with the schedule as the only suspect, and the gap was never decomposed against a reference that had them.

Phase 4: the verified trainer

Everything above this point is the phase-2 Lean\(\to \)JAX trainer. The phase-4 peer, which trains the same network through the proof-rendered StableHLO that §6.1 ran at Imagenette scale, exists and builds. It is apps/imagenette/MainMobileNetV2Imagenet.lean, built as lake build mobilenetv2-imagenet-verified:

-- Same inverted-residual backbone, 1000 classes instead of 10.
def mobilenetv2ImagenetVerified : VerifiedNetSpec where
  name       := "MobileNetV2 (ImageNet-1k)"
  slug       := "mobilenetv2in"
  nClasses   := 1000
  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 phase-2 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, both for the reasons §5.7 gives. 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 JAX 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 selects the optimizer, and one of them is the paper’s. At rms64 the driver loads mobilenetv2in_rms64_train_step.mlir, which has the recipe tabulated below rendered into it: RMSProp at \(\rho = 0.9\), \(\mu = 0.9\), \(\varepsilon = 1.0\) and coupled weight decay \(4 \times 10^{-5}\), at peak LR \(0.045\) with a 5-epoch warmup and \(\times 0.98\) per epoch after it. Those constants live in mnv2RmsSchedule and mnv2RmsHyper, shared with the Imagenette peer so the two cannot carry different values. At adam64 it is AdamW at \(10^{-3}\) instead, which is the default because it matches the other four ImageNet drivers.

What it has not done is run. No phase-4 ImageNet result exists for this network, which is why every accuracy above comes from phase 2. What is measured is the throughput, and the schedule follows from it:

Box

ms/step

min/epoch

350 epochs

Val top-1

4\(\times \) 4060 Ti (CUDA)

\(166 \to 140\)

\(13.8 \to 11.7\)

\(\sim \)84 \(\to \) \(\sim \)72 h  (\(3.5 \to 3.0\) d)

TBD

4\(\times \) 3060 (CUDA)

TBD

The columns are the two arms of rmsdp64 — RMSProp with exponential decay, the recipe mnv2-default-4gpu names. fp32 is the default; the run will be bf16, so read the right-hand column. Exactly half of that \(140\) ms step is no longer the graph, which times as \(69.6\) ms on its own, so this net is now as much data pipeline as compute — and that is also why the fp32 arm is the least repeatable number in this book, spanning \(159\)–\(196\) ms across three consecutive runs on an idle box where the bf16 arm held \(140 \pm 1\). A step that is half producer inherits the producer’s variance.

[TODO: run mobilenetv2-imagenet-verified.]

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:

Knob

ResNet-34 (2018)

MobileNetV2 (2018)

What the MNv2 choice buys

Backbone

basic block, 21.8 M

inverted residual, 3.5 M

depthwise \(3\times 3\) at the wide point between \(1\times 1\)s, at \(1/6\) the parameters

Optimizer

SGD + momentum \(0.9\)

RMSProp

\(\rho {=}0.9\), \(\mu {=}0.9\), \(\varepsilon {=}1.0\), where the huge \(\varepsilon \) keeps the adaptive step damped (B0 shrinks it to \(10^{-3}\) and pays for it)

Peak LR

\(0.1\)

\(0.045\)

RMSProp’s scale, not comparable to SGD’s \(0.1\)

LR schedule

cosine to zero

exponential, \(\times 0.98\) per epoch

the TF-lineage staircase the whole family inherits

Epochs

90

\(\sim \)300–400

the long-haul schedule, and the row the exponential decay above makes non-negotiable

Weight decay

\(10^{-4}\), all params

\(4\times 10^{-5}\)

depthwise filters are 9 weights per channel, and \(10^{-4}\) over-regularizes them

Label smoothing

\(0.1\)

none

Sandler et al. used no smoothing at all

Augmentation

RRC + hflip

RRC + hflip

unchanged, with no learned policy and no mixing, the leanest recipe in this book

What carries over unchanged: batch 256, train and eval both at \(224\), running-BN evaluation, bf16 matmul and convolution, and our 5-epoch warmup (an addition to both papers).

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

[TODO: run with timm tweaks.] The numbers in this section were measured before the trainers moved onto timm’s validation protocol and its antialiased resampler (appendix A). That change moves the training distribution as well as the evaluation, so it needs a re-run rather than a re-score, and none of it is reflected below.

Why the phase-2 trainer? The ImageNet numbers in this book come from the phase-2 (Lean\(\to \)JAX) trainer: at this scale its job is to validate the framework’s logic end to end, and it is the reference the verified path is measured against. Phase-3 verified-IREE codegen led to phase 4 — PJRT, which is the verified path itself. Where a section has both, it reports both.

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: \(\sim \)9.7M 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:

Knob

MobileNetV2 (2018)

MNv4-Conv-M (2024)

What the v4 choice buys

Block

inverted residual

UIB (+ fused IB stage)

the two optional depthwise slots subsume MNv2’s block, ConvNeXt’s, and an FFN in one template

Parameters

3.5 M

9.7 M

“mobile” now means fits-the-phone, not smallest-possible

Paper top-1

\(72.0\% \)

\(79.9\% \)

non-distilled, and distillation pushes the family higher still

Optimizer

RMSProp (\(\varepsilon {=}1.0\))

AdamW \(\beta (0.9, 0.999)\)

the TF-lineage staircase gives way to the transformer-era default

Peak LR

\(0.045\) @ 256

\(0.004\) @ 4096

effective batch 4096 (run here as \(512 \times 8\) grad-accum)

LR schedule

exponential, \(\times 0.98\) per epoch

cosine, 5-epoch warmup

the anneal every post-2020 recipe in this book uses

Epochs

\(\sim \)300–400

500

the family’s long-haul habit, taken further

Weight decay

\(4\times 10^{-5}\)

\(0.1\), skip norm/bias

\(2500\times \) stronger, applied selectively, because AdamW makes heavy decay usable

Augmentation

RRC + hflip

+ RandAugment (N2, m15, p0.7)

heavy learned augmentation, and still no Mixup/CutMix for Conv-M

Label smoothing

none

\(0.1\)

standard by 2024

Stochastic depth

none

drop-path \(0.075\)

linear ramp over the 21 UIB blocks

EMA

none

decay \(0.9999\), warmup-corrected

weight averaging (with the BN buffers shadowed, per Ch 7’s lesson)

The direction of travel matches the recipe tables closing 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 two tiers. 100 epochs (recipe default) is a reduced-regularization confidence tier at wd \(0.05\), dropout \(0.1\) and RandAugment m9, since the paper’s full pack underfits short schedules. 500 epochs (recipe full) is the paper recipe end to end. The 100-epoch tier is the one run here, and the \(79.9\% \) headline is a 500-epoch number:

GPU

Precision

Per epoch

Epochs

Total

Val top-1

Val top-5

4\(\times \) 4060 Ti (CUDA)

bf16

\(\sim \)9.0 min

100 (reduced-reg)

\(\sim \)16 hr

\(\mathbf{75.48\% }\)

\(\mathbf{92.37\% }\)

4\(\times \) 4060 Ti (CUDA)

bf16

\(\sim \)9.0 min

500 (paper)

\(\sim \)75 hr

The 100-epoch tier lands at \(\mathbf{75.48\% }\) top-1 and \(\mathbf{92.37\% }\) top-5 on the full 50,000-image validation split (EMA weights). Set against this chapter’s other half that is the whole six-year argument in two rows: \(\mathbf{+4.04}\) points over MobileNetV2’s paper-faithful \(71.44\% \), from \(2.8\times \) the parameters and less than half the epochs, \(100\) against \(350\) and \(16\) wall-clock hours against \(45\). 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.

\begin{tikzpicture} 
\begin{axis}[
    width=0.92\linewidth, height=6.5cm,
    xlabel={Epoch}, ylabel={Validation accuracy (\%)},
    xmin=0, xmax=102, ymin=0, ymax=95,
    xtick={0,20,40,60,80,100}, ytick={20,40,60,80},
    legend pos=south east, legend cell align={left},
    grid=major, grid style={gray!18},
    tick label style={font=\small}, label style={font=\small},
    every axis plot/.append style={line width=1pt, mark size=1pt},
]
\addplot[blue, mark=*, mark options={fill=blue}] coordinates {
(5,25.31) (10,52.04) (15,60.37) (20,64.37) (25,67.19) (30,68.99) (35,70.27) (40,71.23) (45,71.92) (50,72.51) (55,73.12) (60,73.61) (65,73.93) (70,74.42) (75,74.56) (80,74.93) (85,75.22) (90,75.41) (95,75.51) (100,75.51)
};
\addlegendentry{top-1}
\addplot[orange, mark=*, mark options={fill=orange}] coordinates {
(5,48.88) (10,76.27) (15,82.86) (20,85.67) (25,87.52) (30,88.60) (35,89.48) (40,89.99) (45,90.47) (50,90.85) (55,91.09) (60,91.37) (65,91.60) (70,91.86) (75,91.97) (80,92.13) (85,92.30) (90,92.29) (95,92.35) (100,92.37)
};
\addlegendentry{top-5}
\end{axis}
\end{tikzpicture}

MobileNetV4-Conv-M / ImageNet-1k, 100-epoch reduced-regularization tier (bf16, 4\(\times \) 4060 Ti, AdamW, effective batch \(4096\) via \(512\times 8\) gradient accumulation).

[TODO, 500 epochs.] mobilenet-v4-imagenet full. The tier above is detuned by construction, at wd \(0.05\) against the paper’s \(0.1\), dropout \(0.1\) against \(0.2\), and RandAugment \(m9\) against \(m15\), so its \(75.48\% \) is not a reproduction of \(79.9\% \). Its in-loop curve is flat over the last ten epochs (\(75.41 \to 75.51 \to 75.51\)). The schedule is exhausted, and the remaining points live in the regularization the full tier restores.

Phase 4: MobileNetV4 on the verified path

Everything above in this side quest is the phase-2 Lean\(\to \)JAX trainer running the faithful Conv-Medium. The phase-4 peer exists as of this writing, and it renders the same network.

The spec is built the way §6.5 built MobileNetV2’s and the way Chapter 5 built ResNet-50’s. 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}\), which is the \(\sim \)9.7M Conv-M is quoted at.

The \(75.48\% \) above is this network’s target. It is not this network’s result. The row below reports a measured wall clock and a TBD accuracy, because the network has been timed and has not been trained to convergence.

def mnv4ImagenetVerified : VerifiedNetSpec where
  name       := "MobileNetV4-Conv-M (ImageNet-1k)"
  slug       := "mnv4in"
  nClasses   := 1000
  data       := .imagenet
  shimScript := "generated_mobilenet_v4_imagenet_shim.py"
  layers     := [ ... the Imagenette UIB stack, 1000-class head ... ]

def main (argv : List String) : IO Unit :=
  mnv4ImagenetVerified.toNet.trainAdamSched
    { mnv4ImagenetConfig with batchSize := bs, epochs := epochs }
    (argv.head?.getD "data") baseLR 0.9 0.999 5 variant

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 nClasses, the batch and the slug, which is why adding them was three #eval lines rather than a new proof chain.

It runs. All three artifacts compile under PJRT, the shim feeds it, and a first step lands at loss \(6.99\) against \(\ln 1000 = 6.91\), with evaluation at \(0.106\% \) top-1 over the 49,920-image validation split, which is \(1/1000\) to two figures. Those are the numbers a correctly wired 1000-class network produces before it has learned anything, and they are the whole of what has been measured about its accuracy so far.

Box

ms/step

min/epoch

100 ep

500 ep

Val top-1

4\(\times \) 4060 Ti, global 256

\(176 \to 122\)

\(14.7 \to 10.2\)

\(25.5 \to 18.0\) h

\(127.5 \to 90.0\) h

TBD

4\(\times \) 3060, global 256

TBD

The columns are the two arms of adamdp64; fp32 is what mnv4-default-4gpu names today. \(100\) epochs is the reduced-regularization tier, \(500\) the paper’s.

[TODO: run mnv4-default-4gpu.] The job config exists and its prechecks pass. Note that it ships the fp32 arm by default — the left-hand column above, \(176\) ms/step and \(\sim \)\(25.5\) hr; one line in the config switches it, and the choice was to move one axis at a time rather than pair a first run with a precision change. The Val top-1 column stays TBD either way.