Query, key, value matrices

What Are Q, K, V Matrices?

• Query matrix $Q = XW^Q$ where $W^Q \in \mathbb{R}^{d_{model} \times d_k}$
• Key matrix $K = XW^K$ where $W^K \in \mathbb{R}^{d_{model} \times d_k}$
• Value matrix $V = XW^V$ where $W^V \in \mathbb{R}^{d_{model} \times d_v}$

These are learned linear transformations that project each token into three different subspaces optimized for different purposes in the attention mechanism.

Why Three Separate Matrices?

WHY separate transformations?

1. Decoupling roles: The same token embedding serves different purposes. When token A attends to token B, we need:

   • A's query: "What am I looking for?"
   • B's key: "What do I represent for matching?"
   • B's value: "What information should I contribute?"

2. Asymetric relationships: The similarity between "cat" and "sat" should be computed differently than the information "cat" provides. Keys handle matching geometry, values handle information geometry.

3. Optimization flexibility: Different weight matrices allow the network to learn that "cat" might query for verbs (high similarity with "sat" in Q-K space) but provide noun information (in V space).

Derivation from First Principles

Starting Point: What Problem Are We Solving?

We want each token to attend to relevant context. Raw embedings $x_i \in \mathbb{R}^{d_{model}}$ contain all semantic information mixed together. We need to:

Step 1: Define "relevance" as a similarity score

For token $i$ to determine relevance of token $j$, we compute: $\text{relevance}_{i,j} = \text{similarity}(x_i, x_j)$

But raw dot product $x_i \cdot x_j$ is too rigid—it uses the same dimensions for similarity that encode position, syntax, semantics, etc.

WHY this fails: The dimensions optimized for "word meaning" aren't necessarily optimized for "should attend to this."

Step 2: Project into a specialized "matching space"

Instead, transform $x_i$ into a query vector $q_i$ that represents "what I'm looking for": $q_i = W^Q x_i$

where $W^Q \in \mathbb{R}^{d_k \times d_{model}}$ is learned to extract "search intent."

WHY: Now $q_i$ lives in a $d_k$-dimensional space specifically optimized for computing relevance (typically $d_k < d_{model}$ for efficiency and preventing overfitting).

Step 3: Project comparison targets into a "key space"

Transform $x_j$ into a key vector $k_j$ that represents "what I offer": $k_j = W^K x_j$

Now relevance is computed as: $\text{relevance}_{i,j} = q_i \cdot k_j = (W^Q x_i)^T (W^K x_j)$

WHY separate $W^Q$ and $W^K$: Attention is directional. "The cat sat" → "cat" querying "sat" is different from "sat" querying "cat." Asymetric matrices encode asymetric relationships.

Step 4: Separate the information retrieval

Once we know relevance, we need to extract information. But the information we want from token $j$ might be different from what made it relevant!

Transform $x_j$ into a value vector $v_j$: $v_j = W^V x_j$

where $W^V \in \mathbb{R}^{d_v \times d_{model}}$ is optimized for "information content" not "matching."

WHY: In "The cat sat on the mat," when "sat" attends to "cat," the key helps decide that "cat" is relevant (subject-verb relationship), but the value determines what information "cat" provides (semantic features: animal, subject, etc.).

Full Matrix Form

For a sequence $X = [x_1, x_2, \ldots, x_n]^T \in \mathbb{R}^{n \times d_{model}}$:

Where:

• Each row of $Q$ is a query vector for one token
• Each row of $K$ is a key vector for one token
• Each row of $V$ is a value vector for one token
• $W^Q, W^K, W^V$ are learned parameters updated via backpropagation

HOW they're used in attention: $\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$

Worked Examples

Suppose $d_k = 2$, and we have learned: $W^Q = \begin{bmatrix} 0.1 & 0.4 \\ 0.2 & -0.1 \\ -0.3 & 0.5 \\0.6 & 0.2 \end{bmatrix}, \quad W^K = \begin{bmatrix} -0.2 & 0.3 \\ 0.5 & 0.1 \\ 0.4 & -0.2 \\ -0.1 & 0.6 \end{bmatrix}$

Compute query: $q_1 = x_1 W^Q = [0.5, -0.3, 0.8, 0.2] \begin{bmatrix} 0.1 & 0.4 \\ 0.2 & -0.1 \\ -0.3 & 0.5 \\ 0.6 & 0.2 \end{bmatrix}$

WHY this step? Projecting from 4D semantic space into 2D "search" space.

Step-by-step multiplication:

• First dimension: $0.5(0.1) + (-0.3)(0.2) + 0.8(-0.3) + 0.2(0.6) = 0.05 - 0.06 - 0.24 + 0.12 = -0.13$
• Second dimension: $0.5(0.4) + (-0.3)(-0.1) + 0.8(0.5) + 0.2(0.2) = 0.20 + 0.03 + 0.40 + 0.04 = 0.67$

$q_1 = [-0.13, 0.67]$

Compute key: $k_1 = x_1 W^K = [0.5, -0.3, 0.8, 0.2] \begin{bmatrix} -0.2 & 0.3 \\ 0.5 & 0.1 \\ 0.4 & -0.2 \ -0.1 & 0.6 \end{bmatrix}$

• First dimension: $0.5(-0.2) + (-0.3)(0.5) + 0.8(0.4) + 0.2(-0.1) = -0.10 - 0.15 + 0.32 - 0.02 = 0.05$
• Second dimension: $0.5(0.3) + (-0.3)(0.1) + 0.8(-0.2) + 0.2(0.6) = 0.15 - 0.03 - 0.16 + 0.12 = 0.08$

$k_1 = [0.05, 0.08]$

WHY different from query? Same token has different roles: "what I search for" vs. "what I offer when searched."

$\text{score}_{1,1} = q_1 \cdot k_1 = [-0.13, 0.67] \cdot [0.05, 0.08]$

WHY dot product? Measures alignment in the attention space. High dot product = high relevance.

$= (-0.13)(0.05) + (0.67)(0.08) = -0.0065 + 0.0536 = 0.047$

Interpretation: Token 1 has slight positive relevance to itself. In practice, this would be scaled by $\frac{1}{\sqrt{d_k}} = \frac{1}{\sqrt{2}} \approx 0.707$ and combined with scores to other tokens before softmax.

WHY scaling? As $d_k$ grows, dot products grow in magnitude, pushing softmax into saturation regions. Scaling maintains gradient flow.

$v_1 = x_1 W^V = [0.5, -0.3, 0.8, 0.2] \begin{bmatrix} 0.3 & -0.2 & 0.5 \\ -0.1 & 0.4 & 0.2 \\ 0.6 & 0.1 & -0.3 \\ 0.2 & -0.5 & 0.4 \end{bmatrix}$

• First: $0.5(0.3) + (-0.3(-0.1) + 0.8(0.6) + 0.2(0.2) = 0.15 + 0.03 + 0.48 + 0.04 = 0.70$
• Second: $0.5(-0.2) + (-0.3)(0.4) + 0.8(0.1) + 0.2(-0.5) = -0.10 - 0.12 + 0.08 - 0.10 = -0.24$
• Third: $0.5(0.5) + (-0.3)(0.2) + 0.8(-0.3) + 0.2(0.4) = 0.25 - 0.06 - 0.24 + 0.08 = 0.03$

$v_1 = [0.70, -0.24, 0.03]$

WHY different dimension $d_v$? Values encode information content, which may need different dimensionality than the matching space. Often $d_v = d_k$ for simplicity, but not required.

HOW used: After computing attention weights $\alpha_{i,j}$ via softmax of relevance scores, the output for token $i$ is: $\text{output}_i = \sum_{j=1}^n \alpha_{i,j} v_j$

Common Mistakes

Why it feels right: We define their dimensions and roles architecturally.

Why it's wrong: $W^Q, W^K, W^V$ are learned parameters. They start random and are optimized via gradient descent to learn what features matter for attention in this specific task.

The fix: Think of them as "learnable lenses" that the model tunes to focus on task-relevant patterns. In a translation model, they learn different patterns than in a summarization model.

Why it feels right: Sometimes you see "QKV" as a single matrix in implementation code.

Why it's wrong: Q, K, V are three different linear projections of the same input $X$. They're not slicing $X$ into parts; they're transforming all of $X$ three different ways.

The fix:

• Not: $[Q | K | V] = X$ (wrong: this is concatenation)
• Instead: $Q = XW^Q$, $K = XW^K$, $V = XW^V$ (correct: three separate transformations)

Sometimes implementations optimize by computing $W^{QKV} = [W^Q | W^K | W^V]$ and doing one matrix multiply, then splitting the result, but conceptually they're separate transformations.

Why it feels right: More dimensions = more information preserved.

Why it's wrong:

1. Computational cost: Attention is $O(n^2 d_k)$. Reducing $d_k$ saves compute.
2. Regularization: Lower-dimensional projections force the model to learn compressed, essential features for attention.
3. Multi-head attention: With $h$ heads each of dimension $d_k$, total dimension is $h \cdot d_k = d_{model}$, maintaining expressiveness across heads.

The fix: Think of $d_k$ as the dimension of the "attention query language"—a compressed, task-optimized space. Typical: $d_k = d_{model} / h$ where $h$ is number of heads.

How Q, K, V Enable Multi-Head Attention

In multi-head attention, we create $h$ sets of $(W^Q_i, W^K_i, W^V_i)$ for $i = 1, \ldots, h$. Each head:

1. Computes $Q_i = XW^Q_i$, $K_i = XW^K_i$, $V_i = XW^V_i$
2. Runs attention: $\text{head}_i = \text{Attention}(Q_i, K_i, V_i)$
3. Heads are concatenated and projected: $\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h)W^O$

WHY multiple heads? Different heads can learn different attention patterns:

• Head 1: syntactic relationships (subject-verb)
• Head 2: semantic similarity (synonyms)
• Head 3: positional patterns (nearby words)

Each head's Q, K, V matrices specialize in different aspects.

Query (Q): This is like your hand raised with your specific question written on a card—"I need help with fractions!"

Key (K): This is like everyone else in class holding up cards saying what they're good at—"I know fractions!" or "I'm great at spelling!" or "I know science!"

Value (V): This is the actual helpful information people can give you. Sarah's KEY says "I know fractions," but her VALUE is the actual explanation she can provide about how to add1/2 and 1/3.

The teacher (attention mechanism) looks at your QUERY card, compares it to everyone's KEY cards to find the best matches, then collects the VALUE information from those people and gives you a combined answer.

The magical part? The transformer learns how to write these cards! At first, it writes random stuff on the Q, K, V cards, but after seeing thousands of examples, it learns:

• How to write QUERY cards that ask the right questions
• How to write KEY cards that advertise relevant help
• How to write VALUE cards with the most useful information

And here's the coolest part: The same student (word token) has all three types of cards. When YOU ask a question, you use your QUERY card. But when SOMEONE ELSE asks, they look at your KEY card to see if you can help, and take information from your VALUE card.

Or: "Query the library, match the Keys on the spines, read the Values inside."

Connections

• Attention Mechanism: Q, K, V matrices are inputs to the attention function
• Scaled Dot-Product Attention: Uses $QK^T$ to compute attention scores
• Multi-Head Attention: Creates multiple sets of Q, K, V projections
• Linear Layers: Q, K, V transformations are linear layers (matrix multiplications)
• Embedding Layer: Input $X$ comes from embedding layer
• Positional Encoding: Added to $X$ before Q, K, V projection
• Self-Attention: Q, K, V all derived from same sequence
• Cross-Attention: Q from one sequence, K and V from another (encoder-decoder)
• Backpropagation: How $W^Q, W^K, W^V$ are learned
• Gradient Descent: Optimization algorithm updating weight matrices

#flashcards/ai-ml

What are the three matrices in attention and what does each represent conceptually? :: Query (Q) = "what I'm searching for", Key (K) = "what I offer for matching", Value (V) = "information I provide". They're learned linear projections of input embedings.

Given input $X \in \mathbb{R}^{n \times d_{model}}$, write the formulas for Q, K, V matrices :: $Q = XW^Q$, $K = XW^K$, $V = XW^V$ where $W^Q, W^K \in \mathbb{R}^{d_{model} \times d_k}$ and $W^V \in \mathbb{R}^{d_{model} \times d_v}$ are learned weight matrices.

Concept Map

Hinglish (regional understanding)