Autograd intuition

What autograd does for you

When you call loss.backward() in PyTorch, something computes gradients for every parameter that contributed to the loss. That something is automatic differentiation (autograd) — and you do not need to implement it from scratch, but you absolutely need to know what it is doing when training misbehaves.

Autograd builds a computation graph during the forward pass. Each operation (add, multiply, matmul, ReLU) is a node. Each node knows how to propagate gradients backward to its inputs. When you call .backward(), the graph is traversed in reverse, applying the chain rule at each step.

Without autograd, you would hand-derive partial derivatives for every layer in your network — impractical for millions of parameters. With autograd, you write forward code; gradients follow automatically.

Callout — autograd is not magic, it is bookkeeping: Every tensor with requires_grad=True registers operations. The backward pass is deterministic given the forward graph. If gradients are wrong, the forward graph or your detach/no_grad usage is usually the culprit.

Computation graphs: a concrete example

Trace this tiny graph by hand:

x = 2.0    (input, requires grad)
w = 3.0    (parameter, requires grad)
b = 1.0    (parameter, requires grad)

# Forward pass
z = w * x      # z = 6.0
y = z + b      # y = 7.0
L = y ** 2     # L = 49.0

The graph looks like:

x ──→ (*) ──→ z ──→ (+) ──→ y ──→ (^2) ──→ L
w ──↗              b ──↗

To compute ∂L/∂w, apply the chain rule backward:

1. ∂L/∂y = 2y = 14
2. ∂y/∂z = 1
3. ∂z/∂w = x = 2

So ∂L/∂w = 14 × 1 × 2 = 28.

Similarly:

• ∂L/∂b = 14 × 1 = 14
• ∂L/∂x = 14 × 1 × 3 = 42

Verify in PyTorch:

import torch
x = torch.tensor(2.0, requires_grad=True)
w = torch.tensor(3.0, requires_grad=True)
b = torch.tensor(1.0, requires_grad=True)

z = w * x
y = z + b
L = y ** 2
L.backward()

print(w.grad)  # 28.0
print(b.grad)  # 14.0
print(x.grad)  # 42.0

This is the entire mechanism scaled up: each operation contributes its local derivative; the chain rule multiplies them along paths from L back to each parameter.

The chain rule is the whole game

For composed functions f(g(h(x))), the derivative is:

df/dx = (df/dg) × (dg/dh) × (dh/dx)

In a neural network, f might be the loss, and the composition might be: loss → softmax → linear layer → ReLU → linear layer → input. Backprop is just the chain rule applied systematically from output to input.

Each operation implements two methods conceptually:

• Forward: compute output from inputs.
• Backward: given gradient of loss w.r.t. output, compute gradient w.r.t. each input.

PyTorch calls these forward and backward on Function objects under the hood. You rarely write them — but when you write a custom layer, you do.

A slightly larger example: two paths

When a tensor feeds into multiple branches, gradients add at that tensor:

a ──→ (+) ──→ c ──→ (×) ──→ L
b ──↗       d ──↗

If c = a + b and L = c × d, then:

• ∂L/∂c = d
• ∂L/∂d = c
• ∂L/∂a = ∂L/∂c × ∂c/∂a = d × 1 = d
• ∂L/∂b = d × 1 = d

Both a and b receive the same gradient from c. This "gradient accumulation at forks" matters when you have residual connections and attention — multiple paths from output back to the same weight.

What PyTorch tracks (and what it ignores)

Tracked: tensors with requires_grad=True, and operations connecting them.

Not tracked:

• Tensors with requires_grad=False (default for data inputs unless you set otherwise).
• Operations inside with torch.no_grad(): — used for inference and evaluation.
• .detach() — breaks the graph; treat the value as a constant for gradient purposes.

Common bugs:

Always check param.grad is not None and param.grad.abs().sum() > 0 on the first training step of a new architecture.

Vectors, matrices, and Jacobian intuition

For scalar loss L, gradients w.r.t. a vector parameter w have the same shape as w. Each element of w.grad tells you how L changes if you nudge that one weight.

For matrix operations (like Y = X @ W), the backward pass involves matrix transposes — the chain rule in matrix form. You do not need to memorize Jacobian formulas; autograd handles them. But knowing that grad shapes match parameter shapes prevents silent bugs when you hand-roll a layer.

If you ever implement a custom autograd function:

class MyFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input):
        ctx.save_for_backward(input)
        return input ** 2

    @staticmethod
    def backward(ctx, grad_output):
        input, = ctx.saved_tensors
        return grad_output * 2 * input

The backward receives grad_output (∂L/∂output) and returns gradients for each forward input.

Autograd vs manual gradients

Why bother with hand traces if PyTorch exists?

1. Debugging — when loss is wrong, hand-tracing a tiny example isolates the bug.
2. Interviews and collaboration — "walk me through backprop for one layer" still appears in ML discussions.
3. Custom operations — CUDA kernels, novel attention variants, and research code sometimes need custom backward.
4. Numerical checks — compare autograd to finite differences:

eps = 1e-5
numeric = (L(w + eps) - L(w - eps)) / (2 * eps)
assert torch.allclose(w.grad, numeric, atol=1e-3)

Finite-difference gradient checks are slow but invaluable when you write custom layers.

Engineering problem (staff framing)

Forward builds graph; backward applies chain rule. Required for NaN hunts and interviews.

Diagram — Forward/backward

sequenceDiagram
  participant X as Input
  participant F as Forward
  participant L as Loss
  participant B as Backward
  participant W as Weights
  X->>F: tensors
  F->>L: pred
  L->>B: dL
  B->>W: grads
  W->>W: step

Precise definitions & mental model

Computational graph, requires_grad, no_grad/detach, shape discipline.

Tradeoffs — when to use what

Eager debug-friendly vs compiled speed.

Failure modes (interview + on-call)

In-place version errors; forgot zero_grad; retaining graph → OOM.

Production & OSS practices

Anomaly detect NaNs; GradScaler; activation checkpointing.

Deep dive (FAANG / OSS bar)

Push «autograd-intuition» past tutorial depth: write the interface contract (inputs/outputs/invariants), list three measurable metrics, and name two degrade modes if the happy path fails. Add a short threat note: what an attacker or noisy tool result could do, and which layer catches it (schema, policy, HITL, or eval gate).

flowchart LR
  Contract[Interface contract] --> Metrics
  Metrics --> Degrade[Degrade modes]
  Degrade --> Threat[Threat + control]

Micro-project: Mini autograd exercises

In m2/autograd/ (or your course-portfolio monorepo under the matching module folder):

Complete these exercises and commit solutions.py plus a brief REPORT.md:

1. Hand trace — Given the graph L = (w1*x + w2*x**2)**2 with x=1, w1=2, w2=3, compute ∂L/∂w1 and ∂L/∂w2 by hand. Verify with PyTorch.

2. Gradient check — Implement a 1-hidden-layer MLP forward pass. Compare autograd gradients to finite differences for one random weight.

3. Broken graph debug — The starter code trains but loss never decreases. Find the line that breaks the graph (hint: common mistakes include detaching tensors, using .item() in the loss path, or freezing the wrong module). Fix it and show before/after loss for 10 steps.

4. Accumulation — Run two backward passes without zero_grad() in between. Show that param.grad is the sum of both. Explain why optimizer.zero_grad() is required each step.

Treat the objectives as acceptance criteria: you should be able to explain each result without looking at the code.

Checklist

• Hand-traced gradients match PyTorch for at least one toy graph
• Finite-difference check passes on a small MLP weight
• Broken-graph exercise fixed with explanation in REPORT.md
• You can explain when to use no_grad and detach