ConvNeXt #
A representative ConvNeXt block and a small end-to-end ConvNeXt VJP in
flattened Vec space — the ConvNeXt analogue of cnn_has_vjp_at.
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 cnn_has_vjp_at — 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 derived frompdiv_mul/pdiv_const/pdiv_id.- the ConvNeXt block body
layerScale ∘ project ∘ gelu ∘ expand ∘ LayerNorm ∘ depthwise, everywhere-differentiable, with a pointwise VJP built by chaining the piece VJPs throughvjp_comp_at; and the full blockresidual (block body)(identity skip, no post-add act). convnext_has_vjp_at/convnext_has_vjp_at_correct— a fixed-depth (two-block) end-to-end network: stem-patchify → stem-LN → block₁ → block₂ → global-avg-pool → head-LN → dense.
LayerNorm representation caveat #
The layerNormForward reused here is the proof's Vec→Vec LayerNorm with
scalar γ, β that normalizes over the whole flattened vector,
whereas true ConvNeXt LayerNorm is per-spatial-position over the channel
axis (LayerNorm over NCHW's C). This is the same representation
simplification the audit flagged for the LN family; a faithful
channel-LN-over-NCHW lift is a follow-up. Every other piece (depthwise
7×7, 1×1 convs, GELU, layer scale, GAP, dense) is exact.
Layer scale — per-channel learnable elementwise multiply by γ.
layerScale γ x i = γ i * x i. A diagonal linear map.
Equations
- Proofs.layerScale γ x i = γ i * x i
Instances For
layerScale γ is differentiable everywhere (diagonal linear).
Jacobian of layerScale — ∂(γ_j x_j)/∂x_i = γ_i δ_{ij}.
Layer scale VJP: back(x, dy)_i = γ i * dy i.
Equations
- Proofs.layerScale_has_vjp γ = { backward := fun (_x dy : Proofs.Vec n) (i : Fin n) => γ i * dy i, correct := ⋯ }
Instances For
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.
LN representation caveat. The layerNormForward used here is the
proof's Vec→Vec LayerNorm with scalar γ_n, β_n that normalizes
over the whole flattened c·h·w vector, whereas true ConvNeXt LN is
per-spatial-position over the channel axis (LayerNorm over NCHW's C).
This is the same representation simplification the audit flagged for
the LN family; a faithful channel-LN-over-NCHW is a follow-up. Every
other piece is exact. 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
The block body is differentiable everywhere (composition of everywhere-differentiable maps).
ConvNeXt block body VJP (global) — built by chaining the
everywhere-differentiable piece VJPs through vjp_comp. 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
ConvNeXt block body VJP at a point — the global witness restricted
to a point. Kept for downstream _at consumers.
Equations
- Proofs.convNextBlockBody_has_vjp_at Wdw bdw εn hεn γn βn Wex bex Wpr bpr γls v = (Proofs.convNextBlockBody_has_vjp Wdw bdw εn hεn γn βn Wex bex Wpr bpr γls).toHasVJPAt v
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 Wdw bdw εn γn βn Wex bex Wpr bpr γls = Proofs.residual (Proofs.convNextBlockBody Wdw bdw εn γn βn Wex bex Wpr bpr γls)
Instances For
The full ConvNeXt block is differentiable everywhere (residual of an everywhere-differentiable body).
ConvNeXt block VJP (global) — residual_has_vjp 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
ConvNeXt block VJP at a point — the global witness restricted to a
point. Kept for downstream _at consumers.
Equations
- Proofs.convNextBlock_has_vjp_at Wdw bdw εn hεn γn βn Wex bex Wpr bpr γls v = (Proofs.convNextBlock_has_vjp Wdw bdw εn hεn γn βn Wex bex Wpr bpr γls).toHasVJPAt v
Instances For
Forward ConvNeXt (representative, fixed block count = 2):
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. Same LN representation
caveat as convNextBlockBody applies to the stem-LN and head-LN.
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
cnn_has_vjp_at. Chained entirely through the global vjp_comp, so the
VJP holds at every input, not just a fixed point — putting ConvNeXt
alongside vit_full_has_vjp 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 convnext_has_vjp (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 vit_full_has_vjp_correct.
Public correctness theorem for convnext_has_vjp_at — 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 cnn_has_vjp_at_correct.