Self-Attention Step by Step

Self-attention lets each token gather information from the other tokens. In "the order was late, so refund it", the vector for "it" should pull in information from "order".

1. Each token vector x is turned into a query q = x·Wq, a key k = x·Wk and a value v = x·Wv.
2. Score every pair: score = q · k / √d.
3. Softmax the scores of each row so they add up to 1 (the attention weights).
4. The output for a token is the weighted sum of all value vectors.

The names come from search. Think of a WhatsApp group where each person (token) has three things. The query is the question they are asking ("which noun am I referring to?"). The key is a label describing what they can offer ("I am a noun, an order"). The value is the actual information they will share if picked. Each token compares its query with every key; the better the match, the more of that token's value it receives.

The match is a dot product: multiply two vectors number by number and add. It is large when the vectors point in a similar direction and small or negative when they do not. With q = [0.9, 0.1] for "it" and k = [1, 0] for "order", the dot product is 0.9; with k = [0, 1] for "late" it is only 0.1. Dividing by √d keeps scores from growing too large in big vectors, which would make softmax too sharp.

Softmax then turns raw scores into weights between 0 and 1 that add up to 1: exp(score) ÷ sum of exp(all scores). The exponential keeps the order of the scores but exaggerates the gaps, so the best match gets most of the weight while others still contribute a little. The result for each token is a blend of the values it found relevant: a contextual vector. Stack dozens of layers of this, each followed by a small feed-forward network, and the vectors come to encode grammar, references and meaning.

import math

def matmul(A, B):
    return [[sum(a * b for a, b in zip(row, col)) for col in zip(*B)] for row in A]

def softmax(row):
    m = max(row)
    e = [math.exp(x - m) for x in row]
    return [x / sum(e) for x in e]

tokens = ["order", "late", "it"]
X  = [[1.0, 0.0], [0.0, 1.0], [0.9, 0.1]]       # "it" starts out similar to "order"
Wq = [[1.0, 0.0], [0.0, 1.0]]
Wk = [[1.0, 0.0], [0.0, 1.0]]
Wv = [[0.5, 1.0], [1.0, 0.0]]

Q, K, V = matmul(X, Wq), matmul(X, Wk), matmul(X, Wv)
d = len(Q[0])
scores = [[sum(q * k for q, k in zip(Q[i], K[j])) / math.sqrt(d) for j in range(3)] for i in range(3)]
weights = [softmax(r) for r in scores]
out = matmul(weights, V)

print("attention weights (row = token looking, column = token looked at)")
print("         " + "  ".join(f"{t:>6}" for t in tokens))
for t, row in zip(tokens, weights):
    print(f"{t:>6}   " + "  ".join(f"{w:6.3f}" for w in row))
for t, o in zip(tokens, out):
    print(f"output for {t!r}: {[round(x, 3) for x in o]}")

attention weights (row = token looking, column = token looked at)
          order    late      it
 order    0.412   0.203   0.384
  late    0.244   0.494   0.262
    it    0.398   0.226   0.376
output for 'order': [0.621, 0.758]
output for 'late': [0.76, 0.479]
output for 'it': [0.632, 0.736]

Read the last row: "it" gives the largest weight to "order" (0.398), a similar weight to itself, and the least to "late" (0.226), because its vector [0.9, 0.1] is closest to "order". Its output [0.632, 0.736] is now close to the output for "order": information has flowed from "order" into "it". Here Wq and Wk are identity matrices to keep the numbers simple; in a real model they are learned, which is how the model learns what to look for.

Every token scores every other token, so a prompt of n tokens needs about n × n scores per layer and head. Double the prompt and attention work roughly quadruples. This is one reason long prompts are slower and why trimming context matters.

• Summary: each token makes a query, key and value; dot products of queries and keys measure relevance.
• Softmax turns scores into weights that sum to 1; the output is a weighted blend of values.
• This lets every token pull in context from every other token in one step.
• Cost grows with n², so long prompts are slower.