ParamGradNodes — each batched parameter gradient node is a loss derivative #
GradNodesB states each emitted parameter gradient node as its layer's parameter Jacobian
contracted with an arbitrary output cotangent. Here the cotangent is the gradient of a scalar G
at the layer's output (HasGradAt), and the node becomes ∂G/∂θ with the layer's parameter
varied: one lemma per node kind, shared by every net, each concluding HasGradAt in the
parameter (the loss is differentiable there and the node is its gradient).
The BatchNorm γ/β nodes are stated in the transposed [C, N·H·W] layout at the reassocB
index; their lemmas re-sum the Jacobian over that permutation (bnLAPerm). A conv or depthwise
bias that the render reads with the β op (EfficientNet-B0's) is biasBeta_hasGradAt: the bias enters
as a channel broadcast (*_bias_split), so the channel sum the β node computes is its derivative.
cStridedInB — the emitted strided-conv input-cotangent — is the batched strided conv VJP's
backward, at any saved input.
dInB — the emitted depthwise input-cotangent — is the batched depthwise VJP's backward.
dStridedInB — the emitted (symmetric) strided depthwise input-cotangent — is the batched
strided depthwise VJP's backward, at any saved input.
dStridedXlaInB — the emitted XLA-SAME strided depthwise input-cotangent — is the batched
strided depthwise VJP's backward, at any saved input.
Back through batch BN: the gradient at its input is bnInB of the gradient at its output.
Back through u ↦ u + v with the skip v held fixed: the gradient passes unchanged. This is
how a residual block's body sees the loss once one of its parameters varies.
At a residual layer the loss read at the BODY output is u ↦ G (u + v): the skip is a constant
once a body parameter varies, so its gradient there is still the block-output cotangent.
Back through a relabelling Fin n ≃ Fin m along n = m (MobileNetV4's head reads [N, c] as
[N, c, 1, 1]): the cotangent is read back along the same cast.
Back through batch BN, at the certified backward's own spelling bnBackB (EfficientNet-B0's
tie threads this form; the ResNets and MobileNets emit bnInB).
Back through a batched conv: cInB.
Back through a batched symmetric strided conv: cStridedInB.
Back through a batched depthwise: dInB.
Back through a batched symmetric strided depthwise: dStridedInB.
Back through a batched XLA-SAME strided depthwise: dStridedXlaInB (MobileNetV2's and
EfficientNet-B0's stride-2 bodies).
Conv weight node = ∇_W G.
Conv bias node = ∇_b G.
Stride-2 (symmetric) conv weight node = ∇_W G.
Stride-2 (symmetric) conv bias node = ∇_b G.
XLA-SAME strided conv weight node = ∇_W G (MobileNetV2's and B0's stem).
XLA-SAME strided conv bias node = ∇_b G.
Depthwise weight node = ∇_W G.
Depthwise bias node = ∇_b G.
Symmetric strided depthwise weight node = ∇_W G (MobileNetV4's dw_mid at its three
downsampling rows).
XLA-SAME strided depthwise weight node = ∇_W G.
XLA-SAME strided depthwise bias node = ∇_b G.
Dense weight node = ∇_W G, at the flat Mat.flatten index finProdFinEquiv (i, j).
Dense bias node = ∇_b G. The node's statement carries one activation x₀ for every
example; the bias Jacobian is the identity whatever the activation, so any x₀ serves.
The permutation bnBatchLA reads its per-channel core through: network index J ↦ the
[C, N·H·W] cell bnchwBackIdx (J at the mul_assoc cast). bnBatchLA … v J is
bnPerChannelFlat … (bnchwFwd … (reassocB … v)) (bnLAPerm … J) by rfl.
Equations
- One or more equations did not get rendered due to their size.
Instances For
BatchNorm γ node = ∇_γ G, at the reassocB index the render's node reads.
BatchNorm β node = ∇_β G, at the reassocB index.
A channel-broadcast bias's Jacobian is the channel indicator.
A channel-broadcast bias, read by the β node, = ∇_b G. For any per-example op whose
bias enters as per 0 y + broadcast b, the emitted bnBetaGradB on the op's output cotangent
is the loss gradient in the bias.
A depthwise conv's bias is a channel broadcast.
A symmetric strided depthwise conv's bias is a channel broadcast: decimation keeps channels.
An XLA-SAME strided conv's bias is a channel broadcast.
The stride-4 patchify conv's bias is a channel broadcast: both decimations keep channels.
The stride-4 stem's bias Jacobian is a stride-1 conv's, at any input. Both are the channel
indicator (pdiv_bias_of_split), so a node the render emits as a stride-1 bias reduce at the
output resolution is the patchify conv's bias gradient at the real image x.