Self-Attention: Unlocking Context in Transformers

Introduction

In our last lesson, we explored tokenization, the process of converting raw text into a sequence of numerical Token IDs that a model can understand. We ended on a crucial question: once we have these tokens, how does a model grasp their meaning in the context of a sentence? A single token for "bank" has the same ID whether it refers to a riverbank or a financial institution. This is where the magic of the Transformer architecture truly begins.

This lesson delves into self-attention, the mechanism at the heart of the Transformer. You will learn how it enables a model to dynamically weigh the importance of different words in an input sequence, allowing it to build context-aware representations. We'll dissect this process from its high-level intuition down to the underlying matrix mathematics, giving you the deep theoretical grounding necessary to understand how LLMs "think."

Why Embeddings Need Context

We begin with a sequence of token embeddings—high-dimensional vectors that represent the initial, context-free meaning of each token. The core task of the Transformer's encoder is to enrich these static embeddings with contextual information from the entire sequence.

A brilliant visualization from Grant Sanderson of 3Blue1Brown fame can help build a strong intuition for this process.

The key insight is that the model needs a way for the vector representing "one" to "talk to" the vectors representing "two roads" to understand its specific meaning in that context. Self-attention provides this communication channel.

To see this in action, consider the sentence: "The animal didn't cross the street because it was too tired." What does "it" refer to? For a human, the answer is obviously "the animal." Self-attention is the mechanism that allows the model to make this same connection, linking the representation of "it" to the representation of "animal."

Self-Attention in Detail: Query, Key, and Value

So, how does this "talking" between vectors work? The self-attention mechanism is elegantly modeled on the concept of a database retrieval system. For every input embedding, we create three distinct vectors:

1. Query (Q): What is this token looking for? It's a question about other parts of the sequence. For the token "it," the query might be, "Which noun in this sentence could be tired?"
2. Key (K): What does this token offer? It acts as a label or identifier for the token's content. For the token "animal," the key might announce, "I am a noun that can possess the state of being tired."
3. Value (V): What information should be passed on? It's the actual content or meaning of the token.

The core idea is to match a Query from one token with the Keys of all other tokens in the sequence. The strength of the match determines how much of each token's Value should be incorporated into the new, updated representation of the querying token.

Jay Alammar's "The Illustrated Transformer" is a classic resource that provides an exceptionally clear, step-by-step walkthrough of this process.

Let's summarize the six steps from the article for calculating the new representation, $z_1$, for the first token in a sequence:

1. Create Q, K, V: For each input embedding (e.g., $x_1,x_2,...$), create a Query, Key, and Value vector by multiplying it with three distinct weight matrices ($W^Q,W^K,W^V$). These matrices are the learnable parameters of the model.
2. Calculate Scores: To determine the relevance of other tokens to token 1, calculate the dot product of its query vector, $q_1$, with the key vector of every token in the sequence ($k_1,k_2,...$). So, the first score is $q_1\cdot k_1$, the second is $q_1\cdot k_2$, and so on. The dot product is large when the vectors are aligned, indicating high relevance.
3. Scale: Divide the scores by the square root of the dimension of the key vectors, $\sqrt{d_k}$. This is a normalization step that helps stabilize the training process, especially with large vector dimensions.
4. Softmax: Apply the softmax function to the scaled scores. This converts the scores into a set of positive numbers that sum to 1, effectively creating a probability distribution. These are the attention weights. A high weight means the corresponding token is very important for understanding the current token.
5. Multiply by Value: Multiply each Value vector ($v_1,v_2,...$) by its corresponding attention weight. This amplifies the Values of relevant tokens and diminishes the Values of irrelevant ones.
6. Sum: Sum up the weighted Value vectors. The result is the output vector, $z_1$, which is a new, context-aware representation of the first token.

This entire process is performed in parallel for every token in the sequence.

The Matrix Calculation

While the step-by-step vector process is great for intuition, in practice, self-attention is implemented using matrix operations for massive parallelization on GPUs. Given your background, this is the most precise representation of the mechanism.

The entire calculation can be condensed into a single, elegant formula:

$\text{Attention}(Q,K,V)=\text{softmax}\left(\frac{Q{K}^T}{\sqrt{d_k}}\right)V$

Here:

• $Q$, $K$, and $V$ are now matrices, where each row represents a token in the sequence.
• $Q{K}^T$ computes the dot product of every query with every key, yielding a matrix of attention scores.
• The softmax is applied row-wise (or column-wise depending on convention) to get the attention weights.
• The final multiplication with $V$ produces the matrix of context-aware output vectors.

Multi-Head Attention: Broadening the Focus

A single self-attention mechanism might learn to focus on one type of relationship, like how "it" refers back to "animal." But language is complex; relationships can be syntactic, semantic, positional, and more. To capture this diversity, Transformers use multi-head attention.

Instead of having just one set of $W^Q,W^K,W^V$ matrices, multi-head attention has multiple sets (e.g., 8 or 96 "heads"). Each head performs the self-attention calculation independently, in parallel. This allows each head to learn to focus on different kinds of contextual relationships.

As the article explains, the output vectors from each of the attention heads are concatenated and then projected back to the original embedding dimension using one more learned weight matrix, $W^O$. This final projection integrates the insights from all the different heads, producing a rich, multi-faceted contextual representation for each token.

Seeing the Math in Action

To solidify these concepts, it's incredibly helpful to see them translated into code. The following video walks through a barebones implementation of self-attention using NumPy, making the abstract matrix operations concrete.

Conclusion

In this lesson, we've pulled back the curtain on self-attention, the engine of contextual understanding in Transformers. You've learned that it's a mechanism for each token to assess the relevance of all other tokens in a sequence. The key takeaways are:

• Self-attention operates on Query, Key, and Value vectors, which are learned projections of the input embeddings.
• It calculates attention weights by taking the scaled dot product of Queries and Keys, followed by a softmax function. These weights represent the "importance" of other tokens.
• The final output for a token is a weighted sum of all Value vectors, allowing it to incorporate information from across the sequence.
• Multi-head attention enhances this by running multiple attention mechanisms in parallel, each learning to focus on different types of contextual relationships.

We now understand the core component of a Transformer block. But how do the weight matrices ($W^Q,W^K,W^V,W^O$) learn the correct values to perform these amazing feats of contextualization? Our next lesson will address this by exploring the pre-training objectives and the immense scale of data used to train these foundational models from scratch.