Documentation

LeanMlir.Proofs.Architectures.LayerNorm

LayerNorm #

A quick chapter that extends the normalization family to what ViT needs: a structural footnote to BatchNorm, not new territory — which is itself the point. The smooth activations ViT and EfficientNet use (GELU, Swish, sigmoid) are in Architectures.Activations, which this file re-exports.

LayerNorm: BatchNorm on a different axis #

BatchNorm reduces over (batch, H, W) for each channel. LayerNorm reduces over the feature dimension for each (batch, token). The 1D normalization primitive is literally the same function. What differs is the axis you slice along before applying it.

Concretely, for a 4D activation x : Tensor4 B C H W:

BN computes C means/variances, each over B · H · W elements.
LN computes B · H · W means/variances, each over C elements.

The mean/var/istd/xhat/affine math is identical. The consolidated three-term backward is identical. Only the index being summed over changes. In our Vec n formalism, BN and LN collapse to the same function. This file just renames it to tell the reader "yes, really, it's the same thing."

Also in this file #

Layer scale: layerScale, layerScaleHasVJP.
Vector-[D] LayerNorm (per-feature affine): layerNormVec, layerNormVecHasVJP, and its per-token lift layerNormVecPerTokenHasVJPMat.
Its γ/β parameter gradients: vecLNGradGamma, vecLNGradBeta, the bridges layerNormVec_gamma_grad_bridge / layerNormVec_beta_grad_bridge, and their SGD forms layerNormVec_gamma_sgd_certified / layerNormVec_beta_sgd_certified.

Imported directly by SE, Attention, DropPath and Nets/ConvNeXt/ConvNeXt; the rest of the corpus reaches it through them.

References #

Ba, Kiros, Hinton 2016, Layer Normalization. https://arxiv.org/abs/1607.06450
Touvron et al. 2021, Going deeper with Image Transformers (CaiT, layer-scale). https://arxiv.org/abs/2103.17239

noncomputable def Proofs.layerNormForward (n : ℕ) (ε γ β : ℝ) (x : Vec n) :

LayerNorm forward — renamed bnForward to make the book's claim unambiguous: this is the same function, operating on a different slice of the tensor.

For a single "token's feature vector" x : Vec n:

mu = (1/n) sum_i x_i — mean across features
sigma^2 = (1/n) sum_i (x_i - mu)^2 — variance across features
istd = 1/sqrt(sigma^2 + eps)
xhat_i = (x_i - mu) * istd — normalized
y_i = gamma * xhat_i + beta — affine

Here γ and β are scalars, exactly as in bnForward. The per-feature [D] affine that LN uses in practice is layerNormVec below, built as γ ⊙ layerNormForward D ε 1 0 x + β.

MLIR (MlirCodegen.emitLayerNormForward): identical reduction structure to BN, just across a different axis.

Equations

Proofs.layerNormForward n ε γ β x = Proofs.bnForward n ε γ β x

Instances For

noncomputable def Proofs.layerNormHasVJP (n : ℕ) (ε γ β : ℝ) (hε : 0 < ε) :

HasVJP (layerNormForward n ε γ β)

LayerNorm input gradient — identical closed form to BN.

dx_i = (1/n) * istd * (n * dxhat_i - sum_j dxhat_j - xhat_i * sum_j xhat_j * dxhat_j)

where dxhat_i = gamma * dy_i.

layerNormForward is bnForward by definition, so layerNormHasVJP is bnHasVJP.

Equations

Proofs.layerNormHasVJP n ε γ β hε = id (Proofs.bnHasVJP n ε γ β hε)

Instances For

Why this isn't a new chapter #

The practical differences between BN and LN (batch dependence, inference vs training, running statistics) are engineering concerns, not VJP concerns. The backward pass is the same three-term formula either way. This is a general lesson about formal work: engineering distinctions often dissolve at the math level, and that's worth making explicit. A reader who assumed BN and LN needed separate proofs learns that the separation was an implementation artifact.

BN and LN share one HasVJP (layerNormForward is bnForward). GroupNorm and InstanceNorm apply the same 1D primitive to other slices of the tensor; RMSNorm drops the mean centering and is a different function. None of those three is formalized here.

theorem Proofs.layerNormHasVJP_correct (n : ℕ) (ε γ β : ℝ) (hε : 0 < ε) (x dy : Vec n) (i : Fin n) :

(layerNormHasVJP n ε γ β hε).backward x dy i = ∑ j : Fin n, pdiv (layerNormForward n ε γ β) x i j * dy j

Public correctness theorem for layerNormHasVJP: LayerNorm reuses the BN proof template (LayerNorm is BN on a different axis), so the contract is identical — backward equals the pdiv-contracted Jacobian of layerNormForward.

noncomputable def Proofs.layerScale {n : ℕ} (γ x : Vec n) :

Layer scale — per-channel learnable elementwise multiply by γ. layerScale γ x i = γ i * x i. A diagonal linear map.

Equations

Proofs.layerScale γ x i = γ i * x i

Instances For

theorem Proofs.layerScale_differentiable {n : ℕ} (γ : Vec n) :

Differentiable ℝ (layerScale γ)

layerScale γ is differentiable everywhere (diagonal linear).

theorem Proofs.pdiv_layerScale {n : ℕ} (γ x : Vec n) (i j : Fin n) :

pdiv (layerScale γ) x i j = if i = j then γ i else 0

Jacobian of layerScale — ∂(γ_j x_j)/∂x_i = γ_i δ_{ij}.

noncomputable def Proofs.layerScaleHasVJP {n : ℕ} (γ : Vec n) :

HasVJP (layerScale γ)

Layer scale VJP: back(x, dy)_i = γ i * dy i.

Equations

Proofs.layerScaleHasVJP γ = { backward := fun (_x dy : Proofs.Vec n) (i : Fin n) => γ i * dy i, correct := ⋯ }

Instances For

noncomputable def Proofs.layerNormVec (D : ℕ) (ε : ℝ) (γv βv x : Vec D) :

Vector-[D] LayerNorm: per-token normalize (the scalar LN at γ=1, β=0 — pure x̂), then the per-channel affine γ ⊙ x̂ + β. The committed ViTRender LN form.

Equations

Proofs.layerNormVec D ε γv βv x k = γv k * Proofs.layerNormForward D ε 1 0 x k + βv k

Instances For

theorem Proofs.layerNormVec_differentiable (D : ℕ) (ε : ℝ) (γv βv : Vec D) (hε : 0 < ε) :

Differentiable ℝ (layerNormVec D ε γv βv)

theorem Proofs.pdiv_id_add_const {m : ℕ} (C x : Vec m) (i j : Fin m) :

pdiv (fun (p : Vec m) (k : Fin m) => p k + C k) x i j = if i = j then 1 else 0

Identity-plus-constant Jacobian: ∂(p_k + C_k)/∂p_i = δ_(i,k).

theorem Proofs.pdiv_maskGather_add_const {m D : ℕ} (mask : Vec m) (σ : Fin m → Fin D) (C : Vec m) (x : Vec D) (i : Fin D) (j : Fin m) :

pdiv (fun (cl : Vec D) (k : Fin m) => mask k * cl (σ k) + C k) x i j = mask j * if i = σ j then 1 else 0

Masked-gather-plus-constant Jacobian: ∂(mask_k·cl_(σ k) + C_k)/∂cl_i = mask_k·δ_(i,σ k).

noncomputable def Proofs.biasAddHasVJP {n : ℕ} (βv : Vec n) :

HasVJP fun (z : Vec n) (k : Fin n) => z k + βv k

The bias translation's VJP — backward is the identity (dx = dy).

Equations

Proofs.biasAddHasVJP βv = { backward := fun (_z dy : Proofs.Vec n) => dy, correct := ⋯ }

Instances For

noncomputable def Proofs.layerNormVecHasVJP (D : ℕ) (ε : ℝ) (γv βv : Vec D) (hε : 0 < ε) :

HasVJP (layerNormVec D ε γv βv)

Vector-LN VJP — (+β) ∘ layerScale γ ∘ LN(1,0), three proven pieces glued by vjpComp. Only 0 < ε.

Equations

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

Instances For

noncomputable def Proofs.layerNormVecPerTokenHasVJPMat (N D : ℕ) (ε : ℝ) (γv βv : Vec D) (hε : 0 < ε) :

HasVJPMat fun (X : Mat N D) (r : Fin N) => layerNormVec D ε γv βv (X r)

Per-token vector-LN across a sequence — the rowwise lift.

Equations

Proofs.layerNormVecPerTokenHasVJPMat N D ε γv βv hε = Proofs.rowwiseHasVJPMat (Proofs.layerNormVecHasVJP D ε γv βv hε) ⋯

Instances For

theorem Proofs.rowwise_flat_differentiable {N D P : ℕ} (g : Vec D → Vec P) (hg : Differentiable ℝ g) :

Differentiable ℝ fun (v : Vec (N * D)) => Mat.flatten ((fun (X : Mat N D) (n : Fin N) => g (X n)) (Mat.unflatten v))

Generic flat differentiability of a rowwise lift — each output coordinate is a coordinate of the per-row map applied to one row of the input.

theorem Proofs.layerNormVec_per_token_flat_differentiable (N D : ℕ) (ε : ℝ) (γv βv : Vec D) (hε : 0 < ε) :

Differentiable ℝ fun (v : Vec (N * D)) => Mat.flatten ((fun (X : Mat N D) (n : Fin N) => layerNormVec D ε γv βv (X n)) (Mat.unflatten v))

theorem Proofs.pdiv_vecLN_gamma {N D : ℕ} (ε : ℝ) (βv : Vec D) (X : Mat N D) (γ : Vec D) (i : Fin D) (o : Fin (N * D)) :

pdiv (fun (gv : Vec D) => Mat.flatten fun (r : Fin N) => layerNormVec D ε gv βv (X r)) γ i o = layerNormForward D ε 1 0 (X (finProdFinEquiv.symm o).1) (finProdFinEquiv.symm o).2 * if i = (finProdFinEquiv.symm o).2 then 1 else 0

Jacobian of the rowwise vector-LN site w.r.t. γv — ∂y_(r,k)/∂γv_i = δ_(i,k)·x̂_r(k).

theorem Proofs.pdiv_vecLN_beta {N D : ℕ} (ε : ℝ) (γv : Vec D) (X : Mat N D) (β : Vec D) (i : Fin D) (o : Fin (N * D)) :

pdiv (fun (bv : Vec D) => Mat.flatten fun (r : Fin N) => layerNormVec D ε γv bv (X r)) β i o = if i = (finProdFinEquiv.symm o).2 then 1 else 0

Jacobian of the rowwise vector-LN site w.r.t. βv — ∂y_(r,k)/∂βv_i = δ_(i,k).

noncomputable def Proofs.vecLNGradGamma (N D : ℕ) (ε : ℝ) (X dY : Mat N D) :

The rendered vector-LN γ gradient: per-channel, the batch+token reduce dγ_k = Σ_r dY_(r,k)·x̂_r(k) (KEEPS the channel axis — ViTRender's form).

Equations

Proofs.vecLNGradGamma N D ε X dY i = ∑ r : Fin N, dY r i * Proofs.layerNormForward D ε 1 0 (X r) i

Instances For

noncomputable def Proofs.vecLNGradBeta (N D : ℕ) (dY : Mat N D) :

The rendered vector-LN β gradient: dβ_k = Σ_r dY_(r,k).

Equations

Proofs.vecLNGradBeta N D dY i = ∑ r : Fin N, dY r i

Instances For

theorem Proofs.layerNormVec_gamma_grad_bridge {N D : ℕ} (ε : ℝ) (βv γ : Vec D) (X : Mat N D) (dy : Vec (N * D)) (i : Fin D) :

vecLNGradGamma N D ε X (Mat.unflatten dy) i = ∑ o : Fin (N * D), pdiv (fun (gv : Vec D) => Mat.flatten fun (r : Fin N) => layerNormVec D ε gv βv (X r)) γ i o * dy o

Vector-LN γ-gradient bridge.

theorem Proofs.layerNormVec_beta_grad_bridge {N D : ℕ} (ε : ℝ) (γv β : Vec D) (X : Mat N D) (dy : Vec (N * D)) (i : Fin D) :

vecLNGradBeta N D (Mat.unflatten dy) i = ∑ o : Fin (N * D), pdiv (fun (bv : Vec D) => Mat.flatten fun (r : Fin N) => layerNormVec D ε γv bv (X r)) β i o * dy o

Vector-LN β-gradient bridge.

theorem Proofs.layerNormVec_gamma_sgd_certified {N D : ℕ} (ε : ℝ) (βv γ : Vec D) (X : Mat N D) (dy : Vec (N * D)) (lr : ℝ) (i : Fin D) :

γ i - lr * vecLNGradGamma N D ε X (Mat.unflatten dy) i = γ i - lr * ∑ o : Fin (N * D), pdiv (fun (gv : Vec D) => Mat.flatten fun (r : Fin N) => layerNormVec D ε gv βv (X r)) γ i o * dy o

Vector-LN γ output, certified. γvⁿ_k = γv_k − lr·(Σ_tokens dy·x̂)_k denotes the certified rowwise vector-LN ∂/∂γv contraction: a rewrite by layerNormVec_gamma_grad_bridge, stated at one LN site with generic N, D.

theorem Proofs.layerNormVec_beta_sgd_certified {N D : ℕ} (ε : ℝ) (γv β : Vec D) (X : Mat N D) (dy : Vec (N * D)) (lr : ℝ) (i : Fin D) :

β i - lr * vecLNGradBeta N D (Mat.unflatten dy) i = β i - lr * ∑ o : Fin (N * D), pdiv (fun (bv : Vec D) => Mat.flatten fun (r : Fin N) => layerNormVec D ε γv bv (X r)) β i o * dy o

Vector-LN β output, certified.