Understanding Self-Attention: Queries, Keys, Values, and Scaled Dot-Product

Hello! Welcome to your second lesson in the "LLM and Transformer Fundamentals" module.

In our last lesson, we established how word embeddings provide a powerful way to represent individual words as dense vectors that capture semantic meaning. However, as we noted, the meaning of a sentence is far more than the sum of its parts. The context—the surrounding words and their relationships—is critical. "The mole on his face" and "a mole of carbon dioxide" use the same word "mole", but its meaning is disambiguated entirely by context.

Today, we dive into the core mechanism that allows models like those at Anthropic to understand this context: self-attention. This lesson directly addresses the learning outcome: Describe the self-attention mechanism step-by-step: creating queries, keys, and values, and computing scaled dot-product attention. By the end, you'll understand how a model dynamically updates a word's representation by letting it "look at" and "draw information from" other words in the sequence.

1. The Core Idea: What is Self-Attention?

At a high level, self-attention is a mechanism that allows each word in a sequence to weigh the importance of all other words (including itself) and update its own representation accordingly. It's like a networking session for words. For each word, we want to answer three questions:

1. What information am I looking for? (My Query)
2. What information can I offer? (My Key)
3. What information will I actually provide if someone pays attention to me? (My Value)

The model learns to generate these query, key, and value vectors for every word. Then, for each word's query, it computes a similarity score with every other word's key. These scores are turned into weights, which are then used to create a weighted sum of all words' values. The result is a new, contextually-rich vector for that word.

The following video from 3Blue1Brown provides a brilliant intuitive walkthrough of this process.

2. The Mechanics: Scaled Dot-Product Attention

The specific type of self-attention used in the original Transformer paper is called Scaled Dot-Product Attention. Its entire operation can be captured in a single, elegant formula:

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

This might look dense, but it's just a sequence of matrix operations that perfectly mirrors the conceptual steps. Let's break it down.

Step 1: Create Queries, Keys, and Values

We start with the sequence of input word embeddings. For each word's embedding vector $x_i$, we need to generate its query, key, and value vectors. We do this by multiplying the embedding by three distinct weight matrices ($W^Q$, $W^K$, $W^V$) that are learned during the model's training process.

• Query vector: $q_i=x_i{W}^Q$
• Key vector: $k_i=x_i{W}^K$
• Value vector: $v_i=x_i{W}^V$

From your background, you can think of these weight matrices as implementing learnable linear transformations. They project the initial embedding space into three separate, specialized subspaces: a "query space" optimized for asking questions, a "key space" optimized for advertising relevance, and a "value space" optimized for conveying content.

Let's watch how these are generated and what their conceptual roles are.

In practice, we don't process one word at a time. We stack the embedding vectors for the whole sequence into a matrix $X$, and compute all queries, keys, and values in parallel with matrix multiplications:

• $Q=X{W}^Q$
• $K=X{W}^K$
• $V=X{W}^V$

Step 2: Calculate Attention Scores (Dot Product)

To determine how much attention a word should pay to others, we calculate a similarity score between its query and every other word's key. The dot product is a simple and effective measure of vector similarity.

The score for how much word $i$ should attend to word $j$ is: ${\text{score}}_{ij}=q_i\cdot k_j$.

Again, we can do this for all pairs at once by multiplying the query matrix $Q$ by the transpose of the key matrix $K^T$. This gives us a matrix of attention scores.

$$\text{Scores}=Q{K}^T$$

Step 3: Scale the Scores

This step is a crucial detail for stabilizing the training process. The authors of the Transformer paper found that for large vector dimensions ($d_k$), the dot products could grow very large. When fed into a softmax function (our next step), large values can push its output into regions with extremely small gradients, making learning slow or unstable.

To counteract this, we scale the scores by dividing by the square root of the dimension of the key vectors, $\sqrt{d_k}$.

$$\text{Scaled Scores}=\frac{Q{K}^T}{\sqrt{d_k}}$$

This is a practical engineering choice to ensure numerical stability—a trade-off between raw mathematical expression and trainability.

Step 4: Normalize with Softmax (and Masking)

We now have scaled scores, but they are just arbitrary real numbers. To turn them into a useful probability distribution of weights, we apply the softmax function to each row of the scaled score matrix. This ensures that for each query word, the attention weights it assigns to all key words are positive and sum to 1.

$$\text{AttentionWeights}=\mathrm{softmax}\left(\frac{Q{K}^T}{\sqrt{d_k}}\right)$$

A quick but critical note on Masking: In decoder models like GPT that generate text one word at a time, we must prevent a word from "seeing into the future". During training, we want the prediction for word $i$ to only depend on words $1,...,i$. We achieve this by "masking" the attention scores. Before the softmax step, we add $-\infty$ to all scores that correspond to a word attending to a future word. After softmax, these positions become 0, effectively blocking information flow from the future.

Step 5: Compute the Output (Weighted Sum of Values)

Finally, we compute the output vector for each word. This new vector, $z_i$, is a weighted sum of all the value vectors in the sequence, where the weights are the attention weights we just calculated.

$$z_i=\sum _j{\text{AttentionWeight}}_{ij}\cdot v_j$$

Words that are highly relevant (high attention weight) contribute more of their "value" to the new representation. In matrix form, this is a single multiplication:

$$Z=\text{AttentionWeights}\cdot V$$

The resulting matrix $Z$ contains the new, context-aware embeddings for every word in the sequence. Each row $z_i$ has incorporated relevant information from its surrounding context.

Let's watch the final steps of this process.

The image below provides a fantastic, compact visualization of the entire calculation for a single word.

3. A Practical Implementation in PyTorch

Given your background, seeing how this translates directly into code can be very clarifying. The entire attention formula is a sequence of highly optimized matrix operations, which is why GPUs are so effective at running these models.

The following resource provides a clean, well-commented Python function that implements exactly what we have just discussed.

The beauty of this implementation is its directness. There is no complex control flow; it is a straightforward data-flow graph of tensor operations, which maps beautifully to modern hardware and aligns with the principles of building efficient, scalable systems.

Conclusion

In this lesson, we have demystified the self-attention mechanism, the engine of contextual understanding in Transformers. We've moved beyond static word embeddings to see how models can dynamically build richer representations based on the surrounding text.

Key Takeaways:

• Self-attention allows a model to weigh the importance of different words in a sequence when representing a single word.
• The process involves projecting each input embedding into a Query, a Key, and a Value.
• The step-by-step calculation for Scaled Dot-Product Attention is:
  1. Compute similarity scores via dot product: $Q{K}^T$.
  2. Scale for numerical stability: divide by $\sqrt{d_k}$.
  3. Optionally mask future positions.
  4. Normalize scores into weights via softmax.
  5. Compute the output as a weighted sum of the Values: $\text{softmax}(\dots )V$.
• The entire mechanism is just a series of matrix multiplications, making it highly parallelizable and efficient on modern hardware like GPUs.

Next Lesson Preview

We've just seen how a single attention "head" works. It learns to focus on one type of relationship within the text (e.g., how verbs relate to their subjects). But language is complex, with many overlapping relationships (adjective-noun modification, pronoun-antecedent links, etc.). To capture this richness, Transformers don't use just one attention mechanism.

In the next lesson, we will explore multi-head attention, which runs several attention mechanisms in parallel, allowing the model to jointly attend to information from different representational subspaces at different positions. This is the final piece we need to understand the complete attention block as it's used in models like GPT.