B-preview writes when the script runsDictionary of Applied Machine Learning · attention
Numerical companion to the entry attention: it recomputes what the entry states and prints one line per check
Illustrate attention as a learned *associative memory*. Each token turns its own embedding into a query and searches, via inner products against the keys of the other tokens, for the most relevant tokens; the softmax-weighted values are then read out. Training makes queries and keys align so that a token can be reconstructed from the tokens it is associated with.
Run it with python3 attention.py, from any directory — it writes its output files into the current directory. Requires NumPy and Matplotlib only, and uses fixed seeds, so the printed numbers reproduce exactly. Download attention.py
One cell per block of the script: the code, and what that code printed when it last ran here
"""
attention.py — train a single self-attention layer on sentences from the
Universal Declaration of Human Rights.
Purpose
-------
Illustrate attention as a learned *associative memory*. Each token turns its
own embedding into a query and searches, via inner products against the keys of
the other tokens, for the most relevant tokens; the softmax-weighted values are
then read out. Training makes queries and keys align so that a token can be
reconstructed from the tokens it is associated with.
Corpus
------
Sentences and clauses excerpted from the Universal Declaration of Human Rights
(UN General Assembly resolution 217 A, 1948; public domain), stored verbatim
below so the demo is self-contained and reproducible (no network access, fixed
seed). The display sentence is the first sentence of Article 1, which is also
the example sentence of the scaled-dot-product figure in the `attention`
glossary entry.
Model
-----
One bidirectional self-attention head with learned projections W_Q, W_K, W_V
and a linear read-out W_O to one score per distinct token. Embeddings E are learned too.
The self-supervised objective is a cloze / associative-recall task: for every
position i the diagonal of the attention matrix is masked out (a token may not
attend to itself), and the read-out must predict the held-out token at i from a
softmax-weighted combination of the *other* tokens' values. This forces the
attention weights to encode which tokens are mutually predictive — the essence
of content-addressable memory.
Manual forward/backward (numpy only):
Q = X W_Q, K = X W_K, V = X W_V (per sentence)
S = Q K^T / sqrt(d_k), S_ii = -inf (mask self)
A = softmax(S, axis=1) (rows sum to 1)
Z = A V, scores = Z W_O
loss = mean negative log probability assigned to token i at position i
The whole model is implemented by hand in numpy — no autograd — so every line
of the forward pass has an explicit matching line in the backward pass below.
That is the point of the demo: to expose the arithmetic that a deep-learning
framework would otherwise hide.
Blocks
------
[B-corpus] the excerpted sentences of the Universal Declaration of Human
Rights, stored verbatim so the demo needs no network
[B-tokens] tokenization into lowercase word tokens, and the vocabulary
[B-model] the learned parameters: embeddings E and the projections
W_Q, W_K, W_V, W_O, plus the Adam state
[B-forward] one sentence forward and backward by hand, every backward line
matching a forward one
[B-train] 4000 full-corpus iterations; the loss falls well below its
starting value
[B-display] the attention matrix of the display sentence, the same sentence
the entry's scaled-dot-product figure uses
[B-csv] the two CSVs the entry's pgfplots figures read
[B-preview] the matplotlib preview: training curve and attention heat map
Outputs
-------
pythondemos/attention_loss.csv — columns: iter, loss (training curve)
pythondemos/attention_weights.csv — columns: q, k, w (learned attention
matrix of one display sentence; q = query position, k = key position)
pythondemos/attention.png — matplotlib preview (loss + heat map)
The display sentence and its token order are printed and asserted so the tick
labels hard-coded in the `attention` glossary entry stay in sync.
"""
import re
import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path
OUT_DIR = Path(__file__).parent
# Fix the RNG so the learned embeddings, projections, and hence the heat map
# committed to the repo are reproducible bit-for-bit on every run.
np.random.seed(0)
the excerpted sentences of the Universal Declaration of Human Rights, stored verbatim so the demo needs no network
# One sentence/clause per line (Articles 1, 3, 4, 5, 6, 7, 9, 13, 17,
# 18, 19, 24, 26). Short, thematically overlapping sentences share many
# tokens ("everyone", "right", "no one shall", "freedom"), which gives the
# single head recurring associations to latch onto during training.
CORPUS = """
All human beings are born free and equal in dignity and rights.
They are endowed with reason and conscience and should act towards one another in a spirit of brotherhood.
Everyone has the right to life, liberty and security of person.
No one shall be held in slavery or servitude.
Slavery and the slave trade shall be prohibited in all their forms.
No one shall be subjected to torture or to cruel, inhuman or degrading treatment or punishment.
Everyone has the right to recognition everywhere as a person before the law.
All are equal before the law and are entitled without any discrimination to equal protection of the law.
No one shall be subjected to arbitrary arrest, detention or exile.
Everyone has the right to freedom of movement and residence within the borders of each state.
Everyone has the right to own property alone as well as in association with others.
No one shall be arbitrarily deprived of his property.
Everyone has the right to freedom of thought, conscience and religion.
Everyone has the right to freedom of opinion and expression.
Everyone has the right to rest and leisure, including reasonable limitation of working hours and periodic holidays with pay.
Everyone has the right to education.
"""
tokenization into lowercase word tokens, and the vocabulary
# Keep only runs of letters; this drops punctuation and case so that, e.g.,
# "rights." and "rights" map to the same token.
def tokenize(sentence):
return re.findall(r"[a-z]+", sentence.lower())
sentences = [tokenize(s) for s in CORPUS.strip().split("\n")]
# A self-attention head needs at least two tokens (a token attends to the
# *others*), so discard any degenerate one-token line.
sentences = [s for s in sentences if len(s) >= 2]
# The distinct tokens are collected and sorted; stoi maps each token to
# its integer id (its row in the embedding matrix E and its column in the
# read-out scores). Sorting makes the id assignment deterministic.
vocab = sorted({tok for s in sentences for tok in s})
stoi = {w: i for i, w in enumerate(vocab)}
V = len(vocab)
print(f"[B-tokens] {len(sentences)} sentences, {V} distinct tokens")
[B-tokens] 16 sentences, 98 distinct tokens
the learned parameters: embeddings E and the projections W_Q, W_K, W_V, W_O, plus the Adam state
# The query, key, and value vectors all have the embedding width here for simplicity; in a real
# transformer d_k and d_v are typically smaller than d (per head).
d = 24 # embedding dimension
d_k = 24 # query/key dimension
d_v = 24 # value dimension
def randn(*shape):
# small random init, scaled by fan-in for stable gradients
return 0.1 * np.random.randn(*shape)
# The learnable parameters. E is the input embedding table; W_Q, W_K, W_V are
# the projection matrices that turn an embedding into a query, key, and value;
# W_O reads the attention output back out to a score for every distinct token.
E = randn(V, d) # token embeddings (learned)
W_Q = randn(d, d_k)
W_K = randn(d, d_k)
W_V = randn(d, d_v)
W_O = randn(d_v, V) # read-out to one score per distinct token
params = {"E": E, "W_Q": W_Q, "W_K": W_K, "W_V": W_V, "W_O": W_O}
# Adam optimizer state: a first-moment (m) and second-moment (v_) running
# average per parameter tensor. Adam adapts the step size per coordinate,
# which makes this tiny hand-written model train in a few thousand iterations.
m = {k: np.zeros_like(v) for k, v in params.items()}
v_ = {k: np.zeros_like(v) for k, v in params.items()}
beta1, beta2, eps, lr = 0.9, 0.999, 1e-8, 0.02
one sentence forward and backward by hand, every backward line matching a forward one
def softmax_rows(S):
# Numerically stable row-wise softmax: subtracting the row max before
# exponentiating avoids overflow and leaves the result unchanged.
S = S - S.max(axis=1, keepdims=True)
P = np.exp(S)
return P / P.sum(axis=1, keepdims=True)
def forward_backward(idx, grads):
"""One sentence: accumulate gradients into `grads`, return (loss, n_pred).
`idx` is the list of token ids for one sentence. The forward pass computes
the attention output and the prediction loss; the backward pass then walks
the same operations in reverse, accumulating the gradient of the loss w.r.t.
every parameter into the shared `grads` dictionary (summed over sentences).
"""
n = len(idx)
# ── forward ──
X = E[idx] # (n, d) embeddings of this sentence
Q = X @ W_Q # (n, d_k) queries: what each token asks for
K = X @ W_K # (n, d_k) keys: what each token advertises
Vv = X @ W_V # (n, d_v) values: what each token contributes
S = (Q @ K.T) / np.sqrt(d_k) # (n, n) scaled query-key match scores
mask = np.eye(n, dtype=bool) # forbid self-attention
# Setting the diagonal to a large negative number makes softmax assign it
# ~zero weight: token i must reconstruct itself from the *other* tokens,
# turning the task into associative recall rather than trivial copying.
S = np.where(mask, -1e9, S)
A = softmax_rows(S) # (n, n), rows sum to 1: retrieval weights
Z = A @ Vv # (n, d_v) retrieved content per query token
logits = Z @ W_O # (n, V) one score per distinct token
P = softmax_rows(logits) # predicted token distribution
# loss: negative log probability of the (held-out) token at each position
# targets[i] is the true token id at position i; the loss rewards putting
# probability mass on it. The +1e-12 guards log(0).
targets = np.array(idx)
loss = -np.log(P[np.arange(n), targets] + 1e-12).sum()
# ── backward ──
# Each block below is the derivative of the matching forward line, applied
# in reverse order (chain rule). Shapes are annotated to make the matrix
# multiplications self-checking.
# d(loss)/d(scores) for softmax + negative log probability is simply P - onehot(target).
dlogits = P.copy()
dlogits[np.arange(n), targets] -= 1.0 # (n, V)
# logits = Z @ W_O -> gradients w.r.t. W_O and Z.
grads["W_O"] += Z.T @ dlogits
dZ = dlogits @ W_O.T # (n, d_v)
# Z = A @ Vv -> split the gradient between the attention weights A and
# the values Vv.
dA = dZ @ Vv.T # (n, n)
dVv = A.T @ dZ # (n, d_v)
# softmax backward per row: Jacobian of a row-softmax applied to dA.
dS = A * (dA - (dA * A).sum(axis=1, keepdims=True))
# No gradient flows through the masked diagonal (its score was a constant).
dS = np.where(mask, 0.0, dS)
# S = (Q @ K.T)/sqrt(d_k) -> gradients w.r.t. the queries and keys.
dQ = (dS @ K) / np.sqrt(d_k)
dK = (dS.T @ Q) / np.sqrt(d_k)
# Q = X @ W_Q, K = X @ W_K, Vv = X @ W_V -> projection-matrix gradients.
grads["W_Q"] += X.T @ dQ
grads["W_K"] += X.T @ dK
grads["W_V"] += X.T @ dVv
# X feeds all three projections, so its gradient is the sum of the three
# paths back through W_Q, W_K, W_V.
dX = dQ @ W_Q.T + dK @ W_K.T + dVv @ W_V.T # (n, d)
# Scatter-add into the embedding table: a token may occur several times in
# one sentence (e.g. "and", "for", "the"), so all of its positions
# accumulate into the single shared embedding row. Plain E[idx] += dX would
# drop the duplicates; np.add.at sums them correctly.
np.add.at(grads["E"], idx, dX)
return loss, n
# Pre-convert every sentence to its list of token ids once.
sent_idx = [[stoi[w] for w in s] for s in sentences]
4000 full-corpus iterations; the loss falls well below its starting value
# The corpus is tiny, so each iteration sums the gradient
# over all sentences, then take a single Adam update.
EPOCHS = 4000
loss_curve = []
for epoch in range(1, EPOCHS + 1):
grads = {k: np.zeros_like(v) for k, v in params.items()}
total_loss, total_pred = 0.0, 0
for idx in sent_idx:
l, n = forward_backward(idx, grads)
total_loss += l
total_pred += n
# Report the loss per predicted token so the curve is comparable across
# sentences of different lengths.
mean_loss = total_loss / total_pred
# Adam update: bias-corrected first/second moments give a per-coordinate
# adaptive step. The gradient is averaged over sentences (÷ len) so the
# step size does not scale with corpus size.
for k in params:
g = grads[k] / len(sent_idx)
m[k] = beta1 * m[k] + (1 - beta1) * g
v_[k] = beta2 * v_[k] + (1 - beta2) * (g * g)
mhat = m[k] / (1 - beta1 ** epoch)
vhat = v_[k] / (1 - beta2 ** epoch)
params[k] -= lr * mhat / (np.sqrt(vhat) + eps)
# Subsample the curve (every 20th iteration) to keep the committed CSV small.
if epoch == 1 or epoch % 20 == 0:
loss_curve.append((epoch, mean_loss))
print(f"[B-train] initial loss = {loss_curve[0][1]:.4f}, final loss = {loss_curve[-1][1]:.4f}")
# Guard against a silently broken gradient: a correct implementation drives the
# loss well below its starting value.
assert loss_curve[-1][1] < loss_curve[0][1] - 0.3, "training did not reduce the loss"
[B-train] initial loss = 4.5846, final loss = 0.0000
the attention matrix of the display sentence, the same sentence the entry's scaled-dot-product figure uses
# First sentence of Article 1 — the same sentence used in the
# scaled-dot-product figure of the `attention` glossary entry.
DISPLAY = ["all", "human", "beings", "are", "born", "free",
"and", "equal", "in", "dignity", "and", "rights"]
# This sentence must appear verbatim (as a token list) in the corpus, so the
# heat map shows weights the head actually trained on (not an out-of-sample one).
assert DISPLAY in sentences, "display sentence not found in the corpus"
print("[B-display] display tokens:", " ".join(DISPLAY))
# Recompute the forward attention weights for the display sentence only (no
# gradient needed here), reusing the trained embeddings and projections.
d_idx = [stoi[w] for w in DISPLAY]
X = E[d_idx]
S = (X @ W_Q) @ (X @ W_K).T / np.sqrt(d_k)
S = np.where(np.eye(len(d_idx), dtype=bool), -1e9, S)
A = softmax_rows(S)
[B-display] display tokens: all human beings are born free and equal in dignity and rights
the two CSVs the entry's pgfplots figures read
# Training curve: consumed by the loss plot below (and available to the entry).
np.savetxt(OUT_DIR / "attention_loss.csv", np.array(loss_curve),
delimiter=",", header="iter,loss", comments="", fmt="%.6f")
print(f"[B-csv] saved {OUT_DIR / 'attention_loss.csv'}")
# Attention matrix in long/tidy form (one row per (query, key) cell) so that
# pgfplots can read it directly with `matrix plot*` in the glossary figure.
rows = []
for q in range(A.shape[0]):
for k in range(A.shape[1]):
rows.append((q, k, A[q, k]))
np.savetxt(OUT_DIR / "attention_weights.csv", np.array(rows),
delimiter=",", header="q,k,w", comments="", fmt="%.6f")
print(f"[B-csv] saved {OUT_DIR / 'attention_weights.csv'}")
[B-csv] saved /Users/junga1/dictionaryappliedml/pythondemos/attention_loss.csv [B-csv] saved /Users/junga1/dictionaryappliedml/pythondemos/attention_weights.csv
the matplotlib preview: training curve and attention heat map
# A local sanity-check PDF (not committed as the paper figure): the training
# curve on the left, the learned attention heat map on the right.
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(8.5, 3.6))
ep, ls = zip(*loss_curve)
ax1.plot(ep, ls, "-", color="tab:blue")
ax1.set_xlabel("iteration")
ax1.set_ylabel("loss per predicted token")
ax1.set_title("training loss")
# Rows = query token (searching), columns = key token (searched); a dark cell
# means the query in that row retrieves strongly from the key in that column.
im = ax2.imshow(A, cmap="Blues", vmin=0.0)
ax2.set_xticks(range(len(DISPLAY)))
ax2.set_yticks(range(len(DISPLAY)))
ax2.set_xticklabels(DISPLAY, rotation=90, fontsize=7)
ax2.set_yticklabels(DISPLAY, fontsize=7)
ax2.set_xlabel("key token (searched)")
ax2.set_ylabel("query token (searching)")
ax2.set_title("learned attention weights")
fig.colorbar(im, ax=ax2, fraction=0.046)
fig.tight_layout()
out = OUT_DIR / "attention.png"
fig.savefig(out, bbox_inches="tight", dpi=110)
print(f"[B-preview] saved {out}")
plt.close()
[B-preview] saved /Users/junga1/dictionaryappliedml/pythondemos/attention.png
B-preview writes when the script runs