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.

You have a sequence of token embeddings, each a vector in $\mathbb{R}^{d_\text{model}}$. At this stage, an embedding depends on its own token ID, not its neighbors. The token "bank" has the same embedding whether 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. Scaled dot-product attention, introduced in the 2017 transformer, remains a core operation in models such as Llama. Later architectures change its projections, masks, and execution.

Q and K determine the weights; V supplies the information to mix. First, we need to see where those inputs come from and where the result goes: the residual stream.

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. In the pre-norm layout we'll use here, attention and feed-forward blocks update this vector through residual additions. They read from it, compute something, and add their result back.

Each sub-layer applies $x \leftarrow x + f(x)$. The $+f(x)$ is the "residual connection" inherited from ResNet (He et al., 2015), which helped train deep networks by adding an identity term to each block’s derivative. 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. It helps explain how blocks communicate:

• The token embedding remains a term in the residual sum. Later updates can reinforce or cancel its features; an additive path does not guarantee that all input information stays recoverable.
• 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.

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.

Every unmasked key with a finite score participates in the idealized softmax 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.

In Bahdanau attention (Post 2.7), queries come from the decoder and keys/values come from the encoder. In self-attention, all three come from the same sequence.

Parameters vs. activations

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). Changing the input can change Q, K, and V. Identical prefixes in a causal model retain the same prefix activations; the weight matrices remain fixed.

$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$.

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.

The division by $\sqrt{d_k}$ controls how score magnitude grows with dimension; the next section derives it.

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:

Why divide by $\sqrt{d_k}$?

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}$.

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 keeps its spread roughly constant as $d$ grows.

Under these assumptions, when $d_k = 64$, the dot products will typically land somewhere in the range $[-16, 16]$ (within two standard deviations). That's a problem for softmax. When score gaps become large, softmax saturates: one element gets pushed toward 1 and everything else toward 0. The softmax Jacobian becomes small, which can make it harder to change how the head mixes values.

Under the same independence assumptions, dividing by $\sqrt{d_k}$ restores unit variance, reducing this source of softmax saturation. This is what Vaswani et al. call "scaled dot-product attention."

Suppose $d_k = 64$ and one row of your raw score matrix is:

The softmax of that is approximately $[0.999970, \; 0.000020, \; 0.000000, \; 0.000010]$. Almost a one-hot distribution, with little mixing of the other values.

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

Softmax of that is roughly $[0.616, \; 0.160, \; 0.078, \; 0.146]$. Much softer. The softmax is less saturated, so small score changes have a larger effect on the mixture.

Worked 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.

In Step 4, look at how each row sums to 1.00. 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 the sum contains 0.35 times "sat"'s value vector. The other coefficients sum to 0.65. Those coefficients do not describe fractions of the output’s norm: vector contributions can reinforce or cancel.

Implementation

Here's single-head scaled dot-product attention in PyTorch.

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

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 to keep the example focused on the matrices. Bias conventions vary by implementation and model.

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 scaled dot-product operation stays the same.

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."

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 hand-constructed 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). A trained model can learn such relationships through its Q and K projections; these particular weights are illustrative.

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." Softmax exponentiates and then normalizes. Its weight ratio is $A_i/A_j=e^{s_i-s_j}$, so large score gaps produce concentrated distributions. Dividing raw scores by their sum is a different operation and may not even produce probabilities when scores are negative. The $1/\sqrt{d_k}$ factor controls the scale of those gaps under the assumptions above.

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 avoids storing the full score matrix in GPU high-bandwidth memory.
• 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.