The lecture's three-key example with every number, the scaling, and the causal mask.
The notebook carries the source, the install step and every run below, so nothing has to be on your machine. Open it, then choose Runtime › Run all.
Colab opens it read only. Click Copy to Drive to keep
your changes. You can also download
attention.ipynb, or run it
locally: python3 attention.py.
python3 attention.py
All five sections.
python3 attention.py --one
One query worked end to end, every number the chapter prints.
python3 attention.py --scale
Without the root d_k, entropy over 8 keys collapses from 1.78 bits at d = 4 to 0.13 bits at d = 1024, before any training.
python3 attention.py --mask
Mask then softmax and every row sums to 1. Mask after and row one sums to 0.4551, with the blocked mass thrown away instead of redistributed.
Download attention.py
· served verbatim at https://nlp.jcrlabz.com/code/attention.py
"""Self-attention, worked as Chapter 16 works it.
Five things, all reproducing the chapter's tables:
one query the lecture's three-key example, every number
when it selects a query that actually picks a key, and one that cannot
why root d_k what the scaling stops the softmax from doing
the causal mask mask then softmax, and why never the other way round
the parameter count where a transformer layer's numbers sit
python3 attention.py # all five
python3 attention.py --one # the lecture example
python3 attention.py --select # attention that discriminates
python3 attention.py --scale # the square root, tested
python3 attention.py --mask # causal masking
python3 attention.py --params # one layer, counted
Install: pip install numpy
"""
import argparse
import math
try:
import numpy as np
except ImportError:
raise SystemExit("this one needs numpy: pip install numpy")
def softmax(z):
e = np.exp(z - np.max(z, axis=-1, keepdims=True))
return e / e.sum(axis=-1, keepdims=True)
def entropy(p):
"""In bits. Maximum is log2(n) for a uniform distribution."""
p = p[p > 0]
return float(-(p * np.log2(p)).sum())
# ---------------------------------------------------------------- one query
def show_one():
q = np.array([0.1, 0.2, 0.3])
keys = {"is": [0.4, 0.5, 0.6], "the": [0.2, 0.1, 0.3],
"best": [0.6, 0.7, 0.8]}
K = np.array(list(keys.values()))
V = K.copy() # values equal keys, as the lecture has it
d = q.shape[0]
print("\nSTEP 1 one query against three keys\n")
print(f" query 'who' q = {q}")
for w, k in keys.items():
label = f"key '{w}'"
print(f" {label:<13} k = {np.array(k)}")
print(f"\n Values equal keys here, to keep the arithmetic short.")
print(f" d_k = {d}, so the scale is sqrt({d}) = {math.sqrt(d):.4f}\n")
raw = K @ q
scaled = raw / math.sqrt(d)
ex = np.exp(scaled)
alpha = ex / ex.sum()
print(f" {'key':<7}{'q . k':>9}{'/ sqrt(d)':>12}{'exp':>9}{'alpha':>9}")
print(" " + "-" * 48)
for w, r, s, e, a in zip(keys, raw, scaled, ex, alpha):
print(f" {w:<7}{r:>9.4f}{s:>12.4f}{e:>9.4f}{a:>9.4f}")
print(" " + "-" * 48)
print(f" {'sum':<7}{'':>9}{'':>12}{ex.sum():>9.4f}{alpha.sum():>9.4f}")
z = alpha @ V
print(f"\n output zeta = sum of alpha_i v_i")
for w, a, v in zip(keys, alpha, V):
print(f" {a:.4f} x {v}")
print(f" = ({z[0]:.4f}, {z[1]:.4f}, {z[2]:.4f})")
print(f"\n Now look at the alpha column: {alpha[0]:.2f},"
f" {alpha[1]:.2f}, {alpha[2]:.2f}.")
print(f" That is almost uniform. Its entropy is {entropy(alpha):.4f} bits")
print(f" against a maximum of {math.log2(3):.4f}.")
print(f"\n So this query attends to everything about equally, and the")
print(f" output is close to the plain average of the values. The")
print(f" mechanism ran correctly and selected nothing, because these")
print(f" three keys all point the same way. Step 2 fixes that.")
# ------------------------------------------------------------ when it selects
def show_select():
print("\n\nSTEP 2 attention that actually discriminates\n")
print(" Same machinery, keys that disagree. Three dimensions standing")
print(" for (animal, vehicle, verb).\n")
keys = {"dog": [0.9, 0.0, 0.1], "truck": [0.0, 0.9, 0.1],
"barked": [0.1, 0.0, 0.9]}
K = np.array(list(keys.values()))
V = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=float)
d = 3
for name, q in [("an animal query", np.array([5.0, 0.0, 0.0])),
("a verb query", np.array([0.0, 0.0, 5.0])),
("an undecided query", np.array([2.0, 2.0, 2.0]))]:
alpha = softmax(K @ q / math.sqrt(d))
z = alpha @ V
cells = " ".join(f"{w} {a:.4f}" for w, a in zip(keys, alpha))
print(f" {name:<20}q = {q}")
print(f" alpha {cells}")
print(f" output ({z[0]:.4f}, {z[1]:.4f}, {z[2]:.4f})"
f" entropy {entropy(alpha):.4f} bits\n")
print(f" The first two queries commit. The third cannot, and its")
print(f" entropy sits at the {math.log2(3):.4f} bit maximum for three keys.")
print("\n Low entropy means the head has made a choice. High entropy")
print(" means it is averaging. Both are legitimate, and reading which")
print(" one a head is doing is how attention maps get interpreted.")
# ------------------------------------------------------------- the scaling
def show_scale():
print("\n\nSTEP 3 why divide by the square root of d_k\n")
print(" Take random query and key vectors with unit variance entries.")
print(" Their dot product is a sum of d terms, so its standard")
print(" deviation grows like sqrt(d). Left alone, the scores get")
print(" larger and larger as the model gets wider.\n")
rng = np.random.default_rng(0)
n = 8
print(f" {'d_k':>6}{'sd of q.k':>12}{'entropy raw':>14}"
f"{'entropy scaled':>16}{'max alpha raw':>15}")
print(" " + "-" * 66)
for d in (4, 16, 64, 256, 1024):
sd, e_raw, e_scaled, top = [], [], [], []
for _ in range(200):
q = rng.normal(0, 1, d)
K = rng.normal(0, 1, (n, d))
s = K @ q
sd.append(s.std())
a_raw = softmax(s)
a_sc = softmax(s / math.sqrt(d))
e_raw.append(entropy(a_raw))
e_scaled.append(entropy(a_sc))
top.append(a_raw.max())
print(f" {d:>6}{np.mean(sd):>12.3f}{np.mean(e_raw):>14.4f}"
f"{np.mean(e_scaled):>16.4f}{np.mean(top):>15.4f}")
print(f"\n Maximum possible entropy over {n} keys is"
f" {math.log2(n):.4f} bits.")
print("\n Read the raw column. By d = 1024 the distribution has")
print(" collapsed onto a single key, before any training has happened.")
print(" A softmax that saturated is a softmax with no gradient, so the")
print(" head cannot learn its way out.")
print("\n The scaled column barely moves. Dividing by sqrt(d_k) cancels")
print(" exactly the growth the second column shows, which is why the")
print(" scale is a square root and not something tuned.")
# ------------------------------------------------------------- causal mask
def show_mask():
print("\n\nSTEP 4 the causal mask, and the order of operations\n")
print(" A language model must not read the future. When scoring")
print(" position i, only positions j <= i may contribute.\n")
print(" Set the blocked scores to minus infinity BEFORE the softmax:\n")
print(" S_ij = (Q K^T)_ij / sqrt(d_k) if j <= i, else -inf\n")
S = np.array([[2.0, 1.0, 0.5, 1.5],
[0.5, 2.0, 1.0, 0.5],
[1.0, 0.5, 2.0, 1.0],
[1.5, 1.0, 0.5, 2.0]])
n = len(S)
tokens = ["the", "dog", "barked", "loudly"]
print(f" raw scores S (rows are queries, columns are keys):\n")
print(" " + " " * 10 + "".join(f"{t:>10}" for t in tokens))
for i, t in enumerate(tokens):
print(f" {t:<10}" + "".join(f"{v:>10.2f}" for v in S[i]))
mask = np.tril(np.ones((n, n), dtype=bool))
Sm = np.where(mask, S, -np.inf)
A = softmax(Sm)
print(f"\n after masking then softmax:\n")
print(" " + " " * 10 + "".join(f"{t:>10}" for t in tokens) + f"{'sum':>8}")
for i, t in enumerate(tokens):
cells = "".join(" ." if not mask[i][j] else f"{A[i][j]:>10.4f}"
for j in range(n))
print(f" {t:<10}{cells}{A[i].sum():>8.4f}")
print("\n Every row still sums to 1. That is the whole reason the mask")
print(" goes before the softmax rather than after.\n")
wrong = softmax(S) * mask
print(" Compare softmax first, then zero the blocked cells:\n")
print(" " + " " * 10 + "".join(f"{t:>10}" for t in tokens) + f"{'sum':>8}")
for i, t in enumerate(tokens):
cells = "".join(" ." if not mask[i][j] else f"{wrong[i][j]:>10.4f}"
for j in range(n))
print(f" {t:<10}{cells}{wrong[i].sum():>8.4f}")
print(f"\n Row 1 now sums to {wrong[0].sum():.4f} instead of 1. The mass")
print(" that belonged to the blocked positions has simply been thrown")
print(" away rather than redistributed.")
print("\n Masking first lets the surviving positions inherit that mass.")
print(" Masking second silently scales the whole row down, and the")
print(" earlier the token, the worse it gets.")
# --------------------------------------------------------- the parameter count
def show_params(V=50000, D=512, L=6, h=8, H=2048):
d_k = d_v = D // h
print(f"\n\nSTEP 5 a transformer layer, counted\n")
print(f" |V| = {V}, D = {D}, layers L = {L}, heads h = {h},")
print(f" d_k = d_v = D/h = {d_k}, feed-forward H = {H}\n")
attn = 4 * D * D # W^Q, W^K, W^V, W^O
ffn = 2 * D * H + H + D
ln = 4 * D # two layer norms, scale and shift
layer = attn + ffn + ln
emb = V * D
print(f" {'component':<26}{'formula':<22}{'parameters':>14}")
print(" " + "-" * 64)
print(f" {'attention Q,K,V,O':<26}{'4 D^2':<22}{attn:>14,}")
print(f" {'feed-forward':<26}{'2 D H + H + D':<22}{ffn:>14,}")
print(f" {'layer norms':<26}{'4 D':<22}{ln:>14,}")
print(" " + "-" * 64)
print(f" {'one layer':<26}{'':<22}{layer:>14,}")
print(f" {'all ' + str(L) + ' layers':<26}{'L x layer':<22}"
f"{L * layer:>14,}")
print(f" {'embeddings':<26}{'V D':<22}{emb:>14,}")
print(" " + "-" * 64)
print(f" {'total':<48}{L * layer + emb:>14,}")
print(f"\n Two things to notice.")
print(f"\n The heads are free. Splitting D into {h} heads of {d_k} costs")
print(f" nothing, because h x d_k = D. Multi-head attention is a")
print(f" reshape, not an extra parameter budget.")
print(f"\n The feed-forward block is larger than the attention block,")
print(f" {ffn:,} against {attn:,}. Attention gets the attention, and")
print(f" most of the parameters sit next door.")
print(f"\n\n WHAT THE SEQUENCE LENGTH COSTS\n")
print(f" No parameter count above mentions sequence length. The")
print(f" computation does: every position attends to every position.\n")
print(f" {'tokens n':>10}{'n^2 scores per head':>22}{'vs n = 512':>13}")
print(" " + "-" * 46)
for nt in (128, 512, 2048, 8192, 32768):
print(f" {nt:>10}{nt * nt:>22,}{nt * nt / 512 ** 2:>12.2f}x")
print(f"\n Quadruple the context and the attention cost goes up")
print(f" sixteen-fold. That is the wall the efficiency chapter climbs.")
def main():
ap = argparse.ArgumentParser(description=__doc__)
for f in ("one", "select", "scale", "mask", "params"):
ap.add_argument(f"--{f}", action="store_true")
a = ap.parse_args()
picked = a.one or a.select or a.scale or a.mask or a.params
if a.one or not picked:
show_one()
if a.select or not picked:
show_select()
if a.scale or not picked:
show_scale()
if a.mask or not picked:
show_mask()
if a.params or not picked:
show_params()
print()
if __name__ == "__main__":
main()