ViT-Tiny's step tie at the batched index, the un-fused gradient and the smoothed loss #

ViTStepTie.lean ties all 200 parameters of the SGD-inline vit_train_step.mlir: each fused θ − lr·g op dens to the certified step at the cotangent the emitted backward chain delivers, per example, at a hard label. This file is that statement re-pointed along the three axes ConvNeXtStepTieGB.lean moved for ConvNeXt, applied to ViT's per-example capstone.

Axis 1 — the gradient node. Every conjunct is at the raw gradient node (*GradB), which is what vit_adam_train_step.mlir and every f32 vitin_* artifact emit; the fused op appears only in the SGD-inline file. The optimizer update that consumes the node (AdamW, the wx/clip variants, EMA) is outside this statement. ViTFoldGB.lean is the fold each conjunct delegates to.

Precision is a flag on the statement. The bf16 artifacts (vitin_adamwxclipdropbf16 and its EMA, data-parallel, ε and exact-GELU twins) emit the bf16 per-token dense weight gradient (rowDenseWeightGradBBf16) at the six denses of every block, and — only if the render's bf16ConvW is set, which no shipped artifact does — the bf16 patch-embed weight gradient (Foundation.Bf16GradNodes). Both are stated on the renderers' switch: ViTFoldGB.RowDenseWTiedBAt bf16 and patchEmbedWeightGradBAt bf16 id … (StableHLO.PrecisionSwitch), false the f32 node, true the bf16 one. The capstone vit_net_tiedGB takes the renderer's two backward flags, bf16 and bf16ConvW, and passes bf16 && bf16ConvW to the embed as vitBackAllB does, so at bf16 := true, bf16ConvW := false it reaches the shipped bf16 artifacts' gradient nodes read over ℝ exactly as it reads the f32 ones (Bf16Erasure: at the identity rounding the bf16 kind denotes what its f32 peer does). The bias, LayerNorm, CLS, position and classifier nodes carry no flag because no render switches them.

Axis 2 — the loss. g is a binder, instantiated at smoothedLossCotGraphDiv — the six-op chain expe → softmaxDiv → subB → scaleB → addVB → shiftB → divConstB this render emits at the plain width N·K, at a GENERAL target arriving as %onehot. The fused file pins it to softmax − oneHot.

Axis 3 — the index. N is a binder. Every activation is a per-example lift of the prefix the fused file threads (patchEmbedFlat, the block forward, the final LN, clsSliceFlat) and every cotangent a per-example lift of the chain (vitCotTowerOutV, the block's input cotangent, the vitCot* family). The lift is exact for this net because no ViT op couples examples — LayerNorm, attention, GELU and the denses are all per-example, and the *B constructors' den arms say so. A block's lift is the INDEXED one (batchMapIdx / batchMapAuxIdx, Foundation.Batched.Indexed): with stochastic depth example n runs the block at its own mask entries. nC is a binder too (10 on Imagenette, 1000 on ImageNet).

Stochastic depth is a binder. sd : Option (Fin 12 → Vec N × Vec N) is the render's drop flag: none the drop-free artifacts' chain — every family constant, so the lifts are batchMap / batchMapAux and the statement is the drop-free one by rfl — and some the *drop* artifacts', the book's vitin_emadp128x4wxclipdropeps0000001erfbf16 among them: block k's two sites (vitSdA / vitSdM, the %dp<2k> / %dp<2k+1> inputs) scale the out-projection's and fc2's outputs by example n's entries (BlockParamsV.fwdOD), the out-projection and fc2 weight and bias nodes read the dropped cotangent (dropPathOpt, where ViTRenderB emits dropPathB on it), and both skip fan-ins read the raw one (BlockParamsV.cotInD, ViTDropBlock).

The CLS-token conjunct. The CLS token is one shared [192] vector; its gradient is the sum of every example's CLS-row cotangent. The fused file's vit_cls_den is at denseBiasSgdB (N := 1) — "sum one thing", correct there because pretty B performed the batch lift outside the AST. vitEmbedTiedGB's third conjunct is ViTFoldGB.clsGrad_denB at the batched node, denseBiasGradB (N := N) on batchMap N clsSliceFlat of the embed cotangent, with the batch sum inside den.

Relation to the per-example file #

Every save, every chain cotangent and every Jacobian witness is ViTStepTie.lean's, lifted; every conjunct's proof is one ViTFoldGB.*_den lemma. The per-example saves and internal cotangents are repackaged as functions of the block INPUT and its two sites (blkSaves, cAtt … cM1) so that the indexed lift has something to lift — the let chains of vitBlockTiedAtMHV and vitBlockCotInAtMHV with the sites in. ViT has no *BackBatchedGraph_faithful family and needs none here: the batchMap lift is the batched statement, as it was for ConvNeXt.

Conventions carried unchanged from the fused file: the VECTOR LayerNorm (γ β : Vec D) at all 25 sites, 3 heads × d_head 64, depth 12, D 192, MLP 768, 16×16 patches, GELU (no kink — no smoothness hypothesis anywhere). Stated at ViT-Tiny's literal dims; S and B are other nets.

Scope. One replica: in vitin_adamdp128x4* (four replicas of 128) every gradient node feeds allReduceMeanF, and DataParallel.Node composes the per-replica statement with the replica mean.

Per-example saves and internal cotangents as functions of a block's INPUT #

vitBlockTiedMHV takes the nine saved activations as arguments; batchMapAux lifts a function of (one saved value, one input), so the batched block tie needs each save as a function of xin and each internal cotangent as a function of (xin, dyOut) — exactly the let chains of ViTTie.vitBlockTiedAtMHV and vitBlockCotInAtMHV, packaged.

The multi-head block — all 16 gradient nodes, batched #

Final LN, classifier and patch embedding — batched #

Every cotangent the capstone threads is a certified VJP backward #

The capstone below threads two per-example constructors: the head's vitCotTowerOutV (batchMapAux N) and each block's cotInD at its sites (batchMapAuxIdx N). Their per-example ties to the certified VJPs are ViTStepTie's vitCotTowerOutV_eq_vjp and ViTDropBlock's cotInD_eq_vjp; the two lemmas below lift them over the batch.

The whole-net capstone — all 200 params through the REAL batched forward + composed cotangent #

The fused file's thread, lifted: ib1 is batchMap N of the patch embedding, ib_{k+1} is batchMapIdx N of the multi-head block forward at block k's sites, the final LN / CLS slice / dense head are batchMap N of theirs, g is the smoothed loss cotangent at a general target, and every cotangent is the lift of the per-example chain — vitCotTowerOutV at the top, then twelve cotInD attention-residual fan-ins down to the embed-output cotangent.