Example Causal Attention matrix

Attention

Imagine you were tasked with predicting the next word after "Washington" - you might think "D.C." or "state" but without context, you'd be reaching in the dark. If you knew that the preceding 10 paragraphs were filled with words like "revolution", "British", and "independence", you might be much more confident predicting the next words might be "ordered", "fought", or "led" than the previous options. Put another way, think about learning new words when reading - how do we go about this? We learn from context: the terms and meanings that surround a given word give the best indication of the definition of that word. Usually, the meaning of a word can be reasonably inferred (and/or predicted) from its immediate neighbors. "It was a perfect summer day: the sky was \ and not a single cloud was in sight." - it's mostly clear from immediate context what should go in \. But sometimes a word relies on context well outside of its immediate surroundings. If we are to learn how to predict the next word in a sequence most accurately, we're going to need a mechanism that takes context into account, and is robust to long range dependencies (related words that are far away in text space) where the immediate context is not entirely helpful.

The importance of context has been understood for some time. The literature on Recurrent Neural Networks (RNNs) extends back to the 80s, where model architectures that consumed some version of the previous output state as an input were first proposed. Long Short-Term Memory (LSTMs) networks, which pass forward and carefully update a global state vector, produced some promising results, however failed to generalize and scale sufficiently. Up until the seminal 2017 paper Attention is All You Need, no reliable and scalable method of context integration/utilization had been found.

The Attention mechanism did what other context mechanisms before it could not - it scaled well. What's more (or maybe more aptly, deeply and fundamentally related to its scalability), it exposed every token to the full context of previous tokens, thereby allowing for much more direct long range dependencies.

The most important fundamental intuition of Attention is this: allowing the model to have direct access to a wide context window of the tokens preceding the prediction target, and allowing the model to learn intermediate representations of those tokens and their contextual relationships is the essential ingredient that has enabled LLMs.

Architecture of the Attention mechanism

The architecture of the Attention mechanism is characterized by 3 conceptual components - Keys, Queries, and Values. Each token passes through 3 sets of learned weights to produce a key, query, and value. The Query of each token is then compared to the Keys of all other tokens - how similar is this key to that query? The result of this comparison is then used to produce a weighted combination of the Values. To return the the previous example of Washington, maybe the query would be something like "Am I a city, state, or general?" (this is a gross simplification to provide a conceptual foothold - in reality key, queries, and values are all vector representations without such clear meaning), and the keys of tokens like 'revolution', 'British', and 'independence' might be something akin to "battle, wars, and colonies" - when combined, the product would point to the value of the next token more closely relating to tokens relevant to 'revolution', 'British', and 'independence' than 'D.C' or 'State'.

An important note here is that inevitably analogies like the one above are imperfect, anthropomorphized explanations for what is undoubtedly a much messier, complex, and more inscrutable mechanism. Attempting to logically explain how Attention works on an abstract level is a worthy endeavor, and this explanation may largely be correct, but it's important to remember that there's a lot of matrix soup between this conceptual understanding and reality. And in some sense, that's all you actually need to know - Attention adds the token's context into the matrix soup in a sufficient dosage, and gradient descent does the rest. As soon as the model has access to sufficient gradient signal, it will learn. Everything else is just a clever optimization or hack, and we know what the Bitter Lesson says about such things.

Transformers

Transformers combine the above attention mechanism with a standard feed-forward multilayer perceptron (MLP), aka a neural net, in series. These Transformer blocks are, in turn, placed in series to form the body of a language model. In theory, Attention and Transformer blocks learn to capture separate and complimentary representations and relationships, and at differing levels of abstractions. Each Transformer block transforms the representation it receives into a modified intermediate representation, which is passed through successive Transformer blocks until eventually being transformed into the predicted next token, thus the namesake of the Transformer. At the core of the Transformer, and the component that distinguishes it from prior NN architectures, is the Attention mechanism.