Batch Normalization VJP

This is the first layer where the casual "stare and differentiate" approach breaks down. In dense and conv layers, every output cell yⱼ depends on its input independently of the other inputs. In batch norm, every output depends on every input through the mean and variance reductions, so the Jacobian is dense and the chain rule has to do real work.

The famous result we'll derive: the input gradient collapses to a single three-term closed form that doesn't expose the individual contributions from mean and variance. This is the "consolidated" BN backward formula that every ML framework hard-codes (because deriving it on the fly is a pain). It's what MlirCodegen.lean emits at line 799:

%cbg_t5 = istd * (N * d_xhat - sum(d_xhat) - xhat * sum(d_xhat * xhat))
%cbg_dconv = (1/N) * %cbg_t5

This file:

1. Defines BN forward step by step (mean → var → istd → xhat → affine).
2. States the easy parameter gradients (γ, β).
3. Walks through the derivation of the hard input gradient and states the consolidated formula.

A note on shapes

The actual implementation reduces over [batch, h, w] per channel. For clarity, this file works on a single 1D Vec n (think of n as B · H · W flattened, for one channel). The math is identical; only the indexing changes when you go to 4D.

noncomputable def Proofs.bnMean (n : ℕ) (x : Vec n) : ℝ

Population mean: μ = (1/N) Σᵢ xᵢ

Equations

• Proofs.bnMean n x = (∑ i : Fin n, x i) / ↑n

noncomputable def Proofs.bnVar (n : ℕ) (x : Vec n) : ℝ

Population variance: σ² = (1/N) Σᵢ (xᵢ − μ)²

Equations

• Proofs.bnVar n x = (∑ i : Fin n, (x i - Proofs.bnMean n x) * (x i - Proofs.bnMean n x)) / ↑n

noncomputable def Proofs.bnIstd (n : ℕ) (x : Vec n) (ε : ℝ) : ℝ

Inverse standard deviation: istd = 1 / √(σ² + ε)

Equations

• Proofs.bnIstd n x ε = 1 / √(Proofs.bnVar n x + ε)

noncomputable def Proofs.bnXhat (n : ℕ) (ε : ℝ) (x : Vec n) : Vec n

Normalized output: x̂ᵢ = (xᵢ − μ) · istd

x̂ has mean 0 and variance 1 (up to ε-correction). It's the "centered, unit-scaled" version of x.

Equations

• Proofs.bnXhat n ε x i = (x i - Proofs.bnMean n x) * Proofs.bnIstd n x ε

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

The full BN forward: yᵢ = γ · x̂ᵢ + β

γ and β are learnable per-channel parameters that restore the network's representational freedom that normalization took away. Without them, BN would force every layer's output to have mean 0, variance 1 — too constraining.

MLIR (MlirCodegen.lean lines 723–728): %cbn_g_bc = broadcast %g %cbn_gn = multiply %cbn_norm, %cbn_g_bc %cbn_bt_bc = broadcast %bt %cbn_pre = add %cbn_gn, %cbn_bt_bc

Equations

• Proofs.bnForward n ε γ β x i = γ * Proofs.bnXhat n ε x i + β

noncomputable def Proofs.bn_grad_gamma (n : ℕ) (ε : ℝ) (x dy : Vec n) : ℝ

γ gradient: dγ = Σᵢ dyᵢ · x̂ᵢ

γ is a scalar that multiplies each x̂ᵢ. By the product rule: ∂yᵢ/∂γ = x̂ᵢ. Summing over the output cotangent dy: dγ = Σᵢ dyᵢ · x̂ᵢ.

This is just an inner product of dy with x̂ — no mean/variance chain-rule trickery, because γ doesn't enter the reduction.

MLIR (MlirCodegen.lean lines 766–768): %cbg_gn = multiply %effGrad, %cbn_norm %d_g = reduce add %cbg_gn across dimensions = [0, 2, 3]

Equations

• Proofs.bn_grad_gamma n ε x dy = ∑ i : Fin n, dy i * Proofs.bnXhat n ε x i

noncomputable def Proofs.bn_grad_beta (n : ℕ) (dy : Vec n) : ℝ

β gradient: dβ = Σᵢ dyᵢ

β is added to every output, so ∂yᵢ/∂β = 1 and the gradient is just the sum of the output cotangents. Even simpler than dγ.

MLIR (line 770): %d_bt = reduce add %effGrad across dimensions = [0, 2, 3]

Equations

• Proofs.bn_grad_beta n dy = ∑ i : Fin n, dy i

Why the input gradient is hard

The output yⱼ depends on xᵢ through three paths:

(a) Directly: xⱼ appears in (xⱼ − μ) (only when i = j). (b) Via μ: μ is (1/N) Σₖ xₖ, so changing xᵢ changes μ by 1/N, which shifts every (xⱼ − μ). (c) Via σ²: σ² is (1/N) Σₖ (xₖ − μ)², so changing xᵢ changes σ², which changes istd, which scales every x̂ⱼ.

So ∂yⱼ/∂xᵢ ≠ 0 for every (i, j) pair — the Jacobian is dense. Naively, the VJP costs O(N²); the consolidated form turns it into O(N) by collapsing the cancellations algebraically.

The derivation

Strip off the affine layer first: let dx̂ᵢ := γ · dyᵢ. Then we need the VJP of bnXhat (the normalize step) at the cotangent dx̂.

For x̂ⱼ = (xⱼ − μ) · istd, the chain rule gives:

∂x̂ⱼ/∂xᵢ = (∂xⱼ/∂xᵢ − ∂μ/∂xᵢ) · istd + (xⱼ − μ) · ∂istd/∂xᵢ

We need three sub-derivatives:

∂xⱼ/∂xᵢ  = δᵢⱼ                               (identity)
∂μ/∂xᵢ   = 1/N                                (mean is linear in x)
∂σ²/∂xᵢ  = (2/N) · (xᵢ − μ) · (1 − 1/N)
          ≈ (2/N) · (xᵢ − μ)                  (the (1−1/N) term
                                               eats into a Σ that
                                               sums to zero, so it
                                               doesn't survive)
∂istd/∂xᵢ = (−1/2) · istd³ · ∂σ²/∂xᵢ
          = −istd³ · (xᵢ − μ) / N
          = −istd · x̂ᵢ / N                    (since x̂ᵢ = (xᵢ−μ)·istd)

Substituting:

∂x̂ⱼ/∂xᵢ = (δᵢⱼ − 1/N) · istd − (xⱼ − μ) · istd · x̂ᵢ / N
        = istd · (δᵢⱼ − 1/N − x̂ⱼ · x̂ᵢ / N)
        = (istd / N) · (N · δᵢⱼ − 1 − x̂ᵢ · x̂ⱼ)

Now contract with dx̂ to get the input cotangent of the normalize step:

dxᵢ = Σⱼ (∂x̂ⱼ/∂xᵢ) · dx̂ⱼ
    = (istd / N) · Σⱼ (N · δᵢⱼ − 1 − x̂ᵢ · x̂ⱼ) · dx̂ⱼ
    = (istd / N) · (N · dx̂ᵢ − Σⱼ dx̂ⱼ − x̂ᵢ · Σⱼ x̂ⱼ · dx̂ⱼ)

This is the consolidated formula — three terms, two scalar reductions (Σⱼ dx̂ⱼ and Σⱼ x̂ⱼ · dx̂ⱼ), one elementwise broadcast. O(N) work instead of O(N²). And it's exactly what the MLIR emits.

noncomputable def Proofs.bn_grad_input (n : ℕ) (ε γ : ℝ) (x dy : Vec n) : Vec n

The consolidated BN input gradient.

dxᵢ = (1/N) · istd · (N · dx̂ᵢ − Σⱼ dx̂ⱼ − x̂ᵢ · Σⱼ x̂ⱼ · dx̂ⱼ)

where dx̂ᵢ = γ · dyᵢ (gradient pulled back through the affine layer first).

This matches MlirCodegen.lean lines 794–801: %cbg_t1 = N * d_xhat %cbg_t2 = %cbg_t1 - sum(d_xhat) -- subtract mean %cbg_t3 = xhat * sum(xhat * d_xhat) %cbg_t4 = %cbg_t2 - %cbg_t3 -- the three-term combo %cbg_t5 = istd * %cbg_t4 %cbg_dconv = (1/N) * %cbg_t5

Equations

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

Cleaner view: BN as a composition

The BN forward is really two steps glued together:

1. Normalize (bnXhat): the hard part with mean/var/istd reductions. Vec n → Vec n, no parameters.
2. Affine (fun v i => γ · vᵢ + β): elementwise scale-and-shift. The parameters γ, β live here.

If we had a HasVJP instance for each, we could compose them with vjp_comp from Tensor.lean and get the full BN VJP "for free."

The affine VJP is trivial: ∂(γ · vᵢ + β)/∂vⱼ = γ · δᵢⱼ → back(v, dy)ᵢ = γ · dyᵢ

The normalize VJP is the consolidated three-term formula above (with γ = 1, since the affine has been factored out).

We state both as HasVJP instances. Their composition (via vjp_comp) gives the full BN input gradient — and the parameter gradients are collected at the affine layer alongside.

noncomputable def Proofs.bnNormalize (n : ℕ) (ε : ℝ) : Vec n → Vec n

The normalize step as a function Vec n → Vec n (no params except ε).

Equations

• Proofs.bnNormalize n ε = Proofs.bnXhat n ε

noncomputable def Proofs.bnAffine (n : ℕ) (γ β : ℝ) : Vec n → Vec n

The affine step as a function Vec n → Vec n (γ, β as constants).

Equations

• Proofs.bnAffine n γ β v i = γ * v i + β

theorem Proofs.bnForward_eq_compose (n : ℕ) (ε γ β : ℝ) : bnForward n ε γ β = bnAffine n γ β ∘ bnNormalize n ε

BN as the composition of normalize and affine.

theorem Proofs.pdiv_bnAffine (n : ℕ) (γ β : ℝ) (v : Vec n) (i j : Fin n) : pdiv (bnAffine n γ β) v i j = if i = j then γ else 0

The affine Jacobian is diagonal: ∂(γ·vᵢ + β)/∂vⱼ = γ · δᵢⱼ.

Proved from foundation rules: bnAffine decomposes as (γ · v) + (constant β), where the linear term factors further as (constant γ) * (identity) for pdiv_mul. The pieces collapse via pdiv_add + pdiv_mul + pdiv_const + pdiv_id.

The consolidated three-term formula used to be axiomatized directly. Now it's a theorem: we factor bnXhat as the elementwise product of the centered input and the broadcast istd, apply pdiv_mul, and collapse via ring using the x̂ᵢ = (xᵢ - μ) · istd identity.

Both elementary calculus facts are now proved from the foundation:

1. pdiv_bnCentered — ∂(xⱼ - μ(x))/∂xᵢ = δᵢⱼ - 1/n. Proved via Mathlib's HasDerivAt.sub applied to id and (const_mul) ∘ (Finset.sum).

2. pdiv_bnIstdBroadcast — ∂istd(x,ε)/∂xᵢ = -istd³ · (xᵢ - μ) / n. Proved via the centering CLM + HasFDerivAt.sqrt (under bnVar + ε > 0)

   • (hasDerivAt_inv).comp_hasFDerivAt. The centered sum collapses by Σ_k (x_k − μ) = 0. Carries (hε : 0 < ε) hypothesis throughout.

The three-term formula falls out by ring manipulation alone.

noncomputable def Proofs.bnCentered (n : ℕ) : Vec n → Vec n

Centered input: (x - μ(x)) as a Vec n → Vec n function.

Equations

• Proofs.bnCentered n x j = x j - Proofs.bnMean n x

noncomputable def Proofs.bnIstdBroadcast (n : ℕ) (ε : ℝ) : Vec n → Vec n

Broadcast inverse-stddev: istd(x,ε) as a Vec n → Vec n function (constant in the output index, just lifted for pdiv_mul).

Equations

• Proofs.bnIstdBroadcast n ε x x✝ = Proofs.bnIstd n x ε

theorem Proofs.bnXhat_eq_product (n : ℕ) (ε : ℝ) (x : Vec n) : bnXhat n ε x = fun (j : Fin n) => bnCentered n x j * bnIstdBroadcast n ε x j

bnXhat factors as bnCentered · bnIstdBroadcast (elementwise product).

theorem Proofs.pdiv_bnCentered (n : ℕ) (x : Vec n) (i j : Fin n) : pdiv (bnCentered n) x i j = (if i = j then 1 else 0) - 1 / ↑n

Centered-input Jacobian — proved from foundation rules.

∂(xⱼ - μ(x))/∂xᵢ = δᵢⱼ - 1/n

Decomposition: bnCentered y k = y k - (∑ s, y s)/n factors as (id y) k + (-(1/n)) * (∑ s, y s). The first half collapses via pdiv_id; the second factors as (constant) * (sum) and uses pdiv_mul + pdiv_const + pdiv_finset_sum + pdiv_reindex to yield -1/n.

theorem Proofs.bnIstdBroadcast_diff (n : ℕ) (ε : ℝ) (hε : 0 < ε) : Differentiable ℝ (bnIstdBroadcast n ε)

Smoothness of bnIstdBroadcast — proved from Mathlib calculus (planning/archive/VJP.md follow-up C).

bnIstdBroadcast n ε x = 1/√(σ²(x) + ε). Since σ²(x) ≥ 0 (sum of squares ÷ n ≥ 0) and ε > 0, the argument bnVar + ε is everywhere positive, so Real.sqrt is differentiable (Differentiable.sqrt with non-zero hypothesis), and its reciprocal is differentiable too.

theorem Proofs.pdiv_bnIstdBroadcast (n : ℕ) (ε : ℝ) (hε : 0 < ε) (x : Vec n) (i j : Fin n) : pdiv (bnIstdBroadcast n ε) x i j = -bnIstd n x ε ^ 3 * (x i - bnMean n x) / ↑n

Broadcast inverse-stddev Jacobian — proved (was an axiom).

∂istd(x,ε)/∂xᵢ = -istd³(x,ε) · (xᵢ - μ(x)) / n

Derivation:

• istd = 1/√(σ²+ε) → chain rule through Real.sqrt and x ↦ 1/x: ∂istd/∂σ² = -(1/2) · istd³
• ∂σ²/∂xᵢ = (2/n) · (xᵢ - μ) (product rule on (xⱼ - μ)² summed, using Σⱼ (xⱼ - μ) = 0 to cancel a (1 - 1/n) factor)
• Chain together: ∂istd/∂xᵢ = -istd³ · (xᵢ - μ) / n.

Lean proof structure: HasFDerivAt chain through the centering CLM C k = proj k - (1/n) Σ_i proj i (linear in x'), squared via .mul, summed via .fun_sum, scaled by 1/n via .mul_const, .add_const ε, then .sqrt (with bnVar+ε > 0), then (hasDerivAt_inv ·).comp_hasFDerivAt for the reciprocal. The resulting CLM at basisVec i simplifies via the Σⱼ (xⱼ - μ) = 0 identity.

theorem Proofs.pdiv_bnNormalize (n : ℕ) (ε : ℝ) (hε : 0 < ε) (x : Vec n) (i j : Fin n) : pdiv (bnNormalize n ε) x i j = bnIstd n x ε / ↑n * ((↑n * if i = j then 1 else 0) - 1 - bnXhat n ε x i * bnXhat n ε x j)

The BN normalize Jacobian — derived, no longer axiomatized.

pdiv (bnNormalize n ε) x i j = (istd / n) · (n · δᵢⱼ − 1 − x̂ᵢ · x̂ⱼ)

Proof: factor bnXhat = bnCentered · bnIstdBroadcast, apply pdiv_mul, substitute the two elementary Jacobians, then expand x̂ₖ = (xₖ - μ) · istd and collapse with ring.

noncomputable def Proofs.bnAffine_has_vjp (n : ℕ) (γ β : ℝ) : HasVJP (bnAffine n γ β)

Affine VJP (the easy half): back(v, dy)ᵢ = γ · dyᵢ.

Each input enters one output multiplied by γ; the gradient comes back scaled by γ.

Equations

• Proofs.bnAffine_has_vjp n γ β = { backward := fun (_v dy : Proofs.Vec n) (i : Fin n) => γ * dy i, correct := ⋯ }

noncomputable def Proofs.bnNormalize_has_vjp (n : ℕ) (ε : ℝ) (hε : 0 < ε) : HasVJP (bnNormalize n ε)

Normalize VJP (the hard half): the consolidated formula with γ = 1.

back(x, dx̂)ᵢ = (1/N) · istd · (N · dx̂ᵢ − Σⱼ dx̂ⱼ − x̂ᵢ · Σⱼ x̂ⱼ · dx̂ⱼ)

Equations

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

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

The BN VJP from the composition — chain rule glues affine ∘ normalize.

This is the structural payoff: once bnNormalize_has_vjp and bnAffine_has_vjp are in hand, the full BN input gradient comes from one application of vjp_comp. The chain rule mechanically threads dy → dx̂ → dx:

dx̂ᵢ = γ · dyᵢ                           (from bnAffine_has_vjp)
dxᵢ = (1/N · istd) · (N · dx̂ᵢ − …)     (from bnNormalize_has_vjp)

The composition is exactly the two-step backward pass that the MLIR emits: lines 773 (d_norm = grad * gamma_bc) followed by lines 794–801 (the consolidated three-term formula).

Equations

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

theorem Proofs.bnForward_differentiable (n : ℕ) (ε γ β : ℝ) (hε : 0 < ε) : Differentiable ℝ (bnForward n ε γ β)

bnForward is differentiable everywhere (for ε > 0).

Reuses the exact differentiability argument inside bn_has_vjp: bnForward = bnAffine ∘ bnNormalize, where bnNormalize is the product of bnCentered (affine, hence smooth) and bnIstdBroadcast (smooth because bnVar + ε > 0 keeps the Real.sqrt away from its kink — see bnIstdBroadcast_diff), and bnAffine is affine. The ε > 0 hypothesis is what licenses the inverse-sqrt smoothness. This is the differentiability witness vjp_comp_at needs to chain bn into the conv→bn→relu block.

theorem Proofs.bn_input_grad_correct (n : ℕ) (ε γ β : ℝ) (hε : 0 < ε) (x dy : Vec n) (i : Fin n) : bn_grad_input n ε γ x dy i = ∑ j : Fin n, pdiv (bnForward n ε γ β) x i j * dy j

The standalone end-to-end theorem: bn_grad_input is the correct VJP of bnForward. Follows from bn_has_vjp by definitional unfolding.