The FULL ConvNeXt-T — [3,3,9,3], forward + whole-net VJP + graph + faithfulness #
Scales the ch9 representative (1×1 stem + 2 blocks at one scale) to the real ConvNeXt-T
spec, closing the "full-architecture" gap in planning/archive/convnext_close.md:
4×4/s4 patchify stem (3→96, 224→56) → stem-LN → stage1 (3 blocks @96/56²) → downsample (LN + 2×2/s2 conv 96→192) → stage2 (3 @192/28²) → ds (192→384) → stage3 (9 @384/14²) → ds (384→768) → stage4 (3 @768/7²) → GAP → dense.
Per the handoff recipe (planning/archive/convnext_close.md §"Scaling handoff"):
- Depth-k within a stage —
CnxBlockParamsChbundles the block 10-tuple;convNextStageChK (k) (ps : Fin k → CnxBlockParamsCh …)folds blocks head-first with VJP by induction — the ViT depth-k recipe, simpler here (same-shape blocks within a stage). - Downsample boundaries —
cnxDownChW=flatConvStride2(2×2) ∘ channel-LN; both VJPs existed. - 4×4/s4 patchify stem —
flatConvStride4(= decimate ∘ decimateOdd ∘ stride-1 SAME conv,StridedConv.lean: the left-aligned windowx[4i..4i+3]of the paper's pad-0Conv2d(4, s=4)) + theflatConvStride4Ftoken.
GELU/LN/conv are smooth, so the whole-net VJP is GLOBAL (unconditional except the 22 LN
positivities) — ConvNeXt-T joins efficientnetForwardB_full_has_vjp and vitForwardKV.
The ConvNeXtClose/ConvNeXtChainClose param bridges are dim-generic and cover all 18 blocks
verbatim; the downsample conv W/b reuse the proven stride-2 bridges.
⚠ §2n (2026-07-31): the scalar-LN twin of this chain is GONE. Until then every definition
here had a layerNormForward peer — one mean and one variance over the whole c·h·w map with
scalar γ/β — which is what the repo shipped before §2m flipped ConvNeXt to its real channel
LayerNorm. CnxBlockParams, cnxBlockW, convNextStageK, CnxDownParams, cnxDownW,
CnxTWeights, convNextForwardT/TC and their graph section were deleted once the float
bridges (their last live consumers) had …Ch peers. If you are chasing a dangling reference to
one of those names, it was retired, not moved. planning/archive/xla_pjrt_handoff.md §2n has the
checklist and what the drop did and did not touch.
ConvNeXt specifies channel_layer_norm — h·w statistics per example, each over the c
channels at one spatial position, per-channel [c] affine. See ConvNeXtChannelLN.lean for the
primitive and for why Route A needs no new op and no new VJP.
§2m built this as a PARALLEL chain beside the scalar-LN one it superseded, so that flipping the
net could not change what the then-committed convNextForwardTC denoted (MobileNetV2RenderB's
reason, §2f). §2n then DROPPED that scalar chain — once its last live consumers (the float
bridges) had channel-LN peers, a retired chain that still elaborates is one more thing to drift
(§2a's lesson). What is here is what ships. ⛔ The Nets/ConvNeXt/ConvNeXt.lean ch9
representative (convNextForward, convNextBlock, convNextBlockBody) is a different thing and
SURVIVED the drop: it backs the Diderot comparator and a book chapter.
Two deviations §2m found by the parameter count NOT matching, both invisible to a VJP argument: the reference's 22 LN sites are 1 stem + 18 block + 3 downsample, where the pre-§2m net had 18 + 3 + 1 head. So this forward has a stem LN and no head LN. They nearly cancel (+2·768 − 2·96 = +1,344 of 28.6M), which is exactly why a matching parameter count is a decomposition test and not an architecture check.
The ConvNeXt block body with the LN left as a parameter. Written once and instantiated at
chanLNTensor3, so the channel-LN world costs one definition rather than a second copy of
convNextBlockBody's six-piece vjp_comp chain.
Equations
- Proofs.cnxBodyWith LN Wdw bdw Wex bex Wpr bpr γls = Proofs.layerScale γls ∘ Proofs.flatConv Wpr bpr ∘ Proofs.gelu (cExp * h * w) ∘ Proofs.flatConv Wex bex ∘ LN ∘ Proofs.depthwiseFlat Wdw bdw
Instances For
The body VJP, given the LN's. Only the LN carries a hypothesis — gelu is smooth and
conv/layerScale are linear, so this is global exactly as convNextBlockBody_has_vjp is.
Equations
- One or more equations did not get rendered due to their size.
Instances For
One channel-LN ConvNeXt block's 10 parameters. γn/βn are Vec c — 2c floats per LN site,
where the retired scalar-LN spelling had 2 (§2n deleted it).
- Wdw : DepthwiseKernel c kH kW
- bdw : Vec c
- εn : ℝ
- γn : Vec c
- βn : Vec c
- Wex : Kernel4 cExp c 1 1
- bex : Vec cExp
- Wpr : Kernel4 c cExp 1 1
- bpr : Vec c
- γls : Vec c
Instances For
The per-channel layer-scale expanded to the flat map (cnxGls's peer).
Equations
- Proofs.cnxGlsCh p k = p.γls (Proofs.StableHLO.chanIdx c h w k)
Instances For
The packaged ConvNeXt block: residual of the shared body at chanLNTensor3.
Equations
- Proofs.cnxBlockChW p = Proofs.residual (Proofs.cnxBodyWith (Proofs.chanLNTensor3 c h w p.εn p.γn p.βn) p.Wdw p.bdw p.Wex p.bex p.Wpr p.bpr (Proofs.cnxGlsCh p))
Instances For
Equations
- One or more equations did not get rendered due to their size.
Instances For
Depth-k channel-LN stage fold (head recursion — block 0 runs first).
Equations
- Proofs.convNextStageChK 0 x_2 = fun (v : Proofs.Vec (c * h * w)) => v
- Proofs.convNextStageChK k.succ ps = (Proofs.convNextStageChK k fun (i : Fin k) => ps i.succ) ∘ Proofs.cnxBlockChW (ps 0)
Instances For
Equations
- One or more equations did not get rendered due to their size.
- Proofs.convNextStageChK_has_vjp 0 x_3 x_4 = Proofs.identity_has_vjp (c * h * w)
Instances For
Stage-boundary downsample: 2×2/s2 conv ∘ channel-LN (cin@2h×2w → cout@h×w).
Equations
- Proofs.cnxDownChW h w p = Proofs.flatConvStride2 p.W p.b ∘ Proofs.chanLNTensor3 cin (2 * h) (2 * w) p.ε p.γ p.β
Instances For
Equations
- One or more equations did not get rendered due to their size.
Instances For
All ConvNeXt-T parameters. Every LN affine is a Vec, and BOTH the stem LN (sγ/sβ : Vec 96) and the head LN (hγ/hβ : Vec 768) are present — the paper's forward is
patchify → channel_layer_norm → stages → GAP → LN → dense.
⚠⚠ THE HEAD LN CAME BACK 2026-08-30, and the history is the point. The pre-§2m net had
scalar affines, no stem LN and a head LN. §2m/§2n added the stem LN and DELETED the head one,
to match jax/MainConvNeXtImagenet.lean — which was itself missing it. Both references have
both: facebookresearch/ConvNeXt does self.norm(x.mean([-2,-1])) with
nn.LayerNorm(dims[-1], eps=1e-6), and timm's convnext_tiny head is
NormMlpClassifierHead(global_pool → LayerNorm2d(768) → flatten → fc).
▶ The parameter count is the tell and it was sitting in the blurb: ours was 28,587,592 at
K=1000 against timm.create_model('convnext_tiny')'s 28,589,128 — short by exactly
1,536 = 2×768, the head LN's γ and β. The old note "the two nearly cancel in the parameter
count … which is why the count alone never caught it" was right that they nearly cancel and
wrong that the count could not catch it: the residue IS the missing layer, exactly.
⚠ The lesson is §7.2's one net over — we converged on the JAX reference, and the reference was
the thing that was wrong. planning/archive/next_session_execution_and_parity.md §7.1.
- sW : Kernel4 96 3 4 4
- sb : Vec 96
- sε : ℝ
- sγ : Vec 96
- sβ : Vec 96
- s1 : Fin 3 → CnxBlockParamsCh 96 384 56 56 7 7
- d1 : CnxDownParamsCh 96 192
- s2 : Fin 3 → CnxBlockParamsCh 192 768 28 28 7 7
- d2 : CnxDownParamsCh 192 384
- s3 : Fin 9 → CnxBlockParamsCh 384 1536 14 14 7 7
- d3 : CnxDownParamsCh 384 768
- s4 : Fin 3 → CnxBlockParamsCh 768 3072 7 7 7 7
- hε : ℝ
Head LayerNorm, between GAP and the classifier: ε, then the
Vec 768affine. ⚠ PlainlayerNormVec, NOTchanLNTensor3: after GAP the tensor is[768], one row, so the channel LN and the vector LN are the same function and this is the cheaper spelling — it is alsorowLNVecFlat 1 768, i.e. ViT's per-token LN at ONE row, whose_diffand_has_vjpare already proven and whose graph mirror isrowLN_affine_eq— so the head LN needed no new mathematics at all, only a link in the chain. - hγ : Vec 768
- hβ : Vec 768
- Wd : Mat 768 nC
- bd : Vec nC
Instances For
The channel-LN ConvNeXt-T forward (3×224² → nC). Nested-application form, as the scalar
peers, so the graph faithfulness closes by a structural rfl. nC is a binder: the Imagenette
artifacts run it at 10 and the convnextin_* ImageNet artifacts at 1000, and every theorem
about this forward covers both.
Equations
- One or more equations did not get rendered due to their size.
Instances For
The channel-LN ConvNeXt-T has a (correct) VJP — at every input. 23 LayerNorm
positivities: stem + 18 blocks (via the per-stage ∀ i) + 3 downsamples + the head LN,
which this statement composes as rowLNVecFlat 1 768 w.hε w.hγ w.hβ and takes hhε for.
Chain-stated to keep the blocks opaque.
⚠ The count read 22 … no head LN until 2026-09-04 — the pre-2026-08-30 net — and it was
the third place that stale number had been copied to, each copy citing the last as its
justification. ⛔ docstring-checkrefs cannot catch this: it
resolves cited identifiers, and a stale COUNT cites nothing.
Equations
- One or more equations did not get rendered due to their size.
Instances For
The chain is differentiable everywhere (the 23 LayerNorm positivities only) — the
Differentiable peer of convNextForwardTCh_has_vjp, on the same twelve-factor chain, which
batchMap_has_vjp asks for beside the HasVJP when the net is lifted over a batch
(ConvNeXtWholeBackCertifiedTieB.lean).
The nested↔chain bridge (see convNextForwardTCh_eq_chain for why the proof shape matters —
a simp/rfl proof of this statement dies in the kernel on the recursive stage folds).
Correctness on convNextForwardTCh itself (via the bridge).
One channel-LN forward site, mirroring ConvNeXtRender.lnFwdSite at chLN := true
op-for-op: transpose to [h·w, c], normalise each spatial row over its channels at the
scalar identities %one/%zero, apply the real [c] affine, transpose back. The two ▸
transports are the Nat-associativity casts the render spells the same way; den_reassocS
(ConvNeXtChannelLN.lean) is what makes them the math's Mat-split bridge.
Equations
- One or more equations did not get rendered due to their size.
Instances For
The HEAD LN forward site — the paper's norm(x.mean([-2,-1])), restored 2026-08-30.
⭐ It is chanLNGraph with the transposes deleted, and that is not a shortcut: after GAP the
tensor is a single [768] row, so "normalise each spatial row over its channels" and
"normalise the feature vector" are the same function — m = 1. The render emits exactly these
three ops, which is what makes ConvNeXtRender.headLnFwdSite a mirror rather than a peer.
⚠ Indexed SHlo (1 * c); c is the LITERAL 768 at every call site, so 1 * c reduces and no
transport is needed. Do NOT generalise c to a variable without adding one — that is the trap
convNextBackAll's Vec (1 * nClasses) annotations already record.
Equations
- Proofs.StableHLO.headLNGraph gN btN epsStr ε γ β e = Proofs.StableHLO.SHlo.rowBiasF btN β (Proofs.StableHLO.SHlo.rowScaleF gN γ (Proofs.StableHLO.SHlo.lnRowF "%one" "%zero" epsStr ε 1 0 e))
Instances For
The ConvNeXt block graph — the [3,3,9,3] block segment, with chanLNGraph at its LN site
(the retired scalar spelling put a .bnF there).
Equations
- One or more equations did not get rendered due to their size.
Instances For
Depth-k channel-LN stage graph fold — block base+1 first, prefixes b{base+1}_.
Equations
- One or more equations did not get rendered due to their size.
- Proofs.StableHLO.cnxStageChGraphK epsStr x✝¹ 0 x_5 x✝ = x✝
Instances For
Channel-LN downsample graph: channel-LN → 2×2/s2 widening conv.
Equations
- One or more equations did not get rendered due to their size.
Instances For
The channel-LN ConvNeXt-T forward graph (3×224² → nC): patchify stem → stem
channel-LN → the [3,3,9,3] stages with 3 channel-LN + 2×2/s2 downsample boundaries →
GAP → head LN → dense. ⚠ 23 LN sites, not 22: 1 stem + 18 block + 3 downsample + the head
one restored 2026-08-30 (the retired scalar graph had a head LN and no stem LN; §2m/§2n swapped
which one was missing rather than fixing it — the paper has both).
Equations
- One or more equations did not get rendered due to their size.
Instances For
Channel-LN forward faithfulness — the [3,3,9,3] channel-LN graph denotes
convNextForwardTCh. Same rw chain as the scalar apex, with chanLNGraph_faithful where
the bnFs were. The full-architecture apex for the net §2m makes ConvNeXt actually be.