I should start by disclaiming that this project is essentially reimplementation of Karpathy's NanoGPT. I am currently in my first semester of graduate school, and exploring AI inference/training research. I wanted to challenge my understanding the model layer of the AI stack by first, watching Karpathy's video, taking notes with pen and paper, then reimplementing at version from my understanding. In my experience this is a really productive exercise for learning technologies deeply in a few hours.

Another caveat before the technical sections is the extent to which the transformer was implemented "from scratch". I still used PyTorch because there are many convenient functions for tensors, differentiation, and linear layers. Every other component of the transformer was put together directly.

Data

Because GPT I create is going to be small, the language it produces won't be very intelligible to us. Karpathy trains his GPT on all of Shakespeare's works, which is already a corpus of unintelligible text (ba dum tss!). I searched for a more modern corpus to experiment with and came across this South Park diaglogue set on Kaggle. Getting the GPT to produce not only English but also something resembling humor was a more motivating than generating plays I hated in high school.

First I wrote on a small script that pulled out the character and dialogue information from the dataset, since we only care about the characters and words and not the episode information. The resulting text file was ~79,000 lines with five million characters.

Stan: You guys, you guys! Chef is going away.
Kyle: Going away? For how long?
Stan: Forever

The last step in data processing is tokenizing the data. Tokens are the atomic units LLM's are purposed with understanding and generating, and tokenization is how we convert language to tokens. This diagram is pretty funny illustration of what can go wrong.

Optimal tokenization is it's own subproblem that is outside of the scope of my curiosity so I wrote a simple character-level tokenizer, which is exactly what it sounds like; the GPT treats each character as a unit of meaning. The encoder (character -> token) gives the unique integer value of a given character, and the decoder (token -> character) does the reverse. The tradeoff in this simplicity is a burden on the model to learn how to spell words, names, punctuation, etc. As a result, the model likes to make up names and words that sound plausible, which makes the output more fun to read.

I encoded the text and saved the final 10% as a validation set, and the first 90% of the text was used as the training set. Then I created batches at random positions that were 64 sequences of 256 characters. Each sequence gets a target that is itself shifted one character to the left so it's given everything up to the current character and the model is tasked with predicting the next one.

Attention

Now we get to the meat and potatoes of the transformer: causal self-attention. At the highest level, we have a number of attention heads. Each head has it's own weights and calculates it's own set of queries, keys, and values. I like to think of it as delegating the responsibility of attention, where instead of one head trying to understand the meaning of the data, each head is looking for specific patterns in the data. In my GPT I used 64 total heads. Each head then projects the input into three different representations: Queries, Keys, and Values.

Think of the attention head as a librarian. You ask the librarian a question (input) and her brain processes your question (query). She takes her understanding of your query and compares it to her knowledge of the library catalog (key), and gives you the dewey decimals of a couple of books that may interest you. You take those keys and go find your books. The text of the book is like the value matrix, and contains the information pertaining to your query.

The heads multiply the Query by the Key to decide how relevant every character is to the current character. Those scores are scaled by the head size and normalized by the softmax operation to represent probabilities. We then apply a triangular mask to the attention matrix, making every value above the diagonal negative infinity, giving them zero probability. This prevents the model from "cheating" during training by seeing what comes next in the sequence.

Putting It Together.

The transformer isn't created when multi-head attention is complete; it still needs what is called a feed-forward network. The feed forward network simply expands the embeddings to 4 times it's size, applies a ReLU (fancy term for changing every negative value to 0), and projects it back down. I also added residual connections, denoted by the arrows around the MHA and FF blocks, which allows the initial information to pass through the network to prevent vanishing gradients.

I implemented the LayerNorm myself as well, which is just two lines of code representing this equation : $y_i = \frac{x_i - \mu}{\sqrt{\sigma^2 + \epsilon}} * \gamma + \beta$ where $\gamma$ and $\beta$ are learnable parameters.

By combining the LayerNorm, Attention, and feed forward network into a single transformer block we can have a working GPT. Within the GPT, first we have to give each character a token embedding and a positional embedding. Attention doesn't have any idea where in the sequence the token was located so the positional embedding is used to give the ordering information by appending it to the learned information. Then we send the embeddings through a number of transformer blocks, depending on the number of layers we have in our model. Then, one final linear projection and normalization are done. That's the entirety of the working architecture.

Training & Results

I trained my model for 5,000 iterations with AdamW at a learning rate of 3e-4. Every 500 iterations it averages the loss over 200 batches and outputs it. To generate tokens, the model predicts a probability for every possible next character and samples one from a multinomial distribution. That sample is appended to context and the cycle repeats. The maximum context my model accepts is 256 characters, so longer generations use a rolling context window.

Final Results

Manager: Oh, fellas.
Fingerator: Again.
Stan: Oh, of course. Jesus Christ!
Stan: Girls, is Cartman. How about fingerbon?
Cartman: Oh my God.  What are you doing, you dumbass's bush?  We must hang on. I'm sorry. God's notice morning, I don't really have a code, but kids are you now dead, but ya i-it was because Jimmy, in-My friends.
Cartman: See you can, fly are you asking home and go back-off at like this!  It's Block Hammer.I got a code 2 five days.

The generated dialogue doesn't make much sense and doesn't seem to follow any theme of conversation. The funniest result I observed is when the model tries to swear but produces a misspelled homophone of the curse. Still, I was impressed at the amount of structure it was able to pick up, and the majority of the character names end up correct. The grammar is surprisingly accurate as well.

I would encourage anyone who is exploring AI to attempt to recreate a popular technology themselves. For me, this project not only improved my understanding of how LLMs work, but also built a lot of intuition that has aided my research. It's my opinion that having some understanding of how LLMs work through implementation is a huge advantage not only when using them yourself, but developing resistance to vague, non-technical claims about how they operate.