Unit 1.4 · about 5 h
173 / 324Branches, broadcasting and saved values in backward
Real graphs reuse nodes and stretch small tensors across batches. Each reuse adds an accumulation obligation and each stretch adds a reduction obligation; missing either gives a silently wrong gradient.
- derive
- implement
You can now
- Hand-compute vector-Jacobian products and accumulated gradients for a branched, broadcast graph with a repeated operand.
- State the broadcast forward rules and implement the backward unbroadcast as a sum over exactly the broadcast axes.
- Distinguish leaf from intermediate tensors, and name every forward value each backward function saves.
Covers
- Repeated use, branches and gradient accumulation
- Broadcast forward rules and unbroadcast in backward
- Leaf and intermediate tensors and saved forward values
Study, in this order
4 steps
- Accumulation and non-scalar seedsD2L - tensors, autograd, MLP and backpropagation
Section 2.5.1: PyTorch does not reset the gradient buffer, so the new gradient is added and
x.grad.zero_()resets it, which suits sums of objectives; Section 2.5.2: thegradientargument as the seed vector \(v\) - Broadcasting and in-place memoryD2L - tensors, autograd, MLP and backpropagation
Broadcasting:
aof shape(3, 1)plusbof shape(1, 2)giving a \((3 \times 2)\) result by copying along length-1 axes; Saving Memory:Y = Y + Xreallocates (idchanges) whileZ[:] = X + YandX += Yreuse storageDepth · The page shows forward broadcasting only. Derive the backward rule in Focus below; the PyTorch note gives the exact forward contract it inverts.
- Exact broadcast contractPyTorch Broadcasting semantics
General semantics: trailing-dimension rules, prepend-1 alignment, result size as the per-dimension max, with
(5, 1, 4, 1)plus(3, 1, 1)giving(5, 3, 4, 1); In-place semantics: an in-place tensor may not change shape - Saved forward valuesUnderstanding Deep Learning - Chapters 3 to 7
Chapter 7, Section 7.2 Observation 1 (printed pp. 97 to 98): run the network and store the activations of all hidden units in the forward pass for later gradient computation
Focus on
- A node feeding two downstream ops gets the sum of both vector-Jacobian products. PyTorch
.gradalso accumulates acrossbackwardcalls, so a missingzero_()silently adds yesterday's gradient to today's. - Broadcast forward: align trailing dimensions, prepend 1 to the shorter shape, and take the per-dimension max. Backward inverts it: sum the upstream gradient over exactly the axes where the input had size 1 or a prepended dimension, then reshape back.
- For \(L = \sum((X + b)^2)\) with \(X\)
(2, 3)and \(b\)(3,): \(dL/dX = 2(X + b)\) keeps shape(2, 3), while \(dL/db\) = 2 times the sum of \((X + b)\) over rows. This matches centered differences to about 1e-8 infloat64. - Leaves are user-created tensors with
requires_grad; intermediates carry the graph. Backward saves the forward values its local derivatives need (activations for weights, inputs for nonlinearities); detaching cuts provenance so gradients stop at the cut.
Reproduce closed-book
(2, 3) and \(b\) (3,): derive \(dL/db\) as 2 times the row sums of \((X + b)\) and \(dL/dX\) as \(2(X + b)\). Implement broadcast plus accumulation in your blank-file engine and match centered differences.Traps
- Averaging instead of summing at a branch. The \(1/N\) factor belongs to an explicit mean op, never to fan-out: a doubly used node contributes twice.
- Summing over the wrong axes in unbroadcast. Sum exactly the axes where the input was size 1 or prepended; summing a full-size axis silently shrinks a real dimension.
- Assuming
backwardzeroes.grad. It adds; the D2L page states this explicitly for PyTorch, unlike MXNet and TensorFlow which reset.
Practice
TensorTonic: Accumulate Shared-Path Gradients (autograd-l06-shared-path-gradient-accumulation)
Accumulation and Complete Scalar Operators section: Accumulate Shared-Path Gradients
Extend your engine so a doubly used node sums both path contributions, then re-run your earlier gradient checks.
D2L Data Manipulation Exercises: broadcasting with other shapes
Exercises: replace the two broadcast operands with other shapes, including 3-dimensional tensors, and check the result matches the trailing-dimension rules
For each shape pair, also write down the backward summation axes before running anything.
Self-test
-
z = x * uwithu = y.detach()andy = x * x. What is \(dz/dx\), and why not \(3x^2\)?\(u\). Detach wipes provenance, so \(x\) influences \(z\) only directly through the \(x\) factor; the path through \(y\) carries no gradient. D2L verifiesx.grad == u. -
What is the shape of
(5, 1, 4, 1)plus(3, 1, 1), and over which axes is the smaller operand's gradient summed?The result is(5, 3, 4, 1). The smaller operand aligns as(1, 3, 1, 1), so its gradient sums over dims 0 and 2, then reshapes to(3, 1, 1). -
What happens on a second
backwardcall withoutzero_()?The new gradient adds to the stored one. That suits optimizing a sum of objectives but corrupts per-batch training. -
For \(y = XW + b\) inside a deeper network, which values must backward have available?
\(X\), for \(dL/dW = X^\top G\), and \(W\), for \(dL/dX = G W^\top\). \(W\) costs no extra memory because the parameter is alive anyway; \(X\) is the activation memory backward forces you to keep. \(dL/db\) needs only \(G\).
Completion criterion
Stop when
Your engine accumulates shared-node gradients, unbroadcasts a (2, 3) plus (3,) graph, matches centered differences, and you can state the summation axes for any broadcast pair without notes.
Leads to
- 1.5 Gradient checks with finite differences Check the branched engine against finite differences.
- 2.2 Autograd state: views, in-place ops, leaves and accumulation See leaves, views and in-place ops in the framework.
Comes after
- 1.8 Batched shapes, operation costs and conditioningApply the forward shape rules when deriving broadcast gradients.
- 2.5 Forward mode, reverse mode, JVPs, VJPs and Hessian-vector productsBroadcast backward sums cotangents, which is the VJP of a broadcast; read it as this unit's pullback.
- 1.3 Computation graphs and reverse accumulationHandle shared nodes, broadcasting and saved values.
- 2.5 Custom autograd functionsRebuild the same VJP logic in your own engine.
Resources for this unit
- Understanding Deep Learning - Chapters 3 to 7 Read Chapters 3-7, emphasizing Chapter 5 for likelihood-derived losses and Chapter 7 for backpropagation and initialization. Use notebooks 3.1, 4.1, 5.1, 6.1 and 7.1 only as equation and plot companions.
- D2L - tensors, autograd, MLP and backpropagation Use chapter_preliminaries ndarray and autograd, then MLP, backpropagation, and numerical stability and initialization. Recreate scalar operations, affine maps, activation, reduction, stable cross-entropy, topological backward and unbroadcast in a blank file.
- PyTorch Broadcasting semantics Read General semantics for the trailing-dimension rules, prepend-1 alignment and max-size result, with the
(5, 1, 4, 1)plus(3, 1, 1)examples; read In-place semantics for why an in-place tensor may not change shape.
Also in this section
- TensorTonic - Build Micrograd from Scratch Use the 25-problem Micrograd study plan in order for derivatives, gradient checks, value graphs, reverse mode, accumulation, operators, modules, losses and training. Preserve local tests as the evidence instead of a platform score.
- Deep-ML candidate - autograd and loss problems After the five-sample platform check passes, use at most one of 26 Basic Autograd Operations, 997 Manual Cross-Entropy Backprop or 1227 Stable Softmax as a changed-prompt drill.
- CS231n Linear Classification - Softmax classifier and numeric stability Read the Softmax classifier section for the cross-entropy form \(L_i = -f_{y_i} + \log \sum_j e^{f_j}\), then the Practical issues: Numeric stability part for the max-subtraction code
f -= np.max(f). - CS231n Neural Networks 3 - Gradient checks Read the Gradient Checks section only: centered formula with \(O(h^2)\) error, relative error with thresholds, double precision, step size \(h\) near 1e-5, kinks, few datapoints, burn-in, regularization and dropout discipline, and checking a few dimensions per parameter.
- `torch.autograd.gradcheck` API and tolerances Read the
gradchecksignature defaults (eps1e-6,atol1e-5,rtol1e-3) and the note that inputs with overlapping memory, such asexpandoutputs, will likely fail the check. - `torch.nn.functional.cross_entropy` - `reduction` and `ignore_index` Read the
ignore_indexandreductionparameters: ignored targets contribute no gradient, and mean reduction averages over the non-ignored targets only. - PyTorch Tensor Views Read the view and storage sharing example (
data_ptr, mutating the view changes the base), the transpose non-contiguous example, the full view-ops list, and the special notes:reshapemay copy,contiguouscopies only when needed, basic indexing returns views while advanced indexing copies. - `torch.Tensor.stride` - stride and storage offset Read the stride definition: strides are per-dimension element steps, and with
storage_offsetthey locate every element. Combine with the Tensor Views page for transposed and expanded layouts. - PyTorch Autograd mechanics - in-place operations Read In-place operations with autograd (why in-place rewrites the graph and is discouraged) and In-place correctness checks (the version counter and the error on accessing a modified saved tensor).
- `torch.nn.Module` - parameters, buffers and state Read the parameters and buffers documentation: modules collect parameters recursively for the optimizer, while
register_bufferholds persistent non-parameter tensors saved alongside them.