Derivatives and gradients

Problem

A model's loss is a function of millions of parameters. We want the parameters that make it small, but we can't try every combination, and we can't solve for the minimum with algebra. What we can do is stand at the current parameters and ask: if I change each one a little, how does the loss change? Answer that for every parameter at once and you know which direction lowers the loss fastest. Take a small step that way and repeat.

That answer is the gradient, and "step against it and repeat" is gradient descent. Every optimizer on this site, from plain SGD to AdamW, is a variation on that loop.

Intuition

For a function of one number, the derivative is the slope of its graph: how much the output rises per unit the input moves, measured at a single point. Zoom far enough into any smooth curve and it looks like a straight line, and the derivative is that line's slope.

For a function of many numbers, picture a landscape whose height is the loss. At any spot you can measure the slope in the east–west direction and the north–south direction separately; those are partial derivatives. Put them together as a vector and it points uphill, the steepest way up. Its length says how steep. To go down, walk the opposite way.

The derivative is the limit of "rise over run" as the run shrinks to zero. For x² the secant's slope is exactly 6 + h, so the error shrinks in step with h.

Mechanics

The derivative

f′(x) = lim_h→0 (f(x + h) − f(x)) / h

A handful of rules cover almost everything: the derivative of xⁿ is n·xⁿ⁻¹, of eˣ is eˣ, of ln x is 1/x, constants factor out, and the derivative of a sum is the sum of the derivatives. So for f(x) = x², f′(x) = 2x and f′(3) = 6, matching the tangent above. Functions built from other functions need the chain rule, the next page.

The sign says which way is downhill: if f′(x) > 0, decreasing x decreases f. The size says how sensitive f is. Where f′(x) = 0 the curve is flat: a minimum, a maximum or a saddle.

Partial derivatives

For a function of several inputs, a partial derivative ∂f/∂xᵢ varies one input and holds the others fixed. Take the bowl used for the rest of this page,

f(x, y) = x² + 3y²,   ∂f/∂x = 2x,   ∂f/∂y = 6y

At (3, 2): ∂f/∂x = 6 and ∂f/∂y = 12. The bowl is three times steeper in y, so the same move in y changes the loss more.

The gradient

∇f = (∂f/∂x₁, …, ∂f/∂xₙ),   ∇f(3, 2) = (6, 12)

The gradient has one entry per input, so it has the same shape as the parameters. Moving a small step d changes f by about ∇f · d, a dot product. Among all steps of the same length, that is largest when d points along ∇f (cosine 1) and most negative when d points against it. So the gradient is the direction of steepest ascent, −∇f the direction of steepest descent, and any step at right angles to ∇f leaves f unchanged to first order: the gradient is perpendicular to the contour lines.

Gradient descent

θ ← θ − η ∇f(θ)

η is the learning rate. On the bowl with η = 0.1, x moves by −0.1·2x, so it keeps 0.8 of its value each step, and y keeps 0.4. The steep direction converges faster, but it also sets the speed limit: once η exceeds 0.333, y's factor 1 − 6η falls below −1 and every step overshoots further than the last.

The path drops quickly in the steep y direction, then spends most of its steps crawling along the shallow x direction. Raise η to speed up x and y starts to overshoot. That unequal curvature is what momentum and Adam, on the optimizers page, are built to fix.

The Jacobian

A function with m outputs and n inputs has one gradient per output. Stacked as rows they form the m × n Jacobian, J_ij = ∂fᵢ/∂xⱼ, the best linear approximation of the function near a point: f(x + d) ≈ f(x) + J d. For f(x, y) = (xy, x + y²),

J = [[y, x], [1, 2y]],   J(2, 3) = [[3, 2], [1, 6]]

A loss is a scalar, so its Jacobian is a single row, the gradient. Inside the network the Jacobians of layers are huge, a [4096, 4096] block for one linear layer on one token, and autograd never builds them. It only multiplies a vector by them, which the chain rule page explains.

How PyTorch exposes gradients

Mark a tensor with requires_grad=True and PyTorch records every operation on it. Calling .backward() on a scalar result fills the .grad of each leaf tensor with ∂result/∂leaf (Tensor.backward, which calls torch.autograd.backward). Three details catch everyone:

• .grad accumulates. A second backward adds to it, which is why training loops call zero_grad (Tensor.grad).
• Only leaf tensors, the ones you created, get .grad; intermediates need retain_grad() (is_leaf).
• backward on a non-scalar needs a gradient argument, the vector in a vector-Jacobian product. torch.autograd.grad returns gradients instead of storing them (autograd.grad), and torch.no_grad() turns recording off for inference (no_grad).

Worked example

Implementation

import torch

theta = torch.tensor([3.0, 2.0], requires_grad=True)
scale = torch.tensor([1.0, 3.0])
lr = 0.1

for step in range(6):
    loss = (scale * theta ** 2).sum()   # x² + 3y²
    loss.backward()                     # theta.grad = (2x, 6y)
    with torch.no_grad():               # the update itself isn't recorded
        theta -= lr * theta.grad
    theta.grad = None                   # .grad accumulates; clear it
    print(step, loss.item(), theta.tolist())
# 0 21.0 [2.4, 0.8]   1 7.68 [1.92, 0.32]   ...

from torch.autograd.functional import jacobian
f = lambda p: torch.stack([p[0] * p[1], p[0] + p[1] ** 2])
jacobian(f, torch.tensor([2.0, 3.0]))   # tensor([[3., 2.], [1., 6.]])

Without autograd, a numerical gradient by central differences:

def numerical_gradient(f, point, h=1e-6):
    grad = []
    for i in range(len(point)):
        up, down = list(point), list(point)
        up[i] += h
        down[i] -= h
        grad.append((f(up) - f(down)) / (2 * h))
    return grad

numerical_gradient(lambda p: p[0] ** 2 + 3 * p[1] ** 2, [3, 2])   # ≈ [6.0, 12.0]

Tradeoffs

• Numerical versus automatic differentiation. Finite differences need two function evaluations per parameter, so a million-parameter model would take two million forward passes per gradient, and they lose precision as h shrinks. Autograd gets the exact gradient for about the cost of one extra forward pass. Finite differences remain the standard way to test autograd.
• The learning rate. Too small and training crawls; too large and it diverges. The safe size depends on the steepest curvature, which differs across parameters and changes during training. That is why schedules and adaptive optimizers exist.
• Local information only. The gradient describes an infinitesimal neighbourhood. On non-convex losses descent finds a nearby low point, not necessarily the lowest. In large networks most such points turn out to be good enough, and saddle points, where the gradient is also zero, are the more common nuisance.
• Non-differentiable points. ReLU and absolute value have kinks. Frameworks pick a subgradient there (PyTorch uses 0 for ReLU at 0), which works in practice because exact ties are rare.
• Memory. Recording the graph keeps intermediate values alive until backward. Wrap inference in torch.no_grad() or torch.inference_mode() to skip it.

Check yourself

Connections

References

• Goodfellow, Bengio and Courville. Deep Learning, chapter 4.3 (gradient-based optimization). MIT Press, 2016.
• Deisenroth, Faisal and Ong. Mathematics for Machine Learning, chapter 5 (vector calculus). Cambridge University Press, 2020.
• Baydin et al. Automatic Differentiation in Machine Learning: a Survey. JMLR 2018. arXiv:1502.05767
• PyTorch v2.15.0: torch.autograd.backward and Tensor.backward.