Closing the ConvNeXt render — the parameter-gradient close (two small new families)

planning/archive/convnext_close.md Item C, applied to the representative 2-block ConvNeXt (tests/TestConvNeXtTrain.lean; convNextForward, the proven whole-net VJP convnext_has_vjp). The close is generic in the cotangent dy the backward chain delivers at each layer's output (pinning that cotangent to the actual block chain is the optional Item D), batch-1 — LayerNorm is per-example separable, so no batched apparatus (the EfficientNet contrast).

family (render SSA)	forward fn	certified by
stem/expand/project 1×1 conv W/b	conv2d (stride 1)	cnn_render_conv{W,b}_certified (M3, reuse)
depthwise 7×7 W/b	depthwiseConv2d	cnx_render_dw7{W,b}_certified (kernel-general — pinned below)
dense head Wd/bd	matmul / +bias	M2 weight/bias_grad_bridge (reuse)
layer-scale γ (layerScaleF)	γ ⊙ x	cnx_render_lsgamma_certified (new): ∂(γⱼxⱼ)/∂γᵢ = xᵢδᵢⱼ ⇒ dγ = x ⊙ dy
scalar-LN γ/β (the bnF sites)	layerNormForward	cnx_render_ln{gamma,beta}_certified (new): the Vec 1 embedding
gelu	—	no parameters

Two genuinely-new bridge families:

• Layer-scale γ — layerScale γ x = γ ⊙ x is symmetric in (γ, x), so the γ-Jacobian is the same diagonal as the input-Jacobian with the roles swapped: ∂(γⱼxⱼ)/∂γᵢ = xᵢ·δᵢⱼ (pdiv_layerScale_gamma, the mirror of pdiv_layerScale), giving the rendered dγ = x ⊙ dy.
• Scalar-LN γ/β — the proof's LayerNorm (layerNormForward = bnForward) has scalar γ β : ℝ, and BatchNorm.lean deliberately left bn_grad_gamma/bn_grad_beta as definitions ("scalar params don't fit the pdiv/HasVJP framework cleanly"). This file closes that gap by embedding the scalar as Vec 1: as a function of γ' : Vec 1, LN is affine — γ' ↦ fun k => x̂ₖ · γ'(0) + β — so its Jacobian collapses through pdiv_add/pdiv_mul/pdiv_const/pdiv_reindex (the CifarBnClose recipe with the constant channel map Fin n → Fin 1) to ∂yₖ/∂γ = x̂ₖ (resp. 1 for β), certifying the rendered whole-n reduces dγ = Σ dy·x̂, dβ = Σ dy (bn_grad_gamma/bn_grad_beta — now bridged, not just defined). Affine in the params, so no 0 < ε needed (ε only enters the constant x̂).

The 7×7 depthwise is the only kernel size to pin (prior nets used 3×3/5×5; the generic mnv2_render_depthwise{W,b}_certified is kernel-general — stride-1 only, ConvNeXt blocks don't downsample). The 1×1 convs (stem/expand/project) and the dense head are verbatim M3/M2 reuse at the ConvNeXt shapes. 3-axiom clean by inheritance.

theorem Proofs.cnx_render_dw7W_certified {c h w : ℕ} (b : Vec c) (x : Tensor3 c h w) (W : DepthwiseKernel c 7 7) (dy : Tensor3 c h w) (lr : ℝ) (ci : Fin c) (hi wi : Fin 7) : W ci hi wi - lr * (depthwise_weight_grad_has_vjp3 b x).backward W dy ci hi wi = W ci hi wi - lr * ∑ co : Fin c, ∑ ho : Fin h, ∑ wo : Fin w, pdiv3 (fun (W' : DepthwiseKernel c 7 7) => depthwiseConv2d W' b x) W ci hi wi co ho wo * dy co ho wo

7×7 depthwise weight output, certified. The generic depthwise weight bridge at kH=kW=7; covers every ConvNeXt block's depthwise (stride-1 — ConvNeXt blocks keep resolution).

theorem Proofs.cnx_render_dw7b_certified {c h w : ℕ} (W : DepthwiseKernel c 7 7) (x : Tensor3 c h w) (b : Vec c) (dy : Vec (c * h * w)) (lr : ℝ) (cc : Fin c) : b cc - lr * (depthwise_bias_grad_has_vjp W x).backward b dy cc = b cc - lr * ∑ j : Fin (c * h * w), pdiv (fun (b' : Vec c) => (depthwiseConv2d W b' x).flatten) b cc j * dy j

7×7 depthwise bias output, certified.

theorem Proofs.pdiv_layerScale_gamma {n : ℕ} (x γ : Vec n) (i j : Fin n) : pdiv (fun (γ' : Vec n) => layerScale γ' x) γ i j = if i = j then x i else 0

Jacobian of layerScale w.r.t. γ — ∂(γ_j x_j)/∂γ_i = x_i δ_{ij}. The parameter mirror of pdiv_layerScale.

noncomputable def Proofs.layerScale_grad_gamma {n : ℕ} (x dy : Vec n) : Vec n

The rendered layer-scale γ gradient: dγ_i = x_i · dy_i (elementwise multiply of the saved layer input with the cotangent — the layerScaleF-shaped backward).

Equations

• Proofs.layerScale_grad_gamma x dy i = x i * dy i

theorem Proofs.layerScale_gamma_grad_bridge {n : ℕ} (x γ dy : Vec n) (i : Fin n) : layerScale_grad_gamma x dy i = ∑ j : Fin n, pdiv (fun (γ' : Vec n) => layerScale γ' x) γ i j * dy j

Layer-scale γ-gradient bridge. The rendered dγ = x ⊙ dy equals the certified Jacobian of layerScale (as a function of γ) contracted with the cotangent dy.

theorem Proofs.cnx_render_lsgamma_certified {n : ℕ} (x γ dy : Vec n) (lr : ℝ) (i : Fin n) : γ i - lr * layerScale_grad_gamma x dy i = γ i - lr * ∑ j : Fin n, pdiv (fun (γ' : Vec n) => layerScale γ' x) γ i j * dy j

Layer-scale γ output, certified. γⁿ = γ − lr·(x ⊙ dy) denotes γ − lr·(certified ∂(layerScale)/∂γ · cotangent). The multiplicative-bias peer of cnn_render_convb_certified.

theorem Proofs.cnx_lnGamma_grad_bridge (n : ℕ) (ε β : ℝ) (γ : Vec 1) (x dy : Vec n) : bn_grad_gamma n ε x dy = ∑ j : Fin n, pdiv (fun (γ' : Vec 1) => layerNormForward n ε (γ' 0) β x) γ 0 j * dy j

Scalar-LN γ-gradient bridge. The rendered whole-n reduce dγ = Σ_j dy_j·x̂_j (bn_grad_gamma) equals the certified Jacobian of scalar-LN (as a function of γ, embedded Vec 1) contracted with the cotangent dy. Closes the gap BatchNorm.lean documented — bn_grad_gamma is now bridged, not just defined.

theorem Proofs.cnx_lnBeta_grad_bridge (n : ℕ) (ε γ : ℝ) (β : Vec 1) (x dy : Vec n) : bn_grad_beta n dy = ∑ j : Fin n, pdiv (fun (β' : Vec 1) => layerNormForward n ε γ (β' 0) x) β 0 j * dy j

Scalar-LN β-gradient bridge. Likewise the rendered dβ = Σ_j dy_j (bn_grad_beta) is the certified scalar-LN ∂/∂β contraction.

theorem Proofs.cnx_render_lngamma_certified (n : ℕ) (ε β : ℝ) (γ : Vec 1) (x dy : Vec n) (lr : ℝ) : γ 0 - lr * bn_grad_gamma n ε x dy = γ 0 - lr * ∑ j : Fin n, pdiv (fun (γ' : Vec 1) => layerNormForward n ε (γ' 0) β x) γ 0 j * dy j

Scalar-LN γ output, certified. γⁿ = γ − lr·(Σ dy·x̂) denotes γ − lr·(certified ∂(layerNormForward)/∂γ · cotangent). Covers the stem-LN, both block-LNs, and the head-LN of the representative ConvNeXt (each at its own n).

theorem Proofs.cnx_render_lnbeta_certified (n : ℕ) (ε γ : ℝ) (β : Vec 1) (x dy : Vec n) (lr : ℝ) : β 0 - lr * bn_grad_beta n dy = β 0 - lr * ∑ j : Fin n, pdiv (fun (β' : Vec 1) => layerNormForward n ε γ (β' 0) x) β 0 j * dy j

Scalar-LN β output, certified. βⁿ = β − lr·(Σ dy) denotes the certified scalar-LN ∂/∂β contraction.