ViT at the BATCHED index N := B — forward, backward and AdamW renders #
The ViT peer of ConvNeXtRenderB / ResNet34RenderB / MobileNetV2RenderB, and the reason it
exists is the same one: stochastic depth's mask is per-EXAMPLE, and in the per-example-indexed
render (ViTRender.lean) a node denotes ONE example — pretty B lifts it across the batch, so the
node cannot see j. vitTinyImagenetConfig sets dropPath 0.1 (24 sites, two per block).
ON ViT THE INDEX COLLISION IS IN THE SOURCE TEXT, NOT ONLY IN THE SEMANTICS. The per-example
renderer names its TOKEN axis N — denseRowF {N a c}, clsSliceF {N D}, headSliceF {N heads d}
— and a batched constructor names the BATCH N. Both are Nat, both appear multiplied, and the
swap type-checks in either direction. Every batched call below therefore passes the token count as
an EXPLICIT named argument ((N := 197), (tk := 196)) and never positionally; the batch is vbB
and appears only as pretty's argument and as (N := vbB) on the *B constructors.
The sharpest instance is clsSlice, which CONTRACTS (tk+1)*D → D. A render reading the batch as
the token axis takes (N+1)*D → D — it keeps ONE example and drops the rest — and at N = tk it
type-checks and agrees. den_batchOp_clsSlice_per_example is what pins those apart.
What this file is: the writer of every ViT artifact but one. The forward, the backward
traversal vitBackAllB and, through ViTRender.vitAdamTrainStepText, the AdamW/EMA train
steps.
The one exception is vit_train_step.mlir, the SGD-inline step, which ViTRender writes:
vitBackAllB has no fused-SGD arm (it emits the raw gradient only), and ViTStepTie's
200-parameter tie is stated at exactly those bytes.
The denotation differs from the per-example chain at the same bytes. See the CLS-token
emission below: this chain's denseBiasGradB (N := vbB) sums the batch inside den where the
per-example one writes (N := 1) and lets pretty B lift outside the AST.
Proofs.ViTFoldGB.clsGrad_denB (in ViTFoldGB) is the theorem stated at it.
The gate (lake build vit-fwd-b-tie): the committed artifacts are this chain's, so the gate
renders the independent PER-EXAMPLE chain (vitFwdRenderV, vitAdamTrainStepText) and
requires it to emit verified_mlir/vit_fwd.mlir and vit_adam_train_step.mlir byte for byte. That is available because every batched form was built to emit
its per-example peer's text byte-for-byte and tests/TestBatchedEmitTie.lean pins each
batched form individually — so the whole-net claim is the per-form claim composed, and when it fails that file
localises which form did it in one run.
The batched shapes #
Restated, not re-derived — ViTRender's own vEPS/vSCALE/vDEPTH are shared (they stopped
being private for this file), but the shape numbers are spelled here so that a drift between the
two renderers fails the BYTE tie loudly rather than propagating silently. That is the same choice
ConvNeXtRenderB made and for the same reason.
ViT-Tiny: patch 16×16/s16 over 224² ⇒ 196 patch tokens, 197 with CLS; D = 192 = 3 heads × 64;
MLP 768; depth 12.
The sites, one per per-example site in ViTRender.lean #
The depth-12 ViT-Tiny forward at the batched index. Node for node the same chain
vitFwd12 emits — patch embed (16×16/s16, 196 patches + CLS + pos) → 12 blocks → final
vector-LN → CLS slice → dense head.
Every node is a batchOp/*B form, so den is a batchMap/batchMapAux at N := vbB and
the batch is an index of the AST rather than a number only pretty knows. That is the entire
content of the move; the emitted text is unchanged, which the tie checks.
Equations
- One or more equations did not get rendered due to their size.
Instances For
@vit_fwd from the batched chain — the batched-index peer of vitFwdRenderV, same
signature (200 parameters at ViT-Tiny) and same %x. It writes every committed ViT forward:
vit_fwd.mlir, vit_drop_fwd.mlir, vitin_fwd.mlir, vitin_drop_fwd.mlir, vitsin_fwd.mlir,
vitsin_drop_fwd.mlir and vitbin_fwd.mlir. vit-fwd-b-tie renders @vit_fwd with the
per-example vitFwdRenderV and compares it with the committed vit_fwd.mlir byte for byte.
Equations
- One or more equations did not get rendered due to their size.
Instances For
The whole-net batched traversal — forward + cotangent + all 200 parameter gradients, the
batched peer of vitBackAll bs nClasses lrStr true (some …). Returns
(code, gradients-in-func-arg-order, softmaxSSA), the same shape, so the AdamW tail in
ViTRender.lean can consume either.
Equations
- One or more equations did not get rendered due to their size.
Instances For
The ViT-Tiny AdamW train step at the batched index. It is the SAME renderer the
per-example path uses — vitAdamTrainStepText with traversal pointed at vitBackAllB —
not a copy.
That is possible because the AdamW tail is entirely parameter-space: adamMNextF,
adamVNextF, adamWParamF, gradSumSqAccF and clipScaleF are indexed by the parameter's own
size and never see the batch. So "the AdamW tail at the batched index" is no work at all — the
batch is already factored out of it by the ops' own shapes, and the only thing that moves
is which traversal produced the gradients. ConvNeXt's AdamW tail is the same, and it
generalises: the optimizer is where the batch has already been summed away.
Equations
- One or more equations did not get rendered due to their size.