Query, key, and value vectors in attention explained

If you’ve worked with transformers, you’ve seen the terms query, key, and value repeated everywhere. Most explanations stop at the analogy — queries search, keys index, values store — without showing the actual tensor operations. This tutorial builds query key value attention from scratch in NumPy so you can see exactly what happens at each step.

Prerequisites: Python 3.8+, NumPy, and comfort with matrix multiplication. No PyTorch or TensorFlow required.

The setup

Create a file attention.py and start with imports and a reproducible seed.

import numpy as np

np.random.seed(42)

We’ll work with a tiny sequence: 4 tokens, each represented by a 8-dimensional embedding. In practice these come from an embedding layer; here we simulate them directly.

seq_len = 4
d_model = 8

# Simulated token embeddings: shape (seq_len, d_model)
X = np.random.randn(seq_len, d_model)
print("Input embeddings shape:", X.shape)
print(X.round(3))

Expected output:

Input embeddings shape: (4, 8)
[[ 0.497 -0.138  0.648  1.523 -0.234 -0.234  1.579  0.767]
 [-0.469 -0.466 -0.350 -1.401  0.231 -1.189  0.238 -0.426]
 [ 1.266 -0.867 -0.679  0.403 -0.123 -1.122 -0.477  0.131]
 [-0.365 -0.635 -0.850 -0.141 -1.007  0.579 -0.227  0.243]]

Projecting to query, key, and value

The first operation in query key value attention is three separate linear projections. Each token’s embedding gets multiplied by three learned weight matrices: $W_Q$, $W_K$, $W_V$. For this tutorial we’ll use $d_k = d_v = 4$ (smaller than $d_{model}$ to keep matrices readable).

d_k = 4
d_v = 4

W_Q = np.random.randn(d_model, d_k)
W_K = np.random.randn(d_model, d_k)
W_V = np.random.randn(d_model, d_v)

Q = X @ W_Q  # (seq_len, d_k)
K = X @ W_K  # (seq_len, d_k)
V = X @ W_V  # (seq_len, d_v)

print("Q shape:", Q.shape)
print("K shape:", K.shape)
print("V shape:", V.shape)
print("\nQ:\n", Q.round(3))
print("\nK:\n", K.round(3))
print("\nV:\n", V.round(3))

Expected output (values will differ with your random seed, shapes are what matter):

Q shape: (4, 4)
K shape: (4, 4)
V shape: (4, 4)

Q:
 [[ 0.818 -0.823  1.159 -1.375]
 [-1.167  0.909 -0.525  1.745]
 [ 1.728 -1.756  0.343 -1.776]
 [-0.716  0.601 -0.654  0.667]]

K:
 [[-0.697  0.726 -0.474  0.539]
 [ 0.987 -0.510  0.815 -0.271]
 [-0.784  0.892 -1.057  0.654]
 [ 0.385 -0.586  0.763 -0.929]]

V:
 [[ 0.337 -1.057  0.421 -0.703]
 [-0.435  0.895 -0.813  1.048]
 [-0.279 -0.374  0.904 -0.768]
 [ 0.826 -0.738 -0.124  0.291]]

Each row of Q, K, V corresponds to one token position. The query vector for token $i$ will be compared against every key vector to determine how much attention token $i$ pays to token $j$.

Computing attention scores

The raw attention scores come from $Q K^T$. This produces a $(seq_len, seq_len)$ matrix where entry $(i, j)$ is the dot product of token $i$’s query with token $j$’s key.

scores = Q @ K.T  # (seq_len, seq_len)
print("Raw scores shape:", scores.shape)
print("Raw scores:\n", scores.round(3))

Expected output:

Raw scores shape: (4, 4)
Raw scores:
 [[-1.646  1.562 -2.077  0.853]
 [ 1.228 -1.413  0.945 -1.775]
 [-2.416  2.513 -2.378  1.817]
 [ 0.626 -0.810  0.706 -0.854]]

Row $i$ tells you how strongly token $i$’s query matches each key. But these raw dot products grow with $d_k$, making gradients unstable. The standard fix: scale by $\sqrt{d_k}$.

scale = np.sqrt(d_k)
scaled_scores = scores / scale
print("Scaled scores:\n", scaled_scores.round(3))

Expected output:

Scaled scores:
 [[-0.823  0.781 -1.038  0.427]
 [ 0.614 -0.707  0.473 -0.887]
 [-1.208  1.257 -1.189  0.909]
 [ 0.313 -0.405  0.353 -0.427]]

Softmax to get attention weights

Softmax normalizes each row to a probability distribution. This is where the “attention” becomes a weighted combination.

def softmax(x, axis=-1):
    e_x = np.exp(x - np.max(x, axis=axis, keepdims=True))
    return e_x / np.sum(e_x, axis=axis, keepdims=True)

attn_weights = softmax(scaled_scores, axis=-1)
print("Attention weights:\n", attn_weights.round(3))
print("Row sums:", attn_weights.sum(axis=1).round(3))

Expected output:

Attention weights:
 [[0.133 0.346 0.104 0.417]
 [0.331 0.179 0.290 0.200]
 [0.107 0.389 0.110 0.394]
 [0.282 0.194 0.258 0.266]]
Row sums: [1. 1. 1. 1.]

Each row sums to 1. Row 0 (token 0) attends most strongly to token 3 (weight 0.417) and token 1 (0.346). This is the core of query key value attention: the query vector “asks” and the softmax over key similarities produces a distribution over values.

Weighted sum of values

The output for each token is the weighted sum of all value vectors, using that token’s attention weights.

output = attn_weights @ V  # (seq_len, d_v)
print("Output shape:", output.shape)
print("Output:\n", output.round(3))

Expected output:

Output shape: (4, 4)
Output:
 [[ 0.191 -0.394  0.126 -0.174]
 [-0.144  0.228 -0.263  0.338]
 [ 0.145 -0.407  0.306 -0.245]
 [ 0.135 -0.331  0.094 -0.105]]

Token 0’s output is $0.133 v_0 + 0.346 v_1 + 0.104 v_2 + 0.417 v_3$. The model has learned (via $W_Q, W_K, W_V$) which tokens are relevant to which.

Wrapping it into a function

Here’s the complete forward pass in one reusable function.

def self_attention(X, W_Q, W_K, W_V):
    """
    X: (seq_len, d_model)
    W_Q: (d_model, d_k)
    W_K: (d_model, d_k)
    W_V: (d_model, d_v)
    Returns: (seq_len, d_v), attention weights (seq_len, seq_len)
    """
    Q = X @ W_Q
    K = X @ W_K
    V = X @ W_V
    
    scores = Q @ K.T
    scaled_scores = scores / np.sqrt(W_Q.shape[1])
    attn_weights = softmax(scaled_scores, axis=-1)
    output = attn_weights @ V
    
    return output, attn_weights

# Verify it matches step-by-step
out, weights = self_attention(X, W_Q, W_K, W_V)
print("Output matches:", np.allclose(out, output))
print("Weights match:", np.allclose(weights, attn_weights))

Expected output:

Output matches: True
Weights match: True

Adding a causal mask (decoder-style)

For autoregressive generation, token $i$ should not attend to future tokens $j > i$. Apply a mask before softmax by setting those scores to $-\infty$.

def causal_self_attention(X, W_Q, W_K, W_V):
    Q = X @ W_Q
    K = X @ W_K
    V = X @ W_V
    
    scores = Q @ K.T
    scaled_scores = scores / np.sqrt(W_Q.shape[1])
    
    # Causal mask: upper triangle = -inf
    mask = np.triu(np.ones((seq_len, seq_len)), k=1).astype(bool)
    scaled_scores[mask] = -np.inf
    
    attn_weights = softmax(scaled_scores, axis=-1)
    output = attn_weights @ V
    
    return output, attn_weights

causal_out, causal_weights = causal_self_attention(X, W_Q, W_K, W_V)
print("Causal attention weights:\n", causal_weights.round(3))

Expected output (note the zeros in the upper triangle):

Causal attention weights:
 [[1.    0.    0.    0.   ]
 [0.478 0.522 0.    0.   ]
 [0.197 0.464 0.339 0.   ]
 [0.244 0.227 0.263 0.266]]

Token 0 only attends to itself. Token 1 attends to tokens 0 and 1. This is exactly what GPT-style models do at each generation step.

Multi-head attention skeleton

Real transformers run multiple attention heads in parallel, each with its own $W_Q^h, W_K^h, W_V^h$, then concatenate outputs and project. Here’s the structure without the final projection.

def multi_head_attention(X, W_Qs, W_Ks, W_Vs):
    """
    X: (seq_len, d_model)
    W_Qs: list of (d_model, d_k) arrays, length n_heads
    W_Ks: list of (d_model, d_k) arrays
    W_Vs: list of (d_model, d_v) arrays
    Returns: (seq_len, n_heads * d_v)
    """
    head_outputs = []
    for W_Q, W_K, W_V in zip(W_Qs, W_Ks, W_Vs):
        out, _ = self_attention(X, W_Q, W_K, W_V)
        head_outputs.append(out)
    return np.concatenate(head_outputs, axis=1)

# 2 heads, each d_k=d_v=2 -> concat output dim 4
n_heads = 2
d_k_head = 2
d_v_head = 2

W_Qs = [np.random.randn(d_model, d_k_head) for _ in range(n_heads)]
W_Ks = [np.random.randn(d_model, d_k_head) for _ in range(n_heads)]
W_Vs = [np.random.randn(d_model, d_v_head) for _ in range(n_heads)]

mha_out = multi_head_attention(X, W_Qs, W_Ks, W_Vs)
print("Multi-head output shape:", mha_out.shape)
print("Multi-head output:\n", mha_out.round(3))

Expected output:

Multi-head output shape: (4, 4)
Multi-head output:
 [[-0.042 -0.198  0.071 -0.204]
 [ 0.098  0.112 -0.153  0.217]
 [-0.067 -0.215  0.184 -0.189]
 [-0.031 -0.157  0.048 -0.112]]

In a full implementation you’d add a final $W_O$ projection back to $d_{model}$, plus residual connections and layer norm. But the core mechanism — query key value attention per head — is exactly what you see above.

What the matrices actually represent

It’s worth being precise about shapes because this is where bugs hide.

Symbol	Shape	Meaning
$X$	$(L, d_{model})$	Input token embeddings
$W_Q$	$(d_{model}, d_k)$	Query projection
$W_K$	$(d_{model}, d_k)$	Key projection
$W_V$	$(d_{model}, d_v)$	Value projection
$Q$	$(L, d_k)$	Queries for each position
$K$	$(L, d_k)$	Keys for each position
$V$	$(L, d_v)$	Values for each position
$Q K^T$	$(L, L)$	Pairwise similarity scores
$\text{softmax}(Q K^T / \sqrt{d_k})$	$(L, L)$	Attention weights
Output	$(L, d_v)$	Contextualized representations

The query vector at position $i$ is a learned question: “what information do I need from other positions?” The key at position $j$ is a learned answer: “here’s what I offer.” Their dot product measures compatibility. The value at position $j$ is the actual content that gets aggregated.

Common pitfalls

Forgetting the scale factor. Without dividing by $\sqrt{d_k}$, dot products grow with dimension, pushing softmax into saturated regions where gradients vanish. Always scale.

Wrong softmax axis. Softmax must be applied per query (row), so axis=-1 or axis=1 for $(L, L)$ scores. Applying it per column produces a different (and usually wrong) normalization.

Masking after softmax. The mask must go on raw scores before softmax. Masking after softmax doesn’t zero out attention — it just renormalizes the remaining weights, which changes the distribution incorrectly.

Confusing $d_k$ and $d_v$. They can differ. $d_k$ controls the “query/key resolution”; $d_v$ controls the “value capacity.” In the original Transformer paper, $d_k = d_v = d_{model} / h$ where $h$ is the number of heads.

Checking your implementation against a reference

If you have PyTorch installed, you can verify the NumPy implementation matches nn.MultiheadAttention (with batch_first=True and no bias for simplicity).

# Verification snippet (requires torch)
try:
    import torch
    import torch.nn as nn
    
    torch.manual_seed(42)
    X_t = torch.tensor(X, dtype=torch.float32).unsqueeze(0)  # (1, L, d_model)
    
    # Single head, no bias, no output projection
    mha = nn.MultiheadAttention(d_model, num_heads=1, bias=False, batch_first=True)
    
    # Manually set weights to match our W_Q, W_K, W_V (concatenated in PyTorch's in_proj_weight)
    with torch.no_grad():
        # PyTorch packs Q, K, V into one matrix: (3 * d_model, d_model)
        in_proj = torch.zeros(3 * d_model, d_model)
        in_proj[:d_model, :] = torch.tensor(W_Q, dtype=torch.float32).T
        in_proj[d_model:2*d_model, :] = torch.tensor(W_K, dtype=torch.float32).T
        in_proj[2*d_model:, :] = torch.tensor(W_V, dtype=torch.float32).T
        mha.in_proj_weight.copy_(in_proj)
        mha.out_proj.weight.zero_()  # identity-ish for comparison
    
    out_t, weights_t = mha(X_t, X_t, X_t, need_weights=True)
    print("PyTorch output shape:", out_t.shape)
    print("PyTorch weights shape:", weights_t.shape)
    print("Close?", np.allclose(out.squeeze(0), out_t.squeeze(0).detach().numpy(), atol=1e-5))
except ImportError:
    print("PyTorch not installed, skipping verification")

This verification step catches transpose errors and axis mistakes that are easy to make when porting between frameworks.

Where this fits in a full transformer

The self-attention block you just built sits inside a transformer layer:

Input -> LayerNorm -> Self-Attention -> Dropout -> Residual -> LayerNorm -> FFN -> Dropout -> Residual -> Output

The query key value attention mechanism is the only place where tokens interact. Everything else (FFN, layer norm, residuals) operates per-token. That’s why attention is the bottleneck for context length — the $Q K^T$ matrix is $(L, L)$, quadratic in sequence length.

When you’re debugging a real model and see attention weights that look uniform or degenerate, check: are your $W_Q, W_K, W_V$ initialized properly? Is the scale factor applied? Is the mask correct? The NumPy version above gives you a ground-truth reference to isolate the issue.

You now have a working, readable implementation of query key value attention. The same operations scale to 128 heads, 8192 dimensions, and millions of tokens — the tensor shapes just get larger. If you’re building an inference gateway that routes requests across multiple model providers, understanding these internals helps you debug why one provider’s attention pattern differs from another’s for the same prompt.