Once again I am not writing the Colossus post I thought I was going to (yet). As I started to write up the work I’ve done to make my local models useful, I realized that a lot of it depends on an intuition about generative AI that I only loosely had myself, and that many in my audience might not have at all. So let’s take a not-so-brief side quest into the things that make generative AI work. It’s fun stuff, I promise!
Classic Neural Networks
I often say that today’s AI models are just like the ones we had back in my college days — they’re just way, way bigger. And that’s kind of true, but it’s also totally not true. Some really specific innovations changed the game, and understanding them is important if you want to develop the right intuitions for productive, modern generative work.
But it’s still helpful to start with the basics. I wrote Fake Neurons are Cool more than four years ago but it still does that job pretty well — so if you could use a refresher or a baseline, click on over there and then come on back. I’ll wait.
Tokens
OK, the first idea here is the simplest, but since “token” is a shoo-in for Webster’s Word of the Year, it’s good to know what they really are and how they’re created. (Honestly, if you just think about them as “words” you’ll probably do just fine, but humor me anyways).
Tokens are the “vocabulary” of a model. Each token is a word or a fragment of a word that carries, hopefully, some independent meaning. The vocabulary can’t be infinite — its size has a direct impact on the size of the model, because layers at the bottom and top of the network have nodes for every token in the vocabulary — that gets pretty wide / tall!
Generally (and we’re going to see this approach a lot) model designers just pick a target vocabulary size they think will work well. There’s a ton of art to it. Too large and the model becomes too-big-to-run, or worse the tokens become meaningless and the patterns just aren’t there. Too small and you can’t differentiate enough to be useful. To give you a sense of scale, the Gemma 4 series vocabulary is about a quarter of a million tokens.
Once you have a target vocabulary size and your training inputs, you can create the token space. This is pretty neat, actually:
- Create a starting vocabulary by assigning a token to each unique character (or byte) in the training set. E.g., if your input training set is limited to ASCII text, you’ll start with something like 128 tokens (maybe less if, say, the BEL or tilde characters don’t appear).
- Count the occurrences of every unique adjacent token pair. Merge the most common pair into a new token and substitute that for the pairs in the training set.
- Keep doing this until you have either reached your target vocabulary size OR remaining pairs aren’t prevalent enough to be considered significant (by another mostly arbitrary, empirical definition of “significant”).
Of course there’s more to it than this — but that’s way more than enough for government work. That’s the easy one down.
Iterative Generation
Classic neural networks are great at lots of things — but generating long-form content isn’t one of them. A set of inputs goes in the bottom layer, and out the top comes a prediction score on every output node. For a yes/no answer, that’s a single output node with a score that you translate into something binary. For a classifier, you have one output node for every class and the model scores how likely it is that your inputs match each one.
The key is that the output nodes are fixed and finite. If I want to predict who will win the World Series this year, I can create a net with 30 output nodes, one for every team. But if I want to ask when the Mariners will finally win the World Series, I’d need an infinite number of output nodes — one for every year in the future, forever (stop laughing at me).
Conversational (aka “generative”) responses are by definition infinite — There is no output node in any network for the sentence “The truth is that the Mariners may never win the World Series, bro” and all of the other possible replies. So what are we supposed to do?
Well, it turns out that, given a training set of sequences to learn on, and a sequence of input tokens (aka a “prompt”), a neural network can get pretty good at predicting the next single token that would appear based on the “training set” — and that problem actually looks pretty familiar:
- We present our input prompt as a sequence of tokens. This is our context, so if the model context size is 256k, there are 256k nodes in the bottom layer of our network.
- (A lot to unpack here and we’ll get there; for now just think about it as network layers, except to say that this is where the model parameter count lives.)
- The fixed output layer is our token vocabulary; one score for every possible token.
- We pick the “best” token based on the scores. This is called decoding.
A fun side note: decoding is where a bunch of LLM settings work their magic. For example, “Temperature” controls how random the choice is — at 0.0 the model always picks the token with the highest score; at 2.0 it’s a crap shoot. And the “Top-k” and “Top-p” configurations define how many tokens go into the pool to pick from (using absolute numbers and/or percentages respectively). It all starts to make sense!
All this for one token?
Well, yeah. But here’s the cool part. If we add the token we just generated to the end of our prompt, we can do it again! And again! And again! Until the model spits back a special “stop” token and that, boys and girls, is Generative AI.
Except, not really. This model does work, sort of, but it has two enormous flaws that basically make it impractical for anything real.
First of all, the training effort is absurd. You have to train every single “next token” individually. So a single pass for the sentence “I can’t believe it’s not butter” is five distinct trainings (let’s assume words are tokens):
- I → can’t
- I can’t → believe
- I can’t believe → it’s
- I can’t believe it’s → not
- I can’t believe it’s not → butter
The other killer is positionality. Because each input node is distinct and feeds its own weights up through the network, token position really matters. We may be able to complete the sentence “I can’t believe it’s not butter,” but almost for sure we’ll crash and burn presented with “I absolutely can’t believe it’s not…” — because the four tokens after “absolutely” have been shifted to the right, and those nodes weren’t trained for this!
But take heart — some folks at Google in 2017 solved both of these at the same time. Seriously, these folks should headline the Topps Amazing Scientists All-Stars collectible card set: Attention Is All You Need. The Wikipedia article about the paper is also excellent.
Transformers
The invention of the “transformer” shakes up the fully-connected-layers concept we (I) know and love from classic neural networks, replacing them with a whole bunch of transformations designed to pick out important meaning and to do so largely in parallel. I’ll do my best to walk through it with enough detail to build solid intuition, but not so much that we lose everyone along the way. We’ll see.
Before we do this, I want to reinforce something I said earlier, because it was really important in my journey to getting comfortable with it. Many of the parameters we’ll encounter are empirical, made up, try-stuff-and-see-what-sticks. Although their general scale sometimes tracks in the rear view, nobody really has any precise insight into why “4” is better than “7” in some particular scenario. On the other hand, some matter a lot, and I’ll try to differentiate — it helps to know when to keep digging for the “why,” and when to let it slide.
Phase 1: Embedding
The first step is to take our sequence of input tokens and turn it into a 2D matrix. The mechanics are simple — look up each token identifier in a learned matrix that has one column for every token in the vocabulary and concatenate the selected columns together. The number of rows in the “embedding matrix” is called d_model and it very roughly corresponds to the number of features the model can learn for each token.
This result is called the “residual stream;” it will travel with us from layer to layer up the machine until we reach the top.
One thing we’ll encounter in later steps is that processing is largely non-positional — by default each token learns to “attend to” the tokens around it but can’t tell the difference between “man bites dog” and “dog bites man.” Obviously this needs to be dealt with.
The 2017 paper handled it here at the beginning, by defining a T x d_model matrix and adding it to the residual stream, so everything downstream included positional echoes. Later approaches did the same thing but used a learned matrix rather than defining it mechanically. Current models don’t do this at all — they handle position in the “attention” layers with a technique called RoPE that we’ll encounter later.
Phase 2: Transformer Blocks
After embedding is a stack of “transformer blocks;” each takes the residual stream as input and delivers a modified residual stream — same shape — as output. We’ll refer to the count of blocks as num_layers; it’s typically on the order of 12-100, not thousands.
Blocks execute in two stages, which you can broadly think about as “attention” and “consolidation.” Attention runs a number (num_heads) of transformations in parallel, each trained the same way but (thanks to random initial weights) magically assuming a distinct semantic role. Sometimes we can perceive these roles and sometimes not, but it’s cool regardless. “Consolidation” brings the attention output back together into a single matrix and does some classic NN stuff on that.
Phase 2a: Attention Heads
OK, there’s a lot going on here so buckle up — first we’re going to figure out how much each token should “attend” to the others in context (where attention from T1 to T2 is represented by a number stored at those coordinates within a T x T matrix).
Each head uses three learned matrices W_Q, W_K and W_V. Each is d_model x d_head in shape, where d_head is just d_model / num_heads. Step one is to multiply the residual stream by each of these, cancelling out d_model to get three interim matrices of shape T x d_head: Q, K and V:
The order of operations here was designed to give each matrix a distinct role:
- Q = “Query”: what “kinds” of tokens are relevant to each token (“I need a date!”)
- K = “Keys”: what “kind” of token each token is (“I’m a date!”)
- V = “Values”: the actual “stuff” each token has (“August 5, 1969”)
For models that don’t add token position information during embedding (i.e., most current models), it’s mixed in at this stage, usually with rotary positional embeddings. The idea of RoPE is that the token (horizontal) vectors in Q and K are rotated in proportion to their positions in the sequence, so that the angular difference between Q and K is a measure of relative position. This really deserves an entire article — for now just know that we’ve injected positional information into Q and K, so the model can see the difference between man-eating sharks and shark-eating men.
K is transposed, a math trick that lets us multiply it by Q to cancel d_head and end up with a T x T “attention matrix.” This one is pretty easy to understand — it represents, for each token, how important every other token is to its in-context meaning. With this framing what we do next also makes sense — apply a “causal mask” to each slot that is “in the future” in relation to its (horizontal) context. In the first row, we only know the first token in the stream and want to predict the second, so we set all of the tokens to the right with the value -∞. For the next row we know a bit more, and so on down the matrix.
Believe it or not, there are a couple of more things to do before the attention matrix is fit for use. First we divide its values by the square root of d_head, which is one of those “this seems to help” kind of things. Then we apply softmax to juice the machine with non-linearity (the equivalent of our activation function in classic neural networks).
And, finally, we turn the “attention” into actual output for this head; a simple multiplication with V that lands us back into a T x d_head matrix. Whew.
A fun fact — notice the light shading I added to the SW corner of the attention matrix above. The state of the art here is moving fast, and there a ton of tweaks that attempt to save space or time. A neat one is “sliding window attention.” In this model some heads are configured to ignore old as well as future context — a tradeoff of some long-term memory for less computation and cache.
Anyways, there’s more to do — but let’s take a breath because we’ve already seen the two game changers and they’re useful to reiterate:
- There is a TON of parallelism in this machine — first of all we’re running num_head heads at the same time, fully independently. Second, and this continues up the stack too, we’re able to generate predictions for every position in a single run. Our output matrix has a vector for every token! That is pretty sweet.
- We’ve integrated positionality into the machine in an extremely robust way. Rather than depending on the position of the token in the network layers themselves, we teach each one where it sits relative to the others — and allow that to influence learning.
Home stretch (for attention). Once all the parallel heads finish their work, we just stack their outputs together, giving us a T x d_model matrix. We multiply this by another learned matrix W_O and add the results back into the residual stream.
Phase 2b: Consolidation
The nice thing about this section is that it’s going to be short. After the head outputs are brought together and added to the residual stream, a normalized version of the stream is fed through a much more classic single-hidden-layer network (“MLP” or a Multilayer Perceptron because we’re fancy). Where the attention heads helped the token vectors figure out how they’re impacted by other tokens, this phase is all about refining the tokens themselves.
The hidden layer widens the matrix to a width of d_ff and then brings it back to d_model so we end up with the same dimensions. d_ff is — you guessed it — an empirically defined value, usually set by rule of thumb to d_model x 4. Different activation functions are used, but GELU (and it’s “gated” versions) tends to be the most popular at the moment.
The residual stream, which has been accumulating additions over these last two phases, is finally normalized in place so we have solid footing for the last mile.
Phase 3: Unembedding and Decoding
Remember way back during embedding, we turned the token identifiers into vectors. Now we’ve got to get back to vocabulary space, generating possible “next token” scores for each of the tokens in the stream.
In a lot of models, this translation is done using the original embedding matrix, transposed to cancel out d_model and end up where we want to be:
This reuse isn’t a hard-and-fast rule; sometimes (especially larger) models train it separately. But reuse kind of makes sense — the way it has been described to me is: if the embedding matrix maps a token to a “direction in space”, it follows that the closer the residual stream is to that direction, the more “likely” that token should be. Honestly, I’m not sure my brain stretches to really understand that, but another argument is “it seems to work, and it saves a ton of parameter size because this matrix is enormous.” I like that one.
During training, making predictions for every position is hugely valuable. OTOH, during generation we only really care about that last column — what’s the prediction for the next token given all the tokens in context?
That same softmax is applied to the scores in this column, and then one is picked as the winner. We’ve already talked about how we do that part, using parameters like Temperature, Top_k and Top_p to dial up or down the randomness.
Now append that token to the end of context, fire up the GPUs and turn the crank again.
Mic Drop
That was a lot. I hope the pictures are useful to some folks; no way could I understand any of it just looking at equations. PowerPoint was pretty overwhelmed with the sheer number of boxes involved, but we made it happen.
The good news is that now we’re armed with some vocabulary to help us talk about local models — what does “parameter count” really mean; how do attention schemes impact context size and cache space; why can they seem so smart and so stupid at the same time? And much more.
Next time — honest — we’ll start diving into that.