So, there's a lot of buzz about memory in the AI space; this buzz seems to, by the time of me writing this, have overtaken completely the discussion about compute. This is because here lies perhaps the catch of the "Deep Neural Network scale really well when parameters grow" phrase: those parameters need to a) be loaded somewhere, and b) need to be given to a processor to munch on.

From the previous posts, it should be clear(-ish) that to compute attention for a given token, we need to have performed calculations on all previous tokens. Concretely we need to compute their keys ($K$) and values ($V$). Computing these values depends only on the token and the model weights, therefore they can be computed once and just referred to in the future. This is what a KV-cache is/does: As we process a token, we generate its $K$ and $V$ and store them in the cache that we can just read when we need them again.

In a previous post, when explaining attention I did mention the expression "attention heads", and then quickly moved on without explaining what that meant or how it did affect the computation. In this post, I'll try to atone for that by explaining to the best of my current understanding (clearly here current means my understanding after being done with writing this and all the necessary reading).

To prevent this turning into me playing [insert Bethesda game name here], I'll try to have at least one main story post in between side quests. In this one we'll cover the next item I implemented in my project: parsing and loading of models in safetensors format.

At the highest of the high levels, the transformer applies a number of transformations to a token in order to obtain the next token. We'll, for now, hand-wave the fact that these consider results from application of the same operations to all tokens in the input up to the one we are currently processing.

To understand how these operations are made and how information is accumulated across a pass to all the model layers, we need to understand what the literature calls the "residual stream". The residual stream starts as the original embedding (the vector representation of the token), and as we move forward in the block we add all of the results (this can be imagined as the block imbuing the original vector with meaning derived from context and the model's training), and so, it is "residual" because each layer adds its contribution to the running result (always of the form $x_{i+1} = x_i + f(x_i)$).

In the previous post, I talked about the methodology of my learning project. In this one, I'll describe the general plan I built iterating with the AI tool based on what I wanted to achieve. To recap, I want to write the code myself, with AI just providing guidance, reviews, and information.

This being a learning adventure, the progression I arrived at was: first, I'll write a naive implementation that lets me learn the process, and we'll refine from there.

Usually, when I am using AI at work, I rely on it to write code: fix a bug, scaffold a project, sometimes even to write a project end to end. This flow leads to a couple of things: my understanding of the projects can become stale if I'm not careful, and, honestly, I'm not having as much fun as I used to.