PiccoloGPT & FratelloGPT

"Before I can start carving, the idea must be almost complete. I say 'almost' because the really important thing seems to be the sculptor's ability to let his intuition guide him over the gap between conception and realization without compromising the integrity of the original idea; the point being that the material has vitality - it resists and makes demands." - Barbara Hepworth

Creation is a very personal, albeit, messy act. In a nutshell, that is what I learnt when I was finally done training my own GPT-model who I call, lovingly, as 'Piccolo'.

With this post, I wanted to detail my entire, painstaking journey of creating my very own small language model. Indeed, it is an extremely simple project in my eyes, but I thought it best to at least archive this journey in the form of a blogpost (regardless of whether it garners eyes or not) so that I can have something to reminisce on for the project that marks the start of a long journey.

So fair warning, this will be a long read.

THE IDEA

One of the core tenets in mechanistic-interpretability is that any budding researcher (worth their salt) should be able to code up a GPT from scratch, without looking anything up, on a whim. Neel Nanda outlined this in a blogpost and his reasoning was that this skill serves as a base for you to be able to run any ideas for experiments you might have for LLM interpretability, on command.

For me, this was a bit unreasonable. My first reason is in the way I learn. My learning style is that I break whatever it is I'm learning down to the absolute fundamentals, till I reach the 'boiler-plate' concepts (so to speak) which this piece of knowledge is based on. Much of the time, this is a fast way for me to learn, but on entirely novel areas of knowledge that I have no history with, this can be a slow, disorienting process for me, even though it is a sure shot path to mastery. Now, when it comes to anything machine-learning, experienced practitioners will know that though this skill of first-principles thinking is a massive asset (as it helps you cut through a lot of ML-fluff), acquiring it as a newcomer is an uphill battle because of the state of publishing in this field (the rant on which I'll keep to myself).

Second reason: I was a bit rusty at coding in Python and with PyTorch because though I was taking up competitive programming newly at the time and resuscitating (plus fleshing out) my coding chops; that thing was entirely being done in C++, and moreover it's a completely orthogonal discipline to DL. So, where I was implementing (simple-esque) solutions to fun problems in C++ in the former, I'd be wrestling with high-dimensional tensors in PyTorch with the latter!

Nonetheless, my mathematical maturity was really bolstered from my (admittedly, stupid) decisions to take extremely hard courses (some of them being the hardest) in my first and second years and so, tackling this challenge was just a question of 'cutting through the bullsh*t' (as they say) and solving a ton of mini-problems along the way.

RE-GAINING LOST UNDERSTANDING

I began my dive with reading the first few chapters (up until RNNs) of Simon Prince's 'Understanding Deep Learning'. These first few chapters are an absolute masterclass in pinning down the fundamental nature of deep-learning and how it is more of an alchemical science. I highly recommend it if your learning profile is similar to mine. Moreover, the book also introduces each major architecture (asides state-spaces, if I remember correctly) in the same manner, so indeed, it is a tour de force in deep-learning exposition.

After reading through these chapters, I kept in mind that I wanted to get to Transformers because language-model interpretability is what I want to work in anyways. However, due to my learning profile, I cannot stand to have holes in my understanding and so, I needed to know why the Transformer even exists in the first place. At the time, I was only acutely aware that the Transformer architecture was built to address some weaknesses in the RNN and CNN architectures for language-processing tasks and; having taken Andrew Ng's Machine Learning course on Coursera a few years ago and having read Chris Olah's blogpost on RNNs for a now-defunct project, I could scarcely recollect how they both worked. But, knowing me, if I so much as even went anywhere NEAR either of these architectures' chapters in the book to revisit their ideas, I know I'd be stacking unnecessary learning friction.

Nat Friedman has this really good heuristic that I'm a fan of, and it is that you should maximize your 'learning per unit-time'. So, to circumvent this stacking, I had Claude prepare for me an artifact that outlines how RNNs and CNNs tackled natural language-processing tasks and what weaknesses lay in both, both of which are what's addressed by the Transformer architecture. Unfortunately, this artifact is lost to me now, but it shouldn't be much of a problem in asking your Claude to produce a similar artifact nonetheless!

After reading the limitations of both architectures via the artifact, I was ready for my dive into Transformers. I printed out both - Prince's chapter on them and the infamous 'Attention is All You Need' paper. With a pencil and eraser in my hands, plus these papers at my disposal, I was ready.

What a sweet, summer child I was.

At first, when I worked through Prince's chapter (up until 'Transformer for images'), I thought I had the entire architecture and all the minute details loaded in my head. I thought to myself, 'That was easy.'.

Then, to learn from the horse's mouth, I turned to the 'Attention' paper and all hell broke loose. Before this, I had indeed read some deep learning papers (like the AlexNet paper), but maybe by some cosmic fluke, I was protected from the infuriating side of DL literature 'slop'.

For the sake of my sanity, I went through the paper with my pencil and eraser, scribbling away notes, editing the infamous diagram of that paper (it is literally INCOMPLETE and hence, SO MISLEADING for those who don't want to simply gloss over it and be handwavy with their understanding (looking at you, LARPers)) and finally making a mental 'animation' to round off my understanding (which I also 'wrote' in the whitespace behind another page).

There was also the synergy of mathematical-tricks + engineering-tricks, which ML as a whole is famous for, that I was re-introduced to with this paper. I say 're-introduced' because the AlexNet paper had already done that for me before, but I think I might've forgotten about it. After 2 days of wrestling, it was done.

Sigh... atleast it's over now, right? Nope!

I realized early on that there was a massive discrepancy in the architectures between Prince's treatment and the 'Attention' paper, and now, I had to reconcile the two. It was a disaster for me initially, but chaos settled to calm after I realized that a 'Transformer' is an umbrella term for a whole slew of architectures - there's encoder-only, decoder-only and encoder-decoder Transformers. I had to only concern myself with decoder-only architectures.

Alas, taking after Neel Nanda's tip to have an LLM quiz you on papers you read, I was finally done with sifting through the immediate-literature and gaining a theoretical understanding of (decoder-only) Transformers.

Now came coding one up from scratch.

CODING IT UP, PART-1

As I mentioned at the start of this post, I first had to revive my tensor-wrangling skills. Thanks to my efforts last year, I made 3 very good Colab notebooks for this very purpose, chock-full of notes for myself to convey the ideas of each function and method when working with tensors.

Having coded up an entire workflow on my own from scratch for a simple neural network, I thought I was ready. I wanted to take Friedman's heuristic to the extreme and push myself to fully utilize my mathematical maturity and code up a GPT of my own from scratch as fast as I could.

How hasteful I was.

I first jumped to ARENA's exercise on making your Transformer-based model from scratch. The exercises in their Colab notebook were really terse for me, however, and I realized I wasn't up to the mark. Yet, at the time, with the same grit that pushed me through those hardcore math courses I took in my early years, I pushed. After feeling like Sisyphus though, I relented.

I then turned to Andrej Karpathy's video on coding up GPT-2 from scratch. With this resource though, the exposition was too simplistic and his code was a bit unreadable. Okay, let me rephrase that last bit, the notation in his code was too unreadable for me. The notation from the understanding in my head differed quite a lot from his (which is also another source of friction when it comes to ML code implementation as a whole, I learnt through this). I still grit my teeth and pushed through, but eventually, I caved in.

These 2 days were extremely painful, and in the past when I had attempted to break into mechanistic-interpretability, coding up your own model from scratch was always the singular wall that I could never scale. Yet, I wasn't going to give up. If I could make it through those courses, if I could adapt to, and do, measure theory and abstract algebra better than my whole class (ignoring some unfortunate luck factors), this challenge should be nothing for me. I tried looking for a solution still.

Later that second day, I found Josh Starmer's video and Colab notebook on coding up a transformer. Reviewing his code-cells, the notebook was perfect! It contained the exact kind of exposition that I'd make for myself to learn with/from, with notes on the tensor-methods (that I was still coming to terms with wrangling) and also on what you should do if you want to implement bigger architectures.

The 2 biggest selling points for me were that -

a) It was for decoder-only transformers.

b) It was focused a lot on the atomics - it only used a 5-word dataset and showed exactly how this text gets processed throughout the transformer. For someone with my learning profile, making the entire pipeline clear as glass like this was a godsend as it would complement and even reinforce my theoretical understanding. The notebook may seem laughably easy to a lot, but try and prod the understanding of those who do so, and you shall see their tower of cards fall swiftly.

'Slow is smooth, smooth is fast', it's said. Maybe watching that Formula-1 movie wasn't only entertainment after all!

The following morning, I went to town on it and by that afternoon, I finished half of his notebook (right before my week-long break).

Just the day before then (as of writing this post, that is), I returned with fresh eyes and with an even more joyful spirit. I played around with the code of the second half and had a blast breaking it all down and absorbing the code. With this, I was done.

With this then, I was finally ready to hunt my white whale - coding up your own GPT from scratch on a proper corpus of text.

CODING IT UP, PART-2

I first decided to call my model, 'BabyGPT', but that was lame and so, I decided to settle on 'PiccoloGPT' because I thought 'Piccolo' was a really cute word. PiccoloGPT was only meant to be a fun, afternoon project anyways and so, why not keep a name that reflects the same?

Next, I decided on the dataset to train off of. I recalled that the 'TinyStories' dataset was intended for this purpose and so, I ran to HuggingFace and queried a 20,000 story sample using HF's SQL console. I downloaded the data as a parquet-file on my laptop.

I then went and wrote out on paper, my choice of tokenizer, the vocabulary-size I'd fix, the embedding dimension, the number of attention-heads in each multi-head self-attention block, the number of transformer blocks I'd stack, etc. the whole nine yards for hyperparameters and architecture choices.

It came down to the following - which are all my personal notes that I scribbled in:

EDIT: I re-ran my derivation and found the precise formula for parameter-count on the Piccolo-class of models. When I was planning out on training the beefier Fratello-class, this formula helped me with a lot of analysis on parameter-distribution, Chinchilla-optimal token-count for a model of size N, etc.. If/When I'm done training an instance of the Fratello-class, I'll post up a lot of the graphs and plots I used to plan out the optimal data-corpus-size for the Fratello-class, Fratello's optimal parameter count, etc.. It was quite exciting to budget tokens and parameters in this way!

Using the piece-by-piece heuristic I regained from working through Josh's notebook and my own workflow, I coded up the Absolute Positional Embedding, Attention-head, Multi-Head Self-Attention, Embedding Layer and the FFN layer classes. The training script was really similar to my workflow anyways, so I just had to make a few changes here-and-there. My code, like anyone's, was riddled with bugs at the start, but debugging got me into a flow-state that I really enjoyed.

After training for 5000 iterations, I had hunt down my white whale. Piccolo was done training.

Some later analysis then yielded that Piccolo was actually undertrained, but that was fine by me! It was speaking in garbled-up text most of the time, but sometimes, it was forming coherent words and sentences too.

POST-MORTEM & FratelloGPT

When I ran an analysis of parameter-distribution between the transformers-stack versus the embedding-unembedding stack, I found that the latter took up nearly 67% of Piccolo's total parameters and the former, just 33%. This is non-optimal because intuitively, by having a fatter lookup layer and a thinner transformers layer on top, this means that most of the work done by Piccolo is just looking-up stuff, instead of understanding. Obviously, for a transformer-model, this isn't good and the cause was my haste.

Nonetheless, I am already tinkering with training a chunkier version of Piccolo that simply relies on the same architecture, but with bigger blocks. This class of models is what I call as Fratello. With help from Claude, I used the formula I derived above and Claude generated a plot of '% of parameters taken up by transformers v/s the embedding dimension', assuming that Fratello had hyperparameters and architecture choices described below:

Now, I worked with a dangerous heuristic here by working backwards and seeing how many stories from the TinyStories dataset would I need to fit a model for the Fratello-class first. There are exactly 1,990,338 'clean stories' (stories without special characters in them) and so, for the Fratello class having...

A) 4 $\rightarrow$ 8 attention-heads.

B) 4 $\rightarrow$ 5 transformer-blocks.

C) 6,000 $\rightarrow$ 7,000 tokens in the vocabulary.

...we plug these into the equation to solve for the optimal embedding dimension that Fratello should have (assuming $20\frac{parameters}{token}$ and $159\frac{tokens}{story}$, where these numbers were obtained from Piccolo's dataset): $$[60(d_{embed})^2 + 14,047(d_{embed}) + 7000] *\frac{20}{159} = 1,990,338$$ We get that that the optimal embedding dimension of Fratello is 408 and the number of stories to train Fratello with this embedding dimension now is 1,978,115 stories. That is just 99.4% of the clean-stories present in the TinyStories dataset. I have pushed back training the Fratello class because my engineering-skill needs work to code-up the data-loader, tokenization, etc.. of so many more stories and moreover, I'll be (albeit, quickly) picking up training on GPUs as well with Fratello.

Also, from the above plot then, what this means is that 63.6% of Fratello's parameters will be allocated for the transformers-layer and the rest towards the lookup tables - a polar-reversal from Piccolo's case; which is fantastic news!

This last fact means that Fratello will be awesome for doing mechanistic-interpretability experiments on as well!

TAKEAWAYS & SEND-OFF

Though this was an extremely simple project (yes, I am aware of how behind I am compared to some peers I know in the field, but hey, better late than never!), I also recalled that I absolutely love coding and that, through CP and mechanistic interpretability, coding was my next frontier to make strides in and apply all my pent-up math knowledge to.

This project also got me to pick up tinkering and training (bigger) models, which is where some of my future hobby-interests lie - reasoning models, testing scaling laws (partly inspired by Andy Jones' amazing 'Scaling Scaling Laws via Board Games', which is an absolute gem of a paper), and a whole lot more that I want to keep close to my chest for now!

Moreover, reading Anthropic's Circuits work became so much more easier too (to the point that it's common sense now, some of what is written) and I've been breaking open Piccolo's weights to play around and display something as simple as virtual weights.

My head's been absolutely exploding with new ideas to test in language-model interpretability and I can't wait to execute on them + write about these over here. Some of these may come very soon, but I make no promises!!

The future looks extremely ripe!