Documentation

LeanMlir.Proofs.Architectures.BatchNorm

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.emitConvBnBackward emits:

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

This file:

Defines BN forward step by step (mean → var → istd → xhat → affine).
States the easy parameter gradients (γ, β).
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.

References #

Ioffe & Szegedy 2015, Batch Normalization. https://arxiv.org/abs/1502.03167
Chan, Golub, LeVeque 1982, Updating Formulae and a Pairwise Algorithm for Computing Sample Variances (COMPSTAT; the pairwise combine rule). https://doi.org/10.1007/978-3-642-51461-6_3

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

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

Equations

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

Instances For

theorem Proofs.bnMean_eq_expect (n : ℕ) (x : Vec n) :

bnMean n x = Finset.univ.expect fun (i : Fin n) => x i

bnMean is Mathlib's finite average 𝔼 i, x i (Finset.expect), so its reindexing and product lemmas apply.

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

Instances For

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

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

Equations

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

Instances For

theorem Proofs.bnVar_nonneg (n : ℕ) (x : Vec n) :

0 ≤ bnVar n x

The population variance is nonnegative (a mean of squares).

noncomputable def Proofs.bnMeanSq (n : ℕ) (x : Vec n) :

Second moment: E[x²] = (1/N) Σᵢ xᵢ².

At equal shard sizes the second moment of the union is the mean of the shards' second moments (bnMeanSq_shard). bnSyncTensor4 and the sync backward are stated at μ and a second moment m2; the emitted sync-BN forward exchanges μ and σ² (two rounds, Chan's parallel variance, bnVar_shard_chan), and its den supplies m2 := σ² + μ².

Equations

Proofs.bnMeanSq n x = (∑ i : Fin n, x i * x i) / ↑n

Instances For

theorem Proofs.bnMean_shard {R m M : ℕ} (e : Fin R × Fin m ≃ Fin M) (x : Vec M) :

bnMean M x = 1 / ↑R * ∑ r : Fin R, bnMean m fun (k : Fin m) => x (e (r, k))

The mean of a sharded vector is the mean of the shards' means — at EQUAL shard sizes, which is what makes allReduceMeanF (a plain mean over replicas) the right collective for BatchNorm statistics.

Stated at an arbitrary shard e, in the style DataParallel.meanLoss_shard sets: WHICH cells land on which replica never enters, only that together they are the whole. The contiguous cut the DP shim makes is the finProdFinEquiv instance.

theorem Proofs.bnMeanSq_shard {R m M : ℕ} (e : Fin R × Fin m ≃ Fin M) (x : Vec M) :

bnMeanSq M x = 1 / ↑R * ∑ r : Fin R, bnMeanSq m fun (k : Fin m) => x (e (r, k))

…and so is the second moment. Free from bnMean_shard: bnMeanSq n x is definitionally bnMean n (x·x), and squaring commutes with the shard.

The variance does not shard this way: bnVar of the union is not the mean of the shards' bnVar unless every shard mean coincides, and that spread is what a per-replica BatchNorm drops. The variance's shard formula is bnVar_shard_chan, which adds each shard mean's squared offset from the global mean.

theorem Proofs.bnVar_eq_bnMeanSq_sub_sq (n : ℕ) (hn : n ≠ 0) (x : Vec n) :

bnVar n x = bnMeanSq n x - bnMean n x * bnMean n x

σ² = E[x²] − μ² — the identity that lets the sync-BN definitions be stated at the second moment (bnSyncXhat reads m2 − μ²).

theorem Proofs.bnVar_add_mean_mul_mean (n : ℕ) (hn : n ≠ 0) (x : Vec n) :

bnVar n x + bnMean n x * bnMean n x = bnMeanSq n x

bnVar_eq_bnMeanSq_sub_sq solved for E[x²] — the form a sync-BN statistics pack (μ, σ² + μ²) needs to read the batch's own second moment back.

theorem Proofs.bnVar_shard_chan {R m M : ℕ} (hR : R ≠ 0) (hm : m ≠ 0) (e : Fin R × Fin m ≃ Fin M) (x : Vec M) :

bnVar M x = 1 / ↑R * ∑ r : Fin R, ((bnVar m fun (k : Fin m) => x (e (r, k))) + ((bnMean m fun (k : Fin m) => x (e (r, k))) - bnMean M x) * ((bnMean m fun (k : Fin m) => x (e (r, k))) - bnMean M x))

Chan's parallel variance: the variance of the whole is the mean over shards of each shard's OWN variance plus its mean's squared offset from the global mean.

What a synchronised BatchNorm exchanges instead of E[x²]: every term is a two-pass quantity on its shard, so there is no E[x²] − μ² cancellation anywhere in f32. Proved by expanding every variance as E[x²] − μ² (bnVar_eq_bnMeanSq_sub_sq) and collecting with bnMean_shard / bnMeanSq_shard; stated at an arbitrary shard e, like them.

theorem Proofs.bnIstd_pos {n : ℕ} (v : Vec n) (ε : ℝ) (hε : 0 < ε) :

0 < bnIstd n v ε

istd = 1/√(σ²+ε) > 0 (variance ≥ 0, ε > 0).

noncomputable def Proofs.bnXhat (n : ℕ) (ε : ℝ) (x : 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 ε

Instances For

noncomputable def Proofs.bnForward (n : ℕ) (ε γ β : ℝ) (x : 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.emitConvBnTrain): %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 + β

Instances For

theorem Proofs.bnXhat_sq_le {n : ℕ} (ε : ℝ) (hε : 0 < ε) (v : Vec n) (k : Fin n) :

bnXhat n ε v k ^ 2 ≤ ↑n

Each normalized coordinate is bounded: x̂ₖ² ≤ n. Proof: istd² = 1/(σ²+ε) and (vₖ−μ)² ≤ Σⱼ(vⱼ−μ)² = n·σ² ≤ n·(σ²+ε).

theorem Proofs.bnForward_const {n : ℕ} (hn : 0 < n) (ε γ β c : ℝ) :

(bnForward n ε γ β fun (x : Fin n) => c) = fun (x : Fin n) => β

BN of a constant vector is the (constant) shift β — centering zeroes the normalized term, killing the √. The zero-weight blocks of the concrete ResNet-34 and MobileNetV2 witnesses collapse through it.

noncomputable def Proofs.bnGradGamma (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.emitConvBnBackward): %cbg_gn = multiply %effGrad, %cbn_norm %d_g = reduce add %cbg_gn across dimensions = [0, 2, 3]

Equations

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

Instances For

noncomputable def Proofs.bnGradBeta (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 (same function): %d_bt = reduce add %effGrad across dimensions = [0, 2, 3]

Equations

Proofs.bnGradBeta n dy = ∑ i : Fin n, dy i

Instances For

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ᵢ − μ)                  (exactly: the −1/N term
                                               multiplies Σₖ (xₖ − μ) = 0;
                                               `bnVarDeriv_basisVec`)
∂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.bnGradInput (n : ℕ) (ε γ : ℝ) (x dy : 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.emitConvBnBackward's %cbg_t* chain: %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.

Instances For

noncomputable def Proofs.bnSyncXhat (n : ℕ) (ε μ m2 : ℝ) (x : Vec n) :

The normalised activation under HANDED-IN statistics: x̂ᵢ = (xᵢ − μ)·(m2 − μ² + ε)^(−1/2). bnXhat's peer — same function, but reading its statistics rather than reducing x for them, so under DP they can be the all-reduced global ones. Shared by the sync backward and by the dy-statistics it consumes, so the two cannot drift apart.

Equations

Proofs.bnSyncXhat n ε μ m2 x i = (x i - μ) * (1 / √(m2 - μ * μ + ε))

Instances For

theorem Proofs.bnSyncXhat_apply (n : ℕ) (ε μ m2 : ℝ) (v : Vec n) (i : Fin n) :

bnSyncXhat n ε μ m2 v i = (v i - μ) * (1 / √(m2 - μ * μ + ε))

bnSyncXhat is pointwise, and its size index never enters the value — which is what lets a global row's x̂ restrict to a shard's x̂ at the SAME handed-in statistics.

theorem Proofs.bnSyncXhat_at_own_stats (n : ℕ) (hn : n ≠ 0) (ε : ℝ) (x : Vec n) :

bnSyncXhat n ε (bnMean n x) (bnMeanSq n x) x = bnXhat n ε x

bnSyncXhat at its own statistics is bnXhat. The normalised-activation peer of bnEvalForward_at_own_stats; needed wherever a sync node's x̂ has to be recognised as the training x̂ — in particular inside the dy-reductions the sync backward consumes.

noncomputable def Proofs.bnSyncGradInput (n : ℕ) (ε γ μ m2 mdy mdyx : ℝ) (x dy : Vec n) :

The synchronised batch-norm input-VJP — every reduction handed in.

bnGradInput computes all four of its scalars from x and dy: the mean, the inverse standard deviation, and the two sums. This one takes μ, the second moment m2 and the two reduction means mdy, mdyx as arguments, so under data parallelism they can be the all-reduced global ones and a replica can produce the shard-r block of the global-batch gradient.

It is expressible this way because bnGradInput can be written as istd · (dx̂ᵢ − mean(dx̂) − x̂ᵢ · mean(x̂·dx̂)). Both reductions are means, and a mean over equal shards is the mean of the shards' means (bnMean_shard), so both survive an allReduceMeanF, as μ does in the forward. Nothing here needs a sum.

Equations

Proofs.bnSyncGradInput n ε γ μ m2 mdy mdyx x dy i = 1 / √(m2 - μ * μ + ε) * (γ * dy i - mdy - Proofs.bnSyncXhat n ε μ m2 x i * mdyx)

Instances For

theorem Proofs.bnSyncGradInput_at_own_stats (n : ℕ) (hn : n ≠ 0) (ε γ : ℝ) (x dy : Vec n) :

bnSyncGradInput n ε γ (bnMean n x) (bnMeanSq n x) (bnMean n fun (i : Fin n) => γ * dy i) (bnMean n fun (i : Fin n) => bnXhat n ε x i * (γ * dy i)) x dy = bnGradInput n ε γ x dy

R = 1: the sync backward at its own statistics is bnGradInput.

The backward peer of bnEvalForward_at_own_stats: handed the statistics and reductions the batch would itself have computed, the sync backward denotes the existing three-term formula. So a single-device sync render computes the same function as the per-batch bnGradInput.

Cleaner view: BN as a composition #

The BN forward is really two steps glued together:

Normalize (bnXhat): the hard part with mean/var/istd reductions. Vec n → Vec n, no parameters.
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 vjpComp 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 vjpComp) 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 ε

Instances For

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 + β

Instances For

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ⱼ = γ · δᵢⱼ.

bnAffine is the linear map v ↦ γ · v plus the constant β, so pdiv_of_affine reads the Jacobian off the basis vector.

The consolidated three-term formula is 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 proved from the foundation:

pdiv_bnCentered — ∂(xⱼ - μ(x))/∂xᵢ = δᵢⱼ - 1/n. Proved via pdiv_of_linear: centering is a linear map.
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

Instances For

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 ε

Instances For

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

bnCentered is linear, so pdiv_of_linear reads the Jacobian off the basis vector: eᵢ has entry δᵢⱼ and mean 1/n.

theorem Proofs.bnIstdBroadcast_differentiable (n : ℕ) (ε : ℝ) (hε : 0 < ε) :

Differentiable ℝ (bnIstdBroadcast n ε)

Smoothness of bnIstdBroadcast — proved from Mathlib calculus.

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.

noncomputable def Proofs.bnCenterCLM (n : ℕ) (k : Fin n) :

Vec n →L[ℝ ] ℝ

The centering map y ↦ y k − μ(y) as a continuous linear functional: coordinate k minus the mean. The variance is the mean of its squares, so this is what its derivative is built from.

Equations

Proofs.bnCenterCLM n k = ContinuousLinearMap.proj k - (↑n)⁻¹ • ∑ i : Fin n, ContinuousLinearMap.proj i

Instances For

theorem Proofs.bnCenterCLM_apply (n : ℕ) (k : Fin n) (y : Vec n) :

(bnCenterCLM n k) y = y k - bnMean n y

theorem Proofs.bnCenterCLM_basisVec (n : ℕ) (k i : Fin n) :

(bnCenterCLM n k) (basisVec i) = (if k = i then 1 else 0) - (↑n)⁻¹

theorem Proofs.sum_sub_bnMean (n : ℕ) (hn : n ≠ 0) (x : Vec n) :

∑ k : Fin n, (x k - bnMean n x) = 0

The centered entries sum to zero.

noncomputable def Proofs.bnVarDeriv (n : ℕ) (x : Vec n) :

Vec n →L[ℝ ] ℝ

The derivative of the variance at x: (1/n) Σ_k 2 (x_k − μ) · center_k.

Equations

Proofs.bnVarDeriv n x = (↑n)⁻¹ • ∑ k : Fin n, (2 * (x k - Proofs.bnMean n x)) • Proofs.bnCenterCLM n k

Instances For

theorem Proofs.bnVar_hasFDerivAt (n : ℕ) (x : Vec n) :

HasFDerivAt (fun (x' : Vec n) => bnVar n x') (bnVarDeriv n x) x

The variance's derivative. σ² = (1/n) Σ_k (center_k)², and each center_k is linear, so the product rule gives 2 center_k(x) · center_k per term.

theorem Proofs.bnVarDeriv_basisVec (n : ℕ) (x : Vec n) (i : Fin n) :

(bnVarDeriv n x) (basisVec i) = 2 * (x i - bnMean n x) / ↑n

∂σ²/∂xᵢ = 2 (xᵢ − μ) / n — the derivative evaluated on a basis vector; the (1 − 1/n) factor cancels because the centered entries sum to zero.

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.

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

istd = (√(σ²+ε))⁻¹: the chain rule through Real.sqrt and x ↦ x⁻¹ on top of bnVar_hasFDerivAt, evaluated at basisVec i by bnVarDeriv_basisVec.

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, not 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.bnAffineHasVJP (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.bnAffineHasVJP n γ β = { backward := fun (_v dy : Proofs.Vec n) (i : Fin n) => γ * dy i, correct := ⋯ }

Instances For

noncomputable def Proofs.bnNormalizeHasVJP (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.

Instances For

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

HasVJP (bnForward n ε γ β)

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

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

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

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

Equations

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

Instances For

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

Differentiable ℝ (bnForward n ε γ β)

bnForward is differentiable everywhere (for ε > 0).

Reuses the exact differentiability argument inside bnHasVJP: 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_differentiable), and bnAffine is affine. The ε > 0 hypothesis is what licenses the inverse-sqrt smoothness. This is the differentiability witness vjpCompAt 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) :

bnGradInput n ε γ x dy i = ∑ j : Fin n, pdiv (bnForward n ε γ β) x i j * dy j

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

theorem Proofs.bnForward_chan_diff_γ {n : ℕ} (ε γ β : ℝ) (z : Vec n) (k₀ k₁ : Fin n) :

bnForward n ε γ β z k₀ - bnForward n ε γ β z k₁ = γ * (z k₀ - z k₁) * bnIstd n z ε

BN acts on coordinate differences by γ·istd — the exact identity that propagates the carrier undamped through every BN, and the reason no BN-variance derivative is ever taken: the difference is γ · (difference) · istd with istd evaluated at the same activation.

theorem Proofs.bnForward_abs_sub_le {n : ℕ} (ε γ β : ℝ) (hε : 0 < ε) (v : Vec n) (k : Fin n) :

|bnForward n ε γ β v k - β| ≤ |γ| * √↑n

The two-sided BN margin |bn − β| ≤ |γ|·√n, with no mean/variance computation (bnXhat_sq_le). This is what makes a large β keep a relu off its kink at every input, so the structural net needs no eventually-argument for its relus.

theorem Proofs.sqrt_lt_param (n : ℕ) (β : ℝ) (hβ : 0 ≤ β) (h : ↑n < β ^ 2) :

√↑n < β

√n < β from n < β² — the margin check at each of the witness's BN widths.

theorem Proofs.bnIstd_continuous {n : ℕ} (ε : ℝ) (hε : 0 < ε) :

Continuous fun (v : Vec n) => bnIstd n v ε

bnIstd is continuous in the activation (ε > 0); a fun_prop atom.