P-alt writes when the script runsDictionary of Applied Machine Learning · transformer
Numerical companion to the entry transformer: it recomputes what the entry states and prints one line per check
One block per paragraph of the entry (marked [P...]): each block verifies numerically what the corresponding statement asserts. Self-contained (numpy/matplotlib only), fixed seed.
Run it with python3 transformer.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 transformer.py
One cell per block of the script: the code, and what that code printed when it last ran here
"""
transformer.py — numerical companion to the glossary entry 'transformer'.
One block per paragraph of the entry (marked [P...]): each block verifies
numerically what the corresponding statement asserts. Self-contained
(numpy/matplotlib only), fixed seed.
Blocks
------
[P-shape] A transformer composes layers that map the matrix X in
R^{n x d} of stacked token feature vectors to one of the SAME
shape — which is exactly what lets layers compose freely: a
stack of 6 alternating layers applies without any shape
bookkeeping.
[P-block] The two layer types differ in information flow: a
token-mixing layer (softmax self-attention) lets a
perturbation of token j change the output rows of OTHER
tokens, while the position-wise MLP is row-local — perturbing
token j leaves every other output row untouched. Residual
connections preserve the shape as well.
[P-alt] Token mixing need not be attention: a fixed Fourier mixing
layer (FNet-style, real part of the FFT across the token
axis) is a same-shape, cross-token-propagating replacement
that computes NO pairwise query-key scores, whereas the
attention layer forms an n x n score matrix (n^2 pairwise
comparisons) by construction.
Outputs
-------
transformer.png : preview figure (checking only).
Data generated by pythondemos/transformer.py.
"""
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
rng = np.random.default_rng(42)
report = []
def check(name, ok):
report.append((name, bool(ok)))
print(f" [{'ok' if ok else 'FAIL'}] {name}")
n, dm = 8, 16 # tokens x features
softmax = lambda S: np.exp(S - S.max(1, keepdims=True)) / \
np.exp(S - S.max(1, keepdims=True)).sum(1, keepdims=True)
def attention_layer(X, Wq, Wk, Wv):
S = (X @ Wq) @ (X @ Wk).T / np.sqrt(dm) # n x n scores
return softmax(S) @ (X @ Wv), S
def mlp_layer(X, W1, W2): # position-wise
return np.maximum(X @ W1, 0) @ W2
A transformer composes layers that map the matrix X in R^{n x d} of stacked token feature vectors to one of the SAME shape — which is exactly what lets layers compose freely: a stack of 6 alternating layers applies without any shape bookkeeping.
print("[P-shape] same-shape layers compose freely")
params = [(rng.normal(size=(dm, dm)) / np.sqrt(dm),
rng.normal(size=(dm, dm)) / np.sqrt(dm),
rng.normal(size=(dm, dm)) / np.sqrt(dm),
rng.normal(size=(dm, dm)) / np.sqrt(dm),
rng.normal(size=(dm, dm)) / np.sqrt(dm)) for _ in range(3)]
X = rng.normal(size=(n, dm))
Z = X.copy()
for Wq, Wk, Wv, W1, W2 in params: # 3 blocks = 6 layers
A, _ = attention_layer(Z, Wq, Wk, Wv)
Z = Z + A # residual
Z = Z + mlp_layer(Z, W1, W2) # residual
check("layer output keeps the n x d shape", Z.shape == (n, dm))
check("the 6-layer composition ran without shape bookkeeping",
Z.shape == X.shape)
[P-shape] same-shape layers compose freely [ok] layer output keeps the n x d shape [ok] layer output keeps the n x d shape [ok] layer output keeps the n x d shape [ok] the 6-layer composition ran without shape bookkeeping
The two layer types differ in information flow: a token-mixing layer (softmax self-attention) lets a perturbation of token j change the output rows of OTHER tokens, while the position-wise MLP is row-local — perturbing token j leaves every other output row untouched. Residual connections preserve the shape as well.
print("[P-block] token mixing vs position-wise transformation")
Wq, Wk, Wv, W1, W2 = params[0]
Xp = X.copy()
Xp[3] += 0.5 * rng.normal(size=dm) # perturb token 3
A0, S0 = attention_layer(X, Wq, Wk, Wv)
A1, _ = attention_layer(Xp, Wq, Wk, Wv)
other = [i for i in range(n) if i != 3]
check("token mixing: perturbing token 3 changes OTHER tokens' outputs",
np.max(np.abs(A1[other] - A0[other])) > 1e-6)
M0, M1 = mlp_layer(X, W1, W2), mlp_layer(Xp, W1, W2)
check("position-wise MLP: other rows are untouched (row-local)",
np.allclose(M1[other], M0[other]))
check("position-wise MLP: only row 3 changed",
not np.allclose(M1[3], M0[3]))
[P-block] token mixing vs position-wise transformation [ok] token mixing: perturbing token 3 changes OTHER tokens' outputs [ok] position-wise MLP: other rows are untouched (row-local) [ok] position-wise MLP: only row 3 changed
Token mixing need not be attention: a fixed Fourier mixing layer (FNet-style, real part of the FFT across the token axis) is a same-shape, cross-token-propagating replacement that computes NO pairwise query-key scores, whereas the attention layer forms an n x n score matrix (n^2 pairwise comparisons) by construction.
print("[P-alt] Fourier mixing replaces attention without pairwise scores")
fourier_layer = lambda X: np.real(np.fft.fft(X, axis=0)) # mix tokens
F0, F1 = fourier_layer(X), fourier_layer(Xp)
check("Fourier layer keeps the n x d shape", F0.shape == (n, dm))
check("Fourier layer propagates information across tokens",
np.max(np.abs(F1[other] - F0[other])) > 1e-6)
check("attention forms an n x n score matrix (n^2 = 64 pairwise "
"comparisons) by construction; Fourier mixing forms none",
S0.shape == (n, n))
stack = fourier_layer(X) + mlp_layer(fourier_layer(X), W1, W2)
check("the alternative mixing layer composes in the same block "
"structure", stack.shape == (n, dm))
# ------------------------------------------------------------ preview
fig, ax = plt.subplots(1, 2, figsize=(7.6, 3.0))
ax[0].imshow(softmax(S0), cmap="gray")
ax[0].set_title("[P-block] attention weights (n x n)")
ax[1].imshow(np.abs(A1 - A0) > 1e-9, cmap="gray", aspect="auto")
ax[1].set_title("rows changed by perturbing token 3")
fig.tight_layout()
fig.savefig("transformer.png", dpi=110)
print(f"\n{sum(ok for _, ok in report)}/{len(report)} checks passed")
assert all(ok for _, ok in report)
[P-alt] Fourier mixing replaces attention without pairwise scores [ok] Fourier layer keeps the n x d shape [ok] Fourier layer propagates information across tokens [ok] attention forms an n x n score matrix (n^2 = 64 pairwise comparisons) by construction; Fourier mixing forms none [ok] the alternative mixing layer composes in the same block structure 11/11 checks passed
P-alt writes when the script runs