ViT backward-graph faithfulness (per-token Mat VJP) #
The ViT analogue of the other four nets' *BackB0 capstones: backward
StableHLO graphs (over the ch10 backward tokens — denseRowBack/geluBack/
lnRowBack/softmaxRowBack/matmulF/transposeF/scaleF/addV) whose
denotations ARE the proven VJPs, from the MLP and MHSA pieces up to the multi-head
vector-LN block the shipped net runs (transformerBlockVHasVJPMat) and the
depth-k whole net (vitNetBackGraph).
Unlike the conv nets — whose blocks live natively as Vec → Vec (HasVJP) — a
transformer block lives in the per-token matrix framework HasVJPMat (Mat N D → Mat N D). The block VJP's .backward is therefore Mat-valued, while den is
Vec-valued; the faithfulness statements bridge the two through Mat.flatten
(the convention the forward-graph faithfulness theorems use):
den (…BackGraph A e) = Mat.flatten ((… HasVJPMat …).backward A (Mat.unflatten (den e)))
The pieces, bottom-up:
- MLP —
transformerMlpBackGraph↔transformerMlpHasVJPMat(denseRowBack(Wfc1) ∘ geluBack ∘ denseRowBack(Wfc2)at the saved activations). - MHSA — a clean witness
mhsaCleantied tomhsaHasVJPMat, its multi-head collapsemhsa_backward_collapseMH(a sum over heads of per-headsdpaBack{Q,K,V}), and the MHSA backward graphmhsaBackGraphMHoversdpaBack{Q,K,V}Graph. - Vector-LN block, tower, whole net —
transformerBlockVBackGraphMH,vitBodyBackGraphKMHV(by induction onk) andvitNetBackGraph.
The framework .backward rules (Tensor.lean):
vjpMatComp F G … .backward A dY = hF.backward A (hG.backward (F A) dY)biPathMatHasVJP F G … .backward A dY i j = hF.backward A dY i j + hG.backward A dY i j; the identity skip's.backward A dY = dY.
Each ch10 backward den helper is a flattened rowwise op; here we tie it to
the HasVJPMat .backward of the corresponding framework piece. These are the
ViT-backward analogues of ViTFwdGraph's rowDenseFlat_flat/rowLNFlat_flat
forward bridges.
Per-token dense input-VJP: the flat rowDenseBackFlat IS the flatten of the
rowwise densePerTokenHasVJPMat.backward (which ignores the saved
activation A, dense being affine — denseHasVJP.backward _ dy = Mat.mulVec W dy).
The transformer MLP backward graph (reverse-order chain of
transformerMlpHasVJPMat = dense2 ∘ gelu ∘ dense1 per-token; outermost
backward token = earliest forward op = dense1):
denseRowBack(Wfc1) ∘ geluBack(@ pre-GELU = dense1(Y)) ∘ denseRowBack(Wfc2)
where Y is the MLP input (= LN₂ h). The GELU backward reads its saved
pre-activation m1 = dense1 Y.
Equations
- One or more equations did not get rendered due to their size.
Instances For
The proven mhsaHasVJPMat witness is built via by rw [mhsaLayer_eq_compose]; exact vjpMatComp …, so its .backward field does NOT reduce by rfl (the
Eq.mpr transport blocks whnf). We instead build a clean witness mhsaClean
for the same factored function (whose .backward unfolds transparently), tie it
to mhsaHasVJPMat by VJP determinism; the multi-head collapse below runs on it.
A clean HasVJPMat witness for the factored MHSA — the same vjpMatComp
chain mhsaHasVJPMat's body uses, but stated for the explicit composition
Wo-dense ∘ colSlabApply mhsaG ∘ qkv-dense so its .backward reduces by
rfl (no mhsaLayer_eq_compose transport in the way).
Equations
- One or more equations did not get rendered due to their size.
Instances For
The colSlab-lifted MHSA backward collapses to a SUM over heads: each head h slices the dense Q/K/V projections
and the Wo-back cotangent to head h's columns, runs sdpaBack{Q,K,V} at
d_head, and the qkv-stack dense-back contracts head h's SDPA backward against
the finProdFinEquiv (h, ·) columns of Wq/Wk/Wv.
The collapsed general-heads MHSA backward: for each head h, slice the
dense Q/K/V projections and the Wo-back cotangent to head h's columns,
run sdpaBack{Q,K,V} at d, then contract per-head against the
finProdFinEquiv (h, ·) columns of Wq/Wk/Wv, summed over heads.
Equations
- One or more equations did not get rendered due to their size.
Instances For
MHSA backward general-heads collapse. The clean MHSA witness's backward
equals the per-head sum fan-in mhsaBackCollapsedMH. Each fused-slab column of the
colSlabwise backward is sdpaBackSel at its head's dense Q/K/V; the fan-in lemma splits the
mulVec per head and per third.
The proven MHSA VJP's backward at general heads IS the per-head collapse.
The whole multi-head MHSA backward graph (heads = hm1 + 1): for each head h,
the three-way LN₁-fan-in over the per-head SDPA backward subgraphs (fed the
head-h slice of the Wo-back), padded into head h's columns and contracted
against Wq/Wk/Wv; summed over heads (headsSumG). Saved (per head): the dense
Q/K/V projections, the scaled pre-softmax scores ss, the post-softmax weights
p; plus the block cotangent dh.
Equations
- One or more equations did not get rendered due to their size.
Instances For
MHSA backward-graph faithfulness (multi-head, heads = hm1 + 1). The
per-head fan-in graph denotes the proven mhsaHasVJPMat.backward (flattened)
at general heads, with the saved per-head dense projections, scaled scores, and
post-softmax weights.
The committed verified_mlir/vit_train_step.mlir ViT-Tiny render uses VECTOR
γ/β per LN site, decomposed as (+βv) ∘ layerScale γv ∘ LN(1,0) (layerNormVec,
ViTVecLN.lean). The forward/backward MHSA and MLP-body are IDENTICAL to the scalar
block; the only difference is the two LN sites. So we REUSE mhsaBackGraphMH and
transformerMlpBackGraph verbatim and swap the LN-back fragment to the vec-LN one:
apply rowScaleF γv to the incoming cotangent, then lnRowBack at γ=1 (the bias
backward is the identity, so β drops out of the input cotangent).
The new structural facts:
layerNormVec_per_token_backward_eq— the vec-LN per-token backward IS the normalize-only (LN at γ=1,β=0) backward of the rowwise-layerScale γvcotangent (the bias backward = id collapses).rowVecLNBack_eq_backward— the flat compositionlnRowBack(γ=1) ∘ rowScaleF γvdenotes that vec-LN per-token backward (the crux bridge). Then the sublayer + whole-block capstones mirror the scalar/MH templates exactly, re-targeting theHasVJPMatto the…V…(vec-LN) versions.
Vec-LN LN-back bridge (Stage 1 crux). The flat composition lnRowBack(γ=1) of (rowScaleF γv applied to the cotangent) denotes the vec-LN per-token VJP's
.backward (flattened) at the saved pre-LN input X. The rowScaleF realizes
the rowwise layerScale γv on the incoming cotangent; the lnRowBack at γ=1 is
the normalize-only backward; the bias backward (identity) has dropped out.
The whole vec-LN MLP-sublayer backward graph (inner arm + identity skip).
The incoming cotangent is a subgraph ecot, so this sublayer can sit
downstream of another backward graph — which is what the attention sublayer
below does with it, and what CertLayer composition needs. At a fixed
dz : Vec it is ecot := .operand "%dz" dz.
Equations
- Proofs.StableHLO.mlpSublayerVBackGraph gf ε γ2v Wfc1 bfc1 Wfc2 h Y ecot = (Proofs.StableHLO.mlpSublayerVInnerBackGraph✝ gf ε γ2v Wfc1 bfc1 Wfc2 h Y ecot).addV ecot
Instances For
Vec-LN MLP sublayer backward-graph faithfulness (Stage 2 capstone), at general
(hm1+1)*d. Denotes the proven transformerMlpSublayerVHasVJPMat backward.
The vec-LN attn-sublayer non-trivial arm (mhsa ∘ LNᵥ₁), multi-head. REUSES
mhsaBackGraphMH verbatim, after the vector-LN back fragment lnRowBack(γ=1) ∘ rowScaleF γ1v.
Equations
- One or more equations did not get rendered due to their size.
Instances For
The whole vec-LN attention sublayer backward graph (inner arm + identity skip),
over an incoming-cotangent SUBGRAPH ecot (the old dh : Vec statement is this
one at ecot := .operand "%dh" dh).
Equations
- Proofs.StableHLO.attnSublayerVBackGraphMH ε γ1v Wq Wk Wv Wo bq bk bv bo x X ecot = (Proofs.StableHLO.attnSublayerVInnerBackGraphMH ε γ1v Wq Wk Wv Wo bq bk bv bo x X ecot).addV ecot
Instances For
The vector-LN attention sublayer's VJP backward decomposes (biPathMat unfold, rfl):
the residual skip passes the cotangent through, and the non-trivial arm is
LNᵥ₁-back ∘ mhsa-back at the saved LNᵥ₁ output.
Vec-LN attention sublayer backward-graph faithfulness (Stage 3 capstone, MH).
The whole vec-LN transformer-block backward graph (multi-head). transformerBlockV = mlpSublayerV ∘ attnSublayerV, so block.backward A dY = attn.backward A (mlp.backward (attn A) dY). Saved: A (block input), h = attnSublayerV A.
The MLP sublayer's graph is fed to the attention sublayer as a subgraph,
not as den (…) re-wrapped as an operand: the two sublayers compose symbolically,
and the block itself takes a cotangent subgraph so it can sit downstream of the
next block. This is what makes the depth-k tower a real composite term.
Equations
- One or more equations did not get rendered due to their size.
Instances For
The vector-LN block VJP backward unfolds: block.backward A dz = attn.backward A (mlp.backward (attn A) dz). The outer vjpMatComp's projection, rfl.
Whole vec-LN transformer-block backward-graph faithfulness (Stage 4 capstone, MH).
The production-parity capstone — multi-head (heads = hm1+1) + vector-LN, matching
the committed verified_mlir/vit_train_step.mlir ViT-Tiny config. Wires the vec-LN
MLP-sublayer backward (at the saved attn-sublayer output h = attnSublayerV A) into
the vec-LN attention-sublayer backward (at the saved block input A), per
block.backward A dY = attn.backward A (mlp.backward (attn A) dY).
The whole-net vitForwardKVHasVJP (ViTDepthK.lean) is a Vec → Vec
HasVJP built by three vjpComp steps:
classifierFlat ∘ [finalLN] ∘ vitBodyKVFlat(k) ∘ patchEmbedFlat
so its .backward x dy chains the four stages in REVERSE (head first, image last):
PE.back x (BODY.back (PE x) (LNF.back (BODY (PE x)) (classifier.back (… ) dy)))
We mirror this with a backward graph over the ch10 backward tokens. The blocks
live in HasVJPMat and are bridged to HasVJP by HasVJPMat.toHasVJP, so each
block-back faithfulness statement is stated through Mat.flatten/Mat.unflatten
(exactly the transformerBlockVBackGraphMH_faithful convention). Four stages:
- Stage 1 —
classifierBackGraph↔classifierFlatHasVJP.backward:clsPadF (dotOut Wcls (%dy))(dense-back into row 0 of a zero[N+1,D]). - Stage 2 —
finalLNBackGraph↔ the bridged per-token vec-LN VJP back: the vec-LN fragmentlnRowBack(γ=1) ∘ rowScaleF γF, chained onto Stage 1. - Stage 3 —
vitBodyBackGraphKMHV, a depth-kreverse fold oftransformerBlockVBackGraphMH(last forward block = first backward block), threading each block's saved forward activation (its faithfulness isViTBackNet's fold,vitTrunkV_graph). - Stage 4 —
patchEmbedBackGraph↔patchEmbedFlatHasVJP.backward: the strided-patchify conv input-VJP (thepatchEmbedBacktoken).
The classifier-head backward graph: clsPadF (dotOut Wcls (%dy)). The
dotOut Wcls is the dense head's input-VJP (Mat.mulVec Wcls), scattered by
clsPadF into row 0 of a zero [N+1,D] (the CLS-slice's input-VJP).
Equations
- Proofs.StableHLO.classifierBackGraph N D nClasses Wcls ecot = (Proofs.StableHLO.SHlo.dotOut "%Wcls" Wcls ecot).clsPadF
Instances For
Classifier backward-graph faithfulness (Stage 1). Denotes the proven
classifierFlatHasVJP.backward at any input v (dense + CLS-slice are
both linear, so the saved activation is irrelevant).
Over a cotangent subgraph, like patchEmbedBackGraph_faithful at the other end of the
net, so the head can be a CertLayer and the whole net one comp chain.
The final (pre-head) vector-LN backward graph over (N+1) tokens. REUSES the
vec-LN LN-back fragment (lnRowBack(γ=1) ∘ rowScaleF γF). The bias backward (identity)
drops out; X is the saved pre-LN input (the body output).
The classifier back is a separate graph (classifierBackGraph), composed at the call site,
so the final LN and the head are each their own CertLayer.
Equations
- Proofs.StableHLO.finalLNBackGraph N D ε γF X ecot = Proofs.StableHLO.SHlo.lnRowBack "%gF" "%XF" "ε" ε 1 X (Proofs.StableHLO.SHlo.rowScaleF "%gFv" γF ecot)
Instances For
Final vec-LN backward-graph faithfulness (Stage 2). Denotes the bridged
per-token vec-LN VJP back of whatever cotangent subgraph it is fed, at the saved
body output X.
transformerBlockVBackGraphMH at a bundled BlockParamsV block (the backward
analogue of vitBlockGraphMHVP). Saved: the block input A and its
attn-sublayer output h = attnSublayerV A.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Bundled per-block backward faithfulness. The block back graph at a bundled
BlockParamsV block (saved input A, cotangent dY) denotes the flatten of
the proven blockV's VJP backward — transformerBlockVHasVJPMat.backward,
spelled ONCE here over a generic p (so the depth-k induction never re-spells
the 16-field tuple). The block VJP is bundled as transformerBlockVHasVJPMatP.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Depth-k tower backward graph — the REVERSE fold of
transformerBlockVBackGraphMHP. The forward body runs block 0 first
(vitBodyKV (k+1) ps = vitBodyKV k (ps∘succ) ∘ blockV (ps 0)), so the
backward runs block 0 LAST: the incoming cotangent flows through the tail
(blocks 1..k, at the post-block-0 activation blockV (ps 0) A), then
through block 0 (at the saved block input A). Mirrors
vitBodyKVFlatHasVJP's vjpComp chain.
Equations
- One or more equations did not get rendered due to their size.
- Proofs.StableHLO.vitBodyBackGraphKMHV gf ε 0 x_4 x✝¹ x✝ = x✝
Instances For
The patch-embedding input-backward graph: the patchEmbedBack token (the
strided-patchify conv's input-VJP) on the patch-embed-output cotangent.
Equations
- Proofs.StableHLO.patchEmbedBackGraph ic H W P N D Wc e = Proofs.StableHLO.SHlo.patchEmbedBack "%Wp" Wc e
Instances For
patchEmbed input-backward-graph faithfulness (Stage 4). Denotes the proven
patchEmbedFlatHasVJP.backward at any saved image img (linear — the
activation is irrelevant).
Whole-net depth-k multi-head vector-LN ViT backward graph. Mirrors the
forward vitFwdGraphKMHV in REVERSE: classifier-back → final-vec-LN-back →
depth-k tower-back (reverse fold) → patchEmbed-back. Each stage is fed the
saved forward activation it differentiates at:
xin— the saved image (input to patchEmbed; patchEmbed-back is linear).bodyOut— the saved body output (input to the final LN).- the tower-back fold threads each block's saved input internally from the
patchEmbed output
embOut(= the body input).
Equations
- One or more equations did not get rendered due to their size.