5. Calculus for learning

Why this layer exists

Training a model means choosing its weights. Chapter 4 fixed what “good weights” means: low cross-entropy on the training data, equivalently high likelihood. For a categorical distribution the best weights had a closed form, the normalised counts. For a network they do not, because every output probability depends on all the weights at once, and there are far too many weights to search. What is cheap to compute, at the current weights, is how the loss would change if each weight moved a little. That vector of sensitivities is the gradient, and it points in the direction in which the loss rises fastest. Gradient descent takes a small step the other way and repeats.

Every model trained in this book is trained this way: the bigram, MLP and transformer language models of Part II, the chat fine-tune of chapter 13, and the QLoRA fine-tune of chapter 46 on dixie’s 12 GB RTX 3060, where the base weights stay frozen and only a small set of added low-rank weights receives the steps. The optimizers used on real models rescale and smooth the step, per coordinate and over time; the step is still built from the gradient. The one knob this chapter explains fully is the learning rate, the step size. Too small and training crawls; too large and the loss oscillates and then grows without bound until it overflows to inf and turns into nan. The threshold between the two is a number you can compute, and the lab computes it for a real problem and watches the runs on either side of it.

Chapter 6 builds an engine that computes gradients automatically. The only way to trust such an engine is to compare its answers with a slow, independent estimate, and the numerical gradient built here is that estimate.

Mechanism

The derivative is a local linear approximation

Zoom in far enough on the graph of a smooth function and it looks like a straight line. The slope of that line at $x$ is the derivative $f'(x)$, and the line itself is the approximation

$f(x + h) \approx f(x) + f'(x)\,h ,$

good for small $h$. Precisely, $f'(x)$ is the number that makes the error of this approximation small compared with $h$ itself:

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

when the limit exists. The ratio is the slope of the chord from $x$ to $x + h$; the limit is the slope the chords approach as the second point slides in. Write the error as

$r(h) = f(x + h) - f(x) - f'(x)\,h.$

The definition says exactly that $r(h)/h \to 0$.

The reading that matters for learning: if $f'(x) > 0$, a small decrease in $x$ lowers $f$; if $f'(x) < 0$, a small increase does; and the size of $f'(x)$ says how much $f$ changes per unit of $x$, to first order. $f'(x) = 0$ says nothing to first order, which is what happens at a minimum.

Derivatives from the definition

The functions that appear in the losses and layers of this book are built from powers, the exponential and the logarithm. Each derivative follows from the limit.

Squares. Expand and cancel:

$\frac{(x + h)^2 - x^2}{h} = \frac{x^2 + 2xh + h^2 - x^2}{h} = 2x + h \;\to\; 2x .$

Positive integer powers. Multiplying out $(x + h)^n$ gives $x^n$, then $n$ terms with one factor of $h$ (each $x^{n-1}h$), then terms with $h^2$ or higher powers. This is the binomial expansion:

$(x + h)^n = x^n + n\,x^{n-1}h + \frac{n(n-1)}{2}\,x^{n-2}h^2 + \cdots + h^n .$

Subtract $x^n$ and divide by $h$: the first remaining term is $n\,x^{n-1}$ and every other term still carries at least one factor of $h$, so

$\frac{d}{dx}\,x^n = n\,x^{n-1} .$

The exponential. $e^{x + h} = e^x e^h$, so

$\frac{e^{x + h} - e^x}{h} = e^x \cdot \frac{e^h - 1}{h} .$

Everything rests on the second factor. Take the power series $e^h = 1 + h + h^2/2! + h^3/3! + \cdots$ as the definition of the exponential; then

$\frac{e^h - 1}{h} = 1 + \frac{h}{2!} + \frac{h^2}{3!} + \cdots ,$

and for $\lvert h \rvert \le 1$ the terms after the 1 add up to at most $\lvert h \rvert\,(1/2! + 1/3! + \cdots) = \lvert h \rvert\,(e - 2)$, which goes to 0 with $h$. The limit is 1, and

$\frac{d}{dx}\,e^x = e^x .$

$e$ is the base for which the exponential’s slope at 0 is exactly 1; for any other base $a$, $a^x = e^{x \log a}$ has slope $\log a$ there. At $h = 0.001$ the ratio is 1.0005.

The logarithm. $y = \log x$ means $e^y = x$, for $x > 0$. Let $k = \log(x + h) - \log x$, the change in the log; then $x + h = e^{y + k}$, so $h = e^{y + k} - e^y$, and

$\frac{\log(x + h) - \log x}{h} = \frac{k}{e^{y + k} - e^{y}} = 1 \Big/ \frac{e^{y + k} - e^{y}}{k} .$

As $h \to 0$, $k \to 0$ too, since the logarithm is continuous, and the denominator is a chord slope of the exponential tending to $e^y = x$. So

$\frac{d}{dx}\,\log x = \frac{1}{x} .$

Sum, product and chain rules

Sums and constant multiples. The difference quotient of $a f + b g$ is $a$ times that of $f$ plus $b$ times that of $g$, and limits respect sums, so $(af + bg)' = af' + bg'$. Differentiation is linear.

Products. Write $f(x + h) = f + f'h + r_f$ and $g(x + h) = g + g'h + r_g$, with everything on the right evaluated at $x$ and both remainders small compared with $h$. Multiply:

$f(x + h)\,g(x + h) = fg + (f'g + fg')\,h + \big[\,f'g'h^2 + (f + f'h)\,r_g + (g + g'h)\,r_f + r_f r_g\,\big] .$

Every term in the brackets, divided by $h$, goes to 0: $h^2/h = h$, and $r_f/h$, $r_g/h$ vanish by definition. What is left is the product rule, $(fg)' = f'g + fg'$.

Composition. For $g(f(x))$, first approximate the inner function, then the outer one at the point $u = f(x)$. The inner function moves by $k = f(x + h) - f(x) = f'(x)\,h + r_f(h)$, and the outer function, moved from $u$ by $k$, changes by $g'(u)\,k + r_g(k)$ with $r_g(k)/k \to 0$ as $k \to 0$. Substituting $k$:

$g(f(x + h)) = g(f(x)) + g'(f(x))\,f'(x)\,h + \big[\,g'(f(x))\,r_f(h) + r_g(k)\,\big] .$

The bracket, divided by $h$, goes to 0. Its first term does because $r_f(h)/h \to 0$. For the second, $k \to 0$ as $h \to 0$, and $k/h = f'(x) + r_f(h)/h$ stays bounded, so $r_g(k)/h = \big(r_g(k)/k\big)\,(k/h) \to 0$; when $k = 0$ the term is $r_g(0) = 0$ outright. The coefficient of $h$ is the derivative:

$\big(g \circ f\big)'(x) = g'\big(f(x)\big)\,f'(x) .$

This is the chain rule. A composition of local linear approximations is a local linear approximation whose slope is the product of the slopes. Chapter 3 showed that composing linear maps multiplies their matrices; the chain rule for functions of many variables is the same statement with each slope replaced by a matrix, and chapter 6 is built on it.

Two consequences come at once. Any real power: for $x > 0$, $x^r = e^{r \log x}$, so the chain rule gives $e^{r \log x} \cdot r/x = r\,x^{r - 1}$, extending the integer rule. And the quotient rule is not needed separately: $1/g = g^{-1}$, whose derivative is $-g'/g^2$ by the power rule and the chain rule.

A worked composite. Take

$f(w) = \log\big(1 + e^{-w}\big) .$

This is the cross-entropy of a two-outcome model whose probability for the observed outcome is $1/(1 + e^{-w})$, since $-\log \frac{1}{1 + e^{-w}} = \log(1 + e^{-w})$. Peel it from the outside. The outer function is $\log u$ with $u = 1 + e^{-w}$, contributing $1/u$. The inner function’s derivative is the derivative of $e^{-w}$, itself a composite: outer $e^v$, inner $v = -w$, so $e^{-w} \cdot (-1)$. Multiply:

$f'(w) = \frac{1}{1 + e^{-w}} \cdot \big(-e^{-w}\big) = -\frac{e^{-w}}{1 + e^{-w}} = -\frac{1}{1 + e^{w}} ,$

the last step multiplying top and bottom by $e^{w}$. At $w = 0$ this is $-1/2$, and the chord slope $(f(0.001) - f(-0.001))/0.002$ also comes out $-0.5$. The derivative is always negative and tends to 0 as $w$ grows: raising the logit of the correct outcome always lowers the loss, by less and less once the model is confident. Exercise (a) meets the same function from the other side.

Partial derivatives

A loss depends on many numbers. For $f(x_1, \ldots, x_d)$, the partial derivative $\partial f / \partial x_i$ is the ordinary derivative in $x_i$ with every other coordinate held fixed. For $f(x, y) = x^2 y + e^{y}$:

$\frac{\partial f}{\partial x} = 2xy , \qquad \frac{\partial f}{\partial y} = x^2 + e^{y} ,$

treating $y$ as a constant in the first and $x$ in the second.

Partials combine into a local linear approximation in $d$ dimensions. Move from $\mathbf{x}$ by a small step $\delta$, shape (d,), one coordinate at a time. The first move changes $f$ by about $(\partial f / \partial x_1)\,\delta_1$; the second by about $(\partial f / \partial x_2)\,\delta_2$, with that partial evaluated at the already-shifted point, which differs from its value at $\mathbf{x}$ by an amount that vanishes as $\delta \to \mathbf{0}$ as long as the partials are continuous, as they are for every function in this book. Summing the $d$ moves,

$f(\mathbf{x} + \delta) \approx f(\mathbf{x}) + \sum_{i=1}^{d} \frac{\partial f}{\partial x_i}\,\delta_i .$

The gradient points uphill

Collect the partials into a vector, the gradient:

$\nabla f(\mathbf{x}) = \left( \frac{\partial f}{\partial x_1}, \ldots, \frac{\partial f}{\partial x_d} \right) ,$

shape (d,), the same shape as $\mathbf{x}$. The sum above is a dot product, so the local approximation is

$f(\mathbf{x} + \delta) \approx f(\mathbf{x}) + \nabla f(\mathbf{x}) \cdot \delta .$

This is the form chapter 4 used, ahead of time, to find the maximum-likelihood distribution. Now ask which direction increases $f$ fastest. Take a unit vector $\mathbf{u}$ and step a distance $s$ along it, $\delta = s\,\mathbf{u}$. The rate of change of $f$ along $\mathbf{u}$, the directional derivative, is

$D_{\mathbf{u}} f = \lim_{s \to 0} \frac{f(\mathbf{x} + s\,\mathbf{u}) - f(\mathbf{x})}{s} = \nabla f \cdot \mathbf{u} = \lVert \nabla f \rVert \cos\varphi ,$

where $\varphi$ is the angle between $\nabla f$ and $\mathbf{u}$, by chapter 2’s identity $\mathbf{a} \cdot \mathbf{b} = \lVert \mathbf{a} \rVert\,\lVert \mathbf{b} \rVert \cos\varphi$ with $\lVert \mathbf{u} \rVert = 1$. Only $\cos\varphi$ depends on the direction. It is largest, 1, at $\varphi = 0$: $f$ rises fastest along $\nabla f / \lVert \nabla f \rVert$, at rate $\lVert \nabla f \rVert$. It is smallest, $-1$, at $\varphi = \pi$: $f$ falls fastest along $-\nabla f$. At $\varphi = \pi/2$ the rate is 0: directions perpendicular to the gradient are, to first order, level, and the gradient is perpendicular to the contour through $\mathbf{x}$.

Gradient descent

Given a loss $L(\theta)$ over parameters $\theta$, shape (P,), gradient descent repeats

$\theta_{t+1} = \theta_t - \eta\, \nabla L(\theta_t) ,$

with a learning rate $\eta > 0$. The local approximation with $\delta = -\eta \nabla L$ says what one step does:

$L(\theta_{t+1}) \approx L(\theta_t) - \eta\, \lVert \nabla L(\theta_t) \rVert^2 ,$

a decrease whenever the gradient is nonzero and $\eta$ is small enough for the approximation to hold. Descent stops making progress where $\nabla L = \mathbf{0}$, a stationary point. In general that can be a minimum, a maximum or a saddle, and a minimum need not be the lowest one; for the loss this chapter fits, there is a single stationary point and it is the minimum. The only question left is what “small enough” means.

How large a step: one dimension

Take the simplest function with a minimum, $f(x) = \frac{a}{2} x^2$ with $a > 0$. Its derivative is $a x$, so one step of gradient descent is

$x_{t+1} = x_t - \eta\, a\, x_t = (1 - \eta a)\,x_t , \qquad \text{hence} \qquad x_t = (1 - \eta a)^t\, x_0 .$

The iterates go to the minimum at 0 exactly when $\lvert 1 - \eta a \rvert < 1$, that is, $-1 < 1 - \eta a < 1$, which for positive $\eta$ and $a$ is

$0 < \eta < \frac{2}{a} .$

Within that range there are two regimes, split at $\eta = 1/a$, where one step lands exactly on the minimum. Below it the factor is between 0 and 1: every step moves toward the minimum without passing it, and a small $\eta a$ means many short steps. Above it the factor is between $-1$ and 0: each step overshoots to the other side, landing closer than it started, and the iterates zig-zag inward. At $\eta = 2/a$ the factor is $-1$ and the iterate bounces between $x_0$ and $-x_0$ forever; beyond it each overshoot lands farther out, and $\lvert x_t \rvert$ grows geometrically.

The constant $a$ is the curvature, the second derivative $f''(x) = a$: how fast the slope changes. Any smooth function looks like this near a minimum $x^*$. There $f'(x^*) = 0$, so the local approximation applied to $f'$ gives $f'(x^* + e) \approx f''(x^*)\,e$, and a gradient step from $x^* + e$ moves the error to $e - \eta f''(x^*)\,e = (1 - \eta f''(x^*))\,e$: the same recursion with $a = f''(x^*)$. The stability limit is 2 divided by the curvature.

Many parameters: curvature by direction

The loss of this chapter’s lab is exactly quadratic in its parameters, of the form

$L(\theta) = L^* + \frac{1}{2}\, \mathbf{e}^\top \mathbf{H}\, \mathbf{e} , \qquad \mathbf{e} = \theta - \theta^* ,$

with $\theta$, $\theta^*$ and $\mathbf{e}$ of shape (P,), and $\mathbf{H}$ a symmetric (P, P) matrix. The double sum is $\frac12 \sum_{i}\sum_{j} e_i H_{ij} e_j$; its partial derivative in $e_k$ collects the terms with $i = k$ and those with $j = k$, giving $\frac12 \big(\sum_j H_{kj} e_j + \sum_i e_i H_{ik}\big) = (\mathbf{H}\mathbf{e})_k$ by symmetry. So $\nabla L = \mathbf{H}\mathbf{e}$, and differentiating once more, $\partial^2 L / \partial \theta_j \partial \theta_k = H_{jk}$: $\mathbf{H}$ is the matrix of second partial derivatives, the Hessian. A gradient step subtracts $\eta \mathbf{H} \mathbf{e}$ from the error:

$\mathbf{e}_{t+1} = (\mathbf{I} - \eta \mathbf{H})\, \mathbf{e}_t .$

In general $\mathbf{H}\mathbf{e}$ points in a different direction from $\mathbf{e}$, so the coordinates are coupled. But some directions $\mathbf{v}$ are only stretched: $\mathbf{H}\mathbf{v} = \lambda \mathbf{v}$ for a scalar $\lambda$. Such a $\mathbf{v}$ is an eigenvector of $\mathbf{H}$, and $\lambda$ its eigenvalue. Along an eigenvector the step is $(\mathbf{I} - \eta \mathbf{H})\mathbf{v} = (1 - \eta\lambda)\,\mathbf{v}$, which is the one-dimensional recursion with $a = \lambda$: $\lambda$ is the curvature in that direction. A symmetric (P, P) matrix has $P$ mutually perpendicular eigenvectors, the spectral theorem (exercise (c) finds them by hand for $P = 2$), so any error splits into $P$ components that evolve independently, each multiplied by its own $1 - \eta\lambda_i$ per step.

Three facts follow. The loss has a unique minimum at $\theta^*$ only if every $\lambda_i > 0$, a bowl in every direction; a zero eigenvalue is a direction along which the loss is flat. The iteration converges only if every factor has magnitude below 1, so

$0 < \eta < \frac{2}{\lambda_{\max}} :$

the steepest direction sets the largest usable learning rate. And the flattest direction sets the pace: its factor is $1 - \eta\lambda_{\min}$, and with $\eta$ below $2/\lambda_{\max}$ that is at least $1 - 2\lambda_{\min}/\lambda_{\max}$. When the ratio $\lambda_{\max}/\lambda_{\min}$, the condition number, is large, the safe step is tiny for the flat directions and plain gradient descent is slow. np.linalg.eigvalsh returns the eigenvalues of a symmetric matrix in ascending order.

The analysis is exact only for a quadratic, but near a minimum every smooth loss is close to one, for the reason the one-dimensional case gave. A learning rate that makes a real training run blow up is, locally, above $2/\lambda_{\max}$ of the loss surface where the run happens to be.

The linear-regression gradient

The lab fits the model $\hat{y}_i = \mathbf{x}_i \cdot \mathbf{w} + b$ to $n$ examples. $\mathbf{X}$, shape (n, d), holds the inputs as rows $\mathbf{x}_i$; $\mathbf{w}$ has shape (d,), $b$ is a scalar and the targets $\mathbf{y}$ have shape (n,). The loss is the mean squared error

$L(\mathbf{w}, b) = \frac{1}{n} \sum_{i=1}^{n} r_i^2 , \qquad r_i = \mathbf{x}_i \cdot \mathbf{w} + b - y_i = \sum_{j=1}^{d} X_{ij} w_j + b - y_i ,$

where $r_i$ is the residual of example $i$. In index notation first. $r_i$ is linear in the parameters, with $\partial r_i / \partial w_j = X_{ij}$ and $\partial r_i / \partial b = 1$. By the chain rule with outer function $u^2$, $\partial (r_i^2) / \partial w_j = 2 r_i X_{ij}$, and by linearity

$\frac{\partial L}{\partial w_j} = \frac{2}{n} \sum_{i=1}^{n} r_i\, X_{ij} , \qquad \frac{\partial L}{\partial b} = \frac{2}{n} \sum_{i=1}^{n} r_i .$

The first sum runs over the row index of $X_{ij}$, which is how an entry of $\mathbf{X}^\top \mathbf{r}$ is formed: $(\mathbf{X}^\top \mathbf{r})_j = \sum_i (\mathbf{X}^\top)_{ji}\, r_i = \sum_i X_{ij}\, r_i$. In matrix form,

$\nabla_{\mathbf{w}} L = \frac{2}{n}\, \mathbf{X}^\top (\mathbf{X}\mathbf{w} + b - \mathbf{y}) , \qquad \frac{\partial L}{\partial b} = \frac{2}{n}\, \mathbf{1} \cdot (\mathbf{X}\mathbf{w} + b - \mathbf{y}) .$

Shapes: $\mathbf{X}\mathbf{w}$ is (n, d) @ (d,), giving (n,); the scalar $b$ broadcasts across it; the residual vector $\mathbf{r}$ is (n,); $\mathbf{X}^\top$ is (d, n), so $\mathbf{X}^\top\mathbf{r}$ is (d,), the shape of $\mathbf{w}$, as a gradient must be. Read as a sum, the gradient is the average over examples of $2 r_i \mathbf{x}_i$: each example pulls $\mathbf{w}$ along its own input, in proportion to how wrong its prediction is.

The bias fits the same pattern if it is treated as a weight on an input that is always 1. Append a column of ones to get $\tilde{\mathbf{X}}$, shape (n, d+1), and stack the parameters into $\theta = (\mathbf{w}, b)$, shape (d+1,). Then $\mathbf{r} = \tilde{\mathbf{X}}\theta - \mathbf{y}$ and both results above are one formula, $\nabla L = \frac{2}{n} \tilde{\mathbf{X}}^\top (\tilde{\mathbf{X}}\theta - \mathbf{y})$. It is linear in $\theta$ with matrix

$\mathbf{H} = \frac{2}{n}\, \tilde{\mathbf{X}}^\top \tilde{\mathbf{X}} , \qquad \text{shape } (d+1, d+1) ,$

which is symmetric, since $(\tilde{\mathbf{X}}^\top \tilde{\mathbf{X}})^\top = \tilde{\mathbf{X}}^\top \tilde{\mathbf{X}}$, and is the Hessian. The minimum $\theta^*$ is where the gradient vanishes, $\tilde{\mathbf{X}}^\top \tilde{\mathbf{X}}\, \theta^* = \tilde{\mathbf{X}}^\top \mathbf{y}$, and subtracting that from the gradient at any $\theta$ gives $\nabla L(\theta) = \mathbf{H}(\theta - \theta^*)$: the previous section’s quadratic exactly, with no approximation. The largest stable learning rate for linear regression is $2/\lambda_{\max}$ of this matrix, which depends only on the inputs. For one input feature,

$\mathbf{H} = 2 \begin{pmatrix} \overline{x^2} & \overline{x} \\ \overline{x} & 1 \end{pmatrix} ,$

with bars denoting averages over the $n$ examples: the entries of $\tilde{\mathbf{X}}^\top \tilde{\mathbf{X}}$ are $\sum_i x_i^2$, $\sum_i x_i$ and $\sum_i 1 = n$. Exercise (c) evaluates it for the lab data.

Numerical gradients

The definition of the derivative suggests computing it directly: pick a small $h$ and evaluate the forward difference $(f(x + h) - f(x))/h$. How wrong is it? Extend the local approximation to more terms, $f(x + h) = c_0 + c_1 h + c_2 h^2 + c_3 h^3 + \cdots$. Setting $h = 0$ gives $c_0 = f(x)$; differentiating $k$ times in $h$ and then setting $h = 0$ leaves only $k!\,c_k$ on the right, so $c_k = f^{(k)}(x)/k!$, the $k$-th derivative over $k!$. This is the Taylor expansion:

$f(x + h) = f(x) + f'(x)\,h + \frac{f''(x)}{2}\,h^2 + \frac{f'''(x)}{6}\,h^3 + \cdots .$

Subtract $f(x)$ and divide by $h$:

$\frac{f(x + h) - f(x)}{h} = f'(x) + \frac{f''(x)}{2}\,h + O(h^2) .$

The error shrinks in proportion to $h$: $O(h)$. The central difference evaluates on both sides. Replacing $h$ by $-h$ flips the sign of the odd-power terms:

$f(x + h) - f(x - h) = 2 f'(x)\,h + \frac{f'''(x)}{3}\,h^3 + O(h^5) ,$

the even-power terms cancelling in the subtraction. Divide by $2h$:

$\frac{f(x + h) - f(x - h)}{2h} = f'(x) + \frac{f'''(x)}{6}\,h^2 + O(h^4) .$

The error is now $O(h^2)$: dividing $h$ by 10 divides it by 100. For the same two evaluations per coordinate, the central difference is far more accurate, and it is what numerical_grad uses.

Taylor’s error says a smaller $h$ is always better. Floating point says otherwise. A float64 stores about 16 significant decimal digits; the unit roundoff is $\epsilon_{\text{mach}} = 2^{-53} \approx 1.1 \times 10^{-16}$, and each computed value of $f$ carries an error of order $\epsilon_{\text{mach}} \lvert f \rvert$. For small $h$, $f(x + h)$ and $f(x - h)$ agree in most of their leading digits, and subtracting them cancels those digits and leaves the rounding errors behind at the same absolute size. Dividing by $2h$ then magnifies them: the rounding part of the error is of order $\epsilon_{\text{mach}} \lvert f \rvert / h$ and grows as $h$ shrinks. The total error of the central difference is roughly

$\frac{\lvert f''' \rvert}{6}\, h^2 + \frac{\epsilon_{\text{mach}} \lvert f \rvert}{h} ,$

minimised where the two terms are of comparable size, around $h \approx \epsilon_{\text{mach}}^{1/3} \approx 5 \times 10^{-6}$ when $f$ and $f'''$ are of order 1. For the forward difference the same balance, $h$ against $\epsilon_{\text{mach}}/h$, puts the best $h$ near $\epsilon_{\text{mach}}^{1/2} \approx 10^{-8}$ and the best error near $10^{-8}$ as well. numerical_grad defaults to $h = 10^{-6}$, near the central-difference optimum; exercise (d) measures both curves.

A numerical gradient costs two evaluations of $f$ per parameter. For a single (1000, 1000) weight matrix that is two million evaluations of the loss, which rules it out for training. Chapter 6 computes exact gradients by the chain rule instead, and uses numerical_grad on small inputs to check them; chapter 10 checks its attention layer the same way.

Walkthrough

The reference module is py/tinygpt/descent.py.

def numerical_grad(f, x, eps=1e-6):
    """Central-difference estimate of the gradient of scalar f at x: (f(x+h) - f(x-h)) / 2h per coordinate."""
    x = np.array(x, dtype=np.float64)  # a private copy we can perturb
    grad = np.zeros_like(x)
    for idx in np.ndindex(x.shape):
        orig = x[idx]
        x[idx] = orig + eps
        hi = f(x)
        x[idx] = orig - eps
        lo = f(x)
        x[idx] = orig
        grad[idx] = (hi - lo) / (2 * eps)
    return grad

np.array always copies, so the function perturbs a private array and the caller’s x is never modified; converting to float64 also matters, because adding $10^{-6}$ to an entry of an integer array would be truncated away. np.ndindex(x.shape) yields every index tuple of an array of any shape, so the same function differentiates with respect to a vector, a bias, or a (in, out) weight matrix. Each coordinate is perturbed up, then down, then restored before the next, so every partial derivative is taken with the other coordinates at their original values, which is what a partial derivative is. f must return a scalar; grad[idx] = ... assigns one number per entry.

def mse(pred, y):
    """Mean squared error."""
    r = np.asarray(pred, dtype=np.float64) - np.asarray(y, dtype=np.float64)
    return float((r * r).mean())


def linreg_grad(w, b, x, y):
    """Gradient of mse(x @ w + b, y) with respect to w (shape (d,)) and b (scalar)."""
    n = x.shape[0]
    r = x @ w + b - y          # (n,)
    dw = (2.0 / n) * (x.T @ r)  # (d, n) @ (n,) -> (d,)
    db = (2.0 / n) * r.sum()
    return dw, float(db)

mse returns a Python float rather than a numpy scalar, like the functions of chapters 2 and 4. linreg_grad is the matrix form of the derivation line for line: the residual r, shape (n,), then $\frac{2}{n}\mathbf{X}^\top\mathbf{r}$ for the weights and $\frac{2}{n}\sum_i r_i$ for the bias. x @ w + b relies on the broadcasting of chapter 3 to add the scalar bias to every prediction. The lab test test_linreg_grad_matches_numerical checks this function against numerical_grad at a random point to a relative tolerance of $10^{-5}$, the check chapter 6 uses to verify its engine.

def fit_linreg(x, y, lr, steps):
    """Plain gradient descent from w = 0, b = 0. losses[t] is the loss before step t, plus the final loss."""
    w, b = np.zeros(x.shape[1]), 0.0
    losses = [mse(x @ w + b, y)]
    for _ in range(steps):
        dw, db = linreg_grad(w, b, x, y)
        w, b = w - lr * dw, b - lr * db
        losses.append(mse(x @ w + b, y))
    return w, b, losses

The tuple assignment w, b = w - lr * dw, b - lr * db evaluates both right-hand sides before either name is rebound, and linreg_grad computed both gradients at the old parameters, so the step is the simultaneous update of the mechanism. losses records the loss at the starting point and after each step, steps + 1 values in all, so losses[0] is the loss of the all-zero model and losses[-1] the loss of the returned one.

Exercises

(a) Differentiate the logistic sigmoid $\sigma(x) = 1/(1 + e^{-x})$ and show that $\sigma'(x) = \sigma(x)\,(1 - \sigma(x))$.

(b) Express the derivative of $\tanh x = \frac{e^{x} - e^{-x}}{e^{x} + e^{-x}}$ in terms of $\tanh x$ itself.

(c) The lab tests generate data as $x \sim N(0, 1)$, $n = 100$, seed 0, with $y = 3x - 2$ plus small noise. Estimate the largest stable learning rate for fit_linreg on this data from the quadratic analysis, and check it against test_fit_linreg_diverges_when_lr_too_large, which expects divergence at 2.0.

(d) For $f(x) = e^x$ at $x = 1$, compute the error of the forward and central differences for $h = 10^{-1}, 10^{-2}, \ldots, 10^{-12}$, plot both against $h$ on log–log axes, and explain both slopes and the floor.

import math

x = 1.0
exact = math.exp(x)
print("   eps     forward     central")
for k in range(1, 13):
    h = 10.0 ** -k
    fwd = (math.exp(x + h) - math.exp(x)) / h
    cen = (math.exp(x + h) - math.exp(x - h)) / (2 * h)
    print(f"1e-{k:02d}  {abs(fwd - exact):.2e}  {abs(cen - exact):.2e}")

   eps     forward     central
1e-01  1.41e-01  4.53e-03
1e-02  1.36e-02  4.53e-05
1e-03  1.36e-03  4.53e-07
1e-04  1.36e-04  4.53e-09
1e-05  1.36e-05  5.86e-11
1e-06  1.36e-06  1.63e-10
1e-07  1.40e-07  5.86e-11
1e-08  6.60e-09  6.60e-09
1e-09  2.15e-07  6.60e-09
1e-10  1.55e-06  6.73e-07
1e-11  3.26e-05  1.04e-05
1e-12  4.32e-04  2.10e-04

(e) Why does the gradient of the mean loss not grow with the size of the dataset? What changes if the loss is the sum of squared residuals instead of their mean?

Lab

Implement every function in py/labs/ch05/starter.py, then run

make lab CH=05 IMPL=mine

until it passes. Read py/labs/ch05/test_lab.py first. It pins what the prose leaves open: numerical_grad must return $2\mathbf{x}$ for $\sum_i x_i^2$ to within $10^{-6}$ and must leave its input unchanged, including a 2-D input; mse of [1, 2] against [1, 4] must be exactly 2.0; linreg_grad must agree with numerical_grad at a random point; and fit_linreg must recover $w = 3$ and $b = -2$ to within 0.01 in 500 steps at $\eta = 0.1$, produce a non-increasing loss at $\eta = 0.05$, and end with a higher loss than it started with at $\eta = 2.0$. Derive linreg_grad on paper before writing it; the numerical test will say whether the derivation is right.

Then sweep the learning rate. Save this script outside the repository, for example as /tmp/lr_sweep.py:

import numpy as np

from tinygpt.descent import fit_linreg

rng = np.random.default_rng(0)
x = rng.normal(size=(100, 1))
y = 3.0 * x[:, 0] - 2.0 + 0.01 * rng.normal(size=100)

xb = np.hstack([x, np.ones((100, 1))])  # (100, 2): the bias is a weight on a constant input
hessian = (2 / 100) * xb.T @ xb          # (2, 2)
curvatures = np.linalg.eigvalsh(hessian)
print("curvatures", curvatures.round(4), " stability bound 2/max =", round(2 / curvatures.max(), 4))

for lr in [0.01, 0.1, 0.5, 0.9, 1.1]:
    w, b, losses = fit_linreg(x, y, lr=lr, steps=300)
    hit = next((t for t, loss in enumerate(losses) if loss < 1e-3), None)
    if hit is not None:
        print(f"lr={lr:<4}  loss < 1e-3 after {hit:3d} steps   w={w[0]:.4f}  b={b:.4f}")
    else:
        low = int(np.argmin(losses))
        print(f"lr={lr:<4}  diverged: lowest loss {losses[low]:.3g} at step {low}, {losses[-1]:.3g} at step 300")

and run it from py/ with PYTHONPATH=. uv run python /tmp/lr_sweep.py. The data is the test data. Each run takes 300 steps; the script reports the first step at which the loss is below $10^{-3}$, with the parameters after all 300 steps, or, for a run that never gets there, its lowest loss and its final loss. Change the import to from labs.ch05.starter import fit_linreg to run yours; the output should be the same: