LLMs From Scratch: Day 6

I am finally back! Astute readers might have noticed that it has been almost two weeks since I last posted about any progress on LLMs From Scratch, thanks in part to a week-long trip to the UK. Truthfully, I worked on the content for today gradually over that time and am only finishing it today. The majority of the work centered around implementing the AdamW optimizer from "Decoupled Weight Decay Regularization". The vast majority of the AdamW algorithm is also, unsurprisingly, present in the Adam optimizer from "Adam: A Method for Stochastic Optimization", which AdamW adds to. Optimizers -- like the tokenizers we saw earlier -- are a rich field on their own, so I didn't want to divert too far into the literature. But, the changes Adam introduces are important enough that the paper(s) were worth reading. The biggest mechanism that Adam introduced was the inclusion of the $m$ and $v$ moments into the optimization process. These moments are moving averages of the gradient and gradient squared, respectively, and are exponentially decayed using two parameters, $\beta_1$ and $\beta_2$. As a physicist, I would liken the moments to a "momentum" term and a "velocity" term in our optimizer. The first moment, $m$, establishes a moving average of the overall "direction" that the parameters are moving toward. This prevents any one batch of data from drastically change how the parameters are updated. The second moment, $v$, establishes a velocity term so that each parameter can have a regulated step size. If the gradient of one parameter is consistently large, we want the step size reduced so as to not overshoot the optimal value, and conversely, if a parameter has consistently small gradients, we want to increase the step size to make sure we converge in a reasonble time scale. The analogy is not perfect, in particular calling $v$ velocity-like, but I think it's close enough to illustrate a bit of the intuition. AdamW adds to Adam by including a weight decay parameter. Weight decay is a technique used to regularize model parameters, and like "Decoupled Weight Decay Regularization" notes, is equivalent to $L_2$ regularization in Stochastic Gradient Decent (SGD), but not in Adam, which necessitated the novel work. In AdamW, the basic idea of the weight decay is that you subtract a small fraction of a parameter's value from itself at each step, which prevents prevents parameters from growing too large.

Practically, the Adam and AdamW algorithms are not difficult to implement. However, one challenge I encountered when writing the code for it was how to efficiently apply AdamW to each parameter in my model. When I created the feed forward network on day 1, I used simple SGD, and stored each parameter in a model dictionary. At the time, I used jax.tree.map to pair up each parameter with its corresponding gradient because when I tried to manually loop through the dictionary, I got an error (something like pytree objects are immutable). I didn't fully appreciate just how powerful jax.tree.map was until I started implementing AdamW. Assuming our $m$, $v$, and $\nabla f$ all have the same topology as our model dictionary, we can use jax.tree.map to iterate through all the nested structures without having to worry about which ones are lists, dicts, tuples, etc. And we can ensure $m$ and $v$ are constructed correctly by using jax.tree.map to create a zero-like tensor for each of our parameters! Once I rediscovered jax.tree.map and understood how to leverage it, the code became almost trivial.

The last point to make is that training models requires a lot of memory, and looking at the AdamW code, it's clear why: When you perform SGD, you need to store the model parameters, activations, and gradients; so the memory footprint is about 2 times the model size, plus the activations. With Adam and AdamW, you need the parameters, activations, gradients, and both moments; so the memory footprint is 4 times the model size, plus activations. That means that using Adam and AdamW requires about 3x more memory relative to SGD for the same model (1 model-sized tensor vs 3, not counting the model itself). That almost all modern models are trained with AdamW, despite its cost, is a testament to its effectiveness!