Dependencies

assume:

1. \(\varepsilon {\gt} 0\) [h\(\varepsilon \)]

prove: \(\operatorname {bnIstdBroadcast}= x \mapsto 1/\sqrt{\sigma ^2(x) + \varepsilon }\) is \(\mathsf{Differentiable}\). This is the sqrt/recip smoothness that the product rule needs inside the normalize Jacobian.

Concrete \(\sum _{c,kh,kw}\) cross-correlation with SAME padding. Codegen emits stablehlo.convolution; proofs reason about it via the explicit definition.

Sum the cotangent over spatial dims per channel. prove: \(\mathsf{HasVJP}\, \bigl(b \mapsto \mathrm{flatten} (\mathrm{conv2d}(W, b)\, x)\bigr)\), with backward \(db_o = \sum _{h_i, w_i} dy_{\mathrm{flat}(o, h_i, w_i)}\).

\(\mathsf{HasVJP3}\, \bigl(\mathrm{conv2d}(W, b)\bigr)\), with the input backward

No hypotheses on the input or the kernel: conv2d is affine in its input. For odd \(k_H\), \(k_W\) this sum is the reversed-kernel, channel-swapped convolution of \(dy\) against \(W\) that the emitted backward runs (conv_back_bridge, §3.4); at an even kernel the reversed-kernel convolution is the adjoint of a convolution shifted by one pixel; the emitter pads such a kernel asymmetrically, which is the kernel zero-extended to the next odd size with symmetric padding, and the ties are stated through that zero-extension (padOdd, Chapter 8).

The transpose-trick formula. As with the dense weight Jacobian (Theorem 18), differentiate in the flattened kernel: let \(F(v) := \mathrm{flatten}\bigl(\mathrm{conv2d} (\mathrm{unflatten}\, v,\, b)\, x\bigr)\) for \(v \in \mathbb {R}^{o_c \cdot i_c \cdot k_H \cdot k_W}\). prove: \(\mathsf{HasVJP}\, F\), with backward \(dW_{o, c, k_h, k_w} = \sum _{h_i, w_i} [\text{pad-valid}]\; x_{c,\, h_i + k_h - p_H,\, w_i + k_w - p_W} \cdot dy_{o, h_i, w_i}\) — a convolution of the saved input against the upstream gradient.

prove: \(\mathsf{HasVJP}\, \bigl(b \mapsto \mathrm{flatten} (\mathrm{depthwiseConv2d}(W, b)\, x)\bigr)\), with backward \(db_c = \sum _{h_i, w_i} dy_{\mathrm{flat}(c, h_i, w_i)}\).

Concrete per-channel cross-correlation.

prove: \(\mathsf{HasVJP3}\, \bigl(\mathrm{depthwiseConv2d}(W, b)\bigr)\), with the per-channel backward \(dx_{c, h_i, w_i} = \sum _{h_o, w_o} [\text{pad-valid}]\; W_{c,\, \hat{k}_h,\, \hat{k}_w} \cdot dy_{c, h_o, w_o}\), \(\hat{k}_h = h_i + p_H - h_o\), \(\hat{k}_w = w_i + p_W - w_o\): Theorem 35’s sum without the channel sum. No hypotheses. For odd kernels this is the per-channel reversed-kernel convolution \(\sum _{k_h, k_w} W_{c,\, k_H - 1 - k_h,\, k_W - 1 - k_w} \cdot dy_{c,\, h + k_h - p_H,\, w + k_w - p_W}\) that the emitted backward runs (depthwiseFlatBack_eq_vjp_backward); every depthwise kernel in this book is odd — \(3\times 3\), \(5\times 5\) or ConvNeXt’s \(7\times 7\).

DepthwiseKernel is definitionally \(\mathsf{Tensor3}\), so \(\mathsf{HasVJP3}\) applies to the kernel directly — no flatten bijection needed, unlike the conv2d weight VJP (Theorem 36). prove: \(\mathsf{HasVJP3}\, \bigl(W \mapsto \mathrm{depthwiseConv2d}(W, b)\, x\bigr)\), with the per-channel transpose-trick backward \(dW_{c, k_h, k_w} = \sum _{h_o, w_o} [\text{pad-valid}]\; x_{c,\, k_h + h_o - p_H,\, k_w + w_o - p_W} \cdot dy_{c, h_o, w_o}\).

GELU as Hendrycks and Gimpel define it:

with \(\Phi \) the standard-normal CDF and \(\varphi \) its density. This is PyTorch’s nn.GELU and jax.nn.gelu(approximate=False), the form the ConvNeXt and DeiT recipes train, and what the ImageNet pairs of this chapter and the next compute on both paths. The CDF is written as a closed form over the interval integral, so that \(\Phi ' = \varphi \) is the fundamental theorem of calculus at a continuous integrand (hasDerivAt_gaussPhi) rather than a fact about a measure.

Mathlib’s \(\mathrm{deriv}\, \mathrm{geluErfScalar}\), as the tanh form’s is. The closed form is the theorem geluErfScalarDeriv_eq: \(\mathrm{geluErfScalar}'(x) = \Phi (x) + x\, \varphi (x)\), the product rule on \(x \cdot \Phi (x)\) with \(\Phi ' = \varphi \).

The tanh approximation, jax.nn.gelu’s default and the form every Imagenette-tier run in this chapter trained: \(\mathrm{geluScalar}(x) = 0.5\, x\, \bigl(1 + \tanh (\sqrt{2/\pi }\, (x + 0.044715\, x^3))\bigr)\). The exact \(x \cdot \Phi (x)\) is Definition 105; the codegen emits either, and Theorem 111 is the VJP of both.

Defined as Mathlib’s \(\mathrm{deriv}\, \mathrm{geluScalar}\) — so the connection to the forward is automatic, not asserted. The closed form (what the verified geluBack emitter renders) is the separate theorem geluScalarDeriv_eq, proved by assembling HasDerivAt for the inner polynomial, \(\tanh \) (via \(\tanh = \sinh /\cosh \) and the quotient rule), and the outer product.

Concrete four-way max over 2x2 windows.

The max pool’s VJP record, with the \(\operatorname {pdiv}\)-derived backward: at a point where every window has a strict maximum, each window’s cotangent goes to that cell; at a point where any window ties at its maximum the map is not differentiable, so the backward is the zero tensor in every window, strict ones included. The emitted backward routes a tied window to its first maximal cell.

prove: \(\mathsf{HasVJPMat}\, (\mathrm{mhsaLayer})\).

\(\mathsf{Differentiable}\) sibling of Theorem 137.

The MLP’s VJP record as one function, with the \(\operatorname {pdiv}\)-derived backward. The composed per-layer backward of Theorem 29 agrees with it at smooth inputs: any two pointwise witnesses have the same backward.

\(\mathsf{Differentiable}\) sibling of Theorem 139.

The concrete def unfolds to conv-with-stride + CLS prepend + positional embed. prove: \(\mathsf{HasVJP}\, (\mathrm{patchEmbedFlat})\), with backward the deconvolution formula: a sum over patches with reconstructed kernel offsets, the CLS row contributing nothing to the image gradient.

The partial derivative function. For \(f : \mathbb {R}^{m} \to \mathbb {R}^{n}\), \(\operatorname {pdiv}\, f\, x\, i\, j\) is the \((i, j)\)-entry of the Jacobian at \(x\), defined as \(\operatorname {fderiv}_{\mathbb {R}}\, f\, x\, (\mathbf{e}_i)\, j\) — the \(j\)-th coordinate of Mathlib’s Fréchet derivative applied to the \(i\)-th standard basis vector.

assume:

1. \(f, g : \mathbb {R}^{m} \to \mathbb {R}^{n}\) are both differentiable at \(x\) [hf, hg]

prove: for all \(i, j\):

For \(\mathrm{bnAffine}(\gamma , \beta )\, v = \lambda i.\; \gamma v_i + \beta \): prove: \(\operatorname {pdiv}\bigl(\mathrm{bnAffine}(\gamma , \beta )\bigr)\, v\, i\, j = \gamma \, \delta _{ij}\).

For \(\mathrm{bnCentered}\, x = \lambda j.\; x_j - \mu (x)\), where \(\mu (x) = \frac{1}{n}\sum _s x_s\): prove: \(\operatorname {pdiv}(\mathrm{bnCentered})\, x\, i\, j = \delta _{ij} - 1/n\).

assume:

1. \(\varepsilon {\gt} 0\) [h\(\varepsilon \)]

prove: writing \(s := \mathrm{istd}(x, \varepsilon ) = 1/\sqrt{\sigma ^2(x) + \varepsilon }\):

assume:

1. \(f : \mathbb {R}^{m} \to \mathbb {R}^{n}\) is differentiable at \(x\) [hf]

2. \(g : \mathbb {R}^{n} \to \mathbb {R}^{p}\) is differentiable at \(f(x)\) [hg]

prove: for all \(i, k\):

\(\operatorname {pdiv}(\lambda \_ .\; c)\, x\, i\, j = 0\) for any constant \(c \in \mathbb {R}^{n}\).

For the dense layer \(\mathrm{dense}(W, b)\, x = \lambda j.\; \bigl(\sum _i x_i\, W_{ij}\bigr) + b_j\):

No hypotheses: dense is affine, so the differentiability obligations are discharged inside the proof rather than assumed.

The symmetric counterpart of Theorem 17, differentiating in \(W\) instead of \(x\). Since \(\operatorname {pdiv}\) works on vectors, view the layer as a function of the flattened weights: for \(v \in \mathbb {R}^{m \cdot n}\), let \(F(v) := \mathrm{dense}(\mathrm{unflatten}\, v,\, b)\, x\), and write \(\varphi \) for the index bijection \((i, j) \leftrightarrow \varphi (i, j)\) (finProdFinEquiv). prove: for all \(i, j', j\):

prove: \(\operatorname {pdiv}(\mathrm{gelu})\, x\, i\, j = \delta _{ij} \cdot \mathrm{geluScalar}'(x_i)\) — the diagonal activation Jacobian.

prove: \(\operatorname {pdiv}(\mathrm{geluErf})\, x\, i\, j = \delta _{ij} \cdot \mathrm{geluErfScalar}'(x_i)\).

\(\operatorname {pdiv}(\mathrm{id})\, x\, i\, j = \delta _{ij}\).

assume:

1. \(f, g : \mathbb {R}^{m} \to \mathbb {R}^{n}\) are both differentiable at \(x\) [hf, hg]

prove: for all \(i, j\):

where \((f \odot g)\, y\, k = f\, y\, k \cdot g\, y\, k\) is the elementwise product.

For \(\sigma : \mathrm{Fin}\, b \to \mathrm{Fin}\, a\),

Covers permutations, reshapes, slicing. Generalizes pdiv_id (\(\sigma = \mathrm{id}\)). A gather is a continuous linear map (reindexCLM), so the argument is the affine rule’s (Theorem 9): the map is its own derivative, evaluated at \(\mathbf{e}_i\).

assume:

1. \(x\) is a smooth point of \(\mathrm{ReLU}\): every coordinate \(x_k \neq 0\) [h_smooth]

prove: for all \(i, j\):

where \([P]\) is the Iverson bracket (\(1\) if \(P\) holds, else \(0\)).

Writing \(p := \mathrm{softmax}(z)\): prove: \(\operatorname {pdiv}(\mathrm{softmax})\, z\, i\, j = p_j\, (\delta _{ij} - p_i)\).

assume:

1. \(g : \mathbb {R}^{n} \to \mathbb {R}^{p}\) is differentiable everywhere [h_g_diff]

prove: applying \(g\) to each row of a matrix has block-diagonal Jacobian: \(\operatorname {pdivMat}(\text{rowwise } g)\, A\, (i,j)\, (k,l) = [i = k] \cdot \operatorname {pdiv}g\, (A_i)\, j\, l\).

prove: the flattened \(\mathrm{rowSoftmax}\) is \(\mathsf{Differentiable}\).

Write \(p := \mathrm{softmax}(z)\) and \(\mathrm{CE}(z, \ell ) := -\log p_\ell \) (viewed as \(\mathbb {R}^{1}\)-valued so \(\operatorname {pdiv}\) applies; we read off the only output coordinate). prove: for all \(j\):

\(\mathrm{GeluForm}\) has two values, tanh and erf; \(\mathrm{scalar}\) sends them to Definitions 102 and 105, and \(\mathrm{map}\) applies the chosen scalar componentwise, which at each value is the tanh \(\mathrm{gelu}\) and \(\mathrm{geluErf}\) by rfl. Every ConvNeXt and ViT definition and theorem from here on takes a \(\mathrm{GeluForm}\), so the renders that trained the tanh form and the ones that train the exact form are instances of one statement, and no proof outside the GELU’s own three files cases on which. The StableHLO GELU ops carry the same value, and the exact one emits through chlo.erfc (Theorem 112).

For \(G : \mathbb {R}^{m} \to \mathbb {R}^{1}\) and \(x, dy \in \mathbb {R}^{m}\), \(\texttt{HasGradAt}\, G\, x\, dy\) says \(G\) is differentiable at \(x\) and \(\operatorname {pdiv}\, G\, x\, j\, 0 = dy_j\) for every \(j\). The loss read at any activation of a network is such a \(G\), and the chain’s cotangent there is its \(dy\). The differentiability is part of the claim: \(\operatorname {pdiv}\) alone is \(0\) wherever \(G\) is not differentiable.

For \(f : \mathbb {R}^{m} \to \mathbb {R}^{n}\), \(\mathsf{HasVJP}\, f\) bundles a backward function \(B : \mathbb {R}^{m} \to \mathbb {R}^{n} \to \mathbb {R}^{m}\) with its correctness claim: for all \(x\), \(dy\), \(i\),

Exhibiting a \(\mathsf{HasVJP}\, f\) is exactly the statement “this backward function computes the vector–Jacobian product of \(f\).” Some witness exists for every \(f\): take \(B\) to be the contraction itself (HasVJP.canonical). So a witness says something only through what its \(B\) is: a formula such as a layer’s closed-form backward, the per-layer backwards composed in reverse order, or a rendered chain proved equal to one of those.

The rank-3 analogue of Definition 10. For \(f : \mathrm{Tensor3} \to \mathrm{Tensor3}\), define \(\operatorname {pdiv}_3 f\, x\, (c_i, h_i, w_i)\, (c_o, h_o, w_o)\) as \(\operatorname {pdiv}\) of the flattened map \(v \mapsto \mathrm{flatten}\bigl(f(\mathrm{unflatten}\, v)\bigr)\) with both indices decoded through the \((c, h, w) \leftrightarrow \text{flat}\) bijection. Then \(\mathsf{HasVJP3}\, f\) bundles a backward function \(B\) with: for all \(x\), \(dy\), \((c_i, h_i, w_i)\),

The rank-2 analogue of Definition 10. Define \(\operatorname {pdivMat}f\, A\, (i,j)\, (k,l)\) as \(\operatorname {pdiv}\) of the row-major flattened map \(v \mapsto \mathrm{flatten}\bigl(f(\mathrm{unflatten}\, v)\bigr)\) with both index pairs encoded through the \((i, j) \leftrightarrow \text{flat}\) bijection. Then \(\mathsf{HasVJPMat}\, f\) bundles a backward function \(B\) with: for all \(A\), \(dY\), \((i, j)\),

DeepLab v3+’s marquee module (Chen et al. 2018). Five parallel branches emitting oc channels each: (1) 1\(\times \)1 conv, (2–4) 3\(\times \)3 atrous convs at dilation rates 6 / 12 / 18, (5) global avg-pool + 1\(\times \)1 conv + bilinear upsample. Concatenate, then a 1\(\times \)1 fusion conv back to oc. All branches include BN + ReLU. Atrous rates widen the effective receptive field without changing param count; the pool branch supplies image-level context. Signature: asppModule (ic oc : Nat).

CSP (Wang et al. 2019), used by YOLOv4 onward. Splits input into two halves, processes one half through a stack of residual blocks, then concatenates with the untouched half and 1\(\times \)1-projects to oc. The specific inner block varies across YOLO versions (C3 in v5, C2f in v8, C3k2 in v11); this primitive approximates all three at the same abstraction level. Signature: cspBlock (ic oc nBlocks : Nat).

YOLOv3’s Darknet-53 residual stack (Redmon & Farhadi 2018). nBlocks residual blocks at fixed channels, each being 1\(\times \)1 conv (\(c \to c/2\)) \(+\) 3\(\times \)3 conv (\(c/2 \to c\)) \(+\) residual add. Lighter than a standard ResNet bottleneck; heavier than a ResNet-18 basic block. Signature: darknetBlock (channels nBlocks : Nat).

DenseNet’s bundled dense block (Huang et al. 2017). nLayers BN-ReLU-1\(\times \)1-BN-ReLU-3\(\times \)3 sub-stacks, each adding growthRate channels to the running concatenation of preceding outputs. Input has ic channels; output has ic + nLayers \(\cdot \) growthRate. Bundled because each sub-layer reads a growing concat of all preceding sub-layers, which doesn’t fit a linear NetSpec at sub-layer granularity. Signature: denseBlock (ic growthRate nLayers : Nat).

Per-query class head (linear dim \(\to \) nClasses+1 with the “no object” slot) + box head (3-layer MLP to 4 scalars (cx, cy, w, h)). Signature: detrHeads (dim nClasses : Nat).

Dual-representation (MSA + pair) joint update: MSA row-attention with pair bias + MSA column-attention + MSA transition + outer-product-mean (\(\to \) pair) + triangle multiplicative (outgoing, incoming) + triangle self-attention (starting node, ending node) + pair transition. Signature: evoformerBlock (msaChannels pairChannels nBlocks : Nat). The triangulation-aware operations are the key inductive bias.

Lin et al. 2017. Takes the four stage outputs of a CNN backbone (channels c2/c3/c4/c5 at strides \(4/8/16/32\)), projects each to target channels with a \(1\times 1\) lateral conv, merges them top-down (upsample-2\(\times \) then elementwise add), and applies a \(3\times 3\) smoothing conv at each merged level. Output: four feature maps, each target-wide, at the original spatial resolutions. Bundled because the cross-scale add doesn’t fit a linear NetSpec. Standard kit in 2-stage detectors (Mask R-CNN, Cascade R-CNN) and single-stage detectors (RetinaNet). Signature: fpnModule (c2 c3 c4 c5 target : Nat).

The GoogLeNet parallel-branch module (Szegedy et al. 2014). Four branches computed in parallel: 1\(\times \)1 conv (b1 channels); 1\(\times \)1 reduce then 3\(\times \)3 (b2); 1\(\times \)1 reduce then 5\(\times \)5 (b3); 3\(\times \)3 maxPool then 1\(\times \)1 (b4). Concat along channels for b1 + b2 + b3 + b4 outputs. The 1\(\times \)1 dimension reducers on branches 2 and 3 are the paper’s trick — they make the expensive 3\(\times \)3 and 5\(\times \)5 convs operate on reduced channel counts. Signature: inceptionModule (ic b1out b2reduce b2out b3reduce b3out b4out : Nat).

Selective state-space block in the S6 formulation (Gu & Dao 2023). Signature: mambaBlock (dim stateSize expand nBlocks : Nat). Bundles RMSNorm + linear expand + depthwise 1D conv + SiLU + selective-scan SSM + gated product + output projection, for nBlocks stacked layers. The selective scan is the novel primitive; everything else could be decomposed to existing layers if we cared to unpack the bundle.

Hybrid local-conv + patch-level transformer (Mehta & Rastegari 2022). Local 3\(\times \)3 conv \(\to \) 1\(\times \)1 projection to transformer dim \(\to \) unfold into patches \(\to \) \(L\) transformer blocks across patches \(\to \) fold back \(\to \) 1\(\times \)1 projection back \(\to \) concat with input \(\to \) 3\(\times \)3 fusion. Signature: mobileVitBlock (ic dim heads mlpDim nTxBlocks : Nat). The unfold/fold operations are pdiv_reindex-style shape transformations; all the genuinely new math is already covered by the transformer proof chapter.

The whole NeRF network (Mildenhall et al. 2020) bundled as one primitive: 8 hidden ReLU-FC layers of hiddenDim, mid-skip concatenating \(\gamma (x)\) at layer 5, dual output heads (1-dim volume density \(\sigma \) + 3-dim RGB via a direction-conditioned branch). Under 600K parameters at the canonical config. Signature: nerfMLP (encodedPosDim encodedDirDim hiddenDim : Nat).

The AlphaFold-1 / trRosetta stem as one layer: the two per-residue blocks of a crop, \([B, 2 \cdot L \cdot D]\) on the host, through two dense maps \(D \to C\) (two weights, no bias — the convBn after it carries the pair map’s bias) and outer-summed into the \([B, C, L, L]\) pair map; \(K\) host pair planes — a language model’s contact map, a template distogram — ride behind the blocks and are concatenated as the map’s last \(K\) channels. Codegen-backed: the structure demo trains through it, and its two weight VJPs are Proofs/Foundation/PairTile.lean. Signature: pairTile (seqLen inDim outDim : Nat) (pairIn : Nat := 0).

Swin’s 2\(\times \)2 spatial downsample + linear channel projection (inDim \(\to \) outDim). Transformer-side analog of a stride-2 conv.

Sinusoidal frequency basis (Vaswani 2017, reused by NeRF 2020): \(\gamma (p) = (\sin 2^0 \pi p, \cos 2^0 \pi p, \ldots , \sin 2^{L-1} \pi p, \cos 2^{L-1} \pi p)\). Zero trainable parameters — it’s a deterministic lift of a low-dim coordinate into a high-frequency feature space where an MLP has enough wiggle room to represent sharp details. Output dim \(=\) inputDim \(\cdot 2 \cdot \) numFrequencies. Signature: positionalEncoding (inputDim numFrequencies : Nat).

Grouped 1\(\times \)1 conv + channel-shuffle permutation + 3\(\times \)3 depthwise + grouped 1\(\times \)1 conv, for nUnits units (first downsampling, rest residual). The shuffle is parameter-free; grouping reduces 1\(\times \)1 cost by \(g\). Signature: shuffleBlock (ic oc groups nUnits : Nat).

nUnits v2 units (Ma et al. 2018). Basic unit (stride 1): channel-split \([X_1, X_2]\), leave \(X_1\) alone, run \(X_2\) through \(1\times 1 \to 3\times 3\) DW \(\to 1\times 1\) at half-width, concat, then channel-shuffle. Downsample unit (stride 2): both branches see the full input; left does DW-3\(\times \)3-stride-2 \(+\) 1\(\times \)1, right does 1\(\times \)1 \(+\) DW-3\(\times \)3-stride-2 \(+\) 1\(\times \)1; concat doubles channels. v2 throws out v1’s grouped 1\(\times \)1 convs (G2) and skip-add (G4) per the paper’s practical guidelines. Signature: shuffleV2Block (ic oc nUnits : Nat).

Recurrent Invariant Point Attention + backbone frame update + side-chain \(\chi \)-angle prediction. Weights shared across nBlocks rounds — param count does not multiply by nBlocks. Signature: structureModule (singleChannels pairChannels nBlocks : Nat).

Windowed multi-head self-attention at fixed spatial resolution (Liu et al. 2021). Signature: swinStage (dim heads mlpDim windowSize nBlocks : Nat). Internal blocks alternate W-MSA and SW-MSA (shifted-window) to let information cross window boundaries.

nBlocks blocks of self-attention over nQueries learned object queries + cross-attention against an encoder output + FFN. The query embedding is part of the layer’s parameters. Signature: transformerDecoder (dim heads mlpDim nBlocks nQueries : Nat).

Inter-block downsample for DenseNet. BN + 1\(\times \)1 conv (ic \(\to \) oc) + 2\(\times \)2 average-pool stride 2. Halves spatial resolution and (typically) halves channel count via the 1\(\times \)1 conv to compress feature reuse between dense blocks. Signature: transitionLayer (ic oc : Nat).

\(2 \times \) (conv3\(\times \)3 + BN + ReLU) then maxPool-2. Saves its pre-pool activation as a skip for the matching unetUp. Signature: unetDown (ic oc : Nat).

Bilinear \(2\times \) upsample (parameter-free, in place of a transposed conv) + concat with the matching skip + \(2 \times \) (conv3\(\times \)3 + BN + ReLU). Signature: unetUp (ic oc : Nat), where oc is both the output channel count and the expected skip width.

One stack of nLayers dilated causal residual blocks with doubling dilation rates \(2^0, 2^1, \ldots , 2^{\texttt{nLayers}-1}\) (van den Oord et al. 2016). Each block: dilated 2-tap causal conv \(\to \) gated activation \(\tanh (\text{filter}) \odot \sigma (\text{gate})\) \(\to \) 1\(\times \)1 project back to residualCh (residual path) + 1\(\times \)1 skip projection to skipCh. Skip outputs across blocks are summed into the final head. Signature: waveNetBlock (residualCh skipCh nLayers : Nat). Output channels are skipCh: bestiary convention picks the skip path as the "forward" output since it’s what feeds the final classifier.

prove: for the accumulating renders (acc\(k\)x and lambacc\(k\)x; A3 is \(k = 4\)), one graph serves both phases and two runtime scalars decide which:

1. the accumulator \(G_t = \mathrm{keep}\cdot G + g\) is the momentum op at \(\mu = \mathrm{keep}\). At \(\mathrm{keep} = 0\) the cycle resets by arithmetic, and after \(k\) micro-batches \(G_t\) is their sum;

2. on an accumulate micro-batch, \(\beta = 1\) on both moments and \(\mathrm{lr} = 0\): the moments and the parameter pass through unchanged, for AdamW and for LAMB alike;

3. on the apply micro-batch, the scalar pair \((1 - \beta _1)/k\) and \((1 - \beta _2)/k^2\) applied to the sum \(G_t\) gives exactly the plain step of Theorem 63 or 76 at the mean \(G_t/k\), the \(k^2\) because the second moment squares its gradient.

The typed graph reads the two moment scalars as \(\beta \) and \(1 - \beta \); the pair with \(k\) folded in is the one hand-written block of the step, the one its banner names. Item 3 is the identity that lets the faithful ops be read at the mean.

Used for residual connections. assume:

1. \(B_f\) is a correct backward function for \(f\) (\(\mathsf{HasVJP}\, f\)) [hf]

2. \(B_g\) is a correct backward function for \(g\) (\(\mathsf{HasVJP}\, g\)) [hg]

3. \(f\) is differentiable everywhere [hf_diff]

4. \(g\) is differentiable everywhere [hg_diff]

prove: \(\mathsf{HasVJP}\, (f + g)\).

Rank-2 transcription of Theorem 12: given \(\mathsf{HasVJPMat}\) for \(F\) and \(G\) (with flat differentiability), \(\mathsf{HasVJPMat}\, (F + G)\) with backward \(B_F(A, dY) + B_G(A, dY)\).

\(\mathsf{HasVJP}\, \bigl(\mathrm{bnAffine}(\gamma , \beta )\bigr)\): each input feeds one output scaled by \(\gamma \), so the gradient comes back scaled by \(\gamma \).

assume:

1. \(\varepsilon {\gt} 0\) [h\(\varepsilon \)]

prove: on the network’s \([N, C, H, W]\) layout, with each channel normalized over all \(N \cdot H \cdot W\) of its cells, the batch-statistics layer bnBatchTensor4 is differentiable and has a VJP, and its renderable backward, bnBatchTensor4GradInput, is that VJP: at every input and cotangent it equals the layer’s Jacobian contracted with the cotangent. This is the form the batch-statistics steps of this chapter and every BatchNorm net from Chapter 5 on render.

assume:

1. \(\varepsilon {\gt} 0\) [h\(\varepsilon \)]

prove: \(\mathsf{HasVJP}\, \bigl(\mathrm{bnForward}(\varepsilon , \gamma , \beta )\bigr)\).

assume:

1. \(\varepsilon {\gt} 0\) [h\(\varepsilon \)]

prove: \(\mathsf{HasVJP}\, (\mathrm{bnNormalize})\), with the consolidated three-term backward (writing \(s := \mathrm{istd}\), \(\hat{x} := \mathrm{bnXhat}\)):

assume:

1. \(\varepsilon {\gt} 0\) [h\(\varepsilon \)]

prove: \(\mathsf{HasVJP}\, (\texttt{chanLNTensor3}\; c\; h\; w\; \varepsilon \; \gamma \; \beta )\): LayerNorm over the channel axis at every pixel of a \((c, h, w)\) tensor, with per-channel \(\gamma , \beta \). The ConvNeXt-T chain’s twenty-two channel-LayerNorms are this witness; the head’s, after the pool, is the vector LayerNorm (Theorem 142).

assume:

1. odd kernels [hkH, hkW]

2. every BatchNorm’s \(\varepsilon {\gt} 0\) [hq]

3. at \(x\), every ReLU is off its corner, and at each of the four pools every window is entirely \(\le 0\) before its ReLU or has its maximum strictly above every other cell except its twins, cells equal at every weight upstream of that pool [hx]

4. a selection per pool names a maximal cell of each window

5. \(L\) is any loss of the logits, with gradient \(g\) there [hL]

prove: each gradient node of Theorem 57, at the chain cotangent with each pool routed at its selection, is \(\partial L/\partial \theta \) of the forward with that one parameter varied. Instantiated at the softmax cross-entropy with the emitted loss cotangent.

assume:

1. the first seven convolutions are frozen, \(x_1\) their output [hx\(_1\)]

2. the hypotheses of Theorem 43 at that input: the bounds, the margins up to twins, the \(\eta \)-accurate oracle, the small-step and dominance conditions [hx, hT, hm2, hmq, hm3, hm4, hgh, hsmall, h1, h2]

prove: one SGD step on the last convolution’s kernel decreases the CIFAR-8 cross-entropy loss:

prove: all twenty-two parameter updates of cifar8_train_step.mlir — kernels and biases of \(W_1, \dots , W_8\) across the four stages, and weights and biases of the three dense layers — denote the certified step at cifarCnn8Forward, each at the chain cotangent carried back from the softmax cross-entropy gradient, through the dense head and then stage by stage.

prove: every gradient node of the packed steps cifar8w_{sgd,mom,adam} and cifar8w_bn_{sgd,mom,adam} — twenty-two without BatchNorm and thirty-eight with it — denotes the certified per-layer Jacobian contracted with the chain cotangent. The optimizer update that consumes the node (SGD, Nesterov or AdamW) is outside the statement.

assume:

1. every BatchNorm’s \(\varepsilon {\gt} 0\) [h\(\varepsilon \)]

2. at \(x\), every ReLU is off its corner, and at each of the four pools every window of every example is entirely \(\le 0\) before its ReLU or has its maximum strictly above every other cell except its twins, cells equal at every weight upstream of that pool [hx]

3. a selection per pool names, for every example, a maximal cell of each window

4. \(L\) is any loss of the logits, with gradient \(g\) there [hL]

prove: each of the thirty-eight gradient nodes of Theorem 58, at the chain cotangent with each pool routed at its selection, is \(\partial L/\partial \theta \) of the batched forward with that one parameter varied: the loss is differentiable in the parameter, and the node is its gradient. Instantiated at the batch-mean softmax cross-entropy at the graph’s soft target, with the emitted loss cotangent. The render’s own routing, each window’s cotangent to its first maximal cell, is one such selection, so the chain of Theorem 58 is a chain this theorem is about.

prove: every one of the thirty-eight gradient nodes of the packed steps cifar8wb_bbn_* — the eight-convolution net with each BatchNorm reducing over the batch and the spatial dimensions, as ResNet-34’s do — denotes at batch \(N\) the certified batched gradient: the layer’s Jacobian contracted with the cotangent the rendered backward chain delivers, summed over the batch, at either precision, with the batch-mean softmax cross-entropy at any target. The batch is coupled through the BatchNorms, so there is no per-example reading. That the chain delivering those cotangents is the certified whole-net backward is Theorem 59, and that each node is the loss’s derivative in its parameter is Theorem 60.

assume:

1. every BatchNorm’s \(\varepsilon {\gt} 0\) [h\(\varepsilon \)]

2. at \(x\), each of the eight pre-ReLU activations is off zero, each of the four pools has a unique maximum in every window of every example, and the two head pre-activations are off zero [hs]

prove: the backward chain the cifar8wb_bbn steps render — the three dense backs with the head’s two ReLU masks between them, then per stage the pool scatter and twice the ReLU mask, the batch-statistics BatchNorm backward and the convolution input-VJP — is the backward of the certified VJP of the whole batched forward, and so at every batch size \(N\), every logits cotangent \(dy\) and every input coordinate \(i\) equals

the Jacobian of the net contracted with the cotangent. The forward whose Jacobian this is, is the seventeen-layer composite the net’s certified layer computes.

prove: with BatchNorm after each of the eight convolutions, all thirty-eight parameter updates of cifar8_bn_train_step.mlir — kernel, bias, and BatchNorm \((\gamma , \beta )\) per convolution, and weight and bias per dense layer — denote the certified step at cifarCnnBn8Forward: the dense updates at the cotangent the head carries back, the BatchNorm updates at the ReLU-masked cotangent, the convolution updates at the cotangent BatchNorm’s input gradient passes down.

prove: for a per-channel BatchNorm at any spatial shape and any cotangent \(c\), the two emitted updates denote the certified steps \(\gamma - \alpha \sum \operatorname {pdiv}(\mathrm{bn})\cdot c\) and \(\beta - \alpha \sum \operatorname {pdiv}(\mathrm{bn})\cdot c\) — the scale and shift each take the gradient of Theorem 50’s forward.

prove: all fourteen parameter updates cifar_train_step.mlir emits — weights and biases of the three dense layers and of the two-stage CIFAR net’s four convolutions — denote the certified step at the real forward cifarCnnForward, each at the chain cotangent carried back from the softmax cross-entropy gradient. Stated at one example; the committed text is the batch-128 instance read one image at a time, and the eight-convolution ties below are batched.

assume:

1. odd kernels [hkH, hkW]

2. at \(x\), every ReLU is off its corner, and at each of the two pools every window is entirely \(\le 0\) before its ReLU or has its maximum strictly above every other cell except its twins, cells equal at every weight upstream of that pool [hx]

3. a selection per pool names a maximal cell of each window

4. \(L\) is any loss of the logits, with gradient \(g\) there [hL]

prove: each of the fourteen certified gradients the fused updates of Theorem 52 subtract, at the chain cotangent with each pool routed at its selection, is \(\partial L/\partial \theta \) of cifarCnnForward with that one parameter varied. Instantiated at the softmax cross-entropy with the emitted loss cotangent.

prove: all ten parameter updates cnn_train_step.mlir emits — the dense head’s \(W_5, b_5, W_4, b_4, W_3, b_3\) and the convolutions’ \(W_2, b_2, W_1, b_1\) — denote, at a batch of \(N\) examples, the certified SGD step at the real forward: each update is the parameter minus \(\alpha \) times the sum over the batch of its layer’s VJP at example \(k\), contracted with example \(k\)’s chain cotangent. That cotangent starts as the softmax cross-entropy gradient at the logits; the dense head carries it back to the max pool, the max pool to \(\mathrm{conv}_2\), and \(\mathrm{conv}_2\) to \(\mathrm{conv}_1\). At the batched index the chain is the emitted backward subgraph itself — the selects, the select_and_scatter and the reversed-kernel convolution — and its slices are the hand-named chain cotangents of the single-example reading (cnnChainCotW2 outright, cnnChainCotW1 for odd kernels, the artifact’s being \(3\times 3\)). For the output weight \(W_5\) the term is \(\partial L_k/\partial W_5\), the gradient of example \(k\)’s whole loss, as for the MLP. For every \(N\); the artifact is \(N = 128\).

assume:

1. odd kernels [hkH, hkW]

2. at \(x\), every ReLU is off its corner, and every pooling window is entirely \(\le 0\) before its ReLU or has its maximum strictly above every other cell except its twins: cells equal at every value of both convolutions’ weights and biases, which two cells with identical receptive fields are — the tie a blank stretch of an MNIST image produces [hx]

3. \(\sigma \) names a maximal cell of each window

4. \(L\) is any loss of the logits, with gradient \(g\) there [hL]

prove: each of the ten gradient nodes, at the chain cotangent with each window’s cotangent routed to the cell \(\sigma \) names, is \(\partial L/\partial \theta \) of mnistCnnNoBnForward with that one parameter varied. Instantiated at the softmax cross-entropy with the emitted loss cotangent; and at \(\sigma \) the first maximal cell of each window — the routing the emitted select_and_scatter uses — the chain is Theorem 41’s, so the emitted updates step along these gradients, ties included.

assume:

1. the second convolution’s input and the dense weights are bounded, \(|x_1| \le a\) and \(|W_k| \le w_k\) [hx, hW\(_3\), hW\(_4\), hW\(_5\)]

2. a margin at every ReLU after \(\mathrm{conv}_2\) and in the dense head [hm2, hm3, hm4]

3. in every live pooling window the maximum clears the other cells by a margin, up to twins: cells that read the same input patch, so are equal at every kernel — a blank corner of an MNIST image is one [hmq, hT]

4. an \(\eta \)-accurate gradient oracle, and the small-step and dominance conditions in the shape of Theorem 25, with the chain’s own constants: the small step is \(2\, w_5 d_4 w_4 d_3 w_3\, (4hw)\, a D {\lt} 1\), the Lipschitz constant carries the same product squared with the class count, and the first convolution’s adds the kernel’s fan-in [hgh, hsmall, h1, h2]

5. the bounds, \(\alpha \) and \(\eta \) are non-negative [ha, hw\(_3\), hw\(_4\), hw\(_5\), hlr, h\(\eta \)]

prove: one SGD step on the second convolution’s kernel, at one example with the other parameters fixed, decreases the cross-entropy loss:

The same holds at the first convolution and at both convolutions’ biases, and with the binary32 gradient in place of \(g\), its \(\eta \) the rounding budget of §B.5 and its pool margin strict. Witness: a trained reduced-width network and a real MNIST test image meet every hypothesis, at both kernels and both biases.

assume:

1. \(\varepsilon {\gt} 0\) at the LayerNorms [h\(\varepsilon \)]

2. \(L\) is any loss of the logits, with gradient \(g\) there [hL]

prove: for each node of Theorem 115, the node at the chain cotangent from \(g\) is \(\partial L/\partial \theta \) of the batched ConvNeXt-T forward at the same mask (cnxNetB; at none, convNextForwardTCh applied per example) with that one parameter varied, for every batch and number of classes. Instantiated at the label-smoothed cross-entropy for targets summing to one.

prove: each of the 182 parameter gradients of the batched ConvNeXt-T step (\([3, 3, 9, 3]\) blocks, channel-LayerNorm at twenty-two sites, per-channel layer scale, GELU at either form) denotes the certified batch gradient at the chain cotangent of the label-smoothed cross-entropy, with the batch and the number of classes binders: convnext_adam_train_step.mlir and the convnextin_* steps emit these nodes. Stated at one replica, at either precision, with stochastic depth as an optional mask at each block’s one site, between the layer scale and the skip: at none the chain is the drop-free artifacts’, at some the drop renders’, where every node of a block reads the cotangent scaled by its example’s mask and the skip reads it raw.

assume:

1. \(\varepsilon {\gt} 0\) at the stem’s, the four stages’, the three downsamples’ and the head’s LayerNorms, twenty-three in all [hs\(\varepsilon \), h1–h4, hd1–hd3, hh\(\varepsilon \)]

prove: the twelve-stage ConvNeXt-T input-gradient chain convnextInputGradB equals the backward of the batched forward,

for every batch, number of classes and cotangent, at either GELU form.

assume:

1. \(\varepsilon {\gt} 0\) at the stem’s, the four stages’, the three downsamples’ and the head’s LayerNorms, twenty-three in all [hs\(\varepsilon \), h1–h4, hd1–hd3, hh\(\varepsilon \)]

prove: the backward composed from the per-stage backwards of ConvNeXt-T — the \([3,3,9,3]\) stages at widths \(96\) to \(768\) with channel LayerNorm (convNextForwardTCh) — equals the \(\operatorname {pdiv}\)-contracted Jacobian of the forward at every input, cotangent and pixel, at either GELU form; and the typed forward graph denotes exactly that forward, with no hypothesis, and at the committed ten-class spec the function the spec denotes. Stated per example, on the drop-free f32 chain: the batched bf16 steps have no typed forward graph, and their backward is tied by the next rung.

\(db = dy\). prove: for all \(i\):

\(\mathsf{HasVJP}\, \bigl(\mathrm{dense}(W, b)\bigr)\): the input-gradient backward of a dense layer is multiplication by \(W\).

prove: \(\mathsf{HasVJPMat}\) of the row-wise dense map \(X \mapsto (n \mapsto W^{\mathsf T} X_n + b)\) across the \(N\) tokens, one \(W, b\) shared by every row.

\(dW = x \otimes dy\), with \(F\) and \(\varphi \) as in Theorem 18, prove: for all \(i, j\):

assume:

1. every BatchNorm’s \(\varepsilon {\gt} 0\) [h\(\varepsilon \)w]

2. \(L\) is any loss of the logits, with gradient \(g\) there [hL]

prove: for each node of Theorem 99, the node at the chain cotangent from \(g\) is \(\partial L/\partial \theta \) of the forward at the same masks (efficientnetForwardBFullDrop; efficientnetForwardBFull at none) with that one parameter varied, at any batch and number of classes. Instantiated at the label-smoothed cross-entropy.

assume:

1. every BatchNorm’s \(\varepsilon {\gt} 0\), all forty-nine of them [h\(\varepsilon \)w]

prove: each of the 213 emitted parameter gradients of the batched EfficientNet-B0 step (efficientnet_adam_train_step.mlir and the efficientnetin_* steps) denotes the certified batch gradient at the chain cotangent of the label-smoothed cross-entropy — through swish, the SE gate’s fan-in and the batch-BatchNorm backward — with the target and the number of classes binders. Stated at one replica, at either precision, with stochastic depth on the nine skip-carrying blocks and the classifier dropout as optional masks: at none the chain is the drop-free artifacts’, at some the drop, do and dropdo renders’ — efficientnetin_emarmstfdp64dropdowneps0001bf16, the conf’s, among them — where a masked block’s nodes read the cotangent through its site and the skip reads it raw, and the dense weight node reads the dropped activation; the optimizer tails that consume the nodes are outside it. The tie also carries the forty-nine convolution-bias nodes no render emits.

assume:

1. \(R \ge 1\) replicas of \(N \ge 1\) images each, all forty-nine BatchNorms synchronised with \(\varepsilon {\gt} 0\) [hR, hN, h\(\varepsilon \)w]

2. each replica’s loss cotangent is its shard of the global batch’s, scaled by \(R\) [hgs]

prove: the mean over replicas of each of the 213 all-reduced gradient nodes of the data-parallel renders at either precision (efficientnetin_emarmstfdp64dropdowneps0001bf16, the conf’s, among them, at \(\texttt{replicas} {\gt} 1\); loss \(\div B\) per replica; each replica’s stochastic-depth and dropout masks the shards of global ones) equals the single-device node of Theorem 99 at batch \(R \cdot N\) and those global masks, loss \(\div (R \cdot B)\). Instantiated at the label-smoothed cross-entropy. The forward half — each replica’s shard of the data-parallel forward graph denotes the single-device forward at its shard of the masks — is its own statement.

assume:

1. every BatchNorm’s \(\varepsilon {\gt} 0\) [h\(\varepsilon \)w]

prove: the hand-written sixteen-block reverse chain efficientnetInputGradBFull equals

for every batch, cotangent and pixel.

assume:

1. every BatchNorm’s \(\varepsilon {\gt} 0\), forty-nine of them [h\(\varepsilon \)w]

prove: the backward composed from the per-block backwards, for the batched EfficientNet-B0 — sixteen MBConv blocks with squeeze-and-excitation and swish, at training BatchNorm and batch \(N\) (efficientnetForwardBFull) — equals the \(\operatorname {pdiv}\)-contracted Jacobian of the forward at every input, cotangent and pixel; and the emitted forward graph denotes exactly that forward. With stochastic depth on the nine blocks that carry a skip and classifier dropout between the pool and the dense, the drop and dropdo forward graphs denote the masked forward at every mask, at training and at inference BatchNorm, at either precision — the committed efficientnetin_drop_fwd and efficientnetin_dropdo_fwd are the f32 instances. The graph ties take no hypothesis, and the VJP no smooth-point one: swish and the SE gate have no corner.

Used for Squeeze-and-Excitation. assume:

1. \(B_f\) is a correct backward function for \(f\) (\(\mathsf{HasVJP}\, f\)) [hf]

2. \(B_g\) is a correct backward function for \(g\) (\(\mathsf{HasVJP}\, g\)) [hg]

3. \(f\) is differentiable everywhere [hf_diff]

4. \(g\) is differentiable everywhere [hg_diff]

prove: \(\mathsf{HasVJP}\, (f \odot g)\) for \(f, g : \mathbb {R}^{n} \to \mathbb {R}^{n}\), where \((f \odot g)\, x\, i = f\, x\, i \cdot g\, x\, i\).

prove: in every render carrying ema, the one extra op per parameter denotes \(e' = d\, e + (1 - d)\, \theta '\), with \(\theta '\) the parameter the optimizer tail has just produced, not the one it read, and the decay \(d\) a runtime scalar. The shadow’s distance to the live weight shrinks by the factor \(d\) a step; at \(d = 1\) it is frozen and at \(d = 0\) it is the live weight. The warmup-corrected schedule of \(d\) and the second shadow over the BatchNorm buffers are the driver’s, outside the graph.

Mathlib has no error function, so the book defines \(\mathrm{erf}(z) = \tfrac {2}{\sqrt\pi }\int _0^z e^{-t^2}\, dt\) and \(\mathrm{erfc} = 1 - \mathrm{erf}\). prove: with \(z = -x\sqrt{1/2}\),

which are, term for term, the arithmetic of jax.nn.gelu(approximate=False) and of its jax.vjp, and the text the verified emitter prints around its one chlo.erfc. The backward bridge geluErf_back_bridge pins the emitted diagonal to \(\mathrm{geluErfHasVJP}\)’s backward by rfl.

For every \(gf : \mathrm{GeluForm}\), \(\mathsf{HasVJP}\, (gf.\mathrm{map})\), with backward \(B(x, dy)_i = dy_i \cdot gf.\mathrm{scalar}'(x_i)\). At tanh the record is \(\mathrm{geluHasVJP}\) and at erf it is \(\mathrm{geluErfHasVJP}\), each by rfl; the whole-net theorems below and in Chapter 9 use it at the variable \(gf\).

prove: \(\mathsf{HasVJPMat}\) of the row-wise GELU \(X \mapsto (n \mapsto \mathrm{gelu}(X_n))\) across the \(N\) tokens, at either form of Definition 108.

assume:

1. \(G\) has gradient \(dy\) at \(f(x)\) [hG]

2. \(f\) is differentiable at \(x\), with a VJP there [hf, vf]

prove: \(G \circ f\) has gradient \(\mathrm{backward}\, dy\) at \(x\), the VJP’s backward of \(dy\) (HasGradAt.comp). Read at a parameterised op \(\theta \mapsto \mathrm{layer}\, \theta \) whose output gradient is \(dy\), the loss as a function of \(\theta \) has gradient the op’s parameter Jacobian against it, \(\sum _j \operatorname {pdiv}\, \mathrm{layer}\, \theta \, i\, j \cdot dy_j\) (HasGradAt.param); at a batched op it is the \(\sum _n \sum _j\) every batched gradient node denotes (HasGradAt.param_batchMap).

\(\mathsf{HasVJP}\, \mathrm{id}\), with backward \(B(x, dy) = dy\): the identity passes the upstream gradient through unchanged.

\(\mathsf{HasVJPMat}\, (\mathrm{id})\), backward \(B(A, dY) = dY\).

prove: in the RSB-A3 render of §5-8:

1. the clip: the emitted scale op denotes \(g \cdot \min \bigl(1,\; c / (\sqrt{s} + \varepsilon )\bigr)\), where \(s\) is the sum of squares folded across every parameter, one scalar shared by every leaf. The factor never exceeds one, and at threshold \(kc\), \(k\varepsilon \) and \(k^2 s\) it equals the factor at \(c\), \(\varepsilon \), \(s\), which is what lets the clip run after \(k\)-fold accumulation;

2. LAMB: the four emitted ops denote \((\theta - \mathrm{lr}\, \tau \, r,\; m',\; v')\), with \(m', v'\) Adam’s moments, \(r = \frac{m'/b_1}{\sqrt{v'/b_2} + \varepsilon } + \mathrm{wd}\, \theta \) the decayed direction, and \(\tau = \| \theta \| / \| r\| \) the trust ratio, one when either norm is zero, computed from this parameter’s own norm and no other’s. For the no_weight_decay group the render passes a zero norm, and the step is the plain Adam step \(\theta - \mathrm{lr}\, r\).

assume:

1. \(\varepsilon {\gt} 0\) [h\(\varepsilon \)]

prove: \(\mathsf{HasVJP}\, (\mathrm{layerNormForward})\).

assume:

1. \(\varepsilon {\gt} 0\) [h\(\varepsilon \)]

prove: \(\mathsf{HasVJPMat}\) of the row-wise LayerNorm \(X \mapsto (n \mapsto \mathrm{layerNormForward}\, \varepsilon \, \gamma \, \beta \, (X_n))\) across the \(N\) tokens, with scalar affine parameters \(\gamma , \beta \in \mathbb {R}\).

assume:

1. \(\varepsilon {\gt} 0\) [h\(\varepsilon \)]

prove: \(\mathsf{HasVJP}\, (\texttt{layerNormVec}\; D\; \varepsilon \; \gamma \; \beta )\) for per-feature \(\gamma , \beta \in \mathbb {R}^{D}\), and its row-wise lift to \(\mathsf{HasVJPMat}\) across the \(N\) tokens. This is the LayerNorm of the rendered block and of the depth-\(k\) tower (Theorem 152); the scalar-affine Theorem 141 is the LayerNorm of the weight-shared tower.

For a batch of \(N\) examples \(x_k\) with labels \(\ell _k\), write \(L_k(W) := \mathrm{CE}(\mathrm{dense}(W, b)\, x_k,\ \ell _k)\) for example \(k\)’s loss at fixed \(b\), and \(p_k := \mathrm{softmax}(\mathrm{dense}(W, b)\, x_k)\). prove: the two update operations of § 1.7’s linear_train_step.mlir, read as the batched graph that text is the print of, denote

for every \(N\); and at \(\alpha = \eta / N\), \(N \ge 1\), the weight update is one gradient step at rate \(\eta \) on the batch-mean loss \(\tfrac {1}{N}\sum _k L_k\). That is, the emitted step is the certified gradient step on the batch, the forward it differentiates is the proven fwdGraph lifted across the batch, and the artifact is the instance \(N = 128\), \(\alpha = 0.1/128\). (The single-example reading, poc_train_step_tail_certified, is the \(N = 1\) case; the batched graph prints the committed text byte for byte, tests/TestSmallNetsBatchedEmit.lean.)

assume:

1. \(L\) is any loss of the logits, with gradient \(g\) there [hL]

prove: the classifier’s weight and bias gradient nodes, fed \(g\), are \(\partial L/\partial W\) and \(\partial L/\partial b\) of mnistLinear with that one parameter varied: \(L\) is differentiable in it and the node is its gradient. Instantiated at the softmax cross-entropy with the emitted loss cotangent.

assume:

1. every pixel is bounded, \(|x_i| \le a\) [hx]

2. the gradient \(g\) the step uses is within \(\eta \) of \(\nabla L\) in every coordinate [hgh]

3. the small-step condition \(2aD {\lt} 1\) at the step’s \(\ell _1\)-radius \(D := \alpha \, (\| \nabla L\| _1 + mn\, \eta )\) [hsmall]

4. the error is dominated by the gradient: \(\alpha \eta \, \| \nabla L\| _1\) and \(\tfrac {2a^2}{1-2aD}\, D^2\) are each at most \(\alpha \, \| \nabla L\| _2^2/4\) [h1, h2]

5. \(a\), \(\alpha \) and \(\eta \) are non-negative [ha, hlr, h\(\eta \)]

prove: the step on the weights, at one example with the bias fixed, decreases the loss,

The same holds at the bias, and with the binary32 gradient in place of \(g\), its \(\eta \) the rounding budget of § B.5 rather than an assumption, under that forward’s exp-accuracy and logit-drift bounds.

\(\mathsf{HasVJPMat}\, (B' \mapsto C \cdot B')\), backward \(dB = C^{T} \cdot dY\).

\(\mathsf{HasVJPMat}\, (A' \mapsto A' \cdot D)\), backward \(dA = dY \cdot D^{T}\).

prove: mlp_train_step.mlir emits six update operations, the weights and biases of all three dense layers, and at a batch of \(N\) examples each denotes \(\theta - \alpha \sum _k (\cdot )_k\), a sum over the batch of per-example terms: for the output weight \(W_2\) the term is \(\partial L_k/\partial W_2\), the gradient of example \(k\)’s whole loss; for the other five (\(b_2\) and both hidden layers) it is the certified per-layer Jacobian at example \(k\) contracted with that example’s chain cotangent, the slice of the emitted select/dot_general backward subgraph. The loss cotangent \(\mathrm{softmax}(z_k) - \mathrm{onehot}(\ell _k)\) is pinned to its proven graph. For every \(N\); the artifact is \(N = 128\), and the single-example reading is mlp_train_step_tied_certified.

assume:

1. both hidden pre-activations are off the ReLU corner [h\(_0\), h\(_1\)]

2. \(L\) is any loss of the logits, with gradient \(g\) there [hL]

prove: for each of the six parameters, the gradient node at the chain cotangent of Theorem 30 is \(\partial L/\partial \theta \) of mlpForward with that one parameter varied. Instantiated at the softmax cross-entropy with the emitted loss cotangent.

assume:

1. the pixels and the two upper weight matrices are bounded, \(|x_i| \le a\) and \(|W_k| \le w_k\) [hx, hW\(_1\), hW\(_2\)]

2. a margin at every ReLU of both hidden layers, so the step stays inside the smooth cell [hmargin0, hmargin1]

3. an \(\eta \)-accurate gradient oracle, and the small-step and dominance conditions in the shape of Theorem 25, with the layer’s own constants: the small step is \(2\, w_2 d_2 w_1 a D {\lt} 1\) and the Lipschitz constant \(2\, d_3 d_2^2 w_1^2 w_2^2 a^2 / (1 - 2\, w_2 d_2 w_1 a D)\) in place of \(2a^2/(1-2aD)\) [hgh, hsmall, h1, h2]

4. the bounds, \(\alpha \) and \(\eta \) are non-negative [ha, hw\(_1\), hw\(_2\), hlr, h\(\eta \)]

prove: one SGD step on the input layer’s weights, at one example with the other parameters fixed, decreases the cross-entropy loss:

The same holds at \(W_1\) and at the three biases, each with its own constants, and with the binary32 gradient in place of \(g\) at all three weight layers, its \(\eta \) the rounding budget of § B.5 and a forward rounding margin at every ReLU beside the step margin.

assume:

1. no coordinate of the first pre-activation \(\mathrm{dense}(W_0, b_0)\, x\) is zero [h_smooth_0]

2. no coordinate of the second pre-activation \(\mathrm{dense}(W_1, b_1)\bigl(\mathrm{relu}(\mathrm{dense}(W_0, b_0)\, x)\bigr)\) is zero [h_smooth_1]

prove: the backward composed from the five per-layer backwards in reverse order satisfies, for all \(dy\), \(i\):

Off the two conditions the composed backward is not defined; Definition 28 is the object that stands in.

assume:

1. no coordinate of either convolution’s pre-activation, nor of either hidden dense layer’s, is zero [h1, h2, h3, h4]

2. every \(2 \times 2\) window the max pool sees has a strict maximum [h_mp]

prove: the backward composed from the six per-layer backwards in reverse order — dense, dense+ReLU, dense+ReLU, max pool, conv+ReLU, conv+ReLU — satisfies, for all \(dy\), \(i\):

assume:

1. every BatchNorm’s \(\varepsilon {\gt} 0\) [MNV2PosB]

2. at the batch’s activations, every ReLU6 site is off both its corners [MNV2SmoothAtB]

3. \(L\) is any loss of the logits, with gradient \(g\) there [hL]

prove: for each node of Theorem 86, the node at the chain cotangent from \(g\) is \(\partial L/\partial \theta \) of the forward at the same mask (mobilenetv2ForwardBFullDoOpt; mobilenetv2ForwardBFull at none) with that one parameter varied, at any batch and number of classes. Instantiated at the label-smoothed cross-entropy.

prove: at the structural weights sealW and input sealX, batch 2, there is a class and a pixel at which the backward of the full-width MobileNetV2 is non-zero: the identity of Theorem 82 has content. Its thirty-five ReLU6 clauses are each decided by the weights alone.

prove: each of the 158 emitted parameter gradients of the batched MobileNetV2 step (mobilenetv2_adam_train_step.mlir at \(N = 32\), mobilenetv2in_rmsdp64 at \(N = 64\)) denotes the certified batch gradient at the chain cotangent of a loss cotangent left as a binder, generic in \(N\), at either precision as for ResNet-34, and with the classifier dropout as an optional mask: at none the chain is the drop-free artifacts’, at some the do renders’ — mobilenetv2in_emarmstfdp64wxdwdoeps0001bf16, the run’s, among them — where the dense weight node reads the dropped activation and the cotangent handed down to the pooling passes through the mask. Instantiated at the label-smoothed cross-entropy for targets summing to one. The optimizer that consumes the nodes is outside the tie and is Theorem 87.

assume:

1. \(R \ge 1\) replicas of \(N \ge 1\) images each, every BatchNorm synchronised [hR, hN]

2. each replica’s loss cotangent is its shard of the global batch’s, scaled by \(R\) [hgs]

prove: the mean over replicas of each of the 158 all-reduced gradient nodes of mobilenetv2in_rmsdp64, its bf16 twin and their do renders at \(\texttt{replicas} {\gt} 1\) (loss \(\div B\) per replica; each replica’s classifier-dropout mask the shard of one global mask) equals the single-device node of Theorem 86 at batch \(R \cdot N\) and that global mask, with loss \(\div (R \cdot B)\). Instantiated at the label-smoothed cross-entropy.

assume:

1. \(\varepsilon {\gt} 0\) at the stem’s and the head’s BatchNorm, and the stem’s and head’s ReLU6 off their corners [h\(\varepsilon \)s, h\(\varepsilon \)h, h_stem, h_head]

2. seventeen block maps \(b_1, \dots , b_{17}\), each with a VJP and differentiable at its running activation [hb1 … hb17]

prove: the hand-written reverse pass every MobileNetV2 artifact renders, with its blocks’ backwards filled by those witnesses, is the backward of \(\mathrm{head} \circ b_{17} \circ \dots \circ b_1 \circ \mathrm{stem}\), and so equals

for every batch, cotangent and pixel. The blocks are opaque in the statement; that mobilenetv2ForwardBFull is this composition at the seventeen blocks is a separate shape check.

assume:

1. every BatchNorm’s \(\varepsilon {\gt} 0\) [MNV2PosB]

2. at the batch’s activations, every ReLU6 is off both its corners — the stem’s, an expand clause and a depthwise clause per block (the first block has no expand), the head’s; none after a residual add [MNV2SmoothAtB]

prove: the backward composed from the per-block backwards, for the full-width MobileNetV2 at training BatchNorm and batch \(N\) (mobilenetv2ForwardBFull: seventeen blocks, \(3 \times 224 \times 224\)), equals the \(\operatorname {pdiv}\)-contracted Jacobian of the forward at every input, cotangent and index, generic in \(N\) and in the weights.

assume:

1. at the activation its stage sees, every ReLU is off its corner; the BatchNorm positivities are fields of the weights [Mnv4SmoothAt]

2. \(L\) is any loss of the logits, with gradient \(g\) there [hL]

prove: for each node of Theorem 90, the node at the chain cotangent from \(g\) is \(\partial L/\partial \theta \) of mobilenetv4ForwardBFull with that one parameter varied, at any batch and number of classes. Instantiated at the label-smoothed cross-entropy for targets summing to one.

prove: the same for MobileNetV4-Conv-M, thirty-eight clauses, from a grid-constant base.

prove: each of the 233 parameter gradients of the batched MobileNetV4-Conv-M step (mnv4in_adamdp64) denotes the certified batch gradient at the chain cotangent of a loss cotangent left as a binder, generic in \(N\) and in the number of classes, at either precision. Instantiated at the label-smoothed cross-entropy for targets summing to one.

prove: the same for MobileNetV4’s 233 gradients, instantiated at the label-smoothed cross-entropy.

assume:

1. \(\varepsilon {\gt} 0\) at the stem’s and the two head BatchNorms, and the stem’s and both head ReLUs off their corners [h\(\varepsilon \)s, h\(\varepsilon \)h1, h\(\varepsilon \)h, h_stem, h_h1, h_h2]

2. the fused stage and twenty-one block maps, each with a VJP and differentiable at its running activation [hfused, hb1 … hb21]

prove: the same for the hand-written MobileNetV4-Conv-M reverse pass against \(\mathrm{head} \circ b_{21} \circ \dots \circ b_1 \circ \mathrm{fused} \circ \mathrm{stem}\), generic in \(N\), in the number of classes and in the final feature side (\(7\) at \(224^2\), \(8\) at the paper tier’s \(256^2\)); the shape check against mobilenetv4ForwardBFull is again separate, at \(224^2\).

assume:

1. at the activation its stage sees, every ReLU is off its corner; the BatchNorm positivities are fields of the weights [Mnv4SmoothAt]

prove: the same for MobileNetV4-Conv-M (mobilenetv4ForwardBFull: the fused-IB stage and the universal inverted bottlenecks of §6.6): the backward composed from its stages’ backwards at batch \(N\) equals the \(\operatorname {pdiv}\)-contracted Jacobian at every input, cotangent and index. With stochastic depth on the residual branch of its eighteen skip rows and classifier dropout before the dense, the forward graph of the dropdo train steps denotes the masked forward at every pair of masks and at either precision, and at all-ones masks it is this forward.

prove: the per-parameter update operations the renders of §4.3 emit denote the update laws of that ablation, applied to whatever gradient node \(g\) they consume, with the learning rate a runtime argument:

1. Nesterov momentum: the two emitted ops denote \((\theta - \mathrm{lr}\, (g + \mu v'),\; v')\) with \(v' = \mu v + g\). The heavy-ball step \(\theta - \mathrm{lr}\, v'\) differs from it by exactly \(\mathrm{lr}\, \mu \, (v' - v)\) per coordinate.

2. AdamW: the three emitted ops denote \((\theta ', m', v')\) with \(m' = \beta _1 m + (1 - \beta _1)\, g\), \(v' = \beta _2 v + (1 - \beta _2)\, g^2\) and

   \[ \theta ' = \theta - \mathrm{lr}\, \frac{m'/b_1}{\sqrt{v'/b_2} + \varepsilon } - \mathrm{lr}\, \mathrm{wd}\, \theta , \]

   the bias corrections \(b_1, b_2\) host-computed per step. The denominator is positive for every \(\varepsilon {\gt} 0\).

Plain SGD’s single fused op is inside Theorems 55 and 56 already. Neither law here comes with a descent bound; the claim is that the emitted arithmetic is the law, and nothing about where it goes.

assume:

1. \(S\) is a finite index set with a function \(f_s : \mathbb {R}^{m} \to \mathbb {R}^{n}\) for each \(s \in S\)

2. every \(f_s\) is differentiable at \(x\) [hdiff]

prove: for all \(i, j\):

assume:

1. \(f : \mathbb {R}^{m} \to \mathbb {R}^{n}\) is additive, \(f(u + v) = f(u) + f(v)\) [hadd]

2. \(f\) is homogeneous, \(f(a \cdot v) = a \cdot f(v)\) [hsmul]

prove: for every constant \(c \in \mathbb {R}^{n}\), every \(x\) and all \(i, j\):

At \(c = 0\) this is pdiv_of_linear: the Jacobian of a linear map is the map itself, read on the standard basis. There is no differentiability hypothesis, because a linear map on a finite-dimensional space is continuous and so is its own derivative everywhere. This is the rule every layer that is linear in the variable being differentiated goes through — dense in \(x\), in \(W\) and in \(b\); conv, depthwise and the patch embedding in each of their arguments; BatchNorm’s centering and affine steps; matmul against a constant factor — in place of distributing the sum, product and reindex rules through the layer’s definition term by term. Each use is two one-line algebra side conditions, then the entry is read off \(f(\mathbf{e}_i)\) and a Kronecker indicator collapses.

prove: \(\operatorname {pdivMat}(B' \mapsto C \cdot B')\, B\, (i,j)\, (k,l) = [l = j] \cdot C_{ki}\).

prove: \(\operatorname {pdivMat}(A' \mapsto A' \cdot D)\, A\, (i,j)\, (k,l) = [i = k] \cdot D_{jl}\).

prove: \(\operatorname {pdivMat}(M \mapsto s \cdot M)\, A\, (i,j)\, (k,l) = [i = k \wedge j = l] \cdot s\).

prove: \(\operatorname {pdivMat}(\mathrm{transpose})\, A\, (i,j)\, (k,l) = [j = k \wedge i = l]\).

The ReLU’s VJP record, with the \(\operatorname {pdiv}\)-derived backward. At a point with no zero coordinate the cotangent is masked by \([x_j {\gt} 0]\) (Theorem 26), which is what the emitted mask computes (relu_codegen_matches_canonical). At a point with any zero coordinate ReLU is not differentiable, \(\operatorname {fderiv}\) is the zero map, and the backward is the zero vector in every coordinate, where the emitted mask still passes \(dy\) on the positive ones. The record and the text agree exactly off that set.

assume:

1. \(B_f\) is a correct backward function for \(f\) (\(\mathsf{HasVJP}\, f\)) [hf]

2. \(f\) is differentiable everywhere [hf_diff]

prove: \(\mathsf{HasVJP}\, (\mathrm{residual}\, f)\), where \(\mathrm{residual}\, f\, x = f(x) + x\), with backward \(B(x, dy) = B_f(x, dy) + dy\).

assume:

1. every BatchNorm’s \(\varepsilon {\gt} 0\) [R34PosB]

2. at the batch’s activations, every ReLU is off its corner and every stem-pool window is entirely zero or has one maximal cell, up to twins — cells that read identical input patches [R34LossSmoothAtB]

3. \(L\) is any loss of the logits, with gradient \(g\) there [hL]

prove: for each gradient node of Theorem 68, the node at the chain cotangent from \(g\) is \(\partial L/\partial \theta \) of resnet34ForwardBFull with that one parameter varied, at any batch size and number of classes. Instantiated at the batched label-smoothed cross-entropy, and at the stem with each tied window’s cotangent routed to any of its maximal cells, as the emitted select_and_scatter routes it.

prove: at the structural weights sealW and input sealX, batch 2, for any positive number of classes, there is an (example, class) coordinate \(j_0\) and a pixel \(i_0\) at which the backward of the full-width ResNet-34 is non-zero: the identity of Theorem 65 has content, since the zero map does not satisfy it.

prove: each of the 110 emitted parameter gradients of the batched ResNet-34 step (resnet34_sgd_train_step.mlir and _adam_ at \(N = 32\), resnet34in_momdp64 at \(N = 64\)) denotes the certified batch gradient \(\sum _n \mathrm{Jacobian} \cdot \mathrm{cotangent}\) at the chain cotangent of the label-smoothed cross-entropy (any \(\alpha _{\mathrm{ls}}\) — \(0.1\) in the renders, \(0\) in the ls0 twins — and any target), generic in \(N\). Stated at either precision: the bf16 render resnet34in_momdp64bf16 emits a *GradBBf16 node at every convolution’s weight gradient, and at the identity rounding the renderers instantiate, each such node denotes its f32 peer.

assume:

1. \(R \ge 1\) replicas of \(N \ge 1\) images each, every BatchNorm synchronised [hR, hN]

2. each replica’s loss cotangent is its shard of the global batch’s, scaled by \(R\) [hgs]

prove: the mean over replicas of each of the 110 all-reduced gradient nodes of the data-parallel render (resnet34in_momdp64 and its bf16 twin at \(\texttt{replicas} {\gt} 1\), loss \(\div B\) per replica) equals the single-device node of Theorem 68 at batch \(R \cdot N\) with loss \(\div (R \cdot B)\): the data-parallel step is one device’s step on the global batch. Instantiated at the label-smoothed cross-entropy.

assume:

1. every BatchNorm’s \(\varepsilon {\gt} 0\) [R34PosB]

2. at the batch’s activations, the stem’s ReLU and both ReLUs of each of the sixteen blocks — the body’s and the one after the residual add — are off their corners, and every stem-pool window is entirely \(\le 0\) or has one maximal cell; the head takes no hypothesis [R34SmoothAtB]

prove: the backward composed from the eighteen stages’ backwards in reverse order, for the full-width ResNet-34 at training BatchNorm and batch \(N\) (resnet34ForwardBFull, \(3 \times 224 \times 224\)), equals the \(\operatorname {pdiv}\)-contracted Jacobian of the forward at every input, cotangent and index, generic in \(N\) and in the weights. The strict pool clause is one real batches do not meet; the parameter gradients of Theorem 69 are stated under one they do.

assume:

1. every BatchNorm’s \(\varepsilon {\gt} 0\) [R50PosB]

2. the ResNet-34 smooth-point condition at the batch’s activations, up to twins at the stem pool [R50LossSmoothAtB]

3. \(L\) is any loss of the logits, with gradient \(g\) there [hL]

prove: for each node of Theorem 70, the node at the chain cotangent from \(g\) is \(\partial L/\partial \theta \) of resnet50ForwardBFull with that one parameter varied, at any batch and number of classes. Instantiated at both losses of Theorem 70, and at the stem at any routing to a maximal cell, as for ResNet-34.

prove: the same construction, with its own witness, for ResNet-50 at both 224 and 160 pixels (\(0 {\lt} q \le 7\)).

prove: each of the 161 parameter gradients of the batched ResNet-50 step denotes the certified batch gradient at the chain cotangent of a loss cotangent \(g\) left as a binder, generic in \(N\) and \(q\). Instantiated at the label-smoothed cross-entropy and at BCE-with-logits over \(N \cdot K\), the two losses the ResNet-50 artifacts train. Stated at either precision, as for ResNet-34.

prove: the same for ResNet-50’s 161 gradients at side \(32q\), \(q {\gt} 0\), instantiated at the label-smoothed cross-entropy for targets summing to one and at BCE with the committed divisors \(N \cdot K\) per replica and \((R \cdot N) \cdot K\) on one device.

assume:

1. \(\varepsilon _s {\gt} 0\), the stem ReLU off its corner and the stem pool’s per-example condition [h\(\varepsilon \)s, h_stem, h_pool]

2. sixteen block maps \(b_1, \dots , b_{16}\), each with a VJP and differentiable at its running activation [hb1 … hb16]

prove: the hand-written eighteen-stage reverse pass every ResNet-50 artifact renders, with its blocks’ backwards filled by those witnesses, is the backward of \(\mathrm{head} \circ b_{16} \circ \dots \circ b_1 \circ \mathrm{stem}\), and so equals

for every batch, cotangent and pixel. The blocks are opaque in the statement; that resnet50ForwardBFull is this composition at the sixteen bottlenecks is a separate shape check.

assume:

1. the ResNet-34 conditions, at ResNet-50’s sixteen bottleneck blocks [R50PosB, R50SmoothAtB]

prove: the same for the full-width ResNet-50 (resnet50ForwardBFull) at every resolution \(q\) — the image side is \(32q\), so \(q = 7\) is the 224-pixel net and \(q = 5\) the 160-pixel one.

prove: the three operations the MobileNetV2 and EfficientNet renders emit per parameter denote the TF 1 step \((\theta - b',\; b',\; s')\) with \(s' = \rho s + (1 - \rho )\, g^2\) and \(b' = \mu b + \mathrm{lr}\, g / \sqrt{s' + \varepsilon }\): the learning rate inside the momentum buffer and \(\varepsilon \) inside the root. The rms renders, with the rate applied on the way out, denote the same step with \(b' = \mu b + g / \sqrt{s' + \varepsilon }\) and \(\theta - \mathrm{lr}\, b'\). The root is positive for \(\varepsilon {\gt} 0\) and a nonnegative mean-square, which the initialization at one supplies from step zero. Two differences the numeric ties are built to see: at a zero second moment the textbook placement \(g / (\sqrt{s'} + \varepsilon )\) is the rendered \(g / \sqrt{s' + \varepsilon }\) divided by \(\sqrt{\varepsilon }\), about \(30\times \) larger at B0’s \(\varepsilon = 10^{-3}\); and a buffer built at rate \(\mathrm{lr}_0\) and stepped at \(\mathrm{lr}_1\) lands \(\mu \, b\, (\mathrm{lr}_1 - \mathrm{lr}_0)\) from the rate-outside step, per coordinate, so the two forms part at every step of a staircase.

\(\mathsf{HasVJPMat}\, (\mathrm{rowSoftmax})\): rows are independent, so the backward applies the softmax backward per row.

assume:

1. \(B_g\) is a correct backward for \(g : \mathbb {R}^{n} \to \mathbb {R}^{p}\) (\(\mathsf{HasVJP}\, g\)) [hg]

2. \(g\) is differentiable everywhere [hg_diff]

prove: \(\mathsf{HasVJPMat}\) of the rowwise map on \(\mathbb {R}^{m \times n} \to \mathbb {R}^{m \times p}\), with backward applying \(B_g\) to each row: \(B(A, dY)_r = B_g(A_r, dY_r)\).

\(\mathsf{HasVJPMat}\, (M \mapsto s \cdot M)\), backward \(dA = s \cdot dY\).

prove: for fixed \(Q, V\): \(\mathrm{sdpaBackK}\) matches the \(\operatorname {pdivMat}\) contraction of \(K' \mapsto \mathrm{sdpa}(Q, K', V)\).

prove: for fixed \(K, V\): \(\mathrm{sdpaBackQ}\) matches the \(\operatorname {pdivMat}\) contraction of \(Q' \mapsto \mathrm{sdpa}(Q', K, V)\).

prove: for fixed \(Q, K\): \(\mathrm{sdpaBackV}\) matches the \(\operatorname {pdivMat}\) contraction of \(V' \mapsto \mathrm{sdpa}(Q, K, V')\).

assume:

1. \(B_g\) is a correct backward function for the gate (\(\mathsf{HasVJP}\, \mathrm{gate}\)) [hg]

2. \(\mathrm{gate}\) is differentiable everywhere [hg_diff]

prove: \(\mathsf{HasVJP}\, (\mathrm{seBlock}\, \mathrm{gate})\), where \(\mathrm{seBlock}\, \mathrm{gate}\, x = x \odot \mathrm{gate}(x)\), with backward

SE multiplies input by a sigmoid-gated channel mask.

\(\mathsf{HasVJP}\, (\mathrm{softmax})\), with the closed-form \(O(c)\) backward

where \(\langle p, dy \rangle = \sum _j p_j\, dy_j\) is one precomputed scalar — the rank-1 structure of the Jacobian is what turns the naive \(O(c^2)\) contraction into a reduction plus a broadcast, the same optimization pattern as BN and max-pool.

assume: \(\varepsilon {\gt} 0\) [h\(\varepsilon \)]
prove: \(\mathsf{HasVJPMat}\) of \(X \mapsto X + \mathrm{MHSA}(\mathrm{LN}_1(X))\).

assume:

1. \(\varepsilon {\gt} 0\) [h\(\varepsilon \)]

prove: \(\mathsf{HasVJPMat}\) of one pre-norm encoder block \(z = x + \mathrm{MHSA}(\mathrm{LN}_1 x)\), \(y = z + \mathrm{MLP}(\mathrm{LN}_2 z)\), both LayerNorms scalar-affine, at either GELU form.

\(\mathsf{HasVJPMat}\) of \(\mathrm{dense}_2 \circ \mathrm{gelu} \circ \mathrm{dense}_1\), all per-token.

assume: \(\varepsilon {\gt} 0\) [h\(\varepsilon \)]
prove: \(\mathsf{HasVJPMat}\) of \(h \mapsto h + \mathrm{MLP}(\mathrm{LN}_2(h))\).

assume:

1. \(\varepsilon {\gt} 0\) [h\(\varepsilon \)]

prove: \(\mathsf{HasVJPMat}\) of the weight-tied \(k\)-block tower — one parameter tuple shared by every block, scalar-affine LayerNorms — for every \(k\). The networks the renders train, one parameter tuple per block and per-feature LayerNorms, are Theorem 152.

\(\mathsf{HasVJPMat}\, (\mathrm{transpose})\), backward \(dA = dY^{T}\).

assume:

1. \(\varepsilon {\gt} 0\) at the LayerNorms [h\(\varepsilon \)]

2. \(L\) is any loss of the logits, with gradient \(g\) there [hL]

prove: for each node of Theorem 156, the node at the chain cotangent from \(g\) is \(\partial L/\partial \theta \) of the batched ViT-Tiny forward at the same mask (vitNetB; at none, vitForwardKV applied per example) with that one parameter varied, for every batch and number of classes. Instantiated at the label-smoothed cross-entropy for targets summing to one.

prove: each of the 200 parameter gradients of the batched ViT-Tiny step (three heads of 64, depth 12, \(16 \times 16\) patches) denotes the certified batch gradient at the chain cotangent of the label-smoothed cross-entropy, with the batch, the number of classes and the GELU form binders; the CLS gradient is summed over the batch inside the node. Stated at one replica, at either precision, with stochastic depth as an optional mask at each block’s two sites: vit_adam_train_step.mlir and the vitin_* steps emit these nodes, and at some the chain is the drop renders’, where a block’s out-projection and fc2 nodes read the cotangent scaled by its example’s mask and both skips read it raw.

assume:

1. \(\varepsilon {\gt} 0\) [h\(\varepsilon \)]

prove: at ViT-Tiny’s dimensions the batched input-gradient chain vitInputGradKB equals the backward of the batched network \(\mathrm{batchMap}\, B\; \texttt{vitForwardKV}\) for every \(B\), number of classes, cotangent and index. The same holds at every depth, and the per-example backward graph the render prints denotes that VJP at every depth, image and cotangent.

assume:

1. \(\varepsilon {\gt} 0\) [h\(\varepsilon \)]

prove: the weight-tied ViT backbone \(\mathrm{finalLN} \circ \mathrm{transformerTower}\) is one \(\mathsf{HasVJPMat}\).

assume:

1. \(\varepsilon {\gt} 0\) [h\(\varepsilon \)]

prove: for the depth-\(k\) transformer tower with distinct parameters per block — patch embedding, CLS and position tables, \(k\) blocks of multi-head self-attention and MLP sublayers under vector LayerNorm, the classifier head — the backward composed from the patch embedding’s, the \(k\) blocks’, the final LayerNorm’s and the head’s backwards equals the \(\operatorname {pdiv}\)-contracted Jacobian of vitForwardKV at every input, cotangent and index, for every \(k\) and at either GELU form. The hand-written chain is tied to this backward in Theorem 154. The weight-tied network of Theorem 151, glued to the same patch embedding and head, has its VJP the same way.

prove: for every depth \(k\) the typed forward graph vitFwdGraphKMHV denotes vitForwardKV; and with stochastic depth on both residual branches of every block, the batched graph of the drop renders denotes, at either precision and at each example, this forward with that example’s two mask entries per block as its branch scales — at all-ones masks, vitForwardKV on every example. The emitted ViT text names each shared intermediate once and is not tied to this graph by a text check; the statement is about the typed graph.

assume:

1. \(\varepsilon {\gt} 0\) [h\(\varepsilon \)]

prove: the input-gradient chain vitInputGradK at ViT-Tiny’s shape — a \(3 \times 224 \times 224\) image, \(16 \times 16\) patches (\(N = 196\) plus the CLS token), \(d = 192 = 3 \times 64\), MLP \(768\), twelve blocks with distinct parameters, any number of classes — equals the \(\operatorname {pdiv}\)-contracted Jacobian of the flat forward vitForwardKV at every input, cotangent and index. The statement is generic in the weights and in the GELU form.

assume:

1. \(B_f\) is a correct backward function for \(f\) (\(\mathsf{HasVJP}\, f\)) [hf]

2. \(B_g\) is a correct backward function for \(g\) (\(\mathsf{HasVJP}\, g\)) [hg]

3. \(f\) is differentiable everywhere [hf_diff]

4. \(g\) is differentiable everywhere [hg_diff]

prove: \(\mathsf{HasVJP}\, (g \circ f)\).

assume:

1. \(B_F\), \(B_G\) are correct backward functions (\(\mathsf{HasVJPMat}\, F\), \(\mathsf{HasVJPMat}\, G\)) [hF, hG]

2. the flattened forms of \(F\) and \(G\) are differentiable everywhere [hF_diff, hG_diff]

prove: \(\mathsf{HasVJPMat}\, (G \circ F)\).