The mechanism that lets each token look at every other token and decide how much each contributes.

Data Flow

x (token embeddings, shape [L, d])

         W_Q     W_K     W_V
          \      |      /
x -------> Q    K      V             (each [L, d_k])

scores = Q . K^T  /  sqrt(d_k)        (shape [L, L])

weights = softmax(scores)             (rows sum to 1)

output = weights @ V                  (shape [L, d_v])

What is it?

For each token, take its Query vector and dot-product with every other token’s Key to get a similarity score. Normalise, then take a weighted sum of all Values.

Why the scale sqrt(d_k)?

Inner products grow with d_k, pushing softmax into saturated regions where gradients vanish. Dividing by sqrt(d_k) normalises the variance back to ~1.

Code

import torch, torch.nn.functional as F
 
def self_attention(Q, K, V):
    d_k = Q.size(-1)
    scores = (Q @ K.transpose(-2, -1)) / (d_k ** 0.5)
    weights = F.softmax(scores, dim=-1)
    return weights @ V

Pitfalls

• Confusing self-attention (Q, K, V from the same source) with cross-attention (Q from one source, K and V from another).
• Forgetting the mask in causal (decoder) attention — the model “cheats” by seeing future tokens.

Analogy

In a meeting room, every speaker writes down how relevant everyone else’s comment is to their own, then drafts their reply as a weighted mix.

Interview tip: Memorise softmax(QK^T / sqrt(d_k)) V. Half of all “explain transformers” questions end here.