Back propagation¶

This notebook is copied from Andrej Karpathy's Backpropagation lecture on youtube. I added a diagram and with the help of Claude, included some formulas in mathjax format.

• Building makemore Part 4: Becoming a Backprop Ninja
• Yes you should understand backprop
• Notebook

Compare manual gradients with PyTorch gradients¶

Model¶

• you will note that I changed the initialization a little bit to be small numbers. so normally you would set the biases to be all zero, here I am setting them to be small random numbers

• I'm doing this because if your variables are initialized to exactly zero, sometimes what can happen is that can mask an incorrect implementation of a gradient. Because when everything is zero, it sort of simplifies and gives you a much simpler expression of the gradient than you would otherwise get. so by making it small numbers, I'm trying to unmask those potential errors in these calculations.

• You also notice that I'm using b1 in the first layer. I'm using a bias, despite batch normalization right afterwards so this would typically not be what you do because we talked about the fact that you don't need a bias, but I'm doing this here just for fun because we're going to have a gradient with respect to it and we can check that we are still calculating it correctly, even though this bias is spurious.

1. Calculate gradients manually¶

1.1 - dlogprobs¶

What is inside logprobs? the shape of this is 32x27. so it's not going to surprise you that dlogprobs should also be an array of size 32x27 because we want the derivative loss with respect to all of its elements. so the sizes of those are always going to be equal.

• Now, how does logprobs influence the loss? Loss is negative logprobs indexed with range(n), YB and then the mean of that.

• Now, just as a reminder, Yb is just basically an array of all the correct indices. So what we're doing here is we're taking the logprobs array of size 32x27 and then we are going in every single row, and in each row we are plucking out the index 8 and then 14 and 15 and so on.

• So we're going down the rows That's the iterator range of n and then we are always plucking out the index at the column specified by this tensor Yb. So in the 0th row, we are taking the 8th column. In the first row, we're taking the 14th column, etc.

Calculating loss in Python¶

Cross entropy loss¶

• n - batch size
• K - number of classes
• i - index of the sample in the batch
• k - index of the class
• L - loss

is the true label (0 or 1) for the i-th example and k-th class. This assumes one hot encoding. $$\left(a_k^{[i]}\right)$$ is the predicted probability for the i-th example and k-th class

Derivative of mean¶

From this matrix, we're picking a single entry for each example and taking mean of them.

# Derivative of "mean".
# sample
# l = -(a + b + c) / 3
# l = -a/3 -b/3 - c/3
# dl/da = -1/3
# dl/db = -1/3
# dl/dc = -1/3

You see that logprobs shape is 32x27. But only 32 of them participate in the loss calculation. So what's the derivative of all the other, most of the elements that do not get plucked out here? Well, their gradient intuitively is zero. And that's because they did not participate in the loss.

So most of these numbers inside this tensor does not feed into the loss and so if we were to change these numbers, then the loss doesn't change, which is the equivalent of way of saying that the derivative with respect to them is zero. They don't impact it.

1.2 - dprobs¶

• So if probes is very, very close to 1, that means your network is currently predicting the character correctly, then 1/probs will become 1 over 1, and dlogprobs just gets pass-through.
• But if probabilities are incorrectly assigned, so if the correct character here is getting a very low probability, then 1.0 / probs will be higher and then multiplied by dlogprobs
• So basically what this line is doing intuitively is taking to the examples that have a very low probability currently assigned and it's boosting their gradient.

1.3 - dcounts_sum_inv¶

probs = counts * counts_sum_inv

• counts is 32 x 27
• counts_sum_inv is 32 x 1

This operation is composed of two steps:

• A Broadcast occurs (Each row of counts_sum_inv is replicated 27 times along the column - virtually making a 32 x 27 matrix)
• Element wise multiplication occurs

Derivatives¶

1.4 - dcounts - [1]¶

• counts is used to evaluate two values
  1. probs
  2. counts_sum
• so, when finding dcounts, these two expressions need to be considered.

1.5 - counts_sum¶

1.6 - dcounts [2]¶

• This derivative is a 32x27 matrix filled with ones.
• Each element in this matrix represents the partial derivative of the sum with respect to the corresponding element in the original matrix A. The derivative is 1 for each element because changing any element in A by a small amount δ will change the corresponding row sum by exactly δ.

1.7 - dnorm_logits¶

1.8 - dlogits - [1]¶

Derivatives

1.9 - dlogit_maxes¶

NOTE:

logit_maxes = logits.max(1, keepdim=True).values
norm_logits = logits - logit_maxes  # subtract max for numerical stability
counts = norm_logits.exp()

We've talked previously in the previous lecture that the only reason we're doing this is for the numerical stability of the soft max that we are implementing here and we talked about how if you take these logits for any one of these examples so one row of this logits tensor if you add or subtract any value equally to all the elements then the value of the probs will be unchanged. you're not changing the soft max! the only thing that this is doing is it's making sure that exp() doesn't overflow and the reason we're using a max() is because then we are guaranteed that each row of logits, the highest number is zero. And so this will be safe.

And so basically that has repercussions. If it is the case that changing logit_maxes does not change the probs and therefore there's not change the loss, then the gradient on logit_maxes should be zero, right? Because saying those two things is the same. So indeed, we hope that this is very, very small numbers. Indeed, we hope this is zero. Now, because of floating point sort of wonkiness, this doesn't come out exactly zero, only in some of the rows it does. But we get extremely small values, like 1e -9 or -10.

And so this is telling us that the values of logit_maxes are not impacting the loss as they shouldn't. It feels kind of weird to back propagate through this branch, honestly, because if you have any implementation of like F.cross_entropy() and PyTorch, and you block together all these elements and you're not doing the back propagation piece by piece, then you would probably assume that the derivative through here is exactly zero.

So you would be sort of skipping this branch because it's only done for numerical stability. But it's interesting to see that even if you break up everything into the full atoms and you still do the computation as you'd like with respect to numerical stability, the correct thing happens and you still get a very, very small gradients here, basically reflecting the fact that the values of these do not matter with respect to the final loss.

1.10 - dlogits - [2]¶

Forward operation

# For each row, identify the max value.
logit_maxes = logits.max(1, keepdim=True).values

Derivative

• Each row of the derivative matrix will have exactly one 1 and twenty-six 0s.
• The position of the 1 in each row corresponds to the position of the maximum value in that row of the original logits matrix.

1.11 - dh, dw2, db2¶

1. The simplified equation:

2. The expanded matrix form:

3. The resulting individual equations:

Derivatives

1. The equations represent partial derivatives of L with respect to different elements of matrix a:

• dL/dd11 is the global derivative (for chian rule)
• It is multiplied with local derivative of a11, which is b11
• Since a11 is used twice in this multiplication, its derivatives should be summed.

2. These equations are then represented in matrix form:

3. This is then shown to be equivalent to:

4. Finally, this is expressed in a more compact form using matrix notation:

Here, @ represents matrix multiplication, and $b^T$ is the transpose of matrix b.

This set of equations appears to be demonstrating the chain rule for matrix derivatives, specifically how the gradient of L with respect to a is related to the gradient with respect to d and the elements of matrix b.

Similarly

1.12 - dhpreact¶

Batch normalization¶

1.13 - dbngain, dbnraw, dbnbias¶

1.14 - dbndiff [1]¶

1.15 - dbnvar_inv¶

1.16 - dbnvar¶

1.17 - dbndiff2¶

Bessel's correction

• https://math.oxford.emory.edu/site/math117/besselCorrection/

It turns out that there are two ways of estimating variance of an array.

• One is the biased estimate, which is 1/n
• And the other one is the unbiased estimate, which is 1/n-1

1.18 - dbndiff [2]¶

1.19 - dbnmeani¶

1.20 - dhprebn [1]¶

1.21 - dhprebn - [2]¶

1.22 - dembcat, dW1, db1¶

is of the form:

Here, @ represents matrix multiplication, and $b^T$ is the transpose of matrix b.

This set of equations appears to be demonstrating the chain rule for matrix derivatives, specifically how the gradient of L with respect to a is related to the gradient with respect to d and the elements of matrix b.

Similarly

1.23 - demb¶

1.24 - dC¶

2. Derivatives of Softmax and Cross-Entropy Loss¶

# Linear layer 2
logits = h @ W2 + b2  # output layer

# cross entropy loss (same as F.cross_entropy(logits, Yb))
logit_maxes = logits.max(1, keepdim=True).values
norm_logits = logits - logit_maxes  # subtract max for numerical stability
counts = norm_logits.exp()
counts_sum = counts.sum(1, keepdims=True)
counts_sum_inv = counts_sum ** -1  # if I use (1.0 / counts_sum) instead then I can't get backprop to be bit exact...
probs = counts * counts_sum_inv
logprobs = probs.log()

loss = -logprobs[range(n), Yb].mean()

Rather than implementing Softmax, and Loss functions, we can use PyTorch's built-in functions to calculate the loss.

Backward pass¶

Softmax function:

• This converts logits to probabilities.

Cross-entropy loss:

• Where y is the index of the correct class.

Gradient of loss with respect to logits:

For i ≠ y (incorrect class):

For i = y (correct class): $$\frac{\partial \text{loss}}{\partial z_y} = \frac{\partial}{\partial z_y} \left[-\log \frac{e^{z_y}}{\sum_j e^{z_j}}\right] = P_y - 1$$

Derivation for i ≠ y: $$\begin{align*} \frac{\partial \text{loss}}{\partial z_i} &= -\frac{\sum_j e^{z_j}}{e^{z_y}} \cdot \frac{\partial}{\partial z_i}\left[\frac{e^{z_y}}{\sum_j e^{z_j}}\right] &= -\frac{\sum_j e^{z_j}}{e^{z_y}} \cdot \left[0 \cdot \frac{1}{\sum_j e^{z_j}} + e^{z_y} \cdot -\frac{e^{z_i}}{(\sum_j e^{z_j})^2}\right] &= \frac{e^{z_i}}{\sum_j e^{z_j}} = P_i \end{align*}$$

Derivation for i = y: $$\begin{align*} \frac{\partial \text{loss}}{\partial z_y} &= -\frac{\sum_j e^{z_j}}{e^{z_y}} \cdot \frac{\partial}{\partial z_y}\left[\frac{e^{z_y}}{\sum_j e^{z_j}}\right] &= -\frac{\sum_j e^{z_j}}{e^{z_y}} \cdot \left[\frac{e^{z_y}}{\sum_j e^{z_j}} + e^{z_y} \cdot -\frac{e^{z_y}}{(\sum_j e^{z_j})^2}\right] &= -\left[\frac{e^{z_y}}{\sum_j e^{z_j}} - \frac{(e^{z_y})^2}{(\sum_j e^{z_j})^2}\right] &= -\left[P_y - (P_y)^2\right] = P_y - 1 \end{align*}$$

These derivations show that the gradient with respect to each logit is simply the difference between the predicted probability and the true probability (which is 0 for incorrect classes and 1 for the correct class).

3. backprop through batchnorm but all in one go¶

To complete this challenge look at the mathematical expression of the output of batchnorm, take the derivative w.r.t. its input, simplify the expression, and just write it out.

Forward pass

# Linear layer 1
hprebn = embcat @ W1 + b1  # hidden layer pre-activation

# BatchNorm layer
bnmeani = 1 / n * hprebn.sum(0, keepdim=True)
bndiff = hprebn - bnmeani
bndiff2 = bndiff ** 2
bnvar = 1/(n - 1) * (bndiff2).sum(0, keepdim=True)  # note: Bessel's correction (dividing by n-1, not n)
bnvar_inv = (bnvar + 1e-5) ** -0.5
bnraw = bndiff * bnvar_inv
hpreact = bngain * bnraw + bnbias

# Non-linearity
h = torch.tanh(hpreact)  # hidden layer

Backward pass

4. Using custom backprop¶