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.

---

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$.

A computational graph is designed to model the process of computing a complicated function. Assigning values (typically scalars) to the source variables first and then evaluating the primitives in topological order is called the forward pass. The values created between the inputs and output are intermediate values.

A neural network evaluating its loss on one example is exactly such a computational graph. The source nodes hold the inputs and the parameters: the example’s features and target, together with every entry of each weight matrix, bias vector, and embedding matrix. The primitives are the elementary operations the network performs: the multiplications and additions inside each affine layer, the coordinatewise nonlinearity such as $\tanh$, and the exponentials, sums, and divisions inside softmax and the loss. Each layer is a subgraph, composing layers chains those subgraphs together, and the scalar loss is the output node $v_M$. In a neural network the intermediates are often called activations, meaning values produced by units or layers during the forward pass.

It’s also possible to assign vectors, matrices or higher-dimensional tensors to the source variables (a tensor is a multidimensional rectangular array of numbers). Expanding every tensor entry into its own scalar node recovers a more complicated computational graph, so without loss of generality, we can consider assigning scalars to the variables labelling the source nodes.

We will use one expression throughout the lecture as a running example.

$$\mathcal L=(ab+a)\tanh b.$$

Introduce the intermediate values

$$u=ab, \qquad v=u+a, \qquad w=\tanh b, \qquad \mathcal L=vw.$$

The following picture is the corresponding computational graph. The formula is nested, but the graph makes every direct dependence explicit.

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.

At the Project Step 3 test point $(a,b)=(2,-1)$, the forward pass gives

$$\begin{array}{c|c|c} \text{node} & \text{primitive} & \text{value}\\ \hline a & \text{input} & 2\\ b & \text{input} & -1\\ u & ab & -2\\ v & u+a & 0\\ w & \tanh b & -0.761594\ldots\\ \mathcal L & vw & 0. \end{array}$$

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. The local derivatives in this graph are

$$\frac{\partial u}{\partial a}=b, \qquad \frac{\partial u}{\partial b}=a,$$ $$\frac{\partial v}{\partial u}=1, \qquad \frac{\partial v}{\partial a}=1, \qquad \frac{\partial w}{\partial b}=1-w^2,$$

and

$$\frac{\partial\mathcal L}{\partial v}=w, \qquad \frac{\partial\mathcal L}{\partial w}=v.$$

These are deliberately small facts. Automatic differentiation obtains a large derivative by composing local derivatives according to the graph.

Here is the graph with the local derivatives labelling the 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}.$$

The sum is essential: it combines every immediate route from $r$ through the $s_j$ to $q$.

Automatic differentiation (AD) is the algorithmic evaluation of derivatives by applying exact local derivative rules and the chain rule to a numerical computation. 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 define its tangent

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

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.

Forward mode need not seed only one coordinate. For an input direction $r\in\mathbb R^P$, seed $\dot\theta=r$. The output tangent becomes

$$\dot{\mathcal L} = \left.\frac{d}{d\varepsilon} \mathcal L(\theta+\varepsilon r) \right|_{\varepsilon=0} = \nabla\mathcal L(\theta)^\top r.$$

This scalar is the directional derivative of $\mathcal L$ in direction $r$.

To compare algorithmic costs, 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$. These notations ignore fixed multiplicative constants and describe growth as the problem size increases.

Proof

One seed $e_s$, the $s$th standard basis vector, produces only $\partial\mathcal L/\partial\theta_s$. Here a standard basis vector has one coordinate equal to $1$ and every other coordinate equal to $0$. Repeating for $s=1,\ldots,P$ produces all $P$ gradient coordinates. Each sweep performs a constant amount of derivative work beside each forward primitive, so its cost is a constant multiple of evaluation. $\square$

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 sweep 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 sweep 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

Forward mode asks, “Fix an input. How does this node change with respect to that input?” In other words, we compute $\partial v_i / \partial a$ for a fixed input $a$ at each node $v_i$. Reverse mode asks the complementary question, “Fix an output. How does this node affect that output?” In other words, we compute $\partial \mathcal{L} / \partial v_i$ for a fixed output $\mathcal{L}$ at each node $v_i$.

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}.$$

The bar marks sensitivity of the output to that node. Some software calls the same quantity the node’s gradient.

Proof

The output seed is correct because $\partial\mathcal L/\partial\mathcal L=1$. Suppose that the adjoints of all children $i$ of node $j$ are correct. The multivariable chain rule gives

$$\frac{\partial\mathcal L}{\partial v_j} = \sum_{i:j\in\operatorname{pa}(i)} \frac{\partial\mathcal L}{\partial v_i} \frac{\partial v_i}{\partial v_j}.$$

Reverse topological order guarantees that each child’s adjoint is ready before $j$ is processed. The additive updates give exactly the terms in this sum. Induction backward through the ordering proves the claim. $\square$

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.

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

$$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-sweep forward-mode answer

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

but it obtains both coordinates in one reverse sweep.

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 $x$, passes $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\; x \;\overset{\text{affine}}{\longmapsto}\; u \;\overset{\tanh}{\longmapsto}\; h_{\mathrm{ctx}} \;\overset{\text{affine}}{\longmapsto}\; z \;\overset{\text{softmax}}{\longmapsto}\; p \;\longmapsto\; \ell=-\log p_y.$$

Section 4.1 derives 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 $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_z \ell$.

Lecture 2 defined

$$p_j = \operatorname{softmax}(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 $e_y\in\mathbb R^V$ has a $1$ in coordinate $y$ and zeros elsewhere. Define the log-sum-exp function

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

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

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

Proof

For coordinate $j$,

$$\frac{\partial}{\partial z_j}\operatorname{lse}(z) = \frac{e^{z_j}}{\sum_c e^{z_c}} =p_j,$$

while

$$\frac{\partial(-z_y)}{\partial z_j} = -\mathbf1_{\{j=y\}}.$$

Thus $\partial\ell/\partial z_j=p_j-\mathbf1_{\{j=y\}}$, which is the $j$th coordinate of $p-e_y$. Both $p$ and $e_y$ have coordinate sum one, so their difference has sum zero. $\square$

The vector $p-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.

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

$$q=Wr+b,$$

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

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

The matrix $\bar q\,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 $h=\tanh u$ coordinatewise. Since $d(\tanh s)/ds=1-\tanh^2s$,

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

The symbol $\odot$ denotes the Hadamard product, or elementwise product: $(r\odot s)_i=r_is_i$. The $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 $r$ has length $n$, then $W^\top\bar 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:

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

Concatenation joins vectors end to end. The network then computes

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

Here

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

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

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

$$\nabla_z\ell =p-e_y.$$

Applying the tanh and affine rules above gives

$$\nabla_{W_2}\ell = (\nabla_z\ell)h_{\mathrm{ctx}}^\top, \qquad \nabla_{b_2}\ell = \nabla_z\ell,$$ $$\bar h_{\mathrm{ctx}} = W_2^\top\nabla_z\ell,$$ $$\nabla_u\ell = \bar h_{\mathrm{ctx}} \odot (1-h_{\mathrm{ctx}}\odot h_{\mathrm{ctx}}),$$ $$\nabla_{W_1}\ell = (\nabla_u\ell)x^\top, \qquad \nabla_{b_1}\ell = \nabla_u\ell,$$

and

$$\bar x = W_1^\top\nabla_u\ell.$$

Finally, split $\bar 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 Theorem 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 Theorem 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 sweep 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 sweep. 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 $p-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 sweep 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 Theorem 3.1. Then exhibit an invalid reverse order and identify the contribution it loses.
5. Compute the directional derivative of the running example at $(a,b)=(2,-1)$ in direction $r=(3,-2)^\top$ in two ways: one forward-mode sweep and the dot product $\nabla\mathcal L^\top r$.
6. ★ Prove Proposition 4.1. 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 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.