EfficientNet-B0 train step rendered ENTIRELY from the verified AST (batched) #
The Chapter-7 peer of mnv2TrainStepFaithfulVPaper (MobileNetV2Render.lean), for the committed
full-16-MBConv EfficientNet-B0 (262 params, the real [t,c,n,s,k] B0 spec). Unlike MobileNetV2/
ResNet34 (per-example SHlo index, batch threaded only in emit), EfficientNet emits true
batch-norm, which couples the batch — so the whole net lives at the batched index N·(c·h·w)
(StableHLO.batchOp/bnBatchF/the batched backward + param-SGD ops, all Item B).
The SE wrinkle (vs MobileNetV2's relu6 blocks). Each MBConv has a squeeze-excite gate
x ⊙ sigmoid(dense W₂ (swish (dense W₁ (GAP x)))), and the committed trainer trains all 4 SE dense
params. The fused batchOp seBlock / seBackBatched give the forward value + the SE input
cotangent but NOT the SE param grads, so the renderer un-fuses SE: it keeps the fused seBlock
for the forward out and ADDITIONALLY emits the un-fused gate subnet s = batchOp gap → e1 = batchOp dense W₁ → z = batchOp swish → e2 = batchOp dense W₂ (only to expose s/e1/z/e2); the SE param grads chain
seReduceB → sigmoidBack(e2) → denseWeightSgdB/denseBiasSgdB (W₂) → denseRowBack(W₂) → swishBack(e1) → denseWeightSgdB/denseBiasSgdB (W₁), and dx reuses the fused seBackBatched. Activations
are swish (smooth, no relu6 kink), the head GAP-back uses the batched gapBackBatched.
Render is value-independent (skel erases values), so placeholder zeros + lr := 0/ε := 0 are
passed; the emitted lrStr/epsStr literals carry the real values.
The batched index: why every N below is B, not 1 #
This renderer used to build the graph at the batch-unit index N = 1 and let pretty B supply
the real batch. That was a disclosed convention, and it was sound for the ops where the batch is a
parallel index — batchOp's den is batchMap N (denOp op), which at N = 1 is the per-example
op, exactly what the emit applies across the batch. It was not sound for the ops where the batch
is a reduction axis, and there are ten of those in this graph: bnBatchF and bnBatchBack reduce
μ/var over [0,2,3], and the whole *SgdB param family (bnGammaSgdB, bnBetaSgdB,
dense{Weight,Bias}SgdB, conv{,Strided}WeightSgdB, depthwise{,Strided}WeightSgdB) sums the
per-example gradient over Fin N. At N = 1 each of those dens describes a ONE-EXAMPLE function
while the emitted text reduces over all B — the node and its render were different functions, the
op-level form of the two-writers bug.
The fix is to put the whole graph at N := B, where those dens are honest. What blocked that was
not the batch-coupled ops (their emitters discard N and use B, so they render identically at any
N) but the pointwise ones: swishF/swishBack/sigmoidBack/addV/sub carry only the SHlo
index and emit tensor<B×n> from it, so at the batched index N·s they emit tensor<B×(N·s)> —
which does not even typecheck against its own operand. Hence the batched forms, which all separate
the batch N from the per-example emit width n:
BatchableOp.swish/softmaxRow/denseRowBack— descriptors,den = batchMap N (denOp op). Sound because what they lift is a FIXED function: swish and row-softmax carry no data, anddenseRowBackcarriesW, a parameter shared by every example.swishBackB/sigmoidBackB— their own constructors, NOT descriptors, because their VJPfun x dy i => dy i * deriv (x i)depends on the saved pre-activation, which varies per example.batchMap Nof that would denote "every example shares one example's activation". They carry the whole-batchx : Vec (N*n)instead — which is what the emittedxNameactually holds.addVB/subB— pointwise binary, via the binarybatched2tag.
softmaxRow's m and denseRowBack's rows are NOT the batch — they are rows per example (ViT
uses m := 197 tokens; a classifier head has one logit row), and they stay 1 here.
The artifact is byte-identical across this change, which is what proves the EMIT side
behaviour-preserving. It cannot witness the den side — the render is value-independent, so a
descriptor holding the wrong saved activation would render exactly the same bytes. That half is
carried by the rfl faithfulness theorems in StableHLO.lean.
Saved forward SSA names a block's backward + SGD passes reference.
- code : String
- o : String
- ec : String
- en : String
- er : String
- dc : String
- dn : String
- dr : String
- se : String
- s : String
- e1 : String
- z : String
- e2 : String
- pc : String
The block's BN layers in forward order:
(BN-input SSA, channels, spatial side). The AdamW render turns each into abnBatchMeanB/bnBatchVarBpair — the batch statistics a batch-BN train step has to hand back so the host can EMA them into the eval forward's frozen stats. The SGD render has no such outputs and ignores this field.Order is expand-BN → depthwise-BN → project-BN, with the expand entry ABSENT for the no-expand block (b1, t = 1), because that is the layout the driver's
bnChannelsmetadata and@efficientnet_fwd_evalread positionally. Getting it wrong is silent: the arities still match and the wrong layer's statistics simply flow into the wrong eval slot.The second component is the layer's stat prefix —
@efficientnet_fwd_evaltakes%{prefix}mu/%{prefix}varthere, and the AdamW train step hands the matching batch μ/var back from the SAME entry. So the eval signature, the eval BN sites and the train step's stat outputs all come off this one list; there is deliberately no parallel 49-entry table.- stE : String
- stD : String
- stP : String
Instances For
Equations
Every leaf of the backward ends at a parameter, and there are exactly two things it can emit:
the un-fused gradient (adam := true) or the fused SGD update θ − lr·g (false). The
*SgdB_eq_grad theorems say den (xSgdB …) = θ − lr · den (xGradB …) by rfl, and
tests/TestBatchedEmitTie.lean checks the emit side of the same statement: each *GradB render is
a byte-PREFIX of its *SgdB peer's, the tail being exactly the const-lr / multiply / subtract.
These six helpers are what let ONE backward traversal serve both renders. The alternative — a second
copy of the 16-MBConv backward for AdamW — is the double-writer disease one level down, in code
rather than in artifacts, and it is how efficientnet_train_step ended up with two emitters
computing different functions in the first place (§2a-quinquies).
lrStr is threaded but unused in adam mode: the AdamW render's learning rate is the runtime
%lr argument, not a baked literal. The placeholder values are irrelevant either way — skel
erases them, so the render is value-independent.
Total MBConv blocks = the reference's totalDrop, i.e. the ramp DENOMINATOR is this minus 1.
Equations
Instances For
The block indices (0-based) that carry a drop site, in signature order. This is the single
source: enetDropSig maps over it to build the %dp<i> inputs, the traversal passes the same
i at each eFwd call site, and the driver reads it to know how many scales to supply and at
which ramp index. The two routes fail LOUDLY if they disagree — an entry with no call site
leaves an unused input (arity mismatch at the driver), and a call site with no entry emits an
undeclared %dp<i> (rejected by the lowerer). Neither is silent, which is the §2m property.
Instances For
The number of per-example drop-path scale inputs a stochastic-depth render takes.
Instances For
The %dp<i>: tensor<Bxf32> inputs, appended to a render's signature when stochastic depth is
on. Empty when off, which is what keeps gate 1 byte-identical.
Equations
- One or more equations did not get rendered due to their size.
Instances For
The classifier's input width — EfficientNet-B0's head channel count, i.e. the GAP output and
hence the dropout mask's per-example width. Independent of nClasses.
Equations
Instances For
The %do: tensor<B×1280xf32> input, appended when classifier dropout is on. Empty when off,
which is what keeps the inertness gate byte-identical.
⚠ It goes after enetDropSig, i.e. dead last in every signature. Two independent reasons,
and the second is the one that bites: a parameter inserted mid-list captures an existing
positional slot (the mnv2 convBias failure, §2m) and the driver walks these signatures
positionally; and the drop-mask tail is what the DP shim shards by COUNT from the end
(n_shard_tail), so a per-example input placed before them would be counted as one of them.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Names are stored WITHOUT the leading % and shapes as List Nat rather than a rendered
tensor<…> string, because the AdamW render needs both forms: %{nm}/%{nm}m/%{nm}v for the
moment slots, and the raw dimensions for the emitted Adam ops (adamMNextF's ds). The SGD
render's emitted text is unchanged — ty ds reproduces exactly the strings that used to be
stored.
Every channel width EfficientNet-B0 uses as a conv bias — stem 32, each block's mid and
oc off the [t,c,n,s,k] table, head 1280. One list feeding all four zeroBiasPrelude calls,
so a convBias := false render cannot declare the constants in one artifact and not another.
⚠ NOT the SE widths. SE's two 1×1 convs are followed by an ACTIVATION, not BN, so nothing absorbs their biases and the reference carries them (§2m); they stay real parameters and are never bound to a zero constant. The audit's rule — a rank-1 kind-2 param after a rank-4 kernel — excludes them because SE's params are rank-2, which is why enet's +21,008 gap closed exactly on the first attempt.
Equations
Instances For
Every SSA name the EfficientNet-B0 forward produces, plus the 49-entry BN stat layout.
efficientnetFwd{,Eval}FaithfulV return just logits; the train steps additionally consume the
stem/head names and the 16 block records on the way back.
- code : String
- stc : String
- stn : String
- str : String
- hc : String
- hn : String
- hr : String
- gap : String
- cin : String
⭐⭐ THE CLASSIFIER'S ACTUAL INPUT —
gapwith classifier dropout OFF, thedropoutBoutput with it ON. It exists as its own field, rather than every consumer readinggap, because there are TWO consumers and one of them is easy to miss:- the dense forward, which obviously reads it; and
- ⚠⚠ the dense WEIGHT gradient,
∂L/∂W = Σ_b dy_b ⊗ (input_b)— which reads the dense's input, i.e. the DROPPED activation, not the pooled one.
Feeding
dnWthe undroppedgaptype-checks, trains, descends, and is wrong on the one parameter dropout acts through. It is invisible to every ones-mask gate this feature has, because atmask ≡ 1the two values are equal. That is handoff §0.10's LayerScale-γ defect in the same shape, and the reason it is a named field is the carry-forward that record asks for: when an op is spliced into a chain, list every consumer of the value it displaced. - logits : String
The 49 BN layers as
(BN-input SSA, stat prefix, channels, spatial side), stem → blocks in forward order → head. Single source for the eval signature, the eval BN sites and the AdamW train step's returned batch statistics — seeEFwd.bns.- sst : String
- hst : String
Instances For
Equations
@efficientnet_fwd rendered ENTIRELY from the verified AST — 263 inputs (%x plus the 262
params in enetSig order), returning logits [B, nClasses]. Shares enetFwdChain with the
train step, so it is a byte-identical PREFIX of efficientnet_train_step.mlir, ending exactly
where the loss begins. Replaces the hand-written emitter in tests/TestEfficientNetFwd.lean.
Equations
- One or more equations did not get rendered due to their size.
Instances For
@efficientnet_fwd_eval rendered ENTIRELY from the verified AST — the inference forward,
every BN site consuming frozen per-channel running stats (the bnEval descriptor, den =
batchMap N bnPerChannelEvalTensor3) instead of reducing statistics out of its activation.
Same 262 params in the same order, plus the 98 stat inputs (49 BN layers × μ/var, interleaved
per layer in bnChannels order): 361 inputs.
This is the eval partner of efficientnet_adam_train_step, whose returned batch μ/var the
driver EMAs into exactly these slots — and both sides of that contract now come off one
bns list rather than two independently-written ones.
Equations
- One or more equations did not get rendered due to their size.
Instances For
EfficientNet-B0 (full 16-MBConv) SGD train step rendered ENTIRELY from the verified AST, at
the batched index N·(c·h·w). Every emitted line is pretty of a verified SHlo node. Strided
stem 3×3/s2 (3→32, 224→112) → b1 (no-expand) → b2..b16 (4 strided downsamples 112→7, 9 residual
skips, 2 no-skip widenings) → 1×1 conv-bn-swish head (320→1280) → GAP → dense (1280→nClasses).
The cotangent is plain softmax − onehot with the batch mean folded into lrStr — so the
committed lrStr = 0.05 is an effective 1.6 on the mean loss. That is a tuned value, not a
slip (runs/efficientnet_verified_crop_gpu1.log: 40.6% → 87.81% over 80 epochs, matching
README's 87.58%); the AdamW render below spells the mean explicitly instead.
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/efficientnet_<variant>_train_step.mlir, the entry point is
@efficientnet_<variant>_train_step, and LEAN_MLIR_VARIANT selects it.
All three must agree, or the shim refuses the call outright ("entry mismatch") rather than
running the wrong graph — which is exactly what it did the first time R34's DP render kept the
single-device name (§2b-quater). Deriving the name here is what stops it drifting from the
#eval paths below; the #guards at the bottom pin those literal paths against this function.
B = 32 is deliberately unsuffixed, so the two existing artifacts keep their names and bytes.
Same convention as r34AdamVariant.
Equations
- One or more equations did not get rendered due to their size.
Instances For
EfficientNet-B0 AdamW train step rendered from the verified AST. The certified peer of the
hand-written tests/TestEfficientNetTrain.lean render that efficientnet-verified-adam has
been training on.
Same backward as efficientnet_train_step (enetBackAll, one traversal) but taking the
un-fused gradients, each fed to the proven AdamW triple. Two things differ from the SGD
render and both are load-bearing:
- the cotangent is label-smoothed (α = 0.1, K = nClasses) with an explicit ÷B, where the
SGD render is plain CE with the mean folded into
lr. Measured against the hand-written AdamW emitter, this is the same gap ViT had; get it wrong and the tie fails in a way that looks like a bug in the gradient ops. - it returns the BN running statistics — batch μ/var per BN layer,
bnBatchMeanB/bnBatchVarBrecomputed from that layer's BN input — which the host EMAs into@efficientnet_fwd_eval's frozen stats. The SGD render has no such outputs.
Interface: 889 in (%x, 262 θ, 262 m, 262 v, %lr/%bc1/%bc2, 98 running-stat slots,
%onehot) / 887 out (262 θ', 262 m', 262 v', %loss/%bc1/%bc2, 98 batch stats) —
positionally identical to the hand-written render, so trainAdamSched's packed [θ|m|v]
protocol is unchanged.
Unlike ViT's, this tie can pin the forward bit-exactly: EfficientNet has BatchNorm, so the returned batch statistics are a whole-net forward fingerprint no gradient touches.