MobileNetV2's train-step tie at batch BatchNorm — the un-fused, batched whole-net thread #

GradNodesB makes every parameter gradient node of the batched MobileNetV2 train step den-faithful for an arbitrary cotangent. This file removes the "arbitrary": each cotangent is pinned to the one the emitted backward chain delivers, so the whole train step is den-composed forward → loss → backward with no free activation and no symbolic cotangent. It is ResNet34StepTieB.lean's peer.

The block cotangents are not derived here. MobileNetV2FullBVJP's mnv2{ExpOnly,Resid,Strided,NoExp}BHasVJPAt are the certified block backwards, and mnv2{Body,DownBody,ResidBlock}BackBatchedGraph_faithful (MobileNetV2BackB0.lean) prove the emitted backward subgraphs denote them. The four *CotIn_eq_vjp lemmas below (mnv2NoExpCotIn_eq_vjp, mnv2ExpOnlyCotIn_eq_vjp, mnv2ResidCotIn_eq_vjp, mnv2StridedCotIn_eq_vjp) are those statements in this file's vocabulary: each says one block kind's constructed input cotangent is its certified block VJP's backward, under that block's positivity and smoothness hypotheses. The two ends are tied the same way: mnv2HeadCotBlk_eq_vjp (the cotangent the head hands to b17) and mnv2StemCotC_eq_vjp (the stem's conv-output cotangent, through the conv input-VJP the render does not emit).

One parameter-tie bundle covers twelve of the seventeen blocks. A skip block and a stride-1 widening have the SAME parameter cotangents — the identity skip changes only the dx handed to the previous block, which is why MobileNetV2RenderB's irBackStride1GradB is one function with a skip flag rather than two near-copies. So mnv2Stride1TiedB is stated once and instantiated at b3, b5, b6, b8, b9, b10, b12, b13, b15, b16 (skip) and b11, b17 (no skip).

The loss cotangent is the label-smoothed one, at a general target, shared with ResNet-34: Foundation/SmoothedLossCot.lean. MobileNetV2RenderB composes it from the same six kit ops (softmaxRow → subB → scaleB → addVB → shiftB → divConstB) with α at 0.1 and the target arriving as the graph input %onehot — a soft vector under mixup or cutmix.

N is a binder. The single-device artifacts at 32 (mobilenetv2_adam_train_step) or 64 (mobilenetv2in_adam64_train_step) are instances.

Data parallel. In the dp artifacts (e.g. mobilenetv2in_rmsdp64) each gradient node feeds allReduceMeanF and BatchNorm is synchronised across replicas, so a replica's gradient node is not this file's node at its own N. This file describes a data-parallel artifact only at N := R·N: the per-replica statement is MobileNetV2SyncTieB.mnv2_net_syncTiedB, whose right-hand side is mnv2_net_tiedB's node at the global batch.

Precision is a flag on the statement. The ImageNet run the book reports trains from mobilenetv2in_rmsdp64wxdols0eps0001bf16, whose conv and depthwise weight gradients are the bf16 kinds (convWeightGradBBf16, depthwiseWeightGradBBf16, …; Foundation.Bf16GradNodes). Every conv and depthwise weight node below is stated on the renderers' switch — ConvWTiedBAt bf16 / ConvStridedXlaWTiedBAt bf16 / DepthwiseWTiedBAt bf16 / DepthwiseStridedXlaWTiedBAt bf16 (Foundation.GradNodesBAt): bf16 := false is the f32 node, bf16 := true the bf16 one — and the capstone mnv2_net_tiedB takes bf16, so it reaches the bf16 artifacts' gradient nodes read over ℝ exactly as it reads the f32 ones (Bf16Erasure: at the identity rounding the bf16 kind denotes what its f32 peer does). BatchNorm, bias and dense nodes carry no flag because no render switches them.

Classifier dropout is an optional site on the statement. mobilenetv2in_rmsdp64wxdols0eps0001bf16 is a *do* artifact: MobileNetV2RenderB emits dropoutB on the driver's per-element mask %do between the GAP and the dense in the forward, reads the DROPPED activation at the classifier weight gradient (its cin), and applies the same mask to the dense's input-VJP before the GAP backward (Proofs.dropout_vjp_is_self). The head chain, the head bundle and the capstone take the site as cd : Option (Vec (N * 1280)) (Proofs.Foundation.DropSites): none is the drop-free artifacts' chain — the statement exactly as it was — and some m the *do* artifacts', where mnv2HeadCotGapIn is dropout m of the classifier's input-VJP and the dense weight node reads dropout m of the pooled activation (the content of the dropout-tie gate's W, tests/TestDropoutTie.lean). The head is stated at mnv2HeadBDoOpt — mnv2HeadB at none, mnv2HeadBDo at some m, both by rfl — and the net at mobilenetv2ForwardBFullDoOpt, likewise.

bnInB and bnInB_eq_bnBackB are ResNet-34's, imported rather than copied. They are the batched BatchNorm input-cotangent written as the den of the emitted backward op, and its identity with the certified bnBatchLA VJP — both net-agnostic, and they happen to live in the file that first needed them. What MobileNetV2 adds is the TWO-SIDED relu6 mask (relu6MaskB, where r34 threads the one-sided reluMaskB), the depthwise input-VJP and its XLA-SAME strided peer.

The parameter census is 158, and this file states 210 slots #

MobileNetV2RenderB defaults to convBias := false: every conv, depthwise and project bias is folded into the BatchNorm that follows it and bound to zeroBiasPrelude's shared zero constant. So mobilenetv2_adam_train_step.mlir carries 158 updated parameter tensors — stem 3 + b1 6 + 16 blocks × 9 + head 3 + dense 2 — and 210 is the census at convBias := true. The bias conjuncts below are kept (one delegation each, and they cover the flag) and are about ops the committed artifacts do not emit. Nothing here weakens a theorem: every fold is ∀-quantified over op instances, and bias = 0 is one of them.

Conventions, stated because nothing here checks them #

MobileNetV2RenderB.irBackNoExpGradB, node for node, from the block-output cotangent dyOut:

%dpc = bnBatchBack(pg, pc)   %ddr = convBackBatched(pW)   %ddm = selectMidB(dn)
%ddn = bnBatchBack(dg, dc)   %dxb = depthwiseBackBatched(dW)

with dW,db ← %ddn, dg,dbt ← %ddm, pW,pb ← %dpc, pg,pbt ← %dy.

MobileNetV2RenderB.irBackStride1GradB, node for node, from the block-output cotangent dyOut:

%dpc = bnBatchBack(pg, pc)   %ddr = convBackBatched(pW)         %ddm = selectMidB(dn)
%ddn = bnBatchBack(dg, dc)   %der = depthwiseBackBatched(dW)    %dem = selectMidB(en)
%den = bnBatchBack(eg, ec)   %dxb = convBackBatched(eW)
%dx  = if skip then addVB(%dxb, %dy) else %dxb

with eW,eb ← %den, eg,ebt ← %dem, dW,db ← %ddn, dg,dbt ← %ddm, pW,pb ← %dpc, pg,pbt ← %dy. The skip flag touches only %dx, so all twelve parameter cotangents below are shared between the skip blocks and the two stage-first widenings (b11, b17).

MobileNetV2RenderB.irBackStridedGradB: the stride-1 chain with the depthwise replaced by its XLA-SAME strided peer, so the expand half runs at the 2h × 2w input grid and %der upsamples. There is no skip, so %dx is %dxb directly.

MobileNetV2RenderB, after the seventeen block backwards:

%dsm = selectMidB(stn)   %dsn = bnBatchBack(sg, stc)

with sW,sb ← %dsn and sg,sbt ← %dsm. There is no conv-back past %x, and no pool: the stem is one XLA-SAME strided conv, its BatchNorm and one relu6.

MobileNetV2RenderB, from the loss cotangent g:

%dgi = denseRowBack(Wd)   %dgp = gapBackBatched   %dhm = selectMidB(hn)
%dhn = bnBatchBack(hg, hc)   %dhx = convBackBatched(hW)

with Wd,bd ← %dy, hW,hb ← %dhn, hg,hbt ← %dhm, and %dhx the cotangent handed to b17. Unlike ResNet-34's, this head is not batchMap of a smooth per-example map — MobileNetV2 puts a 1x1 conv-BN-relu6 in front of the pool — so the chain is spelled here as dens, exactly as EfficientNetStepTie.enetHeadTied spells B0's identically-shaped head.

Each conjunct is GradNodesB's ∀ cot fold instantiated at the cotangent the render's chain delivers, so nothing here is a new proof: the bundles are the GradNodesB folds with the freedom removed. reassocB bridges the conv/relu6 index N·(c·h·w) to the BatchNorm parameter ops' N·(c·(h·w)).

The bias conjuncts are about conv{,StridedXla}BiasGradB and depthwise{,StridedXla}BiasGradB, which the committed artifacts do NOT emit — MobileNetV2RenderB runs convBias := false and binds every bias operand to zeroBiasPrelude's zero constant. They are kept because they cost one delegation each and they cover the flag.