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:
| piece | from |
|---|---|
reassocFwd/reassocBack + VJPs | PerChannelBN.lean (the per-channel BN layout bridge) |
transpose_has_vjp | Tensor.lean |
layerNormVec + layerNormVec_per_token_has_vjp_mat | ViTVecLN.lean (ViT's [192] LN) |
hasVJPMat_to_hasVJP | Tensor.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.
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.
The inverse direction, from reassocFwdIdx_val through the round-trip.
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
- Proofs.rowLNVecFlat s c ε γ β v = Proofs.Mat.flatten ((fun (X : Proofs.Mat s c) (r : Fin s) => Proofs.layerNormVec c ε γ β (X r)) (Proofs.Mat.unflatten v))
Instances For
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
- Proofs.rowLNVecFlat_has_vjp s c ε γ β hε = Proofs.hasVJPMat_to_hasVJP (Proofs.layerNormVec_per_token_has_vjp_mat s c ε γ β hε)
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.
Instances For
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
Everywhere-differentiable given 0 < ε — four permutations and one LN.
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
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.
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
- Proofs.chanLNRows c h w x = Proofs.StableHLO.transposeFlat c (h * w) (Proofs.reassocFwd c h w x)
Instances For
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
- Proofs.chanRowsIdx c h w j = finProdFinEquiv ((finProdFinEquiv.symm (Proofs.reassocBackIdx c h w j)).2, (finProdFinEquiv.symm (Proofs.reassocBackIdx c h w j)).1)
Instances For
The inverse direction — the row-view index o, read back in the activation layout.
Equations
- Proofs.chanRowsIdxInv c h w o = Proofs.reassocFwdIdx c h w (finProdFinEquiv ((finProdFinEquiv.symm o).2, (finProdFinEquiv.symm o).1))
Instances For
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
- Proofs.chanRowsPerm c h w = { toFun := Proofs.chanRowsIdx c h w, invFun := Proofs.chanRowsIdxInv c h w, left_inv := ⋯, right_inv := ⋯ }
Instances For
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).
And channel-LN is the row-LN read through the permutation.
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.
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.)
The β peer: const + gather β.
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.
The β peer of chanLN_gamma_contract.
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.
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.
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.
The graph's ▸ transport IS the math's Mat-split bridge — den_cast composed with
reassocFwdIdx_val. This is the lemma that keeps ConvNeXtRender's reassoc and
chanLNTensor3 describing one function.