Documentation

LeanMlir.Proofs.Nets.ConvNeXt.ConvNeXtFold

ConvNeXt-T §1 fold — the per-channel layer-scale γ gradient cert (the one new proof) #

The committed ConvNeXt SGD net trains per-channel layer-scale γ : Vec c (the layerScaleChF forward, which broadcasts γ over the c·h·w activation via chanIdx), NOT a per-element Vec n layer-scale. So the §1 fold needs the per-channel version: the γ-gradient w.r.t. the Vec c parameter is the per-channel reduce dγ_c = Σ_{k : chanIdx k = c} x_k · dy_k (the lsGradCh emit: multiply x dy → reduce[0,2,3]), and this is exactly the certified Jacobian of layerScaleChF's forward (as a function of γ : Vec c) contracted with the cotangent.

This is the only genuinely-NEW proof obligation for the ConvNeXt tie (the depthwise-7×7, 1×1-conv, strided-stem/downsample and dense param grads are covered by the existing M2 / M3 certs, the channel-LN γ/β by ConvNeXtChannelLN). It is linear in the parameter (pdiv_of_linear), with the chanIdx reindex (the per-channel broadcast) — ∂(γ'(chanIdx j)·x_j)/∂γ'_c = x_j·[chanIdx j = c].

The layerScaleChGammaSgd core SHlo op (the per-channel layer-scale param-SGD, emitting lsGradCh + the SGD wrap) denotes the LHS here, so den = certified is a one-line delegation to cnx_render_lsgammaCh_certified (layerScaleChGammaSgd_den).

source
theorem Proofs.CnxPoC.pdiv_layerScaleCh_gamma {c h w : ℕ} (x : Vec (c * h * w)) (γ : Vec c) (cc : Fin c) (j : Fin (c * h * w)) :
pdiv (fun (γ' : Vec c) => layerScale (fun (k : Fin (c * h * w)) => γ' (StableHLO.chanIdx c h w k)) x) γ cc j = if cc = StableHLO.chanIdx c h w j then x j else 0

Jacobian of per-channel layer-scale w.r.t. the Vec c parameter — ∂(γ'(chanIdx j)·x_j)/∂γ'_c = x_j·[chanIdx j = c]. The broadcast chanIdx reindex makes the basis vector read through it the channel indicator.

source
theorem Proofs.CnxPoC.cnx_render_lsgammaCh_certified {c h w : ℕ} (x : Vec (c * h * w)) (γ : Vec c) (dy : Vec (c * h * w)) (lr : ℝ) (cc : Fin c) :
(γ cc - lr * ∑ k : Fin (c * h * w), if StableHLO.chanIdx c h w k = cc then x k * dy k else 0) = γ cc - lr * ∑ j : Fin (c * h * w), pdiv (fun (γ' : Vec c) => layerScale (fun (k : Fin (c * h * w)) => γ' (StableHLO.chanIdx c h w k)) x) γ cc j * dy j

Per-channel layer-scale γ output, certified. The rendered per-channel reduce dγ_c = Σ_{k : chanIdx k = c} x_k·dy_k (the lsGradCh emit) equals the certified Jacobian of layerScaleChF's forward (as a function of γ : Vec c) contracted with the cotangent. The den target of the layerScaleChGammaSgd core op.

source
theorem Proofs.CnxPoC.layerScaleChGammaSgd_den {c h w : ℕ} (gN xN lrStr cotN : String) (x : Vec (c * h * w)) (γ : Vec c) (dy : Vec (c * h * w)) (lr : ℝ) (cc : Fin c) :
StableHLO.den (StableHLO.SHlo.layerScaleChGammaSgd gN xN lrStr x γ lr (StableHLO.SHlo.operand cotN dy)) cc = γ cc - lr * ∑ j : Fin (c * h * w), pdiv (fun (γ' : Vec c) => layerScale (fun (k : Fin (c * h * w)) => γ' (StableHLO.chanIdx c h w k)) x) γ cc j * dy j

Per-channel layer-scale γ op denotes the certified step. The emitted layerScaleChGammaSgd (the lsGradCh per-channel reduce + SGD) denotes γ − lr·(certified ∂(layerScaleChF)/∂γ · cot). One-line delegation to cnx_render_lsgammaCh_certified.

The channel-LN γ/β ops — the two the committed render actually emits #

ConvNeXtRender.lnGammaTail/lnBetaTail re-emit the [h·w, c] transposes and then run ViT's veclnGammaSgd / rowDenseBiasSgd on that view, so the op operands below are the transposed views chanLNRows of the saved LN input and of the chain cotangent — the values those SSA names denote. The certified Jacobian on the right is chanLNTensor3's, in the c·h·w activation layout the rest of the block lives in; ConvNeXtChannelLN's permutation argument is what lets one op serve both layouts. They cover every one of the net's 22 spatial LN sites (1 stem + 18 block + 3 downsample); the 23rd, the head, runs after GAP and is ViT's vector-LN at N = 1 (ViTPoC.veclnGammaSgd_den).

source
theorem Proofs.CnxPoC.chanLnGammaSgd_den {c h w : ℕ} (gN xN epsStr lrStr cotN : String) (ε : ℝ) (β : Vec c) (x : Vec (c * h * w)) (γ : Vec c) (cot : Vec (c * h * w)) (lr : ℝ) (k : Fin c) :
StableHLO.den (StableHLO.SHlo.veclnGammaSgd gN xN epsStr lrStr ε (chanLNRows c h w x) γ lr (StableHLO.SHlo.operand cotN (chanLNRows c h w cot))) k = γ k - lr * ∑ j : Fin (c * h * w), pdiv (fun (γ' : Vec c) => chanLNTensor3 c h w ε γ' β x) γ k j * cot j

Channel-LN γ op denotes the certified step. One-line delegation to cnx_render_chlngamma_certified. The free β is the site's LN β (the γ grad is β-free).

source
theorem Proofs.CnxPoC.chanLnBetaSgd_den {c h w : ℕ} (bN lrStr cotN : String) (ε : ℝ) (γ : Vec c) (x : Vec (c * h * w)) (β : Vec c) (cot : Vec (c * h * w)) (lr : ℝ) (k : Fin c) :
StableHLO.den (StableHLO.SHlo.rowDenseBiasSgd bN lrStr β lr (StableHLO.SHlo.operand cotN (chanLNRows c h w cot))) k = β k - lr * ∑ j : Fin (c * h * w), pdiv (fun (β' : Vec c) => chanLNTensor3 c h w ε γ β' x) β k j * cot j

Channel-LN β op denotes the certified step. The β grad is the plain row reduce, so the render uses the same rowDenseBiasSgd op ViT's LN-β does. The free ε/γ carry the LN constants (the β grad is independent of both).

source
def Proofs.CnxPoC.ChanLNGammaSgdTied (h w : ℕ) {c : ℕ} (gN xN epsStr lrStr cotN : String) (ε : ℝ) (β : Vec c) (x : Vec (c * h * w)) (γ : Vec c) (cot : Vec (c * h * w)) (lr : ℝ) :
Prop

A channel-LN γ SGD node, tied (chanLnGammaSgd_den).

Equations
  • One or more equations did not get rendered due to their size.
Instances For
    source
    def Proofs.CnxPoC.ChanLNBetaSgdTied (h w : ℕ) {c : ℕ} (bN lrStr cotN : String) (ε : ℝ) (γ : Vec c) (x : Vec (c * h * w)) (β : Vec c) (cot : Vec (c * h * w)) (lr : ℝ) :
    Prop

    A channel-LN β SGD node, tied (chanLnBetaSgd_den).

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