Documentation

LeanMlir.Proofs.Nets.ConvNeXt.ConvNeXtFullT

The full ConvNeXt-T — `[3,3,9,3]`, forward + whole-net VJP + graph + faithfulness #

The real ConvNeXt-T spec, beside the two-block scalar-LN net of ConvNeXt.lean:

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 → head LN → dense.

Depth-k within a stage — CnxBlockParamsCh 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).
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 window x[4i..4i+3] of the paper's pad-0 Conv2d(4, s=4)) + the .flatConvStride4F graph node.

GELU/LN/conv are smooth, so the whole-net VJP convNextForwardTChHasVJP is global, under the 23 LN positivities (1 stem + 18 block + 3 downsample + head) and no other hypothesis — as are efficientnetForwardBFullHasVJP and vitForwardKVHasVJP. The ConvNeXtFold/ConvNeXtChainClose param bridges are dim-generic and cover all 18 blocks verbatim; the downsample conv W/b reuse the proven stride-2 bridges.

References #

Liu et al. 2022, A ConvNet for the 2020s (ConvNeXt). https://arxiv.org/abs/2201.03545

ConvNeXt specifies channel_layer_norm — h·w statistics per example, each over the c channels at one spatial position, per-channel [c] affine. See ChannelLN.lean for the primitive and for why Route A needs no new op and no new VJP.

The two-block scalar-LN net in ConvNeXt.lean (convNextForward, convNextBlock, convNextBlockBody) is a different net: it backs the comparator and a book chapter. The LN sites here are 1 stem + 18 block + 3 downsample (all channel-LN) + the head LN after GAP.

noncomputable def Proofs.cnxBodyWith (gf : GeluForm) {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 vjpComp chain.

Equations

Proofs.cnxBodyWith gf LN Wdw bdw Wex bex Wpr bpr γls = Proofs.layerScale γls ∘ Proofs.flatConv Wpr bpr ∘ gf.map (cExp * h * w) ∘ Proofs.flatConv Wex bex ∘ LN ∘ Proofs.depthwiseFlat Wdw bdw

Instances For

theorem Proofs.cnxBodyWith_differentiable {gf : GeluForm} {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 gf LN Wdw bdw Wex bex Wpr bpr γls)

noncomputable def Proofs.cnxBodyWithHasVJP (gf : GeluForm) {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 gf 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 convNextBlockBodyHasVJP 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.

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

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.

Equations

Proofs.cnxGlsCh p k = p.γls (Proofs.StableHLO.chanIdx c h w k)

Instances For

noncomputable def Proofs.cnxBlockChW (gf : GeluForm) {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

Proofs.cnxBlockChW gf p = Proofs.residual (Proofs.cnxBodyWith gf (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

theorem Proofs.cnxBlockChW_differentiable {gf : GeluForm} {c cExp h w kH kW : ℕ} (p : CnxBlockParamsCh c cExp h w kH kW) (hε : 0 < p.εn) :

Differentiable ℝ (cnxBlockChW gf p)

noncomputable def Proofs.cnxBlockChWHasVJP (gf : GeluForm) {c cExp h w kH kW : ℕ} (p : CnxBlockParamsCh c cExp h w kH kW) (hε : 0 < p.εn) :

HasVJP (cnxBlockChW gf p)

Equations

One or more equations did not get rendered due to their size.

Instances For

noncomputable def Proofs.convNextStageChK (gf : GeluForm) {c cExp h w kH kW : ℕ} (k : ℕ) :

(Fin k → CnxBlockParamsCh 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

Proofs.convNextStageChK gf 0 x_2 = fun (v : Proofs.Vec (c * h * w)) => v
Proofs.convNextStageChK gf k.succ ps = (Proofs.convNextStageChK gf k fun (i : Fin k) => ps i.succ) ∘ Proofs.cnxBlockChW gf (ps 0)

Instances For

theorem Proofs.convNextStageChK_three {gf : GeluForm} {c cExp h w kH kW : ℕ} (a b d : CnxBlockParamsCh c cExp h w kH kW) (v : Vec (c * h * w)) :

convNextStageChK gf 3 ![a, b, d] v = cnxBlockChW gf d (cnxBlockChW gf b (cnxBlockChW gf a v))

A three-block stage, unrolled.

theorem Proofs.convNextStageChK_nine {gf : GeluForm} {c cExp h w kH kW : ℕ} (a₁ a₂ a₃ a₄ a₅ a₆ a₇ a₈ a₉ : CnxBlockParamsCh c cExp h w kH kW) (v : Vec (c * h * w)) :

convNextStageChK gf 9 ![a₁, a₂, a₃, a₄, a₅, a₆, a₇, a₈, a₉] v = cnxBlockChW gf a₉ (cnxBlockChW gf a₈ (cnxBlockChW gf a₇ (cnxBlockChW gf a₆ (cnxBlockChW gf a₅ (cnxBlockChW gf a₄ (cnxBlockChW gf a₃ (cnxBlockChW gf a₂ (cnxBlockChW gf a₁ v))))))))

A nine-block stage, unrolled.

theorem Proofs.convNextStageChK_differentiable {gf : GeluForm} {c cExp h w kH kW : ℕ} (k : ℕ) (ps : Fin k → CnxBlockParamsCh c cExp h w kH kW) :

(∀ (i : Fin k), 0 < (ps i).εn) → Differentiable ℝ (convNextStageChK gf k ps)

noncomputable def Proofs.convNextStageChKHasVJP (gf : GeluForm) {c cExp h w kH kW : ℕ} (k : ℕ) (ps : Fin k → CnxBlockParamsCh c cExp h w kH kW) :

(∀ (i : Fin k), 0 < (ps i).εn) → HasVJP (convNextStageChK gf k ps)

Equations

One or more equations did not get rendered due to their size.
Proofs.convNextStageChKHasVJP gf 0 x_3 x_4 = Proofs.identityHasVJP (c * h * w)

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.

ε : ℝ
γ : Vec cin
β : Vec cin
W : Kernel4 cout cin 2 2
b : Vec cout

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

Proofs.cnxDownChW h w p = Proofs.flatConvStride2 p.W p.b ∘ Proofs.chanLNTensor3 cin (2 * h) (2 * w) p.ε p.γ p.β

Instances For

theorem Proofs.cnxDownChW_differentiable (h w : ℕ) {cin cout : ℕ} (p : CnxDownParamsCh cin cout) (hε : 0 < p.ε) :

Differentiable ℝ (cnxDownChW h w p)

noncomputable def Proofs.cnxDownChWHasVJP (h w : ℕ) {cin cout : ℕ} (p : CnxDownParamsCh cin cout) (hε : 0 < p.ε) :

HasVJP (cnxDownChW h w 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 (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 official implementation and timm both have the stem and head LNs: 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).

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 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 rowLNVecFlat_differentiable and HasVJP are already proven and whose graph mirror is rowLN_affine_eq.
hγ : Vec 768
hβ : Vec 768
Wd : Mat 768 nC
bd : Vec nC

Instances For

noncomputable def Proofs.convNextForwardTCh (gf : GeluForm) {nC : ℕ} (w : CnxTWeightsCh nC) (x : Vec (3 * 224 * 224)) :

Vec nC

The channel-LN ConvNeXt-T forward (3×224² → nC). Nested-application form, so the graph faithfulness convNextFwdGraphTCh_faithful closes stage by stage. nC is a binder: the Imagenette artifacts run it at 10 and the convnextin_* ImageNet artifacts at 1000. It has no drop-path; the *drop* artifacts compute another function.

Equations

One or more equations did not get rendered due to their size.

Instances For

noncomputable def Proofs.convNextForwardTChHasVJP (gf : GeluForm) {nC : ℕ} (w : CnxTWeightsCh nC) (hsε : 0 < w.sε) (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.hε) :

HasVJP (dense w.Wd w.bd ∘ rowLNVecFlat 1 768 w.hε w.hγ w.hβ ∘ globalAvgPoolFlat 768 7 7 ∘ convNextStageChK gf 3 w.s4 ∘ cnxDownChW 7 7 w.d3 ∘ convNextStageChK gf 9 w.s3 ∘ cnxDownChW 14 14 w.d2 ∘ convNextStageChK gf 3 w.s2 ∘ cnxDownChW 28 28 w.d1 ∘ convNextStageChK gf 3 w.s1 ∘ chanLNTensor3 96 56 56 w.sε w.sγ w.sβ ∘ flatConvStride4 w.sW w.sb)

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.

Equations

One or more equations did not get rendered due to their size.

Instances For

theorem Proofs.convNextForwardTCh_differentiable {gf : GeluForm} {nC : ℕ} (w : CnxTWeightsCh nC) (hsε : 0 < w.sε) (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.hε) :

Differentiable ℝ (dense w.Wd w.bd ∘ rowLNVecFlat 1 768 w.hε w.hγ w.hβ ∘ globalAvgPoolFlat 768 7 7 ∘ convNextStageChK gf 3 w.s4 ∘ cnxDownChW 7 7 w.d3 ∘ convNextStageChK gf 9 w.s3 ∘ cnxDownChW 14 14 w.d2 ∘ convNextStageChK gf 3 w.s2 ∘ cnxDownChW 28 28 w.d1 ∘ convNextStageChK gf 3 w.s1 ∘ chanLNTensor3 96 56 56 w.sε w.sγ w.sβ ∘ flatConvStride4 w.sW w.sb)

The chain is differentiable everywhere (the 23 LayerNorm positivities only) — the Differentiable peer of convNextForwardTChHasVJP, on the same twelve-factor chain, which batchMapHasVJP asks for beside the HasVJP when the net is lifted over a batch (ConvNeXtWholeBackCertifiedTieB.lean).

theorem Proofs.convNextForwardTCh_eq_chain {gf : GeluForm} {nC : ℕ} (w : CnxTWeightsCh nC) (x : Vec (3 * 224 * 224)) :

convNextForwardTCh gf w x = (dense w.Wd w.bd ∘ rowLNVecFlat 1 768 w.hε w.hγ w.hβ ∘ globalAvgPoolFlat 768 7 7 ∘ convNextStageChK gf 3 w.s4 ∘ cnxDownChW 7 7 w.d3 ∘ convNextStageChK gf 9 w.s3 ∘ cnxDownChW 14 14 w.d2 ∘ convNextStageChK gf 3 w.s2 ∘ cnxDownChW 28 28 w.d1 ∘ convNextStageChK gf 3 w.s1 ∘ chanLNTensor3 96 56 56 w.sε w.sγ w.sβ ∘ flatConvStride4 w.sW w.sb) x

The nested↔chain bridge: convNextForwardTCh w x equals the twelve-factor ∘ chain convNextForwardTChHasVJP is stated on.

theorem Proofs.convNextForwardTChHasVJP_correct {gf : GeluForm} {nC : ℕ} (w : CnxTWeightsCh nC) (hsε : 0 < w.sε) (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.hε) (x : Vec (3 * 224 * 224)) (dy : Vec nC) (i : Fin (3 * 224 * 224)) :

(convNextForwardTChHasVJP gf w hsε h1 hd1 h2 hd2 h3 hd3 h4 hhε).backward x dy i = ∑ j : Fin nC, pdiv (convNextForwardTCh gf 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 cnxLnFwdSite 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 reassoc / unassoc casts the render uses; 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])).

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 cnxHeadLnFwdSite a mirror rather than a peer. Note: 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

def Proofs.StableHLO.cnxBlockChGraphW (gf : GeluForm) (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.

Equations

One or more equations did not get rendered due to their size.

Instances For

theorem Proofs.StableHLO.cnxBlockChGraphW_faithful {gf : GeluForm} (pfx epsStr : String) {c cExp h w kH kW : ℕ} (p : CnxBlockParamsCh c cExp h w kH kW) (e : SHlo (c * h * w)) :

den (cnxBlockChGraphW gf pfx epsStr p e) = cnxBlockChW gf p (den e)

def Proofs.StableHLO.cnxStageChGraphK (gf : GeluForm) (epsStr : String) {c cExp h w kH kW : ℕ} (base k : ℕ) :

(Fin k → CnxBlockParamsCh 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

One or more equations did not get rendered due to their size.
Proofs.StableHLO.cnxStageChGraphK gf epsStr x✝¹ 0 x_5 x✝ = x✝

Instances For

theorem Proofs.StableHLO.cnxStageChGraphK_den {gf : GeluForm} (epsStr : String) {c cExp h w kH kW : ℕ} (base k : ℕ) (ps : Fin k → CnxBlockParamsCh c cExp h w kH kW) (e : SHlo (c * h * w)) :

den (cnxStageChGraphK gf epsStr base k ps e) = convNextStageChK gf 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 (gf : GeluForm) (epsStr : String) {nC : ℕ} (w : CnxTWeightsCh nC) (x : Vec (3 * 224 * 224)) :

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: 1 stem + 18 block + 3 downsample + head.

Equations

One or more equations did not get rendered due to their size.

Instances For

theorem Proofs.StableHLO.convNextFwdGraphTCh_faithful {gf : GeluForm} (epsStr : String) {nC : ℕ} (w : CnxTWeightsCh nC) (x : Vec (3 * 224 * 224)) :

den (convNextFwdGraphTCh gf epsStr w x) = convNextForwardTCh gf w x

Channel-LN forward faithfulness — the [3,3,9,3] channel-LN graph denotes convNextForwardTCh. One rw per stage: cnxStageChGraphK_den / cnxDownChGraphW_faithful at the stages and downsamples, chanLNGraph_faithful at the stem LN and headLNGraph_faithful at the head.