Strided convolution (stride-2 SAME) — Chapter 5 Milestone B, the hard new op #
Real ResNet-34 downsamples with stride-2 convolutions, the one genuinely-new
operator the Chapter-5 handoff (planning/archive/verified_r34.md §3.6) flags as gating
the jump from the ch6-A ResNet-style net to a true 34-layer ResNet.
The key identity that makes this tractable. A stride-2 SAME convolution is exactly a stride-1 SAME convolution followed by spatial decimation (keep every other position):
conv_stride2 W b X = decimate2 (conv2d W b X), X : Tensor3 ic (2h) (2w)
because both read x_pad[c, 2·hi+kh−pH, 2·wi+kw−pW] — the stride-1 conv computes
that at every output position, and decimation throws away the odd ones. So we do
not re-derive the ~800-line conv input-VJP / weight-grad with stride arithmetic;
we reuse conv2d_has_vjp3 and conv2d_weight_grad_has_vjp verbatim and only add a
small linear decimation map decimateFlat (a reindex, hence a CLM) with its
VJP (the backward is the "zero-upsampling" / lhs_dilation scatter). The strided
conv's input- and weight-VJPs then fall out of vjp_comp.
Everything closes under [propext, Classical.choice, Quot.sound].
The decimation index map: a small output flat index k ↔ (co, ho, wo) maps to
the even input position (co, 2·ho, 2·wo) in the (2h)×(2w) grid. A pure
reindex Fin (oc·h·w) → Fin (oc·2h·2w); decimateFlat reads through it.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Flat spatial decimation Vec (oc·2h·2w) → Vec (oc·h·w): keep the even
spatial positions. A coordinate reindex fun y k => y (decimateIdx k) — i.e.
reindexCLM decimateIdx — so it is continuous-linear (hence differentiable),
and decimate2 (conv2d …) = conv_stride2 ….
Equations
- Proofs.decimateFlat oc h w y k = y (Proofs.decimateIdx oc h w k)
Instances For
Decimation VJP. decimateFlat is a reindex, so its Jacobian is the sparse
δ(idx = decimateIdx j) (pdiv_reindex); the backward scatters dy back to
the even positions (zero elsewhere) — the "zero-upsampling" that StableHLO
renders as lhs_dilation = [2,2]. Stated in the universal ∑ pdiv · dy form.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Stride-2 SAME convolution, flattened: Vec (ic·2h·2w) → Vec (oc·h·w).
Defined as decimateFlat ∘ flatConv (the stride-1 SAME conv on the 2h×2w
grid, then keep even positions) — provably the genuine stride-2 conv.
Equations
- Proofs.flatConvStride2 W b = Proofs.decimateFlat oc h w ∘ Proofs.flatConv W b
Instances For
Stride-2 conv input-VJP — the centerpiece. By the chain rule
(vjp_comp) on decimateFlat ∘ flatConv, reusing the proven stride-1 conv
input-VJP (conv2d_has_vjp3 via the flatten bridge) and the decimation VJP.
The backward is flatConv.back (decimate.back dy) — i.e. zero-upsample the
cotangent, then run the reversed-kernel conv (StableHLO: lhs_dilation=[2,2]
on the transpose-reverse convolution).
Equations
- One or more equations did not get rendered due to their size.
Instances For
Stride-2 conv input-VJP correctness (the ℝ-carrying audit headline): the
backward equals the pdiv-contracted Jacobian of flatConvStride2.
Conv2d (as a function of its flattened kernel) is differentiable — it is
affine in the weights (b o + ∑ v(idx)·pad-eval x, the pad-eval being a
weight-independent constant). Needed as the vjp_comp hypothesis for the
strided weight-grad.
Stride-2 conv weight-VJP. The same composition decimate ∘ conv viewed
as a function of the kernel (input x fixed): the weight-grad is
conv_weight_grad run on the zero-upsampled cotangent. By vjp_comp,
reusing the proven stride-1 conv2d_weight_grad_has_vjp + decimateFlat_has_vjp.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Stride-2 conv weight-VJP correctness (ℝ-headline): backward = the
pdiv-Jacobian of the strided conv in its kernel.
conv2d (as a function of its bias) is differentiable — affine in b (bias broadcast
plus a b-independent W,x term). The vjp_comp hypothesis for the strided-conv bias-grad.
Stride-2 conv bias-VJP. fun b => flatConvStride2 W b x = decimate ∘ (conv2d-in-b); by
vjp_comp of the proven stride-1 conv2d_bias_grad_has_vjp with decimateFlat_has_vjp. The
bias peer of flatConvStride2_weight_grad_has_vjp.
Equations
- One or more equations did not get rendered due to their size.
Instances For
The odd decimation index map: like decimateIdx but keeping the odd
positions (co, 2·ho+1, 2·wo+1). Composed under an even decimation it reads
4·ho+1 — which for the SAME conv at pad (k-1)/2 = 1 (k = 4) makes the
window exactly the left-aligned x[4i .. 4i+3]: the real (paper/render)
pad-0 stride-4 patchify, never touching the boundary padding.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Flat odd spatial decimation Vec (oc·2h·2w) → Vec (oc·h·w): keep the odd
spatial positions. A coordinate reindex, exactly as decimateFlat.
Equations
- Proofs.decimateOddFlat oc h w y k = y (Proofs.decimateOddIdx oc h w k)
Instances For
Odd-decimation VJP — the same sparse-δ reindex Jacobian as
decimateFlat_has_vjp, at the odd positions.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Stride-4 patchify convolution, flattened: Vec (ic·4h·4w) → Vec (oc·h·w).
decimateFlat ∘ decimateOddFlat ∘ (stride-1 SAME conv) — reads the SAME conv
(pad (k-1)/2) at positions 4i+1, which for the 4×4 stem is the
left-aligned window x[4i .. 4i+3]: the paper's pad-0 Conv2d(k=4, s=4)
and the committed render's patchify, in-bounds at every tap.
Equations
- Proofs.flatConvStride4 W b = Proofs.decimateFlat oc h w ∘ Proofs.decimateOddFlat oc (2 * h) (2 * w) ∘ Proofs.flatConv W b
Instances For
Stride-4 conv input-VJP — two vjp_comp steps over the proven stride-1
conv input-VJP and the two decimation VJPs (backward = zero-upsample twice,
then the reversed-kernel conv).
Equations
- One or more equations did not get rendered due to their size.
Instances For
Stride-4 conv weight-VJP. The kernel-side peer of flatConvStride4_has_vjp, and the
stride-4 analogue of flatConvStride2_weight_grad_has_vjp: the same
decimateFlat ∘ decimateOddFlat ∘ conv composition viewed as a function of the kernel
(input x fixed), so the weight-grad is conv2d_weight_grad run on the twice-zero-upsampled
cotangent. Two vjp_comp steps over the proven stride-1 weight-VJP and the two decimation
VJPs — no new mathematics, only the composition the stride-2 sibling already does once.
This is the cert that ConvNeXt's 4×4/s4 patchify stem (psW) was missing: its forward
(flatConvStride4) and input-VJP (flatConvStride4_has_vjp) were already proven, so the stem's
weight gradient was the last hand-written emitter in that render.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Stride-4 conv weight-VJP correctness (ℝ-headline): backward = the pdiv-Jacobian of the
stride-4 conv in its kernel. The peer of flatConvStride2_weight_grad_has_vjp_correct.
Why a second stride-2 convolution exists, and why it is not a fix to the first.
The repo carries two genuinely different stride-2 conventions, and both are correct references for the nets that use them:
flatConvStride2(above) pads symmetrically,(k-1)/2on each side. This is He et al. / torchvision —nn.Conv2d(padding=k//2)— and it is what ResNet-34, ResNet-50 and ConvNeXt's references do (jax/Jax/Codegen.lean:462, symmetric ON PURPOSE since 2026-08-04).flatConvStride2Xla(here) pads the way XLA'SAME'does: at an even input the total padding isk-2, split asymmetrically as((k-2)/2, k/2)—(0,1)atk=3,(1,2)atk=5,(2,3)atk=7. This is what the TF-origin ports do — MobileNetV2, MobileNetV4, EfficientNet — wherepadding='SAME'is the reference and must not be "fixed".
⭐ The identity that makes this nearly free. A stride-2 XLA-SAME conv is the same
symmetric stride-1 conv the even-decimation op already uses, decimated at the odd offsets:
convXlaSame_s2 W b X = decimateOddFlat (flatConv W b X), X : Tensor3 ic (2h) (2w)
because output ho then reads x[2·ho + 1 + kh − (k−1)/2] = x[2·ho + kh − ((k−2)/2)], and
(k−2)/2 is exactly XLA's pad_low at an even input. So the whole asymmetry is a phase shift
in the decimation, not new padding arithmetic: flatConv is reused verbatim, and so is every
one of its VJPs. decimateOddFlat and decimateOddFlat_has_vjp already exist above (they were
added for ConvNeXt's 4×4/s4 patchify stem), so this section adds no new proof obligation —
only compositions of results already closed under the three standard axioms.
⚠ This holds at EVEN inputs only, which is the only case any net in this repo hits (224, 112,
56, 28, 14 — every strided site in mnv2/mnv4/enet). At an odd input XLA SAME pads
((k-1)/2, (k-1)/2) — symmetric — so flatConvStride2 is already the right op there and this one
would be wrong. The type enforces it: the input index is ic*(2*h)*(2*w), structurally even.
Verified against jax.lax.conv_general_dilated(…, 'SAME') over
H ∈ {224,112,56,28,14,32,16,9,7,15,33} × k ∈ {3,5,7} — 33 configs, all agreeing with this rule
(planning/archive/mnv4_verified.md §3e).
Stride-2 XLA-SAME convolution, flattened: Vec (ic·2h·2w) → Vec (oc·h·w).
decimateOddFlat ∘ flatConv — the stride-1 symmetric-SAME conv on the 2h×2w grid, then keep
the odd positions. The asymmetric-pad peer of flatConvStride2.
Equations
- Proofs.flatConvStride2Xla W b = Proofs.decimateOddFlat oc h w ∘ Proofs.flatConv W b
Instances For
Stride-2 XLA-SAME input-VJP. vjp_comp on decimateOddFlat ∘ flatConv, reusing the
proven stride-1 conv input-VJP and the odd-decimation VJP. The backward zero-upsamples the
cotangent onto the odd positions and then runs the reversed-kernel conv — i.e. StableHLO's
lhs_dilation = [2,2] with the transposed padding shifted by one, which is exactly the
asymmetry the forward introduced. ⚠ A symmetric backward against this forward is a silent
wrong-gradient (planning/archive/mnv4_verified.md §3b), and it is this composition that rules it out:
the offset lives in one place and both directions read it.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Stride-2 XLA-SAME input-VJP correctness (the ℝ-carrying audit headline): the backward
equals the pdiv-contracted Jacobian. Peer of flatConvStride2_has_vjp_correct.
Stride-2 XLA-SAME weight-VJP. The same composition viewed as a function of the kernel
(input x fixed): conv2d_weight_grad run on the odd-zero-upsampled cotangent.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Stride-2 XLA-SAME weight-VJP correctness (ℝ-headline).
Stride-2 XLA-SAME bias-VJP. fun b => flatConvStride2Xla W b x = the odd decimation of
the stride-1 conv-in-b; by vjp_comp of conv2d_bias_grad_has_vjp with the odd-decimation
VJP. The bias peer of flatConvStride2_bias_grad_has_vjp.
Equations
- One or more equations did not get rendered due to their size.