The ConvNeXt-T backward chains — the ℝ maps the ConvNeXt ties are about #
The hand-composed reverse of the committed ConvNeXt-T forward, as plain defs on the cotangent:
the block body backward cnxBlockBodyBack (the reverse of
layerScale ∘ project ∘ GELU ∘ expand ∘ LN ∘ depthwise), the stage-boundary downsample backward
cnxDownBack (the reverse of flatConvStride2 ∘ LN), the per-example whole-net chain
convnextInputGrad (the reverse of convNextForwardTCh, the [3,3,9,3] ladder at
96→192→384→768 and 56→28→14→7, nC classes) and its batched form convnextInputGradB (the same
chain at each of B examples, every saved-activation slot a StableHLO.batchMapAux B lift). The
LayerNorm backwards, the GELU and layer-scale diagonal scales, the four stage backwards and the
three downsample backwards enter as supplied maps; only the depthwise, the two 1×1 convs, the
4×4/s4 patchify stem, the GAP and the dense are spelled, so that the certified ties
(ConvNeXtBackCertifiedTie, ConvNeXtWholeBackCertifiedTie, ConvNeXtWholeBackCertifiedTieB)
are statements about a NAMED chain of the forward's shape. The channel-LN backward those ties fill
the lnB slots with is chanLNTensor3Back (ChannelLNBack.lean).
Note: padding is symmetric at the downsamples (flatConvStride2Back) and the stem is the
stride-4 patchify (flatConvStride4Back: two scatters, then the reversed-kernel conv); the 7×7
depthwise is odd, and EvenKernelConvBack.lean is where the 4×4 and 2×2 kernels' even-kernel ties
live (the ties pass the kernels through padOdd).
No number is stated about any of these chains.
The ConvNeXt block body input-gradient VJP — the reverse of
convNextBlockBody = layerScale ∘ project ∘ GELU ∘ expand ∘ LN ∘ depthwise:
depthwiseFlatBack Wdw ∘ lnB ∘ convFlatBack Wex ∘ geluB ∘ convFlatBack Wpr ∘ lsB
lsB = diagBack γls (the per-channel layer-scale backward); convFlatBack Wpr the project back;
geluB = diagBack (gelu'(saved)); convFlatBack Wex the expand back; lnB the LayerNorm back
(the ties fill it with chanLNTensor3Back); depthwiseFlatBack Wdw the depthwise input-VJP.
The full block is residual (body), so the block backward is residual (cnxBlockBodyBack …);
ConvNeXtBackCertifiedTie.cnxBodyWithChanLNBack_eq_vjp fills the LN slot with the concrete
channel-LN chain, pins the other slots to their certified backwards, and shows the chain IS
the certified body VJP.
Equations
- Proofs.cnxBlockBodyBack Wdw Wex Wpr lnB lsB geluB = Proofs.depthwiseFlatBack Wdw ∘ lnB ∘ Proofs.convFlatBack Wex ∘ geluB ∘ Proofs.convFlatBack Wpr ∘ lsB
Instances For
The ConvNeXt downsample input-gradient VJP — the reverse of cnxDownChW = flatConvStride2 W ∘ LN:
lnB ∘ flatConvStride2Back W (run the strided-conv backward, then the LayerNorm back).
Equations
- Proofs.cnxDownBack W lnB = lnB ∘ Proofs.flatConvStride2Back W
Instances For
The whole ConvNeXt-T input-gradient VJP — the exact reverse of
convNextForwardTCh: dense ∘ LN ∘ GAP ∘ stage₄ ∘ down₃ ∘ stage₃ ∘ down₂ ∘ stage₂ ∘ down₁ ∘ stage₁ ∘ LN ∘ stem reversed. The stem (flatConvStride4Back sW ∘ lnBstem, the 4×4/s4 patchify
backward), GAP and dense endpoints are concrete; the head-LN, stem-LN, 4 stage backwards and 3
downsample backwards are supplied (the stages fold residual (cnxBlockBodyBack …), the
downsamples are cnxDownBack). The [3,3,9,3] structure is in the stage maps' depths; the
channel/spatial schedule (96→192→384→768, 56→28→14→7) in their dims; the class count nC is
a binder (10 at Imagenette, 1000 at the convnextin_* artifacts). The tie fills lnBhead
with the head LayerNorm's backward at the GAP output (rowLNVecFlatBack 1 768).
Equations
- One or more equations did not get rendered due to their size.
Instances For
The batched whole-net ConvNeXt-T input gradient — convnextInputGrad at each of B
examples, stage by stage, as the batched render computes it. The three input-independent
leaves (the stem's reversed-kernel conv, GAP, the dense head) are StableHLO.batchMap B of
the per-example leaf. Every saved-activation slot is StableHLO.batchMapAux B of a
per-example FAMILY saved ↦ (cotangent ↦ gradient) at a batched saved activation: lnBstem
at a0 (the stem conv's output), the four stage backwards s1B … s4B at their stage inputs
a1 a3 a5 a7, the three downsample backwards d1B … d3B at a2 a4 a6, and lnBhead at a9
(the GAP output). Where the per-example chain fills a slot at ONE saved value, this one fills
it at B of them, which is the only way a batch enters — LayerNorm is per-example and no
ConvNeXt op couples examples, as in the batched step tie ConvNeXtStepTieGB. B and nC are
variables: this chain carries no batch numeral and no class count.
convnextInputGradB_eq_batchMap_convNextForwardTCh_vjp (ConvNeXtWholeBackCertifiedTieB.lean)
says it is the certified gradient of batchMap B convNextForwardTCh.
Equations
- One or more equations did not get rendered due to their size.