MobileNetV4 — the Universal Inverted Bottleneck render #
The UIB block, batch-BN, at N := B:
optional pre-DW (k×k) → BN -- stride 1, no activation (timm dw_start)
expand 1×1 (ic → mid) → BN → relu
optional post-DW (k×k) → BN → relu -- carries the block's stride (timm dw_mid)
project 1×1 (mid → oc) → BN -- NO activation
+ skip -- iff stride = 1 ∧ ic = oc; NO post-add activation
mid = ic * expand, every conv BN-followed and therefore bias-free — VLayer.toSpecs and
the baseline's Layer.nParams both assume that, and uib-layout-tie pins them to each other.
8,447,322 scalars in 233 slots at nClasses = 10 (9,715,512 at 1000) over Conv-M's 21-row
table, counted off the signatures of mnv4_fwd.mlir and mnv4in_fwd.mlir.
k = 0 omits that depthwise, which is how one block expresses MNv4's four families:
ExtraDW (both), IB/MBConv (post only), ConvNeXt-like (pre only), FFN (neither). Those are ifs
here, not separate functions, because omitting a shape-preserving op does not change any type.
The block needs no new op and no new position: the depthwise VJP is kernel-general
(depthwiseFlatHasVJP, depthwiseBack_faithful; the descriptor carries kH kW), and
a leading depthwise already exists — MobileNetV2RenderB.irFwdNoExpB, the t = 1 inverted residual,
emits .depthwise (c := ic) straight onto the block input. What is new is the composition:
ExtraDW puts a depthwise on both sides of the pointwise expand.
Why two functions and not one #
.depthwise : BatchableOp (c*h*w) (c*h*w) and .depthwiseStrided : BatchableOp (c*(2h)*(2w)) (c*h*w) have different INPUT types, so a stride-polymorphic block cannot typecheck — the same
reason MobileNetV2RenderB splits irFwdStridedB from irFwdSkipB. In a strided block the
stride rides the post-DW (timm's dw_mid) whenever there is one, so the optional pre-DW and the
expand run at the input resolution. Read off the Conv-M table
(jax/MainMobilenetV4.lean):
| stride | ic vs oc | function | |
|---|---|---|---|
| 18 blocks | 1 | ic = oc | uibFwdSkipB |
| 3 blocks (1, 3, 11) | 2 | ic ≠ oc, pre- and post-DW present | uibFwdStridedB |
Note: the stride position (post-DW, not pre-DW), the BN-only pre-DW, the relu fused stage, the
1×1 head conv and the symmetric-pad stem follow timm's mobilenetv4_conv_medium; parameter counts
alone cannot distinguish any of these from the alternatives.
Activation is plain relu, not relu6, after the expand and the post-DW; the pre-DW and
the project are BN only. MobileNetV2's blocks use relu6 and this file sits next to that renderer,
so the wrong one is one keystroke away.
A pre/post-DW swap is invisible to every count. Same k, same channels ⇒ same parameter
shapes, so uib-layout-tie passes on a renderer that swaps them, and so does any arity or op-count
audit. At stride 1 it is invisible to the TYPES too, since both positions are shape-preserving —
which is why the four families are ifs that the compiler cannot check. Only a forward tie against
the reference on shared weights pins the order. Same class as R50's stride-on-the-3×3.
Note: the baseline drops the stride entirely for a stride-2 FFN block (no depthwise to carry it,
MlirCodegen.emitTrainStepBody's .uib case). No such block exists in the table; this file has no function for that
shape, so the case is absent rather than silently wrong.
References #
- Qin et al. 2024, MobileNetV4: Universal Models for the Mobile Ecosystem. https://arxiv.org/abs/2404.10518
The stochastic-depth sites: the block indices (into mnv4Blocks) of the 18 skip rows, in
order. Site k is %dp<k>, and the driver's dropKeeps lists one keep per site in this
order. The three strided rows (0, 2, 10) have no residual, so no site, as in the reference's
uib_block (_drop_branch only under residual.shape == x.shape).
Equations
- One or more equations did not get rendered due to their size.
Instances For
Block bi's drop site, none for a strided row.
Equations
Instances For
The MobileNetV4-Conv-M parameter inputs, in func-arg order: stem (3), the fused stage, the
21 UIB blocks, the TWO head convs, the classifier. Single source for the signature and the return
order, the same role r50ShapeList plays for R50.
This list and VLayer.toSpecs are TWO HAND-WRITTEN READINGS of the same layout — the
renderer cannot import the spec without inverting the dependency, which is the same two-lists
shape as toSpecs == XLayout.specs elsewhere. mnv4-fwd-smoke is the gate that pins them.
Equations
- One or more equations did not get rendered due to their size.
Instances For
The same list as MLIR types. Derived, so the shapes have one definition.
Equations
- Proofs.StableHLO.mnv4SigList nClasses = List.map (fun (x : String × List ℕ) => match x with | (n, ds) => (n, Proofs.StableHLO.ty ds)) (Proofs.StableHLO.mnv4ShapeList nClasses)
Instances For
The 154 BN running-statistic slots = 77 BN layers × (μ, var), in forward-traversal order: stem, the fused stage's two, each UIB block's (2–4 depending on family), the head.
A misaligned stat slot is SILENT: the arities still match and the wrong layer's statistics
simply flow into the wrong @mnv4_fwd_eval slot. That is why the order here and the order the
train step returns them in are both derived from the same block table, and why the eval forward
reads them through mnv4Bn's statP rather than an independently-numbered list.
Equations
- One or more equations did not get rendered due to their size.
Instances For
The stat slots as MLIR types. Derived, so the widths have one definition.
Equations
- Proofs.StableHLO.mnv4StatSigList = List.map (fun (x : String × List ℕ) => match x with | (n, ds) => (n, Proofs.StableHLO.ty ds)) Proofs.StableHLO.mnv4StatShapeList
Instances For
Saved forward SSA names the UIB backward + gradient passes reference. The optional fields are
"" when their depthwise is absent (k = 0), which is how the backward reads off which family
it is looking at without re-deriving it from the kernel sizes.
- code : String
- o : String
- qc : String
- qr : String
- ec : String
- en : String
- er : String
- dc : String
- dn : String
- dr : String
- pc : String
- qst : String
- est : String
- dst : String
- pst : String
Instances For
Equations
Fused inverted bottleneck, stride 2, no SE — MobileNetV4's stage 0
(.fusedMbConv 32 48 4 3 2 1 false) and EfficientNetV2's early stages.
k×k regular conv ic→mid at stride 2 → BN → relu -- 2h×2h → h×h
1×1 project mid→oc → BN -- NO activation
"Fused" means the MBConv expand-1×1 and its depthwise collapse into ONE regular k×k conv, so
despite living in a mobile net there is nothing depthwise here. No skip: ic ≠ oc and stride 2.
ReLU, as timm's EdgeResidual in mobilenetv4_conv_medium (scripts/parity/mnv4_timm_parity.py).
The JAX fused_mbconv_block is shared with EfficientNetV2 (swish) and takes the activation as
an argument.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Block-shape dispatch, in one place. Which of the two forwards a row uses is forced by the
row: stride 1 ⇒ the identity skip (ic = oc, pinned by mnv4-fwd-smoke); stride 2 ⇒ the
post-DW carries it (timm's rule when both depthwise convs exist). The two cannot be one
function — .depthwise and .depthwiseStrided differ in INPUT type.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Everything the whole-net render needs out of one forward traversal: the emitted code, the
logits, and every saved activation the backward reads. inputs is each block's INPUT SSA name
(inputs[i] is what block i consumed), which the weight gradients contract against.
- code : String
- logits : String
- stc : String
- stn : String
- str : String
- f0 : UibFwdB
- h1c : String
- h1n : String
- gap : String
- hc : String
- hn : String
- cin : String
- last : String
- sst : String
- h1st : String
- hst : String
Instances For
Equations
Stem forward: 3×3/s2 conv (3→32), 224→112 → batch BN → relu, on %x.
Symmetric (1,1) padding (.convStridedAt), as timm's conv_stem. XLA-SAME padding
(.convStridedXlaAt, (0,1) for a 3×3/s2 at 224) also gives 112×112, so only a forward on shared
weights sees the difference (scripts/parity/mnv4_timm_parity.py,
scripts/parity/mnv4_forward_tie.py).
Equations
- One or more equations did not get rendered due to their size.
Instances For
Head forward, timm's order: 1×1 (256→960) → BN → relu at 7×7 (Conv-M's cn_r1_k1_s1_c960),
GAP(7×7), then conv_head 1×1 (960→1280) → norm_head BN → relu on the POOLED [B, 960, 1, 1],
then dense. The second BN therefore normalises over the batch alone (N·1·1 cells per
channel). Running conv_head at 7×7 before the pool would be a different function with the
same parameter count: batch BN and relu do not commute with pooling.
Equations
- One or more equations did not get rendered due to their size.
Instances For
The MobileNetV4-Conv-M forward chain, batch BN, at N := B, 224² → nClasses.
Transcribed 1:1 from the Conv-M table in jax/MainMobilenetV4.lean. Spatial ladder:
224 --stem s2--> 112 --fused s2--> 56 --uib s2--> 28 --uib s2--> 14 --uib s2--> 7 --GAP--> 1
Block dispatch follows the table: the three stride-2 blocks (1, 3, 11) are ic ≠ oc and
strided at the post-DW (timm's dw_mid); the eighteen stride-1 blocks are all ic = oc,
hence all skip. The stride/skip split is checked by the types; which depthwise positions a
block has (its family) is an if the types do not check.
Families in order after the fused stage: ExtraDW ×7, ConvNeXt, FFN, ConvNeXt, ExtraDW ×4, FFN,
ConvNeXt, ExtraDW ×2, FFN ×2, ConvNeXt — 13 / 4 / 4 and no IB at all.
Activations, as timm's mobilenetv4_conv_medium:
- stem and head
.convBn→ relu (MlirCodegen.emitConvBnTrain … useRelu := true) - the fused stage → relu (timm's
EdgeResidual) - every UIB block → relu after the expand and the post-DW; the pre-DW (
dw_start) and the project are BN only (timm'sUniversalInvertedResidual)
Returns the full forward RECORD, not just (code, logits): the train step needs every saved
activation the backward reads, and the alternative — a second copy of the chain inside the
train step — is the two-readings defect this file exists to avoid.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Every distinct channel width a bias-free conv in this net binds %zb{c} at: the stem, the fused
stage's two convs, each block's pre-DW (ic), expand and post-DW (mid) and project (oc), and
the head. Derived by hand and pinned by mnv4-fwd-smoke, which fails on an unbound %zb.
Equations
Instances For
@mnv4_fwd — the MobileNetV4-Conv-M forward as one MLIR module.
%x plus mnv4SigList's parameters in VLayer.toSpecs order, logits [B, nClasses]. Every
conv is bias-free, so the proven conv ops' bias operands bind to the %zb{c} zero constants the
prelude declares — same op, bias = 0, and x + 0.0 is exact.
zeroBiasPrelude must cover every width the chain references or the module has an unbound
SSA name. That is a link error at parse time rather than a wrong number, which is the good
direction, but mnv4-fwd-smoke checks it anyway so the failure arrives at lake build and not
at iree-compile.
Equations
- One or more equations did not get rendered due to their size.
Instances For
@mnv4_fwd_eval — the inference forward, every BN site reading frozen running stats.
%x + 233 params + 154 stat inputs = 388 inputs (counted off mnv4_fwd_eval.mlir). This
is what the driver scores through.
It is mnv4FwdChainB at .eval — the SAME traversal @mnv4_fwd and the train step use, so
its BN order matches mnv4StatSigList by construction rather than by a second reading. Its
typed graph is mnv4FwdGraphBFullEval at f = s / 32 (MobileNetV4FullBEval).
Equations
- One or more equations did not get rendered due to their size.
Instances For
The driver's variant slug for a given (B, replicas): the artifact is
verified_mlir/mnv4_<variant>_train_step.mlir, the entry point is
@mnv4_<variant>_train_step, and LEAN_MLIR_VARIANT selects it. All three must agree — the
shim checks the entry name and refuses a mismatch outright rather than running the wrong graph.
B = 32 is deliberately unsuffixed so the Imagenette artifact keeps a stable name.
Equations
Instances For
The variant slug of a recipe render — the JAX reference's recipe renders. opt = none is mnv4AdamVariant, so every committed spelling
is unchanged. Markers in the order the driver's predicates read them (VerifiedVariant):
ema first (emaOn is a PREFIX test), then the optimizer with its k (accK parses it back
out after acc[dp]), the per-replica batch, wx, do, the decay mark, bf16.
Equations
- One or more equations did not get rendered due to their size.
- Proofs.StableHLO.mnv4RecipeVariant B replicas bf16 none ema wdExclude wdStr cd sd = Proofs.StableHLO.mnv4AdamVariant B replicas bf16
Instances For
MobileNetV4-Conv-M AdamW train step, batch BN, rendered from the verified AST at N := B.
858 inputs (%x, 233 θ, 233 m, 233 v, %lr/%bc1/%bc2, 154 running-stat slots,
%onehot) and 856 outputs (233 θ', 233 m', 233 v', %loss/%bc1/%bc2, 154 batch
stats) — counted off mnv4_adam_train_step.mlir. Parameter
order comes from mnv4ShapeList — through zipPs, which builds each block's gradient list from
the very same uibSig slice the signature does — and stat order from mnv4StatShapeList.
Forward, as timm's mobilenetv4_conv_medium: symmetric stem 3×3/s2 (3→32,
224→112) → fused MBConv (32→48, 112→56, ReLU) → the 21 UIB blocks (stride on the post-DW,
BN-only pre-DW; ExtraDW, ConvNeXt and FFN — no IB) → 1×1 conv-BN-relu (256→960) → GAP → 1×1
conv_head (960→1280, BN over the batch at 1×1) → dense.
The cotangent is composed from kit ops (softmaxRow → subB → scaleB → addVB → shiftB → divConstB, α = 0.1, K = nClasses), and %loss is report-only and stays outside the AST — the
same carve-out resnet34/mobilenetv2 take.