CNN + CIFAR render half — conv train-step text as pretty of proven graphs #
The conv-net peers of MlpRender.lean: cnnTrainStepText (Chapter-3 MNIST CNN),
cifarTrainStepText (Chapter-4 CIFAR CNN), the 8-conv cifar8* family (SGD /
Nesterov / AdamW tails, the batched …FaithfulB peers, bf16 / fp8) and cifar8Bn*. Every
line that feeds a returned parameter is pretty of a denoted SHlo node, names threaded as in
the MLP render. Hand-written text: the signatures and constants, the report-only %loss block
appended by cnnTrainStepText and by the packed-optimizer ([θ|m|v]) renders, and the
packed renders' %bc1/%bc2 passthroughs; none of it feeds a parameter. The forward is rendered
flat: each .flatConvF/.maxPoolF token reshapes flat→NCHW internally and back at its boundary
(emitTok, in StableHLO/Pretty.lean), so the names pretty exposes are flat. The
*_train_step_tied_certified capstones in CnnFold / CifarFold / Cifar8StepTie /
Cifar8BnStepTie state every output's den as the certified SGD step at the cotangent the rendered
backward chain feeds it. The #eval writers that produce the verified_mlir/cnn_* / cifar* artifacts are in
CnnArtifacts.lean, which nothing imports.
MNIST-CNN train step rendered from the verified AST. The peer of
mlpTrainStepText for the conv net: the forward, the backward chain
(dotOut/selectPos/maxPoolBack/convBack) and ALL ten parameter SGD updates are
pretty of denoted SHlo nodes —
the dense head via weightSgd/biasSgd, the conv layers via the
convWeightSgd/convBiasSgd ops. So every line that feeds a returned parameter is
pretty of a denoted node; the appended report-only %loss block is hand-written and feeds
nothing. CnnFold.cnn_train_step_tied_certified states each output's den as the certified
SGD step at the rendered chain cotangent; cnn_W5_tied_totalloss folds W₅ to ∂CE/∂W₅.
Cotangents (%dy/dy4/dy3/dac2/dhc2/dac1/dhc1) are rendered once and
shared as operand leaves; operand/lr/weight VALUES are skel-erased, so these
placeholders print identically to the live graphs the den theorems use. Dims: h,w
are the post-pool spatial sizes, the image is 2h × 2w.
Equations
- One or more equations did not get rendered due to their size.
Instances For
CIFAR-CNN (Chapter 4, no-BN) train step rendered ENTIRELY from the verified AST.
The deeper, two-spatial-scale peer of cnnTrainStepText: the forward, the
backward chain (dotOut/selectPos/maxPoolBack/convBack, twice through) and all
14 parameter SGD updates are pretty of denoted SHlo nodes — the dense head via
weightSgd/biasSgd, the four conv layers via the convWeightSgd/convBiasSgd
ops (reused from cnn, NO new ops). Every line that feeds a returned parameter is
pretty(provenNode), and CifarFold.cifar_train_step_tied_certified states each output's
den as the certified SGD step at the rendered chain cotangent.
Dims: h,w are the final pooled spatial sizes; image 4h×4w, stage-2 spatial 2h×2w.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Deeper 8-conv CIFAR (cifar8, no-BN) train step rendered ENTIRELY from the verified
AST. The 4-stage peer of cifarTrainStepText ((conv→relu)×2→pool ×4, 3 dense;
22 params). Backward chain (dotOut/selectPos/maxPoolBack/convBack, four stages)
and all 22 param SGD ops are pretty of denoted nodes — conv via convWeightSgd/
convBiasSgd, dense via weightSgd/biasSgd (NO new ops).
Cifar8Tie.cifar8_train_step_tied_certified states each output's den as the certified SGD
step at the rendered chain cotangent. h,w are the final pooled sizes; stage spatials build
up ×2 per pool (s4=2h, s3=4h, s2=8h, s1=16h; image 16h×16w).
Equations
- One or more equations did not get rendered due to their size.
Instances For
Which optimizer tail the cifar8 render emits. All three share ONE forward,
backward and un-fused-gradient body, and one packed [θ|m|v] signature — 71 in / 69 out for
every variant — so the ablation section's "SGD several ways" is genuinely the same net with the
optimizer swapped, and a reader can diff the artifacts to see only the tail move.
- adamw : CifarOpt
θ' = θ − lr·(m̂/(√v̂+ε)) − lr·wd·θ, and both moments live. - sgd : CifarOpt
θ' = θ − lr·g;m/vride through untouched. - nesterov : CifarOpt
v' = μ·v + g,θ' = θ − lr·(g + μ·v'); velocity in thevslot,muntouched.
Instances For
Equations
Equations
- One or more equations did not get rendered due to their size.
Instances For
The optimizer's name in an artifact's banner.
Equations
- Proofs.StableHLO.CifarOpt.adamw.label = "AdamW"
- Proofs.StableHLO.CifarOpt.sgd.label = "SGD"
- Proofs.StableHLO.CifarOpt.nesterov.label = "Nesterov momentum"
Instances For
cifar8 AdamW train step rendered from the verified AST. Identical forward/backward to
cifar8TrainStepText; the 22 fused SGD ops are replaced by 22 un-fused param
gradients (convWeightGrad/convBiasGrad/weightGrad/biasGrad) each feeding the three
proven AdamW ops (adamWParamF/adamMNextF/adamVNextF, denoting Proofs.adamWStep).
Signature matches the packed trainAdamSched protocol byte for byte:
(x, θ×22, m×22, v×22, %lr, %bc1, %bc2, onehot) → (θ'×22, m'×22, v'×22, loss, bc1, bc2).
β₁/β₂/ε/wd are baked as in-body constants (%b1 %ob1 %b2 %ob2 %eps %wd), so the conv
biases are named %cb1…%cb8 — %b1 is β₁.
Two lines are outside the proven surface, both marked in the emitted text: the scalar
%loss (report-only; the kit has no rank-0 loss op and it does not feed the update) and
the %bc1/%bc2 passthroughs. The mean-loss 1/B on the cotangent IS proven — it is
scaleF, not hand-written text. Unlike the SGD render it cannot be folded into lr,
because lr is a runtime scalar here.
Instances For
cifar8AdamTrainStepBText — the batched peer of cifar8AdamTrainStepText.
Same net, same three optimizers, same packed [θ|m|v] signature. The difference is the op
FAMILY: this render carries the batch in the Lean type (SHlo (B*(c*h*w))) and uses the
batched constructors, where the …V render is per-example (SHlo (c*h*w)) with pretty B
broadcasting. Naming follows the ImageNet renderers, where …RenderB is exactly this
migration done once per net (ResNet34RenderB, MobileNetV2RenderB, …).
Why it exists: bf16, and rehearsal. The bf16 ops were built for ImageNet, which is
entirely on the batched family — so bf16 twins exist for convBackBatched/denseRowBack and
do NOT exist for the per-example convBack/dotOut that …V uses. Rather than write two
CIFAR-only bf16 ops (twins of convBack and dotOut) that ImageNet would never run, this moves CIFAR
onto the ops ImageNet already uses. bf16 then drops in for the whole step, forward AND
backward, with zero new verified ops — and CIFAR becomes a real rehearsal for ImageNet
instead of a parallel dialect.
The migration is semantically free, not a re-derivation. Both families denote the same
proven VJP — den's .convBack arm is (conv2dHasVJP3 W b).backward v … and its
.convBackBatched arm is batchMap N (… (conv2dHasVJP3 W b)).backward (fun _ => 0) …. The
only difference is the primal argument, and convBack_faithful records why it is free: conv
is linear, so this is a global VJP — the input-VJP ignores the primal. That is why
.convBack's primal argument is simply dropped below rather than threaded.
Faithful by CONSTRUCTION, like every render here: the AST is built only from verified
constructors, so pretty(provenGraph) needs no new proof. Nothing in this function is
hand-written MLIR except the report-only %loss and the %bc1/%bc2 passthroughs, as in
the …V peer.
Parameters are NOT batched — only activations are. The optimizer tail (optTail) is
therefore untouched and shared verbatim with …V.
Instances For
Deeper 8-conv CIFAR-BN (cifar8-bn) train step rendered ENTIRELY from the verified
AST. The per-channel-BatchNorm peer of cifar8TrainStepText ((conv→BN→relu)×2→pool
×4, 3 dense; 38 params). Pure reuse — NO new ops and NO new proof: conv via
convWeightSgd/convBiasSgd, BN via bnGammaSgd/bnBetaSgd, dense via weightSgd/
biasSgd; Cifar8BnTie.cifar8Bn_train_step_tied_certified states every output's den as
the certified SGD step at the rendered chain cotangent, by the generic lemmas
(SgdNode.conv{W,B}_den, SgdNode.bn{Gamma,Beta}_den, SgdNode.dense{W,B}_den) per layer. Forward + BN-back proof-rendered via bnPerChannelF/
bnPerChannelBack. h,w final pooled; stage spatials s4=2h…s1=16h.
opt selects the optimizer tail, and it changes the INTERFACE, unlike the
no-BN cifar8AdamTrainStepText where all three variants share one packed signature:
opt | entry | interface | tail |
|---|---|---|---|
none | @cifar8_bn_train_step | 40 in / 38 out | the 38 fused *Sgd ops, lr a baked literal |
some o | @cifar8_bn_{adam,mom,sgd}_train_step | 119 in / 117 out | un-fused *Grad + optTail o, packed [θ|m|v], %lr runtime |
Three things branch, not just the tail:
the cotangent (fused folds the batch mean into lrStr; packed cannot, lr is a runtime
arg, so it emits an explicit scaleF invB plus a report-only %loss), the conv bias names
(%b1..%b8 fused, but AdamW bakes β₁/β₂ as %b1/%b2, so packed renames to %cb1..%cb8),
and the signature/return. Everything else — the whole forward, the whole backward, all 38
gradients — is shared verbatim.
Instances For
cifar8BnTrainStepBText — the batched peer of cifar8BnTrainStepText.
Same net, same three optimizers, same packed [θ|m|v] 38-parameter signature. What moves is
the op FAMILY for the convolutions and the dense head: this render carries the batch in the
Lean type (SHlo (B*(c*h*w))) and uses the batched constructors, where the …V render is
per-example with pretty B broadcasting. It is cifar8AdamTrainStepBText with BatchNorm
spliced in, and it exists for one reason: bf16 on the NORMALIZED net.
Why the migration is what unlocks bf16. The 27 bf16 ops were built for ImageNet, which
is entirely batched, so bf16 twins exist for convBackBatched/convWeightGradB and do NOT
exist for the per-example convBack/dotOut the …V render uses. Chapter 4's precision lever
could therefore only measure bf16 on the un-normalized net — the one that NaNs under AdamW at
fp32 — which is the hardest case and the wrong one for licensing Chapter 5. This render puts
bf16 on the net that actually has BatchNorm, with zero new verified ops.
BatchNorm stays PER-EXAMPLE, and that is deliberate. bnPerChannelF/bnPerChannelBack
reduce over [2,3], not [0,2,3]: each example is normalized by its own statistics. Three
consequences, all of them the point:
- It is the same net as
…V. Swapping to the batch-coupledbnBatchFwould change the FUNCTION, not just the emit, and the f32-vs-bf16 comparison would be confounded by a BatchNorm semantics change. The op family is the only variable here. - The eval forward is unchanged. Per-example BN needs no running statistics — train and
eval normalize identically — so
cifar8w_bn_fwd.mliris reused verbatim and cannot drift from what the f32 arm evaluates against. Batch BN would need μ/var returned from the step, a host-side EMA, and a frozen-stat eval render; the packed protocol has nowhere to put them. - It matches mixed precision as everyone practises it. Every bf16 net in this repo —
R34, R50, MNv2, MNv4, EfficientNet — keeps BN in f32 and rounds only the convolutions and
the dense layers (
ResNet34RenderB's block backwards are the pattern: the conv backward branches onbf16, thebnBatchBackabove it does not). A bf16 BatchNorm would measure a recipe Chapter 5 does not use.
The BN nodes are per-example SHlo trees inside a pretty B render, exactly as the head's
rows := 1 ops are in cifar8AdamTrainStepBText: each pretty node is an independent
tree, linked to the next only by the SSA name, so a per-example BN node and a batched conv
node compose in the emitted text without composing in the Lean types. The emitted BN fragment
is byte-identical to the one …V emits.
Faithful by CONSTRUCTION, like every render here: the AST is built only from verified
constructors, so pretty(provenGraph) needs no new proof. Nothing is hand-written MLIR except
the report-only %loss and the %bc1/%bc2 passthroughs, exactly as in both peers.
Parameters are NOT batched — only activations are — so the optimizer tail (optTail) is
untouched and shared verbatim with both peers. Conv biases are %cb1…%cb8: %b1 is β₁.
Instances For
The plain-SGD cifar8 render: cifar8AdamTrainStepText with opt := .sgd, so the
forward, backward and all 22 un-fused gradients are shared verbatim with the AdamW render and
only the tail differs. Entry @cifar8_sgd_train_step, 71 in / 69 out.
Equations
- One or more equations did not get rendered due to their size.
Instances For
The Nesterov cifar8 render (opt := .nesterov), μ baked at 0.9.
Entry @cifar8_mom_train_step, 71 in / 69 out.
Equations
- One or more equations did not get rendered due to their size.