Documentation

LeanMlir.Proofs.Architectures.ChannelLNBack

The channel-LayerNorm backward — the rowwise vector-LN input-VJP and its conjugation #

The backward of the channel LayerNorm (chanLNTensor3, ChannelLN.lean) and of ViT's per-token vector LayerNorm (layerNormVec, LayerNorm.lean), as the ℝ maps the certified ties are stated about. rowLNVecFlatBack is the per-row input-VJP — the per-channel γ scale then the three-term bnGradInput at unit γ, lifted rowwise — and chanLNTensor3Back is that map conjugated by the same four layout permutations the forward uses. ConvNeXtBackCertifiedTie.chanLNTensor3Back_eq_chanLN_vjp proves the conjugation equals the certified chanLNTensor3HasVJP backward, and rowLNVecFlatBack_eq_vecLN_vjp (ViTVecLNBackCertifiedTie.lean) does the same for the row map, so ConvNeXt's LN backward and ViT's share rowLNVecFlatBack. Both ties rest on the three lemmas at the end of this file: bnGradInput is the BN/LN witness's backward, and the vector-LN row and rowwise VJPs reduce to it.

source
noncomputable def Proofs.rowLNVecFlatBack (s c : ℕ) (ε : ℝ) (γ : Vec c) (X : Vec (s * c)) :
Vec (s * c) → Vec (s * c)

The rowwise vector-LN input-VJP. Per spatial row: scale the cotangent by the per-channel γ (layerScale's adjoint is diagBack γ), then the consolidated three-term bnGradInput at γ = 1 over that row's c channels (LN's adjoint = BN's, layerNormForward = bnForward). The +β translation contributes the identity, so it does not appear.

Equations
Instances For
    source
    noncomputable def Proofs.chanLNTensor3Back (c h w : ℕ) (ε : ℝ) (γ : Vec c) (x : Vec (c * h * w)) :
    Vec (c * h * w) → Vec (c * h * w)

    The channel-LN input-VJP (as a function of the cotangent, at a saved input x) — the exact reverse of chanLNTensor3's five factors. A permutation's adjoint is its inverse permutation, so the conjugation comes back unchanged and only the middle flips to rowLNVecFlatBack, read at the TRANSPOSED saved input (the row backward needs its own row's activation). It takes no β: the +β translation's VJP is the identity, and the tie proves the certified backward is β-free too.

    Equations
    • One or more equations did not get rendered due to their size.
    Instances For
      source
      theorem Proofs.bnGradInput_eq_vjp_backward {n : ℕ} (ε γ β : ℝ) (hε : 0 < ε) (x dy : Vec n) :
      bnGradInput n ε γ x dy = (bnHasVJP n ε γ β hε).backward x dy

      The concrete three-term BN/LN input gradient IS the certified VJP backward. bnGradInput is not rfl-equal to (bnHasVJP …).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.layerNormVecHasVJP_backward_eq {D : ℕ} (ε : ℝ) (hε : 0 < ε) (γ β x dy : Vec D) :
      (layerNormVecHasVJP D ε γ β hε).backward x dy = bnGradInput D ε 1 x (diagBack γ dy)

      The vector-LN row backward is bnGradInput 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.rowLNVecFlatHasVJP_backward_eq {s c : ℕ} (ε : ℝ) (hε : 0 < ε) (γ β : Vec c) (X dy : Vec (s * c)) :
      (rowLNVecFlatHasVJP s c ε γ β hε).backward X dy = rowLNVecFlatBack s c ε γ X dy

      The rowwise vector-LN backward is rowLNVecFlatBack. rowLNVecFlatHasVJP is the rowwiseHasVJPMat lift of the row VJP through HasVJPMat.toHasVJP, 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.rowLNBack_affine_eq (s c : ℕ) (ε : ℝ) (γ : Vec c) (X dy : Vec (s * c)) :
      StableHLO.rowLNBackFlat s c ε 1 X (StableHLO.rowScaleFlat s c γ dy) = rowLNVecFlatBack s c ε γ X dy

      The backward peer of rowLN_affine_eq (ChannelLN.lean). Forward, the emitted subtree normalises at the scalar identities %one/%zero and only then applies the real [c] affine, so three denotations collapse onto rowLNVecFlat. Backward it is the same fold one step earlier: the emitted rowScaleF γ applied to the COTANGENT is exactly the per-row diagBack γ that rowLNVecFlatBack folds in, and the LN input gradient then runs at γ = 1.

      β does not appear on either side — the translation's adjoint is the identity, the same β-freeness chanLNTensor3Back_eq_chanLN_vjp has for the certified backward.