Documentation

LeanMlir.Proofs.Nets.ConvNeXt.ConvNeXtFullT

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"):

  1. Depth-k within a stageCnxBlockParamsCh bundles 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).
  2. Downsample boundariescnxDownChW = flatConvStride2(2×2) ∘ channel-LN; both VJPs existed.
  3. 4×4/s4 patchify stemflatConvStride4 (= decimate ∘ decimateOdd ∘ stride-1 SAME conv, StridedConv.lean: the left-aligned window x[4i..4i+3] of the paper's pad-0 Conv2d(4, s=4)) + the flatConvStride4F token.

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_normh·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.

noncomputable def Proofs.cnxBodyWith {c cExp h w kH kW : } (LN : Vec (c * h * w)Vec (c * h * w)) (Wdw : DepthwiseKernel c kH kW) (bdw : Vec c) (Wex : Kernel4 cExp c 1 1) (bex : Vec cExp) (Wpr : Kernel4 c cExp 1 1) (bpr : Vec c) (γls : Vec (c * h * w)) :
Vec (c * h * w)Vec (c * h * w)

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
Instances For
    theorem Proofs.cnxBodyWith_diff {c cExp h w kH kW : } {LN : Vec (c * h * w)Vec (c * h * w)} (hLN : Differentiable LN) (Wdw : DepthwiseKernel c kH kW) (bdw : Vec c) (Wex : Kernel4 cExp c 1 1) (bex : Vec cExp) (Wpr : Kernel4 c cExp 1 1) (bpr : Vec c) (γls : Vec (c * h * w)) :
    Differentiable (cnxBodyWith LN Wdw bdw Wex bex Wpr bpr γls)
    noncomputable def Proofs.cnxBodyWith_has_vjp {c cExp h w kH kW : } {LN : Vec (c * h * w)Vec (c * h * w)} (hLN : Differentiable LN) (vLN : HasVJP LN) (Wdw : DepthwiseKernel c kH kW) (bdw : Vec c) (Wex : Kernel4 cExp c 1 1) (bex : Vec cExp) (Wpr : Kernel4 c cExp 1 1) (bpr : Vec c) (γls : Vec (c * h * w)) :
    HasVJP (cnxBodyWith LN Wdw bdw Wex bex Wpr bpr γls)

    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
      structure Proofs.CnxBlockParamsCh (c cExp h w kH kW : ) :

      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).

      Instances For
        noncomputable def Proofs.cnxGlsCh {c cExp h w kH kW : } (p : CnxBlockParamsCh c cExp h w kH kW) :
        Vec (c * h * w)

        The per-channel layer-scale expanded to the flat map (cnxGls's peer).

        Equations
        Instances For
          noncomputable def Proofs.cnxBlockChW {c cExp h w kH kW : } (p : CnxBlockParamsCh c cExp h w kH kW) :
          Vec (c * h * w)Vec (c * h * w)

          The packaged ConvNeXt block: residual of the shared body at chanLNTensor3.

          Equations
          Instances For
            theorem Proofs.cnxBlockChW_diff {c cExp h w kH kW : } (p : CnxBlockParamsCh c cExp h w kH kW) ( : 0 < p.εn) :
            noncomputable def Proofs.cnxBlockChW_has_vjp {c cExp h w kH kW : } (p : CnxBlockParamsCh c cExp h w kH kW) ( : 0 < p.εn) :
            Equations
            • One or more equations did not get rendered due to their size.
            Instances For
              noncomputable def Proofs.convNextStageChK {c cExp h w kH kW : } (k : ) :
              (Fin kCnxBlockParamsCh c cExp h w kH kW)Vec (c * h * w)Vec (c * h * w)

              Depth-k channel-LN stage fold (head recursion — block 0 runs first).

              Equations
              Instances For
                theorem Proofs.convNextStageChK_diff {c cExp h w kH kW : } (k : ) (ps : Fin kCnxBlockParamsCh c cExp h w kH kW) :
                (∀ (i : Fin k), 0 < (ps i).εn)Differentiable (convNextStageChK k ps)
                noncomputable def Proofs.convNextStageChK_has_vjp {c cExp h w kH kW : } (k : ) (ps : Fin kCnxBlockParamsCh c cExp h w kH kW) :
                (∀ (i : Fin k), 0 < (ps i).εn)HasVJP (convNextStageChK k ps)
                Equations
                Instances For
                  structure Proofs.CnxDownParamsCh (cin cout : ) :

                  The stage-boundary downsample's parameters: the LN affine is Vec cin, over the PRE-downsample channel width.

                  Instances For
                    noncomputable def Proofs.cnxDownChW (h w : ) {cin cout : } (p : CnxDownParamsCh cin cout) :
                    Vec (cin * (2 * h) * (2 * w))Vec (cout * h * w)

                    Stage-boundary downsample: 2×2/s2 conv ∘ channel-LN (cin@2h×2w → cout@h×w).

                    Equations
                    Instances For
                      theorem Proofs.cnxDownChW_diff (h w : ) {cin cout : } (p : CnxDownParamsCh cin cout) ( : 0 < p.ε) :
                      noncomputable def Proofs.cnxDownChW_has_vjp (h w : ) {cin cout : } (p : CnxDownParamsCh cin cout) ( : 0 < p.ε) :
                      Equations
                      • One or more equations did not get rendered due to their size.
                      Instances For
                        structure Proofs.CnxTWeightsCh (nC : ) :

                        All ConvNeXt-T parameters. Every LN affine is a Vec, and BOTH the stem LN (/ : Vec 96) and the head LN (/ : 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 3CnxBlockParamsCh 96 384 56 56 7 7
                        • d1 : CnxDownParamsCh 96 192
                        • s2 : Fin 3CnxBlockParamsCh 192 768 28 28 7 7
                        • d2 : CnxDownParamsCh 192 384
                        • s3 : Fin 9CnxBlockParamsCh 384 1536 14 14 7 7
                        • d3 : CnxDownParamsCh 384 768
                        • s4 : Fin 3CnxBlockParamsCh 768 3072 7 7 7 7
                        • hε :

                          Head LayerNorm, between GAP and the classifier: ε, then the Vec 768 affine. ⚠ Plain layerNormVec, NOT chanLNTensor3: 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 also rowLNVecFlat 1 768, i.e. ViT's per-token LN at ONE row, whose _diff and _has_vjp are already proven and whose graph mirror is rowLN_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
                          noncomputable def Proofs.convNextForwardTCh {nC : } (w : CnxTWeightsCh nC) (x : Vec (3 * 224 * 224)) :
                          Vec nC

                          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
                            noncomputable def Proofs.convNextForwardTCh_has_vjp {nC : } (w : CnxTWeightsCh nC) (hsε : 0 < w.) (h1 : ∀ (i : Fin 3), 0 < (w.s1 i).εn) (hd1 : 0 < w.d1.ε) (h2 : ∀ (i : Fin 3), 0 < (w.s2 i).εn) (hd2 : 0 < w.d2.ε) (h3 : ∀ (i : Fin 9), 0 < (w.s3 i).εn) (hd3 : 0 < w.d3.ε) (h4 : ∀ (i : Fin 3), 0 < (w.s4 i).εn) (hhε : 0 < w.) :

                            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. w. w. 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
                              theorem Proofs.convNextForwardTCh_differentiable {nC : } (w : CnxTWeightsCh nC) (hsε : 0 < w.) (h1 : ∀ (i : Fin 3), 0 < (w.s1 i).εn) (hd1 : 0 < w.d1.ε) (h2 : ∀ (i : Fin 3), 0 < (w.s2 i).εn) (hd2 : 0 < w.d2.ε) (h3 : ∀ (i : Fin 9), 0 < (w.s3 i).εn) (hd3 : 0 < w.d3.ε) (h4 : ∀ (i : Fin 3), 0 < (w.s4 i).εn) (hhε : 0 < w.) :

                              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).

                              theorem Proofs.convNextForwardTCh_has_vjp_correct {nC : } (w : CnxTWeightsCh nC) (hsε : 0 < w.) (h1 : ∀ (i : Fin 3), 0 < (w.s1 i).εn) (hd1 : 0 < w.d1.ε) (h2 : ∀ (i : Fin 3), 0 < (w.s2 i).εn) (hd2 : 0 < w.d2.ε) (h3 : ∀ (i : Fin 9), 0 < (w.s3 i).εn) (hd3 : 0 < w.d3.ε) (h4 : ∀ (i : Fin 3), 0 < (w.s4 i).εn) (hhε : 0 < w.) (x : Vec (3 * 224 * 224)) (dy : Vec nC) (i : Fin (3 * 224 * 224)) :
                              (convNextForwardTCh_has_vjp w hsε h1 hd1 h2 hd2 h3 hd3 h4 hhε).backward x dy i = j : Fin nC, pdiv (convNextForwardTCh w) x i j * dy j

                              Correctness on convNextForwardTCh itself (via the bridge).

                              def Proofs.StableHLO.chanLNGraph (gN btN epsStr : String) {c h w : } (ε : ) (γ β : Vec c) (e : SHlo (c * h * w)) :
                              SHlo (c * h * w)

                              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
                                theorem Proofs.StableHLO.chanLNGraph_faithful (gN btN epsStr : String) {c h w : } (ε : ) (γ β : Vec c) (e : SHlo (c * h * w)) :
                                den (chanLNGraph gN btN epsStr ε γ β e) = chanLNTensor3 c h w ε γ β (den e)
                                def Proofs.StableHLO.headLNGraph (gN btN epsStr : String) {c : } (ε : ) (γ β : Vec c) (e : SHlo (1 * c)) :
                                SHlo (1 * c)

                                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
                                Instances For
                                  theorem Proofs.StableHLO.headLNGraph_faithful (gN btN epsStr : String) {c : } (ε : ) (γ β : Vec c) (e : SHlo (1 * c)) :
                                  den (headLNGraph gN btN epsStr ε γ β e) = rowLNVecFlat 1 c ε γ β (den e)
                                  def Proofs.StableHLO.cnxBlockChGraphW (pfx epsStr : String) {c cExp h w kH kW : } (p : CnxBlockParamsCh c cExp h w kH kW) (e : SHlo (c * h * w)) :
                                  SHlo (c * h * w)

                                  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
                                    theorem Proofs.StableHLO.cnxBlockChGraphW_faithful (pfx epsStr : String) {c cExp h w kH kW : } (p : CnxBlockParamsCh c cExp h w kH kW) (e : SHlo (c * h * w)) :
                                    den (cnxBlockChGraphW pfx epsStr p e) = cnxBlockChW p (den e)
                                    def Proofs.StableHLO.cnxStageChGraphK (epsStr : String) {c cExp h w kH kW : } (base k : ) :
                                    (Fin kCnxBlockParamsCh c cExp h w kH kW)SHlo (c * h * w)SHlo (c * h * w)

                                    Depth-k channel-LN stage graph fold — block base+1 first, prefixes b{base+1}_.

                                    Equations
                                    Instances For
                                      theorem Proofs.StableHLO.cnxStageChGraphK_den (epsStr : String) {c cExp h w kH kW : } (base k : ) (ps : Fin kCnxBlockParamsCh c cExp h w kH kW) (e : SHlo (c * h * w)) :
                                      den (cnxStageChGraphK epsStr base k ps e) = convNextStageChK k ps (den e)
                                      def Proofs.StableHLO.cnxDownChGraphW (pfx epsStr : String) (h w : ) {cin cout : } (p : CnxDownParamsCh cin cout) (e : SHlo (cin * (2 * h) * (2 * w))) :
                                      SHlo (cout * h * w)

                                      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
                                        theorem Proofs.StableHLO.cnxDownChGraphW_faithful (pfx epsStr : String) (h w : ) {cin cout : } (p : CnxDownParamsCh cin cout) (e : SHlo (cin * (2 * h) * (2 * w))) :
                                        den (cnxDownChGraphW pfx epsStr h w p e) = cnxDownChW h w p (den e)
                                        def Proofs.StableHLO.convNextFwdGraphTCh (epsStr : String) {nC : } (w : CnxTWeightsCh nC) (x : Vec (3 * 224 * 224)) :
                                        SHlo nC

                                        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
                                          theorem Proofs.StableHLO.convNextFwdGraphTCh_faithful (epsStr : String) {nC : } (w : CnxTWeightsCh nC) (x : Vec (3 * 224 * 224)) :

                                          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.