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.

[ Read more ]

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.

[ Read more ]