Backpropagating through batch norm
Are we having fun with matrix calculus yet? 🙂
My wife is noticing that I am spending a lot of time “doing math homework” and asked me, “is this applicable to anything?”. Well the answer is that its probably not really necessary in the grand theme of things. I gave a long answer about how modern ML is all about apply these small tricks like batch normalization to neural nets, so really understanding how this math works in a deep way is the key to understanding modern ML.
I think the reality though is just that I am really enjoying these exercises. In comparison, I am blitzing through the coding portions since I feel like they are too easy. Working through these matrix calculus problems though is really scratching that itch for me.
Alright so still following along this wonderful video.
Building makemore Part 4: Becoming a Backprop Ninja
The exercise here is to manually backpropagate through the batch norm layer. Here is Karpathy’s solution. As usual, he goes through the math pretty quickly and does not explain every single step, so its easy to get confused.
First we got the problem setup. Below is the computation graph. The numbers 1 2 3 4 are just the suggest intermediate steps he suggests that we follow.
The key thing to notice is that some of the elements like x and \hat{x} and y are matrices, while others like µ and σ^2 and β and Γ are just single values. This is significant when it comes to chain rule because this means that there are actually many arrows that connect the single values and matrices. When this happens, we need to make sure we sum up all the derivatives instead of just doing it once. You will see below.
Step 1 is pretty easy. Nothing much to say here.
Step 2 is a little bit more interesting. Like I mentioned above, we need to make sure the sum all of the derivatves since σ^2 “fans out” to all the elements of \hat{x}. By now I am pretty familiar with the quotient / chain rule, so can breeze through some of these steps.
Step 3 is a bit hairy. We got two sub derivatives we need to solve for. The first one was pretty straight forward. Quotient rule, easy.
But for some reason I got really tripped up at the second one (maybe its the crying baby the in background), not understanding how we got that negative 2. Isn’t the derivative of x^2 just 2x? Surely this was a typo by Karpathy right??? I even started just googling around to see if I can get some inspiration from other people’s solution to this problem. I couldn’t find any detailed writeups for Karpathy’s course specifically (probably because everyone else finds this so easy) but I did find another person’s notes on the exact same problem.
Deriving the Gradient for the Backward Pass of Batch Normalization
I stared at this line in disbelief for longer than I care to admit. How is this possible!?
I even asked ChatGPT for help (even though I don’t super trust ChatGPT with symbolic math problems).
It was only when I saw this, that I realized my simple mistake. Of course, I forgot to apply the chain rule 🤦. I wrote this lil note for future self in case I get confused about this basic thing again.
Okay step 4 was surprisingly easy, though I did take me a bit of time to write it out. The only tricky thing here is the usage of j in the summation.
Let’s zoom in to that section. Notice that we all of a sudden went from using i to j as the variable in the summation. At first glance it seems like the symbol we decide to use is arbitrary and so maybe Karpathy just decided to do it for fun. But actually he does explain that he switches it up to avoid conflicting with the i in x_i. The i in x_i refers to the ith element of the matrix x whereas the j is the local variable used in the summation — and these are different things!
This distinction is important to understand because in the final line of the derivation here, we use it remove the summation from the expression. As Jonathan on discord explains, when we unroll the sum, only one of the terms where i=j, will matter in our derivative which is with respect to x_i, all the other terms where i{=}\mathllap{/\,}j will have a derivative of 0, so can be ignored.
Hope my lil write up here helps someone!
My final parting thoughts is that I am starting to question my decision to use https://excalidraw.com/. On the one hand its quite nice to be able to easily add colored boxes and position random notes just like pen and paper. On the other hand, its kinda annoying how I have to “draw out” simple things like σ^2 by typing out 2 in a smaller font and just positioning it near the σ. Maybe I will try another tool in the next exercise.