Parameter-gradient bridges for the token ops — per-token dense, patch embedding, position, CLS #
The per-parameter gradient bridges ViT's fold (ViTFold) delegates to, and that the bf16
gradient nodes (Bf16GradNodes) reuse; the ops themselves are Attention's.
Generic in the cotangent dy the backward chain delivers at each site's output
(pinning that cotangent to the actual attention chain is the optional Item D), batch-1 —
everything in a ViT is per-example separable (the EfficientNet contrast).
| family (render SSA) | forward fn | certified by |
|---|---|---|
| Wq/Wk/Wv/Wo, Wfc1/Wfc2 + biases | per-token dense (rowwise) | vit_render_rowdense{W,b}_certified (new family): dW = Σ_tokens xᵣ ⊗ dyᵣ, db = Σ_tokens dyᵣ — the M2 outer-product bridge row-lifted |
classifier Wcls/bcls | dense on the CLS row | M2 weight/bias_grad_bridge (reuse — single-vector dense) |
| LN γ/β (vector, per-token) | rowwise vector LayerNorm | vit_vecln{Gamma,Beta}_grad_bridge (LayerNorm) |
pos_embed | additive (patchEmbed_flat) | vit_render_pos_certified: the pos-Jacobian is the identity ⇒ dPos = dy |
cls_token | row-0 scatter (patchEmbed_flat) | vit_render_cls_certified: masked-gather Jacobian ⇒ dCls = dy row-0 slice |
patch conv Wp/bp | stride-P conv (patchEmbed_flat) | vit_render_patch{W,b}_certified: kernel-linear w/ constant guarded reads ⇒ dWp = Σ_p read·dy_(p+1), dbp = Σ_p dy_(p+1) (CLS row excluded) |
| attention internals (softmax, scale) | — | no parameters |
One genuinely-new bridge family (everything else is reuse or a reindex):
- Per-token dense W/b — every row of
[N,a]through the sameW : [a,c](+b). The W-Jacobian of the flattened rowwise dense is block-sparse —∂y_(r,k)/∂W_(i,j) = X_(r,i)·δ_(k,j)— so the rendered per-token outer-product reducedW_(i,j) = Σ_r X_(r,i)·dY_(r,j)(onedot_generalcontracting the token axis) is the certified contraction. Covers Wq/Wk/Wv/Wo, Wfc1/Wfc2 and their biases at every block.
The classifier head (dense on the CLS vector) is VERBATIM M2 weight/bias_grad_bridge
reuse at [D, nClasses]. The patch-embed conv Wp/bp (§ E) closes over
patchEmbed_flat directly — the kernel is the VARIABLE and the pad-guarded image reads
are CONSTANT coefficients (the mirror of the input-grad case), so the forward is affine in
the kernel and pdiv_of_affine applies with the CLS row masked out. 3-axiom clean by
construction.
Jacobian of the per-token dense w.r.t. the (flattened) shared weight —
∂y_(r,k)/∂W_(i,j') = X_(r,i)·δ_(k,j'). The row-lift of pdiv_dense_W.
The rendered per-token dense weight gradient: the token-axis-contracted
outer product dW_(i,j) = Σ_r X_(r,i)·dY_(r,j) (one dot_general contracting
the token axis — the row-lift of dense_weight_grad = x ⊗ dy).
Equations
- Proofs.rowDense_weight_grad X dY i j = ∑ r : Fin N, X r i * dY r j
Instances For
The rendered per-token dense bias gradient: the token-axis reduce
db_j = Σ_r dY_(r,j).
Equations
- Proofs.rowDense_bias_grad dY j = ∑ r : Fin N, dY r j
Instances For
Per-token dense W-gradient bridge. The rendered token-contracted outer
product equals the certified Jacobian of the rowwise dense (as a function of
the flattened shared W) contracted with the cotangent.
Per-token dense b-gradient bridge. The rendered token-axis reduce equals the certified rowwise-dense ∂/∂b contraction.
Per-token dense W output, certified. Wⁿ = W − lr·(Σ_tokens xᵣ ⊗ dyᵣ) denotes
W − lr·(certified ∂(rowwise dense)/∂W · cotangent). Covers Wq/Wk/Wv/Wo and
Wfc1/Wfc2 at every block of the representative ViT (each at its own [a,c]).
Per-token dense b output, certified. bⁿ = b − lr·(Σ_tokens dyᵣ) denotes the
certified rowwise-dense ∂/∂b contraction. Covers all six per-block biases.
pos-embed output, certified. The pos Jacobian is the identity, so the
rendered dPos = dy (the cotangent itself, batch-summed by the batched
render) is the certified contraction.
The rendered CLS-token gradient: the row-0 slice of the patch-embed
output cotangent (clsSliceF's shape, applied to the embed cotangent).
Equations
- Proofs.cls_token_grad dy i = dy (finProdFinEquiv (0, i))
Instances For
CLS-token output, certified. clsⁿ = cls − lr·(row-0 slice of the embed cotangent) denotes the certified ∂(patchEmbed)/∂cls contraction.
The pad-guarded patch read of patchEmbed_flat, named: input pixel
(c, h'·P + kh, w'·P + kw) of patch p (row-major patch grid of width
W/P), zero out of range. Constant in the kernel.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Jacobian of patchEmbed_flat w.r.t. the (flattened) patch kernel —
∂y_(n,dd)/∂W_(d,c,kh,kw) = [n ≠ 0]·δ_(dd,d)·read(c,kh,kw, patch n−1).
The rendered patch-kernel gradient: for each tap (d,c,kh,kw), the
patch-grid reduce Σ_p read(c,kh,kw,p)·dy_(p+1,d) — the "dilate dy /
valid conv" weight grad, with the CLS row (token 0) excluded.
Equations
- Proofs.patchEmbed_weight_grad ic H W P N D img dy d c kh kw = ∑ n : Fin N, Proofs.patchRead ic H W P img c kh kw ↑n * dy (finProdFinEquiv (n.succ, d))
Instances For
The rendered patch bias gradient: db_d = Σ_p dy_(p+1,d) (the CLS row
excluded — token 0 carries no conv bias).
Equations
- Proofs.patchEmbed_bias_grad N D dy d = ∑ n : Fin N, dy (finProdFinEquiv (n.succ, d))
Instances For
Patch-kernel gradient bridge. The rendered patch-grid reduce equals the certified ∂(patchEmbed)/∂W contraction.
Patch-kernel output, certified. Wpⁿ = Wp − lr·(patch-grid reduce)
denotes the certified ∂(patchEmbed)/∂Wp contraction.
Jacobian of patchEmbed_flat w.r.t. the patch bias — the row-masked
gather ∂y_(n,k)/∂bc_i = [n ≠ 0]·δ_(i,k) (token 0 is the CLS row).
Patch bias gradient bridge. The rendered CLS-row-excluded reduce equals the certified ∂(patchEmbed)/∂bc contraction.