Transformer Attention: A Guide to the Q, K, and V Matrices

Understanding the Transformer Attention Mechanism

Transformers have revolutionized the way machines process language and other sequential data. At the heart of the Transformer architecture is a powerful mechanism called self-attention that was first described in the paper "Attention is All You Need." This self-attention mechanism allows the model to focus on different parts of the input sequence and weigh their importance when making predictions. To fully understand how this works, we need to dive into the matrices that drive it: Q (Query), K (Key), and V (Value). 

But I have found that understanding the Q, K, and V matrices to be the most difficult part of the transformer model. It's not the math that is difficult, but what is difficult is understanding the "why" as much as the "how." Why do these matrices work? What do each of the matrices do? Why are there even three matrices? What is the intuition for all of this?

Okay so let's get started with a simple analogy:

Imagine you’re at a library, searching for books on a particular topic. You have a query in mind (what you're looking for) and the librarian has organized the library catalog by key attributes, such as genre, author, or publication date. Based on how well the attributes match your query, the librarian assigns a score to each book. Once the books are scored, the librarian returns the value—the actual content or summary of the top-scoring books that best match your query.

In this analogy:

• Query (Q) is what you are searching for.
• Key (K) represents the attributes of the books that help in scoring.
• Value (V) is the information or content you get back from the top-matching books.

Now, let’s break down how these ideas translate to the actual self-attention mechanism in Transformers.

Self-Attention: The Basics

In self-attention, each word in a sentence (or token in a sequence) will interact with every other word to figure out how important they are to each other. For each word, a query, key, and value vector is created. The attention mechanism then works by calculating the importance of each word (key) to the word currently being processed (query), and using this information to weigh the corresponding values.

Let's say we have the sentence*:

"The cat sat on the mat."

Each word here will get its own Q, K, and V representation. The goal of the self-attention mechanism is to compute how much each word should attend to other words when making a prediction.

Breaking Down the Q, K, and V Matrices

1. Query (Q): What am I looking for?

The query represents the word we’re focusing on and asks the rest of the sentence, "How relevant are you to me?" Each word generates a different query matrix, and the higher the match with the keys, the more attention it gives to other words.

For example, let’s say our query is the word "cat." We want to know which other words in the sentence provide important information about the word "cat."

2. Key (K): What features do I have?

The key represents the characteristics of each word. Think of the key as each word shouting out, "Here’s what I’m about!" Other words in the sentence will compare their query against these keys to see if they should focus on them.

So, when we look at the key of "mat," it tells us something about the word's identity (perhaps it's an object or a location). Similarly, the key for "cat" might represent something like "animal" or "subject."

3. Value (V): What information do I carry?

The value contains the actual information of each word, like its meaning in the context of the sentence. Once the model has determined which words are important using the query-key matching process, it uses the value to inform the prediction.

For instance, if the query "cat" finds that "sat" is important based on the key, it will give more weight to the value of "sat" to help predict what comes next in the sentence.

Calculating Attention: Putting Q, K, and V Together

The actual attention score is calculated by taking the dot product of the query with all the keys. This gives us a score for how much focus the word (query) should place on each other word (key). The higher the score, the more attention that word receives.

Here’s a high level look at the math we are going to do:

1. Dot product of Q and K: The query matrix of a word is multiplied with the key matrices of all the words in the sequence. This gives a score representing how much each word in the sentence should attend to the current word.

2. Softmax: These scores are then passed through a softmax function, which normalizes them into probabilities (weights) that sum to 1. This step ensures that the attention is distributed in a meaningful way across all words.

3. Weighted Sum of Values: The resulting attention weights are multiplied by the value matrices. This weighted sum gives us the final output for the word, which is used in the next layer of the Transformer model.

Example: "The cat sat on the mat."

Let’s walk through how the word "cat" might process the sentence using self-attention:

1. Query (Q): The model generates a query matrix for "cat," representing what it’s looking for (e.g., context about an action related to the "cat").

2. Key (K): Each word in the sentence has its own key. The word "sat," for instance, might have a key that highlights it as an action verb, making it relevant to the "cat."

3. Dot Product: The query for "cat" is compared (via dot product) with the keys of all the words in the sentence. If "sat" has a high dot product with the query for "cat," it will get a high attention score.

4. Softmax: The scores for all the words are normalized into probabilities, so "sat" might get a large share of the attention.

5. Value (V): The values of the words are then weighted by the attention scores. Since "sat" got a high score, its value (which could include the action or tense) will have a bigger impact on the final representation of the word "cat."

The self-attention mechanism allows the Transformer to look at all parts of a sequence simultaneously and decide which parts are most important to focus on. This is especially powerful for tasks like translation, summarization, and language understanding because it doesn’t rely on processing the input one word at a time. Instead, it lets each word interact with every other word in the sequence, leading to a richer, more flexible understanding of context.

The transformer model is able to "pay attention" to the right information, just like a librarian matching your search with the right books. 

Let's walk through the math:

To make the Transformer self-attention mechanism more concrete, let's work through a simplified example using the sentence:

"The cat sat on the mat."

We'll assign simple numerical values to create embeddings, compute the Q (Query), K (Key), and V (Value) matrices, and see how the attention mechanism operates step by step.

Simplifications for the Example

• Embedding Dimension: We'll use a small embedding size of 2 to keep calculations manageable. In real-world Transformer models, the embedding size is much larger to capture the complex semantic and syntactic nuances of language. These embeddings are learned during the training process, allowing the model to position semantically similar words closer together. Actual embeddings in real models have much larger dimensions (e.g., 512, 768, 2048, and higher) and are learned in a separate process from attention. But by using low-dimensional vectors it will highlight for us how the Query (Q), Key (K), and Value (V) matrices interact during the attention process.
• Weights: We'll define simple weight matrices for Q, K, and V transformations.

Before we get into the step by step walkthrough of how attention is derived, a visual way to think of it is imagining the embeddings as vectors in a high-dimensional space. The weight matrices rotate, scale, or skew these vectors into new configurations (Q, K, V spaces). These transformations adjust the vectors so that the dot products between Query and Key vectors effectively measure the relevance or similarity between tokens. This alignment allows the model to compute attention scores that highlight important relationships, enabling it to determine which tokens are most significant to each other within the sequence. By doing so, the model can accurately capture complex dependencies and contextual nuances, such as grammatical structures and semantic meanings, enhancing its understanding of the input data. 

Step 1: Assign Word Embeddings

First, we assign embeddings to each word in the sentence. Again we are using simple pretend embeddings of size 2. A real embedding for cat might look something like: Embedding (E_cat):  [0.12, -0.03, 0.45, …, 0.07]

Okay, let's define our simple embeddings as follows:

Word	Embedding (E)
The	[1, 0]
cat	[0, 1]
sat	[1, 1]
on	[0, -1]
the	[1, 0]
mat	[0, 1]

(Note: For simplicity, "The" and "the" are treated the same.)

Step 2: Define Weight Matrices for Q, K, and V

We'll define weight matrices that transform embeddings into Q, K, and V matrices. These would be learned during training and would be floating point values. And again we are going to make up some numbers and keep the numbers simple.

Assume the weight matrices are as follows:

• W_Q (2x2 matrix): $$W_Q = \begin{bmatrix}1 & 1 \\ 1 & -1\end{bmatrix}$$
• W_K (2x2 matrix): $$W_K = \begin{bmatrix}1 & 1 \\ 1 & -1\end{bmatrix}$$
• W_V (2x2 matrix): $$W_V = \begin{bmatrix}1 & 1 \\ 1 & -1\end{bmatrix}$$

Step 3: Compute Q, K, and V for Each Word

For each word, we'll compute:

• Q_i = E_i * W_Q
• K_i = E_i * W_K
• V_i = E_i * W_V

Let's compute these for each word.

Word: "The"

Embedding (E_the): [1, 0]

Compute Q_the:

$$Q_{\text{the}} = E_{\text{the}} \times W_Q = [1, 0] \times \begin{bmatrix}1 & 0 \\ 0 & 1\end{bmatrix} = [1, 0]$$

Compute K_the:

$$K_{\text{the}} = E_{\text{the}} \times W_K = [1, 0] \times \begin{bmatrix}0 & 1 \\ 1 & 0\end{bmatrix} = [0, 1]$$

Compute V_the:

$$V_{\text{the}} = E_{\text{the}} \times W_V = [1, 0] \times \begin{bmatrix}1 & 1 \\ 1 & -1\end{bmatrix} = [1, 1]$$

Word: "cat"

Embedding (E_cat): [0, 1]

Compute Q_cat:

$$Q_{\text{cat}} = [0, 1] \times \begin{bmatrix}1 & 0 \\ 0 & 1\end{bmatrix} = [0, 1]$$

Compute K_cat:

$$K_{\text{cat}} = [0, 1] \times \begin{bmatrix}0 & 1 \\ 1 & 0\end{bmatrix} = [1, 0]$$

Compute V_cat:

$$V_{\text{cat}} = [0, 1] \times \begin{bmatrix}1 & 1 \\ 1 & -1\end{bmatrix} = [1, -1]$$

Word: "sat"

Embedding (E_sat): [1, 1]

Compute Q_sat:

$$Q_{\text{sat}} = [1, 1] \times \begin{bmatrix}1 & 0 \\ 0 & 1\end{bmatrix} = [1, 1]$$

Compute K_sat:

$$K_{\text{sat}} = [1, 1] \times \begin{bmatrix}0 & 1 \\ 1 & 0\end{bmatrix} = [1, 1]$$

Compute V_sat:

$$V_{\text{sat}} = [1, 1] \times \begin{bmatrix}1 & 1 \\ 1 & -1\end{bmatrix} = [2, 0]$$

Word: "on"

Embedding (E_on): [0, -1]

Compute Q_on:

$$Q_{\text{on}} = [0, -1] \times \begin{bmatrix}1 & 0 \\ 0 & 1\end{bmatrix} = [0, -1]$$

Compute K_on:

$$K_{\text{on}} = [0, -1] \times \begin{bmatrix}0 & 1 \\ 1 & 0\end{bmatrix} = [-1, 0]$$

Compute V_on:

$$V_{\text{on}} = [0, -1] \times \begin{bmatrix}1 & 1 \\ 1 & -1\end{bmatrix} = [-1, 1]$$

Word: "the" (Again)

Same as before for "The".

Word: "mat"

Embedding (E_mat): [0, 1]

Compute Q_mat:

$$Q_{\text{mat}} = [0, 1] \times \begin{bmatrix}1 & 0 \\ 0 & 1\end{bmatrix} = [0, 1]$$

Compute K_mat:

$$K_{\text{mat}} = [0, 1] \times \begin{bmatrix}0 & 1 \\ 1 & 0\end{bmatrix} = [1, 0]$$

Compute V_mat:

$$V_{\text{mat}} = [0, 1] \times \begin{bmatrix}1 & 1 \\ 1 & -1\end{bmatrix} = [1, -1]$$

Step 4: Compute Attention Scores

Now, we'll compute the attention scores for a target word. Let's focus on the word "cat" and see how it attends to other words in the sentence.

For the word "cat", we have:

• Q_cat = [0, 1]

We will compute the attention scores between "cat" and each word in the sentence by taking the dot product of Q_cat with K_i for each word.

Calculating Dot Products

1. Score between "cat" and "The":

$$\text{Score}_{\text{cat, The}} = Q_{\text{cat}} \cdot K_{\text{the}} = [0, 1] \cdot [0, 1] = (0 \times 0) + (1 \times 1) = 1$$

2. Score between "cat" and "cat":

$$\text{Score}_{\text{cat, cat}} = Q_{\text{cat}} \cdot K_{\text{cat}} = [0, 1] \cdot [1, 0] = (0 \times 1) + (1 \times 0) = 0$$

3. Score between "cat" and "sat":

$$\text{Score}_{\text{cat, sat}} = [0, 1] \cdot [1, 1] = (0 \times 1) + (1 \times 1) = 1$$

4. Score between "cat" and "on":

$$\text{Score}_{\text{cat, on}} = [0, 1] \cdot [-1, 0] = (0 \times -1) + (1 \times 0) = 0$$

5. Score between "cat" and "the":

        Same as with "The":

$$\text{Score}_{\text{cat, the}} = 1$$

6. Score between "cat" and "mat":

$$\text{Score}_{\text{cat, mat}} = [0, 1] \cdot [1, 0] = 0$$

Summary of Scores

Pair	Score
cat & The	1
cat & cat	0
cat & sat	1
cat & on	0
cat & the	1
cat & mat	0

** See note below about scaling these values

Step 5: Apply Softmax to Obtain Attention Weights

Next, we apply the softmax function to these scores to get attention weights.

The softmax function is defined as:

$$\text{softmax}(x_i) = \frac{e^{x_i}}{\sum_{j} e^{x_j}}$$

Compute the exponentials (this is easy and obvious with our numbers):

• $e^1$
• $e^0$

So the exponentials of the scores are:

Pair	Score	Exponential
cat & The	1	2.718
cat & cat	0	1
cat & sat	1	2.718
cat & on	0	1
cat & the	1	2.718
cat & mat	0	1

Compute the sum of exponentials:

$$\text{Sum} = 2.718 + 1 + 2.718 + 1 + 2.718 + 1 = 11.154$$

Compute attention weights:

• Weight(cat, The):

$$\alpha_{\text{cat, The}} = \frac{2.718}{11.154} \approx 0.244$$

• Weight(cat, cat):

$$\alpha_{\text{cat, cat}} = \frac{1}{11.154} \approx 0.090$$

• Weight(cat, sat):

$$\alpha_{\text{cat, sat}} = \frac{2.718}{11.154} \approx 0.244$$

• Weight(cat, on):

$$\alpha_{\text{cat, on}} = \frac{1}{11.154} \approx 0.090$$

• Weight(cat, the):

$$\alpha_{\text{cat, the}} = \frac{2.718}{11.154} \approx 0.244$$

• Weight(cat, mat):

$$\alpha_{\text{cat, mat}} = \frac{1}{11.154} \approx 0.090$$

Summary of Attention Weights

Pair	Weight (α)
cat & The	0.244
cat & cat	0.090
cat & sat	0.244
cat & on	0.090
cat & the	0.244
cat & mat	0.090

Step 6: Compute the Weighted Sum of Values

Now, we use the attention weights to compute the weighted sum of the Value vectors.

Recall the Value vectors:

• V_The: [1, 1]
• V_cat: [1, -1]
• V_sat: [2, 0]
• V_on: [-1, 1]
• V_the: [1, 1]
• V_mat: [1, -1]

Compute the weighted sum:

$$\text{Output}_{\text{cat}} = \sum_{i} \alpha_{\text{cat, i}} \times V_i$$

Compute each term:

1. cat & The:

$$0.244 \times [1, 1] = [0.244, 0.244]$$

2. cat & cat:

$$0.090 \times [1, -1] = [0.090, -0.090]$$

3. cat & sat:

$$0.244 \times [2, 0] = [0.488, 0.000]$$

4. cat & on:

$$0.090 \times [-1, 1] = [-0.090, 0.090]$$

5. cat & the:

$$0.244 \times [1, 1] = [0.244, 0.244]$$

6. cat & mat:

$$0.090 \times [1, -1] = [0.090, -0.090]$$

Add up all these vectors:

$$\begin{align*} \text{Output}_{\text{cat}} &= [0.244, 0.244] + [0.090, -0.090] + [0.488, 0.000] \\ &\quad + [-0.090, 0.090] + [0.244, 0.244] + [0.090, -0.090] \\ &= [(0.244 + 0.090 + 0.488 - 0.090 + 0.244 + 0.090), \\ &\quad (0.244 - 0.090 + 0.000 + 0.090 + 0.244 - 0.090)] \\ &= [1.066, 0.398] \end{align*}$$

So the output vector for "cat" after the attention mechanism is [1.066, 0.398].

Step 7: Interpretation

The output vector [1.066, 0.398] is a context-aware representation of the word "cat". It has incorporated information from other relevant words in the sentence, weighted by their importance as determined by the attention mechanism.

• The higher weights given to "The", "sat", and "the" reflect their relevance to "cat" in this context.
• The contributions from "on" and "mat" are smaller due to lower attention weights.

Generalizing to All Words

In a real Transformer, this process is performed for each word in the sentence, allowing every word to attend to every other word and capture the contextual relationships.

Some Almost Final Words About Attention

Earlier in this post, I said that:

Q represents "What am I looking for?"

K represents "What features do I have?"

V represents "What information do I carry?

But how exactly does Q, K, and V represent these questions?

We can answer the first two questions by considering the dot product. The dot product between Q_i and K_j measures the similarity between  Q_i  and  K_j. A higher dot product indicates a higher relevance or alignment between what token i is seeking and what token j offers. The dot product effectively answers: “How much does what I’m looking for (Q) align with what features you have (K)?”

V_j  is weighted by the attention scores α_ij and aggregated to form the output. These V_j vectors hold the information that is actually used to update or inform token i’s representation - the  V_j  vectors are the actual data that get combined to form the new representation of token i. In other words, after determining which tokens are relevant (via Q and K), the model needs to know what information to extract—this is provided by V.

Conclusion

Through this example, we've illustrated how:

• Embeddings are transformed into Q, K, and V matrices using learned weight matrices.
• Attention scores are computed using the dot product of Q and K.
• Attention weights are derived by applying the softmax function to the scores.
• Weighted sums of the Value vectors produce the output attention representations for each word.

This simplified demonstration shows how the self-attention mechanism enables a word to focus on relevant parts of the input sequence, effectively capturing the context needed for understanding and generating language.

Additional Resources

Here are some other resources beyond the original Attention paper that helped me in my understanding:

• "The Illustrated Transformer" by Jay Alomar is a well known walk through of the entire transformer architecture.
• "All the Math You Need in 15 Minutes" by ritvikmath has a very intuitive walkthrough of the attention mechanism.
• The videos by 3Blue1Brown are fantastic that start with explaining neural networks to how transfromers work:
• "But What is a Neural Network? | Chapter 1: Deep Learning"
• "Gradient Descent: How Neural Networks Learn | Chapter 2: Deep Learning"
• "What is Backpropagation really doing? | Chapter 3: Deep Learning"
• "Backpropagation Calculus | Chapter 4: Deep Learning"
• "How Large Language Models Work - a Visual Intro to Transformers | Chapter 5: Deep Learning"
• "Attention in Transformers: Visually Explained | Chapter 6: Deep Learning"
• "How Might LLMs Store Facts | Chapter 7: Deep Learning"
• "Karpathy" videos, Andrej Karpathy has some great videos on understanding GPT models and embeddings, such as here, here, here and his "makemore" series starting here.
• Build a Large Language Model (from Scratch) by Sebastian Raschka. This book is just coming out. I don't have it yet, but it looks fantastic and I will be getting it. There is a github repository for the code in the book that is really good.

*This sentence, "The cat sat on the mat", I consider to be a well know example going back at least five years to papers on BERT and GPT2. This might be the earliest example of this sentence being used in a paper called "A Multiscale Visualization of Attention in the Transformer Model" by Jesse Vig.

**In high-dimensional vector spaces, which is the norm in transformer models, the dot product of two random vectors tends to have a larger magnitude because each dimension contributes to the total. This can result in attention scores that are large, pushing the softmax function into regions where it outputs very small gradients. Small gradients slow down learning because the model updates are minimal. By scaling down the dot products, we lesson this effect. The scaling factor √d effectively controls the variance of the dot product by scaling the dot product by the square root of the dimensionality of the Key vectors. This keeps the attention scores at a scale where the softmax function operates optimally, and the gradients remain at a magnitude conducive to learning. This isn't a problem in our trivial example here of vectors of size 2 so I chose not to put that in.

Here is the full attention formula where Q_i and K_j are scaled by √d before having the softmax applied: $$\alpha_{ij} = \text{softmax}\left( \frac{Q_i \cdot K_j}{\sqrt{d_k}} \right)$$