Back to all posts
Coffee Codex - Going Through Transformers by Hand
Meshan Khosla ·

Coffee Codex - Going Through Transformers by Hand

Introduction

Today I’m at the library (with a coffee from Story Coffee).

In the last post I worked through transformers conceptually, and I thought I understood them. But then I kept reading Inference Engineering and kept having questions that I knew I should know the answers to, things like How does the final vector get turned into logits? I had forgotten about the unembedding matrix. Not a huge deal, but I started to feel that I didn’t really understand the entire flow. So I thought back to a class I took in college (CS186 - Databases) and how we used to have to count I/Os to understand how a query results in a database lookup, then progressed into things like how B-tree indexes impact the number of I/Os.

I didn’t particularly enjoy that work, but I figured maybe doing this by hand will help solidify my understanding, so here we go.

A Story Coffee cup beside my handwritten notes at the library

Math

I watched 3B1B’s videos on The Essence of Linear Algebra, and they were fascinating. I really liked the idea that a coordinate [x, y] can be thought of as x copies of i-hat plus y copies of j-hat. Even though I can memorize formulas for matrix multiplication like I did in college, this intuition means I can do the multiplications without having to memorize them.

Take the vector [4, 2] and the matrix [[2, 0], [0, 3]]:

|2 0| |4|
|0 3| |2|

The columns of the matrix tell us where the basis vectors land. The first column says i-hat lands at (2, 0), and the second says j-hat lands at (0, 3). Since [4, 2] means four i-hats plus two j-hats, its transformed position is:

4(2, 0) + 2(0, 3) = (8, 0) + (0, 6) = (8, 6)

This works the same way when the matrix is not diagonal. For example:

| 1 3| |4|
|-2 0| |2|

Here, i-hat lands at (1, -2) and j-hat lands at (3, 0), so:

4(1, -2) + 2(3, 0) = (4, -8) + (6, 0) = (10, -8)

Even when it looks like there is no transformation, there still is one. It is:

|1 0|
|0 1|

We call this the identity matrix because i-hat stays at (1, 0) and j-hat stays at (0, 1). Multiplying any vector by it leaves the vector unchanged. I think that’s pretty cool, since these are the “regular” coordinates we learned in school without knowing there was an invisible tranformation going on.

The learned matrices

Obviously, since I’m doing the matrix multiplications by hand, I don’t want large matrices, so I asked ChatGPT to give me a small pedagogical example of learned matrices. It decided on:

  • 3 input tokens
  • d_model = 3 (embedding dimension)
  • 2 attention heads
  • d_head = 2
  • 1 transformer layer
  • MLP hidden size = 4
  • vocab_size = 5

Writing it all out

Here are the pieces of paper I used to get through attention. I haven’t done the MLP yet, but wow!

Handwritten notes explaining linear transformations and basis vectors Handwritten matrix multiplication practice Handwritten transformer weights, embeddings, and positional embedding calculations Handwritten multi-head attention calculations Handwritten residual connection and MLP notes

So that was a lot of work. Is it worth it? Probably not. Is it fun? A little. I do have the urge to programmatically write these steps and create a mini vLLM in Effect just for fun, maybe I’ll do that.

Learnings

It really is a bunch of matrix multiplications, who would have thought. It’s also pretty incredible to consider that one of the most important technologies of our time is moving vectors around in high-dimensional space. And really smart people have figured out how to train weights that represent the transformations applied to those vectors. Pretty neat.

I also understand the point of the KV cache much more now. The process I did was prefill, and the follow-up step would be to keep the generation going with decode. The interesting thing about decode is that you still need the keys and values from the previous tokens when computing attention, but you don’t need to recompute them. You only compute the new key and value for the current token and append them to the cached K and V matrices. With the column-oriented convention I used on paper, each new token adds another column, which is neat.

While I don’t think it’s necessary to go through the transformer by hand, it does make me excited to read through the rest of the Inference Engineering book. I think the optimizations that are discussed in there will make a lot more sense now that my fundamentals are solid.

I know this was a short post (well, it took a while for me), but I’ll see you next time :)

References