The gradient-node fold for ViT-Tiny at the batched index #
ViTFoldG.lean holds the two per-example lemmas specific to this net (posEmbedGrad_den,
clsGrad_den); the other per-example peers are the fused ViTFold lemmas. This file is the batched
peer, at ViTRenderB.vitBackAllB's constructors, from which the committed ViT artifacts render.
Every batched form emits its per-example peer's text, and tests/TestBatchedEmitTie.lean pins each
one individually, so this file is not about different bytes; it is about the AST those bytes are
pretty of.
The CLS token.
The CLS token is ONE shared [192] vector, so its gradient is the sum of every example's CLS-row
cotangent. The per-example render emits denseBiasGradB (N := 1) — "sum one thing", correct there
because pretty B performed the batch lift OUTSIDE the AST — where the batched one emits
denseBiasGradB (N := vbB) and the sum is inside den. clsGrad_denB below is the
statement the per-example clsGrad_den could not make. Same emitted text either way, which is
why the byte tie cannot see it and den_rowDenseBiasGradB_at_one exists to argue the point.
The op table of the committed f32 ViT train steps #
The vector-LN and classifier-head lemmas, and the vector-LN tie clauses, are in GradNodesB under
this namespace, because ConvNeXt's batched tier uses them too.
| emitted node | lemma | per-example peer (fused ViTFold op unless noted) |
|---|---|---|
veclnGammaGradB (25 LN γ: LN1/LN2 × 12 + final) | veclnGammaGradB_den (GradNodesB) | SgdNode.veclnGammaSgd_den |
rowDenseBiasGradB (25 LN β) | rowDenseBiasGradB_den_lnbeta (GradNodesB) | SgdNode.rowDenseBiasSgd_den_lnbeta |
rowDenseWeightGradB (Wq/Wk/Wv/Wo/Wfc1/Wfc2 × 12) | rowDenseWeightGradB_den | ViTFold.rowDenseWeightSgd_den |
rowDenseBiasGradB (bq/bk/bv/bo/bfc1/bfc2 × 12) | rowDenseBiasGradB_den | ViTFold.rowDenseBiasSgd_den |
patchEmbedWeightGradB / patchEmbedBiasGradB | patchEmbedWeightGradB_den / patchEmbedBiasGradB_den | ViTFold.patchEmbedWeightSgd_den / patchEmbedBiasSgd_den |
posEmbedGradB | posEmbedGradB_den | ViTFoldG.posEmbedGrad_den (un-fused) |
denseBiasGradB at N = B (the CLS token) | clsGrad_denB | ViTFoldG.clsGrad_den (un-fused), at N = 1 |
weightGradB / biasGradB (the classifier) | headWGradB_den / headBGradB_den (GradNodesB) | ViTFold.headW_den / headB_den |
No new mathematics: every proof is Finset.sum_congr rfl over the batch and then the
per-example bridge at batchSlice n. That is GradNodeB.denseWGradB_den's shape, and
it is available because each batched den arm is literally ∑_batch of the per-example one — the
constructors were written that way (StableHLO.Basic's own comment on veclnGammaGradB: "TWO-LEVEL:
the outer Σ_n is the batch, the inner Σ_r the rows within one example").
One lemma per f32 node kind. Every lemma is ∀ cot: the f32 AdamW, wx/clip and EMA
artifacts all emit these *GradB kinds, and the optimizer update that consumes the node is outside
these lemmas. The artifact whose accuracy the book quotes, vitin_emadp128x4wxclipdropbf16, is a
bf16 one, so its weight gradients go through the bf16 kinds listed below.
The LayerNorm form is the vector one (γ β : Vec D), which is what vitForwardKV runs.
Scope #
- Every lemma is
∀ cot. The tie at these nodes, with the cotangents the emitted backward chain delivers, isViTTieGB.vit_net_tiedGB; the per-exampleViTStepTie.leanstays at the SGD-inlinevit_train_step.mlir. - The
*bf16artifacts emitrowDenseWeightGradBBf16/patchEmbedWeightGradBBf16, their own kinds;Foundation/Bf16GradNodes.leanfolds them (the row-dense one keeps its f32 result). vitin_*dp128x4*is four replicas: the all-reduce is its ownallReduceMeanFnode after each gradient node (DataParallel.Node), so these lemmas are about the per-replica gradient node it averages.
Batched per-token dense weight GRADIENT denotes the certified Σ_n Σ_tokens x⊗dy. All six
denses in each of the 12 blocks. Note: the emitted dot_general contracts batch AND token in
ONE op; the two sums here are that contraction read apart.
Batched per-token dense bias GRADIENT denotes the certified Σ_n Σ_tokens dy.
bq/bk/bv/bo/bfc1/bfc2.
Batched patch-embed conv weight GRADIENT denotes the certified Σ_n patchify weight
gradient. Note: the emitted convolution contracts the batch axis itself, so the outer sum
here is inside one op rather than across N of them.
Batched patch-embed conv bias GRADIENT denotes the certified Σ_n bias gradient
(Σ_patches dy, the CLS row excluded — the inner p.succ).
Batched positional-embed GRADIENT denotes the certified Σ_n gradient — the summed
cotangent, since the positional table is added to every token and its Jacobian is the
identity. Note: this batch sum is invisible at
N = 1: a render that dropped it type-checks and emits the same bytes.
The CLS-token GRADIENT, and this is the one statement the per-example fold could not make.
The render slices row 0 of every example's embed cotangent (clsSlice, lifted by batchOp)
and reduces the result as an [N, D] batch, so the op is denseBiasGradB at N = B and its
den sums over the batch — which is what a shared [192] parameter's gradient IS.
ViTFoldG.clsGrad_den is the same theorem at N = 1, where the batch lift lived in
pretty B outside the AST. The bytes are identical (biasGrad's emitted reduce takes the B
axis either way) and the functions are not; den_rowDenseBiasGradB_at_one is the general form
of the trap.
Note: stated at the committed ViT-Tiny dims rather than generically, for ViTFoldG.clsGrad_den's
reason: the CLS operand's type is Vec (N * (1 * D)), which reduces to Vec (N * D) only at a
literal D.
A batched per-token dense weight gradient node, tied (rowDenseWeightGradB_den).
Equations
- One or more equations did not get rendered due to their size.
Instances For
A batched per-token dense bias gradient node, tied (rowDenseBiasGradB_den).
Equations
- One or more equations did not get rendered due to their size.
Instances For
Each clause holds, every argument implicit (read off the goal by a step tie's constructor).
RowDenseWTiedB on the renderer's switch (SHlo.rowDenseWeightGradBAt bf16 id …): the f32
node at false, its bf16 peer at true — the six per-block dense weight gradients of a bf16
ViT step. ViT is the only net that emits this kind, so the predicate lives here rather than in
GradNodesBAt; it holds at either value through Bf16Fold.den_rowDenseWeightGradBAt_id.
Equations
- One or more equations did not get rendered due to their size.