Self-Attention Explained: The Mechanism Behind Every Transformer

Every major language model — GPT, Claude, BERT, LLaMA — is built on the same core mechanism: self-attention. It's the piece that replaced recurrence (RNNs, LSTMs) with something that can look at an entire sequence at once and decide, for every word, which other words actually matter to it.

This is the mechanism in detail: the intuition, the math, and a working implementation.

1. The problem self-attention solves

Consider this sentence:

"The trophy didn't fit in the suitcase because it was too big."

To understand what "it" refers to, a model needs to relate that word back to "trophy" — potentially many tokens earlier in the sequence, with no fixed positional rule that would always work.

Older architectures (RNNs) processed tokens one at a time, carrying a hidden state forward. That made long-range relationships hard to learn — information from early tokens had to survive many sequential updates before reaching a token far away.

Self-attention takes a different approach entirely: every token looks directly at every other token in a single step, and learns how much weight to assign to each one.

2. Queries, keys, and values

For every token, self-attention computes three vectors by multiplying the token's embedding against three learned weight matrices:

$Q = XW^Q, \quad K = XW^K, \quad V = XW^V$

Where $X$ is the matrix of input embeddings (one row per token), and $W^Q, W^K, W^V$ are learned parameter matrices.

The intuition behind each vector:

• Query ($Q$) — "what am I looking for?" — represents the current token's request for information.
• Key ($K$) — "what do I contain?" — represents what each token has to offer, in a form comparable to a query.
• Value ($V$) — "what do I actually communicate?" — the content that gets passed along once a token is deemed relevant.

This query/key/value framing borrows directly from information retrieval — think of it as every token issuing a search query, comparing it against every other token's key, then retrieving a weighted blend of values based on how well each key matched.

3. Scaled dot-product attention

To find out how relevant every token is to every other token, we compute the dot product between each query and every key — a higher dot product means the query and key point in a more similar direction, i.e. higher relevance.

$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^\top}{\sqrt{d_k}}\right)V$

Breaking this down step by step:

1. $QK^\top$ — every query dotted against every key, producing a matrix of raw relevance scores of shape (sequence length × sequence length).
2. $\sqrt{d_k}$ — a scaling factor, where $d_k$ is the dimensionality of the key vectors. Without this, dot products grow large as dimensionality increases, which pushes softmax into regions with extremely small gradients. Dividing by $\sqrt{d_k}$ keeps the scores in a numerically stable range.
3. softmax — converts each row of scores into a probability distribution that sums to 1, so each token ends up with a weighted average over all other tokens' values.
4. $\cdot V$ — the softmax weights are used to compute a weighted sum of the value vectors, producing the final output for each token.

A minimal implementation

import numpy as np
 
def softmax(x, axis=-1):
    x = x - np.max(x, axis=axis, keepdims=True)  # numerical stability
    exp_x = np.exp(x)
    return exp_x / np.sum(exp_x, axis=axis, keepdims=True)
 
def scaled_dot_product_attention(Q, K, V):
    d_k = Q.shape[-1]
    scores = Q @ K.transpose(-2, -1) / np.sqrt(d_k)
    weights = softmax(scores, axis=-1)
    output = weights @ V
    return output, weights
 
# Toy example: 4 tokens, embedding dimension 8
seq_len, d_k = 4, 8
Q = np.random.randn(seq_len, d_k)
K = np.random.randn(seq_len, d_k)
V = np.random.randn(seq_len, d_k)
 
output, attention_weights = scaled_dot_product_attention(Q, K, V)
print("Output shape:", output.shape)              # (4, 8)
print("Attention weights shape:", attention_weights.shape)  # (4, 4)
print("Each row sums to 1:", attention_weights.sum(axis=-1))

4. Multi-head attention

A single attention operation forces the model to compress all types of relationships — syntactic, semantic, positional — into one weighting scheme. Multi-head attention runs several attention operations in parallel, each with its own learned $W^Q, W^K, W^V$ matrices, so different heads can specialize in different kinds of relationships.

$\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \dots, \text{head}_h)W^O$

$\text{where head}_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)$

In practice, research on trained Transformers has found heads that specialize — some track subject-verb agreement, others attend to the previous token, others track coreference (like "it" → "trophy" in the example above).

import torch
import torch.nn as nn
 
class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        assert d_model % num_heads == 0
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
 
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)
 
    def forward(self, x, mask=None):
        batch_size, seq_len, _ = x.shape
 
        # Project and reshape into (batch, heads, seq_len, d_k)
        Q = self.W_q(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        K = self.W_k(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        V = self.W_v(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
 
        scores = Q @ K.transpose(-2, -1) / (self.d_k ** 0.5)
 
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float("-inf"))
 
        weights = torch.softmax(scores, dim=-1)
        attended = weights @ V
 
        # Concatenate heads back into (batch, seq_len, d_model)
        attended = attended.transpose(1, 2).contiguous().view(batch_size, seq_len, -1)
        return self.W_o(attended)

5. Causal masking: attention for generation

Encoder self-attention (like in BERT) lets every token attend to every other token, forward and backward. But a model generating text one token at a time — GPT, Claude, and every autoregressive language model — can't be allowed to see future tokens during training, or it would trivially "cheat" by looking ahead at the answer.

Causal masking solves this by setting attention scores for future positions to $-\infty$ before the softmax, so their weight collapses to zero:

def causal_mask(seq_len):
    # Lower-triangular matrix: 1 where attention is allowed, 0 where blocked
    return torch.tril(torch.ones(seq_len, seq_len)).unsqueeze(0).unsqueeze(0)
 
mask = causal_mask(seq_len=5)
print(mask[0, 0])
# tensor([[1., 0., 0., 0., 0.],
#         [1., 1., 0., 0., 0.],
#         [1., 1., 1., 0., 0.],
#         [1., 1., 1., 1., 0.],
#         [1., 1., 1., 1., 1.]])

Each row shows which positions a given token is allowed to attend to — token 3 (row index 2) can see tokens 1–3, but nothing after.

6. Why self-attention won over recurrence

That last row is the catch: self-attention's compute and memory scale quadratically with sequence length, since every token attends to every other token. This is exactly why long-context handling is an active research area — techniques like sparse attention, sliding-window attention, and linear attention variants all exist to soften that quadratic cost for very long sequences.

Summary

• Self-attention lets every token directly weigh the relevance of every other token, in one parallel step.
• Query, key, and value vectors reframe this as a retrieval problem: compare queries against keys, retrieve a weighted blend of values.
• Scaling by $\sqrt{d_k}$ keeps softmax gradients well-behaved as dimensionality grows.
• Multi-head attention lets different heads specialize in different kinds of relationships.
• Causal masking is what makes autoregressive generation — the way GPT and Claude actually produce text — possible to train correctly.
• The tradeoff for all of this is quadratic compute cost in sequence length, which is the central constraint shaping long-context model design today.

Reference links

• Attention Is All You Need (original Transformer paper)
• The Illustrated Transformer
• PyTorch nn.MultiheadAttention documentation