Self-Attention: The Heart of the Transformer

Hello! Welcome to the fifth lesson in our module on the fundamentals of Transformers.

In our last lesson, we saw how positional encodings are added to token embeddings to give the model a sense of word order. Our input is now a sequence of vectors, each rich with both semantic meaning and positional information.

This lesson dives into the computational heart of the Transformer: self-attention. This is the mechanism that takes those enriched vectors and allows the model to dynamically weigh the importance of every other word in the sequence when processing a specific word. We will be covering the learning outcome: Explain the self-attention mechanism as the core of the transformer architecture.

By the end of this lesson, you will understand how a model can figure out that in the sentence "The animal didn't cross the street because it was too tired," the word "it" refers to "animal" and not "street."

1. The Core Idea: Creating Context-Aware Representations

At a high level, self-attention is a process where the representation of each word is updated by incorporating information from all other words in the same sequence. Words that are more relevant get higher weight, and less relevant words get lower weight. This creates a new, "context-aware" representation for each token.

Let's begin with a video from 3Blue1Brown that provides an excellent intuition for the problem that self-attention solves and a high-level overview of how it works.

As the video explains, the initial embedding for a word like "mole" is generic. Self-attention refines this embedding by "looking" at the surrounding words and pulling in relevant information to create a more specific, contextualized meaning.

Here is a classic visualization showing the result of this process. When the model processes the word "it," self-attention has correctly assigned most of its focus to "The" and "animal".

Now, let's unpack the mechanism that produces these attention scores.

2. The Mechanism: Queries, Keys, and Values

To decide which words are relevant to which others, self-attention uses an analogy from databases and information retrieval systems, which will be familiar from your software engineering background. It creates three special vectors for each input token: a Query, a Key, and a Value.

• Query (Q): This vector represents the current word's "question." Think of it as the word asking, "What in this sentence is relevant to me?" For the word "it," the query might be something like, "I am a pronoun; I am looking for a noun I might refer to."
• Key (K): This vector is like a word's "label" or "advertisement." It says, "This is what I am." For the word "animal," the key might signal, "I am a noun."
• Value (V): This vector contains the actual semantic content of the word. It's the information that will be passed on if the word is deemed relevant.

These three vectors are generated by multiplying the input embedding (which already includes positional encoding) by three distinct weight matrices: $W_Q$, $W_K$, and $W_V$. These matrices are learned during the model's training process.

The entire self-attention calculation can be broken down into a few steps:

Let's walk through the steps shown in the diagram.

The Step-by-Step Calculation

1. Generate Q, K, V: For every input token, multiply its embedding by the $W_Q$, $W_K$, and $W_V$ matrices to get its Q, K, and V vectors. Since this is just a matrix multiplication, it can be done in parallel for all tokens.

2. Calculate Scores: To find the relevance of word 2 to word 1, we calculate the dot product of word 1's Query vector with word 2's Key vector: $Score=Q_1\cdot K_2$. This is done for all pairs of words, creating a score matrix. A higher dot product means the key is a better "match" for the query.

3. Scale: The scores are divided by the square root of the dimension of the key vectors, $\sqrt{d_k}$. This is a technical but important step for stabilizing the training process. Without it, for large values of $d_k$, the dot product scores could become very large, pushing the softmax function (our next step) into regions with extremely small gradients, which stalls learning.

4. Normalize with Softmax: We apply a softmax function to the scaled scores. This converts the scores into a set of positive weights that sum to 1.0, which we can think of as a "focus distribution." A word with a high score will get a weight close to 1.0, while irrelevant words will get weights near 0.0.

5. Compute the Output: The final representation for a word is a weighted sum of all the Value vectors in the sequence. The weights used in this sum are the softmax scores we just calculated. In this way, the model "pays attention" only to the words that are relevant, effectively drowning out the noise from others.

To see this process in action with a more detailed walkthrough, the next video explains the roles of Query, Key, and Value, and how they combine to produce the final output.

3. From Vectors to Matrices: Efficient Computation

Performing these calculations vector-by-vector in loops would be inefficient. One of the key strengths of the Transformer is that the entire self-attention mechanism can be expressed as a set of highly optimized matrix multiplications. This is what allows it to leverage GPUs for massive parallelization.

If we pack all our input embeddings into a single matrix $X$, we can compute the Q, K, and V matrices for the entire sequence at once:
$Q=X{W}_Q$
$K=X{W}_K$
$V=X{W}_V$

The full self-attention calculation can then be elegantly expressed in a single line:

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

This formula directly corresponds to the steps we just discussed and is how self-attention is implemented in practice. The article below provides a fantastic, code-oriented walkthrough that connects these concepts directly to a PyTorch implementation.

The article's PyTorch class, SelfAttention, is a perfect summary of what we've learned. It takes the input embeddings, creates the Q, K, and V projections using linear layers, computes the scaled dot-product attention, and returns the final context vectors.

4. A Glimpse into Multi-Head Attention

So far, we've discussed a single attention mechanism. This is called a single "attention head." In practice, what if there are multiple types of relationships in a sentence? For example, in "The fluffy blue creature happily ate the berry," we might want to link:

• "fluffy" and "blue" to "creature" (adjective-noun)
• "happily" to "ate" (adverb-verb)
• "ate" to "creature" and "berry" (subject-verb-object)

A single attention head might learn to focus on one of these relationship types but struggle to capture all of them simultaneously. The solution is Multi-Head Attention. The model simply runs multiple self-attention processes (heads) in parallel, each with its own set of learned $W_Q$, $W_K$, and $W_V$ matrices.

Each head can learn to focus on different types of relationships. The outputs from all heads are then concatenated and combined to produce the final, rich output for that layer.

This idea of running multiple attention heads in parallel gives the model a much broader capacity to understand complex linguistic patterns.

Conclusion

In this lesson, we unpacked the self-attention mechanism, the engine that powers the Transformer's contextual understanding. You learned how it uses Queries, Keys, and Values to calculate relevance scores and produce context-aware word representations.

Key Takeaways:

• Self-attention updates each token's representation by calculating a weighted sum of all other tokens in the sequence.
• The process uses three learned projections of the input: Query (what a word is looking for), Key (what a word offers), and Value (the content a word provides).
• Relevance scores are calculated via a scaled dot-product between Queries and Keys.
• A softmax function converts these scores into attention weights, which determine how much "focus" to place on each word's Value.
• The entire mechanism is implemented with matrix operations, making it highly efficient and parallelizable on modern hardware.

We concluded with a preview of our next topic. While a single attention head is powerful, running multiple heads in parallel is even better. In the next lesson, we will explore multi-head attention in more detail and see how it allows the model to capture diverse and complex patterns within the text.