Batched backward links — the backward graphs and cotangent steps at the batched index #
The vocabulary every batched whole-net backward proof and every T3 tie is written in, at the flat
batched index N·(c·h·w). Each graph or cotangent here denotes the .backward of a proven VJP
(BatchedStages), so a chain built from them is the loss-driven backward, not a free cotangent.
| what | names | namespace |
|---|---|---|
| residual fan-in backward graph | residualBackGraph / _faithful | StableHLO |
batched op backwards: true BN, conv, strided conv, depthwise (stride 1 / symmetric / XLA-SAME stride 2), SE | bnBatchBack_faithful, bnBatchLABack_faithful, convBackBatched_faithful, … seBackBatched_faithful | StableHLO |
| stage backward graphs | cbsBackBatchedGraph, dwbsBackBatchedGraph, dwbsSBackBatchedGraph, projBackBatchedGraph (+ _faithful) | StableHLO |
cotangent steps: BN, swish, sigmoid, conv / depthwise input-VJPs, GAP, SE, SE gate; the c·h·w ↔ c·(h·w) reindex | bnBackB, swBackB, sigBackB, cInB, dInB, dStridedInB, gapInB, seInB, gateCotB, reassocB | EnetTiePoC |
| cotangent steps: relu mask, strided conv input-VJP, BN as emitted, 3×3/s2 max-pool; the one-row head casts | reluMaskB, cStridedInB, bnInB, mpInB, rowB / unrowB | ResNet34TieB |
Namespaces are the net that first needed each piece; the names are cited by every conv net's tie.
Backward graph for a residual block x ↦ x + f x, given a subgraph
fBack that renders the body f's input-cotangent. The identity skip
contributes the cotangent verbatim (%dy); addV sums the two paths.
This is the renderable image of residual_has_vjp's biPath backward.
Equations
- Proofs.StableHLO.residualBackGraph fBack ecot = fBack.addV ecot
Instances For
Residual additive-fan-in backward faithfulness (general).
If fBack denotes the body's VJP backward (den fBack = hf.backward x dy),
then the residual backward graph denotes the proven residual_has_vjp
backward, which is f.backward x dy + dy. The proof is structural — the
only definitional facts are den (addV a b) = den a + den b and the
identity skip's backward = dy — so it composes without a whole-net
terminal rfl.
The renderable batch-norm input-grad IS the certified backward — both equal the
pdiv-contracted Jacobian (bnBatchTensor4_grad_input_correct,
bnBatchTensor4_has_vjp_correct).
bnBatchBack (true batch-norm backward) faithfulness. The first
batched-backward primitive: bnBatchBack denotes the proven
bnBatchTensor4 VJP backward (batch-COUPLED batch-norm on [N,C,H,W],
reduce over [0,2,3] per channel) via the renderable three-term
bnBatchTensor4_grad_input. This is the genuinely-new op the batched MBConv
stages need (their bn is bnBatchLA, not a per-example batchMap); the
other batched stages (conv/depthwise/SE) are batchMap of the per-example
backwards already proven above. The bnBatchLA layout-reindex wrapper to the
network's N·(oc·h·w) index is a thin remaining layer.
Batched conv input-VJP faithfulness. convBackBatched denotes the proven
VJP of the batched conv batchMap N (flatConv W b) — i.e. the per-example
conv input-grad applied independently across the batch. Conv is linear, so
its backward ignores the forward activation; the batched backward is a plain
batchMap of the per-example backward, matching batchMap_has_vjp. The
second batch-separable stage brick (after seB); together with bnBatchBack
these are the batched MBConv's per-stage backward pieces.
Batched STRIDE-2 conv input-VJP faithfulness. The stride-2 analogue of
convBackBatched_faithful: convStridedBackBatched denotes the proven VJP of
the batched strided conv batchMap N (flatConvStride2 W b) — i.e. the
per-example strided-conv input-grad (flatConvStride2_has_vjp = zero-upsample
the cotangent then the reversed-kernel conv) applied independently across the
batch. Strided conv (decimate ∘ conv) is linear, so its backward ignores the
forward activation; the batched backward is a plain batchMap of the
per-example backward, matching batchMap_has_vjp. The downsample basic-block's
stride-2 conv1 backward brick.
Batched STRIDE-2 depthwise input-VJP faithfulness. The stride-2 analogue
of depthwiseBackBatched_faithful (and the depthwise analogue of
convStridedBackBatched_faithful): depthwiseStridedBackBatched denotes the
proven VJP of the batched strided depthwise batchMap N (depthwiseStride2Flat W b)
— i.e. the per-example strided-depthwise input-grad (depthwiseStride2Flat_has_vjp
= zero-upsample the cotangent then the reversed-kernel per-channel depthwise)
applied independently across the batch. Strided depthwise (decimate ∘ depthwise)
is linear, so its backward ignores the forward activation; the batched backward
is a plain batchMap of the per-example backward, matching batchMap_has_vjp.
The EfficientNet downsample MBConv's stride-2 depthwise backward brick.
Batched XLA-SAME STRIDE-2 depthwise input-VJP faithfulness. The odd-phase peer of
depthwiseStridedBackBatched_faithful: depthwiseStridedXlaBackBatched (pad [p+1, p-1],
the token MobileNetV2's Adam render emits at its four strided depthwises) denotes the proven
VJP of batchMap N (depthwiseStride2FlatXla W b). Same proof: a scatter onto the odd
positions is as linear as one onto the even ones.
Batched depthwise input-VJP faithfulness. The depthwise analogue of
convBackBatched_faithful: depthwiseBackBatched denotes the proven VJP of
the batched depthwise batchMap N (depthwiseFlat W b). Depthwise conv is
linear, so its backward is activation-independent and the batched backward is
a plain batchMap of the per-example backward. The MBConv depthwise stage's
batch-separable backward brick.
bnBatchLA backward = reindex-conjugated bnBatchTensor4 backward.
The network indexes at N·(oc·h·w) (left-assoc) but the proven true-BN
bnBatchTensor4 lives at N·(oc·(h·w)); bnBatchLA bridges by conjugating
with the associativity-cast reindexes (bnBatchLA_eq_comp). Its VJP backward
is therefore: scatter the cotangent into [N,C,(H·W)], run the renderable
three-term bnBatchTensor4_grad_input at the reindexed activation, scatter
back. This is what a network-layout bnBatchLABack op denotes.
bnBatchLABack (network-layout true batch-norm backward) faithfulness.
The den (inline scatter-conjugated bnBatchTensor4_grad_input) equals the
proven bnBatchLA_has_vjp backward — the bn backward at the network's
N·(oc·h·w) index. This is the
layout wrapper that lets bnBatchBack compose with convBackBatched /
depthwiseBackBatched (all on the left-assoc index) into batched stages.
seBackBatched (batched squeeze-excite backward) faithfulness. The den
(rowwise application of the proven per-example seBlockFull VJP) equals the
proven batched seB_has_vjp backward. SE is non-linear, so — unlike the
linear convBackBatched/depthwiseBackBatched — the backward threads each
example's forward activation v; the rowwise batchMap_has_vjp structure
handles that. The fourth (and last) MBConv stage's batch-separable backward.
Batched conv → bn → swish stage backward graph (MBConv expand), at the
network layout: convBackBatched ∘ bnBatchLABack ∘ swishBack, each at its
cumulative forward activation.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Batched depthwise → bn → swish stage backward graph (MBConv depthwise).
Equations
- One or more equations did not get rendered due to their size.
Instances For
Batched STRIDE-2 depthwise → bn → swish stage backward graph (the
EfficientNet downsample MBConv's depthwise). The stride-2 analogue of
dwbsBackBatchedGraph: the bn/swish run at the OUTPUT spatial h×w, then
depthwiseStridedBackBatched maps the bn-cotangent back to the larger input
c·(2h)·(2w) (zero-upsample + reversed-kernel per-channel depthwise).
Equations
- One or more equations did not get rendered due to their size.
Instances For
Batched conv → bn stage backward graph (MBConv project, no swish).
Equations
- One or more equations did not get rendered due to their size.
Instances For
Chain-cotangent helpers — the per-op batched backward steps (built fresh, HasVJP-style) #
EfficientNetChainClose proves the per-block VJPs by vjp_comp of the per-op VJPs but exposes no
explicit cotangent-vector defs (unlike mnv2's invresCot*). So the tie BUILDS the chain cotangents
from the proven per-op backwards: bnBackB (true-BN, the batch-coupled bnBatchLA VJP), swBackB
(swish, smooth), cInB/dInB (the batched conv/depthwise input-VJP = den convBackBatched/
depthwiseBackBatched), seInB (the fused SE input-cot = den seBackBatched), gateCotB (the SE
gate cotangent = den seReduceB), sigBackB, rowDenseBackFlat (the SE excite/reduce backs). Every
helper IS a .backward of a proven VJP (or the exact den of the emitted backward op), so the
cotangents are the genuine loss-driven backward, not a free ∀c. reassocB bridges the conv/swish
index (oc·h·w) to the BN param-op index (oc·(h·w)).
(oc·h·w) → (oc·(h·w)) batched reassociation reindex — bridges the conv/swish chain index to the
BN γ/β + conv-bias op index (EnetPoC.bn{Gamma,Beta}B_den consume Vec (N·(oc·(h·w)))).
Equations
- Proofs.EnetTiePoC.reassocB N oc h w v i = v (Fin.cast ⋯ i)
Instances For
Batched true-BN input-cotangent (bnBatchLA VJP — batch-coupled).
Equations
- Proofs.EnetTiePoC.bnBackB N oc h w ε hε γ β x dy = (Proofs.bnBatchLA_has_vjp N oc h w ε hε γ β).backward x dy
Instances For
⭐ The tie's BN node and the emitted BN node denote one map. Every batched render emits
.bnBatchBack, typed at N·(oc·(h·w)); the ties state the BN input cotangent at
.bnBatchLABack, its network-layout N·(oc·h·w) twin (ResNet34TieB.bnInB). The two print
the same text, and their dens differ only by the associativity relabelling reassocB: the
two scatters inside bnBatchLABack's den collapse because Fin.cast is a bijection.
The certified BN input cotangent every batched T3 tie threads IS the emitted bnBatchBack
node's den, read back through reassocB. This is the missing half of
bnBatchLABack_faithful: that lemma certifies the tie's node, this one says the render's node
computes the same numbers.
Batched swish mask-back (smooth, no kink).
Equations
- Proofs.EnetTiePoC.swBackB n x dy = (Proofs.swish_has_vjp n).backward x dy
Instances For
Batched sigmoid back (the SE gate excite-dense output cotangent).
Equations
- Proofs.EnetTiePoC.sigBackB n x dy = (Proofs.sigmoid_has_vjp n).backward x dy
Instances For
Batched 1×1/conv input-VJP (= den convBackBatched; conv is linear, x unused).
Equations
- Proofs.EnetTiePoC.cInB N W b dy = Proofs.StableHLO.batchMap N (fun (d : Proofs.Vec (oc * h * w)) => (Proofs.flatConv_has_vjp W b).backward (fun (x : Fin (ic * h * w)) => 0) d) dy
Instances For
Batched depthwise input-VJP (= den depthwiseBackBatched).
Equations
- Proofs.EnetTiePoC.dInB N W b dy = Proofs.StableHLO.batchMap N (fun (d : Proofs.Vec (c * h * w)) => (Proofs.depthwiseFlat_has_vjp W b).backward (fun (x : Fin (c * h * w)) => 0) d) dy
Instances For
Batched strided depthwise input-VJP (= den depthwiseStridedBackBatched; upsamples h→2h).
Equations
- One or more equations did not get rendered due to their size.
Instances For
Batched GAP input-VJP (= den gapBackBatched; the head's GAP backward, broadcast÷(h·w)).
Equations
- Proofs.EnetTiePoC.gapInB N c h w dy = Proofs.StableHLO.batchMap N (fun (d : Proofs.Vec c) => (Proofs.globalAvgPoolFlat_has_vjp c h w).backward (fun (x : Fin (c * h * w)) => 0) d) dy
Instances For
Batched fused SE input-cotangent (= den seBackBatched, the x⊙gate VJP).
Equations
- Proofs.EnetTiePoC.seInB N W₁ b₁ W₂ b₂ x dy = (Proofs.seB_has_vjp N W₁ b₁ W₂ b₂).backward x dy
Instances For
Batched SE gate cotangent dgate[n,c] = Σ_{h,w}(x⊙dy) (= den seReduceB, the broadcast-adjoint
of x ⊙ dy — the FIRST step of the SE gate backward, feeding the SE dense param grads).
Equations
- One or more equations did not get rendered due to their size.
Instances For
The relu backward mask — den (.selectPosB _ pre e) = fun i => if pre i > 0 then e i else 0.
r34 applies it twice per block (the body's mid-relu and the post-residual outer one) and once at
the stem.
Instances For
Batched STRIDED conv input-VJP (= den convStridedBackBatched; upsamples h → 2h). The
strided peer of EfficientNet's cInB. ⚠ SYMMETRIC padding — flatConvStride2, not the
XLA-SAME twin.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Batched true-BN input-cotangent, as the EMITTED backward computes it. Written as the den
of the backward op rather than as the certified VJP's .backward, because that is the form the
render's chain is in and den ignores the name strings — so every cotangent below is literally
what the artifact's bytes compute. ⚠ The render's node is .bnBatchBack, typed at
N·(oc·(h·w)); this is its network-layout twin, and bnInB_eq_den_bnBatchBack below says
the two denote one map up to reassocB. ⭐ It takes no β: the BatchNorm input-gradient does not
depend on the shift, which bnInB_eq_bnBackB records by holding for every β.
Equations
- Proofs.ResNet34TieB.bnInB N oc h w ε γ x dy = Proofs.StableHLO.den (Proofs.StableHLO.SHlo.bnBatchLABack "" "" "" ε γ x (Proofs.StableHLO.SHlo.operand "" dy))
Instances For
…and it IS the certified bnBatchLA VJP, for every β and every 0 < ε. This is
bnBatchLABack_faithful, and it is the only step in this file's cotangent chain that is not
rfl — everything else (the relu masks, the conv and strided-conv input-VJPs, the pool
backward) denotes its certified backward definitionally.
…and it IS the den of the node the render emits, .bnBatchBack at the N·(oc·(h·w))
index, read back through reassocB (EnetTiePoC.den_bnBatchLABack_eq_bnBatchBack).
Batched 3×3/s2 max-pool backward (= den maxPool3s2BackB): the select_and_scatter
denotation, per example on that example's own saved activation — which is why it is
batchMapAux and not batchMap.
Equations
- Proofs.ResNet34TieB.mpInB N c h w x dy = Proofs.StableHLO.batchMapAux N (Proofs.StableHLO.maxPool3s2BackFlat c h w) x dy
Instances For
Vec (N·(1·K)) → Vec (N·K): the loss chain runs at one ROW per example (softmaxRow needs a
row index) and the dense parameter ops at the plain per-example width. The render writes one
SSA name for both, because 1 * K = K as an emitted shape; in Lean the two indices are
propositionally but not definitionally equal, so the cast is explicit.
Equations
- Proofs.ResNet34TieB.unrowB N K v i = v (Fin.cast ⋯ i)