9.2 Computation Graphs, Backpropagation, and Autograd: Gradients Add at Branches and Multiply Along Paths

A forward pass can create millions of intermediate values. Computing every parameter derivative from scratch would repeat most of that work; backpropagation records local dependencies and reuses them in reverse from the scalar loss.

It's not a new calculus theorem, but rather an algorithm for executing reverse-mode differentiation on directed acyclic computation graphs.

This lesson's objectives

• Express forward as a computation graph and local Jacobian;
• Manually compute the backward pass of a two-layer scalar network;
• Explain why gradients are accumulated at shared/branch nodes;
• Correctly use PyTorch .backward() and .grad with graph control;
• Use finite difference / gradcheck to verify custom gradients.

1. Computation Graph Records Dependencies

Scalar network:

$$z_1=w_{1}x+b_1,a_1=\mathrm{ReLU}(z_1),$$$$z_2=w_2{a}_1+b_2,L=\frac{1}{2}(z_2-y{)}^2.$$

Forward computes values in topological order. Reverse starts from $\overline{L}=\partial L/\partial L=1$ and propagates upstream by multiplying the adjoint with the local derivative:

$${\overline{z}}_2=z_2-y,{\overline{w}}_2={\overline{z}}_2{a}_1,{\overline{b}}_2={\overline{z}}_2,$$$${\overline{a}}_1={\overline{z}}_2{w}_2,{\overline{z}}_1={\overline{a}}_1\mathbf{1}[z_1>0],$$$${\overline{w}}_1={\overline{z}}_{1}x,{\overline{b}}_1={\overline{z}}_1.$$

The bar notation is just "the derivative of the final loss with respect to this intermediate value." Each edge performs a vector-Jacobian product without explicitly constructing the entire massive Jacobian.

2. Why Reverse Mode Is Suitable for Training

A training network typically has:

• a scalar/low-dimensional loss;
• Millions to billions of parameters.

Reverse mode computes the gradient of a scalar loss with respect to all upstream parameters in a single pass, with computational cost typically on par with several forward passes. Forward mode is better suited for scenarios with few inputs and many outputs, such as Jacobian-vector products; modern frameworks also support combining both approaches for computing higher-order derivatives or Jacobians.

“Backward counts each intermediate gradient only once” must be qualified: shared nodes receive multiple downstream contributions, and the framework must first accumulate these contributions before propagating upward.

3. Why Are Branches Added Together?

If:

$$u=x^2,v=3x,L=u+v,$$

Then $x$ has two downstream paths:

$$\frac{dL}{dx}=\frac{dL}{du}\frac{du}{dx}+\frac{dL}{dv}\frac{dv}{dx}=2x+3.$$

The residual connection, parameter sharing, and repeated use of the same embedding all depend on this rule. If you manually implement the backward pass and forget to accumulate branch contributions, you'll get gradients with the correct shape but wrong values.

Broadcasting also affects the backward pass: bias is broadcast along the batch dimension, so the bias gradient must be summed along that dimension.

4. The Reverse Shape of Matrix Layer

Set batch forward:

$$Z=X{W}^T+b,G=\frac{\partial L}{\partial Z}.$$

Then:

$$\frac{\partial L}{\partial X}=GW,\frac{\partial L}{\partial W}=G^{T}X,\frac{\partial L}{\partial b}=\sum _{i=1}^B{G}_i.$$

Shape check:

G: [B, out]
W: [out, in]
dX = G @ W: [B, in]
dW = G.T @ X: [out, in]
db = G.sum(dim=0): [out]

Shape is only a necessary condition; a transpose error might still silently pass in a square matrix, so a numerical gradient check is still needed.

5. Use PyTorch to Check Manual Gradient Calculations

This runnable Python 3 / PyTorch example uses float64 without changing the process-wide default dtype:

import torch

x = torch.tensor(2.0, dtype=torch.float64)
y = torch.tensor(1.0, dtype=torch.float64)
w1 = torch.tensor(0.5, dtype=torch.float64, requires_grad=True)
b1 = torch.tensor(0.1, dtype=torch.float64, requires_grad=True)
w2 = torch.tensor(-1.2, dtype=torch.float64, requires_grad=True)
b2 = torch.tensor(0.3, dtype=torch.float64, requires_grad=True)

z1 = w1 * x + b1
a1 = torch.relu(z1)
z2 = w2 * a1 + b2
loss = 0.5 * (z2 - y) ** 2
loss.backward()

print({
    "loss": loss.item(),
    "dw1": w1.grad.item(),
    "db1": b1.grad.item(),
    "dw2": w2.grad.item(),
    "db2": b2.grad.item(),
})

The expected values are loss=2.0402, dw1=4.848, db1=2.424, dw2=-2.222, and db2=-2.02, up to rounding. In normal grad mode, an operation is recorded if at least one input requires gradients. Non-leaf results are tracked too, but .grad is populated by default only for gradient-requiring leaf tensors; call retain_grad() to inspect an intermediate result.

6. .grad Default Accumulation

Reusing the torch import, these two losses have separate graphs that share the same leaf:

x = torch.tensor(2.0, requires_grad=True)
loss1 = x ** 2
loss2 = 3 * x
loss1.backward()
loss2.backward()
assert x.grad.item() == 7.0

The parameter’s .grad accumulates both contributions. This is useful for gradient accumulation, but means regular training must clear it before each update:

The training fragment below requires an optimizer and a fresh forward computation of loss; the already-consumed loss from above cannot be reused:

optimizer.zero_grad(set_to_none=True)
loss.backward()
optimizer.step()

set_to_none=True may have subtle differences in memory/optimizer behavior compared to filling gradients with zeros, and when using custom optimization logic, you should verify grad is None.

The default graph releases some saved tensors after backward; to perform another backward on the same graph, retain_graph=True might be needed, but frequent reliance on it often indicates problems with the graph lifecycle or loss organization.

7. Graph Control: detach, no_grad, inference_mode

• tensor.detach(): Returns a tensor associated with the original storage but not tracking the current history; in-place modifications should still be used with caution.
• torch.no_grad(): No reverse operations are recorded in the context, commonly used for validation/inference.
• torch.inference_mode(): Stronger inference optimization and constraints; generated tensors subsequent to participation in autograd are subject to additional constraints.

Detaching errors will truncate the path that should be trained; forgetting to disable gradients wastes memory during validation. Don't use .data to bypass autograd version checks.

8. In-place Operations and Saved Tensors

Backward might require inputs/outputs that were saved by forward. If these tensors are modified in place, the framework might throw a version mismatch error; more dangerous custom code could silently break mathematical semantics.

ReLU with inplace=True, in-place parameter updates, and views or shared storage require balancing memory efficiency against verifiability. Prioritize out-of-place expressions during initial learning or debugging.

9. Non-differentiable Points and Numerical Precision

ReLU is non-differentiable at 0, the absolute value at 0, and at max ties. The framework uses a subgradient or convention, so finite differences at these points may not align with autograd.

Float32 finite differences are still susceptible to cancellation and rounding errors; gradient checks typically use float64, a well-chosen ε, and avoid kinks.

10. Finite Difference Gradient Check

For the scalar parameter $\theta$:

$$g_{num}=\frac{L(\theta +ϵ)-L(\theta -ϵ)}{2ϵ}.$$

The relative error can be written as:

$$\frac{|g_{ana}-g_{num}|}{max(1,|g_{ana}|,|g_{num}|)}.$$

Check small models/low parameter counts, fix randomization, disable dropout, and avoid batch statistics changes. PyTorch torch.autograd.gradcheck uses numerical approximation to verify custom functions, with inputs typically requiring double precision and requires_grad=True.

A gradient check cannot prove the forward pass is correct; the forward and backward passes might independently implement the wrong formula, so reference cases and invariants are still needed.

11. Vanishing and Exploding Gradients

The chain rule in deep networks includes a Jacobian product:

$$\frac{\partial L}{\partial h_0}=J_1^T{J}_2^T\cdots J_L^T\frac{\partial L}{\partial h_L}.$$

If typical singular values are long-term less than 1, gradients decay; if greater than 1, gradients amplify. Factors affecting this include weight initialization, activation derivative, network depth, normalization, residual paths, and data scale.

ReLU eliminates the sigmoid saturation in the positive region but doesn't constrain the weight Jacobian. Residual connections add identity-like paths for gradients, normalization and initialization control scale, and gradient clipping bounds the gradient norm; adaptive optimizers and weight decay can still produce updates with a different scale; each addresses a different aspect.

12. Debugging Gradient Flow

Record it rather than guess:

• Parameter/gradient norm per layer;
• activation mean/std, zero/saturation fraction;
• Is loss finite;
• update-to-weight ratio;
• The first operation where NaN/Inf occurs;
• train/eval mode and dtype/device

You can use hooks, anomaly detection, and profilers to diagnose issues, but hooks can affect performance and the graph's lifecycle. First, overfit on a very small batch, verify that the loss decreases, then scale up training.

Common Misconceptions

• Backprop distributes the error evenly across parameters: it computes the loss derivative with respect to each parameter / VJP.
• Shared node gradients are computed only once and then finished: contributions from multiple downstream paths must be summed.
• Calling backward twice will overwrite .grad: by default, accumulation occurs.
• Autograd confirms the model is correct only indicates that the framework computed gradients in graph mode.
• No gradient vanishing problem with ReLU: The Jacobian chain can still vanish or explode.

Practice

1. Manually compute all forward values and gradients for a scalar network, then compare them with PyTorch.
2. Draw the branched computation graph for $L=x^2+3x$ and observe the accumulation of x.grad.
3. Derive the dX/dW/db for the batch affine layer and verify the shape.
4. Intentionally leave .grad uncleared and perform two backward passes, explain the results.
5. Perform a float64 central-difference check on custom activation functions.

Summary

Backpropagation traverses the computation graph in reverse topological order from a scalar loss, multiplying link contributions and adding branch contributions. Autograd automatically records and executes these local rules, but graph lifecycle, gradient accumulation, non-differentiable points, and numerical precision still require developer understanding.

Next lesson: turn correct gradients into stable training, control signal scales with initialization, the optimizer determines updates, and normalization and regularization constrain training and generalization.