Documentation

LeanMlir.Proofs.Nets.ConvNeXt.ConvNeXtChannelLN

ConvNeXt's REAL channel LayerNorm — the math side of §2m Route A #

convnextVerified's LN was bnForward over the whole flattened c·h·w map with a scalar γ/β. ConvNeXt specifies channel_layer_norm: h·w statistics per example, each over the c channels at ONE spatial position, with a per-channel [c] affine. That is a different function on 21 of the net's 22 sites (the 22nd is the head, which runs after GAP where there is no spatial extent left, so reducing "everything" already IS reducing over channels).

Route A — no new SHlo op, and no new VJP. ConvNeXt's channel-LN is ViT's row-LN under a transpose: view one example as [c, s] with s = h·w, transpose to [s, c], and each row is one spatial position holding its c channels — exactly what ViT's layerNormVec normalises. Every piece below is already proven and shipping:

piecefrom
reassocFwd/reassocBack + VJPsPerChannelBN.lean (the per-channel BN layout bridge)
transpose_has_vjpTensor.lean
layerNormVec + layerNormVec_per_token_has_vjp_matViTVecLN.lean (ViT's [192] LN)
hasVJPMat_to_hasVJPTensor.lean

Settled on device before any of this was written (lake build channel-ln): the composition ties the closed form at rel 0 forward and on all three backward pieces, the incumbent .bnF control fires at rel 0.82, and the transposes measure free (Δ 0.00 ms on 16.1 ms of whole-net LN).

⚠ The seam this file closes #

Nat multiplication is not definitionally associative: the ambient activation index is c*h*w = (c*h)*w while the transpose needs c*(h*w). The render spells that with a transport (ConvNeXtRender.reassoc); the math spells it with PerChannelBN's finProdFinEquiv re-association, whose "row c is channel c" reading is what makes the composition legibly a channel LN. Nothing forces those two to be the same map, and if they are not, the math and the artifact are different functions with no gate between them — §2k's own sin in a new place.

They ARE the same map, and reassocFwdIdx_val proves it: row-major finProdFinEquiv sends both ((c,hi),wi) and (c,(hi,wi)) to the same linear offset, so the bridge preserves the underlying natural and is therefore exactly the type-level cast. den_reassocS lifts that to the graph.

theorem Proofs.reassocFwdIdx_val (oc h w : ) (k : Fin (oc * (h * w))) :
(reassocFwdIdx oc h w k) = k

The Mat-split bridge is the Nat.mul_assoc cast. finProdFinEquiv is row-major, so ((c,hi),wi) ↦ wi + w·hi + w·h·c and (c,(hi,wi)) ↦ wi + w·hi + h·w·c are the same offset; the re-association therefore preserves Fin.val. This is what lets the proof-side graph transport its index with while the denotation stays on reassocFwd.

theorem Proofs.reassocBackIdx_val (oc h w : ) (k : Fin (oc * h * w)) :
(reassocBackIdx oc h w k) = k

The inverse direction, from reassocFwdIdx_val through the round-trip.

noncomputable def Proofs.rowLNVecFlat (s c : ) (ε : ) (γ β : Vec c) :
Vec (s * c)Vec (s * c)

Rowwise vector-LN on the flat [s, c] layouts spatial rows, each normalised over its c channels and then given the per-channel affine. Literally ViT's per-token LN with "token" read as "spatial position"; that re-reading is the whole of Route A.

Equations
Instances For
    theorem Proofs.rowLNVecFlat_diff (s c : ) (ε : ) (γ β : Vec c) ( : 0 < ε) :
    noncomputable def Proofs.rowLNVecFlat_has_vjp (s c : ) (ε : ) (γ β : Vec c) ( : 0 < ε) :
    HasVJP (rowLNVecFlat s c ε γ β)

    ViT's per-token LN VJP, bridged to the flat layout. No new proof — layerNormVec_has_vjp is (+β) ∘ layerScale γ ∘ LN(1,0) and needs only 0 < ε.

    Equations
    Instances For

      transposeFlat is a coordinate permutation, hence a reindexCLM.

      transposeFlat's VJP is Tensor.lean's transpose_has_vjp through the flatten bijection — the flat form is definitionally the bridged Mat form, so this is a re-typing, not a proof.

      Equations
      Instances For
        noncomputable def Proofs.chanLNTensor3 (c h w : ) (ε : ) (γ β : Vec c) :
        Vec (c * h * w)Vec (c * h * w)

        ConvNeXt's channel LayerNorm on the activation layout the convolutions use (Vec (c*h*w)): re-associate to the Mat-split [c, h·w], transpose to [h·w, c] so each row is one spatial position, normalise that row over its c channels with the per-channel [c] affine, then transpose and re-associate back.

        Contrast the incumbent layerNormForward (c*h*w) ε γ β, which takes ONE mean and ONE variance over all c·h·w values and applies two scalars — for a stage-1 site that is one statistic where ConvNeXt wants 3,136 of them.

        Equations
        • One or more equations did not get rendered due to their size.
        Instances For
          theorem Proofs.chanLNTensor3_diff (c h w : ) (ε : ) (γ β : Vec c) ( : 0 < ε) :

          Everywhere-differentiable given 0 < ε — four permutations and one LN.

          noncomputable def Proofs.chanLNTensor3_has_vjp (c h w : ) (ε : ) (γ β : Vec c) ( : 0 < ε) :
          HasVJP (chanLNTensor3 c h w ε γ β)

          Channel-LN VJP (global)vjp_comp over the five proven pieces. The only hypothesis is the LN positivity 0 < ε, exactly as the scalar layerNorm_has_vjp it replaces.

          Equations
          • One or more equations did not get rendered due to their size.
          Instances For
            theorem Proofs.rowLN_affine_eq (s c : ) (ε : ) (γ β : Vec c) (u : Vec (s * c)) :
            StableHLO.rowBiasFlat s c β (StableHLO.rowScaleFlat s c γ (StableHLO.rowLNFlat s c ε 1 0 u)) = rowLNVecFlat s c ε γ β u

            The emitted three-op affine tail IS the per-token vector-LN. The chain normalises with lnRowF at scalar γ=1/β=0 and then applies the REAL [c] affine with rowScaleF/rowBiasF — ViT's spelling, and the reason ConvNeXt needs no new op. This is the lemma that lets the graph's five denotations collapse onto chanLNTensor3's three.

            noncomputable def Proofs.chanLNRows (c h w : ) (x : Vec (c * h * w)) :
            Vec (h * w * c)

            The saved activation as the row backward sees it: the [h·w, c] view of x, one row per spatial position holding its c channels. Naming this keeps the backward's operating-point hypotheses (bnIstd/bnXhat per row) readable — it is chanLNTensor3's own first two factors, and it is the value the render's re-emitted transposeF pair denotes.

            Equations
            Instances For
              noncomputable def Proofs.chanRowsIdx (c h w : ) (j : Fin (c * h * w)) :
              Fin (h * w * c)

              chanLNTensor3's conjugation as ONE index map: the activation index j of the c·h·w layout, read in the [h·w, c] row view. Composition of the Mat-split re-association and the transpose, both pure reindexes.

              Equations
              Instances For
                noncomputable def Proofs.chanRowsIdxInv (c h w : ) (o : Fin (h * w * c)) :
                Fin (c * h * w)

                The inverse direction — the row-view index o, read back in the activation layout.

                Equations
                Instances For
                  theorem Proofs.chanRowsIdxInv_chanRowsIdx (c h w : ) (j : Fin (c * h * w)) :
                  chanRowsIdxInv c h w (chanRowsIdx c h w j) = j
                  theorem Proofs.chanRowsIdx_chanRowsIdxInv (c h w : ) (o : Fin (h * w * c)) :
                  chanRowsIdx c h w (chanRowsIdxInv c h w o) = o
                  noncomputable def Proofs.chanRowsPerm (c h w : ) :
                  Fin (c * h * w) Fin (h * w * c)

                  The conjugation is a permutation. Both directions are finProdFinEquiv round-trips, so the [c·h·w] ↔ [h·w, c] relabeling is a genuine bijection — the fact the two certs below turn into "the adjoint is the inverse".

                  Equations
                  Instances For
                    theorem Proofs.chanLNRows_eq_reindex (c h w : ) (v : Vec (c * h * w)) :
                    chanLNRows c h w v = fun (o : Fin (h * w * c)) => v ((chanRowsPerm c h w).symm o)

                    The row view IS the reindex by the inverse permutation (definitional; stated so the two spellings are visibly one map rather than two that happen to agree).

                    theorem Proofs.chanLNTensor3_eq_rows (c h w : ) (ε : ) (γ β : Vec c) (x : Vec (c * h * w)) :
                    chanLNTensor3 c h w ε γ β x = fun (j : Fin (c * h * w)) => rowLNVecFlat (h * w) c ε γ β (chanLNRows c h w x) ((chanRowsPerm c h w) j)

                    And channel-LN is the row-LN read through the permutation.

                    theorem Proofs.pdiv_reindexOut_contract {m n n' : } (f : Vec mVec n) (x : Vec m) (hf : DifferentiableAt f x) (σ : Fin n' Fin n) (i : Fin m) (cot : Vec n') :
                    j : Fin n', pdiv (fun (y : Vec m) (k : Fin n') => f y (σ k)) x i j * cot j = o : Fin n, pdiv f x i o * cot (σ.symm o)

                    An output-side permutation moves onto the cotangent as its inverse. Generic: for any differentiable f and any bijection σ of output indices, contracting the Jacobian of σ-reindexed f with a cotangent is contracting f's own Jacobian with the σ⁻¹-reindexed cotangent. pdiv_comp against pdiv_reindex's indicator, then Equiv.sum_comp.

                    theorem Proofs.rowLNVecFlat_gamma_diffAt (s c : ) (ε : ) (β : Vec c) (X : Vec (s * c)) (γ : Vec c) :
                    DifferentiableAt (fun (γ' : Vec c) => rowLNVecFlat s c ε γ' β X) γ

                    As a function of γ the row-LN is x̂ ⊙ gather γ + β — a masked gather plus a constant, hence differentiable. (ViT proves the Jacobian of this shape; the differentiability is what pdiv_comp needs and what it does not export.)

                    theorem Proofs.rowLNVecFlat_beta_diffAt (s c : ) (ε : ) (γ : Vec c) (X : Vec (s * c)) (β : Vec c) :
                    DifferentiableAt (fun (β' : Vec c) => rowLNVecFlat s c ε γ β' X) β

                    The β peer: const + gather β.

                    theorem Proofs.chanLN_gamma_contract {c h w : } (ε : ) (β γ : Vec c) (x cot : Vec (c * h * w)) (k : Fin c) :
                    j : Fin (c * h * w), pdiv (fun (γ' : Vec c) => chanLNTensor3 c h w ε γ' β x) γ k j * cot j = o : Fin (h * w * c), pdiv (fun (γ' : Vec c) => rowLNVecFlat (h * w) c ε γ' β (chanLNRows c h w x)) γ k o * chanLNRows c h w cot o

                    The γ contraction, moved to the row layout. The activation-layout Jacobian against the activation-layout cotangent equals the row-layout Jacobian against the TRANSPOSED cotangent — which is the operand lnGammaTail actually emits.

                    theorem Proofs.chanLN_beta_contract {c h w : } (ε : ) (γ β : Vec c) (x cot : Vec (c * h * w)) (k : Fin c) :
                    j : Fin (c * h * w), pdiv (fun (β' : Vec c) => chanLNTensor3 c h w ε γ β' x) β k j * cot j = o : Fin (h * w * c), pdiv (fun (β' : Vec c) => rowLNVecFlat (h * w) c ε γ β' (chanLNRows c h w x)) β k o * chanLNRows c h w cot o

                    The β peer of chanLN_gamma_contract.

                    theorem Proofs.cnx_render_chlngamma_certified {c h w : } (ε : ) (β γ : Vec c) (x cot : Vec (c * h * w)) (lr : ) (k : Fin c) :
                    γ k - lr * vecLN_grad_gamma (h * w) c ε (Mat.unflatten (chanLNRows c h w x)) (Mat.unflatten (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 γ output, certified. The rendered per-channel reduce — ViT's vecLN_grad_gamma on the two transposed views the tail emits — equals the certified Jacobian of chanLNTensor3 in its Vec c γ, contracted with the activation-layout cotangent. The Vec c peer of ConvNeXtClose.cnx_render_lngamma_certified, and the den target of the render's veclnGammaSgd LN tail.

                    theorem Proofs.cnx_render_chlnbeta_certified {c h w : } (ε : ) (γ β : Vec c) (x cot : Vec (c * h * w)) (lr : ) (k : Fin c) :
                    β k - lr * vecLN_grad_beta (h * w) c (Mat.unflatten (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 β output, certified. The β grad is the plain reduce Σ_rows dy, so the same rowDenseBiasSgd op ViT's LN-β uses denotes it here too.

                    theorem Proofs.StableHLO.den_cast {m n : } (heq : m = n) (e : SHlo m) :
                    den (heq e) = fun (k : Fin n) => den e (Fin.cast k)

                    den commutes with a type-level index transport. Transporting the GRAPH along m = n reindexes its denotation by the val-preserving Fin.cast. Stated at variable m/n so subst applies — at c*h*w = c*(h*w) neither side is a variable and it would not.

                    theorem Proofs.StableHLO.den_reassocS {c h w : } (e : SHlo (c * h * w)) :
                    den ( e) = reassocFwd c h w (den e)

                    The graph's transport IS the math's Mat-split bridgeden_cast composed with reassocFwdIdx_val. This is the lemma that keeps ConvNeXtRender's reassoc and chanLNTensor3 describing one function.

                    theorem Proofs.StableHLO.den_unassocS {c h w : } (e : SHlo (c * (h * w))) :
                    den ( e) = reassocBack c h w (den e)