AI ENGINEERING: DEEP DIVE · MODULE 02 — TRANSFORMERS, THE REAL MECHANISM · 04 A hands-on series · math + code

Module 2.1

The real attention equation

Module 1.3 attended directly on raw embeddings. Real attention never does that — every token is projected into a Query, a Key, and a Value first. Pick a token below and watch the full equation run.

QUERY: "Why did my Nimbus export fail?" — TOKENS BELOW
Q(export) · K(j) / √d, for every j
OUTPUT = Σ weight × V(j)

The math

Three steps, same formula every layer

Projections
Q, K, V  =  XWq, XWk, XWv
X is the token's embedding; Wq/Wk/Wv are learned
Q(export) = [0.35, 0.22]
Scaled score
score  =  (Q · K) / √d
Q(export) · K(fail) / √2
= 0.36
Weighted output
out  =  Σⱼ softmax(scores)ⱼ Vⱼ
one weight per token, always summing to 1
Σ weights = 1.00

The scale-by-√d isn't decorative — without it, larger embedding dimensions push scores so large that softmax saturates into a near one-hot pick, and gradients vanish everywhere else.

The code

~20 lines, the whole mechanism

self_attention.py
import numpy as np

def softmax(x):
    e = np.exp(x - np.max(x))
    return e / e.sum()

# X: one row per token embedding, shape (n_tokens, d_model)
Q = X @ Wq   # (n_tokens, d_k)
K = X @ Wk   # (n_tokens, d_k)
V = X @ Wv   # (n_tokens, d_v)

d_k = K.shape[-1]
scores  = (Q @ K.T) / np.sqrt(d_k)     # (n_tokens, n_tokens)
weights = np.array([softmax(row) for row in scores])
output  = weights @ V                   # (n_tokens, d_v) — one output per token

That's the entire mechanism. Every transformer's attention sublayer, at every layer, is these same three matrix multiplies plus one softmax.

The Nimbus thread

Same tokens as Module 1.3 — now with real projections

Module 1.3 let "export" attend directly using its raw embedding as both query and key. Real self-attention never does that — it learns three separate projections first, so "how similar am I" (query) and "what do I offer" (key) can differ from "what do I actually contribute" (value). Toggle through each token above; the shape of the story survives (export and fail still pull toward each other) even though every number changed.

So what

Every token asks a question (Q), every token offers an answer (K), and the output is a weighted blend of what's actually useful (V) — that's attention, full stop.