Documentation

LeanMlir.Proofs.Nets.ConvNeXt.ConvNeXtFoldGB

The gradient-node fold for ConvNeXt-T at the batched index #

ConvNeXtFoldG.lean holds the three per-example lemmas specific to this net (layer-scale γ, channel-LN γ/β); the other node kinds are shared (GradNodesB). This file is the batched peer, at ConvNeXtRenderB.convNextBackAllB's constructors, from which every ConvNeXt Adam artifact renders.

Against the per-example chain, the batched convBackBatched emits the conv input-VJP's transpose/reverse in the other order from the per-example convBack — commuting ops on disjoint axes — and tests/TestConvNeXtFwdBTie.lean allows exactly that pair and nothing else.

The op table of convnext_adam_train_step.mlir and every convnextin_* train step #

emitted nodelemmaper-example peer it batches
layerScaleChGammaGradB (18 block γ)layerScaleChGammaGradB_denCnxFoldG.layerScaleChGammaGrad_den
convWeightGradB / convBiasGradB (18 expand + 18 project 1×1, + the stem bias)GradNodeB.convWGradB_den / convBGradB_den (GradNodesB)—
depthwiseWeightGradB / depthwiseBiasGradB (18 × 7×7)GradNodeB.depthwiseWGradB_den / GradNodeB.depthwiseBGradB_den (GradNodesB)—
convStridedWeightGradB / convStridedBiasGradB (3 × 2×2/s2 downsample)GradNodeB.convStridedWGradB_den / convStridedBGradB_den (GradNodesB)—
convStride4WeightGradB (patchify stem)GradNodeB.psWGradB_den (GradNodesB)flatConvStride4WeightGradHasVJP, per example
veclnGammaGradB / rowDenseBiasGradB at R = h·w (22 spatial LN sites)chanLnGammaGradB_den / chanLnBetaGradB_denCnxFoldG.chanLnGammaGrad_den / chanLnBetaGrad_den
veclnGammaGradB / rowDenseBiasGradB at R = 1 (the head LN, after GAP)GradNodeB.veclnGammaGradB_den / rowDenseBiasGradB_den_lnbeta—
weightGradB / biasGradB (the classifier)GradNodeB.headWGradB_den / headBGradB_den—
convWeightGradBBf16 / depthwiseWeightGradBBf16 / convStridedWeightGradBBf16 / convStride4WeightGradBBf16 (the bf16 artifacts)Bf16Fold.convWGradBBf16_den and its siblings, Foundation/Bf16GradNodes.leannone — a bf16 node is its own op kind

No new mathematics. Every proof is Finset.sum_congr rfl over the batch and then the per-example bridge at batchSlice n — GradNodeB.denseWGradB_den's shape — because each batched den arm is literally the per-example one under a batch sum. The channel-LN sites add one step: the batched render hands the LN ops batchMap N (chanLNRows c h w) of the saved input and of the cotangent (the [h·w, c] transposed views, lifted per example), and batchSlice_batchMap peels the lift so ChannelLN's permutation argument applies at each slice.

The bf16 artifacts (convnextin_adamwxclipdropbf16, the S/B twins) emit *GradBBf16 constructors, not these nodes: their den rounds the operands and the result once, outside the batch sum. Those are their own op kinds, folded once for every net in Foundation/Bf16GradNodes.lean; CnxTieGB.cnx_net_tiedGB and cnx_net_lossGrad reach them through the renderers' switch (GradNodesBAt, Bf16Erasure: at the identity rounding each bf16 kind denotes its f32 peer here).

One lemma per f32 node kind. Every lemma is ∀ cot, so it holds at whichever cotangent the chain delivers: the f32 AdamW, wx/clip, EMA, drop-path and data-parallel artifacts all emit these *GradB kinds, and the optimizer update that consumes the node is outside these lemmas.

Scope #

source
theorem Proofs.CnxFoldGB.layerScaleChGammaGradB_den {N c h w : ℕ} (xN cotN : String) (x : Vec (N * (c * h * w))) (γ : Vec c) (dy : Vec (N * (c * h * w))) (cc : Fin c) :
StableHLO.den (StableHLO.SHlo.layerScaleChGammaGradB xN x (StableHLO.SHlo.operand cotN dy)) cc = ∑ n : Fin N, ∑ j : Fin (c * h * w), pdiv (fun (γ' : Vec c) => layerScale (fun (k : Fin (c * h * w)) => γ' (StableHLO.chanIdx c h w k)) (StableHLO.batchSlice N (c * h * w) x n)) γ cc j * StableHLO.batchSlice N (c * h * w) dy n j

Batched per-channel layer-scale γ GRADIENT denotes the certified Σ_n gradient. The emitted reduce contracts batch and spatial in one op; den reads it as the batch sum of the per-example dγ_c = Σ_{k : chanIdx k = c} x_k·dy_k. All 18 blocks.

source
theorem Proofs.CnxFoldGB.chanLnGammaGradB_den {N c h w : ℕ} (xN epsStr cotN : String) (ε : ℝ) (β : Vec c) (x : Vec (N * (c * h * w))) (γ : Vec c) (cot : Vec (N * (c * h * w))) (k : Fin c) :
StableHLO.den (StableHLO.SHlo.veclnGammaGradB xN epsStr ε (StableHLO.batchMap N (chanLNRows c h w) x) (StableHLO.SHlo.operand cotN (StableHLO.batchMap N (chanLNRows c h w) cot))) k = ∑ n : Fin N, ∑ j : Fin (c * h * w), pdiv (fun (γ' : Vec c) => chanLNTensor3 c h w ε γ' β (StableHLO.batchSlice N (c * h * w) x n)) γ k j * StableHLO.batchSlice N (c * h * w) cot n j

Batched channel-LN γ GRADIENT denotes the certified Σ_n γ gradient. All 22 spatial sites (1 stem + 18 block + 3 downsample). Two levels: the outer sum is the batch, the inner the h·w rows within one example.

source
theorem Proofs.CnxFoldGB.chanLnBetaGradB_den {N c h w : ℕ} (cotN : String) (ε : ℝ) (γ : Vec c) (x : Vec (N * (c * h * w))) (β : Vec c) (cot : Vec (N * (c * h * w))) (k : Fin c) :
StableHLO.den (StableHLO.SHlo.operand cotN (StableHLO.batchMap N (chanLNRows c h w) cot)).rowDenseBiasGradB k = ∑ n : Fin N, ∑ j : Fin (c * h * w), pdiv (fun (β' : Vec c) => chanLNTensor3 c h w ε γ β' (StableHLO.batchSlice N (c * h * w) x n)) β k j * StableHLO.batchSlice N (c * h * w) cot n j

Batched channel-LN β GRADIENT denotes the certified Σ_n β gradient. The β gradient is the plain two-level row reduce, so the render uses the same rowDenseBiasGradB op ViT's LN β does.

source
def Proofs.CnxFoldGB.ChanLNGammaTiedB (N h w : ℕ) {c : ℕ} (xN epsStr cotN : String) (ε : ℝ) (β : Vec c) (x : Vec (N * (c * h * w))) (γ : Vec c) (cot : Vec (N * (c * h * w))) :
Prop

A batched channel-LN γ gradient node, tied (chanLnGammaGradB_den).

Equations
  • One or more equations did not get rendered due to their size.
Instances For
    source
    def Proofs.CnxFoldGB.ChanLNBetaTiedB (N h w : ℕ) {c : ℕ} (cotN : String) (ε : ℝ) (γ : Vec c) (x : Vec (N * (c * h * w))) (β : Vec c) (cot : Vec (N * (c * h * w))) :
    Prop

    A batched channel-LN β gradient node, tied (chanLnBetaGradB_den).

    Equations
    • One or more equations did not get rendered due to their size.
    Instances For

      Each clause holds, every argument implicit (read off the goal by a step tie's constructor).

      source
      theorem Proofs.CnxFoldGB.chanLNGammaTiedB_holds {N h w c : ℕ} {xN epsStr cotN : String} {ε : ℝ} {β : Vec c} {x : Vec (N * (c * h * w))} {γ : Vec c} {cot : Vec (N * (c * h * w))} :
      ChanLNGammaTiedB N h w xN epsStr cotN ε β x γ cot
      source
      theorem Proofs.CnxFoldGB.chanLNBetaTiedB_holds {N h w c : ℕ} {cotN : String} {ε : ℝ} {γ : Vec c} {x : Vec (N * (c * h * w))} {β : Vec c} {cot : Vec (N * (c * h * w))} :
      ChanLNBetaTiedB N h w cotN ε γ x β cot