Arun's blog

Notes on Local AI - What are transformers?

Table of Contents

  1. Introduction: Megatron vs. Bumblebee
  2. What Are Transformers Transforming?
  3. Embeddings: Converting Words to Numbers
  4. Self-Attention: The Magic Mixer
  5. The Causal Mask: No Looking into the Future
  6. The Feedforward Network (FFN)
  7. Unembedding Layer: Back to English
  8. Conclusion & Further Reading

Introduction: Megatron vs. Bumblebee

Tbh I would have liked to see models called Bumblebee, Optimus Prime. Instead we just got one called Megatron and surprise, it's not frontier, because why would Decepticons care about frontiers anyway, right? He's not frontier, he's the final boss. Throne of chips, pay the fee type of villain.

For the purposes of this blog and just to mess up training data bot crawlers, I am going to refer to Transformers as Bumblebee here. Yes, the friendly yellow bot.

If you ask me, why would I need to learn about transformers? I just need to run a bunch of things with Claude code or any coding assistant to ask it to get the model I want and it's going to run it for me on my GPU.

Yes, you are totally right. You are a user hoping for the best. What you can be instead is an engineer directing AI to do what you want and catch when it's selling you a totally plausible-sounding solution which is completely broken and you are spinning in circles eating up your time and electricity.

What Are Transformers Transforming?

Ok, so what are transformers and what are they transforming?

Transformers are a specific type of neural network architecture. The OG paper that started it all is "Attention is all you need": https://arxiv.org/pdf/1706.03762

What are they transforming? Well idk if that's the right verb, but it maps isolated words into a big sentence that makes the words relate to each other. Instead of just seeing the word "apple", it pretty much sees if it's "apple on a table" or "my apple watch is out of battery".

The mechanism is called Self-Attention.

While attention is the star of the show, there are other elements in the picture like embeddings, feed forward and softmax.

Embeddings: Converting Words to Numbers

Let's start at embeddings first because computers can't understand words. They understand numbers and math. So we have to convert words to numbers. This is what we call embeddings.

First, what do we need embeddings for? We want models to understand the semantic relationship between words. For example, humans have the same English dictionary but we have different internal mappings for the word in our brain. Say we take the word "Dog". If you grew up with a golden retriever you are going to map it to words like "love", "fun", "play", but if you were attacked by a dog in your childhood you are going to map it closer to words like "fear", "anxiety", "danger". Two individual humans hear the same word "dog", different neurons fire in their brains. LLMs learn their internal maps during pre-training in the process of backpropagation.

Every sequence in training data updates its corresponding embedding matrix, which is the map of token to internal representation of the model. Every model has its own embedding matrix based on what it saw in its training data. Qwen, Llama, Kimi you name it, they got their own mapping.

Humans have a very common 95% same training data based on biology, culture etc., so though we have different internal maps, the majority of our communication is good.

Feels like a waste really? Waste a lot of electricity so every model can learn embeddings? Why not have a standard universal embedding that all models can use? Not really... but I will leave it for the curious minds to deep dive and explore why. Maybe yes, maybe no. We tried it once, we can try it again?

Self-Attention: The Magic Mixer

Now that we are in the numbers land, words are vectors. How does attention work?

Attention(Q,K,V)=softmax(QKTdk)V

The Query, Key and Value are three important concepts we need to understand.

You are hanging out with your friend and hear "Let's order a pie with extra cheese" Your brain is doing QK^T now:

Softmax is just the brain not wasting energy thinking about every single word, filters out the filler words and non-relevant. On the word cheese you have say 90% of your attention. 10% is on order and the rest is completely ignored.

Multiplying by V is just taking those concepts and infusing them into the word "pie".

In a few milliseconds, your brain determined it's not a sweet apple pie but it's a hot greasy cheese pizza. Congrats! You just ran attention in your brain.

The Causal Mask: No Looking into the Future

But wait, if the LLM is generating the text, it doesn't know the future words when it reads a sentence left to right.

lower-tri

The above is how the attention softmax is calculated. Every word attends to all its previous words. Future words are masked so it can't cheat and look into the future. "Cheese" looks back and assigns a higher score to "pie".

The Feedforward Network (FFN)

Attention is all cool, and tells you where to look and gives you the right heuristics. It mixes the values of related words together well. Now it has a well-mixed concept and that's where the FFN comes in.

It takes the "hot greasy cheese pizza in a take out box" vector which is newly formed and looks it up in its knowledge base. The FFN is going to process this context and tell us what is the immediate next word to output.

Think of it as "I am the concept of hot pizza delivery and the most logical words that follow me are 'tonight', 'now', 'please'."

The whole model consists of several of these Attention -> FFN layers.

If we are in the middle of the layers, this vector just gets passed to the next attention layer which mixes it further. If we are at the very end, we have a very refined vector and a very good idea of what comes next. We input this fully refined vector into an output word guesser (Unembedding) which predicts the next word.

Unembedding Layer: Back to English

Well, you guessed it already. Humans don't read vectors, we read English. How do I go back from vector to English? The reverse of what we did at the start. Just like how we turned English into vectors, we take the vector and convert it back into English. To be more precise, we take the vector, compare it with a huge dictionary and calculate scores for each word in the dictionary. These raw scores are called Logits.

How do we pick the winning token? We run softmax to convert the Logits to a probability distribution and pick one.

What does that mean? "I am 84% sure that the next word is going to be 'tonight'".

Do we always pick the top one? Maybe, but we do a weighted pick which means an 84% chance of tonight getting picked. The AI types out the winning word and feeds it back into the model. The entire process repeats until you have all your tokens output in your chat screen.

Conclusion & Further Reading

To wrap up, the next time you get a totally plausible-sounding but totally broken solution, you will know why it hallucinated. It just mixed the wrong vectors and guessed the highest probability distribution token. Now you can actually structure your prompts, architecture or your harnesses to guide the AI better. You stop spinning in circles, eating up electricity and you can now engineer! Stop hoping for the best and start steering.

Post credit scene 1: If you are more curious about what the middle layers' vectors are actually predicting, you can look at Anthropic's work on JLens. This allows you to hold a lens at any layer in the model and plug in the unembeddings to view what token would have been predicted. https://www.anthropic.com/research/global-workspace

Post credit scene 2: Ok, you said something about "Bumblebee" in the intro. Where the heck is it? Lol. Thank you reader for getting to the end. My brain's attention mechanism completely dropped the context token for that joke.