One smooth curve will stay with us while the model learns a new geometric move: the bend.
A straight line has one slope. Our target changes direction. We will give a model several simple hinges, watch them combine, and then let gradient descent move them. The point is not merely to obtain a small error. It is to see what the network is made of.
If the badge does not open, use this direct Colab link or upload the downloaded notebook with File → Upload notebook.
This notebook follows five steps:
Expected reading and execution time: 75–85 minutes. Run the cells in order and pause at each Checkpoint before revealing the next figure.
By the end, you should be able to:
The experiment uses only PyTorch, NumPy, and Matplotlib. It is deterministic, CPU-only, and small enough to rerun during class.
import math
import time
from pathlib import Path
import matplotlib.pyplot as plt
from matplotlib.patches import Circle, FancyArrowPatch
import numpy as np
import torch
import torch.nn as nn
started = time.perf_counter()
SEED = 2026
DEVICE = torch.device("cpu")
torch.manual_seed(SEED)
np.random.seed(SEED)
torch.use_deterministic_algorithms(True)
print(f"device={DEVICE} | seed={SEED}")device=cpu | seed=2026We observe input–target pairs \((x_i,y_i)\) from one smooth function:
\[ y_i=f(x_i),\qquad f(x)=\sin(\pi x)+0.25x,\qquad x\in[-1,1]. \]
The sine term bends down and up; the small linear term tilts the whole curve. Every model in this chapter sees the same 241 points.
def target_function(x):
return torch.sin(torch.pi * x) + 0.25 * x
x_train = torch.linspace(-1.0, 1.0, 241, device=DEVICE).reshape(-1, 1)
y_train = target_function(x_train)
x_plot = x_train.clone()
y_target = y_train.clone()
assert x_train.shape == y_train.shape == (241, 1)
assert torch.isfinite(y_train).all()
print("x shape:", tuple(x_train.shape), "| y shape:", tuple(y_train.shape))x shape: (241, 1) | y shape: (241, 1)The mathematical and training code remains in the story below. Plot construction lives here so that each later implementation cell stays short. These cells are collapsed in Colab and hidden in the rendered chapter; expand them only if you want to change the visual design. No external drawing package is required—even the network diagram is ordinary Matplotlib.
# @title Auxiliary: style and figure export
ROOT = Path.cwd().parent if Path.cwd().name == "tutorials" else Path.cwd()
FIG_DIR = ROOT / "assets/figures/shared/week04"
FIG_DIR.mkdir(parents=True, exist_ok=True)
TARGET_COLOR = "#202A44"
MODEL_COLOR = "#E45756"
HINGE_COLORS = ["#2F6BFF", "#00A896", "#F2A541", "#8F5DA2"]
plt.rcParams.update({"figure.dpi": 110, "axes.spines.top": False,
"axes.spines.right": False, "axes.titleweight": "bold",
"svg.hashsalt": str(SEED)})
figure_names = []
def finish_figure(fig, name):
svg_path = FIG_DIR / f"{name}.svg"
fig.savefig(svg_path, dpi=180, bbox_inches="tight", metadata={"Date": None})
svg_path.write_text("\n".join(line.rstrip() for line in svg_path.read_text().splitlines()) + "\n")
fig.savefig(FIG_DIR / f"{name}.png", dpi=180, bbox_inches="tight")
figure_names.append(name)
plt.show()# @title Auxiliary: Scene 1 visualisation
def plot_affine_attempt(x, target, prediction, mse):
fig, ax = plt.subplots(figsize=(8.2, 4.6))
ax.plot(x[:, 0], target[:, 0], color=TARGET_COLOR, lw=3, label="target f(x)")
ax.plot(x[:, 0], prediction[:, 0], color=MODEL_COLOR, lw=2.5, label="best affine fit")
ax.fill_between(x[:, 0], target[:, 0], prediction[:, 0], color=MODEL_COLOR, alpha=0.12, label="structured residual")
ax.axhline(0, color="0.75", lw=1)
ax.set(xlabel="x", ylabel="value", title="A line can tilt, but it cannot change its slope")
ax.text(0.02, 0.05, f"MSE = {mse:.4f}", transform=ax.transAxes, color="0.35")
ax.legend(frameon=False, ncol=3, loc="upper center")
finish_figure(fig, "01-a-straight-line-tries")# @title Auxiliary: Scene 2 visualisation
def plot_one_hinge(x, target, z, h, contribution, w, b, a, breakpoint):
fig, axes = plt.subplots(1, 2, figsize=(11.0, 4.2), sharex=True)
axes[0].plot(x[:, 0], z, ls="--", color="0.45", lw=2, label="preactivation z=wx+b")
axes[0].plot(x[:, 0], h, color=HINGE_COLORS[0], lw=3, label="h=ReLU(z)")
axes[0].axhline(0, color="0.75", lw=1)
axes[0].set(title=f"Inside: w={w:g}, b={b:g}", xlabel="x", ylabel="value")
axes[0].legend(frameon=False)
axes[1].plot(x[:, 0], target[:, 0], color=TARGET_COLOR, lw=2.5, alpha=0.45, label="unchanged target")
axes[1].plot(x[:, 0], contribution, color=MODEL_COLOR, lw=3, label="a ReLU(wx+b)")
axes[1].axhline(0, color="0.75", lw=1)
axes[1].set(title=f"Outside: a={a:g} scales or reverses it", xlabel="x", ylabel="value")
axes[1].legend(frameon=False)
if breakpoint is not None and -1 <= breakpoint <= 1:
for ax in axes:
ax.axvline(breakpoint, color="black", ls=":", lw=1.5)
axes[0].annotate(r"breakpoint $-b/w$", xy=(breakpoint, 0), xytext=(breakpoint + 0.12, -0.55), arrowprops={"arrowstyle": "->"})
fig.suptitle("One neuron contributes one movable bend", fontsize=14, fontweight="bold")
fig.tight_layout()
finish_figure(fig, "02-one-neuron-creates-a-bend")# @title Auxiliary: hidden-layer network diagram
def plot_hidden_layer_network(w, b, a, output_bias):
fig, ax = plt.subplots(figsize=(10.5, 5.2))
ax.set_xlim(0, 1); ax.set_ylim(0, 1); ax.axis("off")
input_xy, output_xy = (0.10, 0.50), (0.90, 0.50)
hidden_y = np.linspace(0.82, 0.18, len(w))
hidden_xy = [(0.48, float(y)) for y in hidden_y]
for xy in hidden_xy:
ax.add_patch(FancyArrowPatch(input_xy, xy, arrowstyle="->", mutation_scale=12, lw=1.2, color="0.68", zorder=0))
ax.add_patch(FancyArrowPatch(xy, output_xy, arrowstyle="->", mutation_scale=12, lw=1.2, color="0.68", zorder=0))
ax.add_patch(Circle(input_xy, 0.055, fc=TARGET_COLOR, ec="none"))
ax.text(*input_xy, "x", color="white", ha="center", va="center", fontsize=14, fontweight="bold")
for j, xy in enumerate(hidden_xy):
ax.add_patch(Circle(xy, 0.046, fc=HINGE_COLORS[j], ec="none"))
ax.text(*xy, f"h{j + 1}", color="white", ha="center", va="center", fontsize=10, fontweight="bold")
ax.text(0.35, xy[1], f"w={w[j]:+.2f}\nb={b[j]:+.2f}", ha="right", va="center", fontsize=8)
ax.text(0.64, xy[1], f"a={a[j]:+.2f}", ha="left", va="center", fontsize=8)
ax.add_patch(Circle(output_xy, 0.065, fc=MODEL_COLOR, ec="none"))
ax.text(*output_xy, r"$\hat f$", color="white", ha="center", va="center", fontsize=14, fontweight="bold")
ax.text(0.10, 0.38, "input", ha="center", color="0.35")
ax.text(0.48, 0.07, r"hidden layer: $h_j=\mathrm{ReLU}(w_jx+b_j)$", ha="center", color="0.30")
display_bias = 0.0 if abs(float(output_bias)) < 0.005 else float(output_bias)
ax.text(0.90, 0.38, f"weighted sum\noutput bias={display_bias:.2f}", ha="center", color="0.35")
ax.set_title("The same equation, drawn as a one-hidden-layer network", fontsize=14, pad=12)
finish_figure(fig, "03a-hidden-layer-network")# @title Auxiliary: Scene 2 network and component visualisations
def plot_scene2_network(w, b, a):
fig, axes = plt.subplots(1, 2, figsize=(12.0, 4.2))
for ax in axes:
ax.set(xlim=(0, 1), ylim=(0, 1)); ax.axis("off")
def arrow(ax, start, end):
ax.add_patch(FancyArrowPatch(start, end, arrowstyle="->", mutation_scale=13,
lw=1.6, color="0.55", zorder=0))
left = axes[0]
positions = [(0.10, 0.55), (0.39, 0.55), (0.68, 0.55), (0.92, 0.55)]
labels = ["x", f"multiply\nby w={w:g}", f"add\nb={b:g}", "z"]
colors = [TARGET_COLOR, HINGE_COLORS[0], HINGE_COLORS[2], "#555555"]
for start, end in zip(positions[:-1], positions[1:]):
arrow(left, start, end)
for xy, label, color in zip(positions, labels, colors):
left.add_patch(Circle(xy, 0.075, fc=color, ec="none"))
left.text(*xy, label, ha="center", va="center", color="white",
fontsize=9 if "\n" in label else 13, fontweight="bold")
left.text(0.53, 0.25, "nn.Linear(1, 1) computes z = wx + b", ha="center", fontsize=11)
left.text(0.53, 0.16, "It does not apply ReLU.", ha="center", color="0.35")
left.set_title("Inside nn.Linear", fontsize=13)
right = axes[1]
xs = [0.07, 0.29, 0.51, 0.72, 0.94]
for x0, x1 in zip(xs[:-1], xs[1:]):
arrow(right, (x0, 0.55), (x1, 0.55))
right.text(xs[0], 0.55, "x", ha="center", va="center", color="white",
fontsize=13, fontweight="bold", bbox={"boxstyle": "circle,pad=0.55", "fc": TARGET_COLOR, "ec": "none"})
blocks = [(xs[1], "Linear\nw, b", HINGE_COLORS[0]),
(xs[2], "ReLU", HINGE_COLORS[1]),
(xs[3], f"multiply\nby a={a:g}", HINGE_COLORS[2])]
for xpos, label, color in blocks:
right.text(xpos, 0.55, label, ha="center", va="center", color="white",
fontsize=9, fontweight="bold",
bbox={"boxstyle": "round,pad=0.55", "fc": color, "ec": "none"})
right.text(xs[4], 0.55, "a h(x)", ha="center", va="center", color="white",
fontsize=10, fontweight="bold",
bbox={"boxstyle": "circle,pad=0.45", "fc": MODEL_COLOR, "ec": "none"})
right.text(xs[1], 0.30, "z = wx + b", ha="center", color="0.35")
right.text(xs[2], 0.30, "h = ReLU(z)", ha="center", color="0.35")
right.text(xs[3], 0.30, "output weight", ha="center", color="0.35")
right.set_title("The complete Scene 2 network", fontsize=13)
fig.suptitle("One neuron is a sequence of simple operations", fontsize=15, fontweight="bold")
fig.tight_layout()
finish_figure(fig, "02a-linear-layer-and-scene2-network")
def plot_component_sum(x, target, components, total, output_bias, breakpoints, mse, prefix=""):
fig, (ax_parts, ax_sum) = plt.subplots(2, 1, figsize=(9.0, 6.8), sharex=True, gridspec_kw={"height_ratios": [1.0, 1.15]})
for j in range(components.shape[1]):
ax_parts.plot(x[:, 0], components[:, j], color=HINGE_COLORS[j], lw=1.5, alpha=0.72, label=f"neuron {j + 1}")
ax_parts.axhline(float(output_bias), color="0.50", ls=":", lw=1.5, label="output bias")
ax_parts.set(ylabel="contribution", title="The pieces: four weighted hinges")
ax_parts.legend(frameon=False, ncol=5, fontsize=9, loc="upper center")
ax_sum.plot(x[:, 0], total[:, 0], color=MODEL_COLOR, lw=3.2, label="sum of contributions")
ax_sum.plot(x[:, 0], target[:, 0], color=TARGET_COLOR, lw=3.2, ls="--", label="target f(x)")
for point in breakpoints:
if -1 <= point <= 1:
ax_parts.axvline(float(point), color="0.86", lw=1, zorder=0)
ax_sum.axvline(float(point), color="0.86", lw=1, zorder=0)
ax_sum.set(xlabel="x", ylabel="value", ylim=(-1.35, 1.45), title="The result: their sum changes slope at each hinge")
ax_sum.text(0.02, 0.08, f"MSE = {mse:.4f}", transform=ax_sum.transAxes, color="0.35")
ax_sum.legend(frameon=False, ncol=2, loc="upper center")
fig.suptitle(f"{prefix}Simple hinges add up to a new shape", fontsize=15, fontweight="bold", y=0.98)
fig.subplots_adjust(top=0.88, hspace=0.42)
finish_figure(fig, "03-several-bends-become-a-function")# @title Auxiliary: training-progress visualisation
def plot_training_story(x, target, snapshots, epochs, components, total, output_bias, breakpoints):
fig = plt.figure(figsize=(12.0, 8.0))
grid = fig.add_gridspec(2, 4, height_ratios=[1.0, 1.25], hspace=0.38, wspace=0.28)
for column, epoch in enumerate(epochs):
ax = fig.add_subplot(grid[0, column])
frame = snapshots[epoch]
ax.plot(x[:, 0], target[:, 0], color=TARGET_COLOR, lw=2.2, label="target")
ax.plot(x[:, 0], frame["prediction"][:, 0], color=MODEL_COLOR, lw=2.2, label="network")
ax.set_title(f"step {epoch}\nMSE {frame['loss']:.4f}", fontsize=10)
ax.set(xlim=(-1, 1), ylim=(-1.35, 1.35), xlabel="x")
if column == 0:
ax.set_ylabel("value"); ax.legend(frameon=False, fontsize=8)
ax = fig.add_subplot(grid[1, :])
for j in range(components.shape[1]):
ax.plot(x[:, 0], components[:, j], color=HINGE_COLORS[j], lw=1.5, alpha=0.72, label=f"learned neuron {j + 1}")
ax.axhline(float(output_bias), color="0.55", ls=":", lw=1.5, label="learned output bias")
ax.plot(x[:, 0], total[:, 0], color=MODEL_COLOR, lw=3.2, label="learned sum")
ax.plot(x[:, 0], target[:, 0], color=TARGET_COLOR, lw=3.0, ls="--", label="target f(x)")
for point in breakpoints:
if -1 <= point <= 1:
ax.axvline(float(point), color="0.82", ls="--", lw=1)
ax.set(xlabel="x", ylabel="value", title="After training: learned contributions and their sum")
ax.legend(frameon=False, ncol=4, fontsize=8.5, loc="upper center")
fig.suptitle("Gradient descent arranges the bends", fontsize=15, fontweight="bold")
finish_figure(fig, "04-gradient-descent-arranges-the-bends")
def plot_xor_wall(X, y, A1):
fig, axes = plt.subplots(1, 2, figsize=(11.0, 4.6))
class_colors = [TARGET_COLOR, MODEL_COLOR]
for label in (0, 1):
mask = y == label
axes[0].scatter(X[mask, 0], X[mask, 1], s=150, color=class_colors[label],
edgecolor="white", linewidth=1.5, label=f"class {label}", zorder=3)
axes[0].plot([0, 1], [0, 1], color=class_colors[0], ls="--", alpha=0.65)
axes[0].plot([0, 1], [1, 0], color=class_colors[1], ls="--", alpha=0.65)
axes[0].scatter([0.5], [0.5], marker="x", s=90, color="black", zorder=4)
axes[0].annotate("class convex hulls intersect", (0.5, 0.5), (0.10, 0.67),
arrowprops={"arrowstyle": "->"}, fontsize=9)
axes[0].set(title="Original input space: no affine separator", xlabel="$x_1$", ylabel="$x_2$")
axes[0].legend(frameon=False, loc="upper center", ncol=2)
grid = np.linspace(-0.12, 1.12, 200)
boundary = 0.5 - grid
axes[1].fill_between(grid, np.clip(boundary, -0.12, 1.12), 1.12,
color=MODEL_COLOR, alpha=0.08)
for label in (0, 1):
mask = y == label
axes[1].scatter(A1[mask, 0], A1[mask, 1], s=150, color=class_colors[label],
edgecolor="white", linewidth=1.5, label=f"class {label}", zorder=3)
axes[1].plot(grid, boundary, color="black", lw=2, label="$A_{1,1}+A_{1,2}=0.5$")
axes[1].annotate("two class-0 samples", (0, 0), (0.10, 0.18),
arrowprops={"arrowstyle": "->"}, fontsize=9)
axes[1].set(title="Hidden space: one line is enough", xlabel="$A_{1,1}$", ylabel="$A_{1,2}$")
axes[1].legend(frameon=False, loc="upper center", ncol=2, fontsize=9)
for ax in axes:
ax.set(xlim=(-0.12, 1.12), ylim=(-0.12, 1.12), aspect="equal", xticks=[0, 0.5, 1], yticks=[0, 0.5, 1])
fig.suptitle("XOR becomes simple after a change of representation", fontsize=15, fontweight="bold")
fig.tight_layout()
finish_figure(fig, "05-xor-hidden-representation")
def plot_forward_pass_diagram(x0, W1, b1, Z1, A1, W2, contributions, b2, y_hat, target):
fig, ax = plt.subplots(figsize=(12.0, 6.4))
ax.set(xlim=(0, 1), ylim=(0, 1)); ax.axis("off")
input_xy, output_xy = (0.08, 0.50), (0.91, 0.50)
hidden_y = np.linspace(0.82, 0.18, A1.shape[1])
hidden_xy = [(0.48, float(y_pos)) for y_pos in hidden_y]
active = A1[0] > 0
for j, xy in enumerate(hidden_xy):
path_color = HINGE_COLORS[j] if active[j] else "0.78"
path_style = "-" if active[j] else ":"
ax.add_patch(FancyArrowPatch(input_xy, xy, arrowstyle="->", mutation_scale=12,
lw=1.6, ls=path_style, color=path_color, zorder=0))
ax.add_patch(FancyArrowPatch(xy, output_xy, arrowstyle="->", mutation_scale=12,
lw=1.6, ls=path_style, color=path_color, zorder=0))
ax.add_patch(Circle(input_xy, 0.057, fc=TARGET_COLOR, ec="none"))
ax.text(*input_xy, f"$x_0$\n{x0:.2f}", color="white", ha="center", va="center",
fontsize=11, fontweight="bold")
for j, xy in enumerate(hidden_xy):
face = HINGE_COLORS[j] if active[j] else "#E4E4E4"
text_color = "white" if active[j] else "0.35"
ax.add_patch(Circle(xy, 0.058, fc=face, ec="none"))
ax.text(*xy, f"{j + 1}\n$z$={Z1[0,j]:+.2f}\n$A_1$={A1[0,j]:.2f}",
color=text_color, ha="center", va="center", fontsize=8, fontweight="bold")
ax.text(0.34, xy[1], f"$w$={W1[j,0]:+.2f}\n$b$={b1[j]:+.2f}",
ha="right", va="center", fontsize=8, color="0.30")
displayed_contribution = 0.0 if abs(contributions[0,j]) < 0.005 else contributions[0,j]
ax.text(0.65, xy[1], f"$a$={W2[0,j]:+.2f}\n$aA_1$={displayed_contribution:+.2f}",
ha="left", va="center", fontsize=8, color="0.30")
ax.add_patch(Circle(output_xy, 0.071, fc=MODEL_COLOR, ec="none"))
ax.text(*output_xy, f"$\\hat y$\n{y_hat:.3f}", color="white", ha="center", va="center",
fontsize=11, fontweight="bold")
ax.text(0.91, 0.35, f"$b_2$={b2:+.2f}\ntarget={target:.3f}", ha="center", color="0.30")
ax.text(0.48, 0.06, "solid = active path dotted gray = inactive ReLU path",
ha="center", color="0.35")
ax.set_title("One input moving through the trained 1 → 4 → 1 network", fontsize=15, pad=12)
finish_figure(fig, "06-forward-pass-values")
def plot_scalar_graph(values, view):
fig, ax = plt.subplots(figsize=(12.0, 4.8))
ax.set(xlim=(-0.06, 1.06), ylim=(-0.16, 1.12)); ax.axis("off")
positions = {
"x": (0.02, 0.82), "w": (0.02, 0.48), "b": (0.18, 0.18),
"z": (0.31, 0.62), "h": (0.48, 0.62), "a": (0.48, 0.18),
"yhat": (0.66, 0.62), "c": (0.66, 0.18), "y": (0.84, 0.18),
"L": (0.91, 0.62),
}
edges = [("x", "z", "× w"), ("w", "z", "× x"), ("b", "z", "+"),
("z", "h", "ReLU"), ("h", "yhat", "× a"),
("a", "yhat", "× h"), ("c", "yhat", "+"),
("yhat", "L", "ŷ − y"), ("y", "L", "−")]
if view == "forward":
edge_color, title = "0.55", "Forward view: values flow from parameters to one loss"
for source, target_name, label in edges:
ax.add_patch(FancyArrowPatch(positions[source], positions[target_name], arrowstyle="->",
mutation_scale=12, lw=1.4, color=edge_color, zorder=0))
midpoint = np.mean([positions[source], positions[target_name]], axis=0)
ax.text(midpoint[0], midpoint[1] + 0.035, label, fontsize=8, color="0.35", ha="center")
labels = {name: f"{name}\n{values[name]:+.3g}" for name in positions}
labels["yhat"] = f"ŷ\n{values['yhat']:+.3g}"
labels["L"] = f"L\n{values['L']:.4g}"
stem = "07-scalar-graph-forward"
else:
edge_color, title = "#7B2CBF", "Backward view: sensitivities flow from the loss"
for source, target_name, label in edges:
ax.add_patch(FancyArrowPatch(positions[target_name], positions[source], arrowstyle="->",
mutation_scale=12, lw=1.8, color=edge_color, zorder=0))
labels = {name: f"{name}\n∂L/∂{name}={values['grads'][name]:+.3g}" for name in positions}
labels["yhat"] = f"ŷ\n∂L/∂ŷ={values['grads']['yhat']:+.3g}"
labels["L"] = "L\nincoming = 1"
local = [(0.78, 0.88, "square: ŷ−y = −1.5"), (0.57, 0.88, "add: 1"),
(0.395, 0.88, "ReLU′(z): 1"), (0.15, 0.98, "z=wx+b: w, x, 1")]
for px, py, text_value in local:
ax.text(px, py, text_value, ha="center", fontsize=9, color=edge_color)
stem = "08-scalar-graph-backward"
for name, (px, py) in positions.items():
face = edge_color if view == "backward" else (MODEL_COLOR if name in {"z", "h", "yhat", "L"} else "#E8EEF5")
text_color = "white" if face in {"#7B2CBF", MODEL_COLOR} else "0.20"
ax.add_patch(Circle((px, py), 0.062, fc=face, ec="white", lw=1.5, zorder=2))
ax.text(px, py, labels[name], ha="center", va="center", color=text_color,
fontsize=8.3, fontweight="bold", zorder=3)
ax.set_title(title, fontsize=15, pad=10)
ax.text(0.5, -0.08, "Same computation graph; the arrows show which quantity is moving.",
ha="center", color="0.35")
finish_figure(fig, stem)
def plot_relu_cutoff(active_values, inactive_values):
fig, axes = plt.subplots(1, 2, figsize=(11.0, 3.8), sharey=True)
for ax, values, active in zip(axes, [active_values, inactive_values], [True, False]):
ax.set(xlim=(-0.05, 1.05), ylim=(-0.22, 0.30)); ax.axis("off")
points = [(0.10, 0), (0.37, 0), (0.64, 0), (0.91, 0)]
names = [f"z={values['z']:+.1f}", f"h={values['h']:.1f}",
f"ŷ={values['yhat']:+.1f}", f"L={values['L']:.3f}"]
path_color, style = (MODEL_COLOR, "-") if active else ("0.65", ":")
for left, right in zip(points[:-1], points[1:]):
ax.add_patch(FancyArrowPatch(left, right, arrowstyle="->", mutation_scale=13,
color=path_color, ls=style, lw=2))
for point, name in zip(points, names):
ax.add_patch(Circle(point, 0.073, fc=path_color, ec="white", lw=1.5))
ax.text(*point, name, ha="center", va="center", color="white",
fontsize=9, fontweight="bold")
derivative = 1 if active else 0
ax.text(0.37, 0.18, f"ReLU′(z) = {derivative}", ha="center",
color=path_color, fontweight="bold")
ax.text(0.37, -0.17, f"∂L/∂z = {values['grads']['z']:+.2f}", ha="center")
ax.set_title("z > 0: path open" if active else "z < 0: path cut off", fontsize=13)
fig.suptitle("The ReLU gate controls both the forward and backward path", fontsize=15, fontweight="bold")
fig.tight_layout()
finish_figure(fig, "09-relu-zero-gradient-path")
def plot_batch_relu_mask(X, Z1, mask):
fig, axes = plt.subplots(1, 2, figsize=(10.5, 3.8), constrained_layout=True)
z_max = float(np.abs(Z1).max())
images = [
axes[0].imshow(Z1, cmap="RdBu_r", vmin=-z_max, vmax=z_max, aspect="auto"),
axes[1].imshow(mask, cmap=plt.matplotlib.colors.ListedColormap(["#D9D9D9", MODEL_COLOR]),
vmin=0, vmax=1, aspect="auto"),
]
for row in range(Z1.shape[0]):
for col in range(Z1.shape[1]):
axes[0].text(col, row, f"{Z1[row, col]:+.2f}", ha="center", va="center",
color="white" if abs(Z1[row, col]) > 0.55 * z_max else "0.15", fontweight="bold")
axes[1].text(col, row, "pass" if mask[row, col] else "block", ha="center", va="center",
color="white" if mask[row, col] else "0.30", fontweight="bold")
labels = [f"x={value:+.2f}" for value in X[:, 0]]
for ax, title in zip(axes, ["Preactivations $Z_1$", "ReLU mask $\mathbf{1}[Z_1>0]$"]):
ax.set(title=title, xlabel="hidden neuron $j$", xticks=range(Z1.shape[1]),
xticklabels=range(1, Z1.shape[1] + 1), yticks=range(Z1.shape[0]), yticklabels=labels)
for spine in ax.spines.values(): spine.set_visible(False)
fig.colorbar(images[0], ax=axes[0], shrink=0.82, label="preactivation")
fig.suptitle("Three samples open different paths through the same hidden layer", fontsize=15, fontweight="bold")
finish_figure(fig, "10-mini-batch-relu-mask")
def plot_gradient_comparison(manual, automatic, errors):
names = ["W1", "b1", "W2", "b2"]
fig, axes = plt.subplots(2, 2, figsize=(10.0, 6.2), constrained_layout=True)
for ax, name in zip(axes.flat, names):
m = manual[name].reshape(-1).cpu().numpy()
a = automatic[name].reshape(-1).cpu().numpy()
indices = np.arange(len(m))
ax.plot(indices, m, "o", ms=8, color=TARGET_COLOR, label="manual")
ax.plot(indices, a, "x", ms=9, mew=2, color=MODEL_COLOR, label="autograd")
ax.axhline(0, color="0.75", lw=1)
ax.set(title=f"{name} max error={errors[name]:.1e}", xlabel="flattened entry", ylabel="gradient")
ax.set_xticks(indices)
axes[0, 0].legend(frameon=False)
fig.suptitle("Manual backpropagation and autograd land on the same gradients", fontsize=15, fontweight="bold")
finish_figure(fig, "11-manual-versus-autograd")An affine model has two adjustable numbers:
\[\hat f(x)=mx+c. \]
Changing \(c\) moves the line vertically. Changing \(m\) rotates it. But its derivative is constant:
\[\frac{d\hat f}{dx}=m. \]
Checkpoint. Can an affine function reproduce our curve? What kind of error must remain regardless of training time?
# Fit the best affine model on the displayed samples
affine_design = torch.cat([x_train, torch.ones_like(x_train)], dim=1)
affine_solution = torch.linalg.lstsq(affine_design, y_train).solution
affine_slope, affine_intercept = affine_solution[:, 0]
affine_prediction = affine_design @ affine_solution
affine_mse = torch.mean((affine_prediction - y_train) ** 2).item()
print(f"best line: {affine_slope:.3f}x {affine_intercept:+.3f} | MSE={affine_mse:.4f}")best line: 1.193x -0.000 | MSE=0.1990plot_affine_attempt(x_plot, y_target, affine_prediction, affine_mse)
The shaded residual has structure: the line misses one way, then the other, then switches again. This is not an optimization failure. It is a representation failure. The family of affine functions contains no curve with a changing slope.
We therefore need a component that can change slope somewhere.
A ReLU neuron performs two operations. First it creates a preactivation \(z(x)=wx+b\). Then it clips the negative side:
\[ h(x)=\operatorname{ReLU}(wx+b)=\begin{cases} 0, & wx+b\le 0,\\ wx+b, & wx+b>0. \end{cases} \]
When \(w \neq 0\), the two pieces meet at
\[x_{\text{break}}=-\frac{b}{w}. \]
Here the three parameters have different jobs:
In PyTorch, nn.Linear(1, 1) computes only \(z=wx+b\). Despite its name, it implements an affine map because it includes \(b\). ReLU is a separate operation, and multiplication by \(a\) belongs to the following output layer.
Play. Change only \(w\), \(b\), and \(a\). Can you move the breakpoint to \(x=0.5\)? Can you make the contribution bend downward?
w = 1.4
b = 0.35
a = -0.8# Follow one input through the neuron's three visible quantities
play_z = w * x_plot[:, 0] + b
play_h = torch.relu(play_z)
play_contribution = a * play_h
play_breakpoint = -b / w if w != 0 else None
print(f"breakpoint: {play_breakpoint:.3f}" if play_breakpoint is not None else "no finite breakpoint when w=0")breakpoint: -0.250plot_one_hinge(x_plot, y_target, play_z, play_h, play_contribution, w, b, a, play_breakpoint)
plot_scene2_network(w, b, a)
The neuron is not a tiny copy of the target. It offers one simple geometric piece: zero on one side and a line on the other. The output weight can stretch or flip that piece, but one hinge is still only one hinge.
What changes when several neurons work in parallel?
A hidden layer evaluates several hinges at the same input. The output layer forms their weighted sum:
\[ h_j(x)=\operatorname{ReLU}(w_jx+b_j),\qquad \hat f(x)=b_2+\sum_{j=1}^{H}a_jh_j(x). \]
Each symbol has one geometric job:
| Quantity | Geometric role |
|---|---|
| \(w_j,b_j\) | orient and place hinge \(j\) |
| \(a_j\) | scale or reverse its contribution |
| \(b_2\) | move the complete sum vertically |
Predict before computing. If each neuron changes slope at one location, where can the sum change slope?
# Place four hinges; solve only for their output weights and output bias
HIDDEN_WIDTH = 4
manual_breakpoints = torch.tensor([-0.75, -0.25, 0.25, 0.75], device=DEVICE)
manual_w = torch.tensor([1.0, -1.0, 1.0, -1.0], device=DEVICE)
manual_b = -manual_w * manual_breakpoints
manual_activations = torch.relu(x_plot * manual_w + manual_b)
manual_design = torch.cat([manual_activations, torch.ones_like(x_plot)], dim=1)
manual_solution = torch.linalg.lstsq(manual_design, y_target).solution
manual_a = manual_solution[:HIDDEN_WIDTH, 0]
manual_output_bias = manual_solution[HIDDEN_WIDTH, 0]
manual_component_values = manual_activations * manual_a
manual_sum = manual_output_bias + manual_component_values.sum(dim=1, keepdim=True)
manual_mse = torch.mean((manual_sum - y_target) ** 2).item()
print(f"four fixed hinges | MSE={manual_mse:.4f}")four fixed hinges | MSE=0.0122Yes—a hidden-layer neural network can be drawn directly in Colab. The diagram below uses Matplotlib and the actual parameters above, not a separate diagramming library. Every path from \(x\) to \(\hat f\) corresponds to one term \(a_j\operatorname{ReLU}(w_jx+b_j)\) in the sum.
plot_hidden_layer_network(manual_w, manual_b, manual_a, manual_output_bias)
The diagram shows connections; the next figure shows the functions carried by those connections. Notice that individual contributions can be larger than the final curve. Positive and negative pieces may cancel. That cancellation is part of the computation, not a plotting artifact.
plot_component_sum(x_plot, y_target, manual_component_values, manual_sum,
manual_output_bias, manual_breakpoints, manual_mse)
A-ha. A hidden layer does not approximate the function mysteriously. Each neuron contributes a simple bend; the output layer combines them.
Between neighboring breakpoints, every active contribution is affine. A sum of affine pieces is still affine there. At a breakpoint, one neuron’s status changes and the total slope can change. Therefore this network produces a continuous piecewise-linear function.
The target is smooth while the approximation has corners. More well-placed hinges can make those corners less visually obvious, but the network output remains piecewise linear. We do not need a universal-approximation theorem to understand this example: the figure exposes the mechanism directly.
So far, we chose the hinge geometry and solved only the final linear combination. Now every \(w_j\), \(b_j\), \(a_j\), and \(b_2\) may change. Their shared score is mean squared error:
\[ \mathcal L(\theta)=\frac{1}{N}\sum_{i=1}^{N}\left(\hat f(x_i;\theta)-f(x_i)\right)^2. \]
Gradient descent repeatedly applies
\[\theta\leftarrow\theta-\eta\nabla_\theta\mathcal L. \]
We will use that update now, but deliberately postpone the question of how PyTorch obtains \(\nabla_\theta\mathcal L\).
Before training. If the loss falls, what visible changes should happen to the breakpoints, the individual contributions, and their sum?
The PyTorch module mirrors the equation exactly: one affine transformation creates four preactivations, ReLU clips them, and a second affine transformation combines the four activations. Tensor shapes are shown in comments.
class BendNet(nn.Module):
def __init__(self, hidden_width):
super().__init__()
self.hidden = nn.Linear(1, hidden_width)
self.activation = nn.ReLU()
self.output = nn.Linear(hidden_width, 1)
def forward(self, x): # x: (N, 1)
h = self.activation(self.hidden(x)) # h: (N, H)
return self.output(h) # output: (N, 1)A randomly initialized ReLU can be inactive over the entire interval and receive no useful movement in this tiny demonstration. To keep all four hinges visible from the first frame, we start from the same spread-out breakpoints and alternating orientations used above. The output weights begin small and random.
This is a teaching choice, not a claim that initialization is unimportant.
torch.manual_seed(SEED)
model = BendNet(HIDDEN_WIDTH).to(DEVICE)
with torch.no_grad():
model.hidden.weight[:, 0].copy_(manual_w)
model.hidden.bias.copy_(manual_b)
model.output.weight.normal_(mean=0.0, std=0.05)
model.output.bias.zero_()
assert model(x_train).shape == y_train.shape
print(model)BendNet(
(hidden): Linear(in_features=1, out_features=4, bias=True)
(activation): ReLU()
(output): Linear(in_features=4, out_features=1, bias=True)
)Nothing is hidden in a trainer. Each iteration predicts, measures one scalar loss, clears old gradients, asks for new gradients, and updates the parameters. Four snapshots preserve the changing curve and the parameters that produced it.
loss_fn = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.10)
snapshot_epochs = (0, 50, 500, 4000)
training_snapshots = {}
for epoch in range(snapshot_epochs[-1] + 1):
prediction = model(x_train)
loss = loss_fn(prediction, y_train)
if epoch in snapshot_epochs:
training_snapshots[epoch] = {
"prediction": prediction.detach().clone(),
"loss": float(loss.detach()),
"parameters": {name: p.detach().clone() for name, p in model.named_parameters()},
}
if epoch == snapshot_epochs[-1]:
break
optimizer.zero_grad()
loss.backward()
optimizer.step()
print({epoch: round(frame["loss"], 6) for epoch, frame in training_snapshots.items()}){0: 0.638471, 50: 0.134761, 500: 0.012615, 4000: 0.001953}The trained model is still the same sum of hinges. We now extract those pieces explicitly and verify that their sum equals the module’s forward pass. These named values will later support a circles-and-arrows forward-pass diagram.
trained_parameters = {name: p.detach().clone() for name, p in model.named_parameters()}
with torch.no_grad():
learned_preactivations = model.hidden(x_plot)
learned_activations = model.activation(learned_preactivations)
learned_output_weights = model.output.weight[0].detach().clone()
learned_output_bias = model.output.bias[0].detach().clone()
learned_component_values = learned_activations * learned_output_weights
learned_sum = learned_output_bias + learned_component_values.sum(dim=1, keepdim=True)
learned_hidden_weights = model.hidden.weight[:, 0].detach().clone()
learned_hidden_biases = model.hidden.bias.detach().clone()
learned_breakpoints = -learned_hidden_biases / learned_hidden_weights
final_mse = loss_fn(learned_sum, y_target).item()
assert learned_component_values.shape == (241, HIDDEN_WIDTH)
assert torch.allclose(learned_sum, model(x_plot), atol=1e-6)
assert torch.isfinite(learned_sum).all() and math.isfinite(final_mse)
print("learned breakpoints:", [round(v, 3) for v in learned_breakpoints.tolist()])
print(f"final training MSE: {final_mse:.6f}")learned breakpoints: [-0.68, -0.399, 0.4, 0.68]
final training MSE: 0.001953plot_training_story(x_plot, y_target, training_snapshots, snapshot_epochs,
learned_component_values, learned_sum, learned_output_bias,
learned_breakpoints)
At step 0 the small output weights nearly cancel everything. By step 50 the network has found the global direction of the curve. By step 500 the principal turns are visible. The final network places pairs of breakpoints around the valley and peak:
\[x_{\text{break}}\approx -0.68,\;-0.40,\;0.40,\;0.68. \]
The fit is close, but not mathematically smooth: zooming in would reveal straight segments meeting at those locations. The MSE is evidence about these sampled points, not a generalization result; this scene uses no validation or test set.
The loss is one scalar. How did every weight and bias learn its own direction of change?
To approach that question, we first take a short detour through XOR: a tiny problem that makes the purpose of a hidden representation impossible to miss.
XOR assigns opposite labels to neighboring corners of a square:
\[ (0,0)\mapsto0,\qquad (0,1)\mapsto1,\qquad (1,0)\mapsto1,\qquad (1,1)\mapsto0. \]
An affine classifier would choose one side of a line. But the convex hull of class 0 is the diagonal from \((0,0)\) to \((1,1)\), while the convex hull of class 1 is the other diagonal. The two segments intersect at the centre. No line can put the two intersecting sets on opposite sides.
Checkpoint. Can one affine boundary separate the four points? Would more training time change the geometry?
xor_X = torch.tensor([[0.0, 0.0],
[0.0, 1.0],
[1.0, 0.0],
[1.0, 1.0]], device=DEVICE)
xor_y = torch.tensor([0, 1, 1, 0], device=DEVICE)
assert xor_X.shape == (4, 2) and xor_y.shape == (4,)
print("XOR inputs:", tuple(xor_X.shape))XOR inputs: (4, 2)We reuse the trained \(1\rightarrow4\rightarrow1\) function approximator from Part I—no new model and no retraining. With PyTorch’s stored weight orientation,
\[ Z_1=XW_1^\top+b_1,\qquad A_1=\operatorname{ReLU}(Z_1),\qquad \hat Y=A_1W_2^\top+b_2, \]
where \(W_1:[H,1]\), \(W_2:[1,H]\), and \(H=4\) is the hidden width. For one neuron, \(z_j=w_jx_0+b_j\), \(A_{1,j}=\operatorname{ReLU}(z_j)\), and its output contribution is \(a_jA_{1,j}\).
Choose \(x_0=0.20\). It lies away from every learned breakpoint and activates some—but not all—hidden neurons.
Checkpoint. Which hidden neurons will be active at this input? Which paths will contribute exactly zero?
# Reassemble the trained forward pass using PyTorch's stored parameter shapes
H = HIDDEN_WIDTH
x0 = torch.tensor([[0.20]], device=DEVICE) # X: [1, 1]
W1 = model.hidden.weight.detach().clone() # [H, 1]
b1 = model.hidden.bias.detach().clone() # [H]
W2 = model.output.weight.detach().clone() # [1, H]
b2 = model.output.bias.detach().clone() # [1]
Z1 = x0 @ W1.T + b1 # [1, H]
A1 = torch.relu(Z1) # [1, H]
forward_contributions = A1 * W2 # [1, H]
y_hat_explicit = forward_contributions.sum(dim=1, keepdim=True) + b2
y_hat_model = model(x0)
y0 = target_function(x0)
forward_pass_equality_error = (y_hat_explicit - y_hat_model).abs().max().item()
active_mask = A1[0] > 0
assert 0 < int(active_mask.sum()) < H
assert torch.allclose(y_hat_explicit, y_hat_model, atol=1e-7)
print("active neurons:", [j + 1 for j in range(H) if active_mask[j]])
print(f"explicit/model equality error: {forward_pass_equality_error:.3e}")active neurons: [1, 4]
explicit/model equality error: 0.000e+00The table records every scalar on the four hidden paths. Here \(a_j\) is entry \(j\) of the stored output matrix \(W_2\).
from IPython.display import Markdown, display
table_lines = [
"| neuron $j$ | $w_j$ | $b_j$ | $z_j$ | $A_{1,j}$ | $a_j$ | $a_jA_{1,j}$ | path |",
"|---:|---:|---:|---:|---:|---:|---:|:---|",
]
for j in range(H):
state = "active" if active_mask[j] else "inactive"
displayed_contribution = 0.0 if abs(forward_contributions[0,j].item()) < 0.0005 else forward_contributions[0,j].item()
table_lines.append(
f"| {j+1} | {W1[j,0]:+.3f} | {b1[j]:+.3f} | {Z1[0,j]:+.3f} | "
f"{A1[0,j]:.3f} | {W2[0,j]:+.3f} | {displayed_contribution:+.3f} | {state} |"
)
summary = (
f"\n\n$x_0={x0.item():.2f}$, $b_2={b2.item():+.3f}$, "
f"$\\hat y={y_hat_explicit.item():.4f}$, and $f(x_0)={y0.item():.4f}$."
)
display(Markdown("\n".join(table_lines) + summary))| neuron \(j\) | \(w_j\) | \(b_j\) | \(z_j\) | \(A_{1,j}\) | \(a_j\) | \(a_jA_{1,j}\) | path |
|---|---|---|---|---|---|---|---|
| 1 | +1.365 | +0.928 | +1.201 | 1.201 | +1.088 | +1.307 | active |
| 2 | -1.951 | -0.779 | -1.169 | 0.000 | +1.832 | +0.000 | inactive |
| 3 | +1.952 | -0.780 | -0.390 | 0.000 | -1.835 | +0.000 | inactive |
| 4 | -1.361 | +0.926 | +0.653 | 0.653 | -1.080 | -0.706 | active |
\(x_0=0.20\), \(b_2=-0.010\), \(\hat y=0.5910\), and \(f(x_0)=0.6378\).
The same numbers now return to the circles-and-arrows view. Colored nodes and solid arrows are active; gray nodes and dotted arrows are paths cut off by ReLU. The contribution beside each output edge is exactly \(a_jA_{1,j}\).
plot_forward_pass_diagram(
x0.item(), W1.numpy(), b1.numpy(), Z1.numpy(), A1.numpy(), W2.numpy(),
forward_contributions.numpy(), b2.item(), y_hat_explicit.item(), y0.item()
)
\[ X:[N,1] \longrightarrow Z_1:[N,H] \longrightarrow A_1:[N,H] \longrightarrow \hat Y:[N,1]. \]
These are exactly the shapes produced by nn.Linear(1, H), nn.ReLU(), and nn.Linear(H, 1). PyTorch stores each linear layer’s output weights by rows, which is why the forward equations contain \(W_1^\top\) and \(W_2^\top\).
assert x0.shape == (1, 1)
assert W1.shape == (H, 1) and b1.shape == (H,)
assert Z1.shape == A1.shape == (1, H)
assert W2.shape == (1, H) and b2.shape == (1,)
assert y_hat_explicit.shape == y_hat_model.shape == (1, 1)
forward_pass_values = {
"x0": x0.detach().clone(), "W1": W1, "b1": b1, "Z1": Z1,
"A1": A1, "W2": W2, "b2": b2,
"contributions": forward_contributions,
"prediction": y_hat_explicit.detach().clone(), "target": y0.detach().clone(),
}
print("X → Z1 → A1 → Y_hat:", tuple(x0.shape), "→", tuple(Z1.shape),
"→", tuple(A1.shape), "→", tuple(y_hat_explicit.shape))X → Z1 → A1 → Y_hat: (1, 1) → (1, 4) → (1, 4) → (1, 1)Takeaway. Forward propagation is evaluation: the network applies a known sequence of operations to known values.
The table, diagram, matrix equations, and model(x0) all describe the same calculation. Forward propagation explains the prediction. It does not yet explain how responsibility for the error is distributed backward.
Before returning to matrices, isolate one scalar path. This makes the mechanism of credit assignment visible without the bookkeeping of a full layer.
A network architecture groups units into layers. The earlier \(1 \to 4 \to 1\) drawing answered: which units are connected?
A computation graph expands one calculation into elementary operations. It answers: which values produced this particular loss? For one scalar path we will use
\[ z=wx+b,\qquad h=\operatorname{ReLU}(z),\qquad \hat y=ah+c,\qquad L=\frac12(\hat y-y)^2. \]
This is deliberately smaller than the trained curve model. The mechanism we expose here will later apply to every path in that model.
# Small scalar example: every forward value is easy to inspect
scalar_values = {"x": 2.0, "w": 1.5, "b": -1.0, "a": -0.5, "c": 1.0, "y": 1.5}
scalar_values["z"] = scalar_values["w"] * scalar_values["x"] + scalar_values["b"]
scalar_values["h"] = max(scalar_values["z"], 0.0)
scalar_values["yhat"] = scalar_values["a"] * scalar_values["h"] + scalar_values["c"]
scalar_values["L"] = 0.5 * (scalar_values["yhat"] - scalar_values["y"]) ** 2
forward_rows = [
("affine", "z = wx + b", scalar_values["z"]),
("ReLU", "h = max(z, 0)", scalar_values["h"]),
("output", "ŷ = ah + c", scalar_values["yhat"]),
("loss", "L = ½(ŷ − y)²", scalar_values["L"]),
]
print(f"{'operation':<10} {'calculation':<20} {'value':>7}")
print("-" * 39)
for operation, calculation, value in forward_rows:
print(f"{operation:<10} {calculation:<20} {value:>7.3f}")operation calculation value
---------------------------------------
affine z = wx + b 2.000
ReLU h = max(z, 0) 2.000
output ŷ = ah + c 0.000
loss L = ½(ŷ − y)² 1.125The forward pass is a sequence of substitutions:
\[z=(1.5)(2)-1=2,\quad h=2,\quad \hat y=(-0.5)(2)+1=0,\quad L=\tfrac12(0-1.5)^2=1.125.\]
plot_scalar_graph(scalar_values, view="forward")
Here \(\hat y=0\) is below \(y=1.5\). Before calculating, predict the signs of \(\frac{\partial L}{\partial w}\), \(\frac{\partial L}{\partial b}\), \(\frac{\partial L}{\partial a}\), and \(\frac{\partial L}{\partial c}\).
A positive gradient means increasing that parameter increases the loss locally; gradient descent therefore moves the parameter in the opposite direction.
Each operation knows only its own local derivative:
\[ \frac{\partial L}{\partial \hat y}=\hat y-y,\qquad \frac{\partial L}{\partial y}=-(\hat y-y),\qquad \frac{\partial \hat y}{\partial h}=a,\quad \frac{\partial \hat y}{\partial a}=h,\quad \frac{\partial \hat y}{\partial c}=1, \]
\[ \frac{\partial h}{\partial z}=1\;\text{ because }z>0,\qquad \frac{\partial z}{\partial w}=x,\quad \frac{\partial z}{\partial b}=1. \]
Backpropagation multiplies an incoming sensitivity by a local derivative and accumulates contributions whenever paths meet.
For example, the sensitivity arriving at \(z\) is
\[ \frac{\partial L}{\partial z}= \underbrace{(\hat y-y)}_{-1.5} \underbrace{a}_{-0.5} \underbrace{\operatorname{ReLU}'(z)}_{1}=0.75. \]
# Manual reverse pass: incoming sensitivity × local derivative
g = {}
g["L"] = 1.0
g["yhat"] = scalar_values["yhat"] - scalar_values["y"]
g["y"] = -g["yhat"]
g["c"] = g["yhat"]
g["a"] = g["yhat"] * scalar_values["h"]
g["h"] = g["yhat"] * scalar_values["a"]
g["z"] = g["h"] * 1.0 # ReLU'(z)=1 because z>0
g["b"] = g["z"]
g["w"] = g["z"] * scalar_values["x"]
g["x"] = g["z"] * scalar_values["w"]
scalar_values["grads"] = g
gradient_rows = [(name, g[name]) for name in ("w", "b", "a", "c")]
print(f"{'parameter':<10} {'∂L / ∂parameter':>18}")
print("-" * 29)
for name, value in gradient_rows:
print(f"{name:<10} {value:>18.3f}")parameter ∂L / ∂parameter
-----------------------------
w 1.500
b 0.750
a -3.000
c -1.500The prediction was too low, but the parameters do not all receive the same instruction:
\[ \frac{\partial L}{\partial a}=(-1.5)(2)=-3,\qquad \frac{\partial L}{\partial c}=-1.5, \]
\[ \frac{\partial L}{\partial w}=(0.75)(2)=1.5,\qquad \frac{\partial L}{\partial b}=0.75. \]
The chain rule assigns each parameter a responsibility determined by its position in the graph.
plot_scalar_graph(scalar_values, view="backward")
Keep \(x,a,c,y\) fixed, but choose \(w=-1\) and \(b=0.5\). Now \(z=-1.5\) and \(h=0\). The output side still receives a loss signal, but
\[\operatorname{ReLU}'(-1.5)=0\quad\Longrightarrow\quad \frac{\partial L}{\partial z}=\frac{\partial L}{\partial w}=\frac{\partial L}{\partial b}=0.\]
inactive_values = {"x": 2.0, "w": -1.0, "b": 0.5, "a": -0.5, "c": 1.0, "y": 1.5}
inactive_values["z"] = inactive_values["w"] * inactive_values["x"] + inactive_values["b"]
inactive_values["h"] = max(inactive_values["z"], 0.0)
inactive_values["yhat"] = inactive_values["a"] * inactive_values["h"] + inactive_values["c"]
inactive_values["L"] = 0.5 * (inactive_values["yhat"] - inactive_values["y"]) ** 2
inactive_values["grads"] = {"z": 0.0, "w": 0.0, "b": 0.0}
plot_relu_cutoff(scalar_values, inactive_values)
Takeaway. The forward graph records what happened. The backward graph records how sensitive the loss is to every intermediate value. ReLU acts as a gate in both directions: an active path passes values and sensitivities; an inactive path blocks both.
We have calculated these derivatives by hand for one scalar path. Next, we scale the same bookkeeping to tensors and compare it with PyTorch autograd.
# Numerical check of the hand-derived parameter gradients (finite differences, not autograd)
def scalar_loss(w, b, a, c):
z = w * scalar_values["x"] + b
h = max(z, 0.0)
y_hat = a * h + c
return 0.5 * (y_hat - scalar_values["y"]) ** 2
epsilon = 1e-6
manual_parameter_gradients = {name: g[name] for name in ("w", "b", "a", "c")}
finite_difference_gradients = {}
base_parameters = {name: scalar_values[name] for name in ("w", "b", "a", "c")}
for name in base_parameters:
plus, minus = base_parameters.copy(), base_parameters.copy()
plus[name] += epsilon
minus[name] -= epsilon
finite_difference_gradients[name] = (scalar_loss(**plus) - scalar_loss(**minus)) / (2 * epsilon)
derivative_verification_error = max(
abs(manual_parameter_gradients[name] - finite_difference_gradients[name])
for name in manual_parameter_gradients
)
assert derivative_verification_error < 1e-8
print(f"Largest manual/finite-difference discrepancy: {derivative_verification_error:.2e}")Largest manual/finite-difference discrepancy: 1.23e-10Return to the trained \(1\rightarrow4\rightarrow1\) curve model. Instead of one input, place three inputs in the rows of \(X\):
\[X=\begin{bmatrix}-0.75\\0\\0.75\end{bmatrix},\qquad B=3,\qquad H=4.\]
PyTorch stores each output unit’s weights in a row. Therefore the forward equations are
\[ Z_1=XW_1^\top+b_1,\qquad A_1=\operatorname{ReLU}(Z_1),\qquad \hat Y=A_1W_2^\top+b_2. \]
The biases are broadcast across the \(B\) sample rows. Here is the complete shape ledger.
B, H = 3, HIDDEN_WIDTH
X_batch = torch.tensor([[-0.75], [0.0], [0.75]], dtype=torch.float64, device=DEVICE)
Y_batch = target_function(X_batch)
W1_batch = trained_parameters["hidden.weight"].to(torch.float64) # [H, 1]
b1_batch = trained_parameters["hidden.bias"].to(torch.float64) # [H]
W2_batch = trained_parameters["output.weight"].to(torch.float64) # [1, H]
b2_batch = trained_parameters["output.bias"].to(torch.float64) # [1]
shape_rows = [
("$X$", "batch inputs", X_batch.shape),
("$W_1$", "hidden weights", W1_batch.shape),
("$b_1$", "hidden biases", b1_batch.shape),
("$Z_1$", "hidden preactivations", (B, H)),
("$A_1$", "hidden activations", (B, H)),
("$W_2$", "output weights", W2_batch.shape),
("$b_2$", "output bias", b2_batch.shape),
("$\hat Y$", "predictions", (B, 1)),
]
shape_table = ["| tensor | meaning | shape |", "|:--|:--|:--|"]
shape_table += [f"| {name} | {meaning} | `{list(shape)}` |" for name, meaning, shape in shape_rows]
display(Markdown("\n".join(shape_table)))| tensor | meaning | shape |
|---|---|---|
| \(X\) | batch inputs | [3, 1] |
| \(W_1\) | hidden weights | [4, 1] |
| \(b_1\) | hidden biases | [4] |
| \(Z_1\) | hidden preactivations | [3, 4] |
| \(A_1\) | hidden activations | [3, 4] |
| \(W_2\) | output weights | [1, 4] |
| \(b_2\) | output bias | [1] |
| \(\hat Y\) | predictions | [3, 1] |
# Manual forward pass: identical stored values, now in float64
Z1_batch = X_batch @ W1_batch.T + b1_batch
relu_mask = Z1_batch > 0
A1_batch = torch.relu(Z1_batch)
Y_hat_batch = A1_batch @ W2_batch.T + b2_batch
batch_loss = torch.mean((Y_hat_batch - Y_batch) ** 2)
assert Z1_batch.shape == A1_batch.shape == (B, H)
assert Y_hat_batch.shape == Y_batch.shape == (B, 1)
active_sets = [[j + 1 for j in range(H) if relu_mask[i, j]] for i in range(B)]
print("active neurons by sample:", active_sets)
print(f"three-sample MSE: {batch_loss.item():.8f}")active neurons by sample: [[2, 4], [1, 4], [1, 3]]
three-sample MSE: 0.00081532plot_batch_relu_mask(X_batch.cpu().numpy(), Z1_batch.cpu().numpy(), relu_mask.cpu().numpy())
The three rows open different routes: \(\{2,4\}\), \(\{1,4\}\), and \(\{1,3\}\). A zero in the mask means that entry passes neither an activation forward nor a sensitivity backward.
For mean squared error,
\[L=\frac1B\sum_i(\hat y_i-y_i)^2,\]
so the sensitivity arriving at the predictions is
\[G_{\hat Y}=\frac{2}{B}(\hat Y-Y).\]
The remaining reverse pass repeats the scalar rule—incoming sensitivity times local derivative—for whole arrays:
\[ G_{W_2}=G_{\hat Y}^{\top}A_1,\qquad g_{b_2}=\sum_iG_{\hat Y,i},\qquad G_{A_1}=G_{\hat Y}W_2, \]
\[ G_{Z_1}=G_{A_1}\odot\mathbf 1[Z_1>0],\qquad G_{W_1}=G_{Z_1}^{\top}X,\qquad g_{b_1}=\sum_iG_{Z_1,i}. \]
The transpose aligns sample contributions before multiplication. Matrix multiplication and the bias sums then accumulate responsibility across all \(B\) samples. We never need to materialize a full Jacobian: reverse-mode differentiation propagates vector–Jacobian products.
# Manual matrix backpropagation
G_Y_hat = (2.0 / B) * (Y_hat_batch - Y_batch) # [B, 1]
G_W2 = G_Y_hat.T @ A1_batch # [1, H]
g_b2 = G_Y_hat.sum(dim=0) # [1]
G_A1 = G_Y_hat @ W2_batch # [B, H]
G_Z1 = G_A1 * relu_mask.to(torch.float64) # [B, H]
G_W1 = G_Z1.T @ X_batch # [H, 1]
g_b1 = G_Z1.sum(dim=0) # [H]
manual_matrix_gradients = {"W1": G_W1, "b1": g_b1, "W2": G_W2, "b2": g_b2}
assert G_Y_hat.shape == (B, 1) and G_A1.shape == G_Z1.shape == (B, H)
assert G_W1.shape == W1_batch.shape and g_b1.shape == b1_batch.shape
assert G_W2.shape == W2_batch.shape and g_b2.shape == b2_batch.shape
for name, gradient in manual_matrix_gradients.items():
print(f"{name}:\n{gradient}\n")W1:
tensor([[-0.0192],
[-0.0317],
[ 0.0324],
[ 0.0187]], dtype=torch.float64)
b1:
tensor([-0.0256, 0.0423, 0.0432, -0.0249], dtype=torch.float64)
W2:
tensor([[-0.0460, 0.0158, -0.0161, 0.0449]], dtype=torch.float64)
b2:
tensor([-0.0005], dtype=torch.float64)
Notice what the formulas hide compactly. Each row of \(G_{Z_1}\) is filtered by that sample’s ReLU mask. Then \(G_{Z_1}^{\top}X\) adds the surviving contributions from the three samples into one gradient for each stored hidden weight. The bias gradient uses an explicit row sum because the same bias was reused for every sample.
The comparison is only meaningful if both routes use the same values, dtype, data, loss reduction, and ReLU convention. We therefore make leaf copies of the trained parameters in float64, run a fresh forward pass, and clear gradients before calling backward().
# Leaf parameters: detached copies start a fresh autograd graph
W1_auto = W1_batch.clone().detach().requires_grad_(True)
b1_auto = b1_batch.clone().detach().requires_grad_(True)
W2_auto = W2_batch.clone().detach().requires_grad_(True)
b2_auto = b2_batch.clone().detach().requires_grad_(True)
auto_parameters = [W1_auto, b1_auto, W2_auto, b2_auto]
for parameter in auto_parameters:
parameter.grad = None # clear old gradients
Z1_auto = X_batch @ W1_auto.T + b1_auto
A1_auto = torch.relu(Z1_auto)
Y_hat_auto = A1_auto @ W2_auto.T + b2_auto
loss_auto = torch.mean((Y_hat_auto - Y_batch) ** 2)
Z1_auto.retain_grad() # inspect one non-leaf
loss_auto.backward()
autograd_gradients = {
"W1": W1_auto.grad.detach().clone(), "b1": b1_auto.grad.detach().clone(),
"W2": W2_auto.grad.detach().clone(), "b2": b2_auto.grad.detach().clone(),
}
gradient_max_errors = {
name: float((manual_matrix_gradients[name] - autograd_gradients[name]).abs().max())
for name in manual_matrix_gradients
}
gradient_allclose = {
name: torch.allclose(manual_matrix_gradients[name], autograd_gradients[name], atol=1e-12, rtol=1e-12)
for name in manual_matrix_gradients
}
for name in manual_matrix_gradients:
print(f"{name}: allclose={gradient_allclose[name]}, max error={gradient_max_errors[name]:.3e}")W1: allclose=True, max error=0.000e+00
b1: allclose=True, max error=0.000e+00
W2: allclose=True, max error=0.000e+00
b2: allclose=True, max error=0.000e+00plot_gradient_comparison(manual_matrix_gradients, autograd_gradients, gradient_max_errors)
Parameters such as \(W_1\) are leaves: they were created directly with gradient tracking enabled, so PyTorch stores their gradients in .grad. Intermediates such as \(Z_1\) are non-leaves: PyTorch uses their sensitivities during backward propagation but normally releases them instead of filling .grad. Call retain_grad() before backward() only when an intermediate is worth inspecting.
# The same graph without retain_grad(): the non-leaf gradient is not stored
import warnings
W1_probe = W1_batch.clone().detach().requires_grad_(True)
b1_probe = b1_batch.clone().detach().requires_grad_(True)
Z1_probe = X_batch @ W1_probe.T + b1_probe
probe_loss = torch.relu(Z1_probe).square().mean()
probe_loss.backward()
with warnings.catch_warnings():
warnings.simplefilter("ignore")
nonleaf_grad_without_retain = Z1_probe.grad
leaf_parameter_status = {
name: {"is_leaf": parameter.is_leaf, "grad stored": parameter.grad is not None}
for name, parameter in zip(("W1", "b1", "W2", "b2"), auto_parameters)
}
leaf_grad_is_stored = W1_probe.grad is not None
print("autograd parameter status:", leaf_parameter_status)
retained_nonleaf_matches_manual = torch.allclose(Z1_auto.grad, G_Z1, atol=1e-12, rtol=1e-12)
print("leaf W1 gradient stored:", leaf_grad_is_stored)
print("non-leaf Z1 gradient without retain_grad():", nonleaf_grad_without_retain)
print("retained Z1 gradient matches manual G_Z1:", retained_nonleaf_matches_manual)autograd parameter status: {'W1': {'is_leaf': True, 'grad stored': True}, 'b1': {'is_leaf': True, 'grad stored': True}, 'W2': {'is_leaf': True, 'grad stored': True}, 'b2': {'is_leaf': True, 'grad stored': True}}
leaf W1 gradient stored: True
non-leaf Z1 gradient without retain_grad(): None
retained Z1 gradient matches manual G_Z1: TrueCalling backward() adds new contributions to an existing .grad; it does not overwrite them. That behavior is useful when several losses should contribute to one update, but accidental repetition doubles the gradient.
first_backward_gradients = {name: value.clone() for name, value in autograd_gradients.items()}
repeat_prediction = torch.relu(X_batch @ W1_auto.T + b1_auto) @ W2_auto.T + b2_auto
repeat_loss = torch.mean((repeat_prediction - Y_batch) ** 2)
repeat_loss.backward() # no clearing: contributions accumulate
accumulation_errors = {
"W1": float((W1_auto.grad - 2 * first_backward_gradients["W1"]).abs().max()),
"b1": float((b1_auto.grad - 2 * first_backward_gradients["b1"]).abs().max()),
"W2": float((W2_auto.grad - 2 * first_backward_gradients["W2"]).abs().max()),
"b2": float((b2_auto.grad - 2 * first_backward_gradients["b2"]).abs().max()),
}
print("largest error from exactly doubled gradients:", max(accumulation_errors.values()))
for parameter in auto_parameters:
parameter.grad = None # essential before the next independent updatelargest error from exactly doubled gradients: 0.0Which earlier parameters will lose their gradients if the hidden activation is detached?
Detaching \(A_1\) declares it to be a constant for everything that follows. The output layer remains connected to the loss; the hidden layer does not.
W1_det = W1_batch.clone().detach().requires_grad_(True)
b1_det = b1_batch.clone().detach().requires_grad_(True)
W2_det = W2_batch.clone().detach().requires_grad_(True)
b2_det = b2_batch.clone().detach().requires_grad_(True)
A1_detached = torch.relu(X_batch @ W1_det.T + b1_det).detach()
detached_prediction = A1_detached @ W2_det.T + b2_det
detached_loss = torch.mean((detached_prediction - Y_batch) ** 2)
detached_loss.backward()
detach_gradient_status = {
"W1": W1_det.grad is not None, "b1": b1_det.grad is not None,
"W2": W2_det.grad is not None, "b2": b2_det.grad is not None,
}
print("gradient exists after detaching A1:", detach_gradient_status)gradient exists after detaching A1: {'W1': False, 'b1': False, 'W2': True, 'b2': True}The five lines from the original curve experiment now have a precise meaning:
prediction = model(x_train) # 1. forward prediction
loss = loss_fn(prediction, y_train) # 2. one scalar objective
optimizer.zero_grad() # 3. clear accumulated gradients
loss.backward() # 4. reverse-mode differentiation
optimizer.step() # 5. update every parameterThe optimizer performs step 5. Autograd performs step 4. The model, data, and loss chosen by us determine what those gradients mean.
Autograd does not decide what the network should learn. It computes how the chosen loss changes with every connected parameter.
The network’s geometry, the manual gradients, and autograd now describe the same learning process from different views.
# Final verification
expected_figures = [
"01-a-straight-line-tries",
"02-one-neuron-creates-a-bend",
"02a-linear-layer-and-scene2-network",
"03a-hidden-layer-network",
"03-several-bends-become-a-function",
"04-gradient-descent-arranges-the-bends",
"05-xor-hidden-representation",
"06-forward-pass-values",
"07-scalar-graph-forward",
"08-scalar-graph-backward",
"09-relu-zero-gradient-path",
"10-mini-batch-relu-mask",
"11-manual-versus-autograd",
]
assert figure_names == expected_figures
assert tuple(training_snapshots) == snapshot_epochs
assert set(trained_parameters) == {"hidden.weight", "hidden.bias", "output.weight", "output.bias"}
assert torch.equal(xor_decision, xor_y)
assert forward_pass_equality_error <= 1e-7
assert derivative_verification_error < 1e-8
assert inactive_values["grads"]["z"] == inactive_values["grads"]["w"] == inactive_values["grads"]["b"] == 0.0
assert active_sets == [[2, 4], [1, 4], [1, 3]]
assert all(gradient_allclose.values())
assert max(gradient_max_errors.values()) < 1e-12
assert all(item == {"is_leaf": True, "grad stored": True} for item in leaf_parameter_status.values())
assert leaf_grad_is_stored and nonleaf_grad_without_retain is None and retained_nonleaf_matches_manual
assert max(accumulation_errors.values()) < 1e-12
assert detach_gradient_status == {"W1": False, "b1": False, "W2": True, "b2": True}
assert all(torch.isfinite(value).all() for value in forward_pass_values.values())
for name in expected_figures:
for extension in ("svg", "png"):
path = FIG_DIR / f"{name}.{extension}"
assert path.is_file() and path.stat().st_size > 1000, path
notebook_runtime_seconds = time.perf_counter() - started
print(f"Verified {len(expected_figures)} figures in SVG and PNG.")
print(f"Notebook execution including exports: {notebook_runtime_seconds:.2f} seconds")Verified 13 figures in SVG and PNG.
Notebook execution including exports: 5.57 seconds