Calculus for ML

Training works by nudging millions of numbers in whichever direction makes the model slightly less wrong, and a derivative is what says which direction that is. This section covers the idea as a slope, gradients as the same idea in many dimensions at once, and the chain rule that carries a correction back through a stack of layers. You will not compute these by hand in practice, but when a training run refuses to improve the explanation is almost always here.

A derivative is a rate of change#

If you nudge the input slightly, how much does the output move? That is the derivative. In training, the question is always: if I nudge this parameter, how much does the loss move?

Make that concrete with a function you can compute in your head. Take f of x equals x squared, and stand at x equals 3, where f is 9. Move a whole step to x equals 4 and f becomes 16, so the output rose by 7 for an input rise of 1: a rate of 7. Move a tenth of a step instead, to 3.1, and f becomes 9.61, a rise of 0.61 over 0.1, which is a rate of 6.1. A hundredth of a step gives 6.01. The smaller the nudge, the closer the measured rate gets to 6, and 6 is the derivative of x squared at 3. The formula 2x is a shortcut for that measurement, not a separate idea.

def f(x):
    return x ** 2

for h in (1.0, 0.1, 0.01, 0.00001):
    print(h, (f(3 + h) - f(3)) / h)

# 1.0      7.0
# 0.1      6.100000000000012
# 0.01     6.009999999999849
# 1e-05    6.000009999951316

The sign carries as much information as the number. A derivative of 6 says the output goes up when the input goes up. A derivative of -6 says it goes down. Zero says that, right here, moving the input a little does nothing — which is what being at the bottom of a valley feels like, and also the top of a hill, and also a flat stretch. That ambiguity is worth holding on to: it is why training can stall in places that are not the answer.

Gradients#

With millions of parameters, the derivative with respect to each one is a partial derivative, and the collection of all of them is the gradient. It points in the direction of steepest increase, so training steps in the opposite direction.

A partial derivative is the same measurement with one difference: nudge one input and hold every other one still. Take g of x and y equals x squared plus 3 y squared, at the point x equals 1, y equals 2, where g is 13. Nudge x alone and the rate is 2x, which is 2. Nudge y alone and the rate is 6y, which is 12. So the gradient at that point is the pair [2, 12]: the function is six times more sensitive to y than to x here. That is all a gradient is — one number per input, each one answering how much does the output move if I nudge only this.

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

eps = 1e-5
x, y = 1.0, 2.0
print(g(x, y))                                      # 13.0
print((g(x + eps, y) - g(x - eps, y)) / (2 * eps))  # 2.0000000000131024
print((g(x, y + eps) - g(x, y - eps)) / (2 * eps))  # 12.000000000078613

Now the claim that makes training work: of all the directions you could step, the gradient points where the function rises fastest, so its negative is where it falls fastest. Test it rather than believing it. Take a step of the same tiny length in four different directions from that point and compare how much g drops.

import numpy as np

grad = np.array([2.0, 12.0])
here = np.array([1.0, 2.0])
step = 0.01

directions = {
    "against the gradient": -grad / np.linalg.norm(grad),
    "along -x only": np.array([-1.0, 0.0]),
    "along -y only": np.array([0.0, -1.0]),
    "some other way": np.array([-0.6, -0.8]),
}
for name, d in directions.items():
    print(name, round(g(*(here + step * d)) - g(*here), 6))

# against the gradient  -0.121361
# along -x only         -0.0199
# along -y only         -0.1197
# some other way        -0.107772

loss = 2.4

gradient tells you, for every parameter:
  increase this one and loss goes up      -> decrease it
  increase this one and loss goes down    -> increase it
  changing this one does nothing          -> leave it

Take a small step against the gradient. Repeat a few million times.

The chain rule#

A network is functions composed with functions. The chain rule says the derivative of a composition is the product of the derivatives along the path. That is what lets a loss at the output be attributed back to a weight in the first layer.

input -> layer1 -> layer2 -> layer3 -> loss

effect of a layer1 weight on the loss
  = its effect on layer1's output
  x layer1's effect on layer2
  x layer2's effect on layer3
  x layer3's effect on the loss

Work one all the way through, because the rule only becomes obvious once you have. Start with a single weight w equal to 0.5, an input x equal to 2, and a target of 1. Three steps: a equals w times x, which is 1.0; b equals tanh of a, which is 0.7616; and the loss is b minus the target, squared, which is 0.0568. Now go backwards, one local derivative per step. The loss changes with b at a rate of 2 times b minus target, which is -0.4768. b changes with a at a rate of 1 minus b squared, which is 0.4200. And a changes with w at a rate of x, which is 2. Multiply the three: -0.4768 times 0.4200 times 2 is -0.4005.

So increasing w decreases the loss here, at a rate of about 0.4 per unit of w. That number was produced with no calculus beyond three memorised derivatives and a multiplication, and the whole of backpropagation is that procedure with more steps and matrices in place of single numbers.

import math

w, x, target = 0.5, 2.0, 1.0

a = w * x                          # 1.0
b = math.tanh(a)                   # 0.7615941559557649
loss = (b - target) ** 2           # 0.056837346474444175

dloss_db = 2 * (b - target)        # -0.4768116880884703
db_da = 1 - b * b                  #  0.41997434161402614
da_dw = x                          #  2.0
print(dloss_db * db_da * da_dw)    # -0.4004973495576554

def loss_of(w_value):              # the numerical check, always worth doing
    return (math.tanh(w_value * x) - target) ** 2

print((loss_of(w + 1e-6) - loss_of(w - 1e-6)) / 2e-6)    # -0.4004973495549957

The step, and how big it should be#

The gradient gives a direction, not a distance. Gradient descent multiplies it by a learning rate and subtracts. Everything that can go wrong with that choice shows up in a function small enough to follow by hand: take the loss to be 5 w squared, whose gradient is 10 w, and start at w equals 1 where the loss is 5. The update is w minus learning rate times 10 w, so each step multiplies w by 1 minus ten times the learning rate.

loss(w) = 5 w^2      gradient = 10 w      start at w = 1

lr = 0.05   multiplier 0.5      1, 0.5, 0.25, 0.125, 0.0625, 0.0312     converging
lr = 0.19   multiplier -0.9     1, -0.9, 0.81, -0.729, 0.6561, -0.5905  bouncing, still shrinking
lr = 0.21   multiplier -1.1     1, -1.1, 1.21, -1.331, 1.4641, -1.6105  bouncing, growing: diverged
lr = 0.002  multiplier 0.98     1, 0.98, 0.9604, 0.9412, 0.9224         technically working

The multiplier is 1 - 10 x lr. It flips sign at lr = 0.1 and exceeds
1 in size at lr = 0.2, which is exactly where divergence begins.

for lr in (0.05, 0.19, 0.21):
    w = 1.0
    history = [w]
    for _ in range(8):
        gradient = 10 * w              # d/dw of 5 w^2
        w = w - lr * gradient
        history.append(round(w, 4))
    print(lr, history)

# 0.05 [1.0, 0.5, 0.25, 0.125, 0.0625, 0.0312, 0.0156, 0.0078, 0.0039]
# 0.19 [1.0, -0.9, 0.81, -0.729, 0.6561, -0.5905, 0.5314, -0.4783, 0.4305]
# 0.21 [1.0, -1.1, 1.21, -1.331, 1.4641, -1.6105, 1.7716, -1.9487, 2.1436]

Two lessons generalise from those nine lines. The threshold depends on the curvature of the loss and not on the learning rate alone: the 5 in 5 w squared is what put the limit at 0.2, and a loss ten times as curved would diverge at a tenth of the rate. That is why a learning rate copied from another project with different data or a different model size so often fails. And the symptom is diagnostic: a loss that oscillates with growing amplitude means the rate is above the threshold, while a loss that falls smoothly but barely means it is far below.

What this explains#

The chain rule being a product explains the two classic training failures directly, which is the practical payoff of understanding it.

• Vanishing gradients: many factors below one multiply towards zero, so early layers stop learning. This is why deep networks needed residual connections and better activations.
• Exploding gradients: factors above one multiply upwards, the update is enormous, and the loss becomes NaN. This is why gradient clipping exists.

How far to go#

You need comprehension, not fluency in manual differentiation. Frameworks compute gradients automatically, and nobody in industry differentiates by hand. The point is being able to reason about why training is misbehaving.

You know that training follows the gradient downhill. This section is about the surface it is walking on, because the picture most people carry from two dimensions is misleading in millions: nearly every flat point is a saddle rather than a local minimum, which makes getting permanently stuck far less of a worry than it sounds. Once you can see the landscape, the optimiser settings stop being magic numbers. Momentum, learning rate and the schedules around them are each an answer to a specific feature of that surface, and knowing which feature is causing trouble is how you rescue a run that will not converge.

Picture a landscape#

The loss is a surface over millions of dimensions, one per parameter. Training walks downhill. Intuition from two dimensions is partly wrong here, and the most useful correction is that in high dimensions most flat points are saddle points rather than local minima.

Local minimum   downhill in no direction      rare in high dimensions
Saddle point    downhill in some directions   common, and they slow training
Plateau         nearly flat in all directions small gradients, slow progress

The reason saddles dominate is a counting argument worth carrying around. At a flat point, each direction is either curving up or curving down. A local minimum requires every one of them to curve up. With two parameters that is a coin landing heads twice; with a million parameters it is a coin landing heads a million times. Somewhere in all those directions there is almost always one that still goes down, and a point with at least one downhill direction is a saddle, not a minimum.

The smallest saddle is the surface x squared minus y squared. It curves up along x and down along y, so the origin is flat but not a minimum. Start near it, slightly off the ridge, and watch the two coordinates behave completely differently: x is pulled straight to zero, while y, which started at 0.01, takes a long time to make anything of the downhill direction it is on.

import numpy as np

def saddle_run(momentum, lr=0.1, steps=200):
    p, v = np.array([1.0, 0.01]), np.zeros(2)
    for i in range(1, steps + 1):
        g = np.array([2 * p[0], -2 * p[1]])      # gradient of x^2 - y^2
        v = momentum * v + g
        p = p - lr * v
        if abs(p[1]) > 1.0:
            return i, p.round(4)
    return None, p.round(4)

print(saddle_run(0.0))     # (26, array([0.003 , 1.1448]))
print(saddle_run(0.9))     # (13, array([0.4992, 1.2975]))

Momentum#

Plain gradient descent zigzags across narrow valleys and stalls on flat regions. Momentum keeps a running average of past gradients, so consistent directions build up speed and oscillating ones cancel out. It carries the optimiser through saddle points and plateaus that would otherwise trap it.

Build the valley and count the steps. A loss of one half of x squared plus twenty y squared is a long trough: gently sloped along x, steeply walled along y. The ratio of those two curvatures, 20 to 1, is called the condition number, and it is the single number that decides how badly plain gradient descent will behave. The step size is limited by the steep direction, at 2 divided by 20, which is 0.1 — above that it diverges — but progress has to be made in the gentle direction, where a step of that size is tiny.

def valley(lr, momentum, steps=2000):
    p, v = np.array([1.0, 1.0]), np.zeros(2)
    for i in range(1, steps + 1):
        g = np.array([1.0 * p[0], 20.0 * p[1]])         # gradient of 0.5(x^2 + 20y^2)
        v = momentum * v + g
        p = p - lr * v
        if 0.5 * (p[0] ** 2 + 20 * p[1] ** 2) < 1e-4:
            return i
    return None

print(valley(0.02, 0.0), valley(0.02, 0.9))      # 211 83
print(valley(0.09, 0.0), valley(0.09, 0.9))      # 46 98

That second line is the part people are not told. Momentum is not free speed; it multiplies the effective step size by roughly 1 divided by 1 minus the momentum coefficient, which at 0.9 is ten. Turning momentum on without lowering the learning rate can take a run that was stable and make it oscillate, and the symptom looks identical to a learning rate that was simply set too high.

Learning rate is step size#

On a curved surface, too large a step overshoots the valley and bounces up the far side; too small and progress is glacial. The right step depends on local curvature, which varies across the surface and across parameters. Adaptive optimisers like Adam exist to set the step per parameter.

The condition number is why the per-parameter part matters. In the valley above, one learning rate has to serve a direction that is twenty times steeper than the other, so whatever you choose is wrong for one of them. Real networks have condition numbers far worse than 20, and the parameters are not conveniently aligned with the axes. Adam's per-parameter scaling is an attempt to give the gentle directions larger steps and the steep ones smaller, without ever computing the curvature it would need to do this properly.

The size of that expense is worth knowing rather than assuming. The curvature of a loss with n parameters is an n by n matrix, so for a model with a billion parameters it has a billion squared entries — an amount of memory that does not exist. Newton's method would need to solve a system with that matrix at every step. Every practical second-order method is therefore an approximation that never forms the matrix, and the cheapest of those approximations, a diagonal estimate, is essentially what Adam maintains.

What to do when a run stalls#

symptom                                  most likely cause         first thing to try
loss flat from step 0                    bug, not landscape        overfit one batch
loss falls then sits still for a while   saddle or plateau          wait; check gradient norms
loss oscillates, amplitude steady        learning rate near limit   halve it
loss oscillates, amplitude growing       learning rate over limit   halve it, add warmup
loss falls smoothly but very slowly      rate far below limit       raise it, or add momentum
loss fell, then jumped and stayed high   one bad batch or spike     clip gradients

What is worth memorising from this section is short. Flat does not mean finished. The stable learning rate is set by the steepest curvature in the problem, which is why a rate that worked on one model fails on another. Momentum lengthens the effective step by a factor of about ten at the usual setting. And nothing here requires you to compute a second derivative by hand, ever.

Training a model means nudging numbers until a loss goes down, and a derivative is what tells you which way to nudge and by how much. This section takes that idea from something you can feel — change the input a little, see how far the output moves — through to the derivatives of the particular functions that keep appearing in machine learning, and the chain rule that strings them together. Every result here is checked numerically in code, so you never have to take one on faith.

Nudge and measure#

def f(x):
    return x ** 2

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

print(numerical_derivative(f, 3.0))     # about 6.0, and by hand d/dx x^2 = 2x = 6

A derivative says how much the output changes per unit change in the input, for a tiny change. You can always estimate it by nudging the input both ways and measuring. That makes every hand-derived formula checkable.

Nudging both ways rather than one way matters, and so does the size of the nudge. Too large and you are measuring the average slope across a curve rather than the slope at a point; too small and the two function values are so close that floating point subtraction throws away the digits that carried the answer. The usable window is wide but it is not unlimited.

for h in (1e-2, 1e-5, 1e-8, 1e-12, 1e-15):
    print(h, (f(3 + h) - f(3 - h)) / (2 * h))

# 0.01    5.999999999999872
# 1e-05   6.000000000039306
# 1e-08   5.999999963535174
# 1e-12   6.000533403494046
# 1e-15   5.329070518200751      three digits of the answer are gone

The derivatives ML uses constantly#

x^2              ->  2x
e^x              ->  e^x
log(x)           ->  1 / x
sigmoid(x)       ->  sigmoid(x) x (1 - sigmoid(x))
relu(x)          ->  1 if x > 0 else 0

softmax followed by cross-entropy, with respect to the logits:
                 ->  p - y       (predicted probabilities minus the one-hot target)

The last line is worth memorising. The gradient of the standard classification loss with respect to the logits is simply the predicted probabilities minus the true labels. The model is pushed up on the right class by how much it underestimated it, and down on the others.

Confirm it on three numbers rather than trusting it. Logits of 2.0, 1.0 and 0.1 become probabilities 0.659, 0.242 and 0.099 after a softmax. If the correct class is the second one, the target vector is [0, 1, 0], so p minus y is [0.659, -0.758, 0.099]. Positive on the two wrong classes, meaning push those logits down; negative on the right one, meaning push it up, and by the largest amount because that is where the model was most wrong.

import numpy as np
import torch
import torch.nn.functional as F

logits = np.array([2.0, 1.0, 0.1])
e = np.exp(logits - logits.max())          # subtract the max: no overflow
probs = e / e.sum()
target = np.array([0.0, 1.0, 0.0])         # the true class is index 1

print(probs.round(6))                       # [0.659001 0.242433 0.098566]
print((probs - target).round(6))            # [ 0.659001 -0.757567  0.098566]

t = torch.tensor(logits, requires_grad=True)
loss = F.cross_entropy(t.unsqueeze(0), torch.tensor([1]))
loss.backward()
print(t.grad.numpy().round(6))               # [ 0.659001 -0.757567  0.098566]
print(loss.item())                           # 1.4170300162778335

The chain rule through two steps#

import math

x, w, target = 2.0, 0.5, 1.0
z = w * x                     # step 1
p = 1 / (1 + math.exp(-z))    # step 2: sigmoid
loss = (p - target) ** 2      # step 3

dloss_dp = 2 * (p - target)
dp_dz = p * (1 - p)
dz_dw = x
dloss_dw = dloss_dp * dp_dz * dz_dw     # multiply along the chain
print(z, p, loss)              # 1.0 0.7310585786300049 0.07232948812851325
print(dloss_dw)                # -0.21150837113706686

Read the three factors as a sentence. The loss is 0.0723 and the prediction is 0.731 against a target of 1, so the first factor is negative: the prediction needs to go up. The second factor, 0.197, is the sigmoid's derivative at this point, and it is a dampener — no matter how wrong the prediction is, this factor can never exceed 0.25. The third is the input, 2, which is why an input of zero gives that weight no gradient at all. Increase w and the loss falls at a rate of 0.21.

Check against autograd#

w_t = torch.tensor(0.5, requires_grad=True)
p_t = torch.sigmoid(w_t * 2.0)
((p_t - 1.0) ** 2).backward()
print(w_t.grad.item())        # -0.2115083485841751

That seven-digit agreement is the standard you want from a gradient check. Order-of-magnitude agreement means you have the shape of the derivative right and a factor wrong somewhere. A sign flip means you subtracted in the wrong order. Agreement on some entries and not others usually means an axis is being summed that should not be.

Why sigmoid ruins deep networks, in one line of arithmetic#

The sigmoid's derivative is largest at the middle, where it is exactly 0.25, and it falls away fast on either side. At an input of 2 it is 0.105; at 4 it is 0.0177; at 8 it is 0.000335. Since the chain rule multiplies one of these per layer, ten sigmoid layers at their very best multiply the gradient by 0.25 to the power of ten, which is 0.00000095. The early layers of such a network receive a signal a million times weaker than the last layer, before any of the other factors are taken into account.

import math

for z in (0, 1, 2, 4, 8):
    s = 1 / (1 + math.exp(-z))
    print(z, round(s, 6), round(s * (1 - s), 8))

# 0  0.5       0.25         the best it ever is
# 1  0.731059  0.19661193
# 2  0.880797  0.10499359
# 4  0.982014  0.01766271
# 8  0.999665  0.00033524   effectively no gradient at all

print(0.25 ** 10)      # 9.5367431640625e-07
print(0.25 ** 20)      # 9.094947017729282e-13

What to memorise, and what to skip#

Five derivatives are worth having by heart because they turn up in every explanation: x squared goes to 2x, e to the x is its own derivative, log x goes to 1 over x, sigmoid goes to sigmoid times one minus sigmoid, and softmax with cross-entropy goes to p minus y. Add the chain rule, which is multiply the local derivatives along the path, and the numerical check, which is nudge and measure. That is the entire working set.

Skip the rest of a calculus course without hesitation: integration by parts, trigonometric derivatives, limits and epsilon-delta proofs, series expansions, differential equations. None of them appears in an engineering interview for these roles and none is needed to reason about training. If a paper uses one, it will be in a derivation you are reading rather than implementing.

Once you know which way to nudge a number, you need a rule for how far to step and when to stop. That rule is the optimiser, and this section builds the three you meet in every training script — plain gradient descent, momentum, and Adam — one on top of the other, in a few lines each. Writing them yourself is what turns the learning rate from a knob you copy off a blog post into something you can reason about when a run refuses to converge.

Step downhill#

import numpy as np
rng = np.random.default_rng(0)
x = rng.uniform(0, 10, 200)
y = 3 * x + 2 + rng.normal(0, 1, 200)
X = np.c_[x, np.ones_like(x)]

w = np.zeros(2)
lr = 0.01
for step in range(2000):
    grad = 2 / len(y) * X.T @ (X @ w - y)     # gradient of mean squared error
    w -= lr * grad
print(w)                                        # [2.97471834 2.04697713]

Two things are worth naming in those seven lines. The gradient expression is the derivative of the mean squared error, and it has a readable shape: X @ w minus y is the vector of errors, and multiplying by X transposed asks how much each feature was implicated in them. And the loop has no stopping rule; it runs a fixed 2,000 steps. That is deliberate, because the honest way to decide when to stop is to watch the loss, which the next few blocks do.

def mse(w):
    return ((X @ w - y) ** 2).mean()

def gradient(w):
    return 2 / len(y) * X.T @ (X @ w - y)

best = np.linalg.lstsq(X, y, rcond=None)[0]      # the exact answer, for comparison
print(best, mse(best))       # [2.97469853 2.04711649] 1.042141861405242

The learning rate#

Try lr = 0.0001: it crawls. Try lr = 0.05: it overshoots, the loss grows and the weights explode. There is a range that works, and it depends on the curvature of the loss, which is why the learning rate is the most important setting in all of deep learning.

for lr in (0.0001, 0.001, 0.01, 0.03):
    w = np.zeros(2)
    for step in range(2000):
        w = w - lr * gradient(w)
        if not np.isfinite(w).all():
            break
    ok = np.isfinite(w).all()
    print(lr, np.round(w, 4) if ok else "diverged", round(mse(w), 4) if ok else "inf")

# 0.0001  [3.1808 0.5967]   1.5423    still moving; the intercept is nowhere near 2
# 0.001   [3.0638 1.4204]   1.1355    closer, still short after 2000 steps
# 0.01    [2.9747 2.047 ]   1.0421    at the optimum
# 0.03    [-3.02308336e+253 -4.29660063e+252]  inf   diverged

Look carefully at the 0.0001 row before moving on, because it is the failure that wastes the most time in practice. The slope, 3.18, is almost right; the intercept, 0.60, is a third of the way to 2. The loss curve for that run is smooth and decreasing and looks entirely healthy. Nothing in it says the rate is wrong. The only way to know is to have something to compare against: a closed-form answer, a longer run, or a larger rate.

Mini-batches#

Computing the gradient over the whole dataset every step is slow. Stochastic gradient descent uses a small random batch instead. Each step is noisier and far cheaper, and the noise often helps escape poor regions.

w = np.zeros(2)
lr = 0.01
index = np.arange(len(y))
shuffler = np.random.default_rng(1)

for epoch in range(50):
    shuffler.shuffle(index)
    for start in range(0, len(index), 32):          # batches of 32
        batch = index[start:start + 32]
        g = 2 / len(batch) * X[batch].T @ (X[batch] @ w - y[batch])
        w = w - lr * g

print(np.round(w, 4), round(mse(w), 4))     # [2.9494 1.7267] 1.2568

Momentum#

w = np.zeros(2)
v = np.zeros_like(w)
lr, steps = 0.01, 2000

for step in range(steps):
    g = gradient(w)
    v = 0.9 * v + g          # a running direction, like a ball rolling downhill
    w -= lr * v
print(np.round(w, 4), round(mse(w), 4))     # [2.9747 2.0471] 1.0421

In a long narrow valley, plain gradient descent zigzags across. Momentum averages recent gradients, so the zigzags cancel and progress along the valley accumulates.

The same answer is not the point; the number of steps is. Counting how many updates each method needs to get within one per cent of the best achievable loss on this data gives plain descent 435, momentum 68, and Adam at a learning rate of 0.5 just 81 — and Adam at 0.1 needs 307, which is the reminder that these methods still have a learning rate and it still has to be roughly right.

Adam#

w = np.zeros(2)
m = np.zeros_like(w); v = np.zeros_like(w)
b1, b2, eps = 0.9, 0.999, 1e-8
lr, steps = 0.5, 2000

for t in range(1, steps + 1):
    g = gradient(w)
    m = b1 * m + (1 - b1) * g           # average gradient (momentum)
    v = b2 * v + (1 - b2) * g ** 2      # average squared gradient (scale)
    m_hat = m / (1 - b1 ** t)           # correct the start-up bias towards zero
    v_hat = v / (1 - b2 ** t)
    w -= lr * m_hat / (np.sqrt(v_hat) + eps)
print(np.round(w, 4), round(mse(w), 4))     # [2.9747 2.0471] 1.0421

Adam combines momentum with a per-parameter step size: parameters with consistently large gradients take smaller steps, rarely updated ones take larger ones. It works well with little tuning, which is why it, and its weight decay variant AdamW, are the default for transformers.

The bias correction is the line people skip and it is the reason Adam works from step one. Both running averages start at zero, so after the first update m is only a tenth of the gradient it just saw and v is a thousandth of the squared gradient. Dividing by 1 minus the decay to the power of t undoes exactly that: at t equal to 1 it divides m by 0.1 and v by 0.001, recovering the true values, and by a few hundred steps the correction has faded to nothing. Without it the first few steps are far too small in a way that does real damage on short runs.

Where Adam earns its keep#

On well-scaled data Adam is a convenience. Change one thing — put the first feature in units a thousand times larger, which is exactly what happens when one column is in metres and another in millimetres — and it becomes the difference between training and not training.

X2 = np.c_[x * 1000, np.ones_like(x)]          # the same data, one column rescaled

def grad2(w):
    return 2 / len(y) * X2.T @ (X2 @ w - y)

def mse2(w):
    return ((X2 @ w - y) ** 2).mean()

print(np.linalg.lstsq(X2, y, rcond=None)[0], mse2(np.linalg.lstsq(X2, y, rcond=None)[0]))
# [0.0029747 2.04711649] 1.0421418614052422   -- the same fit, slope divided by 1000

for lr in (1e-8, 1e-7):
    w = np.zeros(2)
    for t in range(5000):
        w = w - lr * grad2(w)
        if not np.isfinite(w).all():
            break
    ok = np.isfinite(w).all()
    print(lr, np.round(w, 6) if ok else "diverged", round(mse2(w), 4) if ok else "inf")
# 1e-08  [3.264e-03 4.900e-05]  2.0383    the intercept has barely moved from zero
# 1e-07  diverged               inf

w = np.zeros(2)
m = np.zeros_like(w); v = np.zeros_like(w)
for t in range(1, 5001):
    g = grad2(w)
    m = b1 * m + (1 - b1) * g
    v = b2 * v + (1 - b2) * g ** 2
    w -= 0.01 * (m / (1 - b1 ** t)) / (np.sqrt(v / (1 - b2 ** t)) + eps)
print(np.round(w, 6), round(mse2(w), 4))    # [0.002975 2.045165] 1.0421

That is the whole argument for adaptive optimisers in one experiment, and it is also an argument for standardising features when you can. Adam rescues a badly conditioned problem; it does not make conditioning stop mattering.