Documentation

LeanMlir.Proofs.Nets.ConvNeXt.ConvNeXtBackCertifiedTie

§B: the ConvNeXt block-body backward chain IS the certified VJP #

cnxBlockBodyBack (ConvNeXtBackChains.lean) is the hand-composed reverse of the ConvNeXt block body, written in the per-op backward maps of BackwardMaps.lean. This file closes §B for the net the repo ships, whose LayerNorm is the CHANNEL LN chanLNTensor3: the chain IS the certified input-gradient VJP, in the SAME non-batched vocabulary.

The block body is layerScale ∘ project ∘ GELU ∘ expand ∘ LN ∘ depthwise, whose certified VJP backward applies the reverses in order LS.back → PR.back → GE.back → EX.back → LN.back → D.back. cnxBlockBodyBack is the exact peer chain depthwiseFlatBack ∘ lnB ∘ convFlatBack Wex ∘ geluB ∘ convFlatBack Wpr ∘ lsB. The tie pins the layer-scale and GELU backs (lsB/geluB) to the certified backwards at the exact saved activations, fills the LN slot with the concrete chanLNTensor3Back (tied by chanLNTensor3Back_eq_chanLN_vjp), and ties the two 1×1 convs + the depthwise to their certified input-VJPs via the leaf gates (convFlatBack_eq_vjp_backward, depthwiseFlatBack_eq_vjp_backward). b1-free: the per-example body is the non-batched object the chain reverses, so there is no batchMap reconciliation. 3-axiom-clean.

chanLNTensor3Back (ChannelLNBack.lean) is not an abstract slot — it is a concrete five-factor chain, the row map conjugated by the forward's four layout permutations. So it owes a tie of its own: that the chain IS chanLNTensor3_has_vjp's backward. That is what this section proves.

The proof is piecewise, and every piece is already in the repo:

The tie is β-free: the certified backward does not depend on the LN bias, and neither does the chain — the +β translation's VJP is the identity, which is why chanLNTensor3Back never took a β in the first place.

source
theorem Proofs.bn_grad_input_eq_vjp_backward {n : ℕ} (ε γ β : ℝ) (hε : 0 < ε) (x dy : Vec n) :
bn_grad_input n ε γ x dy = (bn_has_vjp n ε γ β hε).backward x dy

The concrete three-term BN/LN input gradient IS the certified VJP backward. bn_grad_input is not rfl-equal to (bn_has_vjp …).backward — the witness is built through a rw [bnForward_eq_compose] cast — but both reduce to the canonical ∑ pdiv form (bn_input_grad_correct and .correct). The function-level peer of bnBack_faithful_fn.

source
theorem Proofs.transposeFlat_has_vjp_backward_eq (m n : ℕ) (v : Vec (m * n)) (dy : Vec (n * m)) :
(transposeFlat_has_vjp m n).backward v dy = StableHLO.transposeFlat n m dy

The flat transpose's VJP backward is the flat transpose back — transpose_has_vjp's backward is fun i j => dY j i, and hasVJPMat_to_hasVJP reads it at the row-major split, so this is definitional. The permutation adjoint the channel-LN conjugation needs, alongside reassoc{Fwd,Back}_has_vjp_backward_eq.

source
theorem Proofs.layerNormVec_has_vjp_backward_eq {D : ℕ} (ε : ℝ) (hε : 0 < ε) (γ β x dy : Vec D) :
(layerNormVec_has_vjp D ε γ β hε).backward x dy = bn_grad_input D ε 1 x (diagBack γ dy)

The vector-LN row backward is bn_grad_input after the γ scale. layerNormVec is (+β) ∘ layerScale γ ∘ LN(1,0), so its VJP applies: the bias translation's identity backward, then diagBack γ, then the LN input gradient at γ = 1. The +β drops out — this is where the whole channel-LN backward story becomes β-free.

source
theorem Proofs.rowLNVecFlat_has_vjp_backward_eq {s c : ℕ} (ε : ℝ) (hε : 0 < ε) (γ β : Vec c) (X dy : Vec (s * c)) :
(rowLNVecFlat_has_vjp s c ε γ β hε).backward X dy = rowLNVecFlatBack s c ε γ X dy

The rowwise vector-LN backward is rowLNVecFlatBack. rowLNVecFlat_has_vjp is the rowwise_has_vjp_mat lift of the row VJP through hasVJPMat_to_hasVJP, and rowLNVecFlatBack is perRowFlatPR of the row's closed form — the same per-row map at the same row of the saved input, so this is the row lemma read at each (row, col).

source
theorem Proofs.chanLNTensor3Back_eq_chanLN_vjp {c h w : ℕ} (ε : ℝ) (hε : 0 < ε) (γ β : Vec c) (x : Vec (c * h * w)) :
chanLNTensor3Back c h w ε γ x = (chanLNTensor3_has_vjp c h w ε γ β hε).backward x

THE §2n §B TIE: the channel-LN backward chain IS the certified VJP. chanLNTensor3Back — the hand-composed reverse of chanLNTensor3 — equals (chanLNTensor3_has_vjp …).backward at every saved input and cotangent, so the chain is the certified gradient.

Proof: the witness is a term-mode vjp_comp chain, so its backward unfolds to the nested chain; rewrite its five factors (two reassoc collapses, two transposes by rfl, the row map through bn_grad_input). 3-axiom-clean.

source
theorem Proofs.cnxBodyWithChanLNBack_eq_vjp {c cExp h w kHd kWd : ℕ} (hkHd : 2 * ((kHd - 1) / 2) + 1 = kHd) (hkWd : 2 * ((kWd - 1) / 2) + 1 = kWd) (Wdw : DepthwiseKernel c kHd kWd) (bdw : Vec c) (εn : ℝ) (hεn : 0 < εn) (γn βn : Vec c) (Wex : Kernel4 cExp c 1 1) (bex : Vec cExp) (Wpr : Kernel4 c cExp 1 1) (bpr : Vec c) (γls v : Vec (c * h * w)) :
cnxBlockBodyBack Wdw Wex Wpr (chanLNTensor3Back c h w εn γn (depthwiseFlat Wdw bdw v)) ((layerScale_has_vjp γls).backward ((flatConv Wpr bpr ∘ gelu (cExp * h * w) ∘ flatConv Wex bex ∘ chanLNTensor3 c h w εn γn βn ∘ depthwiseFlat Wdw bdw) v)) ((gelu_has_vjp (cExp * h * w)).backward ((flatConv Wex bex ∘ chanLNTensor3 c h w εn γn βn ∘ depthwiseFlat Wdw bdw) v)) = (cnxBodyWith_has_vjp ⋯ (chanLNTensor3_has_vjp c h w εn γn βn hεn) Wdw bdw Wex bex Wpr bpr γls).backward v

The §B channel-LN BODY tie: the block-body backward chain = the certified VJP, for the net the repo ships. cnxBlockBodyBack with its LayerNorm slot filled by the CONCRETE chanLNTensor3Back (at the saved post-depthwise activation) and its layer-scale / GELU slots pinned to the certified backwards equals (cnxBodyWith_has_vjp (chanLNTensor3 …) …).backward.

Note what fills the LN slot: not a certified object but the concrete five-factor chain, which has to go through chanLNTensor3Back_eq_chanLN_vjp to earn its place. The proof rewrites the two 1×1 conv leaves and the depthwise leaf through their gates, rewrites the LN chain through its tie, and the rest matches definitionally. 3-axiom-clean.

source
theorem Proofs.cnxBlockChBack_eq_vjp {c cExp h w kHd kWd : ℕ} (hkHd : 2 * ((kHd - 1) / 2) + 1 = kHd) (hkWd : 2 * ((kWd - 1) / 2) + 1 = kWd) (p : CnxBlockParamsCh c cExp h w kHd kWd) (hε : 0 < p.εn) (v : Vec (c * h * w)) :
residual (cnxBlockBodyBack p.Wdw p.Wex p.Wpr (chanLNTensor3Back c h w p.εn p.γn (depthwiseFlat p.Wdw p.bdw v)) ((layerScale_has_vjp (cnxGlsCh p)).backward ((flatConv p.Wpr p.bpr ∘ gelu (cExp * h * w) ∘ flatConv p.Wex p.bex ∘ chanLNTensor3 c h w p.εn p.γn p.βn ∘ depthwiseFlat p.Wdw p.bdw) v)) ((gelu_has_vjp (cExp * h * w)).backward ((flatConv p.Wex p.bex ∘ chanLNTensor3 c h w p.εn p.γn p.βn ∘ depthwiseFlat p.Wdw p.bdw) v))) = (cnxBlockChW_has_vjp p hε).backward v

The §B channel-LN BLOCK tie (residual-wrapped). cnxBlockChW is residual of the body, so the block backward chain is residual (cnxBlockBodyBack …) and equals (cnxBlockChW_has_vjp …).backward — the additive skip's backward being dy. Immediate from the body tie. With this, the channel-LN net's §B coverage matches the scalar net's: body, block, and (new, because the LN slot is no longer abstract) the LayerNorm op itself.