Fun with low-rank vocab matrices (and a bonus test loss reduction?)
I was nerdsniped! On the HF discussion page for one of the
models I created for my previous post,
AndrewThompson1233
asked if I'd considered trying out factorised embeddings
-- something he's using for his model, Maba, and which was previously used in some other models, like
ALBERT.
It's a really nifty idea; you use a similar trick to LoRA as a way of reducing the number of parameters used for your embeddings and your output head. And for small models, those can be a disproportionate number of the total.
For example, with my 163 million parameter GPT-2 small-style models, the embeddings are about 39 million of them; the output head is the same size, so with those two taken together, that's half the model just getting stuff in and out rather than actually doing the thinking. Even if you use weight tying, like the original GPT-2 models did, you'll wind up spending 30% of your "budget" on the single matrix shared between embeddings and the output head.
Similarly, while larger models have a smaller percentage spent that way, even medium-sized MoE models can wind up spending a lot of the active parameters on embeddings; I calculated that for Qwen 3.6 35B MoE, with 3B active parameters, there were about 1B of them used across both the input embeddings and the output head.
So anything that can reduce that -- so long as it doesn't come at a high cost in terms of the model's capabilities -- is worth considering. You save on your parameter count, which either means smaller models and potentially quicker training, or allows you to "invest" the savings in more thinking parameters -- that is, a wider network with a higher embedding dimensionality, or a deeper one with more layers.
Andrew reported that the impact of using this trick was minimal:
At rank 128 on a 50k vocabulary, the penalty on validation cross-entropy is negligible (typically within +0.02 to +0.04 loss, or <0.5 perplexity delta).
Now, that suggests that he's getting a loss of around 2.5 (that being the point where a loss increase of 0.04 implies a perplexity change of 0.5), which is much lower than I typically get with my GPT-2 architecture (3.5 is more like it for me), but still, perhaps I'd get good numbers too? I decided to train a few models to see what happened.
The results were interesting! I found that using this trick on both the embeddings and the output head caused a somewhat larger increase in loss than Andrew described -- about 0.07 in my first test, 0.10 in the second, and 0.14 in the third. Those are relatively large, though potentially recoverable from further training or larger models.
But more surprisingly, I found that using the trick on the input embeddings only seemed to reduce test loss. Before I'd started, I'd expected it to be less harmful on the input side than it was on the output side, but seeing an improvement was certainly unexpected. Getting loss down by reducing your number of parameters is not what normally happens!
So, it's definitely worth digging in a bit. Let's start with the theory: what are these factorised embeddings, and how do they work? Like LoRA, they rely on low-rank factors, so firstly I'll define those.
Why do OpenAI's GPT-2 weights beat mine? Part five: data quality
When I finished learning how to build an LLM from scratch, I was left with a mystery: my own models were not as good as OpenAI's original GPT-2 models, despite being based on the same architecture.
My models all had 163M parameters, and followed the design from Sebastian Raschka's book "Build a Large Language Model (from Scratch)". That meant that they were pretty much the same as the setup for the OpenAI GPT-2 "small" instance, except that they did not use weight-tying or bias on the QKV matrices. Weight-tying means that you re-use the initial embedding matrix as the output head at the end, and using it means that GPT-2 small saved quite a few parameters -- it was 124M rather than 163M -- at, at least in my own experiments, a cost in quality; similarly, while I found that QKV bias made a tiny improvement in loss terms, I'd felt it was likely within the noise.
But GPT-2 small consistently beat my models on an instruction fine-tuning (IFT) task -- also adapted from Raschka's book. That test fine-tunes the model on a subset of the Alpaca dataset, until validation loss starts rising, and then runs a test set through the resulting model. The responses to the test set questions are stored, and then I run all of the responses from all of the models under test past GPT 5.5 in one go to get an aggregate score; more details here. GPT-2 small always did better than any of my models on this.
Additionally, it did surprisingly well on a simpler eval -- one that just measured the cross entropy loss it got on a test set. It scored close to my own best models, and better than many of them. What made this result particularly interesting was that the test set in question was a split of my own training data; my models would not have seen it when training (at least, in theory), but it seems likely that it would be much more similar to their own training data than it was to OpenAI's.
I've checked two things while probing this mystery:
- It seems very likely that the GPT-2 models were overtrained by modern standards; would overtraining my own models get them closer? It turned out that no, it probably didn't help with the IFT eval (though there might have been some signal there). It did help quite a lot with the test loss eval, though.
- The way I was handling dropout in the IFT test might have been unduly benefiting some models while working against others. I decided to standardise on not using dropout during this eval, as (counter-intuitively for me) it seemed to harm the results of most models, even those that had been pre-trained with dropout. In particular, the OpenAI weights were harmed by using dropout, and making a change that benefited them (along with some of my own models) seemed the most conservative approach to take in investigating this.
The next thing I wanted to look into was the training data.
The exact dataset that the various GPT-2 models were trained on has never been released; all we know about it is from the paper, where they say:
[W]e created a new web scrape which emphasizes document quality. To do this we only scraped web pages which have been curated/filtered by humans. Manually filtering a full web scrape would be exceptionally expensive so as a starting point, we scraped all outbound links from Reddit, a social media platform, which received at least 3 karma. This can be thought of as a heuristic indicator for whether other users found the link interesting, educational, or just funny.
They called it "WebText". There is an OpenWebText that tries to replicate it, but although they tried to follow the same procedure as the original, there's no guarantee that it is all that similar.
By comparison, I'd normally been training against FineWeb. While this is a general web-scraping dataset, without the "curation" provided by using only stuff that was linked from upvoted Reddit posts, it has been refined to remove any obvious junk. I had felt that it was pretty much equivalent.
But what if I were wrong about that? I decided to see if I could get better models by using better data.
Extending Raschka's GPT-2: an MoE trained from scratch on an RTX 3090
Mixture-of-experts models are really nifty. You get inference speed close to a small model's, with a lot of the smarts and knowledge of a large one. While they use as much memory as an equivalently-sized dense (non-MoE) model, they're much faster. The frontier labs don't publish their architectures, but Claude and ChatGPT are widely rumoured to be MoEs these days -- and certainly many large open-weights models like DeepSeek and Kimi K3 are.
In this post, I'll show you how I added MoE support to the GPT-2-style code from Sebastian Raschka's book "Build a Large Language Model (from Scratch)", then used that to train a 446M-parameter model with 220M active parameters from scratch on my RTX 3090 -- essentially, GPT-2 small with 6 experts, 2 active per token. I wanted to understand how MoEs work, and as always, felt that the best way to do that is to build one (and then to write it up like this). Hopefully because it's all fresh in my mind as I write this, I should be able to explain things clearly for others who've finished Raschka's book.
The training run took just less than eight days, and the resulting model got a better loss on my test set than any of the other models I've trained so far -- which was a good thing, given that it took four times longer to train! It was also better than the original OpenAI GPT-2 small (124M parameters), but not as good as GPT-2 medium (345M parameters), although it was close. That second result was interesting, as my model had more total parameters than GPT-2 medium, but fewer active ones.
On an instruction fine-tuning test, it did better than any of my other models, but worse than both OpenAI models (why OpenAI's models are so good on that particular test is a mystery I'm digging into separately).
I'll go into those numbers in more depth later on. Firstly, though, I'll run through the code that I needed to add to GPT-2 to make this work -- not just the code for the experts themselves, but the additional code for training it. With MoEs you can't just train to minimise loss on your training set -- you need to add on an "auxiliary" loss to make sure they actually use all of their experts, and that was where it became most interesting.
Before we get into the weeds, though, let's start off with the basics; how do MoEs work in Transformers-based LLMs?
Putting my JAX-trained models on the Hugging Face Hub
I hadn't uploaded the models that I trained using JAX to the Hugging Face Hub because
Transformers has been PyTorch-only since version 5
(though they say they're working to add interoperability with JAX in the future), so it
would have been tough to get them working natively with AutoModelForCausalLM and the like.
But then it dawned on me that I'd already written a conversion script that could take my JAX safetensors files and convert them into ones compatible with my PyTorch code. It's actually those converted models that I use for my evals -- so I could use my existing PyTorch script to upload them.
So, I've now uploaded PyTorch-compatible versions of all of my JAX-trained models:
"Writing an LLM from scratch, part 34b -- from bigrams to GPT-2, one component at a time (in JAX)"
gpjt/jax-no-mha-bias-no-dropout-- the first full LLM trained in the post, in the "Adding LayerNorm" section.gpjt/jax-no-mha-bias-with-dropout-- the second full LLM trained in the post, in the "Dropout" section.gpjt/jax-with-mha-bias-no-dropout-- the third full LLM trained in the post, in the "Adding bias to the MHA output projections" section.
"Why do OpenAI's GPT-2 weights beat mine? Part three: testing overtraining"
gpjt/jax-with-mha-bias-no-dropout-extended-- the single-epoch, double-Chinchilla-tokens model.gpjt/jax-with-mha-bias-no-dropout-2-epoch-- the model trained on two epochs over the Chinchilla-optimal number of tokens.
"A quick(ish) Chinchilla check"
gpjt/jax-with-mha-bias-larger-chinchilla-1-- theslightly-largermodel.gpjt/jax-with-mha-bias-larger-chinchilla-2-- theslightly-smallermodel.
I've also added links to the posts in question.
Why do OpenAI's GPT-2 weights beat mine? Part four: digging into dropout
I'm still digging into a mystery about the models I've been training; although an increasing number of them beat the OpenAI GPT-2 small weights on the narrow technical measure of the loss they get on a test set, they're not as good at an instruction-fine-tuning test.
While reading about MoE models, I came across this paragraph in the Switch Transformers paper:
Our paper considers the common NLP approach of pre-training on a large corpus followed by fine-tuning on smaller downstream tasks such as summarization or question answering. One issue that naturally arises is overfitting since many fine-tuning tasks have very few examples. During fine-tuning of standard Transformers, Raffel et al. (2019) use dropout (Srivastava et al., 2014) at each layer to prevent overfitting.
So far, when running my IFT test, I'd been aiming to use the same dropout setting for the fine tune as the model concerned had used in its original pre-training. That was just because it seemed natural.
But the goal of dropout is to prevent overfitting when training over multiple epochs -- or, at least, that's how most of what I've read explains why we don't need it on modern single-epoch training runs over large datasets.
If that's the case, though, when we do multiple epochs for a fine-tune with a more restricted dataset -- exactly what I was doing for the IFT test -- it might make sense to use dropout, regardless of whether or not the model was pre-trained with it. The fine-tuning setup already tries to avoid overfitting by bailing out when a validation loss starts rising, but dropout might still help it avoid overfitting prematurely.
On the other hand, something felt a little wrong about fine-tuning a model with dropout if its pre-training had happened without it. A model pre-trained with dropout have been trained on billions of tokens, and so the model will have spent a lot of effort learning to overcome the issues that dropout causes, but one trained without it won't have that benefit. Suddenly exposing it to dropout in a much shorter fine-tuning run felt rather like asking someone who rarely drinks alcohol to take a few shots of whisky; I felt that the models might not be prepared for the effects.
As I looked into this more, I noticed another surprising thing -- there was an error in the configuration that I was using when fine-tuning the OpenAI models, both small and medium. They were originally trained with dropout (or so it's believed -- the paper doesn't say, but "Build a Large Language Model (from Scratch)" says that they were, and this config on the Hugging Face GPT-2 code agrees). But my config was fine-tuning them without!
That actually made my original puzzle of why they outperformed my models on the IFT test seem even more perplexing, at least in the light of this idea. If dropout was a good thing for fine-tuning, then so far they had been penalised by not using it -- that is, they were even further ahead of my own models than I thought they were.
It was time to take a careful look.
Adding diagrams to my static site generator with D2
A lot of the time when I've been writing posts for this blog, I've felt that a diagram would really help. But they're a pain to produce well, and I think I underuse them as a result. I wanted to fix that, and wound up adding D2 support to my static site generator. I think it works pretty well!
Use the built-in GELU, don't roll your own!
Unsurprisingly, PyTorch's own built-in GELU function is faster than the hand-rolled one I've been using to date. But I was surprised at how much faster using it made things when training my models. I discovered this accidentally just now while working on something unrelated, but am logging the details here for anyone else that might find it useful.
The headline numbers: the same code, training the same model on the same data, ran at about:
- 21,000 tokens per second using the hand-rolled GELU from Sebastian Raschka's book "Build a Large Language Model (from Scratch)".
- 25,000 tokens per second using PyTorch's built-in GELU with no arguments.
- 25,000 tokens per second using the built-in GELU with
approximate="tanh", which uses the same maths as Raschka's version under the hood.
That's a 20% increase in throughput for both of the built-in versions -- definitely nothing to be sneezed at.
And what is particularly
interesting is that there aren't that many GELUs going on -- it's a GPT-2 small-style
model, with 12 layers. So that's 12 GELUs handling tensors
shaped (batch_size, seq_len, 4 * d_emb), which is (6, 1024, 3072) for my training setup.
Given that the rest of the model is doing all of the normal full attention stuff for GPT-2, it's
really surprising that the GELUs alone must have been taking up so much of the time. The throughput
numbers mean that we must have been spending about 17% of our time on the extra overhead from the hand-rolled
version, so that sets a lower bound for how much time the GELUs were taking up.
A quick(ish) Chinchilla check
I recently overtrained a couple of GPT-2 style models, training them both on 40 tokens per parameter rather than the 20 per parameter that is generally regarded as "Chinchilla-optimal".
The normal heuristic is that instead of doing that, you should scale up the number of tokens and the number of parameters equally -- so I would have been better off scaling up the model by and the token count by the same amount. By doing that, I should expect to get a better model in terms of loss on my held-back test set than I did with my 40-tokens-per-parameter models.
My training machine poppy wasn't doing anything, so I decided to give that a go.
Would the Chinchilla rule-of-thumb hold up?
As you might expect, it did. But it was a surprisingly close-run thing, and could conceivably have been in the noise. Let's take a look.
How I use AI on this blog
Inspired by this LessWrong post, I thought I'd write about how I use AI here. This is less in the interest of disclosure, more to provide a snapshot of what I'm doing right now so that I can revisit it in the future and see how it changes. And hey, maybe it'll be of interest to you, dear readers.
If I were to summarise my working philosophy in fewer than ten words, it would be: AIs identify problems and I fix them myself. With a very specific kind of exception (which I always flag), the text and code on this blog are human-generated. That's not a moral stand, but more a constraint imposed by what this blog is meant to be -- a place for me to learn in public.
Why do OpenAI's GPT-2 weights beat mine? Part three: testing overtraining
The GPT-2-style models that I've been training work really well, and I've even managed to train some that perform better than the original OpenAI small model in terms of cross entropy loss on a test set. But as I wrote previously, there's a mystery: why do they perform worse on my instruction fine-tuning evaluation?
I had various theories about why that might be, and to me, the most plausible-seeming of them was the amount of data they were trained with. As best I can find out, OpenAI's models were, by modern standards, trained on much more data than they should have been, while I'd used the theoretically optimal amount of training data.
To put it in other words, OpenAI's models were overtrained. If I deliberately overtrained my own models, could I match their performance? This post is a write-up of what happened, but so as not to bury the lede -- it didn't seem to help much, if at all.
Let's see why.