Training a neural network means nudging millions or billions of numbers so that a loss goes down. To know which way to nudge each number, you need the derivative of the loss with respect to it. Frameworks compute those derivatives for you, so why derive anything by hand?

Three reasons. Interviews at every lab ask you to ("derive the gradient of softmax cross-entropy" is close to universal). Every custom kernel, every fused operation, and every new layer you write at work needs a correct backward pass, and the only way to know it is correct is to understand what it should be. And the gradients themselves carry intuition. The gradient of cross-entropy with respect to the logits is "predicted minus true", and once you see that, you understand why classifiers train the way they do.

All of it rests on one idea, that a derivative is the best linear approximation to a function. This section builds from that idea to the matrix gradients of a linear layer and of softmax cross-entropy, and ends with the tool that checks any gradient you will ever write.

The derivative is the best linear approximation

Zoom in on a smooth curve far enough and it looks like a straight line. The derivative is the slope of that line. Precisely, $f$ is differentiable at $x$ if there is a number $f'(x)$ such that

where $o(h)$ means an error that shrinks faster than $h$ itself (its ratio to $h$ goes to 0). In plain words, near $x$, the change in $f$ is the slope times the change in input, up to an error that becomes negligible for small steps.

Take $f(x) = x^2$ at $x = 3$. Expanding, $(3 + h)^2 = 9 + 6h + h^2$. The linear part is $6h$, so $f'(3) = 6$, and the error is exactly $h^2$, which is tiny compared with $h$ when $h$ is small. The cell checks that the error divided by $h^2$ stays at 1.

That framing, "the linear map that best approximates the function nearby", is the one that generalizes. When the input is a vector or a matrix, the derivative is no longer a number. It is a linear map, and the whole of matrix calculus is about writing that map down and composing such maps.

Gradients, for functions that return one number

Let $f: \mathbb{R}^n \to \mathbb{R}$, like a loss as a function of a parameter vector. Its derivative at $\mathbf{x}$ is a linear map from a small step $\mathbf{h} \in \mathbb{R}^n$ to a change in $f$, and every linear map from vectors to numbers is a dot product with some fixed vector. That vector is the gradient,

In plain words, the gradient lists how much $f$ changes per unit step along each coordinate, and the change for any small step is the dot product of the gradient with that step.

The gradient points uphill, as steeply as possible. Here is why. A step of length $\epsilon$ in a unit direction $\mathbf{u}$ changes $f$ by about $\epsilon\, \nabla f^\top \mathbf{u}$. By the Cauchy–Schwarz inequality, $\nabla f^\top \mathbf{u} \le \lVert \nabla f \rVert \lVert \mathbf{u} \rVert = \lVert \nabla f \rVert$, with equality exactly when $\mathbf{u} = \nabla f / \lVert \nabla f \rVert$. So among all directions, the gradient direction increases $f$ fastest, and $-\nabla f$ decreases it fastest. That is the entire justification for gradient descent, which §0.6 develops.

The same idea works when the input is a matrix. If $\mathcal{L}$ is a scalar loss and $\mathbf{W} \in \mathbb{R}^{m \times n}$, then $\partial \mathcal{L} / \partial \mathbf{W}$ is the $m \times n$ matrix of partial derivatives $\partial \mathcal{L} / \partial W_{ij}$, and the first-order change is a sum over all entries,

where $\langle \mathbf{A}, \mathbf{B} \rangle = \sum_{ij} A_{ij} B_{ij} = \operatorname{tr}(\mathbf{A}^\top \mathbf{B})$ is the matrix inner product. We will use one convention throughout the book. The gradient of a scalar with respect to any tensor has the same shape as that tensor. That is what every framework does, and it is what makes the update $\mathbf{W} \leftarrow \mathbf{W} - \eta\, \partial\mathcal{L}/\partial\mathbf{W}$ well-typed, with $\eta$ the learning rate.

Jacobians, for functions that return vectors

Now let $f: \mathbb{R}^n \to \mathbb{R}^m$, like a layer that maps an input vector to an output vector. Its derivative is a linear map from $\mathbb{R}^n$ to $\mathbb{R}^m$, which is a matrix, the Jacobian $\mathbf{J}_f(\mathbf{x}) \in \mathbb{R}^{m \times n}$, with entries

In plain words, row $i$ of the Jacobian is the gradient of output $i$, and multiplying the Jacobian by a small input step predicts the output step.

Two Jacobians come up constantly.

• A linear map $f(\mathbf{x}) = \mathbf{A}\mathbf{x}$ has $\mathbf{J}_f = \mathbf{A}$ everywhere. A linear function is its own best linear approximation.
• An elementwise function $f(\mathbf{x}) = \phi(\mathbf{x})$, applying a scalar function $\phi$ such as tanh or ReLU to each entry, has a diagonal Jacobian $\operatorname{diag}(\phi'(x_1), \dots, \phi'(x_n))$, because output $i$ depends only on input $i$.

The chain rule is matrix multiplication

Compose two functions, $\mathbf{y} = f(\mathbf{x})$ with $f: \mathbb{R}^n \to \mathbb{R}^m$, and $\mathbf{z} = g(\mathbf{y})$ with $g: \mathbb{R}^m \to \mathbb{R}^k$. Near $\mathbf{x}$, a step $\mathbf{h}$ moves $\mathbf{y}$ by about $\mathbf{J}_f \mathbf{h}$, and that moves $\mathbf{z}$ by about $\mathbf{J}_g (\mathbf{J}_f \mathbf{h})$. The composition of two linear maps is the product of their matrices, so

In plain words, the derivative of a composition is the product of the derivatives, evaluated at the right points. The shapes chain like dominoes.

A neural network is a long composition, $\mathcal{L} = \ell(f_L(\cdots f_2(f_1(\mathbf{x}))))$, ending in a scalar loss. Its gradient with respect to the input is a product of Jacobians, and since the final one (from $\ell$) is a $1 \times m$ row, we can transpose the whole product into column form,

where $\mathbf{y}_L$ is the network's output. Evaluate it from the right. Start with the gradient of the loss, a vector. Multiply by $\mathbf{J}_{f_L}^\top$ to get another vector, then by $\mathbf{J}_{f_{L-1}}^\top$, and so on. At every step you multiply a transposed Jacobian by a vector, which is called a vector-Jacobian product (VJP). You never form a full Jacobian, and you never multiply two Jacobians together. That right-to-left sweep is backpropagation, and §0.3 builds an engine that does it automatically.

For the elementwise case the VJP is especially simple. With $\mathbf{y} = \phi(\mathbf{x})$ and upstream gradient $\mathbf{g} = \nabla_{\mathbf{y}} \mathcal{L}$, the diagonal Jacobian gives $\nabla_{\mathbf{x}}\mathcal{L} = \mathbf{g} \odot \phi'(\mathbf{x})$, an elementwise product.

import numpy as np

def grad_linear(a, x):
    raise NotImplementedError

def grad_quadratic(A, x):
    raise NotImplementedError

The gradient of a linear layer

This is the most important gradient in deep learning, because nearly all the parameters and nearly all the compute of a Transformer are in linear layers. A batch of $N$ inputs, stacked as rows of $\mathbf{X} \in \mathbb{R}^{N \times d_{\text{in}}}$, goes through

where the bias row $\mathbf{b}$ is broadcast to every row. Some loss $\mathcal{L}$ is computed from $\mathbf{Y}$, and backpropagation hands us the upstream gradient $\mathbf{G} = \partial \mathcal{L} / \partial \mathbf{Y} \in \mathbb{R}^{N \times d_{\text{out}}}$. We want the three gradients

In plain words, the weight gradient correlates inputs with output errors, the input gradient sends the errors back through the transposed weights, and the bias gradient adds up the errors over the batch.

Let's derive the weight gradient in index notation, the most reliable method there is. Entry $(n, j)$ of the output is $Y_{nj} = \sum_i X_{ni} W_{ij} + b_j$. The weight $W_{ij}$ affects every row $n$ of column $j$, with $\partial Y_{nj} / \partial W_{ij} = X_{ni}$, and nothing else. Summing over every place it appears,

The sum over $n$ is the multivariate chain rule. A variable that influences several outputs collects the gradient from all of them. The input gradient works the same way. $X_{ni}$ affects $Y_{nj}$ for every $j$, with $\partial Y_{nj}/\partial X_{ni} = W_{ij}$, so $\partial \mathcal{L} / \partial X_{ni} = \sum_j G_{nj} W_{ij} = (\mathbf{G}\mathbf{W}^\top)_{ni}$. And $b_j$ adds into $Y_{nj}$ for every $n$ with coefficient 1, so its gradient is $\sum_n G_{nj}$.

Notice what the weight gradient costs. $\mathbf{X}^\top\mathbf{G}$ is a matmul of the same size as the forward $\mathbf{X}\mathbf{W}$, and so is $\mathbf{G}\mathbf{W}^\top$. So the backward pass of a linear layer costs two forward passes. That is where the rule of thumb "training costs about 6 FLOPs per parameter per token" comes from (2 for the forward multiply-add, 4 for the two backward matmuls), which you will use to budget compute in the scaling-law chapter. It holds up at real scale. GPT-3 has $1.75 \times 10^{11}$ parameters and was trained on $3 \times 10^{11}$ tokens, so the rule predicts $6 \times 1.75 \times 10^{11} \times 3 \times 10^{11} \approx 3.15 \times 10^{23}$ floating-point operations, and the paper reports $3.14 \times 10^{23}$ (Brown et al., 2020).

Let's verify all three formulas numerically, using the loss $\mathcal{L} = \sum_{nj} G_{nj} Y_{nj}$, whose gradient with respect to $\mathbf{Y}$ is exactly $\mathbf{G}$ by construction.

import numpy as np

def linear_backward(X, W, b, dY):
    """X: (..., d_in), W: (d_in, d_out), b: (d_out,), dY: (..., d_out) -> dX, dW, db"""
    raise NotImplementedError

The softmax Jacobian

Softmax turns a vector of scores (logits) $\mathbf{z} \in \mathbb{R}^K$ into a probability distribution over $K$ classes,

Every output depends on every input, because they share the denominator, so its Jacobian is full. Differentiate $p_i$ with respect to $z_j$ with the quotient rule, writing $Z = \sum_k e^{z_k}$ and noting $\partial Z/\partial z_j = e^{z_j}$,

where $\delta_{ij}$ is 1 if $i = j$ and 0 otherwise. In plain words, raising your own logit raises your own probability in proportion to $p_i(1 - p_i)$, and raising someone else's logit lowers yours in proportion to $p_i p_j$.

Three facts about this matrix come up in interviews. It is symmetric. Its rows sum to zero, since $\sum_j (p_i\delta_{ij} - p_i p_j) = p_i - p_i \sum_j p_j = 0$, so $\mathbf{J}\mathbf{1} = \mathbf{0}$. That says adding the same constant to every logit changes nothing, the shift invariance we exploit for numerical stability later in this section, and it means the Jacobian has rank at most $K - 1$. And because $\mathbf{J}$ is symmetric with $\mathbf{u}^\top \mathbf{J} \mathbf{u} = \sum_i p_i u_i^2 - (\sum_i p_i u_i)^2$, which is the variance of $\mathbf{u}$ under $\mathbf{p}$, it is positive semidefinite. (Rank, null spaces, and positive semidefinite matrices are developed properly in §0.10. Here you only need the definitions just used.)

The VJP never needs the $K \times K$ matrix. For an upstream gradient $\mathbf{g} = \nabla_{\mathbf{p}}\mathcal{L}$,

In plain words, subtract the probability-weighted average of the upstream gradient, then scale by the probabilities. This costs $O(K)$ instead of $O(K^2)$, and in attention, where softmax runs over every row of a $T \times T$ score matrix, that difference is everything. FlashAttention's backward pass is built on exactly this identity (§12.3).

import numpy as np

def softmax_jacobian(p):
    raise NotImplementedError

def softmax_backward(P, G):
    raise NotImplementedError

Softmax plus cross-entropy gives p minus y

Now put a loss on top. For a classifier with true class $y$, the cross-entropy loss is the negative log-probability of the true class (§0.5 explains why this is the loss). Write the target as a one-hot vector $\mathbf{y} = \mathbf{e}_y$, so

We could chain the two VJPs. The gradient of $\mathcal{L}$ with respect to $\mathbf{p}$ is $g_k = -y_k / p_k$, and then $\mathbf{p} \odot (\mathbf{g} - (\mathbf{p}^\top\mathbf{g})\mathbf{1})$ gives, entry by entry, $p_j(-y_j/p_j) - p_j \sum_k p_k(-y_k/p_k) = -y_j + p_j \sum_k y_k = p_j - y_j$. Everything cancels into one of the most important results in machine learning,

In plain words, the gradient on each logit is the predicted probability minus the true probability. It pushes the true class's logit up by $1 - p_y$ and every other logit down by its own probability.

There is an even shorter route, which also shows how to compute the loss stably. Substitute the softmax into the loss,

The derivative of $-z_y$ with respect to $z_j$ is $-\delta_{jy}$, and the derivative of the log-sum-exp is $e^{z_j} / \sum_k e^{z_k} = p_j$. So $\partial\mathcal{L}/\partial z_j = p_j - \delta_{jy}$, the same answer in one line. You will meet this result again and again. Word2vec's softmax gradient is the same "expected minus observed" (§1.5), and so is the output layer of every language model.

Several things make this gradient special.

• It never saturates. Each entry lies in $[-1, 1]$, and it is large exactly when the prediction is badly wrong. Compare a sigmoid output trained with the half squared error $\tfrac12(\sigma(z) - y)^2$, whose gradient $(\sigma(z) - y)\,\sigma'(z)$ carries an extra factor $\sigma'(z) = \sigma(z)(1 - \sigma(z))$. At $z = -5$ with target 1, $\sigma(z) = 0.0067$, and the cross-entropy gradient is $-0.993$ while the squared-error gradient is $-0.0066$, about 150 times smaller. The model is confidently wrong and squared error barely notices.
• It works for any target distribution. The derivation never used that $\mathbf{y}$ is one-hot, only that its entries sum to 1. So the same $\mathbf{p} - \mathbf{y}$ holds for the soft targets of label smoothing (§0.5) and knowledge distillation.
• It averages over a batch. With a mean over $N$ examples, $\mathcal{L} = \frac1N \sum_n \mathcal{L}_n$, the gradient on the logit matrix is $(\mathbf{P} - \mathbf{Y}) / N$. Forgetting the $1/N$ is a classic bug that silently multiplies your learning rate by the batch size.

Computing it without overflow

The formula $\mathcal{L} = -z_y + \log\sum_k e^{z_k}$ is right, but computed literally it fails on ordinary logits. In float32, $e^{z}$ overflows to infinity once $z$ passes about 88.7, so a logit of 100 makes the sum infinite, the loss infinite, and the softmax $\infty/\infty$, which is NaN. Very negative logits fail the other way, since every $e^{z_k}$ underflows to 0 and the log returns $-\infty$.

The fix is the shift invariance of the softmax, made exact. Let $m = \max_k z_k$ and factor $e^m$ out of the sum,

In plain words, subtract the largest logit before exponentiating and add it back after the log. Every exponent is now at most 0, so nothing overflows, and the largest term is exactly $e^0 = 1$, so the sum is at least 1 and its log is safe. The log-probabilities follow as $\log p_k = z_k - (m + \log\sum_j e^{z_j - m})$, and they never pass through a probability that could round to 0.

This quantity, $\log\sum_k e^{z_k}$, is called the log-sum-exp, and §0.4 studies it further (it is a smooth maximum, it handles masked rows that are entirely $-\infty$, and it carries over to probabilities of whole sequences). For now the shift is all you need.

This is why every framework offers a fused "cross-entropy from logits" function and why you should always use it. It computes the loss through the stable log-sum-exp (never through $\log$ of a softmax that may have underflowed to 0), and its backward pass is the clean $\mathbf{p} - \mathbf{y}$.

import numpy as np

def softmax_cross_entropy(logits, y):
    """logits: (N, K), y: (N,) ints -> (mean loss, dlogits (N, K))"""
    raise NotImplementedError

Checking any gradient with finite differences

You now have a handful of hand-derived gradients, and you will write many more. How do you know they are right? Compare them with a numerical estimate built straight from the definition of the derivative. This is called gradient checking, and it is the single most valuable debugging tool for anyone writing backward passes.

The obvious estimate is the forward difference, $(f(x + h) - f(x))/h$. Taylor's theorem tells us how accurate it is,

So its error shrinks only in proportion to $h$. The central difference does much better. Subtract the expansion of $f(x - h)$, in which the odd powers of $h$ flip sign. The $f$ and $f''$ terms cancel, and

In plain words, the symmetric difference cancels the curvature term, so its error shrinks like $h^2$. Halve the step and the error drops fourfold.

That suggests using a tiny $h$, but floating point fights back. Each evaluation of $f$ carries a relative rounding error of about $\varepsilon_{\text{mach}}$, the machine epsilon ($2.2 \times 10^{-16}$ in float64, $1.2 \times 10^{-7}$ in float32). The numerator subtracts two nearly equal numbers, so its absolute error is about $\varepsilon_{\text{mach}} |f|$, and dividing by $h$ amplifies it to $\varepsilon_{\text{mach}} |f| / h$. For the central difference the total error is roughly

which is smallest near $h^\star \approx \varepsilon_{\text{mach}}^{1/3}$, about $6 \times 10^{-6}$ in float64, giving errors around $10^{-11}$. For the forward difference the same argument gives $h^\star \approx \varepsilon_{\text{mach}}^{1/2} \approx 1.5 \times 10^{-8}$ and errors around $10^{-8}$. Go smaller than $h^\star$ and the error grows. The widget shows the characteristic V shape on a log-log plot. Switch to float32 to see why gradient checks must run in float64.

For $f = \sin$ at $x = 1$, the true derivative is $\cos 1$. With $h = 10^{-5}$, the forward difference is off by $4.2 \times 10^{-6}$ and the central difference by about $10^{-11}$. With $h = 10^{-12}$, both are off by more than $10^{-5}$, because rounding has taken over. In float32 the central difference is typically off by $10^{-6}$ to $10^{-5}$ even at its best steps, around $h \approx 10^{-3}$ to $10^{-2}$, and the error jumps around from one $h$ to the next. That is too crude to catch subtle bugs.

How to read the result

Comparing an analytic gradient $\mathbf{a}$ with a numerical one $\mathbf{n}$ needs a relative error, because gradients can be $10^{-6}$ or $10^{6}$ in scale. A robust choice is

with a tiny floor $\tau$ (say $10^{-12}$) so two zero gradients do not give $0/0$. In float64 with central differences, a correct gradient typically gives $10^{-7}$ or less. Around $10^{-4}$ means a subtle bug or a kink nearby, and $10^{-2}$ or worse means the gradient is wrong. (The CS231n notes linked below give similar thresholds.)

import numpy as np

def numerical_grad(f, x, h=1e-5):
    raise NotImplementedError

def rel_error(a, b):
    raise NotImplementedError

In PyTorch

PyTorch can compute full Jacobians when you need them (rarely, for small functions), and it ships a gradient checker that does exactly what you just built, in float64.

Read more

• Parr & Howard (2018). The Matrix Calculus You Need For Deep Learning. A gentle, complete tour of Jacobians and the vector chain rule, written for practitioners.
• Justin Johnson. Derivatives, Backpropagation, and Vectorization. A short CS231n handout on gradients of tensors with respect to tensors, and why VJPs suffice.
• Petersen & Pedersen. The Matrix Cookbook. The standard reference table of matrix derivatives, for checking your own.
• CS231n course notes. Neural Networks Part 3: Learning and Evaluation. Its gradient-checking section gives practical advice on relative error, kinks, and step sizes.
• Goodfellow, Bengio & Courville (2016). Deep Learning. Section 6.5 states back-propagation and the chain rule carefully.
• Brown et al. (2020). Language Models are Few-Shot Learners. GPT-3, whose reported training compute matches the 6 FLOPs per parameter per token rule.
• Nocedal & Wright (2006). Numerical Optimization. Chapter 8 is the rigorous treatment of finite-difference error and optimal step sizes.

Watch next

▶

▶

▶