Normalizing GPT on the unit-hypersphere (WITH CODE)
SourceHi, NGPT! The paper titled "Normalized Transformer with Representation Learning on the Hypersphere" is fascinating. I’m coming out of a bit of a hiatus to discuss this one with you all, and I have a very special gift along with it. Not only are we going to read and explain this paper in depth today, but we will also go through the code that I replicated from scratch and posted publicly for you to review. The link will be in the description. I have a small 2-million-parameter model trained on Tiny Shakespeare, which is quite a fun little read!
But what is this paper itself, right? Well, basically, it presents some very simple yet clever ideas from NVIDIA. They took the concept of significantly increased normalization by using cosine normalization, placing vectors on a unit hypersphere, which has been recognized for a while in computer vision as useful. They combined this with another idea from the in-context learning sphere that suggests transformers perform gradient descent on their input data to find their answers. Essentially, each layer represents a gradient descent step. It's a very impressive paper with huge implications in terms of actual speed. They used the exact same GPT code with relatively minor edits and achieved, I believe, a 4 to 20 times speedup. So, this one is going to be a good one.
If you're watching my channel, then you probably enjoy learning, which is why you should check out today's sponsor, Brilliant. If you haven't already heard of them, Brilliant is where you learn by doing, with thousands of interactive lessons in math, data analysis, programming, and AI. Their service is designed to be uniquely effective, with a first principles approach that helps you build understanding from the ground up. Each lesson is filled with hands-on problem solving that lets you play with concepts—a method that's six times more effective than watching lecture videos. They wanted me to highlight a specific course or two that you might be interested in.
And last time, I chose their course on how large language models work, of course. But really, all of the content on there is relevant to this channel, whether you're into data analytics, programming, or math. I know a lot of my viewers have a background in only one of those areas, if any at all. So, for example, if you need to get familiar with Python by learning essential coding elements and developing your mindset to think like a programmer, then go check out their general programming course. To try everything Brilliant has to offer for free for a full 30 days, visit brilliant.org/tunadorable or scan the QR code on screen, or you can click on the link in the description.
You'll also get 20% off an annual premium subscription. So, the key contributions here, right? First up is optimizing network parameters on the hypersphere. "Hypersphere" means a high-dimensional sphere. Optimizing them on the hypersphere means that the parameters actually stay on the surface of the hypersphere at all times.
Specifically, all of the vectors of length equal to the model dimension are what remain on the hypersphere, not those of length equal to the head dimension or anything like that. Interestingly, once we do this, keep in mind that it involves both the residual connections of the model and the actual weight matrices. Therefore, whereas in traditional transformers you typically apply layer normalization or RMS normalization only to the residual state, in this case we are applying normalization to both the residual state and all of our weight matrices. What this allows us to do is view matrix multiplications or matrix effect multiplications as dot products that represent cosine similarities.
So, whereas traditionally the magnitude—specifically cosine—has been used, the fact that they have undergone cosine normalization and been placed on the surface of a unit hypersphere means that any dot products are now simply cosine similarity. They are bounded between negative one and one, and suddenly you can interpret them as follows: if the number you get is negative one, that indicates very low similarity; zero means they are unrelated or orthogonal; and one means they are very similar. As I mentioned, this concept has been around for a while in the computer vision field. Last December, I reviewed a variety of papers—mostly in the form of YouTube Shorts and TikToks—where I discussed specifically all the studies that employed this cosine normalization and placed vectors onto the unit hypersphere.
All very interesting and showing good results, but they seem to be forgotten for a while, right? Well, NVIDIA has brought those back to attention. This is an example of one such paper that we’re not going to read today. I did want to show it because it has a good illustration in here, though it's not the actual main one.
The main one, I think, was a 2017 paper. I forget what it's called, but it has a very good display. Here is a three-dimensional hypersphere, also known as a regular sphere, right? When we perform classification on this hypersphere, ideally, you want the vector of every single item in your classification to be as far apart from the others as possible. What you end up with are zones that correspond to each of these vectors. So, here, there are six classification vectors, and we can effectively map out the sphere. For example, in the yellow region at the top, any output result vector of the model indicates that it was trying to give us p0, p0.
Not right? What it does is effectively provide us with very clean, certain, and solid regions for classification, as opposed to what usually happens in transformers with their own embedding vectors. Typically, this phenomenon is called anisotropy. I could be pronouncing that incorrectly—a-n-isotropy—where, if we attempt to draw a 3D coordinate plane, here is what our new model, the NVIDIA proposal, will do: it creates an actual sphere upon which all vectors will be constrained to lie.
Whereas in traditional transformers with RMS normalization, which does not force the vectors onto any specific mean, I won't get into layer normalization quite yet. Basically, what happens in traditional transformers is that you end up with all your vectors of your embedding space roughly clumped together, differing significantly in terms of magnitude. They'll clump into a certain part of the total representation space, which is very inefficient in terms of what’s actually fitting in there. I believe this kind of forces the model to require more numerical representation. Now, for simpler tokens and those more well-trained on classification tasks, what you end up with is actually a bit of isotropy, relatively speaking, but it's especially the rare tokens that clump very close together. The model tends to interpret them as "all the things we don’t quite understand," and simply places them somewhere, using induction heads on them. Now, all of those orange and reddish vectors represent what happens in a usual transformer. We really do not take advantage of the total angular space available, and the model has to interpret terms of magnitude, which results in inefficiency. In contrast, in this new cosine normalization regime introduced by NVIDIA, we find that all of our vectors are distributed evenly on the exact surface of a unit hypersphere, which is quite fascinating. The goal here is to determine whether it is easier to learn on the surface of the hypersphere than in any arbitrary region in space where tokens are usually clumped together. Previous findings in computer vision research have shown that yes, it is indeed super useful. Now, regarding the first introduction of cosine normalization, I would like to mention that I have a video from way back discussing this topic.
I want to say it was about 11 or 12 months ago, during the first month or two when I was trying to teach myself all this AI stuff. In it, I just rambled about a bunch of ideas I had, the main one being what I called the hypersphere layer. It's pretty funny; I just had never read a paper on co-signalinization before regarding the unit hypersphere.
And I was intuitively thinking it would be good to put vectors onto the... Oops, I'm not showing my actual screen. Sorry. The question was asked, "What change may be from orange to yellow?" Orange represents the old transformers that use RMS norm, like the LLaMA type, while yellow refers to the new Nvidia model.
So Nar had asked, "Isn't the volume of a hypersphere past 20-ish dimensions infinitesimal?" Yes, the volume is absurdly small, but again, we are not learning inside the volume of the hypersphere. Actually, we are not using the inside at all. We are learning on the surface of the hypersphere.
We are learning on the surface of the hypersphere, and the actual surface area is very, very large, no matter how small of a sphere you have. High dimensions simply mean very, very large surface area. There is a distinction between a clumped structure and a hypersphere, but what’s the mechanism? The mechanism lies in traditional transformers, or let's just say LAMA, which uses RMS normalization, right? In LAMA, RMS normalization does affect the variance, but it does not actually control the mean of the vector.
Therefore, it allows the magnitude of the vector, and consequently the length of the vector, to be any value. It can take on any length, correct? This is why those orange vectors are positioned far away. Cosine normalization does affect the centering of the actual hypersphere. To be clear, the RMS norm also places vectors onto the surface of a hypersphere with a radius equal to the square root of the model's dimension. The radius is equal to the square root of the actual vector length, but even when it places the vector on the hypersphere, it does not necessarily enforce the use of the same hypersphere centered at (0, 0, 0) with respect to the RMS norm, right?
Meaning, technically, the RMS norm places vectors onto hyperspheres of square radius \(d\), rather than unit hyperspheres. However, if each of these vectors is always centered differently, the RMS norm essentially has no effect. It just allows the vectors to move freely; each of those orange vectors lies on the surface of a different \(d\)-dimensional hypersphere, which is where the issue arises.
And because they are not the same, it might as well just not be on one at all, right? Whereas, with yellow, the cosine normalization is actually forced to sit on the surface of the unit hypersphere, which is the hypersphere with a radius of one, right? So there is only one hypersphere that any vector can be placed on once it goes through cosine normalization—all the exact same hypersphere. Our RMS norm allows for any hypersphere centered at any location. I actually didn't even draw that correctly either.
It's like it could literally be a hypersphere centered way out here or something—a sphere extending in that direction. Hence, it just doesn't really have the same effect. Layer normalization also centers at zero, sort of. We're not going to delve too far into that, though.
I will share a link in the description to a video I made that goes very in-depth on Layer Normalization. RMS and Layer Normalization also do other unusual things regarding the actual distribution on the sphere. I'm nodding. Yeah, thanks, Narhead. They also have other peculiar effects, such as how Layer Normalization forces these vectors onto a hyperplane. There's all this other stuff, and it's just not as ideal. Additionally, depending on the value of your small epsilon clamp, Layer Normalization may not necessarily project things onto the actual surface; they might remain further inside.
It all just wasn't quite as ideal. This cosine normalization is both cool in its simplicity and has shown to be very effective in past computer vision papers. It's like a better version of normalization, effectively. What I was saying a second ago, in this very old video from about 11 or 12 months ago, is that I actually predicted this. I had never learned about cosine normalization; I mean, I know some math, so I could infer it and do the calculations. I just never knew there was a name for it. I was essentially calling it a hypersphere layer.
And in this video, I was insisting that we need to put transformers on the unit hypersphere. I also have my template GPT repository and a little bit of unused code for co-synonymization that I wanted to use one day. So, I kind of preemptively saw this coming. After making this video, I did find the computer vision research about co-synonymization and read through it quite a bit. I probably have about a dozen or so videos on those computer vision papers, which are pretty old now.
But I am very excited to see NVIDIA actually go about using this and showing that it does, in fact, work. The JDen said they were wondering if there are locally Euclidean manifolds that are better than Rn since we could still perform gradient descent on them. Yes, this is just a subset, a very small and relatively speaking subset of Rn, obviously. Again, we are talking to JDen and discovering that these spheres and whatnot are a better manifold for representation than in classification in general, not just in language.
Now, the next big addition that NVIDIA made is something I cannot take credit for at all. I mean, I'm not taking credit for the previous point either, but I didn't anticipate it a year ago. They are now considering transformers as variable metric optimizers. The reference to the hypersphere is just a layover from the actual previous point. But basically, what this means is that you know how we typically train any model using gradient descent on its weights, right?
There has been some suspicion and research in the whole in-context learning field that what is actually happening with transformers is that each layer is essentially performing a gradient descent step, where the input prompt acts as all the conditions and the goal of the gradient descent is to find the actual next token, right? This suspicion has been known in the field for quite some time. What the authors of this paper do is propose that we go a bit out of our way to design NGPT to perform this even better, intentionally designing it as a variable metric optimizer from first principles. The optimizer part I just explained; the variable metric part refers to the idea that with each gradient descent step—or each layer—rather than simply moving in the direction of the gradients with its own magnitude, we adjust the length of the step. What is our eta, in this model?
And they actually add this variable metric, which allows the model to say, "Hey, if the attention mechanism and the MLP are finding gradients to move us towards our final output token, what if we add another parameter that basically just adjusts that gradient step and makes it usually smaller?" It's learnable, but it typically ends up being smaller. This means the model doesn't move too far at a given time. Now, you can think of it this way: if we have an input token, let's say right here, and there’s only one input token, then what the transformer layers do is find the gradients in the direction of, let’s say, this pink dot, which represents the output token. They identify that gradient, and then the variable metric part says, "Hey, don’t go too far." For every gradient step, this variable metric essentially indicates that the actual learning rate should be smaller, so let’s reduce that a bit.
Right? So I just cut that arrow in half there, and we have a smaller piece. Then, for the cosine normalization part, we actually end up instead of keeping this as our new point, which is off the unit hypersphere, what we do instead is we cosine normalize that dot and put it back onto the hypersphere. We repeat that multiple times until we eventually reach our final correct output token. This process turns a transformer into a machine that effectively has a specific vector location for every single input token on the unit hypersphere. All it does during the prediction process is move along the surface of that hypersphere towards the final predicted output token, using a kind of gradient descent-style mechanism that is repeatedly brought back onto the surface of the hypersphere through cosine normalization.
And what are the actual results of this? Well, a faster convergence. This normalized transformer reduces the number of training steps required to achieve the same accuracy by a factor of 4 to 20. Huge. So big, right?
Now, let's get to the actual specifics of this, and along with these specifics, I will be pulling up the actual code and showing you how I did each part. First up, we have our embedding matrices, right? So for every given token, we have to provide an actual vector of model dimension. First up, \(E_{\text{input}}\) is our embedding matrix, and \(E_{\text{output}}\) is our unembedding matrix.
And, of course, we are going to cosine normalize them. That's the big thing. Right here, we have the code for cosine normalization. It's pretty simple. You just calculate the actual norm of the vector and then divide the vector by its norm.
And now, suddenly, your output will be on the surface of the hypersphere. Then, divide the vector by its norm, and your output will be on the surface of the hypersphere with a radius of 1, right? Now, in their models—sometimes in transformers—we take the input embeddings and the output embedding matrix. Gnar had just asked; Gnar thought we were giving the model control over the magnitude of the learning step so it could infer from its experience where it thinks the data will be. It's a bit meta. However, we do not give the model control over the actual gradient descent process during training. That's not what's happening here.
I will show you later. We have a set of parameters that are metaphorically equivalent—if not exactly—to gradient descent in the model itself. We'll get to those later. However, to be clear, we are not actually giving the model control over its own learning rate during the training process. That's not what is happening here.
We'll get more into what's actually happening a bit later. Anyway, back to the embeddings, right? A lot of models pair the input and output embeddings using the exact same matrix. This is less common with larger models because, when you need to parallelize a model across different GPUs during training, that approach becomes less feasible. However, it still holds true for smaller models.
You can tie them to be equivalent. In their model, they did not tie them to be equivalent, which means that we have actual embeddings for each token at input time, right? And separate, different embeddings for each token at output time. What does this mean? Additionally, there are separate, different embeddings for each token at output time.
Meaning that any token can have different embeddings in the input and output sequences. Earlier, I drew this orange, pink, and whatever else was happening. It was a little bit misleading. If you were to tie the input and output embeddings, that would be the case. In reality, though, let's say you take two tokens; let's just say you take one token that may need to be repeated many times.
Like, for example, the space character. This red dot and this pink dot might both actually represent the space character. The input embedding vectors are still on a unit hypersphere, and the output embedding vectors are also still on a unit hypersphere, but now they exist on different hyperspheres, right? This means that the model takes the input vectors and processes them, but ends up comparing against a different hypersphere in the end, right? This isn't necessary.
You can tie them if you want to. I don't believe in the experiments they conduct with tying. I would be curious to see if the model actually gets better or worse after weight tying. I would be very interested in seeing that.
In most models, I believe that where weight tying is implemented, it tends to result in a performance hit, but it is worth the savings in actual parameter count. Here, I would be curious to see if it might even improve performance. I would be very tempted to wonder if that could be the case. As far as the code for our input and output embedding matrices, there is nothing actually new here at all. This is the same way you would set it up in NanoGPT or a similar framework.
And then again, for the output, there's nothing new here yet. We have not gotten to the actual edits of the output. Now, getting back into the idea of cosine similarity, right? The norms of embedding vectors in the original transformer, such as Llama 3, are unconstrained, which can lead to inaccurate similarity estimation. There is this idea that what we're trying to do here in these models can frequently be thought of as similarity estimation.
But remember, cosine normalization requires... oops, I forgot what I was saying. Anyway, what it's saying here is that we normalize the embedding vectors stored in the input and output embedding matrices after each step of the training algorithm. The keywords here are "after each step," right? So any given... oops, I have not been changing my screen, have I? Here’s what I was looking at.
Embedding vectors are stored in the input and output embedding matrices after each step of the training algorithm. That right there is the step where we took our initial gradients and put them back onto the hypersphere, right? That was this piece, moving it back down. The way that actually works is during the training algorithm, not in the actual model definition itself. So, here we have our training loop; nothing too complicated here, very similar to Karpathy's, whatever you've seen before for training algorithms.
But right after we update the parameters, we add this new piece that goes through and cosine-normalizes all of those weights, right? It runs with no gradient and goes through every single module, every single `torch.nn.Module`. If it is an `nn.Linear`, it first finds the correct dimension to normalize, which would be the model's embedding dimension, which I have named `dim`. Once it knows which dimension to normalize, it replaces the actual module's weights with a cosine-normalized version of those weights.
And that is done not just on our embedding input and output; it's done on every weight, every piece of weight throughout the entire model. However, now we encounter this issue. Once we place our embedding matrices onto a unit hypersphere, all of the actual embedding vectors are affected. Here's the problem: suddenly, our logits, our output logits, become dot products that are bounded by the range within which cosine similarity exists, specifically from negative 1 to 1.
Effectively, this means that the actual temperature of the probability distribution of the model is now very, very limited. However, we don't want to force the model into a certain level of entropy for the output with this negative 1 to 1 range. During training, we introduce a new scaling parameter, which we call \(s_z\), that simply scales the logits element-wise. We will get into the specifics of how \(s_z\) works later, as it actually has a couple of parts to it.
For now, just think of it as a single vector of length— I believe either the vocabulary or model dimension. It is of length vocabulary, and it is just entry-wise multiplied with the actual final output logits. We'll get into the actual scaling initialization part later. Next up is the format of the layers in this model, right? Usually, you take a given hidden state, normalize it, using RMS norm, like in Llama 3 style, right?
Then run it through your retention mechanism or your MLP to obtain the output of those mechanisms. After that, simply take that input and add the output to it. That's just a classic residual connection. Nothing new here, right?
What we're going to do in the Normalize Transformer is a bit different. First off, the hidden states are always normalized. You don't even have to worry about it. They are getting normalized so often that your input to the attention mechanism and your MLP just goes straight in as input. It's the output that then gets normalized, right?
Um, which, interpretability-wise, is very cool. I’d love to see a big analysis of what these attention mechanisms and MLPs actually output, now that those outputs are very easily interpretable as actual vectors in the same space as the model’s input embedding space, right? As far as your residual connection goes, though, it’s a little bit different. What we have here is that we actually take the output of your attention MLP and subtract it from the original hidden state. Then, we multiply that by alpha sub a or alpha sub m based on the tangent MLP.
And these kind of act as our metaphorical learning rates, right? There's the equivalent of η (eta), the Greek letter, in gradient descent. Then you take all of that. So h_a minus h and h_m minus m can be thought of as your actual gradient direction. And then these α (alphas) represent your actual gradient direction.
And then these alphas are your gradient step sizes, and you've effectively now taken your original input and adjusted it by your gradient step. Finally, normalize all of that, right? Going back to our iPad drawing earlier, the yellow represents our actual gradient direction, right? Then the black shows how we adjust our gradient direction by our actual step size. This little movement here, putting the output of the attention or MLP mechanism back onto the unit hypersphere, is our final cosine normalization.
This happens in both attention and MLP. It results in a slightly more complicated residual connection, which does have… What do we have here? Previously, there were just two normalizations in the original transformer. Now we have a total of four normalizations happening in each layer, so that's already an increase in the actual computation occurring. Additionally, we previously only had one set of element-wise addition.
Now we have entry-wise addition, entry-wise subtraction, and entry-wise multiplication. These operations are definitely more compute-intensive. I believe that late in the paper, it finally specifies how much more intensive all of this is. Additionally, in the regular transformer, nothing was normalized in the hidden states, so you have to apply RMS before you can obtain the actual final output logits and everything. However, we don't have to do that here because it is already implied in the final step of the layer itself.
Now, delving more into that alpha sub a and alpha sub m, the actual learning rates, the authors call them eigenlearning rates. This is not related to eigenvalues or eigenvectors per se; it probably has the same root meaning, though, as my guess. It comes from the German word "eigen," meaning "own," referring to the internal structure of the transformer.
So these are the transformer's own learning rates in its process, right? And it's just a single learnable parameter. Not too complicated. Nope, sorry, we're still on the iPad screen. Here we have what I was just talking about: the eigenlearning rates.
Now to jump into our actual layer module of the model rate, our attention mechanisms and our MLP are no different from the usual code I always use. However, we do add these new eigenlearning rate vectors during initialization. Again, I'm not going to explain the scale and initialization part quite yet, but basically they are just an nn.Parameter of the length of the configuration, or rather, the length of the model dimension, right? The same applies to both \(a\) and \(m\). Here, you can see the code version of the actual previously described digital connections. We take the input tensor \(h\) and run it through our tensor \(h\) of the actual previously described digital connections.
We take the input tensor \( h \), run it through our attention mechanism, then normalize it. After that, we subtract the original input tensor from the normalized tensor. That's our gradient, effectively, metaphorically. We then multiply by our learning rates—our learned learning rates, our learned metaphorical learning rate. Finally, that gives us our total adjusted gradient step, which we add to the original input. We then cosine normalize it and do the exact same process for your MLP. Also, I saw someone else replicating this paper online; I believe this is the repository, trana.27jk. I could not find any original repository for this paper, but I didn't go into detail on this one. It looks like they actually did include an implementation, although it is not quite finished, it seems.
It looks like they included scaling factors in their norm, which is pretty common to have scales after your norm. However, to be clear, in this paper, they do not use any element-wise scaling factors, whether they are multiplicative or an additive bias, after any normalization in this entire model. Hence, when defining cosine norm, I defined it as a function, not as an nn.Module with its own parameters or anything. Now, let's discuss the attention mechanism.
I said earlier that all weights get normalized, right? So our query, key, value, and output projection weights also get normalized along their embedding dimension. Now, suddenly, this is super interesting: the actual computed dot products with the input tensor and all these weight tensors can now be interpreted as cosine similarity between the unit norm vectors, bounded between negative one and one. Therefore, all attention matrices can be viewed as collections of normalized embedding vectors for comparison.
So, what we're saying here is that the weight tensors have a certain number of length-embedding dimension vectors. The output of the input times the weight tensor can be interpreted as a list of similarities between the input tensor and all the different vectors in the weight tensor. Now, our queries and keys are bounded within -1 to 1, which can be interpreted as cosine similarities. This boundedness is beneficial. However, they themselves are not located on the unit hypersphere. Furthermore, when you use rotary positional encodings, that further distorts the queries and keys.
In this paper, they propose to additionally normalize the queries and keys, ensuring that the dot product of every query and key is under control. This approach is somewhat unusual compared to what I might have expected. It's important to remember that the queries and keys are of size equal to the head dimension, not the model embedding dimension, right? They chose to normalize them here anyway, hopefully for the sake of, normalization is good, right? Additionally, they introduced another vector of trainable scaling factors, each of which is specific to the i-th head.
I, "i" as in there are i or some number of heads, right? Which itself is just a vector of length head dimension. So the way this works in code and in my own self-attention mechanism is as follows: First off, when you initialize the weights, there's actually no difference here at all because, once again, the cosine normalization of the weights happens in the training loop, not in the actual model definition itself. Then we do have our scaling factor to apply to our normalized queries and keys.
When we initialize that, I will define the scale and initialization piece later, but basically, it is a neural network parameter of size corresponding to the number of heads and head dimensions. So, each head of the attention mechanism has its own scale vector, which is just of length equal to the head dimension, representing the x-axis or the actual head size. Then, when you go to implement it, things look pretty normal for the majority of it. You even go through applying ROPE as usual.
However, before you perform your actual attention calculation, we take our queries and our keys and normalize them. We also multiply by our scale, right? Again, I'm not going to explain precisely how the scaling factor works until a bit later, but it is simply entry-wise multiplication right there. Then we proceed.
Apologies if there's noise in the background; I don't have much control over it. I didn't hear it, but it's good to see that. That's good.
Anyways, the attention mechanism itself still works as usual, but now we have logits that are bounded from negative 1 to 1 before applying any form of softmax. It is also worth noting that, in the original transformer, the softmax scaling factor is simply a single floating-point value of 1 divided by the square root of the model dimension. This effectively adjusts the variance of what is expected from the actual output of the attention heads, specifically the dot products of the queries and keys. In the normalized transformer, however, the expected variance of the dot product works in the opposite direction. It is 1 divided by the head dimension. To restore a variance of 1— which is always a goal here— instead of multiplying by 1 over the square root of the head dimension, we multiply by the square root of the head dimension. The code for this adjustment is obviously even simpler.
Instead of putting a negative here, just remove that negative, and your pre-existing code will still work. Next up is our multi-layer perceptron. Again, the weight matrices of these are normalized along the parts that are length and bedding dimension. And again, we introduce more and more scaling factors. So here we're going to call \( u \) the up projection and \( v \) the gate, right?
And so, after we normalize the weights for \( u \), we multiply by our scaling factor. Then, for the gate, we also multiply it by our weights for the gate to compute the initial part of our gate, which we then scale. Additionally, we need to multiply again by the square root of the model dimension. This step is essential to ensure that the Sigmoid Linear Unit (SiLU) non-linearity function remains non-linear. We are using a Swiglu, just like in LLaMA. If we were to do this naively without this floating point value, the values would cluster too closely around zero, preventing us from fully exploiting the non-linearity of the SiLU. The way to think about this is that for sufficiently large values, SiLU will approximate the Rectified Linear Unit (ReLU).
Fundamentally, it's pretty close, but at very small values, it is a bit different because it doesn't have the actual discontinuity or bend point. What happens here is if we do not multiply by our additional floating-point scaling value of the square root multiplied by the model dimension, we end up with a bunch of values that are very clumped right here. Remember, Silo is the pink one in this graph. If you have values that are clumped here, with nothing towards the left and nothing towards the right, this part of the graph—the function—is actually pretty close to just x divided by 2, meaning that you effectively have a linear function. You no longer have a non-linearity.
So, when you multiply by a square root model dimension, what it does is take these small values grouped together and spread them out along the entire range far enough that we can actually see our non-linearity again. If you zoom in enough on this non-linear silo function, the actual tiny range right here is quite close to a linear function and, therefore, not very useful. I hope this helps us improve our little MLP module. For the most part, it's quite similar; I named them "up" and "down," as opposed to "u" and "v" and whatever.
I wasn't very fond of their notation there, but it's a pretty regular MLP module. We do have our scaling factors. Again, I will explain the initializations and the scaling piece later. We just have \(s_u\) and \(s_v\) for now, right? And then our actual floating-point scale, which is just the square root of our input dimension.
And how we do this is we create our update and our gate, which we'll refer to as u and v. We then scale them element-wise by our update and our gate, u and v. After that, we scale them element-wise by our s sub u and s sub v, and also scale v by our floating-point value of our scale. Again, I will describe the actual initialization and scaling process later. Finally, we apply the regular Swiglu operations, use our non-linearity, perform element-wise multiplication, and complete with our downward projection. It’s important to note that all of these weight matrices will be cosine normalized at the end of the training loop, or at every single iteration of the training loop. Now, here's where I can explain what's going on with our scaling parameters, and why they are not just a single `nn.Parameter`, but also include these two separate floating-point values.
So, in normalized GPT, for any trainable vector of scaling parameters, such as \( s_a \), \( s_u \), \( s_v \), or \( s_{query} \) and \( s_{key} \), we use two scalars: an initialization version and a scale version. Not a version, sorry; they're just two scalars. I wasn't very fond of the notation throughout this, actually. I did not like that.
We were keeping that very close to the actual vectors. I would prefer to be more clear about when we are referring to a vector and when we are referring to just a floating-point number, right? I will clarify this: I replicated this part of the model and all parts of the model based on the paper, but this part of the model is based on just one paragraph. For a while, it was quite confusing to me. I just didn't love the way it was phrased.
It was hard to understand. I believe what I have in my code is correct, but please provide accurate information. Let me know in the YouTube comments or on GitHub with a push or by raising an issue. When initializing \( s_a \), and remember \( s_a \) is... actually, who cares which one it is, but when initializing it as a trainable parameter, its initial value is set to scale one, right? However, it's a floating-point number.
So every single entry in \( s_a \) at initialization is set to the floating-point value of whatever \( f_a \) scale is. However, during the forward pass, we restore its actual value by multiplying \( s_{a \text{ init}} \) by \( \frac{1}{s_a \text{ scale}} \). This allows us to control the effective learning rate for \( s_a \) by adjusting \( s_a \text{ scale} \) while keeping the global learning rate unchanged. Essentially, what they're doing here is saying that the whole model has some effective learning rate. Usually, I call it eta, which you define in the model's training loop, right?
What they want to do for the scaling parameters is try to give them different effective learning rates compared to the rest of the model. They achieve this through a simple method of multiplying by floating-point numbers. For example, setting \( s_{a\_net} \) equal to 1 and \( s_{a\_scale} \) equal to \( \frac{1}{\sqrt{\text{model dimension}}} \) ensures that this parameter is updated with the same effective learning rate as other normalized parameters in the network. They are simply adjusting these values. A more thorough explanation of why this is necessary can be found in Section 2.5.
It goes through the Adam algorithm and explains what is actually happening with step sizes. Essentially, the magnitude of the parameter affects the step size, which is quantified in terms of how many steps it takes to double a given value, such as 10 versus 1,000 steps. For instance, if we have a given parameter equal to 0.01 and the global learning rate is 0.001, it would take 10 steps for that parameter to double. However, if that parameter is set to 1.0, then it suddenly takes 1,000 steps to double, using the same learning rate. The main idea here is that we do not want to initialize all of these scaling parameters (n and dot parameters) at 1.0. Instead, we want to keep them closer to the range of the actual cosine-normalized vectors and, therefore, learn within that same range.
And that range is more so in the area of 0.01, right? Now, the final big edit is the removal of weight decay and learning rate warmup. Of course, you remove weight decay because that would actually pull all of the parameters off of the unit hyperspheres and in a seemingly random direction relative to the actual hypersphere surface. Right?
Whereas we no longer need that, and it no longer helps, it would mess things up rather than assist. In general, the model is so well normalized that we don't need it as much anyway. There's no issue of weight exploding because, again, the normalization forces all vectors onto the hypersphere. Finally, we will also remove the learning rate warmup.
A similar reason is that the model is so well normalized that it should not need it. Normally, learning rate warm-up is used to allow for a higher maximum learning rate for as long as possible, but you'll see later that this model is learning so fast that it can use a different linear rate anyway. They trained both the baseline transformer GPT and the normalized transformer NGPT with half a billion and one billion parameter models. Just look at this graph. They didn’t even use as big a budget.
GPT, with a 200,000 iteration budget, I believe it's the iterations that matter. Then there’s NGPT with a 20,000 iteration budget, out of the $200,000 iteration budget. Now, you'll see, first of all, just compare red to black. That's what we're looking for: a lower curve. It's what you want in every experiment of this sort.
The lower curve is fantastic. What's super interesting here is that not only is this a very apples-to-apples comparison with a good gap and a lower curve, it's also performing better. It plateaus at a lower point. More interestingly, when you provide the same model with significantly fewer training iterations, the effective learning rate throughout training is actually higher. As a result, it will decrease even further.
This is very, very cool. If you wanted to, you could take this pink line and keep training it with that different learning rate. After 20,000 iterations, NGBT achieves the same validation loss that NGBT reaches only after 200,000 iterations: approximately 400 billion tokens, demonstrating a ten times speedup in terms of iterations and tokens used. I do agree.
I do think that it is, though, the first time seeing this graph is a bit confusing. I would have liked to still see a GPT with a 20,000 iteration budget. I believe the reason it wasn't included is that such a learning rate setup would have resulted in regular GPT at a 20,000 budget with that different learning rate being just chaotic, not even learning well, right? My impression is that it would have been way up here or something. It wouldn't have been comparable, which is why I think they didn't include it.
If that's incorrect, I would appreciate inclusion, but I think they excluded it for a good reason. Here we have some more charts. The left shows a thousand-token length context window, then four thousand in the middle, and eight thousand tokens on the right. What we see here is that the training of the half-billion and billion-parameter NGBT models is about four times, ten times, and twenty times faster in terms of tokens at the different context lengths.
This is crazy because, as your context length increases, the benefits of NGPT actually increase, right? They don't cap out or diminish, at least as of yet, right? When you have a 1,000-token context window, you get a four times training speed increase. A 4,000-token context length window results in a ten times training speed increase, and an 8,000-token context length window provides a twenty times training speed increase.
I assume this has to hit diminishing marginal returns eventually, either with model size or with context length. It makes no sense to keep getting increasing returns on that speeding train. That would mean that the biggest model ever actually trains faster than a smaller model. That would be wild. So it must hit diminishing returns eventually. But still, this is amazing.
Scaling is progressing so far. Here's one issue: while the time per step of the GPT model is higher—about 80% for the 4,000 context length model and 60% for the 8,000 context length model—this can be reduced after customization. Currently, their code is about 60% to 80% slower because it has not been optimized and lacks fused kernels, among other factors. It's noticeable that when you increase the context length, the overhead becomes smaller. The overhead is less significant for larger networks, right?
So when you get to actual GPT-4 level models, or anything similar, the overhead and time per forward step might be more minimal, and you hopefully still get a training speed increase like this. Now, we're going to do a little inspection of the network parameters and what they look like. The distribution of pairwise dot products between embeddings indicates that even in NGBT, embeddings are not uniformly distributed across the hypersphere. I showed you this paper earlier. This example illustrates a uniform distribution of the actual vectors of interest across the hypersphere.
And I said earlier how transformers tend to have anisotropy, meaning very uneven layouts, so they all get clumped into one area, right? It turns out that even in NGPT, they don't actually all perfectly clump evenly; instead, they do form clusters. This is kind of to be expected, in my opinion. An even distribution of vectors would only make sense in a perfect world if you had classifications—tokens, effectively—that are all completely different and have no relation to each other whatsoever. However, the reality of language is that so many tokens share significant semantic meaning, so you would still expect clustering.
What I want answered, and I don't believe they addressed it here, is whether we still roughly use the whole sphere. Does it still seem pretty wide? It seems to be semantic clusters as opposed to everything being on a single cluster, right? Now, an important contribution of this work is the decoupling of predictions made by the attention and MLP blocks from their impact on the hidden state. These contributions are controlled by the eigenlearning rates, denoted as α and μ. Their interpretation is straightforward.
If a sub a_i for an embedding dimension i is 0.2, then the update follows just that; it is effectively like a metaphorical learning rate, right? Thus, they directly quantify the contribution of the attention mechanism into the actual current residual state vector. Notably, the network learns to take only modest steps in the order of 20% to 30%, so 0.2 to 0.3, in the direction suggested by the attention tension and the MLP. The average magnitude of a sub a decreases from 0.25 in the half-billion network to 0.2 in the one-billion network. Meanwhile, a sub m decreases from 0.37 to 0.32.
The introduced scaling factors seem to compensate for the removal of magnitude information when normalizing matrices and embedding. In the previous RMS norm, we had those crazy vector lengths of all sizes and different magnitudes. Here, it seems like these scaling vectors hopefully compensate for that, and it appears that they do. An important finding is that having fixed, non-learnable values for all those scaling vectors and a single global learning value for the other one, \( s_c \), which I believe was the final one of the output logits, leads to only a slight degradation in accuracy.
You could trim down a bit of the compute here. Here is a cool one: Appendix 8.7 demonstrates that NGPT can handle longer contexts without requiring any modifications to ROPE. Let's skim down to that real quick. Check this out.
So, red is GPT, and black is NGPT. Different shapes are trained on different context lengths. With GPT, the moment it exceeds its training context length, it just falls apart. The perplexity shoots up and off the graph, reaching a very high level. Whereas with NGPT, look at this—the circle, right?
It was trained on a 1,000 context length. It keeps the same perplexity at a 2,000 context length, barely diminishing at 4K and 8K, and then not really suffering until a 16,000 context length. This is crazy. It basically means that if this becomes the norm, the whole field of length extrapolation, the entire sub-discipline, becomes useless. They may have just solved it, right? If this is actually what it claims to be, I'm realizing that earlier, when I was adjusting the scale parameters, I did skip over them.
Let's go find some that are a bit clearer in the way they are set up. Here we go. For our query key scale, remember that it's an `nn.Parameter`. The number of vectors is equal to the number of heads, and each one has a length corresponding to the head dimension.
Basically, when we initialize it, we create a `torch.ones`, right? It's just a bunch of ones, multiplied by the scale factor. The \( s_{qk} \) scale, which is defined here as just 1 divided by the square root of the model dimension. So, we are dividing by the model dimension, or rather, dividing by the square root of the model dimension. Then, when it comes to actual inference time, we also have the initialization, which is just 1.
When it comes to actual inference time, what we do is take our actual vectors and multiply them by the initialization divided by the scale, right? In this case, this ends up being, I believe, just multiplying by the square root of the head dimension or the model dimension, right? So, that takes our actual parameters, which are very small, and makes them larger, I believe, in this case, to a scale that is actually useful at this level. However, because the parameters are small, they receive a different step size and a different learning rate, allowing them to change at the same rate as the rest of the model. I would be very curious to see this technique used on the linear affine and bias of a given layer norm or RMS norm. I believe that not the bias, I guess, but the affine of a given layer of RMS norm itself is initialized to `torch.ones`. I'd be very curious to see how that changes its effectiveness in the future, so there is more to discuss regarding the normalization from the previous discussion.
There is more to discuss. The normalization in the previous equations and the tension mechanism closely resemble the query-key normalization from a prior paper that is, I believe, used relatively frequently. In NGPT, this process can be viewed as restoring the queries and keys of the ith head to a hypersphere of dimension corresponding to the head dimension. The normalization itself, the coastal normalization, places those queries and keys onto the unit hypersphere. However, we then scale them by that \( s_{quadrature} \). The custom normalization also puts those queries and keys onto the unit hypersphere.
But then, of course, we scale them by that f subquery key. Since the input hidden states and the embedding dimensions of the weights of the queries and keys are already normalized, the norms of the queries and keys are also comparable, making their normalization potentially unnecessary. We explored this scenario, also known as omitting the normalization of queries and keys, in our ablation study. We attempted to remove that entire section to see what would happen if we deleted it.
And the results showed only a slight degradation. The performance drop was comparable to the computational time savings per step; therefore, removing the normalization is a viable option. So, if you wanted to replicate my code, you could simply delete these three lines, and it would likely perform similarly, at least well enough given the speedup you gain from removing it. Now, here's a significant point: we believe we are the first to decouple the eigenlearning rates from the rest of the network, recognizing them as trainable parameters.
Saying is back in the actual layer rate how we have these different scaling parameters to affect the actual step made by the attention MLP. So, right here we have \(a_{\text{sub}}\) and \(a_{\text{m}}\); what they're saying is that they have effectively decoupled the size of the attention MLP step from the attention MLP itself. In a regular transformer, the attention mechanism and MLP within their own calculations have to determine how large of an actual step to make on the residual vector. However, now that has been decoupled, and they are just learnable parameters. Now, this has both benefits and costs. The benefit is that this is computation that no longer needs to be done by the attention MLP mechanisms, allowing them to focus on other tasks instead.
However, negative; that means that these steps are potentially negative. These steps are no longer input-dependent in the attentional mechanisms. If we were to do a regular GPT, this idea of how large of a set to create is, in a way, at least potentially input-dependent. So, I would love to see it. It might not be worth the computation, but I would love to consider making these parameters themselves input-dependent, kind of a Mamba style—maybe just a simple linear projection or something. Although that involves a lot of parameters (it's model dimension by model dimension for each of these), I would be very, very curious to see if deciding its own step size makes any sense while still keeping it separate from the actual attention mechanism and MLP.
My guess is that it's not worth it. That's my guess, but I'm not quite confident there. All right, I'm scrolling through some of the appendices now. I didn't mention earlier that we actually define these eigenlearning rates, a sub a and a sub m, as positive. So in the training loop, not only do I enforce cosine normalization, but I also enforce the absolute value of the parameters after our gradient step. However, the authors do mention in this paper that this is for interpretability or convenience, but it doesn't necessarily have to happen. I won't go into many details right now, but it is another potential step that could be removed from my actual implementation and would probably still end up being just fine. If you're curious, here are the actual parameters of the two models they used. Again, no wait.
Decay is the interesting factor here, along with the absence of warm-up steps. To elaborate on the time cost per step issue, the time cost per step for NGPT is 80% higher with a 4k context length and 60% higher with an 8k context length. This overhead is not solely due to NGPT having six normalization steps; two of them are applied for the queries and keys, while four are applied in the actual residual connection.
Usually, a residual connection only has two, but also because of NGVS, these normalizations are not yet fully optimized. Unlike GPT, where normalization layers are fused with other operations, training on larger networks is expected to further reduce the performance gap. Therefore, the performance might decrease from 80% to 60%, or even drop to around 40—who knows? The number of layers in this set of normalizations increases only modestly with the number of network parameters. So, yes, I would personally prefer no normalization of the queries and keys. I bet that also removes part of the issue with—um, does it? Never mind.
They do go into detail on their actual parameter initializations, which has a very small deviation from regular GPT. If we go up to my code here, I've got this init weights function or method where the only difference between this and what you would see in Karpathy's NanoGPT is that the standard deviation of the weights being initialized—usually around 0.02 for back since GPT two days, as far as where it's still a thing. Instead, we use the square root of the model dimension, which is quite a bit wider. However, it really doesn't matter at all because these are the initializations for sure, but we end up cosine normalizing all of the weights anyway.
So, it doesn't really matter per se; this does affect the initial distribution and how the values cluster around the vectors on the actual hypersphere. However, I believe that this will be overturned quite quickly in the end. So, it's not the end of the world to avoid doing this. On another note, some specific layers have unusual scaling. Usually, at the end of the MLP and the attention mechanism, we scale the final layers—specifically, the final weight matrices—differently.
I believe the same. From the original. Oh yeah, it's the same as the original. Never mind, it's actually the same as the original, not just GPT-2, but also the original "Attention Is All You Need" paper. But yeah, it says the initialization is not important because it gets normalized afterwards.
Learning Rates. Optimal learning rates for GPT and NGPT tend to be similar for the same values of validation loss. Longer runs are usually associated with the possibility of achieving lower values of validation loss, and this typically requires lower initial learning rates. So, yes, use a lower learning rate here, and it still learns faster. They're pretty cool.
We already went through the length extrapolation ability earlier, so I have already walked you through most of the code. The rest of it is pretty similar. I do have some tests to ensure that we end up with correctly normalized, cosine-normalized, and absolute-valued weights in certain areas. I also have some displays to showcase if I can fix them, it looks like.
So, here is a little output from the model. Interestingly, to make this perform well, I trained a 2 million parameter model, by the way. It's pretty small, but it's trained on Tiny Shakespeare, so no big deal. Here's the actual layout.
When I gave it the classic Romeo or Romeo line, right? Honestly, most of the time when I train a model to size on Tiny Shakespeare, it's not even smart enough to know that Romeo should come later, like he should say something at some points. So, I thought this was pretty good for a model this tiny off such a small dataset. That was cool. It did seem to also train in my practice runs way faster than my regular GPT comparison version.
And then here we have some statistics on the actual parameters of my 2 million parameter and GPT for all the scaling. So right here we have sub A; this represents the eigenlearning rates on the attention mechanism. You can see that the actual mean increases as the layers progress, and the standard deviation does as well. A similar finding can be observed in sub M; both metrics increase as the layers progress.
And then the scaling parameter for queries and keys. Again, this is a bit different. The mean seems to go down, while the standard deviation still goes up. I'm just curious as to why any of this would actually apply.
It's pretty cool, though. We end up with, I believe I may have messed up if this is showing ones as the actual parameters. I'm not quite sure, though, but yeah, I have a couple of things being displayed. The weird thing, mainly in the output, is that I had to use such a crazy low temperature for this to work well. If you use a regular temperature like 1.0 or 0.7 or something, it's just absolute gobbledygook.
I had to use values like 0.1, 0.01, and 0.001 to get really good outputs from this model. I'm wondering why that is. In the future, I'll explore more. But yeah, I believe that's it. Check out this other person's application as well.
It doesn't look like they actually finished it, I don't think. But this is the only other version I could find. As far as I know, the authors have not open-sourced "Threat of Mutation." I believe I have a bug fixer to implement in the actual layer at some point before I post the video, and I also have to fix the README. But yeah, I guess that’s about it for today. Hope you all enjoyed! And yeah, cue the outro. Don’t leave yet; if you got this far, you're going to love this video or this playlist. Don't forget to like, subscribe, comment, and hit the bell—all the YouTube things! Join the Discord server, follow me on Twitter, connect on LinkedIn, and consider supporting me monthly either through Patreon or by hitting that YouTube join button down below. For as low as $1 per month, you can get early access to my videos. Or, if you're a one-time payment kind of person, hit me up on Venmo.
And if you want to borrow my brain for a bit, consider booking a paid consultation video call. All of those links are in my link tree. And, yeah, that’s the end of the video.