ConvNeXt #
A ConvNeXt block and a two-block end-to-end ConvNeXt VJP with scalar-γ, β LayerNorm, in
flattened Vec space — the ConvNeXt analogue of cnnHasVJPAt. The comparator checks this net;
the full ConvNeXt-T at the channel LayerNorm it ships with is convNextForwardTCh
(ConvNeXtFullT.lean).
ConvNeXt is the "ResNet, modernized" architecture: it keeps the residual
skeleton but swaps in the ViT-era ingredients — a large-kernel (7×7)
depthwise conv for spatial mixing, LayerNorm instead of BatchNorm,
GELU instead of ReLU, an inverted-bottleneck MLP (1×1 expand → GELU →
1×1 project), and a learnable per-channel layer scale. Crucially, the
block has no non-smooth activation (GELU is smooth everywhere) and
no max-pool, so — unlike cnnHasVJPAt — the entire VJP composes
with no kink hypotheses: the only side conditions are the LayerNorm
positivity arguments 0 < ε.
This file contributes three genuinely new pieces:
layerScale— per-channel learnable elementwise scale, a diagonal linear map;Differentiable+ aHasVJP(back i = γ i · dy i), its Jacobian read off the basis vector bypdiv_of_linear.- the ConvNeXt block body
layerScale ∘ project ∘ gelu ∘ expand ∘ LayerNorm ∘ depthwise, everywhere-differentiable, with a global VJP built by chaining the piece VJPs throughvjpComp; and the full blockresidual (block body)(identity skip, no post-add act). convnextHasVJP/convnextHasVJP_correct(and the pointwiseconvnextHasVJPAt/convnextHasVJPAt_correct) — a fixed-depth (two-block) end-to-end network: stem-patchify → stem-LN → block₁ → block₂ → global-avg-pool → head-LN → dense.
LayerNorm representation #
The layerNormForward used here is the Vec→Vec LayerNorm with scalar
γ, β that normalizes over the whole flattened vector. ConvNeXt's
LayerNorm is per spatial position over the channel axis; that is
chanLNTensor3 (ChannelLN.lean), and the full ConvNeXt-T at it is
convNextForwardTCh. The stem here is a 1×1 conv. Every other piece
(depthwise 7×7, 1×1 convs, GELU, layer scale, GAP, dense) is the same
map as in ConvNeXt-T.
References #
- Liu et al. 2022, A ConvNet for the 2020s (ConvNeXt). https://arxiv.org/abs/2201.03545
ConvNeXt block body (Vec→Vec, no skip):
layerScale γ ∘ project(1×1) ∘ gelu ∘ expand(1×1) ∘ layerNorm ∘ depthwise(7×7)
Channel/spatial dims are generic Nat params; the depthwise kernel is
c × kH × kW (the 7×7 in ConvNeXt), the expand conv lifts c → cExp
channels (the usual 4× inverted-bottleneck), gelu is applied on the
expanded activation, the project conv brings cExp → c back, and the
per-channel layer scale closes the block. gf is the form of the GELU, the tanh
approximation or the exact x · Φ(x) (GeluForm); the block, the net and every theorem
about them carry it and hold for both.
The layerNormForward used here normalizes over the whole flattened
c·h·w vector with scalar γ_n, β_n. The per-position channel
LayerNorm ConvNeXt-T ships with is chanLNTensor3, and this body at
it is cnxBodyWith (ConvNeXtFullT.lean). Because gelu is smooth everywhere, LN is smooth
given ε>0, and conv/layerScale are linear, the whole body is
differentiable everywhere — no ReLU-style kink hypotheses needed.
Equations
- One or more equations did not get rendered due to their size.
Instances For
ConvNeXt block body VJP (global) — built by chaining the
everywhere-differentiable piece VJPs through vjpComp. Needs only
0 < εn (the LayerNorm positivity); no kink hypotheses since gelu is
smooth and the rest are linear. Because the body is differentiable
everywhere, the VJP is global (HasVJP), not pointwise.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Full ConvNeXt block = residual (block body). ConvNeXt uses an
identity skip (no projection, no post-add activation), so this is the
plain residual of the block body.
Equations
- Proofs.convNextBlock gf Wdw bdw εn γn βn Wex bex Wpr bpr γls = Proofs.residual (Proofs.convNextBlockBody gf Wdw bdw εn γn βn Wex bex Wpr bpr γls)
Instances For
ConvNeXt block VJP (global) — residualHasVJP on top of the
block-body VJP. Needs only 0 < εn. Global since the body is
everywhere-differentiable and the skip is the identity.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Forward ConvNeXt (two blocks, scalar-γ, β LayerNorm):
stem-patchify(1×1 conv) → stem-LN → block₁ → block₂ → globalAvgPool → head-LN → dense
Generic channel/spatial dims; ic→c patchify, two identity-skip
ConvNeXt blocks at c, GAP to Vec c, a final LN over the pooled
Vec c, and a Mat c nClasses linear head. The stem-LN and head-LN
are the scalar-γ, β layerNormForward, as in convNextBlockBody.
Equations
- One or more equations did not get rendered due to their size.
Instances For
End-to-end ConvNeXt VJP (global). Everything is smooth, so the
only hypotheses are the four LayerNorm positivity conditions
(0 < εst, εn₁, εn₂, εhd) — no ReLU/maxpool kink conditions, unlike
cnnHasVJPAt. Chained entirely through the global vjpComp, so the
VJP holds at every input, not just a fixed point — putting ConvNeXt
alongside vitFullHasVJP (the weight-tied ViT) as an unconditional
whole-network VJP.
Equations
- One or more equations did not get rendered due to their size.
Instances For
End-to-end ConvNeXt VJP at a point — the global witness restricted
to a point. Kept for downstream _at consumers and the comparator.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Public correctness theorem for convnextHasVJP (global) — the
end-to-end ConvNeXt's backward equals the pdiv-contracted Jacobian
(Jacobian-transpose applied to the cotangent), at every input x.
The unconditional ConvNeXt analogue of vitFullHasVJP_correct.
Public correctness theorem for convnextHasVJPAt — exposes the
witness's .correct field: the end-to-end ConvNeXt's backward equals
the pdiv-contracted Jacobian (Jacobian-transpose applied to the
cotangent). Analogue of cnnHasVJPAt_correct.