EfficientNet-B0 train-step tie (fused SGD) — efficientnet_net_tied #

Every one of B0's 262 parameters' SGD ops, at the cotangent the rendered net's own backward chain hands it: the real forward (efficientnetForwardBFull) threaded through every parameter op, the loss cotangent composed back through all 16 MBConv blocks, the residual fan-in at each stride-1 skip and the squeeze-excite gate fan-in included. The batched, un-fused peer (tied at the *GradB nodes) is EnetTieG.efficientnet_net_tiedG in EfficientNetStepTieG.

What B0 adds over MobileNetV2's tie #

• swish masks (swBackB) instead of relu6's two-kink selectMid, at every conv-bn-swish and depthwise-bn-swish stage;
• the SE gate fan-in: the cotangent at an SE input is gate ⊙ dyOut (the fused seBackBatched value), and the SE dense parameters' cotangents come from seReduceB (the gate cotangent Σ_{h,w}(x⊙dy)) through sigmoidBack → denseRowBack(W₂) → swishBack;
• true batch-norm backward (bnBackB, the bnBatchLA VJP), batch-coupled;
• the strided depthwise back (dStridedXlaInB); reassocB bridges the conv/swish (oc·h·w) and BN (oc·(h·w)) layouts.

EfficientNetChainClose's whole-net backward is vjpComp of the per-block HasVJPs, so the tie builds explicit chain-cotangent constructors rather than reading them off.

Contents #

• Per block type, every parameter op at the chain cotangent — each a delegation to the fold generics EnetFold.*:
  • enet_exp_tied (16 params) — the stride-1 expand block: the 9 residual blocks and the two widenings b9/b16 (the parameter ops are skip-agnostic; the fan-in lives in the thread);
  • enet_strided_tied (16) — b2/4/6/12: expand at 2h×2w, strided depthwise;
  • enet_noexp_tied (12) — b1 (t = 1: depthwise on ic → SE → project);
  • enet_stem_tied (4) — the 3×3/s2 conv-bn-swish stem;
  • enet_head_tied (6) — the 1×1 conv-bn-swish head and the dense, which ties at the loss cotangent g.
• efficientnet_net_tied — the whole-net thread. Block inputs are the forward's prefixes (a0..a16); the per-block output cotangents (dy0..dy16) are composed top-down by the proven block VJPs (headFwdBHasVJP, mb{Exp,Resid,Strided,NoExp}WHasVJP) from g. The residual fan-in is in mbResidW's own VJP (it includes the + x). One *TiedAt wrapper per block type takes the B0Weights block bundle.

Not done: the dense head's total-loss fold (Wfc → ∂CE/∂Wfc, the batched Σ_n analogue of StableHLO.lossWeightGrad_eq_sum); the head dense ties at g directly.

Residual stride-1 MBConv block — all 16 params tied (expand → dw → SE → project + skip) #

The centerpiece: exercises the genuinely-new content vs mnv2 — swish masks (smooth), the SE gate fan-in (gateCotB → sigBackB → {zW₂,zb₂} → rowDenseBackFlat → swBackB → {zW₁,zb₁}), and true batch-norm backward (bnBackB), all at the batched index N·(c·h·w). Backward from dyOut (cot at project-BN out): project-BN-back → project-conv-back (cot at SE out) → SE backward (fused dx for the depthwise side; un-fused gate-cot for the SE params) → depthwise swish/BN/conv backs → expand swish/BN/conv backs. Residual (ic=oc=c): the block-input cotangent fan-in + dyOut lives in the whole-net thread, not here (the param ops are skip-agnostic — identical to the no-skip widenings).

Strided downsampling MBConv block — all 16 params tied (b2/b4/b6/b12) #

Same as the expand block EXCEPT the expand stage lives at the block-input grid 2h×2w and the depthwise is strided (depthwiseStridedXlaWeightSgdB, the expand-side cotangent cotEr upsamples h→2h via dStridedXlaInB). No skip (spatial+channels change).

No-expand MBConv block (b1, t=1) — all 12 params tied (depthwise on ic → SE → project) #

NO expand conv: the depthwise runs directly on the block input (ic channels). 12 params (4 depthwise+BN, 4 SE, 4 project). The SE squeeze/excite is on ic channels (ic → r → ic).

Stem — the 3×3/s2 conv-bn-swish (4 params), feeding block 1 #

swish(bn(convStride2Xla Ws bs x)), 3→32 at 224→112, at the XLA-SAME phase the shipped stem uses. The cotangent block 1 delivers at the stem swish output (dyStem) lifts through swish-back

• true-BN-back to the conv-out cotangent (the convStridedXlaWeightSgdB consumes it; NO conv-back past %x). 4 params.

Head — the 1×1 conv-bn-swish (4 params) → GAP → dense (Wfc/bfc), + the loss cotangent #

dense(GAP(swish(bn(conv Wh bh))))) (320→1280 conv, GAP, 1280→nClasses dense), then the batched per-row softmax-CE gradient g = rowSoftmax(logits) − onehot. The head conv params tie at the chain cotangent (loss → dense-back → GAP-back → swish/BN-back); the dense Wfc/bfc tie at the loss cotangent g directly.

Bundle-taking *TiedAt wrappers — one per block type, for the whole-net thread #

Each takes the B0Weights block bundle (MBW/MBWNoExp) + its ε-positivity + the block input + the downstream cotangent dyOut, and delegates to the per-block-type tie.

The whole-net thread — all 262 params tied through the REAL efficientnetForwardBFull #

The capstone: efficientnetForwardBFull's prefixes are the block inputs (a0..a16 = stem, then the 16 MBConv blocks), and the per-block output cotangents (dy0..dy16) are composed TOP-DOWN by the proven block VJPs (headFwdBHasVJP, mb{Exp,Resid,Strided,NoExp}WHasVJP) from the loss cotangent g = rowSoftmax(logits) − onehot. Each block's tie then holds at its real input + threaded dyOut. The whole-net tie: the WHOLE 16-MBConv (262-param) EfficientNet-B0 train step is den-composed forward→loss→backward, no free activations, no symbolic cotangent. The residual fan-in at the 9 identity skips is folded into mbResidW's own VJP (it includes the + x), so it is automatic.