ViT — the attention-block cotangent chain #

TokenParamGrad.lean certifies each ViT param output for any cotangent dy at that site's output. This file defines the cotangent the actual backward chain delivers at each site — the ViT analogue of ConvNeXtChainClose — and the step ties (ViTStepTie, ViTStepTieGB) feed those cotangents to those bridges at the real forward. The definitions are per example; everything in a ViT is per-example separable, so the batched tie lifts them with batchMapAux.

The chain composes the rendered backward denotations — exactly the render's backward tokens: per-token dense input-VJP (denseRowBack's denotation rowDenseBackFlat = rowwise dX = W·dy), the GELU mask at either form (dy ⊙ gf.scalarDeriv at the saved pre-GELU), the rowwise scalar-LN input-VJP (lnRowBack's denotation rowLNBackFlat = rowwise bnGradInput), the row-softmax backward (softmaxRowBack's denotation rowSoftmaxBackFlat, recomputing the weights from the saved pre-softmax scores), and the SDPA matmuls spelled with the forward matmulF/transposeF on cotangents (matMulFlat/transposeFlat):

block: bout = h + fc2(gelu(fc1(LN₂ h))), h = x + Wo·SDPA(Wq·LN₁x, Wk·LN₁x, Wv·LN₁x)

The MLP residual passes dyOut straight to the fc2 output AND down the LN₂ branch (the cotangent at h is dyOut + LN₂-back(…)); the attention residual likewise (the block input's is cotH + LN₁-back(…)). The new wrinkle vs all prior nets is the three-way fan-in at LN₁'s output — the Q/K/V dense-backs all read from LN₁ x, so their three cotangents SUM (vitCotLn1), the biPath fan-in at width 3.

The substantive new ties (vitCotD{Q,K,V}_eq_sdpaBack{Q,K,V}): at the pinned saved activations (pre-softmax scores = the scaled Q·Kᵀ, post-softmax weights = sdpaWeights), the matmul-spelled chain segments ARE the proven closed forms sdpaBack{Q,K,V} (Attention.lean) — dP = dO·Vᵀ → softmax-back → ·1/√d → dQ = dS·K / dK = dSᵀ·Q / dV = Pᵀ·dO, flattened. So the rendered attention backward is pinned to the audited SDPA backward suite.