Lesson 1.3: Tensors, gradients and autograd by hand
Build a small automatic-differentiation engine in Python, check its gradients against numerical estimates, and use it to fit a line by gradient descent.
By the end of this lesson you will be able to explain what a gradient is and why a model needs one, apply the chain rule to a computation written as a graph of small steps, write a working automatic-differentiation engine in about sixty lines of Python, check any gradient against a numerical estimate, state the gradients of a softmax with cross-entropy loss and of a matrix product, and use gradient descent to fit a model while choosing the step size sensibly. You will show that an engine we wrote gives the gradient 8.0 for y = x² + 2x at x = 3, that its gradients for a small neuron agree with finite differences to six decimal places, that the gradient of the loss through a softmax is simply the predicted probabilities minus the one-hot target (−0.341, 0.242, 0.099), that the matrix-product gradient formulas match a numerical check to 2 × 10⁻⁹, and that gradient descent fits y = 2x + 1 from a start of zero, with a step size above about 0.24 making it diverge to 10³⁷. You should have read Lessons 1.1 and 1.2. The mathematics is the derivative of a simple function and the chain rule, which we review. This course is education about how language models work, not a guide to any commercial product.
In Lesson 1.1 we found the best table of counts for a bigram model by counting. That worked because a bigram has one cell per pair of words. A real model has millions of numbers that interact in complicated ways, and there is no counting formula for them. What we have instead is a loss, a single number measuring how badly the model does, and a way of asking, for each parameter: if I nudged this number up a little, would the loss go up or down, and by how much? The answer to that question is the gradient, and the machinery that computes it for every parameter at once is the engine of all deep learning. This lesson builds that machinery from nothing.
The derivative, in one paragraph
For a function f of one number x, the derivative df/dx is the rate at which f changes when x changes. If f(x) = x², then a small step h gives f(x + h) = x² + 2xh + h², so the change is about 2x times h, and the derivative is 2x. At x = 3 the slope is 6. If you move x up by 0.01, f rises by about 0.06.
With many inputs, the gradient is the list of derivatives with respect to each one, holding the others fixed. For a loss L that depends on parameters w₁, w₂, …, the gradient is (∂L/∂w₁, ∂L/∂w₂, …). Each entry says how sensitive the loss is to one parameter. The rule for improving a model follows at once: move every parameter a little against its gradient. If ∂L/∂w is positive, increasing w raises the loss, so decrease w; if it is negative, increase w. That rule, applied again and again, is gradient descent:
w ← w − η × ∂L/∂w
The small positive number η, called the learning rate, sets the size of each step. We will see below what happens when it is wrong.
The chain rule and the computation graph
Nearly every model is a long chain of simple operations: multiply, add, take an exponential or a logarithm. The chain rule says how to get the derivative of a chain. If z depends on y and y depends on x, then
dz/dx = (dz/dy) × (dy/dx)
Rates multiply. If y changes 3 times as fast as x, and z changes 2 times as fast as y, then z changes 6 times as fast as x.
To use this on a big computation, we write the computation as a graph: each node is a number, produced by a simple operation from the nodes that feed it. Take f = (a × b + c) × (−2), with a = 2, b = −3 and c = 10. The graph has these steps:
d = a × b = −6
e = d + c = 4
f = e × (−2) = −8
To find how f depends on a, we walk backwards from f and multiply the local derivatives along the way. Start with df/df = 1.
f = e × (−2), so df/de = −2.
e = d + c, so de/dd = 1 and de/dc = 1. Therefore df/dd = −2 × 1 = −2 and df/dc = −2 × 1 = −2.
d = a × b, so dd/da = b = −3 and dd/db = a = 2. Therefore df/da = −2 × (−3) = 6 and df/db = −2 × 2 = −4.
The gradient of f with respect to (a, b, c) is (6, −4, −2). Going backwards from the output to the inputs and multiplying local derivatives is called backpropagation, or the reverse mode of automatic differentiation. Its great strength is cost: one backward pass gives the derivative with respect to every input at once, however many there are. A model with a billion parameters needs one forward pass and one backward pass per step, not a billion separate experiments.
One more rule completes the picture. If a number is used in more than one place, its gradient is the sum of the contributions from every place it is used. For y = x × x + 2x, x feeds a product (twice) and a multiplication by 2. The derivative of the first part is x + x = 2x, the second is 2, and the total is 2x + 2. At x = 3 that is 8.
An engine in sixty lines
We can now write a program that does this for us. Each number in the computation is an object of class Value that remembers its data, its gradient, the values it was computed from, and a small function that knows how to pass gradient back to them.
python
import math
class Value:
def init(self, data, parents=(), op=""):
self.data = float(data)
self.grad = 0.0
self.parents = parents
self.backward = lambda: None
self.op = op
def add(self, other):
other = other if isinstance(other, Value) else Value(other)
out = Value(self.data + other.data, (self, other), "+")
def backward():
self.grad += out.grad # d(out)/d(self) = 1
other.grad += out.grad
out.backward = backward
return out
def mul(self, other):
other = other if isinstance(other, Value) else Value(other)
out = Value(self.data other.data, (self, other), "")
def backward():
self.grad += other.data out.grad # d(out)/d(self) = other
other.grad += self.data out.grad
out.backward = backward
return out
def pow(self, k):
out = Value(self.data k, (self,), f"{k}")
def backward():
self.grad += k self.data (k - 1) out.grad
out.backward = backward
return out
def exp(self):
out = Value(math.exp(self.data), (self,), "exp")
def backward():
self.grad += out.data out.grad # exp is its own derivative
out.backward = backward
return out
def log(self):
out = Value(math.log(self.data), (self,), "log")
def backward():
self.grad += (1 / self.data) out.grad
out.backward = backward
return out
def tanh(self):
t = math.tanh(self.data)
out = Value(t, (self,), "tanh")
def backward():
self.grad += (1 - t t) out.grad
out.backward = backward
return out
def neg(self): return self -1
def sub(self, other): return self + (-other if isinstance(other, Value) else Value(-other))
def radd(self, other): return self + other
def rmul(self, other): return self other
def truediv(self, other):
return self (other if isinstance(other, Value) else Value(other)) -1
def backward(self):
order, seen = [], set()
def build(v): # a topological order: parents before children
if id(v) not in seen:
seen.add(id(v))
for p in v.parents:
build(p)
order.append(v)
build(self)
self.grad = 1.0 # d(output)/d(output) = 1
for v in reversed(order): # walk backwards from the output
v.backward()
Read the pieces in turn.
Each operator creates a new Value and stores a closure backward that holds the local derivative. For a product, the local derivative with respect to one input is the other input, which is why mul multiplies by other.data.
The gradient is added with +=, never assigned with =. This implements the rule for numbers used twice: every use deposits its contribution and they add up.
backward first sorts the graph so that every node comes after the ones it was built from, then visits the nodes in reverse. By the time a node is visited, everything downstream of it has already deposited its contribution, so its gradient is complete.
Let us run it on the two examples above.
python
a, b, c = Value(2.0), Value(-3.0), Value(10.0)
f = (a b + c) Value(-2.0)
f.backward()
print(f.data, a.grad, b.grad, c.grad) # -8.0 6.0 -4.0 -2.0
x = Value(3.0)
y = x x + x 2
y.backward()
print(y.data, x.grad) # 15.0 8.0
The engine returns the gradients we derived by hand, (6, −4, −2), and the value 8 for the repeated variable.
A check you can always run
How do we know the engine is right? There is an independent way to estimate a gradient that needs no calculus: nudge the input and see how the output moves. The central difference is
df/dx ≈ [f(x + h) − f(x − h)] / (2h)
for a small h such as 10⁻⁶. It is slow, because it needs two forward evaluations for every parameter, but it is simple and almost impossible to get wrong. That makes it the standard gradient check. Whenever we write a new operation, we compare its gradient with this estimate.
Take a single neuron: it multiplies two inputs by two weights, adds a bias, and squashes the result with tanh.
python
w1, w2, bias = Value(-3.0), Value(1.0), Value(6.8814)
x1, x2 = Value(2.0), Value(0.0)
out = (w1 x1 + w2 x2 + bias).tanh()
out.backward()
print(out.data) # 0.70712
print(w1.grad, w2.grad, bias.grad, x1.grad, x2.grad)
# 0.9999626 0.0 0.4999813 -1.4999440 0.4999813
f = lambda w1, w2, b, x1, x2: math.tanh(w1 x1 + w2 x2 + b)
vals = [-3.0, 1.0, 6.8814, 2.0, 0.0]
def numeric(i, h=1e-6):
up, down = list(vals), list(vals)
up[i] += h
down[i] -= h
return (f(up) - f(down)) / (2 h)
print([round(numeric(i), 6) for i in range(5)])
# [0.999963, 0.0, 0.499981, -1.499944, 0.499981]
The two lists agree to six decimal places. Notice what they say. The gradient with respect to w2 is exactly zero, because x2 is zero: changing a weight that multiplies zero has no effect on the output. And the gradient with respect to w1 is the largest, because x1 is 2. A gradient is not only a direction in which to move. It is a statement about which parameters matter on this input.
The gradient through a softmax
Language models end with a softmax and a cross-entropy loss. Both are built from exponentials, sums and a logarithm, so our engine can differentiate them. Take three scores 2.0, 1.0 and 0.1 (logits), and suppose the correct token is the first one. We compute the probabilities with the stable form from Lesson 1.2, then the loss −ln p₀.
python
z = [Value(v) for v in (2.0, 1.0, 0.1)]
m = max(v.data for v in z)
exps = [(v - m).exp() for v in z]
total = exps[0] + exps[1] + exps[2]
probs = [e / total for e in exps]
loss = -(probs[0].log())
loss.backward()
print([round(p.data, 4) for p in probs]) # [0.659, 0.2424, 0.0986]
print(round(loss.data, 4)) # 0.417
print([round(v.grad, 4) for v in z]) # [-0.341, 0.2424, 0.0986]
The loss is 0.417, which is −ln 0.659. Now look at the gradients: (−0.341, 0.2424, 0.0986). Compare them with the probabilities (0.659, 0.2424, 0.0986) and the one-hot target (1, 0, 0). The gradient is exactly the probabilities minus the target: 0.659 − 1 = −0.341, 0.2424 − 0, 0.0986 − 0. This is a general result for the softmax with cross-entropy, and it is one of the most useful facts in the course:
∂L/∂zᵢ = pᵢ − yᵢ
It has a lovely reading. The correct score gets a negative gradient, so gradient descent raises it. Every wrong score gets a positive gradient equal to the probability the model wasted on it, so descent lowers it, and lowers it more the more confident the mistake. When the model is certain and right (p = 1 for the correct token) the gradient is zero, and nothing more is learned. The algebra is short, but the engine has verified it for us. We will use the formula in Lesson 1.4 so that we do not have to run the Python engine for every token, which would be far too slow.
Gradients for whole arrays
The Value engine handles one number at a time. That is perfect for understanding and too slow for a model: a single multiplication of two 300 × 300 matrices would create 27 million objects. Real engines work on whole arrays, which are called tensors when they may have any number of axes. The idea is identical. Only the local derivatives are written for arrays.
The most important is the matrix product. Let Y = X W, where X has shape (n, d), W has shape (d, k) and Y has shape (n, k), and suppose the loss L depends on Y. If we already know dY, the gradient of the loss with respect to Y, then the gradients of the inputs are
dW = Xᵀ dY, of shape (d, k)
dX = dY Wᵀ, of shape (n, d)
A rule of thumb helps you to remember them: each gradient has the same shape as the thing it is the gradient of, and there is usually only one way to arrange the available matrices to achieve that. We verify the first formula against the numerical check on a random problem with a squared-error loss:
python
rng = np.random.defaultrng(0)
X = rng.standardnormal((4, 3))
W = rng.standardnormal((3, 2))
T = rng.standardnormal((4, 2))
def loss(W):
return 0.5 ((X @ W - T) 2).sum()
dY = X @ W - T # (4, 2): gradient of the loss with respect to Y
dW = X.T @ dY # (3, 2)
numeric = np.zeroslike(W)
for i in range(3):
for j in range(2):
up, down = W.copy(), W.copy()
up[i, j] += 1e-6
down[i, j] -= 1e-6
numeric[i, j] = (loss(up) - loss(down)) / 2e-6
print(np.abs(dW - numeric).max()) # about 1.9e-09
The largest disagreement is 1.9 × 10⁻⁹, which is rounding error. The array formula is right. Every layer of the language model we build in later modules is made of operations like this, and each has a short, checkable gradient formula. By Module 4 we will have a handful of them and will use them to train the full model.
Gradient descent in action
Everything is now in place to learn something. We have four points, (0, 1), (1, 3), (2, 5) and (3, 7), which lie exactly on the line y = 2x + 1. We pretend not to know that, start with a slope w = 0 and an intercept b = 0, and ask gradient descent to find them. The loss is the mean squared error, and we use the Value engine for the gradients:
python
xs = [0.0, 1.0, 2.0, 3.0]
ys = [1.0, 3.0, 5.0, 7.0]
def train(lr, steps):
w, b = Value(0.0), Value(0.0)
history = []
for step in range(steps + 1):
loss = Value(0.0)
for x, y in zip(xs, ys):
err = w x + b - y
loss = loss + err err
loss = loss (1 / len(xs))
w.grad = 0.0 # clear old gradients before each backward pass
b.grad = 0.0
loss.backward()
history.append((step, loss.data, w.data, b.data))
w.data -= lr w.grad
b.data -= lr b.grad
return history
for row in train(0.05, 200):
if row[0] in (0, 1, 2, 10, 50, 100, 200):
print(row)
Step · Loss · Slope w · Intercept b0 · 21.0 · 0.0000 · 0.0000
1 · 7.0588 · 0.8500 · 0.4000
2 · 2.3730 · 1.3425 · 0.6325
10 · 0.0009 · 2.0093 · 0.9576
50 · 0.00005 · 2.0054 · 0.9885
100 · 0.0000 · 2.0012 · 0.9975
200 · 0.0000 · 2.0001 · 0.9999
Two things are worth noticing. First, the model found the answer without ever being told the formula. All it had was the gradient. Second, the loss fell very fast at first, when the gradients were large, and slowly at the end, when the gradients were small. The step is the gradient times a constant, so as the parameters approach the best values and the gradient shrinks, so do the steps.
Also notice the line w.grad = 0.0 before backward(). The engine adds gradients with +=, so if they are not cleared, the gradient from the previous step is still there. Forgetting this is one of the most common bugs in training code: the model still trains, in a way, but each step uses the sum of all the previous gradients.
The learning rate is not a detail
We kept the learning rate at 0.05. What happens at other values? We run the same problem for 100 steps with six rates:
Learning rate · Loss after 100 steps · What happened0.001 · 3.95 · Far too slow: still far from the answer
0.01 · 0.0003 · Converging, but slowly
0.05 · 0.0000025 · Good
0.1 · 0.000000005 · Faster
0.2 · about 10⁻¹⁴ · Very fast
0.3 · 2.6 × 10³⁷ · Diverged: each step overshoots by more than the last
There is a precise reason for the last row. Near the minimum the loss looks like a bowl, and for a bowl with curvature a the update multiplies the distance from the bottom by (1 − η a) each step. If η a is between 0 and 2, that factor has size below 1 and the parameters approach the bottom. If η a exceeds 2, the factor has size above 1 and every step lands farther away than the last, on the other side. For this problem the steepest curvature is 8.4, so the stable range is η < 2 / 8.4 = 0.238. That is why 0.2 works beautifully and 0.3 explodes.
The lesson carries over to large models, with a complication: a real loss has many directions with different curvatures. The largest sets the limit on the learning rate, and the smallest sets how long the model takes to learn in its flat directions. This is the reason that real training uses more careful optimisers than plain gradient descent, and why the learning rate is the single most important number in a training run. Lesson 5.2 and 5.3 return to both.
What we have built
We now have the three pieces on which every neural network rests.
1. A way of writing a computation as a graph of simple operations.
2. A backward pass that turns the graph into the gradient of the output with respect to every input.
3. A rule for using gradients, gradient descent, and a warning that the step size must be chosen with care.
Everything else is an arrangement of these pieces. A transformer is a very large graph; training is the same loop we just wrote, with a loss from the softmax rule above, applied to more parameters and more data. The next lesson uses these tools to train a bigram language model, in which the table of counts is replaced by a table of numbers learned by gradient descent.
Key takeaways
The gradient of the loss with respect to a parameter says how much the loss changes if the parameter is nudged. Gradient descent moves each parameter against its gradient: w ← w − η × gradient.
The chain rule multiplies local derivatives along a path. Backpropagation applies it from the output backwards, so one backward pass gives every gradient at once.
A number that is used several times gets the sum of the gradients from each use, which is why gradients are accumulated with +=.
A gradient check compares the engine with (f(x + h) − f(x − h)) / 2h. Our neuron's gradients matched to six places; the matrix-product formulas matched to 1.9 × 10⁻⁹.
The gradient of cross-entropy loss through a softmax is the probabilities minus the one-hot target: (−0.341, 0.242, 0.099) in the example.
For Y = XW, dW = XᵀdY and dX = dYWᵀ, and each gradient has the shape of the thing it belongs to.
The learning rate must be below 2 divided by the steepest curvature. On the example, 0.2 converged to a loss of about 10⁻¹⁴ and 0.3 diverged to 10³⁷.
Common mistakes
Not clearing gradients. The engine accumulates. Set the gradients to zero before every backward pass.
Skipping the gradient check. A wrong gradient does not produce an error, only a model that trains poorly. Check each new operation once, with a small input.
Assigning instead of adding. Writing self.grad = ... instead of += gives the right answer on a chain and the wrong one wherever a value is used twice.
Choosing the learning rate by habit. The stable range depends on the problem. Try a few values across a factor of ten.
Forgetting that the loss is a number. Backpropagation needs a single output to start from. The loss is averaged over examples so that it is one number.
Using tiny test sizes to judge speed. The scalar engine is for understanding. Training uses the array version.
Check yourself
1. For f = x² at x = −2, what is the gradient, and which way does gradient descent move x?
df/dx = 2x = −4. Descent moves against the gradient, by η × 4 in the positive direction, towards zero.
2. A value u is used in two places: g = u × 3 and h = u × u, and the output is g + h. What is d(output)/du at u = 2?
d(3u)/du = 3 and d(u²)/du = 2u = 4, so the total is 7.
3. A softmax gives probabilities (0.1, 0.7, 0.2) and the correct token is the second. What is the gradient with respect to the three logits?
Probabilities minus the one-hot target: (0.1, 0.7 − 1, 0.2) = (0.1, −0.3, 0.2).
4. X has shape (32, 128) and W has shape (128, 512). What are the shapes of Y, dW and dX?
Y is (32, 512), dW = XᵀdY is (128, 512) and dX = dYWᵀ is (32, 128).
5. A loss has a curvature of 40 in its steepest direction. What is the largest stable learning rate?
2 / 40 = 0.05.
Try this
Set aside an hour. You need Python and NumPy.
1. Paste the Value class and reproduce the three gradient results: (6, −4, −2), 8, and the neuron's five gradients.
2. Add a relu operation (output is x if x > 0, otherwise 0; its gradient is 1 where x > 0 and 0 otherwise) and check it against the central difference at x = 2 and x = −2.
3. Change the training data to the noisy points (0, 1.2), (1, 2.7), (2, 5.3), (3, 6.9). Does the loss reach zero? What slope and intercept do you get, and why is the loss not zero?
4. Run the learning-rate study for 0.238 and 0.24. Describe what happens to the loss in each case.
5. Remove the lines that clear the gradients and train with a rate of 0.05. Describe the result, and explain it using the rule for accumulation.
6. Compute the gradient of a two-layer computation, tanh(w₂ × tanh(w₁ × x)), by hand and with the engine for w₁ = 0.5, w₂ = −1.5, x = 2. Check it against the central difference.
A good result is an engine whose gradients match the numerical check, a clear sense of why the learning rate has a ceiling, and the understanding that a "neural network" is a computation graph plus a loop. The next lesson uses all three to build the first learnable language model.
Read more
Andrej Karpathy's lecture "The spelled-out intro to neural networks and backpropagation: building micrograd" is the best-known walk-through of the idea in this lesson, and our engine follows the same design. Christopher Olah's essay "Calculus on Computational Graphs" explains the forward and reverse modes with pictures. Goodfellow, Bengio and Courville's textbook Deep Learning covers backpropagation in chapter 6 and optimisation in chapter 8. The classic reference for the algorithm is the 1986 paper by Rumelhart, Hinton and Williams, "Learning representations by back-propagating errors".
Tags: Build an LLM Gradients Backpropagation Autograd Gradient descent