Understanding Self-Attention

Introduction

In our last lesson, we did a deep dive into the scaled dot-product attention formula:

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

We focused exclusively on the scaling factor, $\frac{1}{\sqrt{d_k}}$, and understood its critical role in stabilizing training by preventing vanishing gradients in the softmax function.

However, we left some fundamental questions unanswered:

• What are the Query (Q), Key (K), and Value (V) matrices?
• Where do they come from?
• And why is this mechanism called self-attention?

Today's lesson will answer these questions. Our learning objective is to describe self-attention and how it allows each token to attend to all other tokens in a sequence. We will build the entire mechanism from the ground up, starting with a powerful intuition.

1. The Core Problem: Words Need Context

Imagine the word "bank". Its meaning is ambiguous without context.

• "I swam across the river to get to the other bank." (a riverbank)
• "I drove across the road to get to the other bank." (a financial institution)

A simple word embedding provides a static representation, which struggles to capture these contextual nuances. The goal of self-attention is to create a new, contextualized representation for each word by allowing it to "look at" and "draw information from" other words in the same sentence.

To see this idea explained, let's watch a short, intuitive video.

2. The Intuition: A Weighted Sum of Context

The video frames the core idea perfectly. How can we enhance the embedding of "bank"? We can create a new embedding for it that is a weighted sum of the embeddings of all words in the sentence.

For the sentence "River bank flows," the new, contextualized embedding for "bank" would be:

$\text{new\_embedding(’bank’)}=w_1\cdot \text{emb(’River’)}+w_2\cdot \text{emb(’bank’)}+w_3\cdot \text{emb(’flows’)}$

The key question is: How do we calculate the weights ($w_1,w_2,w_3$)?

Intuitively, the weights should reflect relevance. For the word "bank," the weight for "River" should be high, while the weight for a less relevant word might be low. The most straightforward way to measure relevance or similarity between two vectors is the dot product. A larger dot product implies greater similarity.

This leads to a simple, two-step process for each word in the sentence:

1. Compute Scores: Calculate the dot product of the target word's embedding with every other word's embedding in the sentence. This gives us a set of raw "relevance scores."
2. Normalize to Weights: Pass these scores through a softmax function. This converts the raw scores into a probability distribution, ensuring the weights sum to 1.

The following reading walks through this exact first-principles approach.

3. The Mechanism: Query, Key, and Value

Our simple weighted-sum model is powerful, but it has a limitation: a word's embedding is used for three different purposes:

1. To compare against other words to generate scores.
2. To be compared by other words.
3. As the content that gets summed up in the final representation.

To make the model more flexible and powerful, self-attention introduces a small but crucial modification. Instead of using the same embedding vector for all three roles, it first projects the original embedding into three separate vectors:

• Query ($q$): Represents the current word's "question" as it looks for context. It asks: "Who is relevant to me?"
• Key ($k$): Represents a word's "label" or "identity" that can be matched against a query. It answers: "Here's what I am; see if I'm relevant to you."
• Value ($v$): Represents the actual content or information of the word. This is what gets passed into the final weighted sum if the key is deemed relevant.

The process then becomes:

1. A token's Query vector is compared (via dot product) with every other token's Key vector to compute attention scores.
2. The scores are passed through softmax to get the weights.
3. The weights are used to create a weighted sum of all tokens' Value vectors.

For a great analogy, let's watch the following clip.

Making it Learnable

But where do these Q, K, and V vectors come from? They are generated by multiplying the original input embedding ($x$) by three separate, learnable weight matrices ($W_Q,W_K,W_V$):

$q=x\cdot W_Q$
$k=x\cdot W_K$
$v=x\cdot W_V$

This is the most important step! These weight matrices are the learnable parameters of the self-attention layer. During training, the model learns the best way to project the input embeddings into Q, K, and V spaces to solve the task at hand (e.g., translation, summarization).

The following resources explain this critical step. The video provides an excellent visual walkthrough, while the article gives the same information in a text format with code snippets.

4. Putting It All Together

Now we can assemble the complete picture of self-attention for a sequence of tokens. The process happens for every token in parallel.

For any given token in the input sequence, its final contextualized representation is computed by:

1. Projecting its embedding into a Query vector.
2. Projecting every token's embedding (including its own) into Key and Value vectors.
3. Calculating the dot product of its Query with every other token's Key to get attention scores.
4. Scaling the scores by $\frac{1}{\sqrt{d_k}}$ and applying softmax to get the final attention weights.
5. Multiplying these weights by the corresponding Value vectors and summing them up to get the final context vector.

This diagram provides a fantastic overview of the entire process.

To see this process in action on a matrix-level, the following video provides a comprehensive animation.

Because each token's Query vector is compared against every other token's Key vector, this mechanism fully enables each token to attend to all other tokens in the sequence, fulfilling our learning objective.

Conclusion

In this lesson, we deconstructed the self-attention mechanism, the engine that powers modern Transformers. We went from the high-level need for context to the specific mechanics of Queries, Keys, and Values.

Key Takeaways:

• Purpose: Self-attention creates contextualized word embeddings by allowing each token to dynamically draw information from all other tokens in the sequence.
• Mechanism: It achieves this by projecting each input embedding into three roles: a Query (what I'm looking for), a Key (what I have), and a Value (what I'll give).
• Calculation: For each token, its Query is compared to every other token's Key to compute attention weights. These weights then determine how much of each token's Value contributes to the final output embedding.
• Learnability: The projections into Q, K, and V spaces are done via learnable weight matrices ($W_Q,W_K,W_V$), which allows the model to learn task-specific contextual relationships.

Preview of the Next Lesson:
We've established that self-attention allows a token to look at all other tokens, including those that come after it. For tasks like sentence classification, this is great. But for autoregressive language modeling—where we predict the next word one at a time—allowing a token to "see the future" would be cheating.

In our next lesson, we will explore causal (or masked) self-attention, a modification that prevents tokens from attending to subsequent positions in the sequence, making it suitable for generative tasks.