Lecture 4 · Computation graphs, backpropagation, and autograd

This is the single complete companion to the lecture. It follows the same numbers and notation as the slides:

  1. local backward rules;
  2. one complete scalar graph;
  3. symbolic differentiation and finite differences as independent checks;
  4. the same graph in PyTorch autograd;
  5. gradient accumulation and zero_grad();
  6. branch accumulation and a vector neuron;
  7. one dense-layer vector-Jacobian product;
  8. a three-example mean batch and equivalent microbatches;
  9. one concrete optimizer step and a tiny MLP.

Evidence contract. Every input is a declared teaching construction. Every displayed result is freshly computed by the cells below. There is no dataset download, randomness, or hidden model training.

Checkpoint 1 · Setup and a compact display helper

The notebook uses only PyTorch and standard Colab/Jupyter display utilities. Double precision keeps the small arithmetic easy to compare with the slides.

import html
import math
import torch
from IPython.display import HTML, display

torch.set_default_dtype(torch.float64)

def fmt(value):
    if isinstance(value, torch.Tensor):
        value = value.detach().cpu().tolist()
    return html.escape(str(value))

def show_table(headers, rows, caption=None):
    head = "".join(f"<th>{html.escape(str(h))}</th>" for h in headers)
    body = "".join(
        "<tr>" + "".join(f"<td>{fmt(v)}</td>" for v in row) + "</tr>"
        for row in rows
    )
    cap = f"<caption>{html.escape(caption)}</caption>" if caption else ""
    display(HTML(
        "<style>div.l4-wrap{max-width:100%;overflow-x:auto;-webkit-overflow-scrolling:touch}"
        "table.l4{border-collapse:collapse;margin:.6em 0}"
        "table.l4 th,table.l4 td{border:1px solid #aeb8bd;padding:.32em .6em;"
        "text-align:left}table.l4 th{background:#edf4f4;color:#17383b}"
        "table.l4 caption{text-align:left;font-weight:700;margin-bottom:.35em}</style>"
        f"<div class='l4-wrap'><table class='l4'>{cap}<thead><tr>{head}</tr></thead>"
        f"<tbody>{body}</tbody></table></div>"
    ))

print("torch", torch.__version__)
print("default dtype", torch.get_default_dtype())
torch 2.13.0
default dtype torch.float64

Checkpoint 2 · Three local backward rules

If a later part of the graph sends an upstream gradient \(g_v=\partial L/\partial v\), each primitive returns \(g_v\) times its local derivative.

  • square: \(v=u^2 \Rightarrow g_u=g_v(2u)\);
  • addition: \(v=a+b \Rightarrow (g_a,g_b)=(g_v,g_v)\);
  • multiplication: \(v=ab \Rightarrow (g_a,g_b)=(g_vb,g_va)\).

The square example explicitly assumes the rest of the graph sends \(g_v=7\). That number is not produced by the square; it arrives from downstream.

# Square: the arriving upstream gradient is deliberately supplied as 7.
u = 3.0
g_v_square = 7.0
v_square = u**2
g_u_square = g_v_square * (2*u)

# Addition and multiplication use the same arriving gradient 3.
a_local, b_local, g_v_local = 2.0, 5.0, 3.0
g_a_add, g_b_add = g_v_local, g_v_local
g_a_mul = g_v_local * b_local
g_b_mul = g_v_local * a_local

assert (v_square, g_u_square) == (9.0, 42.0)
assert (g_a_add, g_b_add) == (3.0, 3.0)
assert (g_a_mul, g_b_mul) == (15.0, 6.0)

show_table(
    ["operation", "forward", "arriving gradient", "returned gradient(s)"],
    [
        ["square", "v = 3^2 = 9", "g_v = 7 (given)", "g_u = 7(2·3) = 42"],
        ["addition", "v = 2 + 5 = 7", "g_v = 3", "g_a = 3, g_b = 3"],
        ["multiplication", "v = 2·5 = 10", "g_v = 3", "g_a = 15, g_b = 6"],
    ],
    "Local rules: upstream × local",
)
Local rules: upstream × local
operation forward arriving gradient returned gradient(s)
square v = 3^2 = 9 g_v = 7 (given) g_u = 7(2·3) = 42
addition v = 2 + 5 = 7 g_v = 3 g_a = 3, g_b = 3
multiplication v = 2·5 = 10 g_v = 3 g_a = 15, g_b = 6

Checkpoint 3 · Forward through the complete scalar graph

We decompose

\[L=(wx+b-y)^2\]

into primitive values \(m=wx\), \(a=m+b\), \(e=a-y\), and \(L=e^2\).

x, w, b, y = 3.0, 2.0, 1.0, 10.0
m = w*x
a = m+b
e = a-y
L = e**2

assert (m, a, e, L) == (6.0, 7.0, -3.0, 9.0)
show_table(
    ["stored value", "calculation", "value"],
    [
        ["m", "w·x", m],
        ["a", "m+b", a],
        ["e", "a-y", e],
        ["L", "e^2", L],
    ],
    "Forward pass",
)
Forward pass
stored value calculation value
m w·x 6.0
a m+b 7.0
e a-y -3.0
L e^2 9.0

Checkpoint 4 · Reverse sweep through the same stored values

Seed \(g_L=\partial L/\partial L=1\). Then repeatedly apply upstream × local. When a value fans out, returned contributions add.

gL = 1.0
ge = gL * (2*e)       # L=e^2
ga = ge * 1.0         # e=a-y
gy = ge * (-1.0)
gm = ga * 1.0         # a=m+b
gb = ga * 1.0
gw = gm * x           # m=w*x
gx = gm * w

expected = {"L": 1.0, "e": -6.0, "a": -6.0, "m": -6.0,
            "w": -18.0, "x": -12.0, "b": -6.0, "y": 6.0}
actual = {"L": gL, "e": ge, "a": ga, "m": gm,
          "w": gw, "x": gx, "b": gb, "y": gy}
assert actual == expected

show_table(
    ["stored value q", "q", "g_q = ∂L/∂q"],
    [[name, {"L":L,"e":e,"a":a,"m":m,"w":w,"x":x,"b":b,"y":y}[name], grad]
     for name, grad in actual.items()],
    "All stored values and their gradients",
)
All stored values and their gradients
stored value q q g_q = ∂L/∂q
L 9.0 1.0
e -3.0 -6.0
a 7.0 -6.0
m 6.0 -6.0
w 2.0 -18.0
x 3.0 -12.0
b 1.0 -6.0
y 10.0 6.0

Checkpoint 5 · Two independent checks

Symbolic differentiation transforms a formula into a derivative formula; it can be done on paper or by a computer-algebra system. Finite differences instead probe nearby loss values and approximate one derivative number. Neither is how PyTorch performs reverse-mode autograd.

residual = w*x+b-y
symbolic = {
    "w": 2*residual*x,
    "x": 2*residual*w,
    "b": 2*residual,
    "y": -2*residual,
}

def loss_value(w_, x_, b_, y_):
    return (w_*x_+b_-y_)**2

eps = 1e-5
base = {"w": w, "x": x, "b": b, "y": y}
finite = {}
order = ["w", "x", "b", "y"]
for name in order:
    plus = base.copy(); minus = base.copy()
    plus[name] += eps; minus[name] -= eps
    finite[name] = (
        loss_value(plus["w"], plus["x"], plus["b"], plus["y"])
        - loss_value(minus["w"], minus["x"], minus["b"], minus["y"])
    ) / (2*eps)

manual = {"w": gw, "x": gx, "b": gb, "y": gy}
for name in order:
    assert symbolic[name] == manual[name]
    assert math.isclose(finite[name], manual[name], rel_tol=1e-9, abs_tol=1e-8)

show_table(
    ["q", "manual reverse", "symbolic formula at this point", "central difference"],
    [[name, manual[name], symbolic[name], f"{finite[name]:.8f}"] for name in order],
    "Three routes agree",
)
Three routes agree
q manual reverse symbolic formula at this point central difference
w -18.0 -18.0 -18.00000000
x -12.0 -12.0 -12.00000000
b -6.0 -6.0 -6.00000000
y 6.0 6.0 6.00000000

Checkpoint 6 · One gradient step lowers this loss

Backprop computes gradients; an optimizer later uses them. For this tiny example, one gradient-descent step with \(\eta=0.01\) is easy to audit.

eta = 0.01
w_new = w - eta*gw
b_new = b - eta*gb
prediction_new = w_new*x+b_new
loss_new = (prediction_new-y)**2

assert math.isclose(w_new, 2.18)
assert math.isclose(b_new, 1.06)
assert math.isclose(prediction_new, 7.6)
assert math.isclose(loss_new, 5.76)
assert loss_new < L

print(f"w: {w:.2f} -> {w_new:.2f}")
print(f"b: {b:.2f} -> {b_new:.2f}")
print(f"loss: {L:.2f} -> {loss_new:.2f}")
w: 2.00 -> 2.18
b: 1.00 -> 1.06
loss: 9.00 -> 5.76

Checkpoint 7 · PyTorch records and replays the same scalar graph

retain_grad() lets us inspect intermediate non-leaf gradients for teaching. Ordinary training normally keeps only parameter gradients.

x_t = torch.tensor(3.0, requires_grad=True)
w_t = torch.tensor(2.0, requires_grad=True)
b_t = torch.tensor(1.0, requires_grad=True)
y_t = torch.tensor(10.0, requires_grad=True)

m_t = w_t*x_t
a_t = m_t+b_t
e_t = a_t-y_t
L_t = e_t**2
for node in (m_t, a_t, e_t, L_t):
    node.retain_grad()
L_t.backward()

torch_expected = {
    "L": (L_t, L_t.grad, 9.0, 1.0),
    "e": (e_t, e_t.grad, -3.0, -6.0),
    "a": (a_t, a_t.grad, 7.0, -6.0),
    "m": (m_t, m_t.grad, 6.0, -6.0),
    "w": (w_t, w_t.grad, 2.0, -18.0),
    "x": (x_t, x_t.grad, 3.0, -12.0),
    "b": (b_t, b_t.grad, 1.0, -6.0),
    "y": (y_t, y_t.grad, 10.0, 6.0),
}
for _, (node, grad, value_expected, grad_expected) in torch_expected.items():
    torch.testing.assert_close(node.detach(), torch.tensor(value_expected))
    torch.testing.assert_close(grad, torch.tensor(grad_expected))

show_table(
    ["node", "stored value", ".grad after backward()"],
    [[name, float(node.detach()), float(grad)]
     for name, (node, grad, _, _) in torch_expected.items()],
    "PyTorch reproduces the complete ledger",
)
PyTorch reproduces the complete ledger
node stored value .grad after backward()
L 9.0 1.0
e -3.0 -6.0
a 7.0 -6.0
m 6.0 -6.0
w 2.0 -18.0
x 3.0 -12.0
b 1.0 -6.0
y 10.0 6.0

Checkpoint 8 · Gradients accumulate until we clear them

Each backward() adds into a leaf’s .grad. A training loop therefore starts each intended update window with zero_grad() (or sets gradients to None).

w_acc = torch.tensor(2.0, requires_grad=True)
history = []
for backward_index in range(1, 4):
    loss_i = (w_acc*3.0+1.0-10.0)**2  # fresh graph, same leaf w_acc
    loss_i.backward()
    history.append(float(w_acc.grad))

assert history == [-18.0, -36.0, -54.0]
w_acc.grad.zero_()
assert float(w_acc.grad) == 0.0

show_table(
    ["backward call", "w.grad"],
    [[i, value] for i, value in enumerate(history, 1)],
    "Accumulation is addition, not replacement",
)
print("after zeroing: w.grad =", float(w_acc.grad))
Accumulation is addition, not replacement
backward call w.grad
1 -18.0
2 -36.0
3 -54.0 
after zeroing: w.grad = 0.0

Checkpoint 9 · At a branch, gradient contributions add

The value \(x=4\) is used along two paths:

\[u=x^2,\qquad v=3x,\qquad L=u+v.\]

The two returned contributions are \(2x=8\) and \(3\), so autograd must add them at the shared leaf: \(x.\text{grad}=8+3=11\).

x_branch = torch.tensor(4.0, requires_grad=True)
u_branch = x_branch**2
v_branch = 3*x_branch
L_branch = u_branch+v_branch
L_branch.backward()

assert tuple(float(t.detach()) for t in (u_branch, v_branch, L_branch)) == (16.0, 12.0, 28.0)
torch.testing.assert_close(x_branch.grad, torch.tensor(11.0))
show_table(
    ["path", "local return", "contribution to x.grad"],
    [
        ["u=x^2", "1·2x", 8.0],
        ["v=3x", "1·3", 3.0],
        ["shared x", "add both paths", float(x_branch.grad)],
    ],
    "Autograd accumulates at a shared value",
)
Autograd accumulates at a shared value
path local return contribution to x.grad
u=x^2 1·2x 8.0
v=3x 1·3 3.0
shared x add both paths 11.0

Checkpoint 10 · One vector neuron is still one affine operation

For \(z=w^Tx+b\) and \(L=z^2\), the arriving scalar is \(g_z=2z\). The affine operation returns \(g_w=g_zx\), \(g_x=g_zw\), and \(g_b=g_z\).

x_vec = torch.tensor([2.0, -1.0], requires_grad=True)
w_vec = torch.tensor([1.0, 3.0], requires_grad=True)
b_vec = torch.tensor(0.0, requires_grad=True)
z_vec = w_vec@x_vec+b_vec
L_vec = z_vec**2
L_vec.backward()

torch.testing.assert_close(z_vec, torch.tensor(-1.0))
torch.testing.assert_close(L_vec, torch.tensor(1.0))
torch.testing.assert_close(w_vec.grad, torch.tensor([-4.0, 2.0]))
torch.testing.assert_close(x_vec.grad, torch.tensor([-2.0, -6.0]))
torch.testing.assert_close(b_vec.grad, torch.tensor(-2.0))
show_table(
    ["quantity", "value", "gradient"],
    [
        ["w", w_vec.detach().tolist(), w_vec.grad.tolist()],
        ["x", x_vec.detach().tolist(), x_vec.grad.tolist()],
        ["b", float(b_vec.detach()), float(b_vec.grad)],
        ["z", float(z_vec.detach()), "g_z = 2z = -2"],
    ],
    "Vector dot-product example from the slides",
)
Vector dot-product example from the slides
quantity value gradient
w [1.0, 3.0] [-4.0, 2.0]
x [2.0, -1.0] [-2.0, -6.0]
b 0.0 -2.0
z -1.0 g_z = 2z = -2

Checkpoint 11 · Dense backward is the scalar affine rule repeated by row

For column vectors, \(z=Wx+b\). A later graph sends \(g_z=(4,-2)^T\). Coordinate calculus gives

\[g_x=W^Tg_z,\qquad g_W=g_zx^T,\qquad g_b=g_z.\]

The transpose is not a memorized decoration: the same input feeds both output rows, so their returned contributions add at \(x\).

x_dense = torch.tensor([2.0, -1.0])
W_dense = torch.tensor([[1.0, 3.0], [-2.0, 1.0]])
b_dense = torch.tensor([0.0, 1.0])
g_z = torch.tensor([4.0, -2.0])
z_dense = W_dense@x_dense+b_dense

gx_rows = g_z[:, None]*W_dense
gW_rows = g_z[:, None]*x_dense[None, :]
g_x_dense = gx_rows.sum(dim=0)
g_W_dense = gW_rows
g_b_dense = g_z

torch.testing.assert_close(z_dense, torch.tensor([-1.0, -4.0]))
torch.testing.assert_close(gx_rows, torch.tensor([[4.0, 12.0], [4.0, -2.0]]))
torch.testing.assert_close(g_x_dense, torch.tensor([8.0, 10.0]))
torch.testing.assert_close(g_W_dense, torch.tensor([[8.0, -4.0], [-4.0, 2.0]]))
torch.testing.assert_close(g_b_dense, torch.tensor([4.0, -2.0]))

# Verify the same vector-Jacobian product with autograd.
x_leaf = x_dense.clone().requires_grad_()
W_leaf = W_dense.clone().requires_grad_()
b_leaf = b_dense.clone().requires_grad_()
z_leaf = W_leaf@x_leaf+b_leaf
z_leaf.backward(g_z)
torch.testing.assert_close(x_leaf.grad, g_x_dense)
torch.testing.assert_close(W_leaf.grad, g_W_dense)
torch.testing.assert_close(b_leaf.grad, g_b_dense)

show_table(
    ["recipient", "row-level return", "stacked result", "shape"],
    [
        ["x", gx_rows.tolist(), g_x_dense.tolist(), "(2,)"],
        ["W", gW_rows.tolist(), g_W_dense.tolist(), "(2,2)"],
        ["b", g_z.tolist(), g_b_dense.tolist(), "(2,)"],
    ],
    "Dense VJP: derive by coordinate, then verify shapes",
)
Dense VJP: derive by coordinate, then verify shapes
recipient row-level return stacked result shape
x [[4.0, 12.0], [4.0, -2.0]] [8.0, 10.0] (2,)
W [[8.0, -4.0], [-4.0, 2.0]] [[8.0, -4.0], [-4.0, 2.0]] (2,2)
b [4.0, -2.0] [4.0, -2.0] (2,)

Checkpoint 12 · A batch adds an example axis, not new parameters

PyTorch stores examples as rows \(X\in\mathbb R^{B\times d}\), so \(Z=XW^T+b\). With half-squared error per example,

\[\ell_n=\tfrac12\|z^{(n)}-y^{(n)}\|^2,\qquad L=\tfrac1B\sum_n\ell_n,\]

the upstream batch gradient is \(G_Z=(Z-Y)/B\).

X = torch.tensor([[2.0, -1.0], [-1.0, 2.0], [2.0, 2.0]])
Y = torch.tensor([[-5.0, -2.0], [7.0, 1.0], [7.0, -2.0]])
B = X.shape[0]
Z = X@W_dense.T+b_dense
R = Z-Y
loss_each = 0.5*(R**2).sum(dim=1)
loss_mean = loss_each.mean()
G_Z = R/B

torch.testing.assert_close(Z, torch.tensor([[-1.0,-4.0],[5.0,5.0],[8.0,-1.0]]))
torch.testing.assert_close(R, torch.tensor([[4.0,-2.0],[-2.0,4.0],[1.0,1.0]]))
torch.testing.assert_close(loss_each, torch.tensor([10.0,10.0,1.0]))
torch.testing.assert_close(loss_mean, torch.tensor(7.0))
torch.testing.assert_close(G_Z, R/3)

show_table(
    ["n", "x^(n)", "z^(n)", "y^(n)", "r^(n)", "ell_n"],
    [[n+1, X[n].tolist(), Z[n].tolist(), Y[n].tolist(), R[n].tolist(), float(loss_each[n])]
     for n in range(B)],
    "Three examples reuse the same W and b",
)
print("mean loss L =", float(loss_mean))
print("Z.grad should be R/B =\n", G_Z)
Three examples reuse the same W and b
n x^(n) z^(n) y^(n) r^(n) ell_n
1 [2.0, -1.0] [-1.0, -4.0] [-5.0, -2.0] [4.0, -2.0] 10.0
2 [-1.0, 2.0] [5.0, 5.0] [7.0, 1.0] [-2.0, 4.0] 10.0
3 [2.0, 2.0] [8.0, -1.0] [7.0, -2.0] [1.0, 1.0] 1.0 
mean loss L = 7.0
Z.grad should be R/B =
 tensor([[ 1.3333, -0.6667],
        [-0.6667,  1.3333],
        [ 0.3333,  0.3333]])

Checkpoint 13 · Shared parameter paths add; the mean then scales

Each example proposes one outer product \(r^{(n)}(x^{(n)})^T\) to the same \(W\). The batch gradient adds those proposals and divides by \(B\). The three input rows are distinct, so their gradients remain separate.

gW_each_unscaled = torch.stack([torch.outer(R[n], X[n]) for n in range(B)])
gW_batch = gW_each_unscaled.mean(dim=0)
gb_batch = R.mean(dim=0)
gX_batch = G_Z@W_dense

expected_gW_each = torch.tensor([
    [[8.0,-4.0],[-4.0,2.0]],
    [[2.0,-4.0],[-4.0,8.0]],
    [[2.0,2.0],[2.0,2.0]],
])
torch.testing.assert_close(gW_each_unscaled, expected_gW_each)
torch.testing.assert_close(gW_batch, torch.tensor([[4.0,-2.0],[-2.0,4.0]]))
torch.testing.assert_close(gb_batch, torch.tensor([1.0,1.0]))
torch.testing.assert_close(gX_batch, torch.tensor([
    [8/3,10/3],[-10/3,-2/3],[-1/3,4/3]
]))

# Full-batch autograd verification, including the non-leaf Z gradient.
X_full = X.clone().requires_grad_()
W_full = W_dense.clone().requires_grad_()
b_full = b_dense.clone().requires_grad_()
Z_full = X_full@W_full.T+b_full
Z_full.retain_grad()
L_full = 0.5*((Z_full-Y)**2).sum(dim=1).mean()
L_full.backward()

torch.testing.assert_close(L_full, loss_mean)
torch.testing.assert_close(Z_full.grad, G_Z)
torch.testing.assert_close(W_full.grad, gW_batch)
torch.testing.assert_close(b_full.grad, gb_batch)
torch.testing.assert_close(X_full.grad, gX_batch)

show_table(
    ["quantity", "manual batch rule", "autograd result", "shape"],
    [
        ["Z.grad", G_Z.tolist(), Z_full.grad.tolist(), tuple(Z_full.grad.shape)],
        ["W.grad", gW_batch.tolist(), W_full.grad.tolist(), tuple(W_full.grad.shape)],
        ["b.grad", gb_batch.tolist(), b_full.grad.tolist(), tuple(b_full.grad.shape)],
        ["X.grad", gX_batch.tolist(), X_full.grad.tolist(), tuple(X_full.grad.shape)],
    ],
    "Manual batch arithmetic equals PyTorch autograd",
)
Manual batch arithmetic equals PyTorch autograd
quantity manual batch rule autograd result shape
Z.grad [[1.3333333333333333, -0.6666666666666666], [-0.6666666666666666, 1.3333333333333333], [0.3333333333333333, 0.3333333333333333]] [[1.3333333333333333, -0.6666666666666666], [-0.6666666666666666, 1.3333333333333333], [0.3333333333333333, 0.3333333333333333]] (3, 2)
W.grad [[4.0, -2.0], [-2.0, 4.0]] [[3.9999999999999996, -2.0], [-2.0, 3.9999999999999996]] (2, 2)
b.grad [1.0, 1.0] [1.0, 1.0] (2,)
X.grad [[2.6666666666666665, 3.3333333333333335], [-3.333333333333333, -0.6666666666666667], [-0.3333333333333333, 1.3333333333333333]] [[2.6666666666666665, 3.3333333333333335], [-3.333333333333333, -0.6666666666666667], [-0.3333333333333333, 1.3333333333333333]] (3, 2)

Checkpoint 14 · Microbatches reproduce the same mean gradient

Clear once, backpropagate each \(\ell_n/B\), and update once. This equals one full-batch backward when parameters stay fixed and the reduction is scaled consistently. It need not remain equivalent with batch-coupled operations such as BatchNorm.

W_micro = W_dense.clone().requires_grad_()
b_micro = b_dense.clone().requires_grad_()
snapshots = []

for n in range(B):
    z_n = W_micro@X[n]+b_micro
    scaled_loss_n = 0.5*((z_n-Y[n])**2).sum()/B
    scaled_loss_n.backward()
    snapshots.append(W_micro.grad.detach().clone())

torch.testing.assert_close(W_micro.grad, W_full.grad)
torch.testing.assert_close(b_micro.grad, b_full.grad)

show_table(
    ["after microbatch", "accumulated W.grad"],
    [[n+1, snapshots[n].tolist()] for n in range(B)],
    "The running sum reaches the full mean-batch gradient",
)
print("full-batch W.grad =\n", W_full.grad)
The running sum reaches the full mean-batch gradient
after microbatch accumulated W.grad
1 [[2.6666666666666665, -1.3333333333333333], [-1.3333333333333333, 0.6666666666666666]]
2 [[3.333333333333333, -2.6666666666666665], [-2.6666666666666665, 3.333333333333333]]
3 [[3.9999999999999996, -2.0], [-2.0, 3.9999999999999996]] 
full-batch W.grad =
 tensor([[ 4.0000, -2.0000],
        [-2.0000,  4.0000]])

Checkpoint 15 · A concrete training step uses the batch gradient

Now the familiar loop is executable: clear, forward, reduce to one scalar, backward, and update. With learning rate \(0.01\), the same fixed batch moves from loss \(7\) to \(6.5865\).

W_train = W_dense.clone().requires_grad_()
b_train = b_dense.clone().requires_grad_()
optimizer = torch.optim.SGD([W_train, b_train], lr=0.01)

optimizer.zero_grad()
Z_train = X@W_train.T+b_train
loss_train = 0.5*((Z_train-Y)**2).sum(dim=1).mean()
loss_train.backward()
optimizer.step()

with torch.no_grad():
    loss_after = 0.5*((X@W_train.T+b_train-Y)**2).sum(dim=1).mean()

torch.testing.assert_close(W_train, torch.tensor([[0.96,3.02],[-1.98,0.96]]))
torch.testing.assert_close(b_train, torch.tensor([-0.01,0.99]))
torch.testing.assert_close(loss_after, torch.tensor(6.5865))
assert loss_after < loss_train

show_table(
    ["quantity", "before", "after one SGD step"],
    [
        ["W", W_dense.tolist(), W_train.detach().tolist()],
        ["b", b_dense.tolist(), b_train.detach().tolist()],
        ["mean loss", float(loss_train.detach()), float(loss_after)],
    ],
    "clear → forward → scalar loss → backward → update",
)
clear → forward → scalar loss → backward → update
quantity before after one SGD step
W [[1.0, 3.0], [-2.0, 1.0]] [[0.96, 3.02], [-1.98, 0.96]]
b [0.0, 1.0] [-0.01, 0.99]
mean loss 7.0 6.586499999999998

Checkpoint 16 · A tiny MLP is the same three-block pattern

The slides finish with dense \(\rightarrow\) ReLU \(\rightarrow\) dense. We set deterministic weights so every forward value is auditable, then let PyTorch return one gradient for every parameter. The inactive hidden unit returns zero.

from torch import nn

layer1 = nn.Linear(2, 3)
layer2 = nn.Linear(3, 2)
with torch.no_grad():
    layer1.weight.copy_(torch.tensor([[1.0,-0.5],[0.0,0.0],[1.0,-1.0]]))
    layer1.bias.copy_(torch.tensor([0.0,-0.5,0.0]))
    layer2.weight.copy_(torch.tensor([[0.4,0.3,1.0],[0.7,-0.2,0.0]]))
    layer2.bias.zero_()

x_mlp = torch.tensor([2.0,-1.0])
y_mlp = torch.tensor([1.0,0.0])
preactivation = layer1(x_mlp)
preactivation.retain_grad()
hidden = torch.relu(preactivation)
prediction = layer2(hidden)
loss_mlp = ((prediction-y_mlp)**2).sum()
loss_mlp.backward()

torch.testing.assert_close(preactivation.detach(), torch.tensor([2.5,-0.5,3.0]))
torch.testing.assert_close(hidden.detach(), torch.tensor([2.5,0.0,3.0]))
torch.testing.assert_close(prediction.detach(), torch.tensor([4.0,1.75]))
torch.testing.assert_close(loss_mlp.detach(), torch.tensor(12.0625))
torch.testing.assert_close(preactivation.grad[1], torch.tensor(0.0))
torch.testing.assert_close(layer1.weight.grad[1], torch.zeros(2))
torch.testing.assert_close(layer1.bias.grad[1], torch.tensor(0.0))

for name, parameter in [
    ("layer1.weight", layer1.weight), ("layer1.bias", layer1.bias),
    ("layer2.weight", layer2.weight), ("layer2.bias", layer2.bias),
]:
    assert parameter.grad is not None
    assert parameter.grad.shape == parameter.shape
    assert torch.isfinite(parameter.grad).all()

show_table(
    ["stage", "value"],
    [
        ["dense 1 preactivation", preactivation.detach().tolist()],
        ["ReLU", hidden.detach().tolist()],
        ["dense 2 prediction", prediction.detach().tolist()],
        ["squared-error loss", float(loss_mlp.detach())],
    ],
    "Dense → ReLU → dense",
)
show_table(
    ["parameter", "shape", "gradient shape"],
    [[name, tuple(parameter.shape), tuple(parameter.grad.shape)] for name, parameter in [
        ("layer1.weight", layer1.weight), ("layer1.bias", layer1.bias),
        ("layer2.weight", layer2.weight), ("layer2.bias", layer2.bias),
    ]],
    "Autograd follows the recorded blocks in reverse",
)
Dense → ReLU → dense
stage value
dense 1 preactivation [2.5, -0.5, 3.0]
ReLU [2.5, 0.0, 3.0]
dense 2 prediction [4.0, 1.75]
squared-error loss 12.0625
Autograd follows the recorded blocks in reverse
parameter shape gradient shape
layer1.weight (3, 2) (3, 2)
layer1.bias (3,) (3,)
layer2.weight (2, 3) (2, 3)
layer2.bias (2,) (2,)

Takeaway

  • Reverse-mode backprop repeatedly applies upstream × local and adds at shared values.
  • Symbolic differentiation derives a formula; finite differences approximate a point derivative; autograd records the executed graph and runs the reverse sweep.
  • A dense layer repeats the scalar affine rule row by row.
  • A batch adds an example axis. Shared parameter gradients add across examples; the loss reduction determines whether the final result is a sum or a mean.
  • Multiple backward calls accumulate into .grad; clear once per intended update window and step once after all correctly scaled contributions arrive.
  • A deeper MLP is not a new differentiation idea: each dense or activation block receives one arriving gradient and returns the next one.