Lecture 3 — Backpropagation and Stochastic Optimization

Project connection. Project Step 3 first asks you to build a scalar reverse-mode automatic-differentiation engine. You will validate it on a graph with fan-out, against finite differences, and on the hand-wired diamond network from Step 2. You will then use PyTorch’s tensor-valued version of the same algorithm to train the neural $n$-gram model introduced in Lecture 2.

Chapter overview. The cross-entropy objective of a neural network is a scalar function of many parameters. We will represent its evaluation as a computational graph and use the chain rule in two orders. Forward mode carries one input perturbation toward the output. Reverse mode carries one output sensitivity toward every input, which is exactly the shape needed for neural-network training. We then turn from differentiation to optimization: a minibatch supplies a cheap random estimate of the full-corpus gradient, and stochastic gradient descent uses that estimate to update the parameters. These are distinct algorithms. Backpropagation computes a gradient; stochastic gradient descent decides which loss to differentiate and how to use the result.

0. The computational task at hand

We begin with a review of our setting. Let

$$\theta=(\theta_1,\ldots,\theta_P)^\top\in\mathbb R^P$$

collect all $P$ trainable parameters of a model. A parameter is a number whose value is chosen during training, such as an entry of a weight matrix, bias vector, or embedding matrix. Suppose the training set contains $N$ examples and the loss of example $i$ is $\ell_i(\theta)$. As in Lecture 2, the empirical loss is the average

$$\mathcal L(\theta) = \frac1N\sum_{i=1}^N\ell_i(\theta).$$

The word empirical means that the average is taken over observed data. The function being minimized is also called the objective function, or simply the objective.

The gradient is the column vector of first partial derivatives

$$\nabla_\theta\mathcal L(\theta) = \begin{pmatrix} \partial\mathcal L/\partial\theta_1\\ \vdots\\ \partial\mathcal L/\partial\theta_P \end{pmatrix}.$$

Lecture 2 defined gradient descent. Written with a possibly changing learning rate $\eta_t>0$, one update is

$$\theta^{(t+1)} = \theta^{(t)}-\eta_t\nabla\mathcal L(\theta^{(t)}).$$

The superscript $(t)$ is an index, not an exponent. One execution of the update is an iteration, and $\eta_t$ is the learning rate, or step size. If the gradient uses all $N$ examples, the method is full-batch gradient descent.

There are now two computational problems.

1. Differentiation: given a particular loss computation and a particular value of $\theta$, compute its gradient.
2. Optimization: choose which examples to use, choose a learning rate, and update $\theta$ so that the objective becomes small.

Reverse-mode automatic differentiation solves the first problem. Full-batch gradient descent or stochastic gradient descent solves the second.

Why not differentiate every formula by hand, as we did for logistic and softmax regression in Lecture 2? A neural network may contain millions or billions of parameters, but the more important difficulty is structural: loss along many paths, and the program changes as the architecture changes. We want an algorithm that differentiates the program itself.

---

0. Our goal

Our goal in this section is to create a very efficient way to do the necessary computations for gradient descent. The loss function for a neural network is much more complicated than that for a simple logistic regression, so we don’t want to deal with its formula, which is a function of all the weights and biases. Instead, we break it apart into its constituent pieces, and the gradient computation becomes many applications of chain rule combining these small pieces. To model this, we need to introduce the notion of a computational graph. This graph is not the neural network graph.

1. Computational graphs

A directed graph consists of nodes and directed edges. A directed path is a sequence of directed edges, followed in their arrow directions. A directed cycle is a nonempty directed path that returns to its starting node. A directed graph is acyclic if it has no directed cycle; a finite directed acyclic graph is abbreviated DAG.

A node with no incoming edges is a source. A node with no outgoing edges is a sink. The parents of a node are the nodes with edges into it, and its children are the nodes with edges out of it. A topological ordering of a finite DAG is an ordered list of its nodes so that every parent occurs before its child. Every finite DAG has at least one topological ordering, though it need not have only one.

An edge $u\to v$ says that the value at node $v$ directly depends on the value at node $u$.

In practice, the functions at each node should be ‘simple’ ones, in the sense that calculus students will be ok with differentiating them. A computational graph is designed to model the process of computing a complicated function, namely the output as a function on the source variables. Assigning values (typically scalars) to the source variables first and then evaluating the primitives in topological order to determine the output is called the forward pass (or forward sweep). The values created between the inputs and output are intermediate values.

The loss function for a neural network will be modelled by a computational graph. The source nodes hold the inputs and the parameters: the training example’s features and target, together with every entry of each weight matrix, bias vector, and embedding matrix.

A node has fan-out if it has more than one child. Fan-out means that one value influences the output along more than one directed path. In the figure, both $a$ and $b$ have fan-out.

The local derivative of a primitive is its derivative with respect to one parent while its other parents are held fixed.
The local derivative of the primitive $v$ with respect to parent $u$ is naturally associated to the edge $u \rightarrow v$, so we can think of this as a labelling of the directed edges.

---

2. The chain rule and forward-mode automatic differentiation

The multivariable chain rule says that if a scalar $r$ influences a scalar $q$ through intermediate variables $s_1,\ldots,s_m$, then

$$\frac{dq}{dr} = \sum_{j=1}^m \frac{\partial q}{\partial s_j} \frac{\partial s_j}{\partial r}.$$

Automatic differentiation (AD) is the algorithmic evaluation of derivatives by applying the chain rule. It is termed automatic because a program constructs and traverses the derivative computation.

Our principal application is to compute the gradient of $\mathcal{L}$ with respect to the parameters $\theta = (\theta_1, \ldots, \theta_P)$ for gradient descent.

Fix one parameter coordinate $\theta_s$. For every graph node compute

$$\dot v_i := \frac{\partial v_i}{\partial\theta_s}$$

and define the tangent to be $\dot v_i$ evaluated at $\theta$.

We initialize the input tangents by

$$\dot\theta_r=\mathbf1_{\{r=s\}},$$

where the indicator $\mathbf1_E$ equals $1$ when statement $E$ is true and $0$ otherwise. More concretely, this sets $\dot\theta_r$ to $1$ when $r=s$ and to $0$ otherwise.

This initialization is called a seed. It selects the input variable (in this case $\theta_s$) or combination of variables whose influence we want to follow.

By design, at the output node,

$$\dot v_M = \frac{\partial\mathcal L}{\partial\theta_s}.$$

The rule in the definition, of course, is just the chain rule at each node. Its computational efficiency comes from storing the tangent of every intermediate, so each shared subcomputation is performed exactly once per seed.

We now want to consider the runtime of forward-mode gradient computation. We’ll use some standard runtime terminology from computer science, called big-Oh notation. For positive functions $A(s)$ and $B(s)$ we write $A(s)=O(B(s))$ when $A(s)/B(s)$ is bounded above for all sufficiently large $s$, and $A(s)=\Theta(B(s))$ when it is bounded both above and below by positive constants for all sufficiently large $s$. A mathematician is more likely to write this as $A(s) < c B(s)$ for some constant $c$. These notations ignore fixed multiplicative constants and describe growth as the problem size increases. The constants are called hidden constants.

We measure cost in units of one evaluation: the work of computing $\mathcal L(\theta)$ itself by a single forward pass, evaluating each primitive in the graph once.

Proof

One forward pass with seed $e_s$, the $s$th standard basis vector, produces $\partial\mathcal L/\partial\theta_s$. Repeating for $s=1,\ldots,P$ produces all $P$ gradient coordinates. Each pass performs a bounded amount of derivative work for each node beside the work of the forward pass itself, so its cost is a constant multiple of one evaluation. $\square$

How does this compare with more naive approaches? A finite-difference gradient approximates each coordinate by $\bigl(\mathcal L(\theta+h e_s)-\mathcal L(\theta)\bigr)/h$, which also costs one evaluation per coordinate — the same $\Theta(P)$ — but each coordinate is only approximate, with the truncation-versus-roundoff tension you met in Project Step 2’s float64 warning. Forward mode has the same cost and is exact (up to floating point), because it applies exact derivative rules. A third option, symbolically differentiating a formula for $\mathcal L$, is exact but explosive: the expression for each $\partial\mathcal L/\partial\theta_s$ can grow far larger than the program that computes $\mathcal L$, and the $P$ expressions share no work. So $\Theta(P)$ evaluations looks hard to beat — and the surprise of the next section is that it can be beaten completely: reverse mode gets the whole gradient for $\Theta(1)$ evaluations.

Forward mode is excellent when there are few inputs and many outputs. The proposition above treats a single scalar output, but notice that one forward pass computes the tangent $\dot v_i$ of every node in the graph, not only of $v_M$. If a computation has many outputs — say $Q$ sink nodes rather than one — then a single pass with seed $e_s$ delivers the derivatives of all $Q$ outputs with respect to $\theta_s$ simultaneously. The cost of forward mode therefore scales with the number of inputs that must be seeded, and is independent of the number of outputs.

Neural-network training has the opposite shape: perhaps billions of input parameters and one scalar loss. This mismatch motivates reverse mode, which we will cover shortly.

---

3. Reverse mode and backpropagation

Fix input values (source values) in a computational graph. Suppose the output is $\mathcal{L}$. Suppose also that we have already done a forward pass, so the values at all nodes are computed as a function of the fixed source values.

As we saw, forward mode computes the value of $\partial v_i / \partial a$ at each node $v_i$ in terms of each source $a$. By contrast, reverse mode computes the value of $\partial \mathcal{L} / \partial v_i$ at each node $v_i$. Both of these compute the values iteratively from previously computed values. In the case of reverse mode, or backpropagation, we work from right to left, computing the values of leftward derivatives (closer to source) from those already computed to the right (closer to output).

The end goal is to compute $\partial \mathcal{L} / \partial a$ for each source $a$. In the case of the loss as a function of source parameters $\theta$, these partial derivatives, taken together, form the gradient $\nabla_\theta \mathcal{L}$ needed for gradient descent.

Consider a node $P$ in the graph. Its value is already known from the forward pass. Its children $c_1,\dots,c_n$ are the nodes that depend on $P$ directly, and because we are working backward from the output, each child’s adjoint $\partial\mathcal L/\partial c$ has already been computed.

The chain rule says

$$\frac{\partial\mathcal L}{\partial P} = \sum_{\text{children } c \text{ of } P} \underbrace{\frac{\partial\mathcal L}{\partial c}}_{\text{by previous steps}} \cdot \underbrace{\frac{\partial c}{\partial P}}_{\text{compute now}}$$

Thus, to compute $\frac{\partial\mathcal L}{\partial P}$, we can use the values $\frac{\partial\mathcal L}{\partial c}$ computed previously (further to the right), and compute the derivative $\frac{\partial c}{\partial P}$ at that node as needed.

The backpropagation algorithm performs this computation at each node, from right to left. To initialize, we begin at the sink with $\partial\mathcal L/\partial\mathcal L=1$. All derivatives are evaluated at the node values (computed during the forward pass), so we are storing and computing numbers, not formulas.

We now formalize this with the typical machine-learning language. For the fixed scalar output $\mathcal L$, define the adjoint of node $v_i$ by

$$\bar v_i := \frac{\partial\mathcal L}{\partial v_i}.$$

Some software calls the same quantity the node’s gradient.

The algorithm is called reverse accumulation because derivative contributions accumulate while moving opposite the forward edges. It is called reverse-mode automatic differentiation because it applies AD in that direction. Backpropagation is the customary neural-network name for reverse-mode AD applied to a loss.

Algorithm 3.1 is this equation with the sum spread out: each edge from $P$ to a child $c$ contributes one term, and the accumulation step adds that term into $\bar P$ when the edge is processed.

Let us revisit our running example, the graph from Section 1. At $(a,b)=(2,-1)$, the values are

$$u=-2, \qquad v=0, \qquad w=-0.761594\ldots, \qquad \mathcal L=0.$$

(This is computed by one simple forward pass.)

Seed $\bar{\mathcal L}=1$. We think of each node sending a message back along each of its incoming edges: the message from child $i$ to parent $j$ is the child’s adjoint times the local derivative on that edge, and a node’s adjoint is the sum of all the messages it receives. From the primitive $\mathcal L=vw$, the messages are

$$\bar v = \bar{\mathcal L}\,w =-0.761594\ldots, \qquad \bar w = \bar{\mathcal L}\,v =0.$$

(Each of $v$ and $w$ has only one child, so a single message is the whole adjoint.) The $\tanh$ node sends $b$ the message

$$\bar w(1-w^2)=0.$$

The addition $v=u+a$ has local derivative $1$ toward each parent, so it sends its adjoint unchanged to both: the message $\bar v=-0.761594\ldots$ goes to $u$, and the same message goes to $a$. Since $v$ is the only child of $u$,

$$\bar u =\bar v=-0.761594\ldots.$$

Finally, $u=ab$ sends $a$ the message

$$\bar u\, b =+0.761594\ldots,$$

and sends $b$ the message

$$\bar u\, a =-1.523188\ldots.$$

Both inputs received two messages, and each adjoint is the sum of the messages received. The two contributions to $\bar a$ cancel:

$$\bar a = -0.761594\ldots+0.761594\ldots =0,$$

while

$$\bar b = 0+(-1.523188\ldots) =-1.523188\ldots.$$

Thus reverse mode reproduces the two-pass forward-mode answer

$$\nabla_{(a,b)}\mathcal L(2,-1) = (0,-1.523188\ldots)^\top,$$

but it obtains both coordinates in one reverse pass.

Real neural-network primitives act on tensors rather than scalars. A tensor is just a finite multidimensional array of numbers: a vector is a $1$-tensor, and a matrix is a $2$-tensor. The shape of the tensor isn’t as important as the number of scalar entries, so let us consider nodes to be labelled with vectors. Let $x \mapsto y$ be one directed edge, where $x$ and $y$ are the vector labels of the nodes. We write

$$y=\varphi(x), \qquad x\in\mathbb R^m, \quad y\in\mathbb R^q,$$

with $\varphi$ differentiable. The Jacobian of $\varphi$ at $x$ is the $q\times m$ matrix of partial derivatives

$$J_\varphi(x) = \left[ \frac{\partial\varphi_i}{\partial x_j} \right]_{1\leq i\leq q,\,1\leq j\leq m}.$$

In forward mode, we compute $\dot y$ from $\dot x$. If the input tangent $\dot x\in\mathbb R^m$ is known, the output tangent is

$$\dot y = J_\varphi(x)\,\dot x,$$

called a Jacobian—vector product (JVP). This is exactly the directional derivative of $\varphi$ at $x$ in direction $\dot x$ (Section 2).

In reverse mode, we compute $\bar x$ from $\bar y$. If $\bar y=\nabla_y\mathcal L\in\mathbb R^q$ is known, the chain rule gives

$$\bar x = J_\varphi(x)^\top\bar y.$$

Equivalently, in row-vector notation, $\bar y^\top J_\varphi(x)$ is a vector—Jacobian product (VJP). Reverse-mode libraries implement a VJP rule for each primitive.

An arithmetic circuit is a computational graph whose primitives are arithmetic operations such as addition, multiplication, and division where the denominator is nonzero. Its size is its number of operation nodes. We first state the classical result for such circuits; standard AD systems extend the same idea with derivative rules for primitives such as exponential, logarithm, and $\tanh$.

Proof sketch

Store the value produced by every node during the forward evaluation. In reverse topological order, visit each edge once. Each visit multiplies a child adjoint by one local derivative and adds the result to a parent adjoint. An arithmetic primitive has a fixed number of parents, and computing each local derivative takes a fixed amount of work. The total reverse work is therefore at most a constant times the number of forward operations. The Baur—Strassen theorem makes this construction precise for rational arithmetic circuits, circuits whose operation nodes use addition, subtraction, multiplication, and division. $\square$

---

4. Backpropagation through softmax and a neural $n$-gram

We now write every step of backpropagation for a concrete model: the Bengio-style neural $n$-gram language model of Lecture 2, Section 4.4, which is also the model trained in Project Step 3. The model predicts the next token from the previous $k$ tokens. It looks up a learned embedding vector for each context token, concatenates the $k$ embeddings into one vector $\mathbf x$, passes $\mathbf x$ through a single $\tanh$ hidden layer, and applies one more affine map followed by softmax to produce a probability for each of the $V$ tokens in the vocabulary. The loss is the negative log of the probability assigned to the token that actually came next. Schematically,

$$\text{context tokens} \;\longmapsto\; \mathbf x \;\overset{\text{affine}}{\longmapsto}\; \mathbf u \;\overset{\tanh}{\longmapsto}\; \mathbf h_{\mathrm{ctx}} \;\overset{\text{affine}}{\longmapsto}\; \mathbf z \;\overset{\text{softmax}}{\longmapsto}\; \mathbf p \;\longmapsto\; \ell=-\log p_y.$$

Section 4.1 recalls the reverse rule for the softmax cross-entropy loss, a short interlude records the local rules for affine maps and coordinatewise $\tanh$, and Section 4.2 states the architecture precisely and assembles the rules into the complete backward pass.

4.1 Backpropagation: loss with respect to logit

Let $\mathbf z=(z_1,\ldots,z_V)^\top\in\mathbb R^V$ be the logit vector: one unconstrained score for each of $V$ possible target tokens. Our goal in this subsection is to compute $\nabla_{\mathbf z} \ell$.

Lecture 2 defined

$$p_j = \operatorname{softmax}(\mathbf z)_j = \frac{e^{z_j}}{\sum_{c=1}^V e^{z_c}}.$$

Let $y\in\{1,\ldots,V\}$ be the target-token index. Its one-hot vector $\mathbf e_y\in\mathbb R^V$ has a $1$ in coordinate $y$ and zeros elsewhere. Define the log-sum-exp function

$$\operatorname{lse}(\mathbf z) = \log\sum_{c=1}^V e^{z_c}.$$

The single-example categorical cross-entropy loss (Lecture 2, Definition 3.2) is

$$\ell(\mathbf z,y) = -\log p_y = -z_y+\operatorname{lse}(\mathbf z).$$

We already computed the gradient of this loss with respect to the logits in Lecture 2, Proposition 3.4, for softmax regression. There the logits were an affine function of the features, but the computation used only the dependence of $\ell$ on $\mathbf z$, so it applies verbatim here: with $\mathbf p=\operatorname{softmax}(\mathbf z)$,

$$\boxed{ \nabla_{\mathbf z}\ell=\mathbf p-\mathbf e_y. }$$

Recall that the coordinates of this gradient sum to zero, because a common shift of every logit leaves the loss unchanged.

The vector $\mathbf p-\mathbf e_y$ is often called the output error signal, meaning the adjoint at the logit vector from which the rest of backpropagation starts. It is not the difference between a predicted token and a token number; token indices have no numerical distance. It is a difference between two probability vectors.

Lecture 2 worked this out for the logits $\mathbf z=(0,\log2,\log3)^\top$ of Example 3.3: with the second token as target, the gradient $(\tfrac16,-\tfrac23,\tfrac12)^\top$ lowers each wrong logit by its current probability and raises the target logit by $1-p_2$. The update is largest when the model assigns little probability to the target.

Backpropagation: derivatives of affine maps and tanh

We need only affine maps and coordinatewise $\tanh$ to propagate the output signal through the Step 3 network. Here are some useful equations for relevant derivatives.

First suppose

$$\mathbf q=W\mathbf r+\mathbf b,$$

where $W\in\mathbb R^{m\times n}$, $\mathbf r\in\mathbb R^n$, and $\mathbf b,\mathbf q\in\mathbb R^m$. If the incoming adjoint is $\bar{\mathbf q}=\nabla_{\mathbf q}\ell$, then

$$\boxed{ \bar{\mathbf r}=W^\top\bar{\mathbf q}, \qquad \nabla_W\ell=\bar{\mathbf q}\,\mathbf r^\top, \qquad \nabla_{\mathbf b}\ell=\bar{\mathbf q}. }$$

The matrix $\bar{\mathbf q}\,\mathbf r^\top$ is an outer product: the product of a column vector and a row vector. Its $(i,j)$ entry is $\bar q_i r_j$, as required because $q_i=\sum_jW_{ij}r_j+b_i$.

Next suppose $\mathbf h=\tanh\mathbf u$ coordinatewise. Since $d(\tanh s)/ds=1-\tanh^2s$,

$$\boxed{ \bar{\mathbf u} = \bar{\mathbf h}\odot(\mathbf 1-\mathbf h\odot\mathbf h). }$$

The symbol $\odot$ denotes the Hadamard product, or elementwise product: $(\mathbf r\odot\mathbf s)_i=r_is_i$. The $\mathbf 1$ in the formula denotes the vector of ones of the appropriate length.

These rules illustrate a general principle: the forward shapes determine the backward shapes. If $W$ is $m\times n$, then its gradient must also be $m\times n$; if $\mathbf r$ has length $n$, then $W^\top\bar{\mathbf q}$ has length $n$. Checking shapes catches many backpropagation mistakes before checking numbers.

4.2 One complete backward pass

Let us recall the notation of the neural $n$-gram model of Lecture 2 (Example 4.7). Let $C\in\mathbb R^{V\times d}$ be an embedding matrix. An embedding is a learned vector representation, and row $C_a\in\mathbb R^d$ represents token $a$. For a context $(a_1,\ldots,a_k)$, an embedding lookup reads those rows and concatenates them into one column vector:

$$\mathbf x = \operatorname{concat}(C_{a_1},\ldots,C_{a_k})^\top \in\mathbb R^{kd}.$$

Concatenation joins vectors end to end. The network then computes

$$\mathbf u=W_1\mathbf x+\mathbf b_1, \qquad \mathbf h_{\mathrm{ctx}}=\tanh \mathbf u,$$ $$\mathbf z=W_2\mathbf h_{\mathrm{ctx}}+\mathbf b_2, \qquad \mathbf p=\operatorname{softmax}(\mathbf z), \qquad \ell=-\log p_y.$$

Here

$$W_1\in\mathbb R^{h\times kd}, \quad \mathbf b_1\in\mathbb R^h, \quad W_2\in\mathbb R^{V\times h}, \quad \mathbf b_2\in\mathbb R^V.$$

The vector $\mathbf u$ is the hidden preactivation, the value before the nonlinear $\tanh$ is applied. The vector $\mathbf h_{\mathrm{ctx}}$ is the hidden representation of this context.

The logit gradient, also called the output error signal, is

$$\nabla_{\mathbf z}\ell =\mathbf p-\mathbf e_y.$$

Applying the tanh and affine rules above gives

$$\nabla_{W_2}\ell = (\nabla_{\mathbf z}\ell)\mathbf h_{\mathrm{ctx}}^\top, \qquad \nabla_{\mathbf b_2}\ell = \nabla_{\mathbf z}\ell,$$ $$\bar{\mathbf h}_{\mathrm{ctx}} = W_2^\top\nabla_{\mathbf z}\ell,$$ $$\nabla_{\mathbf u}\ell = \bar{\mathbf h}_{\mathrm{ctx}} \odot (\mathbf 1-\mathbf h_{\mathrm{ctx}}\odot\mathbf h_{\mathrm{ctx}}),$$ $$\nabla_{W_1}\ell = (\nabla_{\mathbf u}\ell)\mathbf x^\top, \qquad \nabla_{\mathbf b_1}\ell = \nabla_{\mathbf u}\ell,$$

and

$$\bar{\mathbf x} = W_1^\top\nabla_{\mathbf u}\ell.$$

Finally, split $\bar{\mathbf x}\in\mathbb R^{kd}$ into $k$ consecutive blocks of length $d$. The block for position $r$ is added to the gradient of row $C_{a_r}$. If the same token occurs in two context positions, both blocks are added to the same row. This operation is often called a scatter-add: values associated with indices are added back into the corresponding locations of a larger array.

The embedding matrix exhibits parameter sharing. Parameter sharing means that the same parameter participates in several parts of the computation. Every occurrence of token $a$ reads the same row $C_a$, so the row has fan-out across positions and examples. Its gradient is the sum of all returning contributions.

For a minibatch, each example acquires a leading batch coordinate. The same equations then use batched matrix multiplication, and parameter gradients sum or average over the batch coordinate.

---

5. Autograd in practice

Autograd is the common name for a software system that records differentiable operations and automatically computes their derivatives. It is an implementation of automatic differentiation, not a different calculus algorithm.

5.1 What a scalar Value object stores

Project Step 3 wraps every scalar in a Value object. Conceptually, each object stores four things:

1. data: the node’s forward value;
2. grad: its adjoint, initialized to zero;
3. _parents: the nodes on which it directly depends; and
4. _backward: the local VJP rule that adds this node’s contribution to its parents.

A leaf node is a source node created directly by the user, such as a parameter or input. A root for a backward pass is the output node from which reverse traversal begins—normally the scalar loss. A topological sort is an algorithm that constructs a topological ordering. A depth-first search follows one unvisited dependency path as far as it can before returning to try another. The small engine performs such a search from the loss, appends a node after visiting its parents, and reverses the resulting list for backpropagation.

For multiplication

$$q=rs,$$

the closure stored on $q$ implements

$$r.\mathrm{grad}\mathrel{+}=q.\mathrm{grad}\,s.\mathrm{data}, \qquad s.\mathrm{grad}\mathrel{+}=q.\mathrm{grad}\,r.\mathrm{data}.$$

A closure is a function that retains access to values from the scope in which it was created. Here the local backward closure remembers $r$, $s$, and $q$ even after the forward multiplication has returned.

Calling loss.backward() then performs exactly Algorithm 3.1:

1. topologically order the reachable graph;
2. set the loss gradient to $1$;
3. traverse the order backward; and
4. call each node’s local backward rule.

5.2 Why gradients accumulate and must be cleared

AD systems use addition for two different kinds of accumulation.

• Within one graph, fan-out creates several paths to the same node.
• Across several backward calls, a user may intentionally sum gradients from several losses or microbatches. A microbatch is one small part of a larger logical batch processed separately to save memory.

For this reason, PyTorch normally adds into .grad rather than replacing it. A training loop that wants a fresh minibatch gradient must clear old parameter gradients before the next backward pass. Forgetting this step silently changes the update into a sum over the current and previous minibatches.

The parameter update itself should not become part of the differentiated model computation. PyTorch’s torch.no_grad() context temporarily disables graph recording. A tensor is detached when it shares numerical data with a computation but is treated as having no derivative connection to that computation. Detaching in the middle of a model accidentally cuts the backward path; detaching a diagnostic value or disabling recording during the parameter update is intentional.

5.3 Three independent checks

A derivative implementation is most convincing when independent methods agree.

1. Hand calculation: derive a small graph with the chain rule.
2. Finite differences: compare selected coordinates with nearby loss evaluations.
3. Independent systems: compare your scalar engine, finite differences, and PyTorch on one input to the Step 2 diamond network.

For scalars $g$ and $\widehat g$, the absolute error is $|g-\widehat g|$. A scale-aware comparison often uses the relative error

$$\frac{|g-\widehat g|} {\max(1,|g|,|\widehat g|)}.$$

The denominator avoids declaring two tiny, harmless numbers wildly different. No single numerical error threshold fits every data type and graph, but Step 3’s small float64 checks should agree much more closely than a large low-precision training run.

Common failure patterns are diagnostic:

• correct chain graphs but incorrect fan-out graphs suggest assignment instead of addition;
• gradients missing from an early node suggest a wrong traversal order or an accidental detach;
• gradients exactly multiplied by the number of iterations suggest that old gradient values were not cleared;
• plausible AD gradients that disagree with finite differences only at extreme step sizes suggest numerical, rather than chain-rule, error.

---

6. Connection to Project Step 3

6.1 From the scalar engine to the neural model

Step 3 first validates Algorithm 3.1 on $\mathcal L=(ab+a)\tanh b$. It then rebuilds one forward pass of the Step 2 diamond classifier from scalar $\texttt{Value}$ objects and differentiates its loss with respect to the input. The closed formula, finite differences, your engine, and PyTorch all produce the same gradient. This closes the loop between a network assembled by hand and an algorithm that differentiates its computational graph.

The project next trains the neural $n$-gram model of Section 4 with minibatch SGD. Its parameter count is

$$Vd+(kd)h+h+hV+V.$$

The five terms count the embedding matrix $C$, first-layer weights $W_1$, first-layer bias $b_1$, second-layer weights $W_2$, and second-layer bias $b_2$. The count grows linearly with context length $k$, whereas an unrestricted next-token matrix indexed by every length-$k$ context has $V^{k+1}$ entries.

The training set is the set used to update parameters. A disjoint validation set is held out from parameter updates and used to estimate performance on unseen examples during model development. Training loss and validation loss are the same loss formula averaged over these two sets.

6.2 Reading the context-length sweep

The reference run in Project Step 3 reports:

context length $k$	parameters	training loss	validation loss	bits/character
1	11,601	2.4706	2.4805	3.579
3	15,697	1.9171	1.9821	2.860
5	19,793	1.9054	2.0018	2.888
8	25,937	1.9372	2.0329	2.933

A nat is the information unit obtained when cross-entropy uses natural logarithms. A bit is the corresponding unit for base-$2$ logarithms; dividing loss in nats by $\log2$ gives bits per predicted character.

When $k=1$, the model sees only the immediately preceding character, just as a bigram model does. The hidden network is wide enough to assign an independent logit vector to each of the finitely many input characters, so its model family contains the bigram behavior. Its validation loss should therefore be near the earlier bigram result when optimization succeeds. A model family, or hypothesis class, is the set of functions obtainable as its parameters vary.

Moving from $k=1$ to $k=3$ supplies useful context and sharply improves validation loss. Moving to $k=5$ or $k=8$ does not help at the fixed hidden width and training budget. Model capacity is the range and complexity of functions a model family can represent. Flattening the longer vector spreads the fixed hidden capacity across more inputs and also makes the optimization problem harder.

An optimization bottleneck occurs when the training procedure fails to find a low-loss member of the model family. A representational bottleneck occurs when the family itself cannot express the needed function. Both can produce underfitting, meaning that important patterns remain unfitted even in the training data. Overfitting has a different signature: training loss improves while validation performance worsens because the model has specialized too strongly to the training set.

The $k=8$ run has worse training loss than the $k=3$ run, so its poorer validation result cannot be explained by classic overfitting alone. The small $k=5$ training improvement leaves room for a mixture of effects, but the matrix as a whole points toward representational and optimization bottlenecks rather than a large train—validation gap.

6.3 What to carry forward

The scalar engine is intentionally small, but it exposes principles that remain true in large systems:

• computations form a DAG for one forward pass;
• local VJP rules compose into a global gradient;
• fan-out requires addition;
• reverse order is a dependency requirement, not a coding preference;
• tensor backpropagation avoids materializing giant Jacobians;
• minibatching changes the gradient estimator, not the chain rule; and
• clearing gradients and excluding updates from the graph are explicit parts of the training algorithm.

---

Summary

A differentiable program can be represented as a computational graph. Forward mode propagates a tangent from selected inputs toward all later nodes and naturally computes Jacobian—vector products. Computing every coordinate of a scalar-output gradient this way requires one pass per parameter.

Reverse mode seeds the scalar loss with adjoint $1$ and propagates sensitivities in reverse topological order. Each edge contributes a child adjoint times a local derivative, and fan-out contributions are added. This produces every parameter derivative in one reverse pass. The cheap-gradient principle says the arithmetic cost is a constant multiple of the forward computation, while stored or recomputed activations account for the memory cost.

For softmax cross-entropy, the logit adjoint is $\mathbf p-\mathbf e_y$. Affine, $\tanh$, and embedding-lookup VJPs route this signal through the Bengio neural $n$-gram model. The same rules operate on scalars in your Value engine and on tensors in PyTorch.

In practice, training pairs backpropagation with minibatch stochastic gradient descent: a random batch of $B$ examples supplies a cheap estimate of the full-corpus gradient that is correct on average, and the parameters move against it by a learning rate. Batch size, learning-rate schedule, and hardware throughput are tuned together; Lecture 6 takes up these practical choices.

Exercises (paired with Step 3)

A star marks a Project Step 3 task.

1. ★ For $\mathcal L=(ab+a)\tanh b$ at $(a,b)=(2,-1)$, reproduce Example 2.2 and the reverse pass of Section 3 without looking at the tables. Label every local derivative and every accumulated contribution.
2. ★ Implement scalar Value primitives for addition, multiplication, $\tanh$, exponential, logarithm, and powers. Explain which forward values each backward closure must retain.
3. Construct the smallest computational graph you can find on which replacing addition by assignment during reverse accumulation gives a wrong derivative. Give values for which the wrong answer does not accidentally equal the right one.
4. Give a DAG with two valid topological orderings. Reverse one valid ordering and verify Algorithm 3.1. Then exhibit an invalid reverse order and identify the contribution it loses.
5. Compute the directional derivative $\nabla\mathcal L^\top r$ of the running example at $(a,b)=(2,-1)$ in direction $r=(3,-2)^\top$ in two ways: one forward-mode pass seeded with $(\dot a,\dot b)=r$ (a seed need not be a standard basis vector!), and the dot product of $r$ with the gradient you already know.
6. ★ Re-derive the softmax cross-entropy gradient $\nabla_{\mathbf z}\ell=\mathbf p-\mathbf e_y$ of Lecture 2, Proposition 3.4. Verify it with your scalar engine and with PyTorch for a randomly chosen five-logit vector.
7. Derive the affine-layer VJP coordinate by coordinate. Check the shapes of all three outputs $\bar r$, $\nabla_W\ell$, and $\nabla_b\ell$.
8. For a context $(a,b,a)$, write explicitly how the three blocks of $\bar{\mathbf x}$ scatter-add into the embedding-matrix gradient. Which rows are exactly zero for this one example?
9. Count the activations stored by an $L$-layer width-$h$ MLP on a batch of size $B$. Propose a checkpointing scheme and state what it stores and recomputes.
10. Show that permuting the $h$ hidden coordinates of a one-hidden-layer network, together with the corresponding rows and columns of its weight matrices, preserves the represented function.
11. ★ At $\mathbf{x}=(0.5,0.25)^\top$ and $\gamma=4$, use the Step 2 diamond network with target “inside.” Derive $\partial\ell/\partial x_1=\partial\ell/\partial x_2 =2\gamma p(\mathrm{outside}\mid\mathbf{x})$. Reproduce both derivatives with your scalar engine, central finite differences, and PyTorch autograd, and report the maximum absolute disagreement.
12. In Algorithm 5.1, deliberately omit gradient clearing for three identical forward/backward passes without updating parameters. Predict the result before running it, then explain the observed gradient values.
13. The $k=8$ reference model has worse training and validation loss than the $k=3$ model. Explain why this evidence points toward underfitting rather than overfitting, and name at least two interventions that would distinguish an optimization bottleneck from a representational bottleneck.

Pointers

Baydin, Pearlmutter, Radul, and Siskind, Automatic Differentiation in Machine Learning: a Survey; Griewank and Walther, Evaluating Derivatives; Baur and Strassen (1983); Robbins and Monro (1951); Bengio, Ducharme, Vincent, and Jauvin, A Neural Probabilistic Language Model (2003); and Olah’s backpropagation essay. See the resources page, Project Step 3, and Lecture 2.