{"id":2003,"date":"2024-04-01T03:51:03","date_gmt":"2024-04-01T03:51:03","guid":{"rendered":"https:\/\/timallanwheeler.com\/blog\/?p=2003"},"modified":"2024-04-01T03:51:03","modified_gmt":"2024-04-01T03:51:03","slug":"transformers-how-and-why-they-work","status":"publish","type":"post","link":"https:\/\/timallanwheeler.com\/blog\/2024\/04\/01\/transformers-how-and-why-they-work\/","title":{"rendered":"Transformers – How and Why They Work"},"content":{"rendered":"\n

This month I decided to take a break from my sidescroller project and instead properly attend to transformers (pun intended). Unless you’ve been living under a rock, you’ve noticed the rapid advanced of AI in the last year and the advent of extremely large models like ChatGPT 4<\/a> and Google Gemini<\/a>. These models, and pretty much every other large and serious application of AI nowadays, are based on the transformer architecture first introduced in Attention is All You Need<\/a>. So this post is about transformers, how and roughly why they work, and how to write your own.<\/p>\n\n\n\n

The Architecture<\/h2>\n\n\n\n

The standard depiction of the transformer architecture comes from Figure 1 of Attention is All You Need<\/a>:<\/p>\n\n\n

\n
\"\"<\/figure>\n<\/div>\n\n\n

<\/p>\n\n\n\n

I initially found this image difficult to understand. Now that I’ve implemented my own transformer from scratch, I’m actually somewhat impressed by its concise expressiveness. That’s appropriate for a research paper where information density is key, but I think we can do better if we’re explaining the architecture to someone.<\/p>\n\n\n\n

Transformers at base operate on tokens<\/em>, which are just discrete items. In effect, if you have 5 tokens, then you’re talking about having tokens 1, 2, 3, 4, and 5. Transformers first became popular when used on text, where tokens represents words like “potato” or fragments like “-ly”, but they have since been applied to chunks of images<\/a>, chunks of sound<\/a>, simple bits<\/a>, and even discrete states, actions, and rewards<\/a>. I think words are the most intuitive to work with, so let’s go with that. <\/p>\n\n\n

\n
\"\"<\/figure>\n<\/div>\n\n\n

<\/p>\n\n\n\n

Transformers predict the next token some given all of the tokens that came before:<\/p>\n\n\n\n

\\[P(x_{t} \\mid x_{t-1}, x_{t-2}, \\ldots)\\]\n\n\n\n

For example, if we started a sentence with “the cat sat”, we might want it to produce a higher likelihood for “on” than “potato”. Conceptually, this looks like:<\/p>\n\n\n

\n
\"\"<\/figure>\n<\/div>\n\n\n

<\/p>\n\n\n\n

The output probability distribution is just a Categorical distribution over the \\(n\\) possible tokens, which we can easily produce with a softmax layer. <\/p>\n\n\n\n

You’ll notice that my model goes top to bottom whereas academia for some unfathomable reason usually depicts deep neural nets bottom to top. We read top to bottom so I’m sticking with that.<\/p>\n\n\n\n

We could use a model like this to recursively generate a full sentence:<\/p>\n\n\n

\n
\"\"<\/figure>\n<\/div>\n\n\n

<\/p>\n\n\n\n

So far we’ve come up with a standard autoregressive model. Those have been around forever. What makes transformers special is that they split the inputs from the outputs such that the inputs need only be encoded once, that they use attention to help make the model resilient to how the tokens are ordered, and that they solve the problem of vanishing gradients.<\/p>\n\n\n\n

Separate Inputs and Outputs<\/h3>\n\n\n\n

The first weird thing about the transformer architecture figure is that it has a left side and a right side. The left side receives “inputs” and the right side receives “outputs”. What is going on here?<\/p>\n\n\n\n

If we are trying to generate a sentence that starts with “the cat sat”, then “the cat sat” is the inputs and the subsequent tokens are the outputs. During training we’d know what the subsequent tokens are (our training set would split sentences into (input, output) pairs, such as randomly in the middle or via question\/answer), and during inference we’d sample the outputs sequentially.<\/p>\n\n\n\n

Conceptually, we’re now thinking about an architecture along these lines:<\/p>\n\n\n

\n
\"\"<\/figure>\n<\/div>\n\n\n

<\/p>\n\n\n\n

Is this weird? Yes. Why do they do it? Because it’s more efficient during training, and during inference you only need to run the input head once.<\/p>\n\n\n\n

During training, we know all of the tokens and so we just stick a loss function on this to maximize the likelihood of the correct output token:<\/p>\n\n\n

\n
\"\"<\/figure>\n<\/div>\n\n\n

<\/p>\n\n\n\n

During inference, we run the encoder once and start by running the model with an empty output:<\/p>\n\n\n

\n
\"\"<\/figure>\n<\/div>\n\n\n

<\/p>\n\n\n\n

We only use the probability distribution over the first token, sample from it, and append the sampled token to our output. If we haven’t finished out sentence yet, we continue on:<\/p>\n\n\n

\n
\"\"<\/figure>\n<\/div>\n\n\n

<\/p>\n\n\n\n

We only use the probabilistic distribution over the next token, sample from it, etc. This iterative process is auto-regressive, just like we were talking about before. (aside 1)<\/mark><\/span><\/span>This sampling typically ends when we reach a special token dedicated to stopping, called the end token<\/em>. Its just one of the n tokens just like any other except it has this special significance and doesn't get printed out when you show the sentence to the user.<\/span><\/span><\/span><\/span> (aside 2)<\/mark><\/span><\/span>Sometimes we can literally pass nothing into the output head. However, transformers often have a fixed size and we have to fill in those output values with something. In that case we can use special padding tokens and masking to prevent the model from using that information. More on that later.<\/span><\/span><\/span><\/span>\n\n\n\n

The reason I think this approach is awkward is because you’d typically want the representation of \\(P(x_{t} \\mid x_{t-1}, x_{t-2}, \\ldots)\\) to not care about whether the previous tokens were original inputs or not. That is, we want \\(P(\\text{the} | \\text{the cat sat on})\\) to be the same irrespective of whether “the cat sat on” are all inputs or whether “the cat sat” is the input and we just sampled “on”. In a transformer, they are different.<\/p>\n\n\n\n

Robust to Token Ordering<\/h3>\n\n\n\n

The real problem that the transformer architecture solves is a form of robustness to token ordering. Let’s consider the following two input \/ output pairs:<\/p>\n\n\n\n