# Module 2: Calculus and Optimization

> Derivatives, gradients, the chain rule and backpropagation from scratch, with gradient descent and why too large a learning rate makes training diverge.

- Part I: Foundations · core track · Foundational · 4 hours
- By the end you can compute gradients through a small computation graph by hand, and explain why a learning rate can be too large.
- Builds on: [Module 1: Shapes and Linear Algebra](https://transformers-and-math.drmhse.com/md/01-shapes-and-linear-algebra.md)
- Web page (interactive figures, labs and equation explainers): https://transformers-and-math.drmhse.com/modules/02-calculus-and-optimization
- From *Transformers & Math* by Mike Chumba: https://transformers-and-math.drmhse.com

A model learns by nudging its numbers so that its mistakes shrink. To know *which way* to nudge each one, it needs the derivative: how much an output changes when an input is nudged. This module builds from derivatives to the two algorithms that train every model in the course: gradient descent, which takes the steps, and backpropagation, which finds the direction of every step at once.

If you have taken a first calculus course, sections 2.1 and 2.2 are revision.

## 2.1 Derivatives: Slopes and the Rules for Finding Them

The *derivative* $f'(x)$ or $\frac{df}{dx}$ is the slope of $f$ at $x$: how much $f(x)$ changes per unit change in $x$, for a change small enough that the curve looks straight. Formally it is the limit of "rise over run" as the run shrinks to zero:

$$
f'(x) = \lim_{h \to 0} \frac{f(x + h) - f(x)}{h}
$$

In practice you combine a few known slopes with a few rules.

| function | derivative | where it appears |
|---|---|---|
| $x^n$ | $n x^{n-1}$ | squared errors, variances |
| $e^x$ | $e^x$ | softmax (section 4.3) |
| $\ln x$ | $1/x$ | the log in every loss |
| $\tanh x$ | $1 - \tanh^2 x$ | recurrent networks (section 5.2), soft-capping (section 16.5) |
| $\max(0, x)$ (ReLU) | $1$ for $x > 0$, $0$ for $x < 0$ | feed-forward layers (section 7.4) |

The rules that combine them:

- **Sum rule.** The slope of $f + g$ is $f' + g'$.
- **Constant factor.** The slope of $c f$ is $c f'$.
- **Product rule.** The slope of $f g$ is $f' g + f g'$.
- **Quotient rule.** The slope of $f / g$ is $\frac{f' g - f g'}{g^2}$.

> **Worked example.** Take $f(x) = x^2 e^x$. The product rule with $x^2$ and $e^x$ gives $f'(x) = 2x \cdot e^x + x^2 \cdot e^x = (2x + x^2)e^x$. At $x = 1$ the slope is $3e \approx 8.15$: nudging $x$ up by $0.001$ raises $f$ by about $0.00815$.

That last sentence is the whole use of a derivative in this course. **A derivative is a prediction: change the input by a small $\Delta x$, and the output changes by about $f'(x)\,\Delta x$.** Training uses it to choose which way to move every number in the model.

## 2.2 Partial Derivatives and the Gradient

A model's *loss* $\mathcal{L}$ is one number measuring its mistakes, and it depends on millions of weights, written as one long vector $\theta = [\theta_1, \dots, \theta_n]$. The *partial derivative* $\partial \mathcal{L} / \partial \theta_i$ is the slope with respect to one weight, treating every other weight as a constant. One partial derivative per weight gives the *gradient*:

$$
\nabla_\theta \mathcal{L} = \left[\frac{\partial \mathcal{L}}{\partial \theta_1}, \, \frac{\partial \mathcal{L}}{\partial \theta_2}, \, \dots, \, \frac{\partial \mathcal{L}}{\partial \theta_n}\right]
$$

> **Worked example.** For $\mathcal{L}(a, b) = a^2 b + 3b$, holding $b$ fixed gives $\partial \mathcal{L}/\partial a = 2ab$, and holding $a$ fixed gives $\partial \mathcal{L}/\partial b = a^2 + 3$. At $a = 1, b = 2$ the gradient is $[4, 4]$.

The gradient has the same shape as $\theta$. It points where the loss rises fastest, and its length says how steeply. The opposite direction lowers the loss fastest.

## 2.3 The Chain Rule and the Jacobian

Models feed one function into the next. The *chain rule*: **the slope of a chain is the product of the slopes of its links.** If $y = g(x)$ and $z = f(y)$, then

$$
\frac{dz}{dx} = \frac{dz}{dy} \cdot \frac{dy}{dx}
$$

Each link's slope is taken at the point the chain actually passes through.

> **Worked example.** Let $z = \tanh(3x)$. The inner link is $y = 3x$, with slope $3$. The outer link is $z = \tanh y$, with slope $1 - \tanh^2 y$. So $\frac{dz}{dx} = (1 - \tanh^2(3x)) \cdot 3$. At $x = 0$ that is $3$; at $x = 1$, $\tanh 3 \approx 0.995$ and the slope has collapsed to about $0.03$.

This is the danger the module ends on: one link with a tiny slope makes the whole chain's slope tiny.

### Vectors In, Vectors Out: The Jacobian
Most links take a vector and return a vector, so every output entry has a slope with respect to every input entry. The table of these slopes is the *Jacobian*. For $\mathbf{y} = g(\mathbf{x})$ with $\mathbf{x}$ of length $n$ and $\mathbf{y}$ of length $m$:

$$
J = \frac{\partial \mathbf{y}}{\partial \mathbf{x}} = \begin{bmatrix} \frac{\partial y_1}{\partial x_1} & \cdots & \frac{\partial y_1}{\partial x_n} \\ \vdots & \ddots & \vdots \\ \frac{\partial y_m}{\partial x_1} & \cdots & \frac{\partial y_m}{\partial x_n} \end{bmatrix}
$$

Row $i$ says how output $i$ responds to each input. The chain rule keeps its form: the Jacobian of a chain is the matrix product of the links' Jacobians.

Two Jacobians are worth knowing by heart:

- For a linear map $\mathbf{y} = \mathbf{x}W$, output $j$ is $\sum_i x_i W_{ij}$, so $\partial y_j / \partial x_i = W_{ij}$: the Jacobian is $W^\top$.
- For a function applied entry by entry, such as $\tanh$ or ReLU, output $i$ depends only on input $i$, so the Jacobian is diagonal: $\operatorname{diag}(1 - \tanh^2 x_1, \dots, 1 - \tanh^2 x_n)$.

Section 4.3 works out a Jacobian that is not diagonal, the softmax's, and uses it to explain why attention divides its scores by a constant.

## 2.4 Gradient Descent

To reduce the loss, step a little in the direction opposite to the gradient:

$$
\theta \longleftarrow \theta - \eta \, \nabla_\theta \mathcal{L}
$$

The arrow means "replace $\theta$ with". The step size $\eta$ (the Greek letter eta) is the *learning rate*. Too small and training crawls; too large and each step overshoots, so the loss bounces or blows up.

Computing the loss over the whole training set for every step would be far too slow, so each step uses a small random *batch* of examples, whose gradient is noisy but right on average (Module 3 makes this precise). This is *stochastic gradient descent* (SGD). Section 10.2 covers AdamW, the variant that trains almost every large model.

> **Interactive figure** `gradient-descent`, on the web page: https://transformers-and-math.drmhse.com/modules/02-calculus-and-optimization#figure-gradient-descent
>
> **Try:** The loss is $\theta^2$, whose slope is $2\theta$, so each step multiplies $\theta$ by $1 - 2\eta$. Before you try it, work out the largest learning rate for which $\theta$ still shrinks toward zero.
>
> **What happens:** $\theta$ shrinks while $|1 - 2\eta| < 1$, that is for $0 < \eta < 1$. At $\eta = 0.9$ the factor is $-0.8$: $\theta$ flips sign every step but still settles. At $\eta = 1.1$ the factor is $-1.2$, every step lands further out, and training diverges. For a real model the safe range depends on how sharply the loss curves, which is why learning rates are tuned rather than derived.

> **In words.** The figure runs gradient descent on theta squared from theta equal to 1.6. With eta equal to 0.1 it creeps toward zero; with 0.4 it arrives in a few steps; with 0.9 it zig-zags but settles; with 1.1 every step lands further out and the loss grows without limit.

## 2.5 Backpropagation

*Backpropagation* (Rumelhart, Hinton & Williams, 1986) applies the chain rule from the loss backward, one operation at a time, over a *computation graph*: one node per operation, recorded as the forward pass (input to loss) runs.

Each operation receives $\partial \mathcal{L} / \partial \mathbf{y}$, the gradient with respect to its output, and produces two things: the gradient for its own weights, and the gradient with respect to its input, passed to the operation below. For a linear layer $\mathbf{y} = \mathbf{x} W$:

$$
\frac{\partial \mathcal{L}}{\partial W} = \mathbf{x}^\top \frac{\partial \mathcal{L}}{\partial \mathbf{y}}, \quad \frac{\partial \mathcal{L}}{\partial \mathbf{x}} = \frac{\partial \mathcal{L}}{\partial \mathbf{y}} W^\top
$$

The second formula is the chain rule with the Jacobian $W^\top$ from section 2.3. The first is an outer product (section 1.3): the weight $W_{ij}$ connects input $i$ to output $j$, so its gradient is "how big input $i$ was" times "how much the loss cares about output $j$".

> **Worked example.** Take one input $x = 2$, a weight $w = 0.5$, $y = wx = 1$, then $a = \tanh(y) \approx 0.762$, and the loss $\mathcal{L} = (a - 1)^2 \approx 0.0567$. Go backward:
> 1. $\partial \mathcal{L}/\partial a = 2(a - 1) \approx -0.477$.
> 2. $\partial \mathcal{L}/\partial y = -0.477 \cdot (1 - a^2) \approx -0.477 \cdot 0.420 \approx -0.200$.
> 3. $\partial \mathcal{L}/\partial w = -0.200 \cdot x \approx -0.401$, and $\partial \mathcal{L}/\partial x = -0.200 \cdot w \approx -0.100$.
>
> The negative gradient for $w$ says that raising $w$ lowers the loss, which makes sense: a larger $w$ pushes $a$ toward 1.

> **Interactive figure** `backprop-graph`, on the web page: https://transformers-and-math.drmhse.com/modules/02-calculus-and-optimization#figure-backprop-graph
>
> **Try:** In the figure, set $x$ to $4$. Before stepping backward, guess whether the gradient reaching $w$ gets larger or smaller than at $x = 2$.
>
> **What happens:** Smaller, even though $x$ is larger. At $x = 4$, $y = 2$ and $\tanh$ is nearly flat there ($1 - a^2 \approx 0.07$), so almost nothing survives the $\tanh$ link. The factor $x$ in the last step cannot make up for a near-zero factor earlier in the chain.

> **In words.** The figure draws the example as a chain: x and w feed a multiply, then tanh, then a squared error. Each step moves the gradient one node backward, multiplying by that node's local slope, and labels the edge with the running value. The inputs are editable.

Backpropagation visits each operation once forward and once back, so *every* gradient costs only about twice the loss, however many weights there are. It needs the forward values, which is why training stores activations and uses far more memory than running a model (section 11.1). In PyTorch, `loss.backward()` runs the backward pass for every weight.

> **Browser lab** `grad-check`: an experiment that runs in the reader's browser on the web page: https://transformers-and-math.drmhse.com/modules/02-calculus-and-optimization#lab-grad-check

## 2.6 What Goes Wrong in Long Chains

Backpropagation multiplies one local slope per operation. A deep network, or one that runs over many time steps, makes a long product, and long products behave badly:

- If the typical factor is a little below 1, the product shrinks exponentially. Fifty factors of $0.9$ give $0.9^{50} \approx 0.005$. The weights at the start of the chain receive almost no gradient and stop learning. This is a *vanishing gradient*.
- If the typical factor is a little above 1, the product grows exponentially. Fifty factors of $1.1$ give about $117$. The steps become huge and training blows up. This is an *exploding gradient*.

Flat regions produce the small factors, like the saturated $\tanh$ of section 2.3, or a softmax pushed to a near one-hot output (section 4.3).

Much of a transformer's design keeps this product well behaved: residual connections give the gradient a path with no multiplications (section 7.2), normalization keeps values in range (section 7.3), and attention scales its scores (section 6.3). Module 5 shows the problem at its worst, in the recurrent networks that came before transformers.

## Exercises

**Exercise (check).** In one sentence, what does the derivative $f'(x) = 5$ tell you about $f$ near $x$?

*Solution:* That nudging $x$ up by a small amount $\Delta x$ raises $f$ by about $5\,\Delta x$. The derivative is a local prediction of how the output responds to a small change in the input.

**Exercise (apply).** Using the quotient rule, differentiate $p = \dfrac{e^{a}}{e^{a} + e^{b}}$ with respect to $a$, and simplify the result in terms of $p$.

*Solution:* The top is $e^a$ with slope $e^a$; the bottom is $S = e^a + e^b$ with slope $e^a$ too. The quotient rule gives $\frac{e^a S - e^a e^a}{S^2} = \frac{e^a}{S} - \left(\frac{e^a}{S}\right)^2 = p - p^2 = p(1 - p)$. This is the two-class softmax, and $p(1-p)$ is the factor section 4.3 tracks: it is tiny when $p$ is near 0 or 1.

**Exercise (apply).** For $\mathcal{L}(a, b) = (ab - 1)^2$, compute $\partial \mathcal{L} / \partial a$ at $a = 2$, $b = 1$.

*Answer:* 2 (within 0)

*Solution:* By the chain rule, $\partial \mathcal{L}/\partial a = 2(ab - 1) \cdot b$. At $a = 2, b = 1$: $2(2 - 1) \cdot 1 = 2$.

**Exercise (apply).** Backpropagate through $y = \tanh(wx)$ with loss $\mathcal{L} = y$, at $x = 1$, $w = 0$. What is $\partial \mathcal{L}/\partial w$?

*Answer:* 1 (within 0)

*Solution:* $\partial \mathcal{L}/\partial y = 1$. The $\tanh$ link has slope $1 - \tanh^2(0) = 1$. The multiply link gives $\partial(wx)/\partial w = x = 1$. The product is $1 \cdot 1 \cdot 1 = 1$.

**Exercise (apply).** Gradient descent on the loss $\theta^2$ multiplies $\theta$ by $(1 - 2\eta)$ each step. Above what learning rate does it diverge?

*Answer:* 1 (within 0)

*Solution:* It diverges when $|1 - 2\eta| > 1$. For positive $\eta$ that means $1 - 2\eta < -1$, so $\eta > 1$.

**Exercise (apply).** A chain multiplies 30 factors of $0.8$. What is the product, to three decimal places?

*Answer:* 0.00124 (within 0.00002)

*Solution:* $0.8^{30} = e^{30 \ln 0.8} = e^{-6.694} \approx 0.00124$. A gradient passing through these 30 links keeps about an eight-hundredth of its size.

**Exercise (extend).** Write the Jacobian of $\mathbf{y} = \operatorname{ReLU}(\mathbf{x})$ for $\mathbf{x} = [2, -1, 0.5]$. What does it do to a gradient passed backward through it?

*Solution:* ReLU acts entry by entry, so its Jacobian is diagonal with the slopes on the diagonal: $\operatorname{diag}(1, 0, 1)$. Multiplying a backward gradient by it copies the entries for positive inputs unchanged and sets the entry for the negative input to zero. A ReLU blocks the gradient wherever its input was negative, which is why a unit whose input is always negative stops learning.

### Core Academic References
- **Rumelhart, Hinton & Williams (1986)**: *Learning representations by back-propagating errors*. Nature, 323: 533–536.
- **Goodfellow, Bengio & Courville (2016)**: *Deep Learning*, chapters 4 and 6. MIT Press.
- **Baydin et al. (2018)**: *Automatic Differentiation in Machine Learning: a Survey*. JMLR 18(153): 1–43.