Self-Attention: Q, K, V from First Principles

Each token builds a query, scores it against every other token's key, and mixes their values in proportion to the match.

Here's the situation. You have a sequence of token embeddings, each a vector in $\mathbb{R}^{d_\text{model}}$. Each token knows something about itself, but nothing about the tokens around it. The word "bank" doesn't yet know whether it's followed by "account" or "river."

Self-attention is the mechanism that lets tokens talk to each other. It's how "bank" figures out which other tokens in the sequence are relevant to it, and then gathers information from them. Every transformer you've heard of (GPT-4, Claude, Gemini, Llama) uses this exact operation, essentially unchanged since Vaswani et al. introduced it in 2017.

So how does a token decide what's relevant, and how does it gather information from relevant tokens once it has? That's what Q, K, and V are for. But before we get there, there's one piece of mental scaffolding worth setting up first, because every block in the rest of this arc reads from and writes to it.

The residual stream

A useful way to picture a transformer is as a residual stream. There's a single vector at each token position, of dimension $d_\text{model}$, that flows through the entire network from the input embedding to the final unembedding. Attention blocks and feed-forward blocks don't transform this vector. They read from it, compute something, and add their result back. The stream accumulates.

The atom of the architecture is therefore $x \leftarrow x + f(x)$, repeated for every sub-layer. The $+f(x)$ is the "residual connection" inherited from ResNet (He et al., 2015), which made deep networks trainable by giving gradients an unobstructed path from the loss back to the input. In a transformer it does more than just stabilize training: it sets up a shared bus that every block reads from and writes to. An attention head in layer 7 can, in principle, pick up information that an attention head in layer 1 wrote, because layer 1's contribution is still in the sum.

This framing comes from Elhage et al.'s mathematical framework for transformer circuits. A few consequences worth carrying:

• The token embedding never gets overwritten. It's still in the stream at the final layer, plus every additive contribution from every block.
• The stream is a communication channel between blocks. Heads that write specific patterns into specific dimensions can be "read" by later FFN or attention blocks that learned to pay attention to those dimensions. That's what people mean by a "circuit."
• Each token position has its own residual stream. Attention is what lets positions exchange information; FFN blocks are position-independent.

Keep that picture in mind for the rest of the arc. The rest of this post zooms in on one specific block — self-attention — and what exactly it reads from and writes to the stream.

The soft lookup analogy

I find it helpful to think of self-attention as a database lookup, but a soft one.

In a regular hash table, you have a query and you get back one exact match. Self-attention does something similar but fuzzier. Every token simultaneously plays three roles:

• Query (Q): "What am I looking for?"
• Key (K): "What do I advertise to others?"
• Value (V): "What information do I carry?"

When token $i$ wants to gather information, it broadcasts its query, compares it against every other token's key, and then takes a weighted average of all the value vectors. Tokens whose keys match the query strongly contribute more. Tokens whose keys don't match still contribute, just very little.

That "still contribute" part is important. This isn't a hard lookup. Every key participates in every query result. The softmax function (which we covered in Post 1.3) turns the raw similarity scores into a probability distribution. So each token's output is a convex combination of all value vectors in the sequence.

If you've seen Bahdanau attention from the seq2seq era (Post 2.7), this will feel familiar. The key difference is that in Bahdanau attention, queries come from the decoder and keys/values come from the encoder. In self-attention, all three come from the same sequence. That's the "self" part.

Parameters vs. activations

I got tripped up on this when I first learned it.

Q, K, and V are not parameters. They are activations.

The parameters are three weight matrices: $W_Q$, $W_K$, and $W_V$. These get learned during training and stay fixed at inference time. They live inside the model's checkpoint file.

The activations Q, K, and V are computed fresh for every single input:

where $X$ is the input matrix (one row per token). Feed in a different sentence, you get completely different Q, K, and V matrices. The weight matrices are the same; the activations change every time.

When someone says "the model's Q, K, V," they almost certainly mean the weight matrices $W_Q, W_K, W_V$. The actual Q, K, V vectors are ephemeral, computed per-input, and thrown away after each forward pass.

One more way to think about it. $W_Q$ is the model's learned concept of "what makes a good query." When multiplied by a specific token embedding, it produces that token's query vector. Same idea for $W_K$ and $W_V$.

Here's the concrete version. Two token embeddings (rows of $X$) multiplied by $W_Q$ to produce two query vectors. Click any output cell to see how it's assembled from a single row-column dot product.

The math, step by step

Let's trace the full computation. We have $n$ tokens, each with a $d_\text{model}$-dimensional embedding. Stack them into a matrix $X$ of shape $(n, d_\text{model})$.

Step 1: Project into Q, K, V.

The projection dimensions $d_k$ and $d_v$ are typically equal and often smaller than $d_\text{model}$. In the original transformer, $d_\text{model} = 512$ and $d_k = d_v = 64$. We'll talk about why they're smaller when we get to multi-head attention in Post 4.2.

Step 2: Compute raw attention scores.

This is a matrix of dot products. Entry $S_{ij}$ is the dot product of token $i$'s query with token $j$'s key. High dot product means strong alignment between what token $i$ is looking for and what token $j$ is advertising.

The result is an $n \times n$ matrix. Every token has a score for every other token. This is where the quadratic cost of attention comes from.

Here's what that multiplication looks like with real (small) numbers. Two query vectors times two key vectors, transposed, producing a 2x2 score matrix.

Step 3: Scale.

We'll dig into why in a moment. For now, just note that we divide every score by $\sqrt{d_k}$.

Step 4: Softmax.

Softmax is applied row-wise. Each row of $A$ is a probability distribution: it sums to 1, every entry is non-negative. Row $i$ of $A$ is the attention pattern for token $i$, telling us how much weight to give each token's value vector.

Step 5: Weighted sum of values.

Each output row is a weighted average of the value vectors, where the weights come from the attention pattern. Token $i$'s output is $\sum_j A_{ij} V_j$.

Or, all at once:

That's the complete formula. Everything in the rest of this post is just understanding why each piece is there.

Why divide by $\sqrt{d_k}$?

This is one of those details that seems minor but has a real mathematical reason behind it.

Consider the dot product of two random vectors $q$ and $k$, each with $d_k$ components drawn independently from a standard normal distribution. The dot product is:

Each term $q_i k_i$ has mean 0 and variance 1 (since each factor has variance 1 and they're independent). The sum of $d_k$ such terms has variance $d_k$, which means the standard deviation is $\sqrt{d_k}$.

You can watch this happen directly. The viz below samples 240 random query/key pairs at each $d$ and plots the dot products. Notice how the cloud of red dots widens as $d$ steps up. The blue panel applies the $1/\sqrt{d}$ rescaling and the cloud stays put no matter how big $d$ gets.

So when $d_k = 64$, your dot products will typically land somewhere in the range $[-16, 16]$ (within two standard deviations). That's a problem for softmax. When inputs are large in magnitude, softmax saturates: one element gets pushed toward 1 and everything else toward 0. The gradients in those saturated regions are vanishingly small, so learning stalls.

Dividing by $\sqrt{d_k}$ normalizes the variance back to approximately 1, keeping the scores in a range where softmax is well-behaved. This is what Vaswani et al. call "scaled dot-product attention."

Let me make this concrete. Suppose $d_k = 64$ and one row of your raw score matrix is:

The softmax of that is approximately $[0.9999, \; 0.0001, \; 0.0000, \; 0.0000]$. Basically a hard one-hot. Not useful for learning.

After scaling by $\sqrt{64} = 8$:

Softmax of that is roughly $[0.62, \; 0.16, \; 0.08, \; 0.15]$. Much softer. Gradients can flow through all four positions. The model can learn nuanced attention patterns instead of being forced into hard, binary lookups.

You can feel this effect directly by sweeping $d_k$ in the viz below. The left panel is what softmax does to raw dot products, the right panel is what it does after we divide by $\sqrt{d_k}$.

Worked example

Let me walk through a complete concrete example. We have four tokens: "The", "cat", "sat", "down". I'll use $d_\text{model} = 8$ and $d_k = d_v = 4$, with small hand-picked numbers so the patterns are visible.

The interactive visualization below lets you step through each phase of the computation. Click on any token to highlight its row through the matrices.

In Step 4, look at how each row sums to 1.00. That's the softmax doing its job. Each token distributes its attention budget across all four tokens, and the weights have to add up.

In Step 5, when you click on a specific token, you can see exactly how its output is assembled. If "cat" attends strongly to "sat" (weight 0.35, say), then 35% of "cat"'s output vector comes from "sat"'s value vector. The remaining 65% is a blend of the other three value vectors.

This is the key insight. Attention doesn't copy information from a single source. It blends information from all sources, proportional to relevance. And "relevance" is determined entirely by the learned Q and K projections.

Implementation

Here's single-head scaled dot-product attention in PyTorch. Shorter than you'd expect.

import torch
import torch.nn as nn
import torch.nn.functional as F
import math
 
class SingleHeadAttention(nn.Module):
    def __init__(self, d_model: int, d_k: int):
        super().__init__()
        self.d_k = d_k
        # These are the PARAMETERS — learned during training
        self.W_Q = nn.Linear(d_model, d_k, bias=False)
        self.W_K = nn.Linear(d_model, d_k, bias=False)
        self.W_V = nn.Linear(d_model, d_k, bias=False)
 
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: (batch, seq_len, d_model)
 
        # These are the ACTIVATIONS — different for every input
        Q = self.W_Q(x)  # (batch, seq_len, d_k)
        K = self.W_K(x)  # (batch, seq_len, d_k)
        V = self.W_V(x)  # (batch, seq_len, d_k)
 
        # Scaled dot-product attention
        scores = Q @ K.transpose(-2, -1)  # (batch, seq_len, seq_len)
        scores = scores / math.sqrt(self.d_k)
        weights = F.softmax(scores, dim=-1)
        output = weights @ V  # (batch, seq_len, d_k)
 
        return output

That's it. Twenty-ish lines for the core mechanism of every modern language model.

A few implementation notes. The nn.Linear modules are the weight matrices. PyTorch's nn.Linear(d_model, d_k, bias=False) stores a matrix of shape $(d_k, d_\text{model})$ internally and computes $x W^\top$, which is equivalent to our $X W$ formulation. I've set bias=False because the original transformer doesn't use bias terms in the attention projections, and most modern architectures follow suit.

The @ operator is matrix multiplication. K.transpose(-2, -1) transposes the last two dimensions, giving us $K^\top$ while preserving the batch dimension. The softmax is applied along dim=-1, meaning each row independently becomes a probability distribution.

In a real transformer, you'd add causal masking (which we'll cover in Post 4.3) and wrap this in a multi-head structure (Post 4.2). The fundamental computation is exactly what you see here.

What each projection learns

So we have three weight matrices that each take the same input and produce different vectors. What do they actually learn to extract?

$W_Q$ learns to produce queries: representations of what each token is looking for. For a verb, the query might encode "I need my subject." For an adjective, it might encode "I need the noun I'm modifying."

$W_K$ learns to produce keys: representations of what each token has to offer. A noun might advertise "I'm an entity" or "I'm the subject of this clause." These keys don't have to be semantically meaningful to humans. They just need to produce high dot products with the right queries.

$W_V$ learns to produce values: the information that actually gets passed along when a token is attended to. This is what gets blended into the output. A token's key says "I'm relevant," and its value says "here's my information."

This separation is crucial. A token's role as a query target (what matches against it?) can be completely different from the information it carries (what does it contribute?). The three projections give the model three independent degrees of freedom per token.

Mechanistic interpretability work, particularly the "mathematical framework for transformer circuits" from Anthropic, has shown that these projections often learn clean structure. Certain attention heads learn to implement specific operations: "attend to the previous token," "attend to tokens with matching syntactic role," "attend to the most recent proper noun." The QKV decomposition is what enables this specialization.

The heatmap below shows a realistic self-attention pattern for the sentence "The cat sat down." Notice how each row (query token) distributes its attention budget differently. "cat" attends back to "The" (its determiner) and forward to "sat" (its verb). "sat" attends heavily to "cat" (its subject). These patterns emerge entirely from the learned Q and K projections.

Misconceptions

"Q, K, V are the model's weights." No. $W_Q$, $W_K$, $W_V$ are the weights. Q, K, V are activations computed from the input.

"The attention matrix shows what the model is thinking about." The attention weights $A$ show how information flows between positions for that specific head. But the model has many heads (typically 12-96), and the actual computation depends on what's in the value vectors, not just the attention pattern. A head that attends strongly to position 5 but has near-zero values at position 5 effectively ignores it. Attention weights are one signal, not the whole picture.

"Softmax just normalizes the scores." Calling softmax a normalizer undersells what it does. A plain normalizer would divide each score by the sum of scores. Softmax exponentiates first, which makes the operation a sharpener: small differences in raw score become large differences in weight, because $e^x$ is convex. That's why the $1/\sqrt{d_k}$ rescaling matters so much. Softmax wants its inputs in a moderate range so it has room to amplify a real signal without saturating into a hard one-hot.

What's next

This post covered a single attention head. The next post covers multi-head attention and representation subspaces, where the transformer runs many of these heads in parallel on different projections of the same stream.

Additional reading (and watching)

• Vaswani, A., et al. (2017). Attention Is All You Need. NeurIPS 2017. Introduces scaled dot-product attention and the transformer architecture.
• Bahdanau, D., Cho, K., & Bengio, Y. (2015). Neural Machine Translation by Jointly Learning to Align and Translate. ICLR 2015. The original soft-alignment attention, precursor to QKV.
• Elhage, N., et al. (2021). A Mathematical Framework for Transformer Circuits. Anthropic. Decomposes the transformer into interpretable QK/OV circuits operating on the residual stream.
• Dao, T., et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS 2022. Tiled exact attention that keeps the $n \times n$ matrix in on-chip memory and unlocks long contexts.
• Ainslie, J., et al. (2023). GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. EMNLP 2023. The grouped-query attention formulation used by Llama 2/3, Mistral, and many modern open-weights models.