MobileNetV4 — the batched UIB backward, and the four families as one chain #

The block- and stage-level backward of MobileNetV4-Conv-M at the batched index: the depthwise stages, the UIB body, the stride-2 body, the fused stage and the head, each a CertLayer whose backward graph is faithful to its certified VJP. The identity skip is CertLayer.residual, applied by the caller.

The four families are one chain #

The UIB body is preDW? → expand → postDW? → project, and k = 0 omits a depthwise. Both depthwise positions are channel- and shape-preserving — preDW : ic → ic, postDW : mid → mid — so an absent one is the identity layer (CertLayer.id') in the same slot:

mnv4UibBody takes the two depthwise slots as CertLayer arguments and the caller passes id' where the table says k = 0 — one body, no case split, and no dispatch that could disagree with the forward's (mnv4BodyOfRow reads the slots off the UibSpec row).

The two depthwise stages, and where the activation is #

timm's mobilenetv4_conv_medium (the pinned spec, checked by scripts/parity/mnv4_timm_parity.py): the pre-DW (dw_start) is depthwise → BN with no activation (dwbB, globally certified); the post-DW (dw_mid) is depthwise → BN → relu (dwbReluB, and its strided peer dwbReluBstrided). Relu, not relu6 — MobileNetV2 uses relu6 one file over. The relu stage is bnReluStageHasVJPAt (Foundation/Batched/StageLayers, generic in the op) at depthwiseFlat; its backward graph masks with .selectPos, relu's one-sided mask.

Contents #

• the depthwise stages — BN-only (mnv4DWBnLayer, the pre-DW) and BN-relu, stride-1 and strided (mnv4DWReluLayer, mnv4DWReluStridedLayer, the post-DW) — with backward graphs;
• the family-collapsing body mnv4UibBody and the stride-2 form mnv4UibStridedBody (the post-DW carries the stride, timm's dw_mid rule);
• the fused stage (mnv4FusedStage, stage 0): cbReluStridedLayer then projLayer;
• the head mnv4Head: conv-bn-relu, GAP, conv_head-bn-relu on the pooled features, dense;
• UibParams, the row-typed weight record, and the table-driven k = 0 dispatch.

The net level consumes these block by block: the whole-net forward, graph and VJP are MobileNetV4FullB / MobileNetV4FullBVJP, the train-step tie is MobileNetV4StepTieB, and the input-gradient tie is MobileNetV4WholeBackCertifiedTieB.

The stem is not here and is not a CertLayer: no render emits a gradient into %x, so there is no backward graph for it, and a net-level forward composes the stem's VJP by vjpCompAt rather than CertLayer.comp.

The stride-1 collapse does not extend here, and the reason is the type.

At stride 1 an absent depthwise is id' because the slot is shape-preserving. At stride 2 the depthwise that carries the stride maps (2h, 2w) ↦ (h, w) — a different type — so it cannot be replaced by an identity. timm puts the stride on dw_mid whenever a block has one, so in all three Conv-M downsamples (rows 1, 3, 11, all ExtraDW) the optional pre-DW and the expand run at the INPUT resolution 2h and the strided post-DW takes it to h. The pre-DW is still a slot (at 2h). Putting the stride on the pre-DW instead (expand at h) is a different function with the same parameter shapes.

All three stride-2 blocks change channels (ic ≠ oc), so none has a skip: the block IS the body, with no CertLayer.residual wrapper. Adding one would not typecheck, which is the good case.

MNv4's stage 0 is .fusedMbConv 32 48 4 3 2 1 false .relu (timm's EdgeResidual): a regular k×k conv (not a depthwise) doing expansion and downsampling at once, then a 1×1 project. 32 → mid = 32·4 = 128 → 48, stride 2, symmetric padding, relu — so the strided conv-bn-relu stage is cbReluStridedLayer and the project is projLayer (both Batched.StageLayers).

timm's head: Conv-M's cn_r1_k1_s1_c960 (1×1 256 → 960, BN, relu) at 7×7, then global pool, then conv_head (1×1 960 → 1280) → norm_head (BN) → relu on the POOLED [N, 960, 1, 1], then the classifier. The second conv-bn-relu is cbReluLayer at h = w = 1, so its BatchNorm normalises over the batch alone. Running conv_head at 7×7 before the pool would be a different function: batch BN and relu do not commute with the pool.

What this section fixes. Above, mnv4UibBody takes its depthwise slots as arguments, so passing id' where block 4's real pre-DW belongs is well-typed and still certified — graph and VJP both move with the caller's arguments, so the theorem stays true about the wrong net. Types catch the stride split (resolution is in the type); they catch nothing about k = 0 vs k > 0, because that slot is shape-preserving — the very property that made the collapse possible.

The fix is to make the slot a function of the block table's k, so the proof side runs the same k = 0 dispatch the render does, off the same mnv4Blocks list — one table, not two readings. mnv4-fwd-smoke already pins the render against that table; these #guards pin the table's own shape, so a bad edit fails at lake env lean rather than becoming a silent net.

The gap this closes. A block builder that reads the dispatch from the table but takes its weights as separate arguments lets a caller pair row 4's ks with row 7's widths. The dispatch was table-driven; the wiring was not.

Fix: index the parameter record by the row. Every width in UibParams s is a projection of s — s.ic, s.oc, s.ic * s.expand, s.preDWk, s.postDWk. A record with widths that disagree with its row cannot be constructed, so the block builder needs no side conditions and no #guard: it is impossible by typing rather than checked after the fact.

What this still does not pin. Rows with the same (ic, oc, expand, preDWk, postDWk, h) — {4, 5, 7}, {8, 10}, {12, 18}, {13, 14} and {15, 19, 20} — have records of the same type, so swapping two such blocks' weights typechecks. Typing pins shape (channels, expand ratio, kernel extents, resolution); it cannot pin identity between rows that are shape-identical. Closing that needs the weights to come from one indexed array — a renderer concern, since the render already folds mnv4Blocks in order.

The slot parameters by name, read through DWSlot.params (placeholders at k = 0).