ViT at vector-[D] LayerNorm #

The committed ViT render (ViTRender.lean, the GPU-trained ViT-Tiny) carries a vector γ, β : [D] per LN site, spelled scalar-LN(1,0) ∘ per-channel scale γ ∘ + β. The op itself — layerNormVec, its VJP, its per-token lift and its γ/β parameter gradients (vecLNGradGamma / _beta and the vit_render_vecln*_certified step forms) — is in Architectures/LayerNorm, which ConvNeXt's channel-LN shares. This file is ViT's use of it:

§1 transformerBlockV — the vector-LN block, with its VJP composed through the sublayer recipe (biPathMatHasVJP + vjpMatComp + rowwiseHasVJPMat). UNCONDITIONAL except 0 < ε. The depth-k net built from it is ViTDepthK's.
§2 the row-broadcast flat bridges (rowScaleFlat_flat, rowBiasFlat_flat) for the tokens each vector-LN site is spelled with: lnRowF(1,0) → rowScaleF γ → rowBiasF β.
§3 the chain cotangents at vector LN.

References #

Ba, Kiros, Hinton 2016, Layer Normalization. https://arxiv.org/abs/1607.06450

source

noncomputable def Proofs.transformerAttnSublayerV (N heads d_head : ℕ) (ε : ℝ) (γ1 β1 : Vec (heads * d_head)) (Wq Wk Wv Wo : Mat (heads * d_head) (heads * d_head)) (bq bk bv bo : Vec (heads * d_head)) :

Mat N (heads * d_head) → Mat N (heads * d_head)

Attention sublayer with vector-LN: X ↦ X + MHSA(LNᵥ(X)).

Equations

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

Instances For

source

noncomputable def Proofs.transformerMlpSublayerV (gf : GeluForm) (N heads d_head mlpDim : ℕ) (ε : ℝ) (γ2 β2 : Vec (heads * d_head)) (Wfc1 : Mat (heads * d_head) mlpDim) (bfc1 : Vec mlpDim) (Wfc2 : Mat mlpDim (heads * d_head)) (bfc2 : Vec (heads * d_head)) :

Mat N (heads * d_head) → Mat N (heads * d_head)

MLP sublayer with vector-LN: h ↦ h + MLP(LNᵥ(h)).

Equations

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

Instances For

source

noncomputable def Proofs.transformerBlockV (gf : GeluForm) (N heads d_head mlpDim : ℕ) (ε : ℝ) (γ1 β1 : Vec (heads * d_head)) (Wq Wk Wv Wo : Mat (heads * d_head) (heads * d_head)) (bq bk bv bo γ2 β2 : Vec (heads * d_head)) (Wfc1 : Mat (heads * d_head) mlpDim) (bfc1 : Vec mlpDim) (Wfc2 : Mat mlpDim (heads * d_head)) (bfc2 : Vec (heads * d_head)) :

Mat N (heads * d_head) → Mat N (heads * d_head)

Vector-LN transformer block: MLPᵥ-sublayer ∘ attentionᵥ-sublayer — the ViTRender block form.

Equations

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

Instances For

source

noncomputable def Proofs.transformerAttnSublayerVHasVJPMat (N heads d_head : ℕ) (ε : ℝ) (γ1 β1 : Vec (heads * d_head)) (hε : 0 < ε) (Wq Wk Wv Wo : Mat (heads * d_head) (heads * d_head)) (bq bk bv bo : Vec (heads * d_head)) :

HasVJPMat (transformerAttnSublayerV N heads d_head ε γ1 β1 Wq Wk Wv Wo bq bk bv bo)

Attentionᵥ sublayer VJP.

Equations

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

Instances For

source

noncomputable def Proofs.transformerMlpSublayerVHasVJPMat (gf : GeluForm) (N heads d_head mlpDim : ℕ) (ε : ℝ) (γ2 β2 : Vec (heads * d_head)) (hε : 0 < ε) (Wfc1 : Mat (heads * d_head) mlpDim) (bfc1 : Vec mlpDim) (Wfc2 : Mat mlpDim (heads * d_head)) (bfc2 : Vec (heads * d_head)) :

HasVJPMat (transformerMlpSublayerV gf N heads d_head mlpDim ε γ2 β2 Wfc1 bfc1 Wfc2 bfc2)

MLPᵥ sublayer VJP.

Equations

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

Instances For

source

theorem Proofs.transformerBlockV_flat_differentiable {gf : GeluForm} (N heads d_head mlpDim : ℕ) (ε : ℝ) (γ1 β1 : Vec (heads * d_head)) (hε : 0 < ε) (Wq Wk Wv Wo : Mat (heads * d_head) (heads * d_head)) (bq bk bv bo γ2 β2 : Vec (heads * d_head)) (Wfc1 : Mat (heads * d_head) mlpDim) (bfc1 : Vec mlpDim) (Wfc2 : Mat mlpDim (heads * d_head)) (bfc2 : Vec (heads * d_head)) :

Differentiable ℝ fun (v : Vec (N * (heads * d_head))) => (transformerBlockV gf N heads d_head mlpDim ε γ1 β1 Wq Wk Wv Wo bq bk bv bo γ2 β2 Wfc1 bfc1 Wfc2 bfc2 (Mat.unflatten v)).flatten

Flat Diff of the vector-LN block.

source

noncomputable def Proofs.transformerBlockVHasVJPMat (gf : GeluForm) (N heads d_head mlpDim : ℕ) (ε : ℝ) (γ1 β1 : Vec (heads * d_head)) (hε : 0 < ε) (Wq Wk Wv Wo : Mat (heads * d_head) (heads * d_head)) (bq bk bv bo γ2 β2 : Vec (heads * d_head)) (Wfc1 : Mat (heads * d_head) mlpDim) (bfc1 : Vec mlpDim) (Wfc2 : Mat mlpDim (heads * d_head)) (bfc2 : Vec (heads * d_head)) :

HasVJPMat (transformerBlockV gf N heads d_head mlpDim ε γ1 β1 Wq Wk Wv Wo bq bk bv bo γ2 β2 Wfc1 bfc1 Wfc2 bfc2)

Vector-LN block VJP — one vjpMatComp of the two sublayer witnesses.

Equations

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

Instances For

source

theorem Proofs.StableHLO.rowScaleFlat_flat {m n : ℕ} (γ : Vec n) (A : Mat m n) :

rowScaleFlat m n γ A.flatten = Mat.flatten fun (r : Fin m) => layerScale γ (A r)

Row-broadcast scale commutation bridge.

source

theorem Proofs.StableHLO.rowBiasFlat_flat {m n : ℕ} (β : Vec n) (A : Mat m n) :

rowBiasFlat m n β A.flatten = Mat.flatten fun (r : Fin m) (k : Fin n) => A r k + β k

Row-broadcast bias commutation bridge.

source

noncomputable def Proofs.vitCotHV (gf : GeluForm) {Np1 D mlpDim : ℕ} (ε : ℝ) (γ2 : Vec D) (Wfc1 : Mat D mlpDim) (Wfc2 : Mat mlpDim D) (h : Vec (Np1 * D)) (m1 : Vec (Np1 * mlpDim)) (dyOut : Vec (Np1 * D)) :

Vec (Np1 * D)

Cot at the attention-sublayer output h, vector-LN form: dyOut + the decomposed LN₂ input-VJP (rowScaleFlat γ2 then rowLNBackFlat at γ=1).

Equations

Proofs.vitCotHV gf ε γ2 Wfc1 Wfc2 h m1 dyOut i = dyOut i + Proofs.StableHLO.rowLNBackFlat Np1 D ε 1 h (Proofs.StableHLO.rowScaleFlat Np1 D γ2 (Proofs.vitCotLn2 gf Wfc1 Wfc2 m1 dyOut)) i

Instances For

source

noncomputable def Proofs.vitCotAttV (gf : GeluForm) {Np1 D mlpDim : ℕ} (ε : ℝ) (γ2 : Vec D) (Wo : Mat D D) (Wfc1 : Mat D mlpDim) (Wfc2 : Mat mlpDim D) (h : Vec (Np1 * D)) (m1 : Vec (Np1 * mlpDim)) (dyOut : Vec (Np1 * D)) :

Vec (Np1 * D)

Cot at the SDPA output, vector-LN form.

Equations

Proofs.vitCotAttV gf ε γ2 Wo Wfc1 Wfc2 h m1 dyOut = Proofs.StableHLO.rowDenseBackFlat Np1 D D Wo (Proofs.vitCotHV gf ε γ2 Wfc1 Wfc2 h m1 dyOut)

Instances For

source

noncomputable def Proofs.vitCotXinV {Np1 D : ℕ} (ε : ℝ) (γ1 : Vec D) (Wq Wk Wv : Mat D D) (xin dQ dK dV cotH : Vec (Np1 * D)) :

Vec (Np1 * D)

Cot at the block input, vector-LN form: cotH + the decomposed LN₁ input-VJP of the three-way Q/K/V fan-in.

Equations

Proofs.vitCotXinV ε γ1 Wq Wk Wv xin dQ dK dV cotH i = cotH i + Proofs.StableHLO.rowLNBackFlat Np1 D ε 1 xin (Proofs.StableHLO.rowScaleFlat Np1 D γ1 (Proofs.vitCotLn1 Wq Wk Wv dQ dK dV)) i

Instances For

source

noncomputable def Proofs.vitCotTowerOutV (N D nClasses : ℕ) (ε : ℝ) (γF : Vec D) (Wcls : Mat D nClasses) (towerOut : Vec ((N + 1) * D)) (dy : Vec nClasses) :

Vec ((N + 1) * D)

Cot at the last block's output (the final-LN input), vector-LN form: the decomposed final-LN input-VJP of the classifier-back row-0 scatter.

Equations

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

Instances For

Documentation

LeanMlir.Proofs.Nets.ViT.ViTVecLN

ViT at vector-[D] LayerNorm #

References #