CNN VJP Proofs #
VJP correctness for the convolutional and pooling layers used in
historical/mlir_poc/hand_cnn_train_step.mlir. The architecture there is:
x(1,28,28) → Conv(1→32) → ReLU → Conv(32→32) → ReLU → MaxPool
→ Flatten → Dense(6272→512) → ReLU → Dense(512→512)
→ ReLU → Dense(512→10) → logits
The dense and ReLU layers are inherited from MLP.lean. This file
adds the new operations: conv2d, max-pool, and flatten.
The big idea: conv backward IS conv #
The most pedagogically valuable result is that the VJP of a convolution is itself expressible as convolutions — with appropriate kernel reversal and axis transposition. This is why conv layers train efficiently: there's no special backward operator. The same primitive runs in both directions.
Two specific tricks appear in the MLIR:
- Input-gradient via reversed kernel:
dx = conv(dy, reverse(Wᵀ)) - Weight-gradient via the transpose trick:
dW = conv(xᵀ, dyᵀ), where the spatial dims of the gradient become the "kernel".
We bundle the VJP formulas as HasVJP3 / HasVJP defs whose
.correct fields are proved (the proofs are standard matrix calculus
on cross-correlations), and the commentary explains why each formula
has the form it does.
Kernel4 oc ic kH kW and Vec (oc * ic * kH * kW) are in bijection
by row-major flattening — mirrors Mat.flatten / Tensor3.flatten.
We need this so that the weight-gradient VJP can be stated as a plain
HasVJP (Vec → Vec) on the flattened kernel, reusing the existing
framework instead of introducing a parallel 4D machinery.
Nat multiplication associates left, so `oc * ic * kH * kW` parses as
`((oc * ic) * kH) * kW` — three nested `finProdFinEquiv` calls.
Row-major unflatten: inverse of flatten.
Equations
- Proofs.Kernel4.unflatten v o c kh kw = v (finProdFinEquiv (finProdFinEquiv (finProdFinEquiv (o, c), kh), kw))
Instances For
Conv2d forward (SAME padding, stride 1).
y[o, h, w] = (Σ_{c, kh, kw} x[c, h+kh−p, w+kw−p] · W[o, c, kh, kw]) + b[o]
where p = (kH−1)/2 is the padding offset and out-of-bounds reads
return 0 (zero padding). The output spatial size equals the input.
Note: this is technically cross-correlation, not convolution in the strict signal-processing sense. ML literature uses "convolution" loosely; the difference (kernel flipping) only matters when comparing against classical signal-processing references.
MLIR (hand_cnn_train_step.mlir):
%cv0 = "stablehlo.convolution"(%x, %W0) {
padding = dense<[[1, 1], [1, 1]]>, ...
}
%h0pre = stablehlo.add %cv0, broadcast(%b0)
Equations
- One or more equations did not get rendered due to their size.
Instances For
Conv2d is differentiable everywhere. Each output coordinate
conv2d W b x o hi wi is the affine map
b o + ∑_{c,kh,kw} W o c kh kw · (pad-eval x): a constant bias plus a
finite ℝ-linear combination of input coordinates (the dependent
if-pad-eval being either a projection or the constant 0).
differentiable_pi reduces to per-coordinate differentiability;
DifferentiableAt.fun_sum lifts the triple sum, and each pad-eval
summand is a projection (pad true) or constant (pad false).
Differentiability of an if hpad : P then v(σ hpad) else 0 term.
The dependent-if would otherwise stymie fun_prop. By proof
irrelevance, the chosen branch is a CLM (eval-at-σ if P, the
constant 0 otherwise) — both differentiable in v.
Pdiv of a per-output dependent if-eval-or-zero family.
Given a per-output dependent if fun v k' ↦ if h : P k' then v (σ k' h) else 0,
its pdiv at (idx_in, idx_out) is the indicator
if P holds at idx_out ∧ σ matches idx_in then 1 else 0. The proof uses
fderiv_apply to extract the idx_out-th component, then by_cases on
P idx_out to discharge the dependent-if.
Pdiv of c_const * pad-eval family. Combines pdiv_mul,
pdiv_const, and pdiv_pi_pad_eval for the conv2d per-summand
pattern: a k'-varying constant times the dependent if-eval-or-zero.
Closed-form input gradient for conv2d — direct formula, written as
a sum over output positions (co, ho, wo) with reconstructed kernel
offsets kh_nat = hi + pH − ho, kw_nat = wi + pW − wo. The body is
nonzero only when the reconstructed (kh_nat, kw_nat) lies in
[0, kH) × [0, kW) — i.e., when the input position (hi, wi) is
actually reachable from output (ho, wo) via some valid kernel offset.
Equivalent (under the (ho, wo) ↔ (kh, kw) partial bijection
ho = hi+pH-kh) to the MLIR-aligned "reversed-kernel" formula
dx[c, h, w] = Σ_{o, kh, kw} W[o, c, kH−1−kh, kW−1−kw] · dy[o, h+kh−p, w+kw−p].
Equations
- One or more equations did not get rendered due to their size.
Instances For
Conv2d input-VJP — proved from foundation rules.
The function v ↦ flatten (conv2d W b (unflatten v)) is affine in
v: a constant b o(idx_out) plus a triple sum over (c, kh, kw)
of W o(idx_out) c kh kw * (if pad-cond then v(reindex) else 0).
Each summand factors as (constant W) * (if-pad-conditional in v),
so pdiv_add + pdiv_const + pdiv_finset_sum (×3) + pdiv_mul +
a by_cases on the pad condition (CLM-projection on the pad-true
branch, constant zero on the pad-false branch) collapse the
per-(idx_in, idx_out) pdiv. Reindex Fin (oc*h*w) ↔ Fin oc × Fin h × Fin w
on the sum-over-idx_out, then a triple Finset.sum_eq_single over
(c, kh, kw) (matching idx_in's decoded (ci, hi, wi)) gives the
closed-form input gradient conv2d_input_grad_formula.
The backward function (accessed as (conv2d_has_vjp3 W b).backward,
or via the conv2d_input_grad abbrev below) implements
conv2d_input_grad_formula W dy ci hi wi — a direct sum over
(co, kh, kw) of W co ci kh kw * dy co ho_nat wo_nat for valid
(ho_nat, wo_nat). Equivalent (under kh ↔ kH−1−kh) to the
MLIR-aligned reversed-kernel formula
dx[c, h, w] = Σ_{o, kh, kw} W[o, c, kH−1−kh, kW−1−kw] · dy[o, h+kh−p, w+kw−p].
MLIR emits the reversed-kernel form directly: %W1_t = stablehlo.transpose %W1, dims = [1, 0, 2, 3] -- swap oc↔ic %W1_rev = stablehlo.reverse %W1_t, dims = [2, 3] -- flip spatial %d_h0 = "stablehlo.convolution"(%d_h1pre, %W1_rev) ...
Equations
- Proofs.conv2d_has_vjp3 W b = { backward := fun (_x : Proofs.Tensor3 ic h w) (dy : Proofs.Tensor3 oc h w) => Proofs.conv2d_input_grad_formula W dy, correct := ⋯ }
Instances For
Named accessor for the conv2d input backward — aligns with MLIR
codegen (stablehlo.convolution in the backward pass).
Equations
- Proofs.conv2d_input_grad W b x dy = (Proofs.conv2d_has_vjp3 W b).backward x dy
Instances For
Uniform VJP-correctness wrapper for conv2d — a citable _correct
matching the convention of every other layer (just unfolds the
HasVJP3.correct field of conv2d_has_vjp3).
Flat conv — conv2d bridged into flattened Vec → Vec space:
flatConv W b = flatten ∘ conv2d W b ∘ unflatten. Spatial dims are
preserved (ic h w → oc h w), so this is Vec (ic*h*w) → Vec (oc*h*w).
This is the form the ResNet/CNN VJP composition actually uses
(everything lives in flat Vec space).
Equations
- Proofs.flatConv W b v = (Proofs.conv2d W b (Proofs.Tensor3.unflatten v)).flatten
Instances For
flatConv is differentiable everywhere. Composition of the three
differentiable maps unflatten, conv2d, flatten. This is the
differentiability witness vjp_comp_at needs to chain conv into the
block.
conv → bn → relu block VJP at a smooth point.
The workhorse for composing a ResNet VJP. In flattened Vec space,
the block is relu ∘ bnForward ∘ flatConv : Vec (ic*h*w) → Vec (oc*h*w)
(BatchNorm runs over the oc*h*w flattened activations with scalar
ε, γ, β). We build HasVJPAt at a point v via two vjp_comp_at:
- inner =
bnForward ∘ flatConv— both differentiable everywhere (flatConv_differentiable,bnForward_differentiable), so their bundled VJPs lift through.toHasVJPAt. The conv witness is theHasVJP3-bridgedhasVJP3_to_hasVJP (conv2d_has_vjp3 W b). - outer =
relu— needs the smoothness hypothesish_smooth(no post-BN activation hits the ReLU kink) for bothrelu_differentiableAt_of_smoothandrelu_has_vjp_at.
Mirrors mlp_has_vjp_at (dense→relu→dense), with flatConv/bn in
place of the dense layers.
Equations
- One or more equations did not get rendered due to their size.
Instances For
conv → bn block VJP (no ReLU), everywhere. Just flatConv then
bnForward, both differentiable everywhere, so this is a global
HasVJP (no smoothness needed). This is the building block for the
second conv→bn of a residual body and for the 1×1 projection skip.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Basic-block body VJP at a smooth point.
F := convBn₂ ∘ convBnRelu₁ = bn₂ ∘ conv₂ ∘ relu ∘ bn₁ ∘ conv₁,
the body of a post-activation ResNet basic block (the outer ReLU and
skip-add are applied later). Channels go ic → mid → oc (generic;
set ic = mid = oc = c for the identity-skip block, ic ≠ oc for the
downsample/projection block). Spatial dims h w preserved.
Inner convBnRelu₁ needs the smoothness hyp h_smooth₁ (no post-bn₁
activation hits the ReLU kink); outer convBn₂ is everywhere
differentiable, lifted via .toHasVJPAt. Two vjp_comp_at chain.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Basic-block body is DifferentiableAt at a smooth point. Needed as
the diff witness when feeding the body into the residual/projection
fan-in and the post-add ReLU.
Full basic residual block VJP (identity skip).
relu(x + F(x)) with F the conv→bn→relu→conv→bn body and an
identity skip (so ic = mid = oc = c, spatial preserved). Two
smoothness hyps: h_smooth₁ for the inner block ReLU, and
h_smooth_out for the post-add outer ReLU (F v + v avoids the
kink). Built as relu ∘ residual F via residual_has_vjp_at then a
final vjp_comp_at with relu.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Downsample/projection basic residual block VJP.
relu(proj(x) + F(x)) where the channel/stride change is folded into
the conv dims: body F maps ic → oc (first conv ic → oc, second
oc → oc), and the skip is a 1×1 convBn projection proj : ic → oc
(everywhere differentiable — no ReLU, so its diffAt is immediate).
Built with residualProj_has_vjp_at, then vjp_comp_at with the
post-add relu under h_smooth_out.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Weight gradient (now proved from foundation via unfold + fun_prop) #
The conv weight gradient implements the transpose trick:
`dW[o, c, kh, kw] = Σ_{h, w} x[c, h+kh−p, w+kw−p] · dy[o, h, w]`
Here's the slick observation: this is a convolution, with the input and gradient playing the roles of "input" and "kernel" respectively.
- View the input
x : (ic, H, W)as(ic, 1, H, W)(treat channels as batch). - View the gradient
dy : (oc, H, W)as(oc, 1, H, W)(same trick). - Now do a standard convolution: input shape
(ic, 1, H, W), kernel shape(oc, 1, H, W). The "spatial" dims of the kernel are H×W (the whole image), so the output is the small(ic, oc, kH, kW)weight gradient — produced by sliding the gradient as a giant kernel. - Transpose the output
(ic, oc, kH, kW) → (oc, ic, kH, kW)to match the kernel layout.
This avoids needing a separate "convolution-with-funny-dimension-numbers"
op; we use the same forward conv operator, just with shapes reinterpreted.
Critical for backends like IREE that don't accept non-standard
dimension_numbers (see iree-org/iree#21955).
MLIR (Conv 1 backward — exactly this trick): %x_t = stablehlo.transpose %x, dims = [1, 0, 2, 3] -- (1,128,28,28) %dh0p_t = stablehlo.transpose %d_h0pre, dims = [1, 0, 2, 3] -- (32,128,28,28) %d_W0_raw = "stablehlo.convolution"(%x_t, %dh0p_t) ... -- (1,32,3,3) %d_W0 = stablehlo.transpose %d_W0_raw, dims = [1, 0, 2, 3] -- (32,1,3,3)
Framework. HasVJP3 covered only input→output VJPs. For the
weight gradient we reuse the plain HasVJP on Vec by flattening
both the kernel (Kernel4.flatten : Kernel4 → Vec (oc*ic*kH*kW)) and
the output (Tensor3.flatten : Tensor3 → Vec (oc*h*w)). The bundled
HasVJP def packages a correct backward for the flattened function
together with its proof; the user-facing conv2d_weight_grad wrapper
does the flatten / unflatten housekeeping so callers see the natural
Kernel4 type.
Numerical validation: check_jacobians.py:test_conv2d_weight_grad
gradient-checks the transpose-trick formula against finite differences.
Conv2d weight-VJP — proved from foundation rules.
The function v ↦ flatten (conv2d (unflatten v) b x) is affine in
v: a constant b o(idx_out) plus a triple sum over (c, kh, kw)
of (unflatten v) o(idx_out) c kh kw * x_pad_term. Each summand
factors as (reindex of v) * (x-only constant), so pdiv_add +
pdiv_const + pdiv_finset_sum (×3) + pdiv_mul + pdiv_reindex
collapse the per-(idx_in, idx_out) pdiv. Triple-sum collapse via
Finset.sum_eq_single gives the transpose-trick backward
dW[o', c', kh', kw'] = Σ_{hi, wi} x_pad_term(...) · dy(flat(o', hi, wi)).
Equations
- One or more equations did not get rendered due to their size.
Instances For
Named accessor for the conv2d weight backward — aligns with MLIR
codegen (the "transpose trick" stablehlo.convolution in the backward
pass). Unwraps the flattening so callers see Kernel4 → Kernel4.
Equations
- Proofs.conv2d_weight_grad W b x dy = Proofs.Kernel4.unflatten ((Proofs.conv2d_weight_grad_has_vjp b x).backward W.flatten dy.flatten)
Instances For
Conv2d bias-VJP — proved from foundation rules. Now that conv2d
is a real def, the function b ↦ flatten (conv2d W b x) decomposes
as (channel-reindex from b) + (W,x-only term constant in b). Apply
pdiv_add + pdiv_reindex + pdiv_const, then collapse the
Kronecker over the (c, hi, wi) decomposition of Fin (oc*h*w).
The backward is db[o] = Σ_{hi, wi} dy[o, hi, wi] (matches
conv2d_bias_grad_formula below).
Equations
- One or more equations did not get rendered due to their size.
Instances For
Named accessor for the conv2d bias backward via the VJP framework.
Equations
- Proofs.conv2d_bias_grad W b x dy = (Proofs.conv2d_bias_grad_has_vjp W x).backward b dy.flatten
Instances For
Conv2d bias gradient — closed-form formula (documented, numerically
verified, expected to equal conv2d_bias_grad up to fp precision).
db[o] = Σ_{h, w} dy[o, h, w]
Each output cell adds the same b[o], so its gradient accumulates
the contributions from every spatial position. MLIR emits this as
a stablehlo.reduce across the spatial (and batch) dims.
Equations
- Proofs.conv2d_bias_grad_formula dy o = ∑ y : Fin h, ∑ x : Fin w, dy o y x
Instances For
MaxPool 2×2 stride 2 forward — concrete definition.
Each output cell is the maximum of a 2×2 window of input cells:
y[c, h, w] = max{ x[c, 2h+a, 2w+b] : a, b ∈ {0,1} }. No longer
an axiom — replaced with the explicit four-way max.
MLIR: %pool = "stablehlo.reduce_window"(%h1, %neginf) ({ ^bb0(%a, %b): stablehlo.return (stablehlo.maximum %a, %b) }) {window_dimensions = [1, 1, 2, 2], window_strides = [1, 1, 2, 2]}
Equations
Instances For
MaxPool2 input-VJP — gradient routes only to the argmax positions.
The backward function implements:
dx[c, 2h+a, 2w+b] = dy[c, h, w] · 𝟙[(a,b) is the argmax of the window]
Conceptually, max-pool is a piecewise selection: each output is one specific input. So the Jacobian is a sparse 0/1 matrix and the VJP just routes the gradient to the chosen input.
MLIR uses tile-compare-select — stablehlo.select_and_scatter
is avoided because IREE does not support it (see MlirCodegen.lean's
maxPool backward case for the full emitter):
// Broadcast dy and the pooled output back up to the input shape: %dy_tiled = stablehlo.broadcast_in_dim %d_pool %out_tiled = stablehlo.broadcast_in_dim %pool // Mask the input cells whose value matches the window max: %mask = stablehlo.compare EQ, %out_tiled, %h1 // Route gradient through that mask (zeros elsewhere): %d_h1 = stablehlo.select %mask, %dy_tiled, %zero
Canonical (junk-at-tie) witness. HasVJP3.correct is
satisfied by the canonical pdiv3-derived backward via rfl. At
argmax-tie boundaries maxPool2 is not differentiable, so pdiv3
agrees with fderiv's junk default of 0 and the canonical
witness is also 0 there. The codegen emits the tile-compare-
select formula above instead — at ties, the EQ-mask routes the
gradient to every tied input cell (PyTorch/JAX semantics), not
a single deterministic argmax. See LeanMlir/Proofs/README.md for
the trust-boundary discussion. Smooth-point agreement is formal:
see maxPool2_codegen_matches_canonical below.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Smooth-point codegen bridge for MaxPool2 #
Closes the smooth-point half of the codegen trust boundary at MaxPool2.
At points where every 2×2 window has a unique strict argmax, the
canonical pdiv-derived backward in maxPool2_has_vjp3 collapses to
"route dy to the argmax position, zero elsewhere" — the formula that
MlirCodegen.lean emits via tile-compare-select (broadcast dy and the
pooled output, compare EQ to find the argmax cells, select to
route). Mirrors relu_codegen_matches_canonical in MLP.lean, but
the local linearization is per-2×2-window rather than per-coordinate.
Smoothness: every 2×2 window of x has pairwise-distinct
values (so a unique strict argmax). The natural domain on which
maxPool2 is differentiable.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Input position (ci, hi_in, wi_in) attains the max of its window.
Equations
- Proofs.MaxPool2IsArgmax x ci hi_in wi_in = ∀ (a b : Fin 2), x ci (Proofs.winRowInv (Proofs.winRow hi_in) a) (Proofs.winColInv (Proofs.winCol wi_in) b) ≤ x ci hi_in wi_in
Instances For
A (not necessarily unique) argmax of the 2×2 window at output
position (co, ho, wo). Unique under MaxPool2Smooth.
Equations
- Proofs.maxPool2Argmax x co ho wo = Classical.choose ⋯
Instances For
If (a, b) dominates every other window cell, the max-pool output
equals the value at (a, b). No smoothness needed.
Under smoothness, the argmax of any window is unique: two positions that both dominate the window coincide.
Under smoothness, MaxPool2IsArgmax pins maxPool2Argmax to the
(winRowMod, winColMod) position of the witness.
For each output flat index k_out (decoded to (co, ho, wo)), the
flat index of the argmax's input position in Vec (c * (2*h) * (2*w)).
Used as the carrier of the local-linearization reindexCLM.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Smooth-point local-linearization for max-pool. On a metric ball
around flatten x, the flattened max-pool agrees with the reindex
y ↦ y ∘ σ where σ routes each output position to its argmax's
input position. Promoted via EventuallyEq.
MaxPool2 smooth-point Jacobian. At a smooth point, pdiv3 of
maxPool2 is a sparse 0/1 indicator: 1 exactly when the output
(co, ho, wo) is the window of the input (ci, hi_in, wi_in) AND
that input is the argmax of its window.
Bridge: maxPool2_has_vjp3's canonical backward matches the
codegen formula at smooth points.
At points where every 2×2 window has a unique strict argmax, the
canonical pdiv3-derived backward collapses to "dy at the
window's output position, but only at the argmax input cell" — the
tile-compare-select formula MlirCodegen.lean emits. Closes the
smooth-point half of the codegen trust boundary; what remains is
the kink convention at argmax-tie boundaries (EQ-mask routes the
gradient to every tied cell).
MaxPool2 pointwise VJP — no canonical-witness escape.
HasVJPAt3 maxPool2 x under MaxPool2Smooth x. The backward is
the codegen tile-compare-select formula directly (route dy to
the argmax cell, zero elsewhere); the correct field is
maxPool2_codegen_matches_canonical flipped, not rfl.
Companion of relu_has_vjp_at in MLP.lean — together they let
mlp_has_vjp_at and (future) cnn_has_vjp_at3 discharge the chain
rule through every kinked operator without the global vacuous
witness.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Reshape (flatten / unflatten) #
Flatten is a permutation of indices, so its VJP is just the inverse permutation. No gradient computation needed.
MLIR: %flat = stablehlo.reshape %pool : (tensor<128x32x14x14xf32>) -> tensor<128x6272xf32>
The flatten / unflatten bijection is already defined in
Tensor.lean as Tensor3.flatten / Tensor3.unflatten (used by the
pdiv3 derivation in Phase 5). We reuse those here rather than
duplicating — see Tensor3.flatten_unflatten / unflatten_flatten
for the mutual-inverse proofs.
Summary of derivations in this file #
conv2d,maxPool2— forward operations (black-box forward).maxPool2_has_vjp3— input-path VJP for maxPool2 (argmax-routing subgradient convention).
Derived (not axioms):
conv2d_has_vjp3— Phase 4: input-path VJP, proved from foundation rules using the per-coord pdiv chain plus a custompdiv_pi_pad_evalhelper for the dependent-if hpad : pad then v(σ hpad) else 0pattern. Backward function isconv2d_input_grad_formula(sum over(co, ho, wo)with reconstructed kernel offsetskh = hi+pH-ho,kw = wi+pW-wo).conv2d_weight_grad_has_vjp— Phase 7: the weight-path VJP, bundled as a plainHasVJPon the Kernel4-flattened function. Numerically gradient-checked against the transpose-trick formula incheck_jacobians.py:test_conv2d_weight_grad.conv2d_bias_grad_has_vjp— Phase 9: the bias-path VJP, same bundledHasVJPpattern. The closed-form "sum output cotangent over spatial dims per channel" is expressed asconv2d_bias_grad_formula; the namedconv2d_bias_gradextracts the backward via the VJP.conv2d_input_grad,maxPool2_input_grad,conv2d_weight_grad,conv2d_bias_grad— named accessors, defined as.backward(plus flatten / unflatten housekeeping for the weight / bias variants) of the corresponding VJP.conv2d_input_grad_formula,conv2d_bias_grad_formula— the concrete closed-form formulas (numerically verified to equal the VJP's backward).- 3D reshape (
Tensor3.flatten/Tensor3.unflatten) imported fromTensor.lean; 4D reshape (Kernel4.flatten/Kernel4.unflatten) defined here, both proved bijections.
Public correctness theorem for maxPool2_has_vjp3: the
canonical-witness backward equals the pdiv3-contracted Jacobian
by definition. The codegen substitutes the standard argmax-routing
convention at non-smooth tiebreaks (see LeanMlir/Proofs/README.md's
Codegen Trust Boundary).
Public correctness theorem for maxPool2_has_vjp_at3 — the
pointwise variant under MaxPool2Smooth. The underlying .correct
field is maxPool2_codegen_matches_canonical flipped (a real proof),
not rfl; this wrapper exposes it for comparator re-verification.
The capstone: a whole-network ResNet-style CNN VJP #
cnn_has_vjp_at is the CNN analogue of vit_full_has_vjp — a single
HasVJPAt for an end-to-end forward pass, chained entirely in flattened
Vec space via vjp_comp_at. It first needs global average pooling,
which was previously only referenced in codegen, so we define it here:
globalAvgPool x ci = (∑ hi ∑ wi x ci hi wi) / (h*w) (mean over spatial
per channel), bridge it to flat Vec space (globalAvgPoolFlat), and
prove its linear VJP (globalAvgPoolFlat_has_vjp, backward broadcasts
dy ci / (h*w) to every spatial cell of channel ci) and
differentiability.
Fixed structural choices (a concrete-but-representative pipeline, in
the spirit of hand_cnn_train_step.mlir; prioritising a complete
axiom-clean end-to-end witness over maximal generality):
input : Vec (ic * (2h) * (2w))
stem : convBnRelu ic → c (spatial 2h×2w preserved)
pool : maxPool2 c, 2h×2w → c, h×w
block1 : resblock_has_vjp_at (identity skip, c → c, h×w)
block2 : resblockProj_has_vjp_at (projection skip, c → oc, h×w)
gap : globalAvgPool oc, h×w → Vec oc
head : dense oc → nClasses
So: 1 stem conv, 1 max-pool, exactly two residual blocks (one of EACH
skip type — identity and 1×1 projection — to exercise both code paths),
global-average pool, one dense classifier. Channel/spatial dims stay
implicit Nat params; the block/stage counts are fixed. The bundled
smoothness hypotheses (h_stem, h_mp, h_rb1/h_rb1o, h_rb2/
h_rb2o) are the family of every ReLU + max-pool site's smooth-point
condition, exactly like mlp_has_vjp_at's multiple h_smooth_*.
The differentiability obstacle (max-pool is non-smooth globally, so
vjp_comp_at cannot get DifferentiableAt of a max-pool-containing
prefix from a global lemma) is discharged at the smooth point by
maxPool2_flat_hasFDerivAt (the local linearization already proved for
the max-pool Jacobian) via .differentiableAt.
Flat GAP: Vec (c*h*w) → Vec c.
Equations
Instances For
The channel of a flat index idx : Fin (c*h*w).
Equations
- Proofs.flatChannel c h w idx = (finProdFinEquiv.symm (finProdFinEquiv.symm idx).1).1
Instances For
Global average pool VJP (flattened). Linear map; backward
broadcasts dy ci / (h*w) to every spatial cell of channel ci.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Uniform VJP-correctness wrapper for globalAvgPoolFlat — a citable
_correct matching the convention of every other layer (just unfolds the
HasVJP.correct field of globalAvgPoolFlat_has_vjp).
Equations
- Proofs.maxPoolFlat c h w v = (Proofs.maxPool2 (Proofs.Tensor3.unflatten v)).flatten
Instances For
Equations
- Proofs.maxPoolFlat_has_vjp_at x h_smooth = Proofs.hasVJPAt3_to_hasVJPAt (Proofs.maxPool2_has_vjp_at3 x h_smooth)
Instances For
Max is exact in floating point + 1-Lipschitz. max a b is a
compare-and-select: it returns one of a, b verbatim, rounding nothing.
So a float max over operands within e of the reals stays within e —
the max-peer of relu_close (FloatBridge.lean), with no rounding term
and no amplification. The one genuinely-new fact the MNIST-CNN forward
rounding budget (planning §1b-A) needs beyond the dense/relu machinery.
MaxPool2 is exact in floating point + 1-Lipschitz. Four window cells
through three max-selections, no arithmetic — inherited input error e
passes through with no rounding term and no amplification.
Equations
- Proofs.cbr W b ε γ β = Proofs.relu (oc * h * w) ∘ Proofs.bnForward (oc * h * w) ε γ β ∘ Proofs.flatConv W b
Instances For
Equations
- One or more equations did not get rendered due to their size.
Instances For
The forward CNN: stem(convBnRelu) → maxpool → resblock(id) → resblockProj → gap → dense.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Equations
- One or more equations did not get rendered due to their size.
Instances For
Public correctness theorem for cnn_has_vjp_at — exposes the
witness's .correct field as a top-level proposition: the full
ResNet-style CNN's backward equals the pdiv-contracted Jacobian
(Jacobian-transpose applied to the cotangent). CNN analogue of
vit_full_has_vjp_correct.