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

The chain rule of calculus has finally connected to the production line of the Model Workshop. A single forward pass might generate millions of intermediate values; if we were to compute derivatives for each parameter from scratch, the calculations would involve massive redundancy. Backpropagation preserves local dependencies and then reuses the results 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_1x+b_1,quad a_1=\operatorname{ReLU}(z_1), $$

$$ z_2=w_2a_1+b_2,quad L=\frac12(z_2-y)^2. $$

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

$$ \bar z_2=z_2-y,qquad \bar w_2=\bar z_2a_1,qquad \bar b_2=\bar z_2, $$

$$ \bar a_1=\bar z_2w_2,qquad \bar z_1=\bar a_1\mathbf1[z_1>0], $$

$$ \bar w_1=\bar z_1x,qquad \bar b_1=\bar z_1. $$

The cross-line 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.

Why Reverse Mode Is Suitable for Training

The internet 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.

Why Are Branches Added Together?

If:

$$ u=x^2,qquad v=3x,qquad 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 add the residual, 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=XW^T+b,qquad G=\frac{\partial L}{\partial Z}. $$

Then:

$$ \frac{\partial L}{\partial X}=GW,qquad \frac{\partial L}{\partial W}=G^TX,qquad \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 Counter Manual Gradient Calculations

import torch

torch.set_default_dtype(torch.float64)

x = torch.tensor(2.0)
y = torch.tensor(1.0)
w1 = torch.tensor(0.5, requires_grad=True)
b1 = torch.tensor(0.1, requires_grad=True)
w2 = torch.tensor(-1.2, requires_grad=True)
b2 = torch.tensor(0.3, 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(),
})

Only requires_grad=True and tensors participating in the graph are tracked. Parameters are typically leaf tensors, and their gradients are accumulated into .grad. Intermediate tensors do not necessarily retain .grad by default; call retain_grad() when debugging, at the cost of additional memory.

6. .grad Default Accumulation

Calling backward consecutively:

loss1.backward()
loss2.backward()

The parameter .grad contributes twice. This is useful for gradient accumulation, but means regular training must clear it before each update:

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 padding, 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 reasoning 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+\epsilon)-L(\theta-\epsilon)}{2\epsilon}. $$

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^TJ_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 limits exploding updates; 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. Plot the bifurcation diagram 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 blur .grad 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.