Documentation

LeanMlir.Proofs.Architectures.ChannelLN

The channel LayerNorm — per-position LN over channels, with a per-channel affine #

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 from bnForward over the whole flattened c·h·w map with a scalar γ/β on 22 of ConvNeXt-T's 23 LayerNorm sites (stem, 18 blocks, 3 downsamples, head); the head 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:

piecefrom
reassocFwd/reassocBack + VJPsPerChannelBN.lean (the per-channel BN layout bridge)
transposeHasVJPTensor.lean
layerNormVec + layerNormVecPerTokenHasVJPMatLayerNorm (ViT's [192] LN uses the same op)
HasVJPMat.toHasVJPTensor.lean

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 the index cast reassoc (IndexCast); 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.

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.

References #

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

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

source
noncomputable def Proofs.rowLNVecFlat (s c : ℕ) (ε : ℝ) (γ β : Vec c) :
Vec (s * c) → Vec (s * c)

Rowwise vector-LN on the flat [s, c] layout — s 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
    source
    theorem Proofs.rowLNVecFlat_gamma_differentiable (s c : ℕ) (ε : ℝ) (β : Vec c) (x : Vec (s * c)) :
    Differentiable ℝ fun (γ : Vec c) => rowLNVecFlat s c ε γ β x

    rowLNVecFlat is differentiable in γ (the parameter-gradient nodes vary it).

    source
    theorem Proofs.rowLNVecFlat_beta_differentiable (s c : ℕ) (ε : ℝ) (γ : Vec c) (x : Vec (s * c)) :
    Differentiable ℝ fun (β : Vec c) => rowLNVecFlat s c ε γ β x

    …and in β.

    source
    theorem Proofs.rowLNVecFlat_differentiable (s c : ℕ) (ε : ℝ) (γ β : Vec c) (hε : 0 < ε) :
    Differentiable ℝ (rowLNVecFlat s c ε γ β)
    source
    noncomputable def Proofs.rowLNVecFlatHasVJP (s c : ℕ) (ε : ℝ) (γ β : Vec c) (hε : 0 < ε) :
    HasVJP (rowLNVecFlat s c ε γ β)

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

    Equations
    Instances For
      source
      theorem Proofs.transposeFlat_differentiable (m n : ℕ) :
      Differentiable ℝ (StableHLO.transposeFlat m n)

      transposeFlat is a coordinate permutation — Attention's transpose_flat_differentiable.

      source
      noncomputable def Proofs.transposeFlatHasVJP (m n : ℕ) :
      HasVJP (StableHLO.transposeFlat m n)

      transposeFlat's VJP is Tensor.lean's transposeHasVJP 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
        source
        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
          source
          theorem Proofs.chanLNTensor3_differentiable (c h w : ℕ) (ε : ℝ) (γ β : Vec c) (hε : 0 < ε) :
          Differentiable ℝ (chanLNTensor3 c h w ε γ β)

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

          source
          theorem Proofs.chanLNTensor3_gamma_differentiable (c h w : ℕ) (ε : ℝ) (β : Vec c) (x : Vec (c * h * w)) :
          Differentiable ℝ fun (γ : Vec c) => chanLNTensor3 c h w ε γ β x

          chanLNTensor3 is differentiable in γ (the parameter-gradient nodes vary it).

          source
          theorem Proofs.chanLNTensor3_beta_differentiable (c h w : ℕ) (ε : ℝ) (γ : Vec c) (x : Vec (c * h * w)) :
          Differentiable ℝ fun (β : Vec c) => chanLNTensor3 c h w ε γ β x

          …and in β.

          source
          noncomputable def Proofs.chanLNTensor3HasVJP (c h w : ℕ) (ε : ℝ) (γ β : Vec c) (hε : 0 < ε) :
          HasVJP (chanLNTensor3 c h w ε γ β)

          Channel-LN VJP (global) — vjpComp over the five proven pieces. The only hypothesis is the LN positivity 0 < ε, exactly as the scalar layerNormHasVJP it replaces. A term, not a tactic proof, so its .backward unfolds to the nested chain.

          Equations
          • One or more equations did not get rendered due to their size.
          Instances For
            source
            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.

            source
            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
              source
              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
                source
                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
                  source
                  theorem Proofs.chanRowsIdxInv_chanRowsIdx (c h w : ℕ) (j : Fin (c * h * w)) :
                  chanRowsIdxInv c h w (chanRowsIdx c h w j) = j
                  source
                  theorem Proofs.chanRowsIdx_chanRowsIdxInv (c h w : ℕ) (o : Fin (h * w * c)) :
                  chanRowsIdx c h w (chanRowsIdxInv c h w o) = o
                  source
                  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
                    source
                    theorem Proofs.pdiv_reindexOut_contract {m n n' : ℕ} (f : Vec m → Vec 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.

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

                    source
                    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 β.

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

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