ResNet-50 at TRUE BATCH-NORM — the whole net's forward and graph #
A net-level ℝ forward at the [3,4,6,3] bottleneck ladder, in the world the artifacts run, and the typed graph over it (the second half of this file).
ResNet-50 has only a batched renderer (ResNet50RenderB.lean), so bnBatchLA is the
world of resnet50_fwd, resnet50in160_lambaccdp8x64wxclipbce and everything between.
What is new here, and what is not #
Nothing about the blocks is new. ResNet50BackB0.lean carries all three batched
bottleneck forms at bnBatchLA with their _at VJPs and backward-graph faithfulness. This file
adds the level above, exactly as ResNet34FullB.lean does for r34.
The stem and the head are ResNet-34's, imported rather than re-declared. r34StemB is
generic in {ic oc} and r34HeadB in {c nCls}, and R50's stem (7×7/s2 conv-bn-relu, then He et
al.'s 3×3/s2 max-pool, 3 → 64) and head (GAP then dense) are the same functions at different
widths. So their VJP lemmas are reused verbatim one tier up, the way ResNet50BackB0.lean reuses
all four of ResNet34BackB0.lean's stages. A second r50StemB would be two writers for one fact.
ONE weight record serves both projection forms. The stride-1 projection block (stage 1
block 0) and the strided one (stages 2/3/4 block 0) have identical parameter shapes — 1×1, 3×3,
1×1 and a 1×1 skip — and differ only in which convolutions are strided. R50ProjW is that record;
r50ProjB and r50DownB are the two forwards over it. ResNet-34 needed two records for its two
block kinds.
Conventions this net runs at #
| depth | [3,4,6,3] BOTTLENECKS, 64/256 → 128/512 → 256/1024 → 512/2048 |
| BatchNorm | batch (bnBatchLA, reduce [0,2,3], width N*h*w) |
| activation | relu, THREE kinks per block — two interior and the post-residual one |
| stride-2 padding | symmetric at all seven sites (stem + three downsample 3×3 + three 1×1 skips) |
| stride placement | v1.5: the stride is on the 3×3, not the leading 1×1 |
| stem | 7×7/s2 conv-bn-relu, then 3×3/s2 max-pool |
| head | GAP then dense, generic in the class count |
| artifacts | resnet50_fwd, resnet50in_fwd, resnet50in160_fwd, every resnet50in*_train_step |
THE STRIDE IS ON THE 3×3. r50DownB puts cbReluStridedB on the SECOND convolution, so
the leading 1×1 runs at the INPUT resolution and carries mid channels there until W₂ decimates.
That is ResNet v1.5 / torchvision, which is what jax/MainResnet50Imagenet.lean trains. The v1
placement compiles, trains and descends, and is a different net
(VerifiedSpec.Layer.bottleneckStage's note). Nothing in the types sees the difference.
Stage 1 block 0 is a STRIDE-1 projection, and it is the block with no ResNet-34 analogue.
Channels go 64 → 256 at unchanged resolution, so it needs a projection but not a strided one.
r50DownB cannot be substituted — the halving is in its signature, so that is a shape error and
would be caught. Reaching for the identity form r50IdB is the dangerous one: it is well-typed
only when ic = oc, which is exactly why this block exists.
The spatial size is a BINDER, q, and that is not a stylistic choice. ResNet-50 ships at
TWO resolutions — resnet50in_fwd at 224 and resnet50in160_fwd at 160. The ladder is q, 2q, 4q, 8q with the
input at 32q, so q = 7 is the 224 net and q = 5 the 160 one and ONE statement covers both.
ResNet-34 could pin 56 because it has a single shipped resolution. Every resolution is written as
an explicit nest of 2 * (…) rather than a product like 8 * q: 2 * (4 * q) and 8 * q are
equal Nats and NOT definitionally equal terms at a variable q, and the block signatures demand
the operand at exactly the spelling they name. The render's own q1 … q5 comment records the same
trap on the emitter side.
N stays a variable throughout, as at r34: the batch size is never pinned (the widths are
literals and the resolution is the binder q). On the data-parallel
artifacts the render's N is the PER-REPLICA batch; their BatchNorm is
synchronised, and ResNet50SyncB.lean is this file's twin for them: replica r's forward graph
denotes shard r of resnet50ForwardBFull (R * N) q, this file's forward at the global batch.
The census is 161 updated parameters, which is ResNet50RenderB's own docstring ("161 θ /
161 m / 161 v"): stem 3 (sW, sγ, sβ) + 12 identity bottlenecks × 9 + 4 projection
bottlenecks × 12 + head 2. The records ALSO carry a bias slot per convolution, as ResNet-34's and
MobileNetV2's do: both R50 renders run convBias := false, each conv bias is folded into the
BatchNorm after it and bound to a zeroBiasPrelude zero, so those fields are the convBias := true
census and are ∀-quantified over — bias = 0 is one instance. The 106 running-statistic slots the
render's signature also carries belong to inference and do not appear here, since training-mode
BatchNorm computes its statistics from the batch.
The typed forward graph is the second half of this file — resnet50FwdGraphBFull and
its _faithful, over four per-block-kind graphs at r50FwdChainB's own tokens. Against
the committed bytes: verified_mlir/resnet50_fwd.mlir's signature is 162 arguments = %x + 161
parameters, with 12 projection slots, and every name this file writes (%sW, %sg, %sbt,
%zb64 … %zb2048, %s1b0W1 … %s4b2bt3, %s1b0Wp/%gp/%btp, %Wd, %bd) appears there.
References #
- He, Zhang, Ren, Sun 2016, Deep Residual Learning for Image Recognition. https://arxiv.org/abs/1512.03385
- torchvision. https://github.com/pytorch/vision
Weights of one identity bottleneck: 1×1 (oc → mid), 3×3 (mid → mid), 1×1 (mid → oc),
BatchNorm after each. The third convolution has NO activation — the outer relu comes after
the residual add.
- W₁ : Kernel4 mid oc 1 1
- b₁ : Vec mid
- ε₁ : ℝ
- γ₁ : Vec mid
- β₁ : Vec mid
- W₂ : Kernel4 mid mid 3 3
- b₂ : Vec mid
- ε₂ : ℝ
- γ₂ : Vec mid
- β₂ : Vec mid
- W₃ : Kernel4 oc mid 1 1
- b₃ : Vec oc
- ε₃ : ℝ
- γ₃ : Vec oc
- β₃ : Vec oc
Instances For
Weights of a PROJECTION bottleneck — ic → mid → mid → oc plus the 1×1 option-B skip.
ONE record for BOTH projection forms. Stage 1 block 0 (stride 1, 64 → 256) and stages
2/3/4 block 0 (strided) have identical parameter shapes and differ only in which convolutions
are strided, which is a property of the forward and not of the weights. r50ProjB and
r50DownB are those two forwards.
- W₁ : Kernel4 mid ic 1 1
- b₁ : Vec mid
- ε₁ : ℝ
- γ₁ : Vec mid
- β₁ : Vec mid
- W₂ : Kernel4 mid mid 3 3
- b₂ : Vec mid
- ε₂ : ℝ
- γ₂ : Vec mid
- β₂ : Vec mid
- W₃ : Kernel4 oc mid 1 1
- b₃ : Vec oc
- ε₃ : ℝ
- γ₃ : Vec oc
- β₃ : Vec oc
- Wp : Kernel4 oc ic 1 1
- bp : Vec oc
- εp : ℝ
- γp : Vec oc
- βp : Vec oc
Instances For
Every ResNet-50 parameter: stem (7×7/s2, 3 → 64) + the [3,4,6,3] bottleneck ladder + the dense
head. Generic in the class count, so one statement covers the 10-class Imagenette artifacts and
the 1000-class resnet50in ones.
The field names are ResNet50RenderB's own SSA prefixes (s1b0 … s4b2), so a reader can
match a parameter to its emitted name without a table.
- sW : Kernel4 64 3 7 7
- sb : Vec 64
- sε : ℝ
- sγ : Vec 64
- sβ : Vec 64
- s1b0 : R50ProjW 64 64 256
- s1b1 : R50IdW 64 256
- s1b2 : R50IdW 64 256
- s2b0 : R50ProjW 256 128 512
- s2b1 : R50IdW 128 512
- s2b2 : R50IdW 128 512
- s2b3 : R50IdW 128 512
- s3b0 : R50ProjW 512 256 1024
- s3b1 : R50IdW 256 1024
- s3b2 : R50IdW 256 1024
- s3b3 : R50IdW 256 1024
- s3b4 : R50IdW 256 1024
- s3b5 : R50IdW 256 1024
- s4b0 : R50ProjW 1024 512 2048
- s4b1 : R50IdW 512 2048
- s4b2 : R50IdW 512 2048
- Wd : Mat 2048 nCls
- bd : Vec nCls
Instances For
Batched identity bottleneck relu(F(x) + x), F = projB ∘ cbReluB ∘ cbReluB: 1×1-reduce-relu,
3×3-relu, 1×1-expand with no activation. THREE relu kinks — the two interior ones and the
post-residual outer one — where ResNet-34's basic block has two.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Batched stride-1 projection bottleneck relu(F(x) + proj(x)) — stage 1 block 0 and
nowhere else in the net. The channels change (64 → 256) so a projection is needed; the
resolution does not, so that projection is a plain 1×1 conv-BN. The form with no ResNet-34
analogue: R34's stage 1 runs at ic = oc = 64, where block 0 is an identity block.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Batched strided projection bottleneck — stages 2/3/4 block 0, halving the grid.
v1.5: the stride is on the 3×3 (cbReluStridedB is the SECOND stage), so the leading 1×1
and its BN and relu run at the INPUT resolution 2h × 2w. Both stride-2 sites — the 3×3 and
the 1×1 skip — are SYMMETRIC padding, as the PyTorch-origin convention requires.
Equations
- One or more equations did not get rendered due to their size.
Instances For
The full batch-BN ResNet-50 forward, N*(3*32q*32q) -> N*nCls, at the [3,4,6,3] bottleneck
ladder. Nested-application form, as resnet34ForwardBFull and efficientnetForwardBFull
both are, so a later tie can peel it one block at a time.
q is a binder: q = 7 is resnet50in_fwd and q = 5 is resnet50in160_fwd. The stem and head are ResNet-34's functions at R50's widths.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Identity bottleneck graph: relu(addVB(bn3(conv3(relu(bn2(conv2(relu(bn1(conv1 e))))))), e)).
The skip reuses the block-input subtree e in both addVB operands, as the render does, and
the operand ORDER is the render's (body first) — which is also residual's.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Stride-1 projection bottleneck graph — stage 1 block 0. The skip is a plain 1×1 conv → BN
(.conv, NOT .convStrided), which is the whole point of this form. Both addVB operands are
nontrivial subtrees and both read the block-input subtree e.
The render emits addVB(body, projection) where residualProj proj body adds
proj + body — so this graph is in the RENDER's order and the faithfulness proof carries one
add_comm. The alternative, writing the graph in residualProj's order, would make den
close by rfl and the emitted operand order wrong.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Strided projection bottleneck graph — stages 2/3/4 block 0. v1.5: .convStrided appears at
the 3×3 and at the 1×1 skip, and conv1/bn1/relu1 run at the input resolution
2h × 2w. Both stride-2 sites are SYMMETRIC padding (.convStrided, not .convStridedXla).
Equations
- One or more equations did not get rendered due to their size.
Instances For
Stem graph: 7×7/s2 conv → batch BN → relu → He et al.'s 3×3/s2 max-pool — ResNet-34's
r34StemGraphB, token for token (both name the bias operand biasName false "" oc).
Equations
- Proofs.StableHLO.r50StemGraphB epsStr N h w Ws bs εs γs βs bf16 e = Proofs.StableHLO.r34StemGraphB epsStr N h w Ws bs εs γs βs bf16 e
Instances For
The full batch-BN ResNet-50 forward graph. Block prefixes are ResNet50RenderB's own
(s1b0 … s4b2) and the head's are %Wd/%bd, so the typed graph diffs against
resnet50_fwd and its ImageNet twins name for name. The head graph is ResNet-34's,
unchanged: r34HeadGraphB is generic in {c nCls} and emits the same two tokens. bf16
selects the conv kinds at every site (.convAt / .convStridedAt, StableHLO.PrecisionSwitch),
as resnet34FwdGraphBFull does: false is the f32 artifacts' graph, true the bf16 ones'
(resnet50in_momdp64bf16, resnet50in160_lambaccdp4x128wxclipbcebf16, …).
Equations
- One or more equations did not get rendered due to their size.
Instances For
Forward-graph faithfulness for ResNet-50 at batch BN: the typed graph denotes the
whole-net forward, at either precision (Bf16Erasure: the bf16 kind at the identity rounding
is its f32 peer, so the right-hand side does not move with the flag).