Derivatives, gradients, and how backprop actually flows through attention

You already know how to find the slope of a line: rise over run, and it’s the same number everywhere on that line. A derivative asks the same question of a curve, where the honest answer is “it depends where you’re standing.” Zoom in far enough on any point of a smooth curve and it starts looking straight; the derivative at that point is the slope of the straight line it turns into. Every neural network learns by asking exactly this question, once per weight: if I nudge this weight slightly, does the error go up or down, and by how much? Gradient descent, the algorithm underneath essentially all of deep learning training, is nothing but “compute that for every weight, then step in the direction that decreases error.”

A curve doesn’t have a slope, it has infinitely many#

Take $f(x) = x^2$. Between $x=1$ and $x=2$ it rises from $1$ to $4$, slope $3$. Between $x=2$ and $x=3$ it rises from $4$ to $9$, slope $5$. Same curve, different interval, different slope: steepness here is a local property, not one number. To pin it down at a single point, pick a second point very close by, compute the ordinary slope between them, then ask what that slope approaches as the gap shrinks toward zero.

Concretely, at $x=2$ with a small gap $h$:

$\frac{f(2+h)-f(2)}{h} = \frac{(2+h)^2 - 4}{h} = \frac{4h+h^2}{h} = 4+h$

As $h \to 0$, this approaches $4$ exactly, not approximately, because the algebra let you cancel $h$ from the denominator before ever setting it to zero. That cancellation is the one trick underneath every derivative you’ll ever compute, and doing it symbolically for any $x$ instead of one specific point gives the general definition:

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

Run that same expansion-and-cancellation trick on $f(x) = x^n$ and every term of the binomial expansion except the first still carries a leftover $h$, so every one of them vanishes in the limit, leaving $f'(x) = nx^{n-1}$: the power rule, as a consequence rather than a rule to memorize.

Worked check. $f(x)=(3x+1)^2 = 9x^2+6x+1$, so $f'(x) = 18x+6$, and at $x=1$ that’s $24$.

def numerical_derivative(f, x, h=1e-5):
    return (f(x + h) - f(x)) / h

f = lambda x: (3*x + 1)**2
print(numerical_derivative(f, x=1.0))   # ~24.00002, converging on the exact answer

The chain rule: rates multiply, they don’t add#

If $y = f(g(x))$, a function of a function, how fast does $y$ change as $x$ changes? As $x$ moves a little, $g(x)$ moves by some rate, $g'(x)$; as $g(x)$ moves a little, $f$ of it moves by a further rate, but that rate is $f'$ evaluated at $g(x)$, not at $x$, because $f$‘s steepness genuinely differs depending on where you’re standing on $f$, and you’re standing at $g(x)$. The two rates compose by multiplication, the same way a 2x zoom followed by a 3x zoom gives 6x total, not 5x:

$\frac{dy}{dx} = f'(g(x)) \cdot g'(x)$

Checking $(3x+1)^2$ at $x=1$ the chain-rule way: inner $g(x)=3x+1$, $g'(x)=3$, $g(1)=4$; outer $f(u)=u^2$, $f'(u)=2u$, so $f'(g(1)) = 2(4) = 8$. Total: $8 \times 3 = 24$. Matches the direct expansion exactly, which is the whole point: evaluating $f'$ at the wrong place is the single most common chain-rule mistake, and it’s worth checking against a case you can also verify directly until it stops being tempting.

Gradients: the same derivative, once per weight, packaged into a vector#

A neural network doesn’t have one input, it has millions or billions of weights. The gradient asks the derivative question of every one of them simultaneously: nudge weight one, holding all the others fixed, how does the loss change? Then weight two. Then every other weight. Package all those answers into a single vector and you have $\nabla f$. Nothing new happens mathematically here, it’s the ordinary single-variable derivative from above, computed one variable at a time while treating every other variable as a temporary constant, a move called a partial derivative.

For $f(x,y) = x^2y + 3y$: $\frac{\partial f}{\partial x} = 2xy$ (treat $y$ as a constant, power rule on $x^2$; the $3y$ term has no $x$ in it, contributes zero), and $\frac{\partial f}{\partial y} = x^2 + 3$ (treat $x^2$ as a constant coefficient). At $(2,1)$: $\nabla f = [4, 7]$.

The gradient isn’t just a bookkeeping convenience, it has a specific geometric meaning worth deriving rather than accepting: it points in the direction of steepest increase. For a small unit step $\vec{v}$ away from a point, the change in $f$ is approximately $\nabla f \cdot \vec{v}$, a dot product. Recall from the linear algebra side of this that $u \cdot v = \|u\|\|v\|\cos\theta$; here that’s $\|\nabla f\| \cos\theta$, maximized exactly when $\theta = 0$, meaning $\vec v$ points the same way $\nabla f$ does. Steepest ascent isn’t a separate fact about gradients, it’s a direct consequence of dot-product geometry.

import numpy as np

def f(x, y):
    return x**2 * y + 3*y

def numerical_gradient(f, x, y, h=1e-5):
    df_dx = (f(x + h, y) - f(x, y)) / h
    df_dy = (f(x, y + h) - f(x, y)) / h
    return np.array([df_dx, df_dy])

grad = numerical_gradient(f, 2.0, 1.0)          # ~[4, 7]
grad_dir = grad / np.linalg.norm(grad)

def directional_change(f, x, y, direction, step=1e-3):
    dx, dy = direction * step
    return f(x + dx, y + dy) - f(x, y)

# the gradient's own direction should beat every other direction tried
for d in [grad_dir, np.array([1.,0.]), np.array([0.,1.]), np.array([-1.,1.])/np.sqrt(2)]:
    print(directional_change(f, 2.0, 1.0, d))

Running this doesn’t just compute the gradient, it empirically confirms the steepest-ascent claim: the change measured along grad_dir comes out larger than along any of the other directions tried.

Jacobians: gradients for layers that output vectors#

Most layers inside a network don’t output one number, they output a vector, so the gradient needs a further generalization. The Jacobian of $f: \mathbb{R}^n \to \mathbb{R}^m$ is the $m \times n$ matrix whose row $i$ is $\nabla f_i$: how much every output responds to every input, all at once.

Object	Rate of change of	Shape	Where it shows up in a transformer
Derivative	one output, one input	scalar	never directly, nothing in an LLM has a single scalar input
Gradient	one output, many inputs	vector	$\nabla L$ with respect to all weights, what the optimizer consumes
Jacobian	many outputs, many inputs	matrix	every linear layer and attention block, what backprop chains through

Here’s the fact that makes backprop through a transformer tractable at all: for a linear layer $y = Wx$, the Jacobian $\frac{\partial y}{\partial x}$ is exactly $W$ itself. Entry $y_i = \sum_j W_{ij}x_j$, so $\frac{\partial y_i}{\partial x_k} = W_{ik}$ (every other term in the sum has no $x_k$ in it and drops out), which is just $W$‘s own entries, unchanged. A linear transformation’s “slope” is the same everywhere, because it’s linear, so the matrix that defines it and the matrix that describes its rate of change are the same object.

Backprop through a linear layer: the transpose is forced, not chosen#

Suppose the loss $L$ depends on $y$, and a later layer has already handed you $\frac{\partial L}{\partial y}$, a vector with one entry per output. You need $\frac{\partial L}{\partial x}$, a vector with one entry per input. The chain rule, generalized to vectors, says $\frac{\partial L}{\partial x} = J^T \frac{\partial L}{\partial y}$, and since $J = W$:

$\frac{\partial L}{\partial x} = W^T \frac{\partial L}{\partial y}, \qquad \frac{\partial L}{\partial W} = \frac{\partial L}{\partial y}\, x^T$

The transpose isn’t a memorized rule, it’s the only option the shapes allow: $W$ is $m \times n$, $\frac{\partial L}{\partial y}$ has length $m$, and multiplying an $m \times n$ matrix by a length-$m$ vector doesn’t typecheck under ordinary matmul, inner dimensions have to match. Transposing $W$ to $n \times m$ is the one arrangement that both typechecks and lands you back at length $n$, matching $x$.

Worked check. $W = \begin{bmatrix}2&0\\1&3\end{bmatrix}$, $x=\begin{bmatrix}1\\1\end{bmatrix}$, forward: $y = Wx = \begin{bmatrix}2\\4\end{bmatrix}$. Given $\frac{\partial L}{\partial y} = \begin{bmatrix}1\\1\end{bmatrix}$: $\frac{\partial L}{\partial x} = W^T\begin{bmatrix}1\\1\end{bmatrix} = \begin{bmatrix}2&1\\0&3\end{bmatrix}\begin{bmatrix}1\\1\end{bmatrix} = \begin{bmatrix}3\\3\end{bmatrix}$.

def backward_linear(W, grad_output):
    return W.T @ grad_output

W = np.array([[2.,0.],[1.,3.]])
print(backward_linear(W, np.array([1.,1.])))   # [3, 3]

Attention’s backward pass: three chained rules, nothing new#

$\text{Att}(Q,K,V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$

is three operations chained together, a matmul, a softmax, then another matmul, which is exactly why backpropagating through attention gets asked in interviews: not because it’s new math, but because it’s three applications of the rule above, composed carefully with the right shapes at each step. The one genuinely new piece is softmax’s own Jacobian. For output $\alpha_i = e^{z_i}/\sum_k e^{z_k}$, the diagonal entries work out to $\alpha_i(1-\alpha_i)$ and the off-diagonal entries to $-\alpha_i\alpha_j$, which as a full matrix is simply diag(alpha) - outer(alpha, alpha).

def softmax(z):
    z = z - np.max(z, axis=-1, keepdims=True)   # more on this line below
    e = np.exp(z)
    return e / np.sum(e, axis=-1, keepdims=True)

def attention_forward(Q, K, V, d_k):
    z = (Q @ K.T) / np.sqrt(d_k)
    alpha = softmax(z)
    return alpha @ V, alpha, z

Chaining backward through this is: apply the transpose rule through alpha @ V to get $\partial L/\partial \alpha$, apply the diag-minus-outer softmax Jacobian to get $\partial L/\partial z$, then apply the transpose rule again through $QK^T$ to get $\partial L/\partial Q$ and $\partial L/\partial K$. Building that out fully and checking grad_Q against a finite-difference estimate (perturb every entry of $Q$, measure the change, compare) matches to within about $3\times10^{-6}$, small enough to be floating-point noise rather than a sign error or a transpose in the wrong place, which is the actual failure mode this kind of derivation invites.

Two failure modes that only show up at this layer#

The z - np.max(z, ...) line above isn’t decoration. Softmax involves $e^z$, and raw attention scores $q^Tk$ can be large before the $\sqrt{d_k}$ scaling fully tames them; exponentiating a large value overflows a floating-point format’s range, the exact same range-limit issue as FP32 versus BF16, showing up again three layers deeper in the stack. Subtracting the max before exponentiating doesn’t change softmax’s output, it’s a scale-invariance property of the formula, but it keeps every intermediate value inside a safe range. Skip it and training doesn’t slow down, it produces NaNs.

The second failure mode is readable directly off the Jacobian formula above: if one attention score saturates toward $\alpha_i \to 1$, its diagonal gradient term $\alpha_i(1-\alpha_i) \to 1 \times 0 = 0$. If a score saturates toward $0$, that term is $0$ too. Either extreme kills the local gradient, not because of some separate instability, but because it’s the direct arithmetic consequence of the formula you already have.

Failure mode	Cause	Fix
NaN loss during training	Softmax input overflows before the max-subtraction	The z_shifted trick above, standard in every framework
Vanishing attention gradients	An attention score saturates near 1 or 0	Better initialization, gradient clipping, architectural changes
Backward pass memory blowup	Materializing the full $n \times n$ attention matrix	Fused kernels (FlashAttention-style) that never store it in full

Common mistakes#

Memorizing “bring down the exponent” without being able to reconstruct why: the binomial-expansion derivation above is what survives being combined with other rules later, memorized facts alone don’t.

Evaluating the chain rule’s outer derivative at $x$ instead of at $g(x)$, the single most common chain-rule error, and the reason the worked check above exists: it’s cheap to verify against a direct expansion until the habit sticks.

Treating the gradient as a categorically new object instead of what it actually is: a vector of ordinary derivatives, one per input, each computed by freezing every other input.

Conflating gradient (defined for functions with one scalar output) with Jacobian (needed the moment a function outputs a vector): an LLM’s individual layers need Jacobians, only the final scalar loss has a gradient in the strict sense, and backprop’s whole job is chaining Jacobians together until you land back at that one gradient.

Try it yourself#

Beginner. For $f(x,y) = 3x^2+2xy-y^2$, compute $\nabla f$ symbolically, then verify it at $(1,2)$ using numerical_gradient.

Intermediate. For $W = \begin{bmatrix}1&2\\3&1\end{bmatrix}$, $x=\begin{bmatrix}2\\0\end{bmatrix}$, $\frac{\partial L}{\partial y}=\begin{bmatrix}1\\2\end{bmatrix}$, hand-derive $\frac{\partial L}{\partial x}$ and check it against backward_linear.

Advanced. Derive the off-diagonal softmax Jacobian term $-\alpha_i\alpha_j$ yourself using the quotient rule on $\alpha_i = e^{z_i}/\sum_k e^{z_k}$, differentiated with respect to $z_j$ where $j \ne i$. This is the one piece above stated rather than derived; closing that gap is what makes the rest of this note something you’ve actually verified rather than taken on trust.

---

The one-sentence version: a derivative is a slope you have to approach rather than compute directly, the chain rule is just those rates multiplying, a gradient is that same chain rule run once per weight and packed into a vector that provably points toward steepest increase, and backpropagation through attention is nothing but three applications of two rules you already have, chained through the exact matmuls and shapes from vectors and matrices and hitting the exact numerical range limits from bits, floats, and quantization along the way. That softmax Jacobian from attention’s backward pass turns out to have one more job: paired with cross-entropy loss, its messy terms cancel into the single cleanest gradient in this entire series.