\n<\/p>\n
Writing the post that I wished I’d found when I started learning whatever it was…<\/p>\n
Archives <\/p>\n
Categories <\/p>\n
Blogroll <\/p>\n
Posted on 10 July 2026<\/p>\n
in<\/p>\n
AI<\/p>\n
When I was building my GPT-2 implementation in JAX,
\nI started with just token embeddings for the input, and a separate output head (as I was not using
\nweight tying). It wasn’t an
\nLLM — no Transformer blocks, no attention, no feed-forward networks.<\/p>\n
I was somewhat surprised when I noticed that even that stripped-down model had 77 million parameters
\nwith the “small” settings I was using to train — specifically, an embedding dimension of 768.
\nHowever, I realised I shouldn’t be — with a vocab size of 50,257, each of those components is essentially
\na 768\u00d750,257 matrix, and that is indeed over 38 million numbers.<\/p>\n
But the finished LLM at the end of the project was only 163 million parameters — that meant that the
\ninput and output components alone were almost half of it. That felt
\nlike a surprisingly large percentage.<\/p>\n
I had a similar shock when I was first looking into the feed-forward network,
\nand realised that it had roughly twice as many parameters as the attention layers.<\/p>\n
When we learn about the internals of LLMs, a lot of the focus is on the attention
\nmechanism. This makes sense — it’s the hardest part to get your head around. The
\nrest of the setup, at least for simple GPT-2 type models, is fairly standard stuff.<\/p>\n
But that means that it is easy to overestimate how much of the total
\nparameter count of the model attention uses up — especially for smaller models,
\nwhere the token embeddings and the output head are so large in comparison to
\nthe Transformer layers that make up the actual body of the LLM.<\/p>\n
OpenAI released GPT 5.6 today, so I decided to take its “Sol” variant for a ride
\nin Codex and asked it to write a visualiser.
\nIt shows
\nbreakdowns of how the parameters are split between embeddings, attention, the FFNs, and
\nthe output head
\nfor different sizes of GPT-2 models (or your own custom settings
\nwith the same architecture), and you can also add\/remove weight tying and QKV bias.
\nIt did a really good job — check it out! Here’s a screenshot of what it showed
\nfor GPT-2 small without weight tying.<\/p>\n
It’s well worth a play. In particular, it’s interesting to see what happens
\nas the number of tokens in the vocab gets very large (many modern models have
\nhundreds of thousands). You can very easily create a “tiny” model which is
\nalmost entirely embeddings and the output head.<\/p>\n