maxPool3s2 — the 3×3 stride-2 max pool of He et al.'s ResNet stem #
maxPool2 (CNN.lean) is 2×2, stride 2, non-overlapping. Every ResNet in He et al.
(18/34/50/101/152) specifies a 3×3 stride-2 pool after the stem conv; this file is that op.
Which 3×3 pool: the paper's symmetric padding #
This file implements He et al. / torchvision: nn.MaxPool2d(3, stride=2, padding=1) —
symmetric padding, so window i covers input [2i−1, 2i+1].
XLA's reduce_window(…, 'SAME') would be a different function: on a 112→56 axis it gives
pad_total = max((56−1)·2 + 3 − 112, 0) = 1, split pad_low = 0, pad_high = 1 — window
i = [2i, 2i+2], padded at the end. Measured on device at n = 12: SAME windows peak at
[2,4,6,8,10,11], symmetric at [1,3,5,7,9,11]; the two grids are offset by one input position.
The JAX reference's max_pool2d uses the same symmetric padding (p, p) with p = (size−1)//2.
At 2×2 symmetric (k−1)//2 padding is p = 0, bit-identical to SAME, so the choice only
matters for the 3×3 pools (the ResNet-34 and ResNet-50 ImageNet stems).
The padding needs no extended-reals type #
reduce_window pads with -∞, which Tensor3 _ _ _ = … → ℝ cannot hold. It does not need to:
for max, clamping the index is equivalent to -∞ padding. The only out-of-range read is
2i−1 at i = 0, and Nat's truncated subtraction clamps it to 0 — a cell the window already
contains at offset a = 1. So max over the clamped triple equals max over the unpadded pair,
which is exactly what -∞ padding computes. win3RowInv_first_dup is that statement, and
win3ColInv_first_dup its column half.
The symmetric form needs no min: the upper end 2(h−1)+2−1 = 2h−1 is in range by
construction, so truncated subtraction is the whole story.
The pool is a windowMax instance #
maxPool3s2 is windowMax win3RowInv win3ColInv (WindowMax.lean), and every analytic fact here
is that file's lemma at these window maps. What is particular to this pool is the window
geometry: the clamped duplicate (win3RowInv_first_dup), which is why smoothness is stated over
positions, and the overlap (win3Row_mem_le_two): an input lies in up to four windows, so where
maxPool2's backward is a lookup, this one accumulates. The witness is maxPool3s2HasVJPAt3, with
its flat form maxPool3s2FlatHasVJPAt; the ResNet-34/50 stems, their seals and the float stem
bridge build on it.
Input row of offset a ∈ Fin 3 inside output window hi_out: 2·hi + a − 1, in Nat.
The truncated subtraction is the low pad (see the header) — it is the only clamp needed,
because the high end 2(h−1)+2−1 = 2h−1 is in range by construction.
Instances For
Column peer of win3RowInv.
Instances For
The padding statement. In the FIRST window offset a = 0 duplicates a = 1 rather
than reading out of range — exactly what a -∞ pad contributes to a max.
3×3 stride-2 symmetrically-padded max pool, [c, 2h, 2w] → [c, h, w]: the max over the
window [2i−1, 2i+1] × [2j−1, 2j+1], clamped at the near edge (= -∞ padded, header).
Equations
Instances For
Smoothness: every 3×3 window attains its max at exactly one input position
(WindowSmooth). Stated over positions, so the clamped duplicate in the first window is not a
tie; a window whose max sits at two positions (an all-zero post-ReLU window) does not
qualify.
Instances For
Positional injectivity ⇒ MaxPool3s2Smooth, the discharge used by
BatchSeal.ctConv_pool_smooth for the ResNet-34 and ResNet-50 full-width seals: one
injectivity argument in place of a per-window case split.
Smooth or dead (WindowSmoothOrDead): every 3×3 window has its maximum at one position,
or every cell ≤ 0. The ResNet stems state their pool condition in this form
(StemPoolSmoothAt), and maxPool3s2Flat_relu_eventuallyEq is the lemma that uses it.
Equations
Instances For
A smooth pool input is smooth-or-dead.
Smooth, dead, or tied only between twins (WindowSmoothUpTo): the ResNet stem's twins are
cells reading identical input patches (StemPoolTwinAt).
Equations
Instances For
The overlap fact, stated rather than assumed: an input row lies in at most TWO windows —
p/2 and (p+1)/2. With symmetric padding the shared cell is at ODD p (window (p−1)/2
takes it at offset 2, window (p+1)/2 at offset 0); even p lies in exactly one. So an input
feeds at most 4 outputs and the backward accumulates at most 4 terms — the count maxPool2
does not have.
A (not necessarily unique) argmax of the 3×3 window at output (co, ho, wo).
Equations
- Proofs.maxPool3s2Argmax x co ho wo = Proofs.windowArgmax Proofs.win3RowInv Proofs.win3ColInv x co ho wo
Instances For
If (a, b) dominates every window cell, the pooled value is the value there.
For each output flat index, the flat index of its argmax's input position. Not injective: two overlapping windows may select the same input.
Equations
Instances For
The VJP witness. The backward accumulates dy over every output whose window selects
this input, at most 4 of them (win3Row_mem_le_two squared). The backward is spelled out
rather than taken from windowMaxHasVJPAt3, so that unfolding this name exposes the sum the
backward ties match.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Flattened 3×3/s2 pool, the Vec-level form the codegen denotes. Spelled through
maxPool3s2, which is windowMaxFlat at the 3×3/s2 maps by unfolding.
Equations
- Proofs.maxPool3s2Flat c h w v = (Proofs.maxPool3s2 (Proofs.Tensor3.unflatten v)).flatten
Instances For
The VJP of maxPool3s2Flat at a flattened input satisfying MaxPool3s2Smooth:
maxPool3s2HasVJPAt3 moved to Vec form.
Equations
- Proofs.maxPool3s2FlatHasVJPAt x h_smooth = (Proofs.maxPool3s2HasVJPAt3 x h_smooth).toHasVJPAt
Instances For
maxPool3s2Flat is continuous (a sup' of coordinates).
The 3×3/s2 pool shifts with a uniform offset, at every point of the ray, which is what lets the carrier cross the only real kink in the net.