EfficientNet — MBConv with Squeeze-Excite, end-to-end VJP #

The hardest of the three flagship CNN VJPs in this stack (alongside the ResNet cnnHasVJPAt and the MobileNet depthwise chain), because the squeeze-excite gate is a genuine fan-out sub-network multiplied back into the main path. We reuse seBlockHasVJP (SE.lean), which already carries the product-rule fan-in for x ⊙ gate(x); here we supply the concrete gate (seGate) — a real Vec → Vec differentiable map with its own composed VJP — plus its differentiability.

What this file provides #

• The gate's pieces — sigmoid, broadcastFlat (the adjoint of GAP), the concrete gate seGate = broadcast ∘ sigmoid ∘ dense ∘ swish ∘ dense ∘ GAP and seBlockFull (the full x ⊙ gate(x)) — are op-level and live in SE.lean.
• mbconvBody / mbconvBodyHasVJP — one MBConv block body project(1×1 conv-bn) ∘ SE ∘ depthwise(bn-swish) ∘ expand(1×1 conv-bn-swish), smooth everywhere (global HasVJP).
• efficientnetHasVJP / efficientnetHasVJP_correct — a two-block end-to-end EfficientNet (stem → MBConv-with-SE-and-residual → MBConv-with-SE → globalAvgPool → dense head), a global HasVJP built by vjpComp, exposing the pdiv-contracted Jacobian; efficientnetHasVJPAt / efficientnetHasVJPAt_correct restrict it to a point by .toHasVJPAt. The full sixteen-block B0 is efficientnetForwardBFull (EfficientNetFullB0.lean). Spatial dims held constant (stride-1; the separable striding/pooling plumbing is already in CNN.lean). Only the 0 < ε batch-norm hypotheses are required — swish and sigmoid are smooth, so there are no relu-style kink hypotheses anywhere in the block.

References #

• Tan & Le 2019, EfficientNet: Rethinking Model Scaling for CNNs. https://arxiv.org/abs/1905.11946

Architectural choices (documented).

We assemble a representative EfficientNet, all in flattened Vec space, spatial dims held constant (stride-1 throughout — pooling/striding is a separable concern already covered by maxPoolFlat/strided conv in CNN.lean; the VJP plumbing is identical):

stem (3×3 conv-bn-swish, ic → c) → MBConv₁ with SE, residual (stride-1, c → c identity skip) → MBConv₂ with SE, no skip (channel change c → cout) → globalAvgPool (cout·h·w → cout) → dense head (cout → nClasses)

MBConv₁ is the headline block: a genuine squeeze-excite gate (seBlockFull) inside an identity residual. MBConv₂ exercises the channel-changing path (no skip). Both blocks are smooth everywhere (swish + sigmoid + convs + bn + SE), so only the 0 < ε batch-norm hypotheses are needed — no relu-style kink hypotheses.