Closing the ViT render — the parameter-gradient close (ch10 Item C)

planning/archive/vit_close.md Item C, applied to the representative 2-block / 1-head ViT (vitForward2, the proven whole-net VJP vitForward2_has_vjp; graph vitFwdGraph). Generic in the cotangent dy the backward chain delivers at each site's output (pinning that cotangent to the actual attention chain is the optional Item D), batch-1 — everything in a ViT is per-example separable (the EfficientNet contrast).

family (render SSA)	forward fn	certified by
Wq/Wk/Wv/Wo, Wfc1/Wfc2 + biases	per-token dense (rowwise)	vit_render_rowdense{W,b}_certified (new family): dW = Σ_tokens xᵣ ⊗ dyᵣ, db = Σ_tokens dyᵣ — the M2 outer-product bridge row-lifted
classifier Wcls/bcls	dense on the CLS row	M2 weight/bias_grad_bridge (reuse — single-vector dense)
LN γ/β ×5 sites (scalar, per-token)	rowwise layerNormForward	vit_render_rowln{gamma,beta}_certified (new): the ConvNeXtClose Vec 1 embedding row-lifted — dγ = Σ_{tokens} Σ_D dy·x̂, dβ = Σ Σ dy (affine in the params ⇒ no 0 < ε)
pos_embed	additive (patchEmbed_flat)	vit_render_pos_certified: the pos-Jacobian is the identity ⇒ dPos = dy
cls_token	row-0 scatter (patchEmbed_flat)	vit_render_cls_certified: masked-gather Jacobian ⇒ dCls = dy row-0 slice
patch conv Wp/bp	stride-P conv (patchEmbed_flat)	vit_render_patch{W,b}_certified: kernel-linear w/ constant guarded reads ⇒ dWp = Σ_p read·dy_(p+1), dbp = Σ_p dy_(p+1) (CLS row excluded)
attention internals (softmax, scale)	—	no parameters

Two genuinely-new bridge families (everything else is reuse or a reindex):

• Per-token dense W/b — every row of [N,a] through the same W : [a,c] (+ b). The W-Jacobian of the flattened rowwise dense is block-sparse — ∂y_(r,k)/∂W_(i,j) = X_(r,i)·δ_(k,j) — so the rendered per-token outer-product reduce dW_(i,j) = Σ_r X_(r,i)·dY_(r,j) (one dot_general contracting the token axis) is the certified contraction. Covers Wq/Wk/Wv/Wo, Wfc1/Wfc2 and their biases at every block.
• Row-lifted scalar-LN γ/β — the ConvNeXtClose Vec 1 embedding generalized from one LN site over Vec n to N token rows: as a function of γ' : Vec 1 the rowwise LN is affine, γ' ↦ fun (r,k) => x̂_r(k)·γ'(0) + β, so the Jacobian is the flattened per-row x̂ and the rendered whole-tensor reduce dγ = Σ_r Σ_k dY_(r,k)·x̂_r(k) = Σ_r bn_grad_gamma (X r) (dY r) is certified. Likewise dβ = Σ_r Σ_k dY_(r,k).

The classifier head (dense on the CLS vector) is VERBATIM M2 weight/bias_grad_bridge reuse at [D, nClasses]. The patch-embed conv Wp/bp (§ E) closes over patchEmbed_flat directly — the kernel is the VARIABLE and the pad-guarded image reads are CONSTANT coefficients (the mirror of the input-grad case), so the same const×reindex recipe applies with the CLS row masked out. 3-axiom clean by construction.

theorem Proofs.sum_fin_prod {M : Type u_1} [AddCommMonoid M] (m n : ℕ) (f : Fin (m * n) → M) : ∑ idx : Fin (m * n), f idx = ∑ r : Fin m, ∑ k : Fin n, f (finProdFinEquiv (r, k))

Sum over a flat product index = double sum over the factors (row-major).

theorem Proofs.pdiv_rowDense_W {N a c : ℕ} (bb : Vec c) (X : Mat N a) (W : Mat a c) (i : Fin a) (j' : Fin c) (idx : Fin (N * c)) : pdiv (fun (v : Vec (a * c)) => Mat.flatten fun (r : Fin N) => dense (Mat.unflatten v) bb (X r)) W.flatten (finProdFinEquiv (i, j')) idx = if j' = (finProdFinEquiv.symm idx).2 then X (finProdFinEquiv.symm idx).1 i else 0

Jacobian of the per-token dense w.r.t. the (flattened) shared weight — ∂y_(r,k)/∂W_(i,j') = X_(r,i)·δ_(k,j'). The row-lift of pdiv_dense_W.

noncomputable def Proofs.rowDense_weight_grad {N a c : ℕ} (X : Mat N a) (dY : Mat N c) : Mat a c

The rendered per-token dense weight gradient: the token-axis-contracted outer product dW_(i,j) = Σ_r X_(r,i)·dY_(r,j) (one dot_general contracting the token axis — the row-lift of dense_weight_grad = x ⊗ dy).

Equations

• Proofs.rowDense_weight_grad X dY i j = ∑ r : Fin N, X r i * dY r j

noncomputable def Proofs.rowDense_bias_grad {N c : ℕ} (dY : Mat N c) : Vec c

The rendered per-token dense bias gradient: the token-axis reduce db_j = Σ_r dY_(r,j).

Equations

• Proofs.rowDense_bias_grad dY j = ∑ r : Fin N, dY r j

theorem Proofs.vit_rowDenseW_grad_bridge {N a c : ℕ} (bb : Vec c) (X : Mat N a) (W : Mat a c) (dy : Vec (N * c)) (i : Fin a) (j : Fin c) : rowDense_weight_grad X (Mat.unflatten dy) i j = ∑ o : Fin (N * c), pdiv (fun (v : Vec (a * c)) => Mat.flatten fun (r : Fin N) => dense (Mat.unflatten v) bb (X r)) W.flatten (finProdFinEquiv (i, j)) o * dy o

Per-token dense W-gradient bridge. The rendered token-contracted outer product equals the certified Jacobian of the rowwise dense (as a function of the flattened shared W) contracted with the cotangent.

theorem Proofs.vit_rowDenseb_grad_bridge {N a c : ℕ} (W : Mat a c) (X : Mat N a) (bb : Vec c) (dy : Vec (N * c)) (i : Fin c) : rowDense_bias_grad (Mat.unflatten dy) i = ∑ o : Fin (N * c), pdiv (fun (b' : Vec c) => Mat.flatten fun (r : Fin N) => dense W b' (X r)) bb i o * dy o

Per-token dense b-gradient bridge. The rendered token-axis reduce equals the certified rowwise-dense ∂/∂b contraction.

theorem Proofs.vit_render_rowdenseW_certified {N a c : ℕ} (bb : Vec c) (X : Mat N a) (W : Mat a c) (dy : Vec (N * c)) (lr : ℝ) (i : Fin a) (j : Fin c) : W i j - lr * rowDense_weight_grad X (Mat.unflatten dy) i j = W i j - lr * ∑ o : Fin (N * c), pdiv (fun (v : Vec (a * c)) => Mat.flatten fun (r : Fin N) => dense (Mat.unflatten v) bb (X r)) W.flatten (finProdFinEquiv (i, j)) o * dy o

Per-token dense W output, certified. Wⁿ = W − lr·(Σ_tokens xᵣ ⊗ dyᵣ) denotes W − lr·(certified ∂(rowwise dense)/∂W · cotangent). Covers Wq/Wk/Wv/Wo and Wfc1/Wfc2 at every block of the representative ViT (each at its own [a,c]).

theorem Proofs.vit_render_rowdenseb_certified {N a c : ℕ} (W : Mat a c) (X : Mat N a) (bb : Vec c) (dy : Vec (N * c)) (lr : ℝ) (i : Fin c) : bb i - lr * rowDense_bias_grad (Mat.unflatten dy) i = bb i - lr * ∑ o : Fin (N * c), pdiv (fun (b' : Vec c) => Mat.flatten fun (r : Fin N) => dense W b' (X r)) bb i o * dy o

Per-token dense b output, certified. bⁿ = b − lr·(Σ_tokens dyᵣ) denotes the certified rowwise-dense ∂/∂b contraction. Covers all six per-block biases.

theorem Proofs.pdiv_rowLN_gamma (N D : ℕ) (ε β : ℝ) (X : Mat N D) (γ : Vec 1) (i : Fin 1) (idx : Fin (N * D)) : pdiv (fun (γ' : Vec 1) => Mat.flatten fun (r : Fin N) => layerNormForward D ε (γ' 0) β (X r)) γ i idx = bnXhat D ε (X (finProdFinEquiv.symm idx).1) (finProdFinEquiv.symm idx).2

Jacobian of the rowwise scalar-LN w.r.t. γ — ∂y_(r,k)/∂γ = x̂_r(k) (dense: the shared scalar γ scales every token's every channel).

theorem Proofs.pdiv_rowLN_beta (N D : ℕ) (ε γ : ℝ) (X : Mat N D) (β : Vec 1) (i : Fin 1) (idx : Fin (N * D)) : pdiv (fun (β' : Vec 1) => Mat.flatten fun (r : Fin N) => layerNormForward D ε γ (β' 0) (X r)) β i idx = 1

Jacobian of the rowwise scalar-LN w.r.t. β — ∂y_(r,k)/∂β = 1.

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

The rendered rowwise-LN γ gradient: the whole-tensor reduce dγ = Σ_r Σ_k dY_(r,k)·x̂_r(k) — per-row bn_grad_gamma, summed over tokens.

Equations

• Proofs.rowLN_grad_gamma N D ε X dY = ∑ r : Fin N, Proofs.bn_grad_gamma D ε (X r) (dY r)

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

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

Equations

• Proofs.rowLN_grad_beta N D dY = ∑ r : Fin N, Proofs.bn_grad_beta D (dY r)

theorem Proofs.vit_rowlnGamma_grad_bridge (N D : ℕ) (ε β : ℝ) (γ : Vec 1) (X : Mat N D) (dy : Vec (N * D)) : rowLN_grad_gamma N D ε X (Mat.unflatten dy) = ∑ idx : Fin (N * D), pdiv (fun (γ' : Vec 1) => Mat.flatten fun (r : Fin N) => layerNormForward D ε (γ' 0) β (X r)) γ 0 idx * dy idx

Rowwise scalar-LN γ-gradient bridge. The rendered whole-tensor reduce equals the certified rowwise-LN ∂/∂γ contraction.

theorem Proofs.vit_rowlnBeta_grad_bridge (N D : ℕ) (ε γ : ℝ) (β : Vec 1) (X : Mat N D) (dy : Vec (N * D)) : rowLN_grad_beta N D (Mat.unflatten dy) = ∑ idx : Fin (N * D), pdiv (fun (β' : Vec 1) => Mat.flatten fun (r : Fin N) => layerNormForward D ε γ (β' 0) (X r)) β 0 idx * dy idx

Rowwise scalar-LN β-gradient bridge. The rendered whole-tensor reduce Σ_r Σ_k dY_(r,k) equals the certified rowwise-LN ∂/∂β contraction.

theorem Proofs.vit_render_rowlngamma_certified (N D : ℕ) (ε β : ℝ) (γ : Vec 1) (X : Mat N D) (dy : Vec (N * D)) (lr : ℝ) : γ 0 - lr * rowLN_grad_gamma N D ε X (Mat.unflatten dy) = γ 0 - lr * ∑ idx : Fin (N * D), pdiv (fun (γ' : Vec 1) => Mat.flatten fun (r : Fin N) => layerNormForward D ε (γ' 0) β (X r)) γ 0 idx * dy idx

Rowwise scalar-LN γ output, certified. γⁿ = γ − lr·(Σ_tokens Σ_D dy·x̂) denotes the certified rowwise-LN ∂/∂γ contraction. Covers all five LN sites of the representative ViT (LN1/LN2 per block + the final LN).

theorem Proofs.vit_render_rowlnbeta_certified (N D : ℕ) (ε γ : ℝ) (β : Vec 1) (X : Mat N D) (dy : Vec (N * D)) (lr : ℝ) : β 0 - lr * rowLN_grad_beta N D (Mat.unflatten dy) = β 0 - lr * ∑ idx : Fin (N * D), pdiv (fun (β' : Vec 1) => Mat.flatten fun (r : Fin N) => layerNormForward D ε γ (β' 0) (X r)) β 0 idx * dy idx

Rowwise scalar-LN β output, certified.

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).

theorem Proofs.pdiv_patchEmbed_pos {ic H W P N D : ℕ} (Wc : Kernel4 D ic P P) (bc cls : Vec D) (pos : Mat (N + 1) D) (img : Vec (ic * H * W)) (i j : Fin ((N + 1) * D)) : pdiv (fun (p : Vec ((N + 1) * D)) => patchEmbed_flat ic H W P N D Wc bc cls (Mat.unflatten p) img) pos.flatten i j = if i = j then 1 else 0

Jacobian of patchEmbed_flat w.r.t. the (flattened) position embedding — the identity: pos is broadcast-added to every token.

theorem Proofs.vit_render_pos_certified {ic H W P N D : ℕ} (Wc : Kernel4 D ic P P) (bc cls : Vec D) (pos : Mat (N + 1) D) (img : Vec (ic * H * W)) (dy : Vec ((N + 1) * D)) (lr : ℝ) (i : Fin ((N + 1) * D)) : pos.flatten i - lr * dy i = pos.flatten i - lr * ∑ j : Fin ((N + 1) * D), pdiv (fun (p : Vec ((N + 1) * D)) => patchEmbed_flat ic H W P N D Wc bc cls (Mat.unflatten p) img) pos.flatten i j * dy j

pos-embed output, certified. The pos Jacobian is the identity, so the rendered dPos = dy (the cotangent itself, batch-summed by the batched render) is the certified contraction.

noncomputable def Proofs.cls_token_grad {N D : ℕ} (dy : Vec ((N + 1) * D)) : Vec D

The rendered CLS-token gradient: the row-0 slice of the patch-embed output cotangent (clsSliceF's shape, applied to the embed cotangent).

Equations

• Proofs.cls_token_grad dy i = dy (finProdFinEquiv (0, i))

theorem Proofs.pdiv_patchEmbed_cls {ic H W P N D : ℕ} (Wc : Kernel4 D ic P P) (bc cls : Vec D) (pos : Mat (N + 1) D) (img : Vec (ic * H * W)) (i : Fin D) (j : Fin ((N + 1) * D)) : pdiv (fun (cl : Vec D) => patchEmbed_flat ic H W P N D Wc bc cl pos img) cls i j = (if ↑(finProdFinEquiv.symm j).1 = 0 then 1 else 0) * if i = (finProdFinEquiv.symm j).2 then 1 else 0

Jacobian of patchEmbed_flat w.r.t. the CLS token — the row-0 masked gather: ∂y_(n,k)/∂cls_i = [n = 0]·δ_(i,k).

theorem Proofs.vit_render_cls_certified {ic H W P N D : ℕ} (Wc : Kernel4 D ic P P) (bc cls : Vec D) (pos : Mat (N + 1) D) (img : Vec (ic * H * W)) (dy : Vec ((N + 1) * D)) (lr : ℝ) (i : Fin D) : cls i - lr * cls_token_grad dy i = cls i - lr * ∑ j : Fin ((N + 1) * D), pdiv (fun (cl : Vec D) => patchEmbed_flat ic H W P N D Wc bc cl pos img) cls i j * dy j

CLS-token output, certified. clsⁿ = cls − lr·(row-0 slice of the embed cotangent) denotes the certified ∂(patchEmbed)/∂cls contraction.

noncomputable def Proofs.patchRead (ic H W P : ℕ) (img : Vec (ic * H * W)) (c : Fin ic) (kh kw : Fin P) (p : ℕ) : ℝ

The pad-guarded patch read of patchEmbed_flat, named: input pixel (c, h'·P + kh, w'·P + kw) of patch p (row-major patch grid of width W/P), zero out of range. Constant in the kernel.

Equations

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

theorem Proofs.pdiv_patchEmbed_W {ic H W P N D : ℕ} (Wc : Kernel4 D ic P P) (bc cls : Vec D) (pos : Mat (N + 1) D) (img : Vec (ic * H * W)) (d : Fin D) (c : Fin ic) (kh kw : Fin P) (idx : Fin ((N + 1) * D)) : pdiv (fun (v : Vec (D * ic * P * P)) => patchEmbed_flat ic H W P N D (Kernel4.unflatten v) bc cls pos img) Wc.flatten (finProdFinEquiv (finProdFinEquiv (finProdFinEquiv (d, c), kh), kw)) idx = if (finProdFinEquiv.symm idx).2 = d then if ↑(finProdFinEquiv.symm idx).1 = 0 then 0 else patchRead ic H W P img c kh kw (↑(finProdFinEquiv.symm idx).1 - 1) else 0

Jacobian of patchEmbed_flat w.r.t. the (flattened) patch kernel — ∂y_(n,dd)/∂W_(d,c,kh,kw) = [n ≠ 0]·δ_(dd,d)·read(c,kh,kw, patch n−1).

noncomputable def Proofs.patchEmbed_weight_grad (ic H W P N D : ℕ) (img : Vec (ic * H * W)) (dy : Vec ((N + 1) * D)) : Kernel4 D ic P P

The rendered patch-kernel gradient: for each tap (d,c,kh,kw), the patch-grid reduce Σ_p read(c,kh,kw,p)·dy_(p+1,d) — the "dilate dy / valid conv" weight grad, with the CLS row (token 0) excluded.

Equations

• Proofs.patchEmbed_weight_grad ic H W P N D img dy d c kh kw = ∑ n : Fin N, Proofs.patchRead ic H W P img c kh kw ↑n * dy (finProdFinEquiv (n.succ, d))

noncomputable def Proofs.patchEmbed_bias_grad (N D : ℕ) (dy : Vec ((N + 1) * D)) : Vec D

The rendered patch bias gradient: db_d = Σ_p dy_(p+1,d) (the CLS row excluded — token 0 carries no conv bias).

Equations

• Proofs.patchEmbed_bias_grad N D dy d = ∑ n : Fin N, dy (finProdFinEquiv (n.succ, d))

theorem Proofs.vit_patchW_grad_bridge {ic H W P N D : ℕ} (Wc : Kernel4 D ic P P) (bc cls : Vec D) (pos : Mat (N + 1) D) (img : Vec (ic * H * W)) (dy : Vec ((N + 1) * D)) (d : Fin D) (c : Fin ic) (kh kw : Fin P) : patchEmbed_weight_grad ic H W P N D img dy d c kh kw = ∑ o : Fin ((N + 1) * D), pdiv (fun (v : Vec (D * ic * P * P)) => patchEmbed_flat ic H W P N D (Kernel4.unflatten v) bc cls pos img) Wc.flatten (finProdFinEquiv (finProdFinEquiv (finProdFinEquiv (d, c), kh), kw)) o * dy o

Patch-kernel gradient bridge. The rendered patch-grid reduce equals the certified ∂(patchEmbed)/∂W contraction.

theorem Proofs.vit_render_patchW_certified {ic H W P N D : ℕ} (Wc : Kernel4 D ic P P) (bc cls : Vec D) (pos : Mat (N + 1) D) (img : Vec (ic * H * W)) (dy : Vec ((N + 1) * D)) (lr : ℝ) (d : Fin D) (c : Fin ic) (kh kw : Fin P) : Wc d c kh kw - lr * patchEmbed_weight_grad ic H W P N D img dy d c kh kw = Wc d c kh kw - lr * ∑ o : Fin ((N + 1) * D), pdiv (fun (v : Vec (D * ic * P * P)) => patchEmbed_flat ic H W P N D (Kernel4.unflatten v) bc cls pos img) Wc.flatten (finProdFinEquiv (finProdFinEquiv (finProdFinEquiv (d, c), kh), kw)) o * dy o

Patch-kernel output, certified. Wpⁿ = Wp − lr·(patch-grid reduce) denotes the certified ∂(patchEmbed)/∂Wp contraction.

theorem Proofs.pdiv_patchEmbed_b {ic H W P N D : ℕ} (Wc : Kernel4 D ic P P) (bc cls : Vec D) (pos : Mat (N + 1) D) (img : Vec (ic * H * W)) (i : Fin D) (j : Fin ((N + 1) * D)) : pdiv (fun (b' : Vec D) => patchEmbed_flat ic H W P N D Wc b' cls pos img) bc i j = (if ↑(finProdFinEquiv.symm j).1 = 0 then 0 else 1) * if i = (finProdFinEquiv.symm j).2 then 1 else 0

Jacobian of patchEmbed_flat w.r.t. the patch bias — the row-masked gather ∂y_(n,k)/∂bc_i = [n ≠ 0]·δ_(i,k) (token 0 is the CLS row).

theorem Proofs.vit_patchb_grad_bridge {ic H W P N D : ℕ} (Wc : Kernel4 D ic P P) (bc cls : Vec D) (pos : Mat (N + 1) D) (img : Vec (ic * H * W)) (dy : Vec ((N + 1) * D)) (i : Fin D) : patchEmbed_bias_grad N D dy i = ∑ o : Fin ((N + 1) * D), pdiv (fun (b' : Vec D) => patchEmbed_flat ic H W P N D Wc b' cls pos img) bc i o * dy o

Patch bias gradient bridge. The rendered CLS-row-excluded reduce equals the certified ∂(patchEmbed)/∂bc contraction.

theorem Proofs.vit_render_patchb_certified {ic H W P N D : ℕ} (Wc : Kernel4 D ic P P) (bc cls : Vec D) (pos : Mat (N + 1) D) (img : Vec (ic * H * W)) (dy : Vec ((N + 1) * D)) (lr : ℝ) (i : Fin D) : bc i - lr * patchEmbed_bias_grad N D dy i = bc i - lr * ∑ o : Fin ((N + 1) * D), pdiv (fun (b' : Vec D) => patchEmbed_flat ic H W P N D Wc b' cls pos img) bc i o * dy o

Patch bias output, certified.