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 (transformerBlockV_has_vjp_mat) 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 ((… _has_vjp_mat …).backward A (Mat.unflatten (den e)))
The pieces, bottom-up:
- MLP —
transformerMlpBackGraph↔transformerMlp_has_vjp_mat(denseRowBack(Wfc1) ∘ geluBack ∘ denseRowBack(Wfc2)at the saved activations). - MHSA — a clean witness
mhsaCleantied tomhsa_has_vjp_mat, its multi-head collapsemhsa_backward_collapseMH(a sum over heads of per-headsdpa_back_{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):
vjpMat_comp F G … .backward A dY = hF.backward A (hG.backward (F A) dY)biPathMat_has_vjp 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 dense_per_token_has_vjp_mat.backward (which ignores the saved
activation A, dense being affine — dense_has_vjp.backward _ dy = Mat.mulVec W dy).
Per-token LayerNorm input-VJP: rowLNBackFlat IS the flatten of the rowwise
layerNorm_per_token_has_vjp_mat.backward at the saved pre-LN activation A.
(layerNorm_has_vjp is definitionally bn_has_vjp, whose backward is
bn_grad_input; rowLNBackFlat is the rowwise bn_grad_input.)
Per-token GELU input-VJP: the flat gelu_has_vjp (N*D) backward IS the
flatten of the rowwise gelu_per_token_has_vjp_mat.backward at the saved
pre-GELU activation A (GELU is elementwise, so flat and rowwise agree).
The transformer MLP backward graph (reverse-order chain of
transformerMlp_has_vjp_mat = 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
MLP backward-graph faithfulness. The reverse-order chain denotes the
proven transformerMlp_has_vjp_mat.backward at the saved MLP input Y. The
two dense backs ignore the activation; GELU's reads dense1 Y.
⭐ Stated over an arbitrary incoming-cotangent SUBGRAPH ecot (any graph whose
den is the flattened cotangent), not over a bare Vec wrapped internally as
.operand "%dz". That is what lets this arm sit downstream of another graph;
the old statement is this one at ecot := .operand "%dz" (Mat.flatten dz).
The proven mhsa_has_vjp_mat witness is built via by rw [mhsa_layer_eq_compose]; exact vjpMat_comp …, 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 mhsa_has_vjp_mat by VJP determinism; the multi-head collapse below runs on it.
A clean HasVJPMat witness for the factored MHSA — the same vjpMat_comp
chain mhsa_has_vjp_mat's body uses, but stated for the explicit composition
Wo-dense ∘ colSlabApply mhsa_g ∘ qkv-dense so its .backward reduces by
rfl (no mhsa_layer_eq_compose transport in the way).
Equations
- One or more equations did not get rendered due to their size.
Instances For
The clean witness's backward IS mhsa_has_vjp_mat's backward (both VJPs of
mhsa_layer, tied by determinism).
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 sdpa_back_{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 sdpa_back_{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.
The proven MHSA VJP's backward at general heads IS the per-head collapse.
SDPA dQ-segment subgraph: matmulF(scaleF(softmaxRowBack(matmulF(dAtt, transposeF v))), k)
— denotes vitCotDQ. ss = saved pre-softmax scaled scores; k/v/dAtt saved.
Equations
- One or more equations did not get rendered due to their size.
Instances For
SDPA dV-segment subgraph — denotes vitCotDV. p = saved post-softmax weights.
Equations
- Proofs.StableHLO.sdpaBackVGraph Np1 D p e = (Proofs.StableHLO.SHlo.operand "%p" p).transposeF.matmulF e
Instances For
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 mhsa_has_vjp_mat.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 the_has_vjp_matto the…V…(vec-LN) versions.
The vec-LN per-token backward collapses to normalize-only-of-scaled. The
vec-LN VJP (+βv) ∘ layerScale γv ∘ LN(1,0) has, by vjp_comp, backward
LN(1,0).backward x (layerScale_has_vjp.backward _ (biasAdd.backward _ dy));
biasAdd.backward = id, layerScale_has_vjp.backward _ dy = (γv · * dy ·) = layerScale γv dy. Rowwise-lifted, this is the normalize-only (layerNorm at
γ=1, β=0) per-token backward fed the rowwise layerScale γv of the cotangent.
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 vec-LN MLP-sublayer non-trivial arm backward graph (transformerMlp ∘ LNᵥ₂;
outermost backward token = earliest forward op = LN₂). REUSES transformerMlpBackGraph
verbatim (the MLP body is LN-agnostic), after the vector-LN back fragment
lnRowBack(γ=1) ∘ rowScaleF γ2v. Y = LNᵥ₂ h is the saved MLP input.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Vec-LN MLP-sublayer inner-arm backward-graph faithfulness. Denotes the proven
(vjpMat_comp LNᵥ₂ transformerMlp).backward h ·. Y = LNᵥ₂ h.
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. The old
dz : Vec statement is this one at ecot := .operand "%dz" dz.
Equations
- Proofs.StableHLO.mlpSublayerVBackGraph ε γ2v Wfc1 bfc1 Wfc2 h Y ecot = (Proofs.StableHLO.mlpSublayerVInnerBackGraph ε γ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 transformerMlpSublayerV_has_vjp_mat 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 now 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 vjpMat_comp'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 vitForwardKV_has_vjp (ViTDepthK.lean) is a Vec → Vec
HasVJP built by three vjp_comp steps:
classifier_flat ∘ [finalLN] ∘ vitBodyKVFlat(k) ∘ patchEmbed_flat
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_to_hasVJP, so each
block-back faithfulness statement is stated through Mat.flatten/Mat.unflatten
(exactly the transformerBlockVBackGraphMH_faithful convention). Four stages:
- Stage 1 —
classifierBackGraph↔classifier_flat_has_vjp.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↔patchEmbed_flat_has_vjp.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
classifier_flat_has_vjp.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 — the head is the LAST layer of the forward and therefore the FIRST of the backward, so
this is the one place where a Vec cotangent was genuinely natural. Generalizing it anyway
is what lets the head be a CertLayer and the whole net be 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).
⚠ This used to bundle the classifier back inside it — it took dy : Vec nClasses and
called classifierBackGraph itself, so the final LN and the head were one indivisible
thing. They are now separate graphs composed at the call site, which is what lets each be
its own CertLayer. Bundling two stages into one node is the same mistake as wrapping a
cotangent as an operand: it works exactly until something needs to sit between them.
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 — transformerBlockV_has_vjp_mat.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 transformerBlockV_has_vjp_matP.
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
vitBodyKVFlat_has_vjp's vjp_comp chain.
Equations
- One or more equations did not get rendered due to their size.
- Proofs.StableHLO.vitBodyBackGraphKMHV ε 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
patchEmbed_flat_has_vjp.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.